From 63d879e7fcc6727c2f91c5937542b3f8d502123a Mon Sep 17 00:00:00 2001 From: Aditya Sharma Date: Mon, 5 Jan 2026 11:45:44 -0800 Subject: [PATCH] Implement GTC sample eviction in FullTraceCollector. PiperOrigin-RevId: 852392402 --- tsl/profiler/lib/profiler_factory.cc | 19 ++++++++++++++++++- 1 file changed, 18 insertions(+), 1 deletion(-) diff --git a/tsl/profiler/lib/profiler_factory.cc b/tsl/profiler/lib/profiler_factory.cc index a11a11b50..acb625b59 100644 --- a/tsl/profiler/lib/profiler_factory.cc +++ b/tsl/profiler/lib/profiler_factory.cc @@ -44,10 +44,27 @@ void RegisterProfilerFactory(ProfilerFactory factory) { std::vector> CreateProfilers( const tensorflow::ProfileOptions& options) { + // Create a copy of options to modify if circular buffer is enabled. + tensorflow::ProfileOptions options_to_use = options; + bool circular_buffer_enabled = false; + + // Check if tpu_circular_buffer_tracing is enabled in advanced configuration. + auto it = + options.advanced_configuration().find("tpu_circular_buffer_tracing"); + if (it != options.advanced_configuration().end()) { + circular_buffer_enabled = it->second.bool_value(); + } + + if (circular_buffer_enabled) { + // Disable other tracers by zeroing their levels in the local options copy. + options_to_use.set_host_tracer_level(0); + options_to_use.set_python_tracer_level(0); + } + std::vector> result; absl::MutexLock lock(mu); for (const auto& factory : *GetFactories()) { - auto profiler = factory(options); + auto profiler = factory(options_to_use); // A factory might return nullptr based on options. if (profiler == nullptr) continue; result.emplace_back(