Enable Pathways profiling with jax.profiler.ProfileOptions.

This change allows users to configure Pathways profiling by passing a jax.profiler.ProfileOptions object to the start_trace function. The options are translated into the Pathways profile request, enabling control over a subset of parameters.
Explicitly, `start_timestamp_ms`, `duration_ms`, `host_tracer_level`, `advanced_configuration`, and `python_tracer_level`.

PiperOrigin-RevId: 885730249

Author: lukebaumann

pathwaysutils/profiling.py

@@ -59,10 +59,55 @@ def toy_computation() -> None:
 
 def _create_profile_request(
     log_dir: os.PathLike[str] | str,
+    profiler_options: jax.profiler.ProfileOptions | None = None,
 ) -> Mapping[str, Any]:
   """Creates a profile request mapping from the given options."""
-  profile_request = {}
-  profile_request["traceLocation"] = str(log_dir)
+  if profiler_options is None:
+    profiler_options = jax.profiler.ProfileOptions()
+
+  advanced_config = None
+  if profiler_options.advanced_configuration:
+    advanced_config = {}
+    for k, v in profiler_options.advanced_configuration.items():
+      # Convert python dict to tensorflow.ProfileOptions.AdvancedConfigValue
+      # json-compatible dict
+      if isinstance(v, bool):
+        advanced_config[k] = {"boolValue": v}
+      elif isinstance(v, int):
+        advanced_config[k] = {"intValue": v}
+      elif isinstance(v, str):
+        advanced_config[k] = {"stringValue": v}
+      else:
+        raise ValueError(
+            f"Unsupported advanced configuration value type: {type(v)}. "
+            "Supported types are bool, int, and str."
+        )
+
+  xprof_options: dict[str, Any] = {
+      "traceDirectory": str(log_dir),
+  }
+
+  if profiler_options.host_tracer_level != 2:
+    xprof_options["hostTraceLevel"] = profiler_options.host_tracer_level
+
+  pw_trace_opts: dict[str, Any] = {}
+  if profiler_options.python_tracer_level:
+    pw_trace_opts["enablePythonTracer"] = bool(
+        profiler_options.python_tracer_level
+    )
+
+  if advanced_config:
+    pw_trace_opts["advancedConfiguration"] = advanced_config
+
+  if pw_trace_opts:
+    xprof_options["pwTraceOptions"] = pw_trace_opts
+
+  profile_request: dict[str, Any] = {
+      "xprofTraceOptions": xprof_options,
+  }
+
+  if profiler_options.duration_ms > 0:
+    profile_request["maxDurationSecs"] = profiler_options.duration_ms / 1000.0
 
   return profile_request
 
@@ -104,7 +149,7 @@ def start_trace(
     *,
     create_perfetto_link: bool = False,
     create_perfetto_trace: bool = False,
-    profiler_options: jax.profiler.ProfileOptions | None = None,  # pylint: disable=unused-argument
+    profiler_options: jax.profiler.ProfileOptions | None = None,
 ) -> None:
   """Starts a profiler trace.
 
@@ -133,7 +178,6 @@ def start_trace(
       This feature is experimental for Pathways on Cloud and may not be fully
       supported.
     profiler_options: Profiler options to configure the profiler for collection.
-      Options are not currently supported and ignored.
   """
   if not str(log_dir).startswith("gs://"):
     raise ValueError(f"log_dir must be a GCS bucket path, got {log_dir}")
@@ -144,7 +188,18 @@ def start_trace(
         "features for Pathways on Cloud and may not be fully supported."
     )
 
-  _start_pathways_trace_from_profile_request(_create_profile_request(log_dir))
+  if jax.version.__version_info__ < (0, 8, 2) and profiler_options is not None:
+    _logger.warning(
+        "ProfileOptions are not supported until JAX 0.8.2 and will be omitted. "
+        "Some options can be specified via command line flags."
+    )
+    profiler_options = None
+
+  profile_request = _create_profile_request(log_dir, profiler_options)
+
+  _logger.debug("Profile request: %s", profile_request)
+
+  _start_pathways_trace_from_profile_request(profile_request)
 
   _original_start_trace(
       log_dir=log_dir,

pathwaysutils/test/profiling_test.py

@@ -225,9 +225,18 @@ def test_start_trace_success(self):
 
     self.mock_toy_computation.assert_called_once()
     self.mock_plugin_executable_cls.assert_called_once_with(
-        json.dumps(
-            {"profileRequest": {"traceLocation": "gs://test_bucket/test_dir"}}
-        )
+        json.dumps({
+            "profileRequest": {
+                "traceLocation": "gs://test_bucket/test_dir",
+                "profilingStartTimeNs": 0,
+                "profilingDurationMs": 0,
+                "hostTraceLevel": 2,
+                "pwTraceOptions": {
+                    "advancedConfiguration": {},
+                    "enablePythonTracer": True,
+                },
+            }
+        })
     )
     self.mock_plugin_executable_cls.return_value.call.assert_called_once()
     self.mock_original_start_trace.assert_called_once_with(
@@ -393,7 +402,71 @@ def test_monkey_patched_stop_server(self):
 
   def test_create_profile_request_no_options(self):
     request = profiling._create_profile_request("gs://bucket/dir")
-    self.assertEqual(request, {"traceLocation": "gs://bucket/dir"})
+    self.assertEqual(
+        request,
+        {
+            "traceLocation": "gs://bucket/dir",
+            "profilingStartTimeNs": 0,
+            "profilingDurationMs": 0,
+            "hostTraceLevel": 2,
+            "pwTraceOptions": {
+                "advancedConfiguration": {},
+                "enablePythonTracer": True,
+            },
+        },
+    )
+
+  def test_create_profile_request_default_options(self):
+    options = jax.profiler.ProfileOptions()
+    request = profiling._create_profile_request(
+        "gs://bucket/dir", profiler_options=options
+    )
+    self.assertEqual(
+        request,
+        {
+            "traceLocation": "gs://bucket/dir",
+            "profilingStartTimeNs": 0,
+            "profilingDurationMs": 0,
+            "hostTraceLevel": 2,
+            "pwTraceOptions": {
+                "advancedConfiguration": {},
+                "enablePythonTracer": True,
+            },
+        },
+    )
+
+  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.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
+    )
+    self.assertEqual(
+        request,
+        {
+            "traceLocation": "gs://bucket/dir",
+            "hostTraceLevel": 2,
+            "profilingDurationMs": 2000,
+            "profilingStartTimeNs": 123456789,
+            "pwTraceOptions": {
+                "enablePythonTracer": True,
+                "advancedConfiguration": {
+                    "tpu_num_chips_to_profile_per_task": 3,
+                    "tpu_num_sparse_core_tiles_to_trace": 5,
+                    "tpu_trace_mode": "TRACE_COMPUTE",
+                },
+            },
+        },
+    )
 
   @parameterized.parameters(
       ({"traceLocation": "gs://test_bucket/test_dir"},),