Allow specifying a custom session id instead of always using timestamp.

This change introduces a `session_id` option in `ProfileOptions` and `RemoteProfilerSessionManagerOptions`. When provided, this id will be used as the subdirectory for storing profile data, instead of the default timestamp.

PiperOrigin-RevId: 825216123

third_party/xla/xla/tsl/profiler/rpc/client/capture_profile.cc

@@ -199,8 +199,10 @@ absl::Status CaptureRemoteTrace(const std::string& logdir,
   DCHECK_GT(opts.profiler_options().duration_ms(), 0);
   DCHECK(!opts.service_addresses().empty());
 
-  // Use the current timestamp as the run name.
-  std::string session_id = GetCurrentTimeStampAsString();
+  // Sets the session ID if provided, otherwise uses the current timestamp.
+  std::string session_id = opts.profiler_options().session_id().empty()
+                               ? GetCurrentTimeStampAsString()
+                               : opts.profiler_options().session_id();
   std::string repository_root = GetTensorBoardProfilePluginDir(logdir);
   auto duration_ms = opts.profiler_options().duration_ms();
 

third_party/xla/xla/tsl/profiler/utils/BUILD

@@ -486,6 +486,18 @@ cc_library(
     ],
 )
 
+tsl_cc_test(
+    name = "session_manager_test",
+    srcs = ["session_manager_test.cc"],
+    deps = [
+        ":session_manager",
+        "//xla/tsl/platform:test",
+        "@com_google_absl//absl/container:flat_hash_map",
+        "@com_google_absl//absl/strings:string_view",
+        "@com_google_googletest//:gtest_main",
+    ],
+)
+
 cc_library(
     name = "timestamp_utils",
     srcs = ["timestamp_utils.cc"],

third_party/xla/xla/tsl/profiler/utils/session_manager.cc

@@ -136,6 +136,13 @@ RemoteProfilerSessionManagerOptions GetRemoteSessionManagerOptionsLocked(
             options.set_delay_ms(value);
           },
           nullptr);
+    } else if (key == "session_id") {
+      SetOption<std::string>(
+          key, kw.second,
+          [&options](tensorflow::ProfileOptions*, std::string value) {
+            options.mutable_profiler_options()->set_session_id(value);
+          },
+          nullptr);
     } else {
       LOG(WARNING) << "Unrecognised key: " << key;
     }

third_party/xla/xla/tsl/profiler/utils/session_manager_test.cc

@@ -0,0 +1,50 @@
+/* Copyright 2025 The OpenXLA Authors.
+
+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 "xla/tsl/profiler/utils/session_manager.h"
+
+#include <string>
+#include <variant>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/strings/string_view.h"
+#include "xla/tsl/platform/test.h"
+
+namespace tsl {
+namespace profiler {
+namespace {
+
+using tensorflow::RemoteProfilerSessionManagerOptions;
+
+TEST(SessionManagerTest, OptionsWithSessionIdTest) {
+  absl::string_view logdir = "/tmp/logdir";
+  absl::flat_hash_map<std::string, std::variant<bool, int, std::string>> opts;
+  opts["session_id"] = std::string("test_session_id");
+  RemoteProfilerSessionManagerOptions options =
+      GetRemoteSessionManagerOptionsLocked(logdir, opts);
+  EXPECT_EQ(options.profiler_options().session_id(), "test_session_id");
+}
+
+TEST(SessionManagerTest, OptionsWithoutSessionIdTest) {
+  absl::string_view logdir = "/tmp/logdir";
+  absl::flat_hash_map<std::string, std::variant<bool, int, std::string>> opts;
+  RemoteProfilerSessionManagerOptions options =
+      GetRemoteSessionManagerOptionsLocked(logdir, opts);
+  EXPECT_EQ(options.profiler_options().session_id().empty(), true);
+}
+
+}  // namespace
+}  // namespace profiler
+}  // namespace tsl