Add tracing for complete params

This is to measure the overhead of param resolution, so it should start before the call and end after the callback is invoked.

PiperOrigin-RevId: 339768796
Change-Id: I10b2479047226328e75fa2e3000d4ec2f7256cfe
This commit is contained in:
Ran Chen 2020-10-29 16:22:02 -07:00 committed by TensorFlower Gardener
parent e329815152
commit 2bb87b1d02

View File

@ -354,8 +354,11 @@ void BaseCollectiveExecutor::CompleteParamsAsync(
// timeout callback executes, done_safe will become a no-op and the timeout
// callback is responsible for invoking done() at the end.
const auto is_callback_called = std::make_shared<std::atomic<bool>>(false);
auto done_safe = [this, is_callback_called, cancel_mgr,
auto trace_id =
profiler::TraceMe::ActivityStart("CollectiveExecutor::CompleteParams");
auto done_safe = [this, is_callback_called, cancel_mgr, trace_id,
done](const Status& s) {
profiler::TraceMe::ActivityEnd(trace_id);
bool called = is_callback_called->exchange(true);
if (!called) {
if (!s.ok() && !IsCancelled(cancel_mgr)) {