Use JAX's Pathways profiler instead of monkey patching jax.profiler.

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,

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

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:

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),

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()