diff --git a/tensorflow/c/eager/c_api_experimental.cc b/tensorflow/c/eager/c_api_experimental.cc index f46562a21df..06bbb4ac412 100644 --- a/tensorflow/c/eager/c_api_experimental.cc +++ b/tensorflow/c/eager/c_api_experimental.cc @@ -25,7 +25,7 @@ void TFE_OpConsumeInput(TFE_Op* op, TFE_TensorHandle* h, TF_Status* status) { op->operation.ConsumeInput(h->handle); } -TFE_Profiler* TFE_NewProfiler(TFE_Context* ctx) { +TFE_Profiler* TFE_NewProfiler(TFE_ProfilerContext* ctx) { return new TFE_Profiler(ctx); } @@ -50,17 +50,21 @@ void TFE_ProfilerSerializeToString(TFE_Context* ctx, TFE_Profiler* profiler, }; } -TFE_ProfilerServerOptions* TFE_NewProfilerServerOptions() { - return new TFE_ProfilerServerOptions; +TFE_ProfilerContext* TFE_NewProfilerContext() { + return new TFE_ProfilerContext; } -void TFE_ProfilerServerOptionsSetEagerContext( - TFE_ProfilerServerOptions* options, TFE_Context* ctx) { - options->profiler_context.eager_context = &ctx->context; +void TFE_ProfilerContextSetEagerContext(TFE_ProfilerContext* profiler_context, + TFE_Context* eager_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 // terminating the main thread. - tensorflow::StartProfilerServer(&options->profiler_context, port).release(); + tensorflow::StartProfilerServer(&context->profiler_context, port).release(); } diff --git a/tensorflow/c/eager/c_api_experimental.h b/tensorflow/c/eager/c_api_experimental.h index e6a3d737073..51a5fa0d816 100644 --- a/tensorflow/c/eager/c_api_experimental.h +++ b/tensorflow/c/eager/c_api_experimental.h @@ -25,6 +25,8 @@ extern "C" { TF_CAPI_EXPORT extern void TFE_OpConsumeInput(TFE_Op* op, TFE_TensorHandle* h, TF_Status* status); +typedef struct TFE_ProfilerContext TFE_ProfilerContext; + // 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 // 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. 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 void TFE_DeleteProfiler(TFE_Profiler* profiler); @@ -43,15 +45,16 @@ TF_CAPI_EXPORT extern void TFE_ProfilerSerializeToString(TFE_Context* ctx, TF_Buffer* buf, TF_Status* status); -typedef struct TFE_ProfilerServerOptions TFE_ProfilerServerOptions; - -// Return a new Profiler server options object. -TF_CAPI_EXPORT extern TFE_ProfilerServerOptions* TFE_NewProfilerServerOptions( - void); +// Return a new profiler context object. +TF_CAPI_EXPORT extern TFE_ProfilerContext* TFE_NewProfilerContext(void); // Set the eager context in TFE_ProfilerServerOptions -TF_CAPI_EXPORT extern void TFE_ProfilerServerOptionsSetEagerContext( - TFE_ProfilerServerOptions* options, TFE_Context* ctx); +TF_CAPI_EXPORT extern void TFE_ProfilerContextSetEagerContext( + 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 // 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 // file following // https://cloud.google.com/tpu/docs/cloud-tpu-tools#capture_trace. -TF_CAPI_EXPORT extern void TFE_StartProfilerServer( - TFE_ProfilerServerOptions* options, int port); +TF_CAPI_EXPORT extern void TFE_StartProfilerServer(TFE_ProfilerContext* context, + int port); #ifdef __cplusplus } /* end extern "C" */ diff --git a/tensorflow/c/eager/c_api_experimental_test.cc b/tensorflow/c/eager/c_api_experimental_test.cc index 1994d652245..7b84c5426fa 100644 --- a/tensorflow/c/eager/c_api_experimental_test.cc +++ b/tensorflow/c/eager/c_api_experimental_test.cc @@ -41,9 +41,12 @@ void ExecuteWithProfiling(bool async) { TFE_ContextOptions* opts = TFE_NewContextOptions(); TFE_ContextOptionsSetAsync(opts, static_cast(async)); 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); TFE_DeleteContextOptions(opts); + TFE_DeleteProfilerContext(profiler_context); TFE_TensorHandle* m = TestMatrixTensorHandle(); TFE_Op* matmul = MatMulOp(ctx, m, m); @@ -108,14 +111,18 @@ TEST(CAPI, MultipleProfilerSession) { CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status); 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)); - TFE_Profiler* profiler2 = TFE_NewProfiler(ctx); + TFE_Profiler* profiler2 = TFE_NewProfiler(profiler_context); EXPECT_FALSE(TFE_ProfilerIsOk(profiler2)); TFE_DeleteProfiler(profiler1); TFE_DeleteProfiler(profiler2); + TFE_DeleteProfilerContext(profiler_context); } } // namespace diff --git a/tensorflow/c/eager/c_api_internal.h b/tensorflow/c/eager/c_api_internal.h index b070e1cb1f9..a563e4b8f50 100644 --- a/tensorflow/c/eager/c_api_internal.h +++ b/tensorflow/c/eager/c_api_internal.h @@ -107,20 +107,18 @@ struct TFE_Op { tensorflow::EagerOperation operation; }; +struct TFE_ProfilerContext { + tensorflow::ProfilerContext profiler_context; +}; + struct TFE_Profiler { - TFE_Profiler(TFE_Context* ctx) { - tensorflow::ProfilerContext profiler_context; - profiler_context.eager_context = &ctx->context; - profiler = tensorflow::ProfilerSession::Create(&profiler_context); + TFE_Profiler(TFE_ProfilerContext* ctx) { + profiler = tensorflow::ProfilerSession::Create(&ctx->profiler_context); } std::unique_ptr profiler; }; -struct TFE_ProfilerServerOptions { - tensorflow::ProfilerContext profiler_context; -}; - namespace tensorflow { // Set an AttrValue on the op. Doesn't handle the list types. void SetOpAttrValueScalar(TFE_Context* ctx, TFE_Op* op, diff --git a/tensorflow/python/eager/profiler.py b/tensorflow/python/eager/profiler.py index 7ddd891cdcb..659c0cc6e69 100644 --- a/tensorflow/python/eager/profiler.py +++ b/tensorflow/python/eager/profiler.py @@ -46,7 +46,13 @@ def start(): if _profiler is not None: raise AssertionError('Another profiler is running.') 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): logging.warning('Another profiler session is running which is probably ' 'created by profiler server. Please avoid using profiler ' @@ -93,12 +99,13 @@ def start_profiler_server(port): Args: 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: - pywrap_tensorflow.TFE_ProfilerServerOptionsSetEagerContext( - opts, + pywrap_tensorflow.TFE_ProfilerContextSetEagerContext( + profiler_context, 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): diff --git a/tensorflow/python/pywrap_tfe.i b/tensorflow/python/pywrap_tfe.i index db6c14fb601..3bf775537d2 100755 --- a/tensorflow/python/pywrap_tfe.i +++ b/tensorflow/python/pywrap_tfe.i @@ -37,8 +37,9 @@ limitations under the License. %rename("%s") TFE_ProfilerIsOk; %rename("%s") TFE_DeleteProfiler; %rename("%s") TFE_ProfilerSerializeToString; -%rename("%s") TFE_NewProfilerServerOptions; -%rename("%s") TFE_ProfilerServerOptionsSetEagerContext; +%rename("%s") TFE_NewProfilerContext; +%rename("%s") TFE_ProfilerContextSetEagerContext; +%rename("%s") TFE_DeleteProfilerContext; %rename("%s") TFE_StartProfilerServer; %rename("%s") TFE_OpNameGetAttrType; %rename("%s") TFE_Py_InitEagerTensor;