Add JAX version metadata to xspace

PiperOrigin-RevId: 840138208

Author: Google-ML-Automation

jax/_src/lib/_profiler.pyi

@@ -0,0 +1,43 @@
+from typing import Any
+
+class ProfilerServer: ...
+def start_server(port: int) -> ProfilerServer: ...
+
+def register_plugin_profiler(c_api: Any) -> None: ...
+
+def get_profiled_instructions_proto(tensorboard_dir: str) -> bytes: ...
+def get_instructions_profile(tensorboard_dir: str) -> list[tuple[str, float]]: ...
+def get_fdo_profile(
+    xspace: bytes, as_textproto: bool = ...
+) -> bytes | str: ...
+
+class ProfilerSession:
+  def __init__(self, options: ProfileOptions | None = ...) -> None: ...
+  def stop(self) -> bytes: ...
+  def stop_and_export(self, tensorboard_dir: str) -> None: ...
+  def export(self, xspace: bytes, tensorboard_dir: str) -> None:...
+
+class ProfileOptions:
+  include_dataset_ops: bool
+  host_tracer_level: int
+  python_tracer_level: int
+  enable_hlo_proto: bool
+  start_timestamp_ns: int
+  duration_ms: int
+  repository_path: str
+  raise_error_on_start_failure: bool
+  advanced_configuration: dict[str, Any]
+  session_id: str
+
+def aggregate_profiled_instructions(profiles: list[bytes], percentile: int) -> str: ...
+
+class TraceMe:
+  def __init__(self, name: str, **kwargs: Any) -> None: ...
+  def __enter__(self) -> TraceMe: ...
+  def __exit__(self, exc_type, exc_val, exc_tb) -> None: ...
+  def set_metadata(self, **kwargs: Any) -> None: ...
+  @staticmethod
+  def is_enabled() -> bool: ...
+
+def set_metadata(key: str, value: str) -> None: ...
+def clear_metadata() -> None: ...

jax/_src/profiler.py

@@ -138,14 +138,23 @@ def start_trace(
     # session. Otherwise on Cloud TPU, libtpu may not be initialized before
     # creating the tracer, which will cause the TPU tracer initialization to
     # fail and no TPU operations will be included in the profile.
-    xla_bridge.get_backend()
-
-    if profiler_options is None:
-      _profile_state.profile_session = _profiler.ProfilerSession()
-    else:
-      _profile_state.profile_session = _profiler.ProfilerSession(
-          profiler_options
-      )
+    client = xla_bridge.get_backend()
+
+    options = profiler_options
+    if options is None:
+      options = ProfileOptions()
+    _profiler.clear_metadata()  # pytype: disable=module-attr
+    try:
+      import jax  # type: ignore
+      _profiler.set_metadata("jax_version", jax.__version__)  # pytype: disable=module-attr
+      jaxlib_version = jax.lib.version.__version__
+      if client.platform == "tpu":
+        jaxlib_version += f" ({client.platform_version})"
+      _profiler.set_metadata("jaxlib_version", jaxlib_version)  # pytype: disable=module-attr
+    except (ImportError, AttributeError):
+      _profiler.set_metadata("jax_version", "unknown")  # pytype: disable=module-attr
+      _profiler.set_metadata("jaxlib_version", "unknown")  # pytype: disable=module-attr
+    _profile_state.profile_session = _profiler.ProfilerSession(options)
     _profile_state.create_perfetto_link = create_perfetto_link
     _profile_state.create_perfetto_trace = (
         create_perfetto_trace or create_perfetto_link)
@@ -226,6 +235,7 @@ def stop_trace():
       if _profile_state.create_perfetto_link:
         _host_perfetto_trace_file(abs_filename)
     _profile_state.reset()
+    _profiler.clear_metadata()  # pytype: disable=module-attr
 
 
 def stop_and_get_fdo_profile() -> bytes | str:
@@ -240,6 +250,7 @@ def stop_and_get_fdo_profile() -> bytes | str:
     xspace = _profile_state.profile_session.stop()
     fdo_profile = _profiler.get_fdo_profile(xspace)
     _profile_state.reset()
+    _profiler.clear_metadata()  # pytype: disable=module-attr
     return fdo_profile
 
 

tests/profiler_test.py

@@ -437,6 +437,27 @@ def on_profile(port, logdir, worker_start):
     thread_worker.join(120)
     self._check_xspace_pb_exist(logdir)
 
+  def testDeviceVersionSavedToMetadata(self):
+    with tempfile.TemporaryDirectory() as tmpdir_string:
+      tmpdir = pathlib.Path(tmpdir_string)
+      with jax.profiler.trace(tmpdir):
+        jax.pmap(lambda x: jax.lax.psum(x + 1, 'i'), axis_name='i')(
+            jnp.ones(jax.local_device_count()))
+
+      proto_path = tuple(tmpdir.rglob("*.xplane.pb"))
+      self.assertEqual(len(proto_path), 1)
+      (proto_file,) = proto_path
+      proto = proto_file.read_bytes()
+
+      self.assertIn(b"jax_version", proto)
+      self.assertIn(b"jaxlib_version", proto)
+      if jtu.test_device_matches(["tpu"]):
+        self.assertIn(b"libtpu_version", proto)
+      if jtu.test_device_matches(["gpu"]):
+        self.assertIn(b"cuda_version", proto)
+        self.assertIn(b"cuda_runtime_version", proto)
+        self.assertIn(b"cuda_driver_version", proto)
+
   @unittest.skipIf(
       not (portpicker and _pywrap_profiler_plugin),
     "Test requires xprof and portpicker")