Add runtime trace quality checks

[kyle-hoffmeyer]

Summary

This PR introduces a unified trace quality check system that runs automatically when traces are loaded and after reports are generated. Basic checks run on every trace load. More detailed checks run when TRACELENS_DETAILED_QUALITY_CHECKS=1 is set. All pre-report checks can be skipped with TRACELENS_SKIP_QUALITY_CHECK=1.

The system builds upon examples/custom_workflows/trace_checker.py and moves it into core TraceLens at TraceLens/trace_check.py. To easily see the changes, TraceLens/trace_check.py has been copied back over into examples/custom_workflows/trace_checker.py.

Once this PR is merged, examples/custom_workflows/trace_checker.py can be deleted and replaced by TraceLens/trace_check.py. It is simply left here to make reviewing easier.

Implementation

Previously, every module was loading the raw trace via DataLoader.load_data(filepath) and then manually extracting the trace events. This has been replaced via the new DataLoader.load_trace_events function that extracts the events and metadata in one call.

load_trace_events then calls run_trace_checks(), ensuring that trace quality checks always run when trace events are loaded. run_trace_checks is also called after report generation where it runs post-report checks on the generated DataFrames and the built tree.
run_trace_checks emits errors and warnings to stderr. For report writing tools, a new file called trace_checks.json is also created in the same directory as the report output.

How run_trace_checks works

run_trace_checks is the single entry point for both pre-report and post-report checks. It accepts:

Parameter	Description
events	The traceEvents list from the loaded trace.
capture_trace_filepath	Path to a graph-capture trace folder, if provided.
filepath	Path to the trace file on disk.
framework	"pytorch", "jax", or None. Controls which framework-specific checks run.
pre_report	If True (default), runs pre-report checks on raw events. If False, runs post-report checks on the built tree and generated DataFrames.
tree	A built TraceToTree object.
dfs	Dict of report DataFrames
attn_cfg	Optional AttnCountConfig

The function iterates a registry of all the trace checks that exist. The registry contains tuples tracking the scope of each check:

1. Phase filter — Notes if its a pre-report check (operates on raw trace events) or post-report check (operates on processed trace).
2. Detail filter — Notes if the check is detailed and potentially time consuming. These checks are skipped unless TRACELENS_DETAILED_QUALITY_CHECKS=1 is set.
3. Framework filter — Notes the framework the test is built for (pytorch, jax, rocprof).

Every check has a function. Each function accepts a single ctx dict containing all the parameters, so the dispatch loop has no per-function branching — it just calls fn(ctx) for each matching entry. Each check returns a List[TraceCheck] (empty if the check passes), and the results are collected into a TraceCheckReport. The findings of TraceCheckReport are stderr via the log_findings() function.

For descriptions of each check, see docs/reference/trace-quality.md.

trace_checker.py deletions

Much of trace_checker.py has been deleted to limit the scope of the tool as something that runs during normal TraceLens operation rather than as a standalone CLI tool. The following was removed during refactoring:

• CLI code: The checks are now invoked programmatically via run_trace_checks() during normal TraceLens operation, not as a standalone command-line tool. Code supporting the CLI invocation have been deleted.

• Framework-specific runner functions (run_pytorch_phase1, run_jax_phase1, _pytorch_busy_idle, _jax_busy_idle): These orchestrated which checks to run per framework. Replaced by the framework field in the registry that filters checks automatically.

• run_quality_checks orchestrator: The old top-level function that loaded traces, ran phases, and caught exceptions. Replaced by run_trace_checks() which iterates the registry.

• Phase 2 report generation (generate_report, find_megatron_extension): These generated perf reports as part of the checker's Phase 2. Now the report generators themselves call run_trace_checks(pre_report=False) after they have already generated the report.

Checks removed

Several checks from the original trace_checker.py were removed after evaluation against the test trace corpus (see docs/trace_quality_eval.md):

• _check_kernels_dropped_windowed — bucketed the kernel timeline into 20 windows and flagged interior windows with fewer kernels than 50% of the median. Produced false positives on small traces where natural variance in per-window counts triggered the threshold. Replaced by the correlation-ID check (_check_kernels_dropped_correlation) which directly verifies whether each runtime launch has a matching GPU kernel.

• _check_kernel_counts — reported unique kernel count and warned when no attention/FMHA kernels were detected. Model-specific — non-transformer workloads have no attention kernels.

• _check_sdpa_count — warned when no SDPA/attention ops appeared in the generated report. Same issue as _check_kernel_counts — assumes every workload is a transformer.

• _check_op_count_consistency — computed the GCD of operation counts and warned when GCD=1 (no common per-step multiple). Fires on single-pass inference or short profiling windows. The check is wrong for the common case.

• _check_report_generated — verified expected sheets existed in the generated report DataFrames. Redundant since if report generation succeeds, the sheets are always present. The agent triage system already checks for report files on disk separately.

Changes to TraceLens Agent Triaging

The TraceLens Agent triage system previously implemented several trace quality checks that are now handled by trace_check. Three triage checks have been updated to delegate to trace_check instead of duplicating the logic:

• check_trace_size — previously computed file size and compared against hardcoded thresholds inline. Now calls trace_check._check_trace_file_size() and converts the finding to a FindingDraft.
• check_no_gpu_kernels — previously scanned events for cat == "kernel" inline. Now calls trace_check._check_kernels_present().
• check_missing_cpu_op_shapes — previously counted cpu_ops with Input Dims inline. Now calls trace_check._check_cpu_op_shapes_basic().

The remaining triage checks (check_trace_missing, check_capture_missing, check_high_idle, check_corrupt_json, check_runtime_instability, etc.) are unchanged as they check agent-pipeline-specific artifacts (directory structure, CSV outputs, agent stream logs) that are outside the scope of trace_check.

Tests

tests/test_trace_check.py — 22 tests across 8 test classes:

• TestKernelPresence (2 tests) — verifies that an empty event list produces an error-level finding and that events with GPU kernels produce no finding.
• TestTraceFileSize (3 tests) — verifies that a normal-sized file produces no finding, a small file (< 100 KB) produces a warning, and a missing file produces no finding.
• TestCpuOpShapes (4 tests) — verifies that cpu_ops with Input Dims produce no finding, cpu_ops without shapes produce a warning, traces without cpu_ops produce no finding, and the check is skipped for JAX traces (framework filtering).
• TestKernelDropDetection (4 tests) — builds a TraceToTree with runtime launch events, ac2g flow events, and GPU kernels linked by correlation IDs. Verifies that matched launches produce no finding, a launch with no matching kernel produces a warning, traces without runtime events produce no finding, and the warning message includes the drop count.
• TestCallStackMissing (2 tests) — verifies that python_function events suppress the warning and their absence triggers it.
• TestGraphMode (3 tests) — verifies no finding without graph launches, a warning when graph launches exist without a capture trace, and no finding when a capture trace path is provided.
• TestBypass (2 tests) — verifies that TRACELENS_SKIP_QUALITY_CHECK=1 suppresses all warnings during DataLoader.load_trace_events(), and that unsetting it causes the kernels_present error to raise.
• TestLogFindings (2 tests) — verifies that warn-level findings emit UserWarning and error-level findings raise ValueError.

tests/test_analysis_agent_triage.py — updated existing tests to work with the new trace_check module. Tests that create traces without GPU kernels now set TRACELENS_SKIP_QUALITY_CHECK=1 via monkeypatch since load_trace_events raises on zero kernels. Message assertions updated to match new trace_check message strings.

[codecov-commenter]

Codecov Report

❌ Patch coverage is 97.47706% with 11 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
TraceLens/trace_check.py	98.01%	7 Missing ⚠️
TraceLens/Reporting/reporting_utils.py	71.42%	2 Missing ⚠️
TraceLens/TreePerf/jax_analyses.py	33.33%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

[spandoesai]

LGTM, tracediff changes are very minimal and only involve renaming the import variable

[gabeweisz]

this should be a generic bool parser - take "1", "true", "True", etc. There should be a standard utility function if we don't have one, something like

def parse_bool(value):
return str(value if value is not None else "").lower() in ('true', '1', 't')

[kyle-hoffmeyer]

Thanks. Added a string bool parser

[gabeweisz]

Looks good other than the parser

[kyle-hoffmeyer]

TODO: Add warnings and errors to some kind of JSON as well. Not just stderr