[mesakhcienet]

Description

Background

The "Track Test Duration" CI job checks whether a PR makes tests slower. It compares the PR's test times with the times from the last scheduled run on main. It had two problems:

1. Adding tests could fail the job. It also failed when the number of tests in a group grew by more than 15%, even if the new tests were fast. In 28 PR runs on Oct 29–30, 10 failed only for this reason.
2. Slow tests in some files were never caught. Seven test files were set to "warning only", because their times used to jump depending on which machine ran them. fix(ci): keep TPU unit test classes on one shard and flag partial results #5565 (already merged) stopped most of these jumps by keeping each test class on one machine, but the files were still warning-only.

What this PR changes

1. New tests are judged by time, not by count. A PR fails only if a test group gets more than 20% slower and more than 5 minutes longer, counting the new tests' time. The job log lists the slowest new tests, so the author knows what to fix.
2. Only PRs fail. Scheduled, release and manual runs only warn. Only scheduled runs on main save the new test times and update the dashboard.
3. Slow tests are caught again in 4 files. Their times no longer jump, so a slow test in them fails the PR again. 3 files stay warning-only because their times still jump.
4. All test machines split the tests the same way. Each machine needs the saved test times to know its share of the tests. Before, a machine that couldn't download them quietly used a different split, so some tests ran twice and others not at all. Now that machine fails instead. The saved times are kept 30 days, so re-running a failed machine later still works.

Details for reviewers

• Files: tools/dev/process_test_results.py, tests/utils/class_split.py, .github/workflows/track_performance.yml, run_tests_coordinator.yml, run_tests_against_package.yml, and the tests.
• New check: for each test group (e.g. tpu-unit), it compares the time of all passed tests with the saved time of the same existing tests. The limits are SUITE_RATIO_LIMIT = 1.20 and SUITE_FLOOR_SEC = 300. New scheduled_only tests, CPU groups and incomplete runs are skipped. A renamed test counts as new until the next scheduled run. The check runs inside process_testcase(), in the same pass as the per-test check.
• Enforced again: megablox_gmm_3d_test, megablox_relayout_test, m3.qwen3_vl_comparison_test, tiling_test. Still warning-only: moe_test, attention_test, integration.pipeline_parallelism_test.
• Which runs do what (track_performance.yml): every event except pull_request passes --warn-only. Only schedule on refs/heads/main passes --save-baseline and --output-benchmark, so the baseline push and the dashboard step only run there. The dashboard step comments on an alert and doesn't fail the run.
• Test durations download: the coordinator now tells each shard whether it uploaded .test_durations (test_durations_uploaded). If it did, the shard's download step is required (no more continue-on-error). If the coordinator had no baseline, every shard uses the same fallback split, as before. The artifact is kept 30 days instead of 1 (GitHub's re-run limit): a re-run shard downloads it again, and with 1 day a re-run the next day would now fail.
• Smaller changes:
  • the saved baseline is filtered once when it is loaded (<flavor>::<test> keys with a number ≥ 0);
  • the [class-split] log line prints on its own line.
• Data: normal PR runs stay within about ±5% of the saved times. When [TC sort 4/4] All-gather / combine MoE tokens in the TC kernels' 3D layout #5507 (which added many slow kernel tests) merged, the tpu-unit total rose 1.28x (+1,494 s), well past both limits.

Tests

• Unit tests: all 105 tests in tests/dev pass.
• Full CI run with TPU/GPU tests: run 41765923851. All tests passed. The job reported no slowdown ([SUITE] lines from 0.95x to 1.02x, no new tests) and, as a PR run, skipped the baseline push and the dashboard step.
• Replay on real CI results:
  • the 10 PR runs that failed only because of the test count now pass;
  • on 120 replays of 30 real runs (with added slow tests and new tests), the job gives the same output as the previous version of this PR.
• Workflow step, run locally: it passes the right options for PR, scheduled, release and manual runs.

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 and provided workload links above if applicable. (Note: New PRs get skip-ci-tests so only CPU/lint/docs checks run during review; a bot removes it after 2 approvals on a 1-commit PR to run TPU/GPU suites, or remove it yourself / ask a maintainer on fork PRs to run them sooner.)
• 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.

[github-actions]

🤖 Hi @mesakhcienet, I'm running an automatic MaxText review of this pull request. You can track my progress in the logs. To request another review later, add the gemini-review label or comment @gemini-cli /review.

[gemini-code-assist]

Summary of Changes

Hello, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request updates the CI test regression logic to shift from test-count-based alerts to a time-based suite growth check. This change addresses false positives caused by fluctuations in test counts and improves the reliability of performance monitoring. Additionally, it re-enables per-test alerts for specific modules that have been stabilized by the class-level shard split and introduces stricter data validation for test duration inputs.

Highlights

• Suite-level regression check: Implemented a new suite-level check in process_test_results.py that evaluates PR test growth based on total duration rather than test count, failing PRs if they exceed 1.20x the baseline time and 300s of absolute growth.
• Re-enabled per-test alerts: Re-enabled per-test regression alerts for four modules (megablox_gmm_3d_test, megablox_relayout_test, m3.qwen3_vl_comparison_test, and tiling_test) as the class-level shard split now stabilizes their execution times.
• Improved test duration processing: Enhanced load_durations and process_test_results.py to robustly handle non-numeric, NaN, infinite, or boolean values in duration files, preventing skew in regression metrics.
• Duplicate shard detection: Added logic to detect and warn about overlapping shards where tests are reported multiple times, treating such flavors as missing shards to avoid inaccurate dashboard reporting.

New Features

🧠 You can now enable Memory (public preview) to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Ignored Files

• Ignored by pattern: .github/workflows/** (1)
  • .github/workflows/track_performance.yml

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize the Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counterproductive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for GitHub and other Google products, sign up here.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[github-actions]

📋 Review Summary

Verdict: Spec met; robust implementation with excellent test coverage. One minor edge case found.

Top issues:

1. [Bug] add_suite_sample omits the base_time > 0 validation used elsewhere. If a baseline entry is NaN or negative, it poisons the flavor's base time and silently disables the suite check for that flavor. → Add and base_time > 0 to the condition (see inline comment).

Inferred spec & validation

• Sources: PR description, code before/after.
• Problem: PR test growth was judged by test count, which failed PRs when many quick tests were added, causing false positive regressions on CI.

Spec item	Implementation	Validating test
R1 (stated): PR test growth is judged by time, comparing passed tests to the baseline of existing tests	✅	test_suite_check_fails_a_pr_whose_new_tests_grow_a_flavor_past_both_limits ✅
R2 (stated): PR fails if flavor time grows > 1.20x AND > 300s longer	✅	test_suite_check_fails_a_pr_whose_new_tests_grow_a_flavor_past_both_limits ✅
R3 (stated): New scheduled_only tests are not counted towards PR suite regression	✅	test_suite_check_leaves_out_new_scheduled_only_tests ✅
R4 (stated): CPU, decoupled, partial, corrupt, and overlapping flavors are skipped	✅	test_suite_check_skips_cpu_and_partial_flavors_and_ignores_failed_tests ✅
R5 (stated): Per-test alerts are enforced again for 4 modules	✅	test_process_testcase_regression_enforced_for_megablox_gmm_3d_test ✅
C1 (stated): Dashboard keeps tracking test counts but does not fail PRs	✅ preserved	N/A (Checked in GitHub workflow)

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[SurbhiJainUSC]

Pass --warn-only for all non-PR events so release/manual runs log regressions without failing the pipeline. Also, only generate new_per_test_baseline.json and benchmark-results.json on schedule runs on refs/heads/main.

EXTRA_ARGS=()
if [ "$EVENT_NAME" != "pull_request" ]; then
  EXTRA_ARGS+=("--warn-only")
fi
if [ "$EVENT_NAME" = "schedule" ] && [ "$REF" = "refs/heads/main" ]; then
        EXTRA_ARGS+=(
          "--save-baseline" "new_per_test_baseline.json"
          "--output-benchmark" "benchmark-results.json"
        )
fi
python3 tools/dev/process_test_results.py test-results \
--baseline per_test_baseline.json \
"${EXTRA_ARGS[@]}"

[mesakhcienet]

Done, with your snippet. Every non-PR event now passes --warn-only, and only scheduled runs on main write new_per_test_baseline.json and benchmark-results.json. The dashboard step needs benchmark-results.json, so it now only runs on those runs. I kept fail-on-alert: false there so those runs only warn too; an alert still adds a commit comment.

[SurbhiJainUSC]

Can we simplify this by dropping the duplicate-shard detection?

Shards only overlap when the coordinator uploads the durations artifact and one worker fails to download it while other workers succeed. When that happens, other tests are silently skipped so we shouldn't let that test shard pass with a fallback split and try to filter out the aftermath during result processing.

Instead of adding extra post-processing logic and edge cases here, let's remove the duplicate-test handling from this PR and make the test runner require the durations download whenever the coordinator uploaded it so all shards are guaranteed to use the same split.

[mesakhcienet]

Agreed, removed. The coordinator now outputs test_durations_uploaded (true when its upload step created the artifact) and passes it to run_tests_against_package.yml. Shards download the durations only in that case, and without continue-on-error. So a failed download fails the shard, instead of letting it use a different split. If the coordinator had no baseline, every shard uses the same fallback split, as before.

One related change: a re-run shard downloads the durations again, so with a 1-day retention a re-run the next day would now fail. The artifact is now kept 30 days, which is how long GitHub allows re-runs.

[SurbhiJainUSC]

Can we simplify this by removing the extra checks and tests for unrealistic duration values like booleans, NaN, infinity, negative test times, and 400-digit numbers? Both the XML test reports and the baseline JSON are generated automatically by pytest and this script, so test durations in CI are always normal non-negative numbers. If we still want a basic safety check when reading the baseline JSON, let's just filter out invalid entries once right when the file is loaded rather than adding separate validation helpers and edge-case tests across multiple files.

[mesakhcienet]

Done. I removed baseline_seconds(), testcase_seconds(), the extra filtering in class_split.load_durations and the edge-case tests. Only one check is left: when the baseline JSON is loaded, it keeps the <flavor>::<test> keys whose value is a number ≥ 0. This replaces the old isinstance filter, which only applied to the saved copy. Test times from the XML are read with plain float() again.