Implement __eq__ monkey patch for JAX ProfileOptions.

PiperOrigin-RevId: 919254885

Author: lukebaumann

pathwaysutils/profiling.py

@@ -100,20 +100,29 @@ def toy_computation() -> None:
   x.block_until_ready()
 
 
-def _is_default_profile_options(
-    profiler_options: jax.profiler.ProfileOptions,
-) -> bool:
-  if jax.version.__version_info__ < (0, 9, 2):
-    return True
-
-  default_options = jax.profiler.ProfileOptions()
+def _profile_options_eq(self, other: Any) -> bool:
+  if not isinstance(other, jax.profiler.ProfileOptions):
+    return NotImplemented
   return (
-      profiler_options.host_tracer_level == default_options.host_tracer_level
-      and profiler_options.python_tracer_level
-      == default_options.python_tracer_level
-      and profiler_options.duration_ms == default_options.duration_ms
-      and not getattr(profiler_options, "advanced_configuration", None)
-      and not getattr(profiler_options, "session_id", None)
+      getattr(self, "include_dataset_ops", None)
+      == getattr(other, "include_dataset_ops", None)
+      and getattr(self, "host_tracer_level", None)
+      == getattr(other, "host_tracer_level", None)
+      and getattr(self, "python_tracer_level", None)
+      == getattr(other, "python_tracer_level", None)
+      and getattr(self, "enable_hlo_proto", None)
+      == getattr(other, "enable_hlo_proto", None)
+      and getattr(self, "start_timestamp_ns", None)
+      == getattr(other, "start_timestamp_ns", None)
+      and getattr(self, "duration_ms", None)
+      == getattr(other, "duration_ms", None)
+      and getattr(self, "raise_error_on_start_failure", None)
+      == getattr(other, "raise_error_on_start_failure", None)
+      and getattr(self, "advanced_configuration", None)
+      == getattr(other, "advanced_configuration", None)
+      and getattr(self, "repository_path", None)
+      == getattr(other, "repository_path", None)
+      and getattr(self, "session_id", None) == getattr(other, "session_id", None)
   )
 
 
@@ -128,7 +137,11 @@ def _create_profile_request(
       "maxNumHosts": max_num_hosts,
   }
 
-  if profiler_options is None or _is_default_profile_options(profiler_options):
+  if (
+      profiler_options is None
+      or jax.version.__version_info__ < (0, 9, 2)
+      or profiler_options == jax.profiler.ProfileOptions()
+  ):
     return profile_request
 
   advanced_config = None
@@ -188,6 +201,10 @@ def _start_pathways_trace_from_profile_request(
 
   Args:
     profile_request: A mapping containing the profile request options.
+
+  Raises:
+    RuntimeError: If a trace is already active, or if starting the trace fails
+      due to an incompatible Pathways backend version when using a session ID.
   """
   with _profile_state.lock:
     global _first_profile_start
@@ -217,7 +234,7 @@ def _start_pathways_trace_from_profile_request(
               "the request, likely because the running Pathways server images "
               "do not support the trace session ID option. Please ensure you "
               "are running the latest versions of both Pathways server images "
-              "and the pathwaysutils library."
+              "and the pathwaysutils package."
           ) from e
       _logger.exception("Failed to start trace")
       raise
@@ -318,6 +335,9 @@ def start_server(port: int) -> None:
 
   Args:
     port : The port to start the server on.
+
+  Raises:
+    RuntimeError: If a profiler server is already active.
   """
   def server_loop(port: int):
     _logger.debug("Starting JAX profiler server on port %s", port)
@@ -448,3 +468,7 @@ def stop_server_patch() -> None:
     stop_server()
 
   jax.profiler.stop_server = stop_server_patch
+
+
+if hasattr(jax.profiler, "ProfileOptions"):
+  jax.profiler.ProfileOptions.__eq__ = _profile_options_eq

pathwaysutils/test/profiling_test.py

@@ -24,6 +24,17 @@
 import requests
 
 
+def _make_profile_options(**kwargs) -> jax.profiler.ProfileOptions:
+  options = jax.profiler.ProfileOptions()
+  for k, v in kwargs.items():
+    # Confirm that the attribute is one of the ProfileOptions attributes
+    # that is supported by Pathways. This is not a test function so it cannot
+    # use assertHasAttr.
+    assert hasattr(options, k), f"ProfileOptions does not have attribute {k}"
+    setattr(options, k, v)
+  return options
+
+
 class ProfilingTest(parameterized.TestCase):
   """Tests for Pathways on Cloud profiling."""
 
@@ -305,8 +316,7 @@ def test_stop_trace_success(self):
       "ProfileOptions requires JAX 0.9.2 or newer",
   )
   def test_stop_trace_with_xprof_options_passes_out_avals(self):
-    options = jax.profiler.ProfileOptions()
-    options.duration_ms = 2000
+    options = _make_profile_options(duration_ms=2000)
 
     request = profiling._create_profile_request(
         "gs://test_bucket/test_dir", options
@@ -467,7 +477,7 @@ def test_monkey_patched_stop_server(self):
 
     mocks["stop_server"].assert_called_once()
 
-  @parameterized.parameters(None, jax.profiler.ProfileOptions())
+  @parameterized.parameters(None, _make_profile_options())
   def test_create_profile_request_default_options(self, profiler_options):
     request = profiling._create_profile_request(
         "gs://bucket/dir", profiler_options=profiler_options
@@ -497,17 +507,18 @@ def test_create_profile_request_with_max_num_hosts(self):
       "ProfileOptions requires JAX 0.9.2 or newer",
   )
   def test_create_profile_request_with_options(self):
-    options = jax.profiler.ProfileOptions()
-    options.host_tracer_level = 2
-    options.python_tracer_level = 1
-    options.duration_ms = 2000
-    options.start_timestamp_ns = 123456789
-    options.session_id = "test_session"
-    options.advanced_configuration = {
-        "tpu_num_chips_to_profile_per_task": 3,
-        "tpu_num_sparse_core_tiles_to_trace": 5,
-        "tpu_trace_mode": "TRACE_COMPUTE",
-    }
+    options = _make_profile_options(
+        host_tracer_level=2,
+        python_tracer_level=1,
+        duration_ms=2000,
+        start_timestamp_ns=123456789,
+        session_id="test_session",
+        advanced_configuration={
+            "tpu_num_chips_to_profile_per_task": 3,
+            "tpu_num_sparse_core_tiles_to_trace": 5,
+            "tpu_trace_mode": "TRACE_COMPUTE",
+        },
+    )
 
     request = profiling._create_profile_request(
         "gs://bucket/dir", profiler_options=options
@@ -609,10 +620,62 @@ def test_jax_profiler_trace_calls_patched_functions(self):
       jax.version.__version_info__ < (0, 9, 2),
       "ProfileOptions requires JAX 0.9.2 or newer",
   )
-  def test_is_default_profile_options_with_session_id(self):
-    options = jax.profiler.ProfileOptions()
-    options.session_id = "test_session"
-    self.assertFalse(profiling._is_default_profile_options(options))
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="default_equal",
+          options1=_make_profile_options(),
+          options2=_make_profile_options(),
+      ),
+      dict(
+          testcase_name="session_id_equal",
+          options1=_make_profile_options(session_id="test"),
+          options2=_make_profile_options(session_id="test"),
+      ),
+      dict(
+          testcase_name="host_tracer_level_equal",
+          options1=_make_profile_options(host_tracer_level=3),
+          options2=_make_profile_options(host_tracer_level=3),
+      ),
+      dict(
+          testcase_name="advanced_config_equal",
+          options1=_make_profile_options(
+              advanced_configuration={"foo": "bar"}
+          ),
+          options2=_make_profile_options(
+              advanced_configuration={"foo": "bar"}
+          ),
+      ),
+  )
+  def test_profile_options_equal(self, options1, options2):
+    self.assertEqual(options1, options2)
+
+  @absltest.skipIf(
+      jax.version.__version_info__ < (0, 9, 2),
+      "ProfileOptions requires JAX 0.9.2 or newer",
+  )
+  @parameterized.named_parameters(
+      dict(
+          testcase_name="session_id_diff",
+          options1=_make_profile_options(session_id="test1"),
+          options2=_make_profile_options(session_id="test2"),
+      ),
+      dict(
+          testcase_name="host_tracer_level_diff",
+          options1=_make_profile_options(host_tracer_level=1),
+          options2=_make_profile_options(host_tracer_level=2),
+      ),
+      dict(
+          testcase_name="advanced_config_diff",
+          options1=_make_profile_options(
+              advanced_configuration={"foo": "bar"}
+          ),
+          options2=_make_profile_options(
+              advanced_configuration={"foo": "baz"}
+          ),
+      ),
+  )
+  def test_profile_options_not_equal(self, options1, options2):
+    self.assertNotEqual(options1, options2)
 
   @absltest.skipIf(
       jax.version.__version_info__ < (0, 9, 2),
@@ -623,8 +686,7 @@ def test_start_trace_compatibility_error(self):
         "Bad PluginProgram"
     )
 
-    options = jax.profiler.ProfileOptions()
-    options.session_id = "test_session"
+    options = _make_profile_options(session_id="test_session")
 
     with self.assertRaisesRegex(
         RuntimeError,