Patch internal JAX profiler functions (enabling jax.profiler.trace) and add a test for jax.profiler.trace.

The `jax.profiler.trace` context manager uses internal `jax._src.profiler` functions. This change ensures that these internal functions are also patched by `pathwaysutils.profiling.monkey_patch_jax` to correctly intercept profiling calls. A new test is added to verify that `with jax.profiler.trace(...)` now triggers the patched Pathways profiling functions.

PiperOrigin-RevId: 845135750

Author: lukebaumann

pathwaysutils/profiling.py

@@ -104,6 +104,7 @@ def start_trace(
     *,
     create_perfetto_link: bool = False,
     create_perfetto_trace: bool = False,
+    profiler_options: jax.profiler.ProfileOptions | None = None,  # pylint: disable=unused-argument
 ) -> None:
   """Starts a profiler trace.
 
@@ -131,6 +132,8 @@ def start_trace(
       want to generate a Perfetto-compatible trace without blocking the process.
       This feature is experimental for Pathways on Cloud and may not be fully
       supported.
+    profiler_options: Profiler options to configure the profiler for collection.
+      Options are not currently supported and ignored.
   """
   if not str(log_dir).startswith("gs://"):
     raise ValueError(f"log_dir must be a GCS bucket path, got {log_dir}")
@@ -270,19 +273,27 @@ def monkey_patch_jax():
 
   def start_trace_patch(
       log_dir,
-      create_perfetto_link: bool = False,  # pylint: disable=unused-argument
-      create_perfetto_trace: bool = False,  # pylint: disable=unused-argument
+      create_perfetto_link: bool = False,
+      create_perfetto_trace: bool = False,
+      profiler_options: jax.profiler.ProfileOptions | None = None,  # pylint: disable=unused-argument
   ) -> None:
     _logger.debug("jax.profile.start_trace patched with pathways' start_trace")
-    return start_trace(log_dir)
+    return start_trace(
+        log_dir,
+        create_perfetto_link=create_perfetto_link,
+        create_perfetto_trace=create_perfetto_trace,
+        profiler_options=profiler_options,
+    )
 
   jax.profiler.start_trace = start_trace_patch
+  jax._src.profiler.start_trace = start_trace_patch  # pylint: disable=protected-access
 
   def stop_trace_patch() -> None:
     _logger.debug("jax.profile.stop_trace patched with pathways' stop_trace")
     return stop_trace()
 
   jax.profiler.stop_trace = stop_trace_patch
+  jax._src.profiler.stop_trace = stop_trace_patch  # pylint: disable=protected-access
 
   def start_server_patch(port: int):
     _logger.debug(

pathwaysutils/test/profiling_test.py

@@ -310,6 +310,18 @@ def test_monkey_patch_jax(self):
     original_jax_stop_trace = jax.profiler.stop_trace
     original_jax_start_server = jax.profiler.start_server
     original_jax_stop_server = jax.profiler.stop_server
+    self.addCleanup(
+        setattr, jax.profiler, "start_trace", original_jax_start_trace
+    )
+    self.addCleanup(
+        setattr, jax.profiler, "stop_trace", original_jax_stop_trace
+    )
+    self.addCleanup(
+        setattr, jax.profiler, "start_server", original_jax_start_server
+    )
+    self.addCleanup(
+        setattr, jax.profiler, "stop_server", original_jax_stop_server
+    )
 
     profiling.monkey_patch_jax()
 
@@ -322,7 +334,12 @@ def test_monkey_patch_jax(self):
         profiling, "start_trace", autospec=True
     ) as mock_pw_start_trace:
       jax.profiler.start_trace("gs://bucket/dir")
-      mock_pw_start_trace.assert_called_once_with("gs://bucket/dir")
+      mock_pw_start_trace.assert_called_once_with(
+          "gs://bucket/dir",
+          create_perfetto_link=False,
+          create_perfetto_trace=False,
+          profiler_options=None,
+      )
 
     with mock.patch.object(
         profiling, "stop_trace", autospec=True
@@ -342,12 +359,6 @@ def test_monkey_patch_jax(self):
       jax.profiler.stop_server()
       mock_pw_stop_server.assert_called_once()
 
-    # Restore original jax functions
-    jax.profiler.start_trace = original_jax_start_trace
-    jax.profiler.stop_trace = original_jax_stop_trace
-    jax.profiler.start_server = original_jax_start_server
-    jax.profiler.stop_server = original_jax_stop_server
-
   def test_create_profile_request_no_options(self):
     request = profiling._create_profile_request("gs://bucket/dir")
     self.assertEqual(request, {"traceLocation": "gs://bucket/dir"})
@@ -384,6 +395,7 @@ def test_create_profile_request_no_options(self):
           },
       },),
   )
+
   def test_start_pathways_trace_from_profile_request(self, profile_request):
     profiling._start_pathways_trace_from_profile_request(profile_request)
 
@@ -407,6 +419,29 @@ def test_original_stop_trace_called_on_stop_failure(self):
       profiling.stop_trace()
     self.mock_original_stop_trace.assert_called_once()
 
+  def test_jax_profiler_trace_calls_patched_functions(self):
+    original_jax_start_trace = jax.profiler.start_trace
+    original_jax_stop_trace = jax.profiler.stop_trace
+    self.addCleanup(
+        setattr, jax.profiler, "start_trace", original_jax_start_trace
+    )
+    self.addCleanup(
+        setattr, jax.profiler, "stop_trace", original_jax_stop_trace
+    )
+    mock_pw_start_trace = self.enter_context(
+        mock.patch.object(profiling, "start_trace", autospec=True)
+    )
+    mock_pw_stop_trace = self.enter_context(
+        mock.patch.object(profiling, "stop_trace", autospec=True)
+    )
+    profiling.monkey_patch_jax()
+
+    with jax.profiler.trace("gs://bucket/dir"):
+      pass
+
+    mock_pw_start_trace.assert_called_once()
+    mock_pw_stop_trace.assert_called_once()
+
 
 if __name__ == "__main__":
   absltest.main()