xplane_to_perfetto.py

def _xplane_to_perfetto(xplane_path: str) -> bytes:
    data, _ = raw_to_tool_data.xspace_to_tool_data([xplane_path], "trace_viewer@", {})
    if isinstance(data, str):
        data = data.encode()
    return gzip.compress(data)


def _bundle_xplane(xplane_path: str, logdir: str) -> bytes:
    """Tar.gz the profile run dir, laid out so `xprof <XPROF_DIR>` works after extraction."""
    buf = io.BytesIO()
    with tarfile.open(fileobj=buf, mode="w:gz") as tar:
        for path in glob.glob(f"{os.path.dirname(xplane_path)}/**/*", recursive=True):
            # skip dirs and the derived junk xprof writes next to the xplane,
            # only the raw profiler outputs need to be bundled and uploaded.
            if os.path.isdir(path) or path.endswith((".SSTABLE", ".json.gz")):
                continue
            tar.add(path, arcname=f"{XPROF_DIR}/{os.path.relpath(path, logdir)}")
    return buf.getvalue()


def _upload_to_s3(data: bytes, key: str) -> str:
    # you can implement this yourself
  
@contextlib.contextmanager
def with_perfetto_trace(logdir: str = "/tmp/xplane-traces"):
  options = jax.profiler.ProfileOptions()
  with jax.profiler.trace(logdir, profiler_options=options):
        yield

    xplanes = glob.glob(f"{logdir}/**/*.xplane.pb", recursive=True)
    if not xplanes:
        print(f"no xplane produced under {logdir}")
        return
    xplane = max(xplanes, key=os.path.getmtime)
    print(f"trace captured: {xplane}")

    url = _upload_to_s3(_xplane_to_perfetto(xplane), timestamp)
    print(
        f"perfetto: "
        f"https://ui.perfetto.dev/#!/?url={urllib.parse.quote(url, safe='')}"
    )