From a81774fb4578b8c13014599d28fb45e84532a43a Mon Sep 17 00:00:00 2001 From: Yi Situ Date: Wed, 21 Oct 2020 14:53:54 -0700 Subject: [PATCH] Derive profiling start timestamp from delay_ms. PiperOrigin-RevId: 338345983 Change-Id: Ia6f0ccaba88d5b721807b53c6f1075f0b9719ab1 --- .../integration_test/profiler_api_test.py | 8 +++++++- .../profiler/internal/profiler_wrapper.cc | 18 ++++++++++++++++++ 2 files changed, 25 insertions(+), 1 deletion(-) diff --git a/tensorflow/python/profiler/integration_test/profiler_api_test.py b/tensorflow/python/profiler/integration_test/profiler_api_test.py index 7b9b51ae9c2..4e2a9dfd4e3 100644 --- a/tensorflow/python/profiler/integration_test/profiler_api_test.py +++ b/tensorflow/python/profiler/integration_test/profiler_api_test.py @@ -83,7 +83,7 @@ class ProfilerApiTest(test_util.TensorFlowTestCase): model.fit(x=train_ds, epochs=2, steps_per_epoch=steps) - def test_single_worker_sampling_mode(self): + def test_single_worker_sampling_mode(self, delay_ms=None): """Test single worker sampling mode.""" def on_worker(port): @@ -100,6 +100,7 @@ class ProfilerApiTest(test_util.TensorFlowTestCase): host_tracer_level=2, python_tracer_level=0, device_tracer_level=1, + delay_ms=delay_ms, ) profiler_client.trace('localhost:{}'.format(port), logdir, duration_ms, @@ -115,6 +116,11 @@ class ProfilerApiTest(test_util.TensorFlowTestCase): thread_worker.join(120) self._check_xspace_pb_exist(logdir) + def test_single_worker_sampling_mode_delayed(self): + """Test single worker sampling mode with delay.""" + + self.test_single_worker_sampling_mode(delay_ms=1000) + def test_single_worker_programmatic_mode(self): """Test single worker programmatic mode.""" logdir = self.get_temp_dir() diff --git a/tensorflow/python/profiler/internal/profiler_wrapper.cc b/tensorflow/python/profiler/internal/profiler_wrapper.cc index 485a67dc178..2b513547612 100644 --- a/tensorflow/python/profiler/internal/profiler_wrapper.cc +++ b/tensorflow/python/profiler/internal/profiler_wrapper.cc @@ -146,6 +146,7 @@ RemoteProfilerSessionManagerOptions GetOptionsLocked(absl::string_view logdir, VLOG(2) << "repository_path set to " << options.profiler_options().repository_path(); + int delay_ms = 0; for (const auto& kw : opts) { std::string key = py::cast(kw.first); if (key == "host_tracer_level") { @@ -160,11 +161,28 @@ RemoteProfilerSessionManagerOptions GetOptionsLocked(absl::string_view logdir, auto value = py::cast(kw.second); options.mutable_profiler_options()->set_python_tracer_level(value); VLOG(1) << "python_tracer_level set to " << value; + } else if (key == "delay_ms") { + if (!kw.second.is_none()) { + delay_ms = py::cast(kw.second); + } } else { LOG(WARNING) << "Unrecognised key: " << key; } } + if (delay_ms) { + absl::Time start_timestamp = now + absl::Milliseconds(delay_ms); + tensorflow::int64 start_timestamp_ns = absl::ToUnixNanos(start_timestamp); + options.mutable_profiler_options()->set_start_timestamp_ns( + start_timestamp_ns); + LOG(INFO) << "delay_ms was " << delay_ms << ", start_timestamp_ns set to " + << start_timestamp_ns << " [" << start_timestamp << "]"; + } else { + DCHECK_EQ(options.mutable_profiler_options()->start_timestamp_ns(), 0); + LOG(INFO) << "Profiling will start immediately because delay_ms was unset " + "or zero."; + } + return options; }