[hsuan-lun-chiang]

Description

This PR implements a comprehensive automated mechanism to detect and prevent test suite run-time inflation before code is merged into main, addressing the need for a fast, reliable, and responsive CI pipeline.

How the Workflow Works:

When a PR is opened or updated, the pipeline runs the standard test suites (CPU, GPU, TPU) and generates JUnit XML reports. Once those tests complete, the new track_performance job runs and does the following:

1. Micro-Level Tracking: It runs enforce_test_limits.py to parse the XML reports and ensures no individual test function regresses significantly compared to its stored baseline. The fine-grained baseline (per_test_baseline.json) is maintained directly on the gh-pages branch.
2. Macro-Level Tracking: It runs parse_junit_to_benchmark.py to aggregate the total runtime for each device. It then uses github-action-benchmark to compare these aggregated total runtimes against the baseline stored on the gh-pages branch to prevent overall suite inflation.

Changes Made:

1. Macro-Level Suite Tracking (gh-pages):

   • Integrated github-action-benchmark to track macro-level execution times (suite/shard level).
   • Added tests/utils/parse_junit_to_benchmark.py to convert JUnit XML test results into the JSON format required by the benchmark action.
   • Results are automatically pushed to the gh-pages branch for historical trend visualization.

2. Micro-Level Individual Test Tracking:

   • Uses a custom, highly reliable Python-based baseline tracking mechanism.
   • Saves a fine-grained per_test_baseline.json dictionary containing the execution times of individual test functions to the gh-pages branch.

3. Regression Enforcement:

   • Added tests/utils/enforce_test_limits.py to automatically flag and block any test that significantly regresses compared to its stored per-test baseline.
   • Relative Regression Limits: Any test that is 50% slower (for Train Compile tests) or 20% slower (for Unit and Integration tests) than its baseline AND exceeds its absolute noise threshold (+15s for Unit, +120s for Train Compile, +30s for Integration) will be blocked. This protects against JAX JIT compilation jitter on shared CI runners while keeping non-compile tests tightly constrained.

Tests

• Validated macro-level tracking via github-action-benchmark in CI.
• Validated micro-level test baseline generation, caching, and evaluation in CI.
• Tested enforce_test_limits.py in CI to ensure that test jitter thresholds are accurately tuned.

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[github-actions]

⚠️ Performance Alert ⚠️

Possible performance regression was detected for benchmark 'MaxText Test Execution Times'.
Benchmark result of this commit is worse than the previous benchmark result exceeding threshold 1.15.

Benchmark suite	Current: 83666c2	Previous: 7b61908bf7e0a90e8b8a44b19d9cfa203ae308a6	Ratio
Total CPU-UNIT Tests Duration	9845.751000000002 sec	7476.170000000001 sec	1.32
Total CPU-TORCH-REFERENCE Tests Duration	283.71700000000004 sec	246.48799999999994 sec	1.15

This comment was automatically generated by workflow using github-action-benchmark.

[hsuan-lun-chiang]

Hi @shralex,
I wanted to share an update from our workflow testing. We encountered a Track Test Performance failure that perfectly highlights the volatile runner noise we discussed:

• test_moe_deepseek_with_device_limit spiked from 27.89s to 183.99s (+156s / +559.7%) on the shared CPU runner.
• test_moe_deepseek_scanned_bf16 spiked from 112.11s to 236.90s (+124s / +111.3%).

To proceed, we see two main directions. How would you prefer we approach this?

Option 1: Skip regression tracking for train_compile_test

Option 2: Stabilize compile test runtimes on runners

If we want to keep compile tests in our performance tracking, we need to mitigate XLA's compilation volatility. Here are the main ways we can do this (Suggested by Gemini):

A. Enable Persistent XLA Compilation Caching

• Pros: Skips compilation entirely on cache hits, reducing subsequent test runs to sub-seconds.
• Cons: Cache keys are highly sensitive. Any dependency updates (JAX/jaxlib) or config/shape changes will trigger cache misses, re-exposing the test to compilation spikes. Also introduces large binary storage overhead in GitHub cache.

B. Enforce Runner Thread Limits (Thread-Thrashing Mitigation)

• Pros: Restricts OMP_NUM_THREADS and MKL_NUM_THREADS to 2 to prevent XLA from over-subscribing threads and thrashing CPU caches on multi-tenant hosts.
• Cons: Hardcoding low thread limits will throttle execution speeds on larger self-hosted runners or local developer workstations unless configured dynamically.

C. Measure Pure Execution (Warm Up Step 0)

• Pros: Performs a warmup step first to compile the graph, then times subsequent steps (1 +). This isolates the LLVM/XLA compilation overhead completely from our performance metrics.
• Cons: Requires minor refactoring of the pytest timing wrappers to separate Step 0 from the main loop.
──────
Would love to hear your thoughts on which direction fits best.