From 379dfc891f039c9d9bfdef3a2e149375a96cfa4a Mon Sep 17 00:00:00 2001 From: Pathways-on-Cloud Team Date: Mon, 29 Jun 2026 10:51:00 -0700 Subject: [PATCH] Implement stop_and_get_fdo_profile and monkey-patch in pathwaysutils.profiling PiperOrigin-RevId: 939917877 --- pathwaysutils/profiling.py | 31 ++++++++++++++++++++++++++++ pathwaysutils/test/profiling_test.py | 17 +++++++++++++++ 2 files changed, 48 insertions(+) diff --git a/pathwaysutils/profiling.py b/pathwaysutils/profiling.py index 5a54586..1e9db87 100644 --- a/pathwaysutils/profiling.py +++ b/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" diff --git a/pathwaysutils/test/profiling_test.py b/pathwaysutils/test/profiling_test.py index d6d926b..3acf770 100644 --- a/pathwaysutils/test/profiling_test.py +++ b/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()