diff --git a/pathwaysutils/_initialize.py b/pathwaysutils/_initialize.py index 7a133eb..4b6c65c 100644 --- a/pathwaysutils/_initialize.py +++ b/pathwaysutils/_initialize.py @@ -93,7 +93,8 @@ def initialize() -> None: if is_pathways_backend_used(): _logger.debug("Detected Pathways-on-Cloud backend. Applying changes.") proxy_backend.register_backend_factory() - profiling.monkey_patch_jax() + if not profiling.is_native_profiling_supported(): + profiling.monkey_patch_jax() if _is_persistence_enabled(): type_handlers.register_type_handler( jax.Array, diff --git a/pathwaysutils/experimental/profiling.py b/pathwaysutils/experimental/profiling.py index 0816302..33b80bf 100644 --- a/pathwaysutils/experimental/profiling.py +++ b/pathwaysutils/experimental/profiling.py @@ -14,6 +14,7 @@ """Experimental profiling utilites.""" from collections.abc import Mapping +import json from typing import Any from pathwaysutils import profiling @@ -54,6 +55,23 @@ def start_trace( f" {log_dir}" ) + if profiling.is_native_profiling_supported(): + # Let JAX send the profile request to the Pathways backend. + profiling._original_start_trace( # pylint: disable=protected-access + log_dir=log_dir, + create_perfetto_link=create_perfetto_link, + create_perfetto_trace=create_perfetto_trace, + profiler_options=profiling._with_advanced_configuration( # pylint: disable=protected-access + None, + { + profiling._PROFILE_REQUEST_KEY: json.dumps( # pylint: disable=protected-access + {"profileRequest": profile_request} + ) + }, + ), + ) + return + profiling._start_pathways_trace_from_profile_request(profile_request) # pylint: disable=protected-access profiling._original_start_trace( # pylint: disable=protected-access diff --git a/pathwaysutils/profiling.py b/pathwaysutils/profiling.py index 1f773e3..c8161dd 100644 --- a/pathwaysutils/profiling.py +++ b/pathwaysutils/profiling.py @@ -16,6 +16,7 @@ import asyncio from collections.abc import Iterator, Mapping import contextlib +import copy import dataclasses import datetime import json @@ -36,6 +37,11 @@ _logger = logging.getLogger(__name__) +# Keys of `jax.profiler.ProfileOptions.advanced_configuration` that configure +# how JAX traces Pathways backends. +_MAX_NUM_HOSTS_KEY = "pathways_max_num_hosts" +_PROFILE_REQUEST_KEY = "pathways_profile_request" + class _ProfileState: """Holds the state of an ongoing profiling session. @@ -262,6 +268,45 @@ def _validate_gcs_bucket(log_dir: str) -> None: ) +def is_native_profiling_supported() -> bool: + """Returns whether JAX natively supports tracing Pathways backends. + + If so, JAX traces the Pathways backend of its `proxy` backend in every + profiling session, so the `jax.profiler` APIs do not need to be monkey + patched. + """ + try: + from jax._src.lib import ifrt_version # pylint: disable=g-import-not-at-top + from jax._src.lib import jaxlib_extension_version # pylint: disable=g-import-not-at-top + except ImportError: # Older JAX versions. + return False + # jaxlib registers the profiler for IFRT proxy clients, and JAX tells it where + # sessions write their profiles. + return jaxlib_extension_version >= 515 and ifrt_version >= 83 + + +def _with_advanced_configuration( + profiler_options: jax.profiler.ProfileOptions | None, + advanced_configuration: Mapping[str, Any], +) -> jax.profiler.ProfileOptions: + """Returns a copy of `profiler_options` with additional configuration. + + Args: + profiler_options: The options to copy, if any. + advanced_configuration: Advanced configuration to add, unless + `profiler_options` already sets the respective keys. + """ + if profiler_options is None: + options = jax.profiler.ProfileOptions() + else: + options = copy.copy(profiler_options) + options.advanced_configuration = { + **advanced_configuration, + **options.advanced_configuration, + } + return options + + def start_trace( log_dir: os.PathLike[str] | str, *, @@ -319,6 +364,17 @@ def start_trace( "features for Pathways on Cloud and may not be fully supported." ) + if is_native_profiling_supported(): + # JAX traces the Pathways backend itself. + return _original_start_trace( + log_dir=log_dir, + create_perfetto_link=create_perfetto_link, + create_perfetto_trace=create_perfetto_trace, + profiler_options=_with_advanced_configuration( + profiler_options, {_MAX_NUM_HOSTS_KEY: max_num_hosts} + ), + ) + if jax.version.__version_info__ < (0, 9, 2): if profiler_options is not None: _logger.warning( @@ -376,6 +432,9 @@ def stop_trace() -> Any: `jax.profiler.stop_trace` fails to stop the trace. The JAX trace is stopped even if stopping the Pathways trace fails. """ + if is_native_profiling_supported(): + return _original_stop_trace() + try: with _profile_state.lock: if _profile_state.executable is None: @@ -591,6 +650,9 @@ def collect_profile( def monkey_patch_jax() -> None: """Monkey patches JAX with Pathways versions of functions. + This is not needed if `is_native_profiling_supported()`, in which case JAX + traces the Pathways backend itself. + The signatures in patched functions should match the original. Patched functions are: diff --git a/pathwaysutils/test/initialize_test.py b/pathwaysutils/test/initialize_test.py index e2e33f1..52e62fc 100644 --- a/pathwaysutils/test/initialize_test.py +++ b/pathwaysutils/test/initialize_test.py @@ -13,6 +13,7 @@ # limitations under the License. import os +from unittest import mock from absl.testing import absltest from absl.testing import parameterized @@ -50,6 +51,32 @@ def test_first_initialize(self): "Detected Pathways-on-Cloud backend. Applying changes.", logs.output[1] ) + @parameterized.named_parameters( + ("native_profiling_supported", True, False), + ("native_profiling_unsupported", False, True), + ) + def test_initialize_monkey_patches_jax_only_without_native_profiling( + self, native_profiling_supported, expect_monkey_patch + ): + jax.config.update("jax_platforms", "proxy") + _initialize._initialization_count = 0 + self.enter_context( + mock.patch.object( + _initialize.profiling, + "is_native_profiling_supported", + return_value=native_profiling_supported, + ) + ) + mock_monkey_patch_jax = self.enter_context( + mock.patch.object( + _initialize.profiling, "monkey_patch_jax", autospec=True + ) + ) + + _initialize.initialize() + + self.assertEqual(mock_monkey_patch_jax.called, expect_monkey_patch) + @parameterized.named_parameters( ("initialization_count 1", 1), ("initialization_count 2", 2), diff --git a/pathwaysutils/test/profiling_test.py b/pathwaysutils/test/profiling_test.py index dd6a10f..afe71a2 100644 --- a/pathwaysutils/test/profiling_test.py +++ b/pathwaysutils/test/profiling_test.py @@ -33,6 +33,13 @@ class ProfilingTest(parameterized.TestCase): def setUp(self): super().setUp() + # These tests cover the profiling that pathwaysutils drives itself, i.e., + # without JAX tracing the Pathways backend. + self.enter_context( + mock.patch.object( + profiling, "is_native_profiling_supported", return_value=False + ) + ) self.mock_post = self.enter_context( mock.patch.object(requests, "post", autospec=True) ) @@ -1160,6 +1167,117 @@ def test_profiling_server_passes_duration_ms_to_start_trace(self): self.assertEqual(call_kwargs["profiler_options"].duration_ms, 2500) +class IsNativeProfilingSupportedTest(parameterized.TestCase): + + @parameterized.named_parameters( + ("supported", 515, 83, True), + ("old_jaxlib", 514, 83, False), + ("old_xla", 515, 82, False), + ) + def test_is_native_profiling_supported( + self, jaxlib_extension_version, ifrt_version, expected + ): + with mock.patch.multiple( + jax._src.lib, # pylint: disable=protected-access + jaxlib_extension_version=jaxlib_extension_version, + ifrt_version=ifrt_version, + ): + self.assertEqual(profiling.is_native_profiling_supported(), expected) + + def test_unsupported_without_jaxlib_extension_version(self): + self.enter_context(mock.patch.dict(jax._src.lib.__dict__)) # pylint: disable=protected-access + jax._src.lib.__dict__.pop("jaxlib_extension_version", None) # pylint: disable=protected-access + + self.assertFalse(profiling.is_native_profiling_supported()) + + +class NativeProfilingTest(parameterized.TestCase): + """Tests for Pathways on Cloud profiling when JAX traces Pathways.""" + + def setUp(self): + super().setUp() + self.enter_context( + mock.patch.object( + profiling, "is_native_profiling_supported", return_value=True + ) + ) + self.mock_plugin_executable_cls = self.enter_context( + mock.patch.object( + profiling.plugin_executable, "PluginExecutable", autospec=True + ) + ) + self.mock_original_start_trace = self.enter_context( + mock.patch.object(profiling, "_original_start_trace", autospec=True) + ) + self.mock_original_stop_trace = self.enter_context( + mock.patch.object(profiling, "_original_stop_trace", autospec=True) + ) + + def test_start_trace_uses_jax_profiler(self): + result = profiling.start_trace("gs://test_bucket/test_dir") + + self.assertIs(result, self.mock_original_start_trace.return_value) + self.mock_plugin_executable_cls.assert_not_called() + self.mock_original_start_trace.assert_called_once() + call_args = self.mock_original_start_trace.call_args[1] + self.assertEqual(call_args["log_dir"], "gs://test_bucket/test_dir") + self.assertFalse(call_args["create_perfetto_link"]) + self.assertFalse(call_args["create_perfetto_trace"]) + self.assertEqual( + call_args["profiler_options"].advanced_configuration, + {"pathways_max_num_hosts": 1}, + ) + + def test_start_trace_with_max_num_hosts(self): + profiling.start_trace("gs://test_bucket/test_dir", max_num_hosts=4) + + call_args = self.mock_original_start_trace.call_args[1] + self.assertEqual( + call_args["profiler_options"].advanced_configuration, + {"pathways_max_num_hosts": 4}, + ) + + @absltest.skipIf( + not profiling.is_native_profiling_supported(), + "Requires a JAX version that supports native Pathways profiling, whose" + " ProfileOptions can be copied.", + ) + def test_start_trace_with_profiler_options(self): + options = jax.profiler.ProfileOptions() + options.python_tracer_level = 0 + options.advanced_configuration = { + "pathways_max_num_hosts": 8, + "tpu_trace_mode": "TRACE_ONLY_HOST", + } + + profiling.start_trace( + "gs://test_bucket/test_dir", profiler_options=options, max_num_hosts=4 + ) + + call_args = self.mock_original_start_trace.call_args[1] + session_options = call_args["profiler_options"] + self.assertIsNot(session_options, options) + self.assertEqual(session_options.python_tracer_level, 0) + self.assertEqual( + session_options.advanced_configuration, + {"pathways_max_num_hosts": 8, "tpu_trace_mode": "TRACE_ONLY_HOST"}, + ) + + @parameterized.parameters("/logs/test_log_dir", "cns://test_bucket/test_dir") + def test_start_trace_log_dir_error(self, log_dir): + with self.assertRaisesRegex(ValueError, "Path must be a GCS path"): + profiling.start_trace(log_dir) + + self.mock_original_start_trace.assert_not_called() + + def test_stop_trace_uses_jax_profiler(self): + result = profiling.stop_trace() + + self.assertIs(result, self.mock_original_stop_trace.return_value) + self.mock_original_stop_trace.assert_called_once() + self.mock_plugin_executable_cls.assert_not_called() + + if __name__ == "__main__": absltest.main()