Expose ProfilerSessionOrchestrator in JAX profiler bindings.

tsl/profiler/lib/BUILD

@@ -140,7 +140,7 @@ cc_library(
     ]),
     deps = [
         "//tsl/profiler/protobuf:xplane_proto_cc",
-        "@xla//xla/tsl/platform:status",
+        "@com_google_absl//absl/status",
     ],
 )
 
@@ -378,7 +378,6 @@ cc_library(
         ":profiler_interface",
         "//tsl/profiler/protobuf:xplane_proto_cc",
         "@com_google_absl//absl/status",
-        "@xla//xla/tsl/platform:status",
     ],
 )
 
@@ -391,3 +390,42 @@ cc_library(
         "@com_google_absl//absl/strings:string_view",
     ],
 )
+
+cc_library(
+    name = "profiler_orchestrator",
+    srcs = ["profiler_orchestrator.cc"],
+    hdrs = ["profiler_orchestrator.h"],
+    visibility = internal_visibility([
+        "@xla//xla/python:__pkg__",
+    ]),
+    deps = [
+        ":profiler_session",
+        "//tsl/platform:platform_port",
+        "//tsl/profiler/protobuf:profiler_options_proto_cc",
+        "//tsl/profiler/protobuf:xplane_proto_cc",
+        "@com_google_absl//absl/status",
+        "@com_google_absl//absl/status:statusor",
+        "@xla//xla/tsl/platform:env",
+        "@xla//xla/tsl/platform:errors",
+        "@xla//xla/tsl/platform:logging",
+        "@xla//xla/tsl/profiler/convert:post_process_single_host_xplane",
+        "@xla//xla/tsl/profiler/utils:xplane_builder",
+        "@xla//xla/tsl/profiler/utils:xplane_schema",
+        "@xla//xla/tsl/profiler/utils:xplane_utils",
+    ],
+)
+
+tsl_cc_test(
+    name = "profiler_orchestrator_test",
+    srcs = ["profiler_orchestrator_test.cc"],
+    deps = [
+        ":profiler_factory",
+        ":profiler_interface",
+        ":profiler_orchestrator",
+        ":profiler_session",
+        "//tsl/profiler/protobuf:profiler_options_proto_cc",
+        "@com_google_absl//absl/memory",
+        "@com_google_absl//absl/status",
+        "@xla//xla/tsl/platform:test",
+    ],
+)

tsl/profiler/lib/profiler_collection.cc

@@ -55,5 +55,28 @@ absl::Status ProfilerCollection::CollectData(
   return status;
 }
 
+absl::Status ProfilerCollection::Consume(void* ptr) {
+  absl::Status status;
+  for (auto& profiler : profilers_) {
+    absl::Status s = profiler->Consume(ptr);
+    if (!s.ok() && !absl::IsUnimplemented(s)) {
+      status.Update(s);
+    }
+  }
+  return status;
+}
+
+absl::Status ProfilerCollection::Serialize(
+    void* ptr, tensorflow::profiler::XSpace* output_space) {
+  absl::Status status;
+  for (auto& profiler : profilers_) {
+    absl::Status s = profiler->Serialize(ptr, output_space);
+    if (!s.ok() && !absl::IsUnimplemented(s)) {
+      status.Update(s);
+    }
+  }
+  return status;
+}
+
 }  // namespace profiler
 }  // namespace tsl

tsl/profiler/lib/profiler_collection.h

@@ -19,7 +19,6 @@ limitations under the License.
 #include <vector>
 
 #include "absl/status/status.h"
-#include "xla/tsl/platform/status.h"
 #include "tsl/profiler/lib/profiler_interface.h"
 #include "tsl/profiler/protobuf/xplane.pb.h"
 
@@ -38,6 +37,9 @@ class ProfilerCollection : public ProfilerInterface {
   absl::Status Stop() override;
 
   absl::Status CollectData(tensorflow::profiler::XSpace* space) override;
+  absl::Status Consume(void* ptr) override;
+  absl::Status Serialize(void* ptr,
+                         tensorflow::profiler::XSpace* output_space) override;
 
  private:
   std::vector<std::unique_ptr<ProfilerInterface>> profilers_;

tsl/profiler/lib/profiler_controller.cc

@@ -85,5 +85,14 @@ absl::Status ProfilerController::CollectData(
   return status;
 }
 
+absl::Status ProfilerController::Consume(void* ptr) {
+  return profiler_->Consume(ptr);
+}
+
+absl::Status ProfilerController::Serialize(
+    void* ptr, tensorflow::profiler::XSpace* output_space) {
+  return profiler_->Serialize(ptr, output_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 <cstddef>
 #include <memory>
 
 #include "absl/status/status.h"
@@ -45,6 +46,11 @@ class ProfilerController : public ProfilerInterface {
 
   absl::Status CollectData(tensorflow::profiler::XSpace* space) override;
 
+  absl::Status Consume(void* ptr) override;
+
+  absl::Status Serialize(void* ptr,
+                         tensorflow::profiler::XSpace* output_space) override;
+
  private:
   enum class ProfilerState {
     kInit = 0,

tsl/profiler/lib/profiler_interface.h

@@ -15,7 +15,7 @@ limitations under the License.
 #ifndef TENSORFLOW_TSL_PROFILER_LIB_PROFILER_INTERFACE_H_
 #define TENSORFLOW_TSL_PROFILER_LIB_PROFILER_INTERFACE_H_
 
-#include "xla/tsl/platform/status.h"
+#include "absl/status/status.h"
 #include "tsl/profiler/protobuf/xplane.pb.h"
 
 namespace tsl {
@@ -41,6 +41,17 @@ class ProfilerInterface {
 
   // Saves collected profile data into XSpace.
   virtual absl::Status CollectData(tensorflow::profiler::XSpace* space) = 0;
+
+  // Pulls collected profile data into arbitrary raw memory.
+  virtual absl::Status Consume(void* ptr) {
+    return absl::UnimplementedError("Consume not implemented");
+  }
+
+  // Serializes collected profile data into XSpace.
+  virtual absl::Status Serialize(void* ptr,
+                                 tensorflow::profiler::XSpace* output_space) {
+    return absl::UnimplementedError("Serialize not implemented");
+  }
 };
 
 }  // namespace profiler

tsl/profiler/lib/profiler_orchestrator.cc

@@ -0,0 +1,120 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tsl/profiler/lib/profiler_orchestrator.h"
+
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "xla/tsl/platform/env.h"
+#include "xla/tsl/platform/errors.h"
+#include "xla/tsl/platform/logging.h"
+#include "xla/tsl/profiler/convert/post_process_single_host_xplane.h"
+#include "xla/tsl/profiler/utils/xplane_builder.h"
+#include "xla/tsl/profiler/utils/xplane_schema.h"
+#include "xla/tsl/profiler/utils/xplane_utils.h"
+#include "tsl/platform/host_info.h"
+#include "tsl/profiler/lib/profiler_session.h"
+#include "tsl/profiler/protobuf/profiler_options.pb.h"
+
+namespace tsl {
+namespace profiler {
+
+ProfilerSessionOrchestrator::ProfilerSessionOrchestrator(
+    const tensorflow::ProfileOptions& options)
+    : options_(options) {}
+
+ProfilerSessionOrchestrator::~ProfilerSessionOrchestrator() {
+  Stop().IgnoreError();
+}
+
+absl::Status ProfilerSessionOrchestrator::Start() {
+  if (session_ != nullptr) {
+    return absl::FailedPreconditionError("Session already started.");
+  }
+  global_start_time_ns_ = tsl::Env::Default()->NowNanos();
+  session_ = tsl::ProfilerSession::Create(options_);
+  if (session_ == nullptr) {
+    return absl::InternalError("Failed to create ProfilerSession.");
+  }
+  return session_->Status();
+}
+
+absl::Status ProfilerSessionOrchestrator::Stop() {
+  if (session_ == nullptr) {
+    return absl::OkStatus();  // Already stopped or not started.
+  }
+  session_.reset();
+  return absl::OkStatus();
+}
+
+absl::StatusOr<int> ProfilerSessionOrchestrator::Consume() {
+  if (session_ == nullptr) {
+    return absl::FailedPreconditionError("Session not started.");
+  }
+
+  consume_buffers_.emplace_back(sizeof(std::vector<char>));
+  auto& buffer = consume_buffers_.back();
+  TF_RETURN_IF_ERROR(session_->Consume(buffer.data()));
+
+  uint64_t now = tsl::Env::Default()->NowNanos();
+  consume_stop_times_.push_back(now);
+
+  return consume_buffers_.size() - 1;
+}
+
+absl::Status ProfilerSessionOrchestrator::Serialize(int buffer_index) {
+  if (session_ == nullptr) {
+    return absl::FailedPreconditionError("Session not started.");
+  }
+
+  if (buffer_index < 0 || buffer_index >= consume_buffers_.size()) {
+    return absl::InvalidArgumentError("Invalid buffer index.");
+  }
+  serialize_space_.Clear();
+  auto& buffer = consume_buffers_[buffer_index];
+  TF_RETURN_IF_ERROR(session_->Serialize(buffer.data(), &serialize_space_));
+
+  serialize_space_.add_hostnames(tsl::port::Hostname());
+  profiler::SetXSpacePidIfNotSet(serialize_space_,
+                                 tsl::Env::Default()->GetProcessId());
+
+  uint64_t stop_time = 0;
+  if (buffer_index < consume_stop_times_.size()) {
+    stop_time = consume_stop_times_[buffer_index];
+  }
+  profiler::PostProcessSingleHostXSpace(&serialize_space_,
+                                        global_start_time_ns_, stop_time);
+
+  {
+    profiler::XPlaneBuilder xplane(profiler::FindOrAddMutablePlaneWithName(
+        &serialize_space_, tsl::profiler::kTaskEnvPlaneName));
+    xplane.AddStatValue(
+        *xplane.GetOrCreateStatMetadata(tsl::profiler::GetTaskEnvStatTypeStr(
+            tsl::profiler::kEnvProfileOptions)),
+        options_);
+  }
+
+  return absl::OkStatus();
+}
+
+void ProfilerSessionOrchestrator::ClearConsumeBuffers() {
+  std::vector<std::vector<uint8_t>>().swap(consume_buffers_);
+  consume_stop_times_.clear();
+}
+
+}  // namespace profiler
+}  // namespace tsl

tsl/profiler/lib/profiler_orchestrator.h

@@ -0,0 +1,66 @@
+/* Copyright 2026 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#ifndef TENSORFLOW_TSL_PROFILER_LIB_PROFILER_ORCHESTRATOR_H_
+#define TENSORFLOW_TSL_PROFILER_LIB_PROFILER_ORCHESTRATOR_H_
+
+#include <cstdint>
+#include <memory>
+#include <vector>
+
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "tsl/profiler/lib/profiler_session.h"
+#include "tsl/profiler/protobuf/profiler_options.pb.h"
+#include "tsl/profiler/protobuf/xplane.pb.h"
+
+namespace tsl {
+namespace profiler {
+
+class ProfilerSessionOrchestrator {
+ public:
+  explicit ProfilerSessionOrchestrator(
+      const tensorflow::ProfileOptions& options);
+  ~ProfilerSessionOrchestrator();
+
+  absl::Status Start();
+
+  absl::StatusOr<int> Consume();
+
+  absl::Status Serialize(int buffer_index);
+
+  absl::Status Stop();
+
+  void ClearConsumeBuffers();
+
+  const std::vector<uint8_t>& GetConsumeBuffer(int index) const {
+    return consume_buffers_[index];
+  }
+  const tensorflow::profiler::XSpace& GetSerializeSpace() const {
+    return serialize_space_;
+  }
+
+ private:
+  tensorflow::ProfileOptions options_;
+  std::unique_ptr<tsl::ProfilerSession> session_;
+  std::vector<std::vector<uint8_t>> consume_buffers_;
+  tensorflow::profiler::XSpace serialize_space_;
+  uint64_t global_start_time_ns_;
+  std::vector<uint64_t> consume_stop_times_;
+};
+
+}  // namespace profiler
+}  // namespace tsl
+
+#endif  // TENSORFLOW_TSL_PROFILER_LIB_PROFILER_ORCHESTRATOR_H_