From 6f8ac2157c05d76ed75e6e8c0e93077d7d664457 Mon Sep 17 00:00:00 2001
From: Brennan Saeta <saeta@google.com>
Date: Wed, 7 Mar 2018 17:23:57 -0800
Subject: [PATCH] Add tracing annotations to RemoteCallOp's execution.

PiperOrigin-RevId: 188260984
---
 tensorflow/core/kernels/BUILD           |  1 +
 tensorflow/core/kernels/function_ops.cc | 30 ++++++++++++++++++------------
 2 files changed, 19 insertions(+), 12 deletions(-)

diff --git a/tensorflow/core/kernels/BUILD b/tensorflow/core/kernels/BUILD
index 52be90e..1e2a335 100644
--- a/tensorflow/core/kernels/BUILD
+++ b/tensorflow/core/kernels/BUILD
@@ -1951,6 +1951,7 @@ tf_kernel_library(
         "//tensorflow/core:core_cpu_internal",
         "//tensorflow/core:framework",
         "//tensorflow/core:lib",
+        "//tensorflow/core:lib_internal",
     ],
 )
 
diff --git a/tensorflow/core/kernels/function_ops.cc b/tensorflow/core/kernels/function_ops.cc
index e3c78d6..7c302e2 100644
--- a/tensorflow/core/kernels/function_ops.cc
+++ b/tensorflow/core/kernels/function_ops.cc
@@ -28,6 +28,7 @@ limitations under the License.
 #include "tensorflow/core/graph/gradients.h"
 #include "tensorflow/core/graph/graph_constructor.h"
 #include "tensorflow/core/platform/macros.h"
+#include "tensorflow/core/platform/tracing.h"
 #include "tensorflow/core/util/device_name_utils.h"
 
 namespace tensorflow {
@@ -317,6 +318,8 @@ class RemoteCallOp : public AsyncOpKernel {
       if (cached_entry != handle_cache_.end()) {
         handle = cached_entry->second;
       } else {
+        port::Tracing::TraceMe activity(strings::StrCat(
+            "RemoteCall: Instantiate: ", func_.name(), " on ", target_device));
         OP_REQUIRES_OK_ASYNC(
             ctx,
             lib->Instantiate(func_.name(), AttrSlice(&attr_values),
@@ -344,21 +347,24 @@ class RemoteCallOp : public AsyncOpKernel {
       args.push_back(argument);
     }
     auto* rets = new std::vector<Tensor>;
-    lib->Run(opts, handle, args, rets, [rets, done, ctx](const Status& status) {
-      if (!status.ok()) {
-        ctx->SetStatus(status);
-      } else {
-        for (size_t i = 0; i < rets->size(); ++i) {
-          ctx->set_output(i, (*rets)[i]);
-        }
-      }
-      delete rets;
-      done();
-    });
+    auto* trace = new port::Tracing::TraceMe(strings::StrCat(
+        "RemoteCall: Run: ", func_.name(), " on ", target_device));
+    lib->Run(opts, handle, args, rets,
+             [rets, trace, done, ctx](const Status& status) {
+               if (!status.ok()) {
+                 ctx->SetStatus(status);
+               } else {
+                 for (size_t i = 0; i < rets->size(); ++i) {
+                   ctx->set_output(i, (*rets)[i]);
+                 }
+               }
+               delete rets;
+               delete trace;
+               done();
+             });
   }
 
  private:
-  string target_;
   NameAttrList func_;
 
   mutex mu_;
-- 
2.7.4