From 2f63b65d95de106ba359c2ac456e8dc8597bff9f Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Fri, 15 Nov 2019 16:29:54 -0800
Subject: [PATCH] Add TraceMe to track which step_id is from tf data.

PiperOrigin-RevId: 280760041
Change-Id: I76e41c26e21088c831d24a14201ac1adfb46b268
---
 tensorflow/core/kernels/data/BUILD            |  1 +
 .../core/kernels/data/captured_function.cc    | 26 +++++++++++++++++++
 2 files changed, 27 insertions(+)

diff --git a/tensorflow/core/kernels/data/BUILD b/tensorflow/core/kernels/data/BUILD
index 03f9d26cf0d..78079583ed6 100644
--- a/tensorflow/core/kernels/data/BUILD
+++ b/tensorflow/core/kernels/data/BUILD
@@ -161,6 +161,7 @@ cc_library(
         "//tensorflow/core:lib_internal",
         "//tensorflow/core:protos_all_cc",
         "//tensorflow/core/kernels:variable_ops",
+        "//tensorflow/core/profiler/lib:traceme",
         "@com_google_absl//absl/time",
     ],
 )
diff --git a/tensorflow/core/kernels/data/captured_function.cc b/tensorflow/core/kernels/data/captured_function.cc
index d9fd67c31cf..e345a2778a2 100644
--- a/tensorflow/core/kernels/data/captured_function.cc
+++ b/tensorflow/core/kernels/data/captured_function.cc
@@ -32,6 +32,7 @@ limitations under the License.
 #include "tensorflow/core/lib/random/random.h"
 #include "tensorflow/core/lib/strings/strcat.h"
 #include "tensorflow/core/platform/notification.h"
+#include "tensorflow/core/profiler/lib/traceme.h"
 
 namespace tensorflow {
 namespace data {
@@ -662,6 +663,12 @@ Status InstantiatedCapturedFunction::Run(IteratorContext* ctx,
                            ret_types_);
   Notification n;
   Status s;
+  profiler::TraceMe activity(
+      [&] {
+        return absl::StrCat(
+            "InstantiatedCapturedFunction::Run#id=", f_opts.step_id, "#");
+      },
+      profiler::TraceMeLevel::kInfo);
   lib_->Run(f_opts, f_handle_, &frame, [&n, &s](Status func_status) {
     s.Update(func_status);
     n.Notify();
@@ -700,6 +707,13 @@ Status InstantiatedCapturedFunction::RunWithBorrowedArgs(
   Notification n;
   Status s;
 
+  profiler::TraceMe activity(
+      [&] {
+        return absl::StrCat(
+            "InstantiatedCapturedFunction::RunWithBorrowedArgs#id=",
+            f_opts.step_id, "#");
+      },
+      profiler::TraceMeLevel::kInfo);
   lib_->Run(f_opts, f_handle_, &frame, [&n, &s](Status func_status) {
     s.Update(func_status);
     n.Notify();
@@ -737,6 +751,12 @@ Status InstantiatedCapturedFunction::RunInstantiated(
   Notification n;
   Status s;
 
+  profiler::TraceMe activity(
+      [&] {
+        return absl::StrCat("InstantiatedCapturedFunction::RunInstantiated#id=",
+                            f_opts.step_id, "#");
+      },
+      profiler::TraceMeLevel::kInfo);
   lib_->Run(f_opts, f_handle_, &frame, [&n, &s](Status func_status) {
     s.Update(func_status);
     n.Notify();
@@ -839,6 +859,12 @@ void InstantiatedCapturedFunction::RunAsync(
       std::move(done), ctx, std::move(deregister_fn), prefix,
       std::move(stats_collector), std::placeholders::_1);
 
+  profiler::TraceMe activity(
+      [&] {
+        return absl::StrCat(
+            "InstantiatedCapturedFunction::RunAsync#id=", f_opts.step_id, "#");
+      },
+      profiler::TraceMeLevel::kInfo);
   lib_->Run(f_opts, f_handle_, frame, std::move(callback));
 }