xprof-rs

xprof-rs is a fast backend for the XProf profile viewer. It is one Rust binary. It replaces the XProf server and the XProf agent CLI, and it serves the same user interface. Its responses are the same as the responses of XProf 2.23.2. The tests compare them with the output of XProf.

xprof-rs is an independent project. Google and the OpenXLA project do not maintain or endorse it.

We wrote xprof-rs with Claude, an AI model from Anthropic. The tests compare the output of each tool with the output of XProf 2.23.2.

On an 80 MB TPU v4 profile of a job that trains a CLIP model, the trace viewer opens in 0.30 s (XProf: 5.6 s). After the first request, it opens in approximately 2 ms (XProf: 0.62 s).

The peak memory of the server is 0.42 GB (XProf: 1.42 GB). On a 1.69 GB profile, XProf uses more than 16 GB of memory, and xprof-rs opens the trace viewer in 8.3 s with 7.7 GB.

xprof-rs also reads profiles from S3-compatible stores such as Cloudflare R2. XProf 2.23.2 cannot read them. The performance section has the times of all tools, also on a disk and on R2.

We measured all times in this file with 4 cores of a larger machine, unless the text gives a different number of cores.

XProf converts the .xplane.pb file again for each cold request. xprof-rs does these steps:

1. It reads the file into memory.
2. It converts the events in parallel.
3. It keeps the result in a cache.
4. It renders only the time window that you zoom to. A new window takes about 10 ms, and the same window again takes about 3 ms.

Performance

The tables show the times on an 80 MB TPU v4 profile of a job that trains a CLIP model. Each server and each CLI process used 4 cores.

• Server, cold: the first request to a new server on a new copy of the profile. There is no cache on the disk.
• Server, warm: the mean of 5 more requests to the same server.
• CLI, cold: one process on a new copy of the profile, with no result cache.
• CLI, warm: the same command again.
• Peak memory: the peak resident memory of the server over the cold and the warm requests, or of the cold CLI process. It includes the pages of the profile that the process maps.

Each cell is the mean and the standard deviation of 8 trials, after 1 warmup trial that is not counted, as in pyperf. The trials of XProf and xprof-rs alternate, so a change of the machine load affects the two the same. The mark (!) shows a standard deviation of more than 10% of the mean.

Some warm XProf commands are as fast as xprof-rs or faster. XProf keeps each result in a cache in $TMPDIR, and a second call reads this cache. xprof-rs has no result cache. It reads the profile again for each call.

The table has all the commands that need only a session. get_hlo_neighborhood, verify_numerical_parity, and upload_trace need more input, so the table does not show them. get_xspace_proto writes the 80 MB profile to /tmp on a disk. On a warm call, XProf writes over the file of the cold call, and the write waits until the disk has the data.

The next table shows five profiles of jobs that train CLIP models. The trace viewer time is for a cold server. The peak memory is the peak of the server process after the trace viewer and seven other tools.

These numbers are for the 80 MB profile:

More cores make xprof-rs faster. The next table shows the cold times of xprof-rs on the 80 MB profile. XProf on 4 cores takes 5.4 s for the trace viewer and 3.9 s for the overview page.

To measure the times of the first two tables on your profile, run examples/benchmark.py SESSION_DIR --xprof PATH --cores 0-3 --trials 8.

The next table shows the cold times of the overview page on a disk and on Cloudflare R2. Each time starts with a new server and stops at the end of the first response. XProf 2.23.2 cannot read s3:// URLs, so it reads R2 through rclone mount. The section Times on a disk and on R2 has all the tools and the method.

Large profiles

We made two larger profiles with the same script that trains the CLIP model. 140 steps gave a 798 MB profile, and 300 steps gave a 1.69 GB profile. These tables have 3 trials, after 1 warmup trial. A cold call stops after 120 s, and a process stops when its resident memory is more than 16 GB. The profiles have no .hlo_proto.pb files, so the tables do not show memory_viewer.

On the 798 MB profile, a cold call of XProf takes up to 74 s, and XProf uses up to 14.8 GB of memory. xprof-rs takes at most 3.7 s and uses at most 3.9 GB. On the 1.69 GB profile, XProf uses more than 16 GB for each server tool and for most commands. xprof-rs opens the trace viewer in 8.3 s with a peak of 7.7 GB. All other tools and commands take at most 6.1 s and use at most 6.2 GB.

The 798 MB profile:

** The process used more than 16 GB of memory. Because it used so much, the script stopped it and did not measure the row further.

** The process used more than 16 GB of memory. Because it used so much, the script stopped it and did not measure the row further.

The 1.69 GB profile:

** The process used more than 16 GB of memory. Because it used so much, the script stopped it and did not measure the row further.

** The process used more than 16 GB of memory. Because it used so much, the script stopped it and did not measure the row further.

Install and start

The releases have binaries for Linux on x86_64 and aarch64, and for macOS on Apple silicon. The Linux binaries need glibc 2.28 or newer (RHEL 8, Debian 10, Ubuntu 20.04, or newer). You do not need Python.

version=v0.1.3 target=x86_64-linux   # or aarch64-linux, aarch64-macos
curl -LO https://github.com/Locamage/xprof-rs/releases/download/$version/xprof-rs-$version-$target.tar.gz
tar -xzf xprof-rs-$version-$target.tar.gz
xprof-rs-$version-$target/xprof-rs --logdir ~/logs          # open http://localhost:8791
xprof-rs-$version-$target/xprof-rs get_overview ~/logs/run1 # the same binary runs the XProf agent CLI

The same binaries are on PyPI. pip install xprof-rs puts xprof-rs on the PATH. The wheel has no Python code.

To build from source, you need Rust 1.95 or newer, a C compiler, and a 64-bit Unix system. You do not need a C++ toolchain or protoc.

cargo install --locked xprof-rs

The release binaries are for all CPUs of their architecture. To build only for the CPU of your machine, set RUSTFLAGS="-C target-cpu=native". That binary possibly does not start on a different CPU.

The release profile uses fat LTO and one codegen unit. A full build takes about 80 s on a machine with 240 cores. It uses about 10 CPU minutes. On 8 cores, it takes about 4 minutes. For a quick build, set CARGO_PROFILE_RELEASE_LTO=false and CARGO_PROFILE_RELEASE_CODEGEN_UNITS=16.

Server

xprof-rs [--logdir DIR|URL] [--port 8791] [--host ADDRESS] [--src_prefix PREFIX] [--hide_capture_profile_button] [--enable_tab_name_label]

• The user interface is the prebuilt frontend of the XProf 2.23.2 Python package. It uses the Apache License 2.0. The copyright holder is The TensorFlow Authors.
• The binary includes the files from ui/. To serve other files, set XPROF_STATIC_DIR to a directory.
• Each route answers at the root and under /data/plugin/profile.
• The server listens on 127.0.0.1 port 8791. There is no authentication. To let other computers connect, use --host 0.0.0.0 (all IPv4 interfaces) or --host :: (all interfaces). XProf listens on all interfaces.
• With --logdir, a request can read only the files in the log directory.
• A run, session_path, or module_name that points outside the log directory gets a 400 response.
• The directory walk does not follow symbolic links to directories.
• Without --logdir, the server lists no runs. A session_path or run_path request can still open a directory.
• The server accepts --grpc_port, --worker_service_address, and --max_concurrent_worker_requests. It does not use them.

The server has these tools. They are all written in Rust.

• Trace viewer: trace_viewer@ (zoom, search, details, DMA, format=pb) and trace_viewer.
• Statistics: overview_page, op_profile, hlo_stats, framework_op_stats, input_pipeline_analyzer, roofline_model, kernel_stats, inference_profile.
• Memory and graphs: memory_profile, memory_viewer, graph_viewer, module_list.
• Hardware: pod_viewer, megascale_stats, smart_suggestion, perf_counters, utilization_viewer, kernel_utilization.
• Session: runs, run_tools, hosts, data_csv, version, config, POST /generate_cache, /capture_profile (gRPC client), and the static files.

A machine with more cores uses more memory, because more work runs at the same time.

Remote log directories

--logdir accepts these URL schemes: gs://, s3:// (also S3-compatible stores), az://, abfs://, http(s):// (WebDAV), and file://. The object_store crate does the access.

The credentials come from the environment:

• gs://: GOOGLE_APPLICATION_CREDENTIALS or the application default credentials.
• S3: the AWS_* variables.
• Azure: the AZURE_* variables.

The server copies the store to a local directory, and all tools read this copy. The directory is xprof-rs-<hash of the URL> in XPROF_CACHE_DIR. The default is the temporary directory. Only the user of the server can read the directory. The server does not start if a different user owns it.

• A request that names a session downloads the .xplane.pb and .hlo_proto.pb files of the session.
• The server downloads 4 files at a time. It reads each file in ranges of 8 MiB, and it reads 8 ranges at a time.
• The server does not download a file again when the size and the modification time are the same as in the store.
• The server asks the store at most one time in 5 s for each session.
• Each 30 s, a background poll checks the list of sessions. It also checks the 8 sessions with the newest requests.
• If the copy is larger than XPROF_CACHE_BYTES, the server removes sessions. The default is 20 GiB. The server removes the session with the oldest request first. It keeps a session with a request in the last minute.
• If the store is not reachable, the response is 502 with the error of the store. The response is the same for credentials that are not correct. A list request stops after 30 s.
• The server does not write to the store. There is one exception. /capture_profile uploads the captured files to <URL>/plugins/profile/<session>/.

Set the URL to the directory that holds the profiles. The server visits each directory under it.

Times on a disk and on R2

XProf 2.23.2 does not read s3:// URLs. Its server and its file code know only gs:// and local paths. To read R2, XProf needs a FUSE mount of the bucket, for example rclone mount. xprof-rs can read R2 through the s3:// URL or through the mount.

The next tables show the cold times of a new server with 4 cores. Each time starts with the process and stops at the end of the first response.

• tmpfs: the files are in memory.
• Disk: the persistent disk of a cloud VM. Before each trial, the page cache did not have the files.
• R2 mount: a new rclone mount --vfs-cache-mode writes of Cloudflare R2 for each trial. XProf writes its cache files into the session directory, so the mount must accept writes.
• R2: xprof-rs reads the s3:// URL. Before each trial, the cache directory was empty.

Before each trial, we removed the files that the servers wrote. Each cell is the mean of 3 trials on the 80 MB profile and of 2 trials on the 279 MB profile.

80 MB profile:

279 MB profile:

• After the first request, xprof-rs answers in 0.4 ms to 6 ms from all sources. XProf answers in 55 ms to 2.2 s on the 80 MB profile and in 0.6 s to 6.9 s on the 279 MB profile.
• The responses of xprof-rs were the same for all sources. For 5 tools, the bytes of XProf change from one process to the next. The changes are the order of map keys and of equal items, and the bind_id values of the trace viewer. Without these changes, the responses of XProf and xprof-rs have the same data.
• The disk gave approximately 150 MB/s. This is the limit of the disk, and it changes with the load of the disk.
• A read through rclone mount is approximately 46 MB/s, because the mount reads a file in order with one stream. On a FUSE mount, xprof-rs reads each file in order. On other file systems, it reads parts of the file in parallel.
• With the s3:// URL, the first list of the store takes approximately 0.3 s. The download is approximately 210 MB/s, approximately the same as rclone with 8 to 32 streams. The server downloads all the files of the session, also for the memory viewer, which needs only the HLO file.
• The R2 times change with the network. One XProf trial of the op profile on the mount took 149 s, and we measured it again.
• When the server starts again with the same cache directory, it does not download the files again. The first overview page of the 279 MB profile then takes 1.05 s.

Command line

Use xprof-rs <command> <session> [--flag=value ...] to run the agent CLI of XProf. The stdout and the exit codes are the same as in XProf 2.23.2.

The session is one of these:

• a run directory
• a log directory (the CLI uses the latest run)
• an .xplane.pb file
• a run name, with --logdir

The flags follow XProf (Python Fire):

• --name=value and --name value
• --name and --noname for booleans
• positional parameters in the order of XProf
• the aliases --session_dir, --session_path, and --source

An error prints {"status": "ERROR", "reason": ...}. The exit code is 2 (usage), 3 (path), 4 (value not valid), or 1 (internal). A result larger than 10 MB goes to a file in $TMPDIR. This is the same as in XProf.

The commands are:

• Summary: get_overview, get_kpi_metrics, get_device_information, get_hosts, get_profile_summary, get_avg_step_time, get_step_trace, check_host_boundness.
• Operations: get_roofline_model, get_top_hlo_ops, get_hlo_op_profile, get_hlo_stats, get_kernel_stats.
• Memory: get_memory_profile, get_peak_allocations.
• HLO: list_hlo_modules, get_hlo_module_content, get_hlo_neighborhood, get_hlo_text, get_graph_viewer.
• Utilization: get_utilization_viewer, get_kernel_utilization, compute_utilization.
• Events and files: list_xplane_events, aggregate_xplane_events, get_xspace_proto, upload_trace.
• Not available: get_llo_analysis, get_llo_debug_string, verify_numerical_parity.

The flags of each command are the flags of XProf. Examples:

xprof-rs get_top_hlo_ops ~/logs/run1 --limit=5 --category_filter=fusion
xprof-rs get_hlo_op_profile ~/logs --view=tree --path=by_program --depth=3
xprof-rs list_xplane_events ~/logs/run1 --plane_regex='TPU:0$' --event_regex=all-reduce --max_events=20

More cores make a command faster, up to about 16 cores.

The CLI is different from the Python CLI in these points:

• std_dev_ps of aggregate_xplane_events is the correctly rounded value, as statistics.stdev gives in Python 3.11 or newer. With Python 3.10, XProf can give a different last digit.
• There is no result cache. Each call computes the result again, as with --bypass_cache. There are no __cached__ fields.
• The regex options find text as Python re does, and a bad pattern gives the error message of Python. But \N{name} always gives the error undefined character name, and a pattern with more than approximately 60 levels of nested groups gives an error.
• --help prints a short synopsis.
• get_graph_viewer --output_type=pb returns the bytes of the module file. XProf serializes the module again with hash-ordered maps.
• For a multi-host session, the tools that use combined op statistics report no data. memory_profile fails. XProf needs its worker service for these sessions.

Not available and not supported

xprof-rs replaces the backend. It does not replace the complete XProf distribution.

Inputs

• xprof-rs does not read .xplane.riegeli sessions. The continuous capture mode of the gRPC client is the only source of these files. The jax.profiler.start_trace function writes .xplane.pb. MaxText, Levanter, and marin use this function.
• xprof-rs rejects a profile of 4 GiB or more. A profile file is one protobuf message, and protobuf limits a message to 2 GiB. Thus XProf cannot read a profile of 2 GiB or more.

Server

• There is no distributed worker mode. There is no worker gRPC service. /capture_profile is a client.
• There is no TLS. There is no authentication. There is no TensorBoard plugin loader.
• xprof-rs runs on 64-bit Unix. CI tests Linux on x86_64 and aarch64, and macOS on Apple silicon. There is no Windows build.

Tools

• get_llo_analysis and get_llo_debug_string return the LLO-not-available report of XProf. The source code of the LLO analysis is not public.
• verify_numerical_parity returns the error of XProf for dependencies that are not installed. XProf runs Python callables for this command.
• The detect_* commands from the README of XProf are not available. XProf 2.23.2 does not register them.
• mpmd_pipeline_view=true has no effect. XProf has the same behavior.
• perf_counters with names_only and utilization_viewer support only TPU v6e and v7x. XProf 2.23.2 has counter tables only for these two chips. The tables are in src/data/.

Python part of the XProf package

• These parts are not available: xprof.convert, the ProfileData library, the xprof.api snapshot client, install_and_run.py, the labs, and the agent skills.

Frontend

• This repository does not contain the sources of the user interface.

Output that changes from run to run in XProf

xprof-rs cannot copy these outputs.

• The order of tied rows in hlo_stats, op_profile, memory_viewer, and get_peak_allocations. XProf uses hash order. The seed changes for each process. A node of op_profile shows at most 100 children. When children have the same time, the tie order selects the children that show. The category sums of get_hlo_op_profile add the children that show, so these sums can also change.
• The tied ops of get_top_hlo_ops. This command reads the same children, so its order of tied ops and its list of ops can also change.
• The order of processes and threads in the protobuf.
• The allocator peaks and the fragmentation snapshot of the memory profile.
• The bind_id flow ids of the trace viewer.
• The order of the stackFrames entries of the trace viewer. XProf writes them in hash order.
• The graph node ids. They are heap pointers.
• The hover CSS rules and the caller lists of the graph viewer.
• The Launch Stats line of a GPU device. XProf writes it in hash order. The uid of host events at the same timestamp can differ.
• The statistics of a multi-host session. XProf caches the first op statistics under ALL_HOSTS. It sometimes leaves the duty cycle at 0. xprof-rs answers each host selection as a new XProf process does.
• The tensor pattern rows of inference_profile.
• The order of requests with equal latency in inference_profile.
• The order of the sums in the latency averages.

Test coverage

• tests/data has a demo trace. It is a small synthetic workload on a TPU v4-8. We replaced the host name, the script path, and the host system information.
• tests/data also has the outputs of XProf for the demo trace. It has synthetic GPU profiles and HLO fixtures in tests/data/gpu.
• cargo test compares each tool and the trace viewer with these outputs.
• Other tests used TPU v4 and v6e traces and one 8-GPU trace. They also used synthetic sessions of 2, 4, and 8 hosts.
• The synthetic GPU sessions use 1, 2, or 8 GPUs from NVIDIA or AMD. They cover TensorFlow, JAX, eager steps, NCCL, and CUDA graphs.
• The synthetic inference sessions cover TensorFlow, request batches, TFRT, Pathways, and Orbax. No real inference profile was available.
• These traces are not in the repository.
• The real-world coverage is small. If a response is different from XProf 2.23.2, open an issue. Do not include private data.
• To report a security problem, follow SECURITY.md. Do not open a public issue.

Memory

• Three LRU caches hold converted traces, op statistics, and rendered responses (compressed). Each of the first two holds an eighth of the physical memory, at most 4 GiB. The third holds a quarter of that. On a machine with 32 GiB or more, the total is about 9 GiB.
• The server loads 2 files at once on a machine with 32 GiB or more, and 1 file otherwise. A load needs about four to six times the size of its file. More cores need more memory.
• The server and the CLI use one thread per core, at most 32.
• A cache removes an entry after one hour without use.
• A trace that is larger than its cache is still served. The server caches it alone.
• The server loads a file again when the size, the modification time, or the inode changes. It does not cache a file that changed in the last 2 s.
• A background task loads the 8 newest sessions at startup and when a session changes. It loads only the files that did not change for 30 s and that are smaller than a thirty-second of the physical memory. This task keeps about 0.7 GB more in memory for a 279 MB trace.
• A session of 4 hosts has 279 MB v4 traces. The trace viewer of all hosts takes 1.8 s. overview_page takes 3.8 s and a peak of 4.5 GB.
• For 8 hosts, the times are 3.8 s and 5.2 s. The peak is 7.6 GB.

Development

Run these commands before each commit. Cargo.toml turns on the pedantic lints of clippy.

cargo fmt --check
cargo clippy --release --all-targets -- -D warnings
cargo test --release

A tag that starts with v makes a draft release with x86_64-linux, aarch64-linux, and aarch64-macos archives and SHA-256 files. The steps are in RELEASING.md.

The code is in src/:

• main.rs: flags, routes, and caches.
• xplane/: the .xplane.pb reader, derived lines, event groups, and steps.
• trace/: the trace viewer conversion and its JSON and protobuf output.
• hlo/: the HLO module, its text form, the GPU cost model, the graph viewer, and the memory viewer.
• tools/: one module for each other tool.
• server/: object stores, profile capture, and the tool list of a run.
• cli/: the agent CLI.
• data/: the protobuf descriptors, the TPU counter tables, and the HTML page of the graph viewer.
• tests/: tests from XProf and XLA. tests/inline/ has the unit tests of the modules, in the same folders.

License

Apache-2.0. The algorithms follow Apache-licensed code from OpenXLA XProf and XLA/TSL. See NOTICE.