Support chunked XSpace serialization in jax.stop_trace when enable_continuous_profiling is enabled

tsl/profiler/lib/BUILD

@@ -193,6 +193,7 @@ tsl_cc_test(
         "//tsl/profiler/protobuf:xplane_proto_cc",
         "@com_google_absl//absl/status",
         "@com_google_absl//absl/status:statusor",
+        "@com_google_absl//absl/synchronization",
         "@com_google_absl//absl/time",
         "@com_google_googletest//:gtest_main",
         "@xla//xla/tsl/platform:test",

tsl/profiler/lib/continuous_profiler_orchestrator.h

@@ -101,6 +101,23 @@ class ContinuousProfilerOrchestrator : public ProfilerInterface {
     return status;
   }
 
+  std::vector<tensorflow::profiler::XSpace> SerializeChunks() {
+    std::vector<std::any> chunks = PopBuffer();
+    std::vector<tensorflow::profiler::XSpace> spaces;
+    spaces.reserve(chunks.size());
+    tensorflow::profiler::XSpace space;
+    for (auto& chunk : chunks) {
+      space.Clear();
+      absl::Status status = profiler_->Serialize(std::move(chunk), &space);
+      if (status.ok()) {
+        spaces.push_back(std::move(space));
+      } else {
+        LOG(ERROR) << "Failed to serialize profiler chunk: " << status;
+      }
+    }
+    return spaces;
+  }
+
   // Returns the current polling interval (primarily for testing).
   absl::Duration polling_interval() const {
     absl::MutexLock lock(mutex_);

tsl/profiler/lib/continuous_profiler_orchestrator_test.cc

@@ -18,10 +18,11 @@ limitations under the License.
 #include <atomic>
 #include <memory>
 #include <utility>
+#include <vector>
 
 #include "absl/status/status.h"
 #include "absl/status/statusor.h"
-#include "absl/time/clock.h"
+#include "absl/synchronization/notification.h"
 #include "absl/time/time.h"
 #include "xla/tsl/platform/test.h"
 #include "tsl/profiler/lib/profiler_interface.h"
@@ -31,7 +32,6 @@ namespace tsl {
 namespace profiler {
 namespace {
 
-using ::testing::Invoke;
 using ::testing::Pointee;
 using ::testing::Return;
 
@@ -67,29 +67,28 @@ TEST(ContinuousProfilerOrchestratorTest,
 
   // Setup Consume mock to return sequential chunks with high sizes to shrink
   // the interval
+  absl::Notification consumed_enough;
   std::atomic<int> consume_count(0);
   EXPECT_CALL(*mock, Consume())
-      .WillRepeatedly(Invoke([&]() -> absl::StatusOr<ConsumeResult> {
+      .WillRepeatedly([&]() -> absl::StatusOr<ConsumeResult> {
         int count = ++consume_count;
+        if (count >= 4 && !consumed_enough.HasBeenNotified()) {
+          consumed_enough.Notify();
+        }
         return ConsumeResult{
             .data = std::any(count),
             .estimated_size_bytes = 1000 * 1024 * 1024  // 1000MB (>512MB)
         };
-      }));
+      });
 
   ContinuousProfilerOrchestrator<ProfilerInterface> orchestrator(
       std::move(mock_profiler));
 
   // Start orchestrator (spawns background loop)
   ASSERT_OK(orchestrator.Start());
 
-  // Wait until we have consumed at least 4 chunks.
-  // Due to high watermark scaling, interval shrinks from 1s -> 500ms -> 250ms
-  // -> 125ms -> 100ms. This will happen in less than 1.0 second of real-time
-  // sleep.
-  while (consume_count < 4) {
-    absl::SleepFor(absl::Milliseconds(50));
-  }
+  // Wait until we have consumed at least 4 chunks using notification.
+  consumed_enough.WaitForNotification();
 
   // Stop orchestrator (terminates background loop)
   ASSERT_OK(orchestrator.Stop());
@@ -122,13 +121,17 @@ TEST(ContinuousProfilerOrchestratorTest, DynamicIntervalLowWatermarkScaling) {
   EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus()));
 
   // Consume returns a very small chunk (1MB < 5MB low watermark)
+  absl::Notification first_consumed;
   EXPECT_CALL(*mock, Consume())
-      .WillRepeatedly(Invoke([]() -> absl::StatusOr<ConsumeResult> {
+      .WillRepeatedly([&]() -> absl::StatusOr<ConsumeResult> {
+        if (!first_consumed.HasBeenNotified()) {
+          first_consumed.Notify();
+        }
         return ConsumeResult{
             .data = std::any(1),
             .estimated_size_bytes = 1 * 1024 * 1024  // 1MB (<5MB)
         };
-      }));
+      });
 
   ContinuousProfilerOrchestrator<ProfilerInterface> orchestrator(
       std::move(mock_profiler));
@@ -137,9 +140,8 @@ TEST(ContinuousProfilerOrchestratorTest, DynamicIntervalLowWatermarkScaling) {
   // Start
   ASSERT_OK(orchestrator.Start());
 
-  // Wait a very short time for the first immediate consume to run and adjust
-  // the interval
-  absl::SleepFor(absl::Milliseconds(50));
+  // Wait for the first immediate consume to run and adjust the interval
+  first_consumed.WaitForNotification();
 
   // Stop immediately
   ASSERT_OK(orchestrator.Stop());
@@ -148,6 +150,42 @@ TEST(ContinuousProfilerOrchestratorTest, DynamicIntervalLowWatermarkScaling) {
   EXPECT_EQ(orchestrator.polling_interval(), absl::Seconds(2));
 }
 
+TEST(ContinuousProfilerOrchestratorTest, SerializeChunks) {
+  auto mock_profiler = std::make_unique<MockProfiler>();
+  MockProfiler* mock = mock_profiler.get();
+
+  EXPECT_CALL(*mock, Start()).WillOnce(Return(absl::OkStatus()));
+  EXPECT_CALL(*mock, Stop()).WillOnce(Return(absl::OkStatus()));
+
+  absl::Notification consumed_enough;
+  std::atomic<int> consume_count(0);
+  EXPECT_CALL(*mock, Consume())
+      .WillRepeatedly([&]() -> absl::StatusOr<ConsumeResult> {
+        int count = ++consume_count;
+        if (count >= 2 && !consumed_enough.HasBeenNotified()) {
+          consumed_enough.Notify();
+        }
+        return ConsumeResult{.data = std::any(count),
+                             .estimated_size_bytes = 10 * 1024 * 1024};
+      });
+
+  ContinuousProfilerOrchestrator<ProfilerInterface> orchestrator(
+      std::move(mock_profiler));
+
+  ASSERT_OK(orchestrator.Start());
+  consumed_enough.WaitForNotification();
+  ASSERT_OK(orchestrator.Stop());
+
+  EXPECT_CALL(*mock, SerializeMock(::testing::_, ::testing::_))
+      .WillRepeatedly(Return(absl::InternalError("serialization error")));
+  EXPECT_CALL(*mock, SerializeMock(Pointee(AnyEqInt(1)), ::testing::_))
+      .WillOnce(Return(absl::OkStatus()));
+
+  std::vector<tensorflow::profiler::XSpace> spaces =
+      orchestrator.SerializeChunks();
+  EXPECT_EQ(spaces.size(), 1);
+}
+
 }  // namespace
 }  // namespace profiler
 }  // namespace tsl

tsl/profiler/lib/profiler_controller.cc

@@ -14,9 +14,11 @@ limitations under the License.
 ==============================================================================*/
 #include "tsl/profiler/lib/profiler_controller.h"
 
+#include <any>
 #include <memory>
 #include <utility>
 
+#include "absl/status/status.h"
 #include "xla/tsl/platform/errors.h"
 #include "xla/tsl/platform/logging.h"
 #include "tsl/profiler/lib/profiler_interface.h"
@@ -85,5 +87,20 @@ absl::Status ProfilerController::CollectData(
   return status;
 }
 
+absl::StatusOr<ConsumeResult> ProfilerController::Consume() {
+  if (state_ != ProfilerState::kStart) {
+    return absl::AbortedError("Consume called in the wrong order.");
+  }
+  if (!status_.ok()) {
+    return absl::AbortedError("Previous call returned an error.");
+  }
+  return profiler_->Consume();
+}
+
+absl::Status ProfilerController::Serialize(
+    std::any data, tensorflow::profiler::XSpace* space) {
+  return profiler_->Serialize(std::move(data), space);
+}
+
 }  // namespace profiler
 }  // namespace tsl

tsl/profiler/lib/profiler_controller.h

@@ -15,6 +15,7 @@ limitations under the License.
 #ifndef TENSORFLOW_TSL_PROFILER_LIB_PROFILER_CONTROLLER_H_
 #define TENSORFLOW_TSL_PROFILER_LIB_PROFILER_CONTROLLER_H_
 
+#include <any>
 #include <memory>
 
 #include "absl/status/status.h"
@@ -45,6 +46,11 @@ class ProfilerController : public ProfilerInterface {
 
   absl::Status CollectData(tensorflow::profiler::XSpace* space) override;
 
+  absl::StatusOr<ConsumeResult> Consume() override;
+
+  absl::Status Serialize(std::any data,
+                         tensorflow::profiler::XSpace* space) override;
+
  private:
   enum class ProfilerState {
     kInit = 0,

tsl/profiler/lib/profiler_session.cc

@@ -18,6 +18,7 @@ limitations under the License.
 #include <cstdint>
 #include <memory>
 #include <utility>
+#include <vector>
 
 #include "absl/memory/memory.h"
 #include "absl/status/status.h"
@@ -77,6 +78,53 @@ absl::Status ProfilerSession::Status() {
   return status_;
 }
 
+bool ProfilerSession::IsContinuousProfilingEnabled() const {
+#if defined(IS_MOBILE_PLATFORM)
+  return false;
+#else
+  const auto& advanced_config = options_.advanced_configuration();
+  auto it = advanced_config.find("enable_continuous_profiling");
+  return (it != advanced_config.end()) && it->second.bool_value();
+#endif
+}
+
+absl::Status ProfilerSession::Stop() {
+#if !defined(IS_MOBILE_PLATFORM)
+  absl::MutexLock l(mutex_);
+  if (profilers_ != nullptr) {
+    auto status = profilers_->Stop();
+    stop_time_ns_ = profiler::GetCurrentTimeNanos();
+    return status;
+  }
+#endif
+  return absl::OkStatus();
+}
+
+std::vector<tensorflow::profiler::XSpace> ProfilerSession::SerializeChunks() {
+  std::vector<tensorflow::profiler::XSpace> spaces;
+#if !defined(IS_MOBILE_PLATFORM)
+  absl::MutexLock l(mutex_);
+  if (profilers_ != nullptr) {
+    auto* orchestrator = dynamic_cast<
+        profiler::ContinuousProfilerOrchestrator<profiler::ProfilerInterface>*>(
+        profilers_.get());
+    if (orchestrator != nullptr) {
+      spaces = orchestrator->SerializeChunks();
+    }
+    for (auto& space : spaces) {
+      profiler::SetXSpacePidIfNotSet(space,
+                                     tsl::Env::Default()->GetProcessId());
+      profiler::PostProcessSingleHostXSpace(&space, start_time_ns_,
+                                            stop_time_ns_);
+      SetProfileOptionsIntoSpace(options_, &space);
+    }
+    profilers_.reset();
+  }
+  profiler_lock_.ReleaseIfActive();
+#endif
+  return spaces;
+}
+
 #if !defined(IS_MOBILE_PLATFORM)
 absl::Status ProfilerSession::CollectDataInternal(XSpace* space) {
   absl::MutexLock l(mutex_);

tsl/profiler/lib/profiler_session.h

@@ -15,15 +15,12 @@ limitations under the License.
 #ifndef TENSORFLOW_TSL_PROFILER_LIB_PROFILER_SESSION_H_
 #define TENSORFLOW_TSL_PROFILER_LIB_PROFILER_SESSION_H_
 
-#include <functional>
+#include <cstdint>
 #include <memory>
 #include <vector>
 
 #include "absl/status/status.h"
 #include "absl/synchronization/mutex.h"
-#include "xla/tsl/platform/status.h"
-#include "xla/tsl/platform/types.h"
-#include "tsl/platform/platform.h"
 #include "tsl/platform/thread_annotations.h"
 #include "tsl/profiler/protobuf/profiler_options.pb.h"
 #include "tsl/profiler/protobuf/xplane.pb.h"
@@ -67,6 +64,13 @@ class ProfilerSession {
   absl::Status CollectData(tensorflow::profiler::XSpace* space)
       TF_LOCKS_EXCLUDED(mutex_);
 
+  absl::Status Stop() TF_LOCKS_EXCLUDED(mutex_);
+
+  std::vector<tensorflow::profiler::XSpace> SerializeChunks()
+      TF_LOCKS_EXCLUDED(mutex_);
+
+  bool IsContinuousProfilingEnabled() const;
+
  private:
   // Constructs an instance of the class and starts profiling
   explicit ProfilerSession(const tensorflow::ProfileOptions& options);