456 add support for traces generated by jax 08

[gphuang]

Summary

Closes #456. Migrates TraceLens from tensorboard-plugin-profile to xprof 2.20.7 for JAX trace processing so traces from JAX 0.8+ can be loaded and analyzed.

Why xprof 2.20.7 instead of 2.21.3

• xprof 2.20.8+ / 2.21.x removes hlo_op and correlation_id from GPU kernel events, which breaks TraceLens HLO linking and perf metrics.
• xprof 2.20.7 keeps these fields and loads JAX 0.8 traces without the INT_MAX crash seen with tensorboard-plugin-profile 2.19.

Backward compatibility

• Older JAX 0.6 traces: protobuf version conflict! (See comments below.) (e.g. jax_conv_minimal_legacy, JAX ~0.6): All 10 test_jax_conv_analysis tests pass, including perf model (FLOPS, bytes, shapes).
• JAX 0.8 traces: Trace loading, event categorization, and GPU timelines work; HLO metadata for perf models is limited until Add hlo_op and framework name scope call_stack nodes to Jax.tree #425 lands.

Changes

• Replace tensorboard-plugin-profile with xprof 2.20.7 (with fallback for backward compatibility)
• Switch GPU event detection from PID-based (pid <= 100) to process_name-based (/device:GPU)
• Add _resolve_operand_references for xprof's bare operand references
• Fix custom_call_target and backend_config extraction when metadata= is absent
• Fix get_gpu_event_lists PID classification (use pid in self.gpu_pids instead of pid < 100)
• Rename traces: jax_conv_minimal → jax_conv_minimal_legacy, jax08_conv → jax_conv_minimal_08 in tests/traces/mi300/
• Add JAX 0.8 conv trace tests/traces/mi300/jax_conv_minimal_08 for E2E smoke tests
• Pin protobuf>=6.31.1,<7.0.0 for grpcio-status compatibility (avoids RET_CHECK on JAX 0.8 traces)
• Pin Black to 26.1.0 in lint workflow (avoids CI failing when Black versions changes)
• Code quality: _is_gpu_event() helper, guard name/args against None, test cleanup fixture

[Copilot]

Pull request overview

Migrates TraceLens’ JAX trace ingestion from tensorboard-plugin-profile to xprof 2.20.7 to support JAX 0.8+ xplane traces while preserving legacy perf/HLO linking behavior.

Changes:

• Switch trace conversion to prefer xprof (with fallback import path) and add HLO parsing fixes for xprof output (operand reference resolution; metadata/custom_call parsing).
• Update JAX GPU-event identification to use process.process_name (e.g. /device:GPU) instead of PID heuristics, and pass metadata events into JAX tree building.
• Update/relax certain JAX conv tests and strengthen perf report comparison by asserting missing reference columns; add/rename MI300 JAX trace artifacts + documentation updates.

Reviewed changes

Copilot reviewed 9 out of 22 changed files in this pull request and generated no comments.

Show a summary per file

File	Description
.github/workflows/lint.yml	Pins Black version used in CI lint job.
.gitignore	Ignores *.SSTABLE artifacts.
setup.py	Adds xprof==2.20.7 dependency for JAX trace conversion.
TraceLens/util.py	Prefers xprof for xplane conversion; improves HLO metadata parsing and resolves bare operand references.
TraceLens/TreePerf/gpu_event_analyser.py	Derives GPU PIDs using process_name metadata.
TraceLens/TreePerf/tree_perf.py	Splits metadata vs non-metadata events and threads metadata into JaxTraceToTree.build_tree().
TraceLens/Trace2Tree/trace_to_tree.py	Uses process.process_name to detect GPU events for tree augmentation/categorization.
tests/test_compare_perf_report.py	Adds assertion for missing reference columns before comparing values.
tests/test_jax_conv_analysis.py	Updates legacy JAX conv trace path and loosens event-category assertions.
tests/traces/README.md	Documents MI300 JAX trace naming for legacy vs 0.8 traces.
tests/traces/mi300/jax_conv_minimal_legacy_perf_report.xlsx	Adds/updates reference perf report artifact for MI300 legacy JAX conv.
tests/traces/mi300/jax_conv_minimal_legacy/jit_forward_3d_conv(15).hlo_proto.pb	Adds HLO sidecar for legacy trace perf/HLO linking.
tests/traces/mi300/jax_conv_minimal_legacy/jit_convert_element_type(1).hlo_proto.pb	Adds HLO sidecar for legacy trace perf/HLO linking.
tests/traces/mi300/jax_conv_minimal_legacy/jit__unstack(7).hlo_proto.pb	Adds HLO sidecar for legacy trace perf/HLO linking.
tests/traces/mi300/jax_conv_minimal_legacy/jit__threefry_split(5).hlo_proto.pb	Adds HLO sidecar for legacy trace perf/HLO linking.
tests/traces/mi300/jax_conv_minimal_legacy/jit__threefry_seed(3).hlo_proto.pb	Adds HLO sidecar for legacy trace perf/HLO linking.
tests/traces/mi300/jax_conv_minimal_legacy/jit__normal(11).hlo_proto.pb	Adds HLO sidecar for legacy trace perf/HLO linking.
tests/traces/mi300/jax_conv_minimal_legacy/jit__multi_slice(13).hlo_proto.pb	Adds HLO sidecar for legacy trace perf/HLO linking.

Comments suppressed due to low confidence (3)

TraceLens/Trace2Tree/trace_to_tree.py:435

• name = event.get("name") may be None for some trace events; later logic in this method does substring checks against name, which will raise a TypeError if name is not a string. Consider normalizing to an empty string (or skipping categorization) when name is missing/non-string to keep tree building robust across trace variants.

            if "/device:GPU" in proc_name:

                if event.get("cat") == "kernel":
                    name = event.get("name")

TraceLens/util.py:56

• This change introduces xprof-based conversion (and a new JAX 0.8 trace is added to test data), but the test suite still only exercises JAX parsing against the legacy trace. Add a unit/integration test that loads the jax_conv_minimal_08 .xplane.pb via JaxTreePerfAnalyzer.from_file to ensure the xprof conversion path works and tree building/categorization doesn’t regress.

        if filename_path.endswith("pb"):
            try:
                from xprof.convert import raw_to_tool_data as convert

                converter_lib = "xprof"
            except ImportError:
                from tensorboard_plugin_profile.convert import (
                    raw_to_tool_data as convert,
                )

                converter_lib = "tensorboard-plugin-profile"
                logger.warning(
                    "xprof not available, falling back to tensorboard-plugin-profile "
                    "for trace conversion. Install xprof for JAX 0.8+ support."
                )

            data, _ = convert.xspace_to_tool_data([filename_path], "trace_viewer@^", {})
            if data is None:
                raise RuntimeError(
                    f"Trace conversion using '{converter_lib}' returned None for "
                    f"{filename_path}. Ensure the file exists and the output directory "
                    "is writable (cache files may need to be written)."
                )
            data = data.decode("utf-8")  # we get bytes back from the call above

TraceLens/TreePerf/gpu_event_analyser.py:337

• JaxGPUEventAnalyser now derives gpu_pids from process_name metadata, but get_gpu_event_lists() still classifies GPU vs CPU using pid < 100. If xprof/JAX 0.8 traces don’t preserve the legacy PID scheme, this will silently drop GPU events from metrics. Consider using the same process_name-based check (or pid in self.gpu_pids) inside get_gpu_event_lists() for consistency.

    def __init__(self, events):
        super().__init__(events)  # Call the parent's __init__
        self.gpu_pids = list(
            set(
                event["pid"]
                for event in events
                if "/device:GPU"
                in (
                    event.get("process", {}).get("process_name", "")
                    if isinstance(event.get("process"), dict)
                    else str(event.get("process", ""))
                )
            )

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[gabeweisz]

These hlo_proto.pb files should not be checked in - they are generated when xprof.convert / tensorboard_plugin_profile.convert processes the xplane.pb file

[gphuang]

“xprof’s xspace_to_tool_names() no longer extracts .hlo_proto.pb files as a side effect (unlike tensorboard_plugin_profile). For the legacy trace to work with xprof, these files need to be checked in. They’ve been in the repo since PR #380.”

[gabeweisz]

If that is then case, how do we process new traces from Jax 0.6 starting from scratch if we complete this PR? We can't break Jax 0.6.0 functionality yet

[gabeweisz]

Do we need to install both xprof and tensorboard-plugin-profile 2.19.0 and import the correct library based on the trace? Or does the protobuf version break this?

[gphuang]

The legacy hlo_proto.pb files had caused the tests to pass with xprof. I will remove them.

The protobuf version actually breaks in this situation:

• xprof’s conversion path uses grpcio / grpcio-status, which pins protobuf ≥ 6.31.1
• tbp 2.19 cannot be used safely with protobuf 6.31.1, which is outside the supported range, causing descriptor error.

I did not yet find a safe solution to have both xprof2.20.7 and tbp2.19 in the same environment. What are you suggestions, which is the priority jax0.6 or 0.8? @devalshahamd @gabeweisz

[gphuang]

Current State

Tool	JAX 0.6	JAX 0.8	Protobuf
xprof 2.20.7	Trace loads; no full perf model (FLOPS/bytes) — doesn't extract .hlo_proto.pb	Full support	Requires ≥6.31.1 (grpcio-status)
tbp 2.19	Full support (extracts .hlo_proto.pb)	Crashes (INT_MAX in HLO ids)	Requires <5.0 (3.x–4.x)