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

@@ -35,6 +35,12 @@
 from jax._src.lib import _profiler
 from jax._src.lib import _profile_data
 
+# TODO: remove messy fallback for set_metadata / clear_metadata
+if not hasattr(_profiler, "set_metadata"):
+  _profiler.set_metadata = lambda key, value: None
+if not hasattr(_profiler, "clear_metadata"):
+  _profiler.clear_metadata = lambda: None
+
 ProfileData = _profile_data.ProfileData
 ProfileEvent = _profile_data.ProfileEvent
 ProfilePlane = _profile_data.ProfilePlane
@@ -138,14 +144,24 @@ 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()
+    try:
+      from jax._src.lib import version as version_lib
+      version = ".".join(map(str, version_lib))
+      _profiler.set_metadata("jax_version", version)
+      jaxlib_version = version
+      if client.platform == "tpu":
+        jaxlib_version += f" ({client.platform_version})"
+      _profiler.set_metadata("jaxlib_version", jaxlib_version)
+    except (ImportError, AttributeError):
+      _profiler.set_metadata("jax_version", "unknown")
+      _profiler.set_metadata("jaxlib_version", "unknown")
+    _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 +242,7 @@ def stop_trace():
       if _profile_state.create_perfetto_link:
         _host_perfetto_trace_file(abs_filename)
     _profile_state.reset()
+    _profiler.clear_metadata()
 
 
 def stop_and_get_fdo_profile() -> bytes | str:
@@ -240,6 +257,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()
     return fdo_profile