make sure we capture up to 1 million trace events to avoid crashing the trace viewer. add an unit-test

PiperOrigin-RevId: 298403826
Change-Id: Idfcbf9501e9d2a626828c83501b2c548b8df48dd
This commit is contained in:
A. Unique TensorFlower 2020-03-02 11:35:59 -08:00 committed by TensorFlower Gardener
parent a190eddea4
commit e82714b529
3 changed files with 45 additions and 0 deletions

View File

@ -36,6 +36,27 @@ Device BuildDeviceAndResource(const XPlaneVisitor& plane) {
}
} // namespace
void MaybeDropEventsForTraceViewer(Trace* trace, uint32 limit) {
auto* trace_events = trace->mutable_trace_events();
size_t trace_event_size = trace_events->size();
if (trace_event_size <= limit) return; // Nothing to do.
// Sort the events according to start time.
std::vector<uint64> timestamps;
timestamps.reserve(trace_event_size);
for (const auto& event : *trace_events) {
timestamps.push_back(event.timestamp_ps());
}
std::partial_sort(timestamps.begin(), timestamps.begin() + limit,
timestamps.end(), std::less<uint64>());
uint64 cutoff_timestamp = timestamps[limit - 1];
trace_events->erase(std::remove_if(trace_events->begin(), trace_events->end(),
[&](const TraceEvent& event) {
return event.timestamp_ps() >
cutoff_timestamp;
}),
trace_events->end());
}
void ConvertXSpaceToTraceEvents(const XSpace& xspace, Trace* trace) {
auto* trace_devices = trace->mutable_devices();
@ -69,6 +90,11 @@ void ConvertXSpaceToTraceEvents(const XSpace& xspace, Trace* trace) {
});
});
}
// Trace viewer (non-streaming) has scalability issues, we need to drop
// events to avoid loading failure for trace viewer.
constexpr uint64 kMaxEvents = 1000000;
MaybeDropEventsForTraceViewer(trace, kMaxEvents);
}
} // namespace profiler

View File

@ -26,6 +26,9 @@ namespace profiler {
void ConvertXSpaceToTraceEvents(const XSpace& xspace, Trace* trace);
// Not Public API, Testing only.
void MaybeDropEventsForTraceViewer(Trace* trace, uint32 limit);
} // namespace profiler
} // namespace tensorflow

View File

@ -71,6 +71,22 @@ TEST(ConvertXPlaneToTraceEvents, Convert) {
EXPECT_EQ(trace.trace_events_size(), 3);
}
TEST(ConvertXPlaneToTraceEvents, Drop) {
Trace trace;
for (int i = 0; i < 100; i++) {
trace.add_trace_events()->set_timestamp_ps((100 - i) % 50);
}
MaybeDropEventsForTraceViewer(&trace, 150);
EXPECT_EQ(trace.trace_events_size(), 100); // No dropping.
MaybeDropEventsForTraceViewer(&trace, 50);
EXPECT_EQ(trace.trace_events_size(), 50);
for (const auto& event : trace.trace_events()) {
EXPECT_LT(event.timestamp_ps(), 25);
}
}
} // namespace
} // namespace profiler
} // namespace tensorflow