Fix JaxRuntimeError during profiler stop_trace array verification

When executing a profiler request with Xprof Trace Options, the IFRT proxy outputs a single JAX shaped array. Previously, the Python profiler expected 0 outputs, causing a JaxRuntimeError (`Mismatch between out_handlers and num_results: 0 vs 1`).

This CL updates the profiler state implementation to correctly expect `(1,)` output from `PluginExecutable.call()` when stopping a trace with Xprof options, consuming the URL suffix parameter and fixing the crash.

PiperOrigin-RevId: 888268746

Author: lukebaumann

pathwaysutils/profiling.py

@@ -35,14 +35,17 @@
 
 class _ProfileState:
   executable: plugin_executable.PluginExecutable | None = None
+  profile_request: Mapping[str, Any] | None = None
   lock: threading.Lock
 
   def __init__(self) -> None:
     self.executable = None
+    self.profile_request = None
     self.lock = threading.Lock()
 
   def reset(self) -> None:
     self.executable = None
+    self.profile_request = None
 
 
 _first_profile_start = True
@@ -153,6 +156,7 @@ def _start_pathways_trace_from_profile_request(
     _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()
@@ -233,8 +237,19 @@ def stop_trace() -> None:
       if _profile_state.executable is None:
         raise ValueError("stop_trace called before a trace is being taken!")
       try:
-        _, result_future = _profile_state.executable.call()
-        result_future.result()
+        if (
+            _profile_state.profile_request
+            and "xprofTraceOptions" in _profile_state.profile_request
+        ):
+          out_avals = [jax.core.ShapedArray((1,), jnp.object_)]
+          out_shardings = [jax.sharding.SingleDeviceSharding(jax.devices()[0])]
+        else:
+          out_avals = ()
+          out_shardings = ()
+
+        _profile_state.executable.call(
+            out_avals=out_avals, out_shardings=out_shardings
+        )
       finally:
         _profile_state.reset()
   finally:

pathwaysutils/test/profiling_test.py

@@ -213,10 +213,7 @@ def test_lock_released_on_stop_failure(self):
     """Tests that the lock is released if stop_trace fails."""
     profiling.start_trace("gs://test_bucket/test_dir3")
     self.assertFalse(profiling._profile_state.lock.locked())
-    mock_result = (
-        self.mock_plugin_executable_cls.return_value.call.return_value[1]
-    )
-    mock_result.result.side_effect = RuntimeError("stop failed")
+    self.mock_plugin_executable_cls.return_value.call.side_effect = RuntimeError("stop failed")
     with self.assertRaisesRegex(RuntimeError, "stop failed"):
       profiling.stop_trace()
     self.assertFalse(profiling._profile_state.lock.locked())
@@ -277,6 +274,28 @@ def test_stop_trace_success(self):
     with self.subTest("executable_is_none"):
       self.assertIsNone(profiling._profile_state.executable)
 
+  @mock.patch("jax.version.__version_info__", (0, 9, 2))
+  def test_stop_trace_with_xprof_options_passes_out_avals(self):
+    from jax import numpy as jnp
+    options = jax.profiler.ProfileOptions()
+    options.duration_ms = 2000
+    
+    # Bypass start_trace and explicitly populate profile state
+    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
+    
+    profiling.stop_trace()
+    
+    self.mock_plugin_executable_cls.return_value.call.assert_called_once()
+    _, kwargs = self.mock_plugin_executable_cls.return_value.call.call_args
+    self.assertIn("out_avals", kwargs)
+    self.assertIn("out_shardings", kwargs)
+    self.assertLen(kwargs["out_avals"], 1)
+    # Check that it's an object dtype ShapedArray
+    self.assertEqual(kwargs["out_avals"][0].shape, (1,))
+    self.assertEqual(kwargs["out_avals"][0].dtype, jnp.object_)
+
   def test_stop_trace_before_start_error(self):
     with self.assertRaisesRegex(
         ValueError, "stop_trace called before a trace is being taken!"
@@ -496,10 +515,7 @@ def test_original_stop_trace_called_on_stop_failure(self):
     """Tests that original_stop_trace is called if pathways stop_trace fails."""
     profiling.start_trace("gs://test_bucket/test_dir")
     self.assertFalse(profiling._profile_state.lock.locked())
-    mock_result = (
-        self.mock_plugin_executable_cls.return_value.call.return_value[1]
-    )
-    mock_result.result.side_effect = RuntimeError("stop failed")
+    self.mock_plugin_executable_cls.return_value.call.side_effect = RuntimeError("stop failed")
     with self.assertRaisesRegex(RuntimeError, "stop failed"):
       profiling.stop_trace()
     self.mock_original_stop_trace.assert_called_once()