Implement stop_and_get_fdo_profile and monkey-patch in pathwaysutils.profiling

pathwaysutils/profiling.py

@@ -94,6 +94,9 @@ def call_profile_executable(self) -> None:
 _profile_state = _ProfileState()
 _original_start_trace = jax.profiler.start_trace
 _original_stop_trace = jax.profiler.stop_trace
+_original_stop_and_get_fdo_profile = getattr(
+    jax._src.profiler, "stop_and_get_fdo_profile", None  # pylint: disable=protected-access
+)
 
 
 def toy_computation() -> None:
@@ -326,6 +329,23 @@ def stop_trace() -> None:
     _original_stop_trace()
 
 
+def stop_and_get_fdo_profile() -> Any:
+  """Stops the currently-running profiler trace and exports fdo_profile."""
+  try:
+    with _profile_state.lock:
+      if _profile_state.executable is not None:
+        try:
+          _profile_state.call_profile_executable()
+        finally:
+          _profile_state.reset()
+  finally:
+    if _original_stop_and_get_fdo_profile is not None:
+      return _original_stop_and_get_fdo_profile()
+    raise NotImplementedError(
+        "stop_and_get_fdo_profile is not supported in this version of JAX."
+    )
+
+
 _profiler_thread: threading.Thread | None = None
 
 
@@ -455,6 +475,17 @@ def stop_trace_patch() -> None:
   jax.profiler.stop_trace = stop_trace_patch
   jax._src.profiler.stop_trace = stop_trace_patch  # pylint: disable=protected-access
 
+  def stop_and_get_fdo_profile_patch() -> Any:
+    _logger.debug("jax._src.profiler.stop_and_get_fdo_profile patched")
+    return stop_and_get_fdo_profile()
+
+  if hasattr(jax._src.profiler, "stop_and_get_fdo_profile"):
+    jax._src.profiler.stop_and_get_fdo_profile = (  # pylint: disable=protected-access
+        stop_and_get_fdo_profile_patch
+    )
+  if hasattr(jax.profiler, "stop_and_get_fdo_profile"):
+    jax.profiler.stop_and_get_fdo_profile = stop_and_get_fdo_profile_patch
+
   def start_server_patch(port: int) -> None:
     _logger.debug(
         "jax.profile.start_server patched with pathways' start_server"

pathwaysutils/test/profiling_test.py

@@ -439,6 +439,7 @@ def _setup_monkey_patch(self):
         (jax.profiler, "stop_server"),
         (jax._src.profiler, "start_trace"),
         (jax._src.profiler, "stop_trace"),
+        (jax._src.profiler, "stop_and_get_fdo_profile"),
     ]
     original_jax_funcs = {}
     for module, func_name in targets:
@@ -461,6 +462,11 @@ def _setup_monkey_patch(self):
         "stop_trace": self.enter_context(
             mock.patch.object(profiling, "stop_trace", autospec=True)
         ),
+        "stop_and_get_fdo_profile": self.enter_context(
+            mock.patch.object(
+                profiling, "stop_and_get_fdo_profile", autospec=True
+            )
+        ),
         "start_server": self.enter_context(
             mock.patch.object(profiling, "start_server", autospec=True)
         ),
@@ -515,6 +521,17 @@ def test_monkey_patched_stop_trace(self, profiler_module):
 
     mocks["stop_trace"].assert_called_once()
 
+  @parameterized.named_parameters(
+      dict(testcase_name="jax_profiler", profiler_module=jax.profiler),
+      dict(testcase_name="jax_src_profiler", profiler_module=jax._src.profiler),
+  )
+  def test_monkey_patched_stop_and_get_fdo_profile(self, profiler_module):
+    mocks = self._setup_monkey_patch()
+
+    if hasattr(profiler_module, "stop_and_get_fdo_profile"):
+      profiler_module.stop_and_get_fdo_profile()
+      mocks["stop_and_get_fdo_profile"].assert_called_once()
+
   def test_monkey_patched_start_server(self):
     mocks = self._setup_monkey_patch()