Plumb session_id from ProfileOptions to traceSessionName.

PiperOrigin-RevId: 919283544

Author: lukebaumann

pathwaysutils/profiling.py

@@ -35,6 +35,14 @@
 
 
 class _ProfileState:
+  """Holds the state of an ongoing profiling session.
+
+  Attributes:
+    executable: The `plugin_executable.PluginExecutable` instance used for the
+      profiling session.
+    profile_request: The mapping containing the profile request options.
+    lock: A thread lock to protect access to the state.
+  """
   executable: plugin_executable.PluginExecutable | None = None
   profile_request: Mapping[str, Any] | None = None
   lock: threading.Lock
@@ -48,6 +56,37 @@ def reset(self) -> None:
     self.executable = None
     self.profile_request = None
 
+  def call_profile_executable(self) -> None:
+    """Calls the profiling executable and waits for the result."""
+    if self.executable is None:
+      raise RuntimeError(
+          "_call_profile_executable called with no active executable."
+      )
+    # If the profile request contains xprofTraceOptions, then we need to pass
+    # out_avals and out_shardings to the executable call because the
+    # executable will return a future that needs to be resolved. This is true
+    # for both starting and stopping a trace.
+    if (
+        self.profile_request is not None
+        and "xprofTraceOptions" in self.profile_request
+    ):
+      out_avals = [jax.core.ShapedArray((1,), jnp.object_)]
+      out_shardings = [
+          getattr(
+              jax.sharding,
+              "make_single_device_sharding",
+              jax.sharding.SingleDeviceSharding,
+          )(jax.devices()[0])
+      ]
+    else:
+      out_avals = ()
+      out_shardings = ()
+
+    _, result_future = self.executable.call(
+        out_avals=out_avals, out_shardings=out_shardings
+    )
+    result_future.result()
+
 
 _first_profile_start = True
 _profile_state = _ProfileState()
@@ -74,6 +113,7 @@ def _is_default_profile_options(
       == default_options.python_tracer_level
       and profiler_options.duration_ms == default_options.duration_ms
       and not getattr(profiler_options, "advanced_configuration", None)
+      and not getattr(profiler_options, "session_id", None)
   )
 
 
@@ -128,6 +168,9 @@ def _create_profile_request(
   if pw_trace_opts:
     xprof_options["pwTraceOptions"] = pw_trace_opts
 
+  if getattr(profiler_options, "session_id", None):
+    xprof_options["traceSessionName"] = profiler_options.session_id
+
   profile_request["xprofTraceOptions"] = xprof_options
 
   if profiler_options.duration_ms > 0:
@@ -153,19 +196,30 @@ def _start_pathways_trace_from_profile_request(
       toy_computation()
 
     if _profile_state.executable is not None:
-      raise ValueError(
+      raise RuntimeError(
           "start_trace called while a trace is already being taken!"
       )
-    _profile_state.executable = plugin_executable.PluginExecutable(
-        json.dumps({"profileRequest": profile_request})
-    )
-    _profile_state.profile_request = profile_request
     try:
-      _, result_future = _profile_state.executable.call()
-      result_future.result()
-    except Exception:
-      _logger.exception("Failed to start trace")
+      _profile_state.executable = plugin_executable.PluginExecutable(
+          json.dumps({"profileRequest": profile_request})
+      )
+      _profile_state.profile_request = profile_request
+      _profile_state.call_profile_executable()
+    except Exception as e:
       _profile_state.reset()
+      if (
+          "xprofTraceOptions" in profile_request
+          and "traceSessionName" in profile_request["xprofTraceOptions"]
+      ):
+        if "Bad PluginProgram" in str(e):
+          raise RuntimeError(
+              "Failed to start Pathways trace. The Pathways backend rejected "
+              "the request, likely because the running Pathways server images "
+              "do not support the trace session ID option. Please ensure you "
+              "are running the latest versions of both Pathways server images "
+              "and the pathwaysutils library."
+          ) from e
+      _logger.exception("Failed to start trace")
       raise
 
 
@@ -243,28 +297,9 @@ def stop_trace() -> None:
   try:
     with _profile_state.lock:
       if _profile_state.executable is None:
-        raise ValueError("stop_trace called before a trace is being taken!")
+        raise RuntimeError("stop_trace called before a trace is being taken!")
       try:
-        if (
-            _profile_state.profile_request is not None
-            and "xprofTraceOptions" in _profile_state.profile_request
-        ):
-          out_avals = [jax.core.ShapedArray((1,), jnp.object_)]
-          out_shardings = [
-              getattr(
-                  jax.sharding,
-                  "make_single_device_sharding",
-                  lambda x: jax.sharding.SingleDeviceSharding(x),
-              )(jax.devices()[0])
-          ]
-        else:
-          out_avals = ()
-          out_shardings = ()
-
-        _, result_future = _profile_state.executable.call(
-            out_avals=out_avals, out_shardings=out_shardings
-        )
-        result_future.result()
+        _profile_state.call_profile_executable()
       finally:
         _profile_state.reset()
   finally:
@@ -306,7 +341,7 @@ async def profiling(pc: ProfilingConfig) -> Mapping[str, str]:
 
   global _profiler_thread
   if _profiler_thread is not None:
-    raise ValueError("Only one profiler server can be active at a time.")
+    raise RuntimeError("Only one profiler server can be active at a time.")
 
   _profiler_thread = threading.Thread(target=server_loop, args=(port,))
   _profiler_thread.start()
@@ -318,7 +353,7 @@ def stop_server() -> None:
   Pathways profiling servers are not stoppable at this time.
   """
   if _profiler_thread is None:
-    raise ValueError("No active profiler server.")
+    raise RuntimeError("No active profiler server.")
 
 
 def collect_profile(

pathwaysutils/test/profiling_test.py

@@ -280,7 +280,7 @@ def test_start_trace_no_toy_computation_second_time(self):
 
   def test_start_trace_while_running_error(self):
     profiling.start_trace("gs://test_bucket/test_dir")
-    with self.assertRaisesRegex(ValueError, "trace is already being taken"):
+    with self.assertRaisesRegex(RuntimeError, "trace is already being taken"):
       profiling.start_trace("gs://test_bucket/test_dir2")
 
   def test_stop_trace_success(self):
@@ -308,22 +308,16 @@ def test_stop_trace_with_xprof_options_passes_out_avals(self):
     options = jax.profiler.ProfileOptions()
     options.duration_ms = 2000
 
-    with mock.patch.object(
-        profiling, "_profile_state", autospec=True
-    ) as mock_profile_state:
-      request = profiling._create_profile_request(
-          "gs://test_bucket/test_dir", options
-      )
-      mock_profile_state.profile_request = request
-      mock_profile_state.executable = (
-          self.mock_plugin_executable_cls.return_value
-      )
-      mock_profile_state.lock = mock.MagicMock()
-      mock_profile_state.lock.locked.return_value = True
-      mock_profile_state.lock.__enter__.return_value = None
-      mock_profile_state.lock.__exit__.return_value = None
+    request = profiling._create_profile_request(
+        "gs://test_bucket/test_dir", options
+    )
+    profiling._profile_state.profile_request = request
+    profiling._profile_state.executable = (
+        self.mock_plugin_executable_cls.return_value
+    )
+    self.addCleanup(profiling._profile_state.reset)
 
-      profiling.stop_trace()
+    profiling.stop_trace()
 
     with self.subTest("plugin_executable_called"):
       self.mock_plugin_executable_cls.return_value.call.assert_called_once()
@@ -340,7 +334,7 @@ def test_stop_trace_with_xprof_options_passes_out_avals(self):
 
   def test_stop_trace_before_start_error(self):
     with self.assertRaisesRegex(
-        ValueError, "stop_trace called before a trace is being taken!"
+        RuntimeError, "stop_trace called before a trace is being taken!"
     ):
       profiling.stop_trace()
 
@@ -359,12 +353,12 @@ def test_start_server_twice_raises_error(self):
     )
     profiling.start_server(9000)
     with self.assertRaisesRegex(
-        ValueError, "Only one profiler server can be active"
+        RuntimeError, "Only one profiler server can be active"
     ):
       profiling.start_server(9001)
 
   def test_stop_server_no_server_raises_error(self):
-    with self.assertRaisesRegex(ValueError, "No active profiler server"):
+    with self.assertRaisesRegex(RuntimeError, "No active profiler server"):
       profiling.stop_server()
 
   def test_stop_server_does_nothing_if_server_exists(self):
@@ -508,6 +502,7 @@ def test_create_profile_request_with_options(self):
     options.python_tracer_level = 1
     options.duration_ms = 2000
     options.start_timestamp_ns = 123456789
+    options.session_id = "test_session"
     options.advanced_configuration = {
         "tpu_num_chips_to_profile_per_task": 3,
         "tpu_num_sparse_core_tiles_to_trace": 5,
@@ -525,6 +520,7 @@ def test_create_profile_request_with_options(self):
             "maxNumHosts": 1,
             "xprofTraceOptions": {
                 "traceDirectory": "gs://bucket/dir",
+                "traceSessionName": "test_session",
                 "pwTraceOptions": {
                     "enablePythonTracer": True,
                     "advancedConfiguration": {
@@ -609,6 +605,36 @@ def test_jax_profiler_trace_calls_patched_functions(self):
     mocks["start_trace"].assert_called_once()
     mocks["stop_trace"].assert_called_once()
 
+  @absltest.skipIf(
+      jax.version.__version_info__ < (0, 9, 2),
+      "ProfileOptions requires JAX 0.9.2 or newer",
+  )
+  def test_is_default_profile_options_with_session_id(self):
+    options = jax.profiler.ProfileOptions()
+    options.session_id = "test_session"
+    self.assertFalse(profiling._is_default_profile_options(options))
+
+  @absltest.skipIf(
+      jax.version.__version_info__ < (0, 9, 2),
+      "ProfileOptions requires JAX 0.9.2 or newer",
+  )
+  def test_start_trace_compatibility_error(self):
+    self.mock_plugin_executable_cls.side_effect = RuntimeError(
+        "Bad PluginProgram"
+    )
+
+    options = jax.profiler.ProfileOptions()
+    options.session_id = "test_session"
+
+    with self.assertRaisesRegex(
+        RuntimeError,
+        "likely because the running Pathways server images do not support the"
+        " trace session ID option",
+    ):
+      profiling.start_trace(
+          "gs://test_bucket/test_dir", profiler_options=options
+      )
+
 
 if __name__ == "__main__":
   absltest.main()