From e3ce6536e89cf9c509415de7cf506079f2ff361b Mon Sep 17 00:00:00 2001
From: Sannidhya Chauhan <sannidhya@google.com>
Date: Mon, 22 Jun 2026 03:54:10 -0700
Subject: [PATCH] Enable continuous profiling for HostTracer

PiperOrigin-RevId: 935970674
---
 tsl/profiler/lib/BUILD                              |  1 +
 tsl/profiler/lib/continuous_profiler_orchestrator.h | 11 ++++++++++-
 2 files changed, 11 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..c9a26d9ee 100644
--- a/tsl/profiler/lib/continuous_profiler_orchestrator.h
+++ b/tsl/profiler/lib/continuous_profiler_orchestrator.h
@@ -74,6 +74,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 +109,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 +119,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();