Graph viewer on Xprof missing

[justinjfu]

When running from JAX near head (https://github.com/jax-ml/jax/commits/f0ed696cfcaa54a8f6b7217d63ffc897f72eb6f3/), it seems that Xprof loses the graph viewer.

Reproducer:

# Build JAX at head
git checkout f0ed696cfcaa54a8f6b7217d63ffc897f72eb6f3
python build/build.py build --wheels=jaxlib --python_version=3.12 --verbose
pip install <build jaxlib wheel>

Now run:

import jax
from jax import numpy as jnp

@jax.jit
def foo(x, y):
  return x + y

with jax.profiler.trace("/tmp/jax-trace"):
  x = jnp.ones((5,))
  y = jnp.arange(5)
  jax.block_until_ready(foo(x, y))

Now when running xprof --logdir=/tmp/jax-trace, the graph viewer is missing. If I click on a graph viewer link for a particular op in the trace, it also brings me back to the overview page.

Note that the graph viewer still seems to be present when using jaxlib==0.9.0, so this seems to be a regression introduced recently.

I'm running Xprof version 2.21.6 (this was installed by running pip install xprof)

[justinjfu]

For some more debugging context, the trace directory for older jaxlib (0.9.0) looks like this:

-rw-rw-r-- 1 jujfu jujfu    346 Feb 12 20:20  .cached_tools.json
-rw-rw-r-- 1 jujfu jujfu   1767 Feb 12 20:21  ALL_HOSTS.op_stats.pb
-rw-rw-r-- 1 jujfu jujfu      6 Feb 12 20:20  cache_version.txt
-rw-rw-r-- 1 jujfu jujfu 394320 Feb 12 20:20  ip-10-100-58-215.SSTABLE
-rw-rw-r-- 1 jujfu jujfu 163334 Feb 12 20:20  ip-10-100-58-215.metadata.SSTABLE
-rw-rw-r-- 1 jujfu jujfu 283280 Feb 12 20:19  ip-10-100-58-215.trace.json.gz
-rw-rw-r-- 1 jujfu jujfu 218588 Feb 12 20:20  ip-10-100-58-215.trie.SSTABLE
-rw-rw-r-- 1 jujfu jujfu 435185 Feb 12 20:19  ip-10-100-58-215.xplane.pb
-rw-rw-r-- 1 jujfu jujfu    871 Feb 12 20:20 'jit_broadcast_in_dim(663).hlo_proto.pb'
-rw-rw-r-- 1 jujfu jujfu    390 Feb 12 20:20 'jit_convert_element_type(1).hlo_proto.pb'
-rw-rw-r-- 1 jujfu jujfu   1187 Feb 12 20:20 'jit_foo(1987).hlo_proto.pb'
-rw-rw-r-- 1 jujfu jujfu    541 Feb 12 20:20 'jit_iota(1325).hlo_proto.pb'

However, with the newer jaxlib there are no longer module hlo_proto files, and instead there is just a NO_MODULE.hlo_proto.pb file.

-rw-rw-r-- 1 jujfu jujfu    305 Feb 12 20:39 .cached_tools.json
-rw-rw-r-- 1 jujfu jujfu    510 Feb 12 20:39 ALL_HOSTS.op_stats.pb
-rw-rw-r-- 1 jujfu jujfu      0 Feb 12 20:39 NO_MODULE.hlo_proto.pb
-rw-rw-r-- 1 jujfu jujfu      6 Feb 12 20:39 cache_version.txt
-rw-rw-r-- 1 jujfu jujfu 410829 Feb 12 20:39 ip-10-100-58-215.SSTABLE
-rw-rw-r-- 1 jujfu jujfu 171576 Feb 12 20:39 ip-10-100-58-215.metadata.SSTABLE
-rw-rw-r-- 1 jujfu jujfu 293192 Feb 12 20:38 ip-10-100-58-215.trace.json.gz
-rw-rw-r-- 1 jujfu jujfu 223011 Feb 12 20:39 ip-10-100-58-215.trie.SSTABLE
-rw-rw-r-- 1 jujfu jujfu 450505 Feb 12 20:38 ip-10-100-58-215.xplane.pb