From c6c0175e92c14bbc50614903dd948ef295ded239 Mon Sep 17 00:00:00 2001
From: Sannidhya Chauhan <sannidhya@google.com>
Date: Wed, 10 Jun 2026 02:12:49 -0700
Subject: [PATCH] Introduce the ContinuousProfilerOrchestrator in 1P

PiperOrigin-RevId: 929698357
---
 tsl/profiler/lib/BUILD                              |  1 +
 tsl/profiler/lib/continuous_profiler_orchestrator.h | 12 +++++++++++-
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/tsl/profiler/lib/BUILD b/tsl/profiler/lib/BUILD
index bcca03f6f..9b6787baf 100644
--- a/tsl/profiler/lib/BUILD
+++ b/tsl/profiler/lib/BUILD
@@ -435,5 +435,6 @@ cc_library(
         "@com_google_absl//absl/time",
         "@xla//xla/tsl/platform:env",
         "@xla//xla/tsl/platform:errors",
+        "@xla//xla/tsl/platform:logging",
     ],
 )
diff --git a/tsl/profiler/lib/continuous_profiler_orchestrator.h b/tsl/profiler/lib/continuous_profiler_orchestrator.h
index 33f233e99..711003c61 100644
--- a/tsl/profiler/lib/continuous_profiler_orchestrator.h
+++ b/tsl/profiler/lib/continuous_profiler_orchestrator.h
@@ -30,6 +30,7 @@ limitations under the License.
 #include "absl/time/time.h"
 #include "xla/tsl/platform/env.h"
 #include "xla/tsl/platform/errors.h"  // IWYU pragma: keep
+#include "xla/tsl/platform/logging.h"
 #include "tsl/profiler/lib/profiler_interface.h"
 #include "tsl/profiler/protobuf/xplane.pb.h"
 
@@ -74,6 +75,13 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface {
   // Stops background thread and profiling.
   absl::Status Stop() override {
     StopIngestionThread();
+    auto result = profiler_->Consume();
+    if (result.ok()) {
+      absl::MutexLock lock(mutex_);
+      circular_buffer_.push_back(std::move(result->data));
+    } else if (!absl::IsUnimplemented(result.status())) {
+      LOG(WARNING) << "Final Consume failed during Stop: " << result.status();
+    }
     return profiler_->Stop();
   }
 
@@ -102,7 +110,6 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface {
   ProfilerType* profiler() { return profiler_.get(); }
   const ProfilerType* profiler() const { return profiler_.get(); }
 
- private:
   std::vector<std::any> PopBuffer() {
     absl::MutexLock lock(mutex_);
     std::vector<std::any> chunks;
@@ -113,7 +120,10 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface {
     circular_buffer_.clear();
     return chunks;
   }
+
+ private:
   void IngestionLoop() {
+    LOG(INFO) << "ContinuousProfilerOrchestrator::IngestionLoop started";
     while (true) {
       absl::StatusOr<ConsumeResult> result = profiler_->Consume();