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.

(Note: Re-created from PR #4290 to reset Copybara import status)