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

Add TraceMe to track which step_id is from tf data.

Browse Source 
PiperOrigin-RevId: 280760041
Change-Id: I76e41c26e21088c831d24a14201ac1adfb46b268
This commit is contained in:
A. Unique TensorFlower 2019-11-15 16:29:54 -08:00 committed by TensorFlower Gardener
parent ad9485bd46
commit 2f63b65d95
2 changed files with 27 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
tensorflow/core/kernels/data/BUILD
View File

@ -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",
],
)

26
tensorflow/core/kernels/data/captured_function.cc
View File

@ -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));
}