experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Update profiler API to allow user using it in graph mode.

Browse Source 
PiperOrigin-RevId: 233492307
This commit is contained in:
A. Unique TensorFlower 2019-02-11 15:35:11 -08:00 committed by TensorFlower Gardener
parent d3e43986e2
commit 5a297f5efe
6 changed files with 56 additions and 36 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
tensorflow/c/eager/c_api_experimental.cc
View File

@ -25,7 +25,7 @@ void TFE_OpConsumeInput(TFE_Op* op, TFE_TensorHandle* h, TF_Status* status) {
op->operation.ConsumeInput(h->handle); op->operation.ConsumeInput(h->handle);
} }
TFE_Profiler* TFE_NewProfiler(TFE_Context* ctx) { TFE_Profiler* TFE_NewProfiler(TFE_ProfilerContext* ctx) {
return new TFE_Profiler(ctx); return new TFE_Profiler(ctx);
} }
@ -50,17 +50,21 @@ void TFE_ProfilerSerializeToString(TFE_Context* ctx, TFE_Profiler* profiler,
}; };
} }
TFE_ProfilerServerOptions* TFE_NewProfilerServerOptions() { TFE_ProfilerContext* TFE_NewProfilerContext() {
return new TFE_ProfilerServerOptions; return new TFE_ProfilerContext;
} }
void TFE_ProfilerServerOptionsSetEagerContext( void TFE_ProfilerContextSetEagerContext(TFE_ProfilerContext* profiler_context,
TFE_ProfilerServerOptions* options, TFE_Context* ctx) { TFE_Context* eager_context) {
options->profiler_context.eager_context = &ctx->context; profiler_context->profiler_context.eager_context = &eager_context->context;
} }
void TFE_StartProfilerServer(TFE_ProfilerServerOptions* options, int port) { void TFE_DeleteProfilerContext(TFE_ProfilerContext* profiler_context) {
delete profiler_context;
}
void TFE_StartProfilerServer(TFE_ProfilerContext* context, int port) {
// Release child thread intentionally. The child thread can be terminate by // Release child thread intentionally. The child thread can be terminate by
// terminating the main thread. // terminating the main thread.
tensorflow::StartProfilerServer(&options->profiler_context, port).release(); tensorflow::StartProfilerServer(&context->profiler_context, port).release();
} }

23
tensorflow/c/eager/c_api_experimental.h
View File

@ -25,6 +25,8 @@ extern "C" {
TF_CAPI_EXPORT extern void TFE_OpConsumeInput(TFE_Op* op, TFE_TensorHandle* h, TF_CAPI_EXPORT extern void TFE_OpConsumeInput(TFE_Op* op, TFE_TensorHandle* h,
TF_Status* status); TF_Status* status);
typedef struct TFE_ProfilerContext TFE_ProfilerContext;
// A profiler which will start profiling when creating the object and will stop // A profiler which will start profiling when creating the object and will stop
// when the object is destroyed. It will profile all operations run under the // when the object is destroyed. It will profile all operations run under the
// given TFE_Context. Multiple instance of it can be created, but at most one // given TFE_Context. Multiple instance of it can be created, but at most one
@ -32,7 +34,7 @@ TF_CAPI_EXPORT extern void TFE_OpConsumeInput(TFE_Op* op, TFE_TensorHandle* h,
// Thread-safety: TFE_Profiler is thread-safe. // Thread-safety: TFE_Profiler is thread-safe.
typedef struct TFE_Profiler TFE_Profiler; typedef struct TFE_Profiler TFE_Profiler;
TF_CAPI_EXPORT extern TFE_Profiler* TFE_NewProfiler(TFE_Context* ctx); TF_CAPI_EXPORT extern TFE_Profiler* TFE_NewProfiler(TFE_ProfilerContext* ctx);
TF_CAPI_EXPORT extern bool TFE_ProfilerIsOk(TFE_Profiler* profiler); TF_CAPI_EXPORT extern bool TFE_ProfilerIsOk(TFE_Profiler* profiler);
TF_CAPI_EXPORT extern void TFE_DeleteProfiler(TFE_Profiler* profiler); TF_CAPI_EXPORT extern void TFE_DeleteProfiler(TFE_Profiler* profiler);
@ -43,15 +45,16 @@ TF_CAPI_EXPORT extern void TFE_ProfilerSerializeToString(TFE_Context* ctx,
TF_Buffer* buf, TF_Buffer* buf,
TF_Status* status); TF_Status* status);
typedef struct TFE_ProfilerServerOptions TFE_ProfilerServerOptions; // Return a new profiler context object.
TF_CAPI_EXPORT extern TFE_ProfilerContext* TFE_NewProfilerContext(void);
// Return a new Profiler server options object.
TF_CAPI_EXPORT extern TFE_ProfilerServerOptions* TFE_NewProfilerServerOptions(
void);
// Set the eager context in TFE_ProfilerServerOptions // Set the eager context in TFE_ProfilerServerOptions
TF_CAPI_EXPORT extern void TFE_ProfilerServerOptionsSetEagerContext( TF_CAPI_EXPORT extern void TFE_ProfilerContextSetEagerContext(
TFE_ProfilerServerOptions* options, TFE_Context* ctx); TFE_ProfilerContext* profiler_context, TFE_Context* eager_context);
// Destroy a profiler context object.
TF_CAPI_EXPORT extern void TFE_DeleteProfilerContext(
TFE_ProfilerContext* profiler_context);
// Start a profiler grpc server which listens to specified port. It will start // Start a profiler grpc server which listens to specified port. It will start
// the server on its own thread. It can be shutdown by terminating tensorflow. // the server on its own thread. It can be shutdown by terminating tensorflow.
@ -61,8 +64,8 @@ TF_CAPI_EXPORT extern void TFE_ProfilerServerOptionsSetEagerContext(
// tensorflow/contrib/tpu/profiler/capture_tpu_profile to capture tracable // tensorflow/contrib/tpu/profiler/capture_tpu_profile to capture tracable
// file following // file following
// https://cloud.google.com/tpu/docs/cloud-tpu-tools#capture_trace. // https://cloud.google.com/tpu/docs/cloud-tpu-tools#capture_trace.
TF_CAPI_EXPORT extern void TFE_StartProfilerServer( TF_CAPI_EXPORT extern void TFE_StartProfilerServer(TFE_ProfilerContext* context,
TFE_ProfilerServerOptions* options, int port); int port);
#ifdef __cplusplus #ifdef __cplusplus
} /* end extern "C" */ } /* end extern "C" */

13
tensorflow/c/eager/c_api_experimental_test.cc
View File

@ -41,9 +41,12 @@ void ExecuteWithProfiling(bool async) {
TFE_ContextOptions* opts = TFE_NewContextOptions(); TFE_ContextOptions* opts = TFE_NewContextOptions();
TFE_ContextOptionsSetAsync(opts, static_cast<unsigned char>(async)); TFE_ContextOptionsSetAsync(opts, static_cast<unsigned char>(async));
TFE_Context* ctx = TFE_NewContext(opts, status); TFE_Context* ctx = TFE_NewContext(opts, status);
TFE_Profiler* profiler = TFE_NewProfiler(ctx); TFE_ProfilerContext* profiler_context = TFE_NewProfilerContext();
TFE_ProfilerContextSetEagerContext(profiler_context, ctx);
TFE_Profiler* profiler = TFE_NewProfiler(profiler_context);
CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
TFE_DeleteContextOptions(opts); TFE_DeleteContextOptions(opts);
TFE_DeleteProfilerContext(profiler_context);
TFE_TensorHandle* m = TestMatrixTensorHandle(); TFE_TensorHandle* m = TestMatrixTensorHandle();
TFE_Op* matmul = MatMulOp(ctx, m, m); TFE_Op* matmul = MatMulOp(ctx, m, m);
@ -108,14 +111,18 @@ TEST(CAPI, MultipleProfilerSession) {
CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
TFE_DeleteContextOptions(opts); TFE_DeleteContextOptions(opts);
TFE_Profiler* profiler1 = TFE_NewProfiler(ctx); TFE_ProfilerContext* profiler_context = TFE_NewProfilerContext();
TFE_ProfilerContextSetEagerContext(profiler_context, ctx);
TFE_Profiler* profiler1 = TFE_NewProfiler(profiler_context);
EXPECT_TRUE(TFE_ProfilerIsOk(profiler1)); EXPECT_TRUE(TFE_ProfilerIsOk(profiler1));
TFE_Profiler* profiler2 = TFE_NewProfiler(ctx); TFE_Profiler* profiler2 = TFE_NewProfiler(profiler_context);
EXPECT_FALSE(TFE_ProfilerIsOk(profiler2)); EXPECT_FALSE(TFE_ProfilerIsOk(profiler2));
TFE_DeleteProfiler(profiler1); TFE_DeleteProfiler(profiler1);
TFE_DeleteProfiler(profiler2); TFE_DeleteProfiler(profiler2);
TFE_DeleteProfilerContext(profiler_context);
} }
} // namespace } // namespace

14
tensorflow/c/eager/c_api_internal.h
View File

@ -107,20 +107,18 @@ struct TFE_Op {
tensorflow::EagerOperation operation; tensorflow::EagerOperation operation;
}; };
struct TFE_ProfilerContext {
tensorflow::ProfilerContext profiler_context;
};
struct TFE_Profiler { struct TFE_Profiler {
TFE_Profiler(TFE_Context* ctx) { TFE_Profiler(TFE_ProfilerContext* ctx) {
tensorflow::ProfilerContext profiler_context; profiler = tensorflow::ProfilerSession::Create(&ctx->profiler_context);
profiler_context.eager_context = &ctx->context;
profiler = tensorflow::ProfilerSession::Create(&profiler_context);
} }
std::unique_ptr<tensorflow::ProfilerSession> profiler; std::unique_ptr<tensorflow::ProfilerSession> profiler;
}; };
struct TFE_ProfilerServerOptions {
tensorflow::ProfilerContext profiler_context;
};
namespace tensorflow { namespace tensorflow {
// Set an AttrValue on the op. Doesn't handle the list types. // Set an AttrValue on the op. Doesn't handle the list types.
void SetOpAttrValueScalar(TFE_Context* ctx, TFE_Op* op, void SetOpAttrValueScalar(TFE_Context* ctx, TFE_Op* op,

17
tensorflow/python/eager/profiler.py
View File

@ -46,7 +46,13 @@ def start():
if _profiler is not None: if _profiler is not None:
raise AssertionError('Another profiler is running.') raise AssertionError('Another profiler is running.')
with _profiler_lock: with _profiler_lock:
_profiler = pywrap_tensorflow.TFE_NewProfiler(context.context()._handle) # pylint: disable=protected-access profiler_context = pywrap_tensorflow.TFE_NewProfilerContext()
if context.default_execution_mode == context.EAGER_MODE:
pywrap_tensorflow.TFE_ProfilerContextSetEagerContext(
profiler_context,
context.context()._handle) # pylint: disable=protected-access
_profiler = pywrap_tensorflow.TFE_NewProfiler(profiler_context)
pywrap_tensorflow.TFE_DeleteProfilerContext(profiler_context)
if not pywrap_tensorflow.TFE_ProfilerIsOk(_profiler): if not pywrap_tensorflow.TFE_ProfilerIsOk(_profiler):
logging.warning('Another profiler session is running which is probably ' logging.warning('Another profiler session is running which is probably '
'created by profiler server. Please avoid using profiler ' 'created by profiler server. Please avoid using profiler '
@ -93,12 +99,13 @@ def start_profiler_server(port):
Args: Args:
port: port profiler server listens to. port: port profiler server listens to.
""" """
opts = pywrap_tensorflow.TFE_NewProfilerServerOptions() profiler_context = pywrap_tensorflow.TFE_NewProfilerContext()
if context.default_execution_mode == context.EAGER_MODE: if context.default_execution_mode == context.EAGER_MODE:
pywrap_tensorflow.TFE_ProfilerServerOptionsSetEagerContext( pywrap_tensorflow.TFE_ProfilerContextSetEagerContext(
opts, profiler_context,
context.context()._handle) # pylint: disable=protected-access context.context()._handle) # pylint: disable=protected-access
pywrap_tensorflow.TFE_StartProfilerServer(opts, port) pywrap_tensorflow.TFE_StartProfilerServer(profiler_context, port)
pywrap_tensorflow.TFE_DeleteProfilerContext(profiler_context)
class Profiler(object): class Profiler(object):

5
tensorflow/python/pywrap_tfe.i
View File

@ -37,8 +37,9 @@ limitations under the License.
%rename("%s") TFE_ProfilerIsOk; %rename("%s") TFE_ProfilerIsOk;
%rename("%s") TFE_DeleteProfiler; %rename("%s") TFE_DeleteProfiler;
%rename("%s") TFE_ProfilerSerializeToString; %rename("%s") TFE_ProfilerSerializeToString;
%rename("%s") TFE_NewProfilerServerOptions; %rename("%s") TFE_NewProfilerContext;
%rename("%s") TFE_ProfilerServerOptionsSetEagerContext; %rename("%s") TFE_ProfilerContextSetEagerContext;
%rename("%s") TFE_DeleteProfilerContext;
%rename("%s") TFE_StartProfilerServer; %rename("%s") TFE_StartProfilerServer;
%rename("%s") TFE_OpNameGetAttrType; %rename("%s") TFE_OpNameGetAttrType;
%rename("%s") TFE_Py_InitEagerTensor; %rename("%s") TFE_Py_InitEagerTensor;