jax.profiler.trace fails with Perfetto example

[jwtkeeble]

Hi All,

I'm trying to trace my code to understand where I can improve it, but I can't seem to get the Perfetto trace to work. I'm using the following example script from here below to test the profiling of the code with Perfetto,

import jax

with jax.profiler.trace("/tmp/jax-trace", create_perfetto_link=True, create_perfetto_trace=True):
  # Run the operations to be profiled
  key = jax.random.key(0)
  x = jax.random.normal(key, (5000, 5000))
  y = x @ x
  y.block_until_ready()

However, when I run it I get the following error,

2025-01-30 17:32:57.209764: E external/local_xla/xla/stream_executor/cuda/cuda_fft.cc:477] Unable to register cuFFT factory: Attempting to register factory for plugin cuFFT when one has already been registered
WARNING: All log messages before absl::InitializeLog() is called are written to STDERR
E0000 00:00:1738254777.219665   11671 cuda_dnn.cc:8310] Unable to register cuDNN factory: Attempting to register factory for plugin cuDNN when one has already been registered
E0000 00:00:1738254777.222845   11671 cuda_blas.cc:1418] Unable to register cuBLAS factory: Attempting to register factory for plugin cuBLAS when one has already been registered
Open URL in browser: https://ui.perfetto.dev/#!/?url=http://127.0.0.1:9001/perfetto_trace.json.gz

And, when I go to open the URL above it throws an error saying: Could not load local trace TypeError: Failed to fetch I have the following jax versions, for reference:

jax                          0.4.35
jaxlib                       0.4.35
jaxtyping                    0.2.36

and tensorflow versions (if that's even relevant) here,

tensorboard                  2.18.0
tensorboard-data-server      0.7.2
tensorboard-plugin-profile   2.19.0
tensorflow                   2.18.0
tensorflow-io-gcs-filesystem 0.37.1

Any help would be appreciated!

[quattro]

I've found using viztracer to be a bit better experience compared to the built-in, when trying to profile JAX programs.

[sharadmv]

Out of curiosity, how are you using viztracer with JAX?

[quattro]

with VizTracer(output_file="output.json") as tracer:
  my_jax_function(...)

[sharadmv]

Thanks!