Set up an API to top trace and fdo profile in memory.

PiperOrigin-RevId: 573276173

jax/_src/profiler.py

@@ -30,6 +30,7 @@
 
 from jax._src import xla_bridge
 from jax._src.lib import xla_client
+from jax._src.lib import xla_extension_version
 
 _profiler_server: Optional[xla_client.profiler.ProfilerServer] = None
 
@@ -76,6 +77,13 @@ def __init__(self):
     self.create_perfetto_trace = False
     self.lock = threading.Lock()
 
+  def reset(self):
+    _profile_state.profile_session = None
+    _profile_state.create_perfetto_link = False
+    _profile_state.create_perfetto_trace = False
+    _profile_state.log_dir = None
+
+
 _profile_state = _ProfileState()
 
 
@@ -193,15 +201,33 @@ def stop_trace():
   with _profile_state.lock:
     if _profile_state.profile_session is None:
       raise RuntimeError("No profile started")
-    _profile_state.profile_session.stop_and_export(_profile_state.log_dir)
+    if xla_extension_version > 205:
+      sess = _profile_state.profile_session
+      sess.export(sess.stop(), _profile_state.log_dir)
+    else:
+      _profile_state.profile_session.stop_and_export(_profile_state.log_dir)  # pytype: disable=attribute-error
     if _profile_state.create_perfetto_trace:
       abs_filename = _write_perfetto_trace_file(_profile_state.log_dir)
       if _profile_state.create_perfetto_link:
         _host_perfetto_trace_file(abs_filename)
-    _profile_state.profile_session = None
-    _profile_state.create_perfetto_link = False
-    _profile_state.create_perfetto_trace = False
-    _profile_state.log_dir = None
+    _profile_state.reset()
+
+
+def stop_and_get_fdo_profile() -> bytes:
+  """Stops the currently-running profiler trace and export fdo_profile.
+
+  Currently, this is only supported for GPU.
+  Raises a RuntimeError if a trace hasn't been started.
+  """
+  if xla_extension_version < 206:
+    raise NotImplementedError("stop and get fdo profile is not supported.")
+  with _profile_state.lock:
+    if _profile_state.profile_session is None:
+      raise RuntimeError("No profile started")
+    xspace = _profile_state.profile_session.stop()
+    fdo_profile = xla_client.profiler.get_fdo_profile(xspace)
+    _profile_state.reset()
+    return fdo_profile
 
 
 @contextmanager
@@ -316,7 +342,6 @@ def wrapper(*args, **kwargs):
   return wrapper
 
 
-
 def device_memory_profile(backend: Optional[str] = None) -> bytes:
   """Captures a JAX device memory profile as ``pprof``-format protocol buffer.
 

tests/profiler_test.py

@@ -27,6 +27,7 @@
 import jax.numpy as jnp
 import jax.profiler
 from jax import config
+from jax._src.lib import xla_extension_version
 import jax._src.test_util as jtu
 
 try:
@@ -103,6 +104,19 @@ def testProgrammaticProfiling(self):
         self.assertIn(b"/device:TPU", proto)
       self.assertIn(b"pxla.py", proto)
 
+  def testProfilerGetFDOProfile(self):
+    if xla_extension_version < 206:
+      return
+    # Tests stop_and_get_fod_profile could run.
+    try:
+      jax.profiler.start_trace("test")
+      jax.pmap(lambda x: jax.lax.psum(x + 1, "i"), axis_name="i")(
+          jnp.ones(jax.local_device_count())
+      )
+    finally:
+      fdo_profile = jax._src.profiler.stop_and_get_fdo_profile()
+    self.assertEqual(fdo_profile, b"")
+
   def testProgrammaticProfilingErrors(self):
     with self.assertRaisesRegex(RuntimeError, "No profile started"):
       jax.profiler.stop_trace()