Use rolling window to track goodput metrics over time

[lydhr]

Description

Enables rolling window Goodput tracking in MaxText to monitor performance trends over configurable intervals (1h, 1d, 3d, 5d). This feature allows for identifying performance degradation or improvement trends within specific time windows, supplementing the cumulative goodput metrics.

• Added enable_rolling_window_goodput flag to control this feature (default true). Extended rolling_windows_seconds to support 1h, 1d, 3d, and 5d windows.
• Integrated start_rolling_window_goodput_uploader in maybe_monitor_goodput context manager in goodput.py. Ensured proper lifecycle management (start/stop) to prevent process leaks.

BUGS: b/556291470

Tests

Verified with test runs remotely on TPU VMs. Confirmed metrics are generated and sent to Google Cloud Monitoring for all configured windows.

Command:

python3 -m src.maxtext.trainers.pre_train.train \
    src/maxtext/configs/base.yml \
    run_name=ly_202609270731\
    enable_jax_profiler=True\
    profiler=xplane \
    steps=200 \
    dataset_path=${GS_BUCKET} \
    dataset_name=c4/en:3.0.1 \
    eval_dataset_name=c4/en:3.0.1 \
    base_output_directory=${GS_BUCKET}/output/goodput/202609270731 \
    enable_goodput_recording=true \
    monitor_goodput=true \
    enable_rolling_window_goodput=true \
    enable_checkpoint_cloud_logger=true \

• Log Explorer: Direct Link
• Metrics Explorer: Direct Link
• Dashboards: Direct Link, screenshot of sliding window figures

Checklist

Reference:

• [axlearn/cloud/gcp/measurement.py]

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.

[gemini-code-assist]

Code Review

This pull request introduces rolling window goodput monitoring to MaxText. It adds a new context manager maybe_monitor_rolling_window_goodput in goodput.py and exposes corresponding configuration options (rolling_windows_seconds and enable_rolling_window_goodput) in base.yml and types.py. The feedback recommends tracking whether the rolling window uploader successfully started before stopping it in the finally block to prevent masking initialization exceptions, and suggests using list[PositiveInt] in the Pydantic configuration to enforce positive integer intervals.

[codecov]

Codecov Report

❌ Patch coverage is 14.28571% with 24 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/common/goodput.py	8.33%	22 Missing ⚠️
src/maxtext/common/gcloud_stub.py	50.00%	2 Missing ⚠️

📢 Thoughts on this report? Let us know!

[lydhr]

Re: "Track Test Duration" failure. This isn't caused by this PR.

The failing check only compares TPU unit-test durations against the last scheduled main run, and it flags a different, mostly non-overlapping set of tests every run.

That includes main itself: the scheduled main run (commit 4919941, without this PR) flagged 8 of these regressions. It passed only because scheduled main runs use --warn-only, while PR runs fail on the same alerts.

This PR only touches goodput.py, base.yml and types.py, and goodput monitoring is off in unit tests (monitor_goodput: false). The flagged tests are MoE/attention tests unrelated to this change. They randomly take about 15–110s longer from run to run, which points to TPU runner or compile-time variance.

Test (TPU unit, seconds)	main (Baseline used by this PR)	This PR
tpu-unit tests ran (UTC)	2026-09-30 16:09–16:39 (saved as baseline gh-pages@03db447)	2026-09-30 19:46–20:18
Commit tested	4919941 (main)	dc0ef22 (PR merge)
test_prefuse_moe_weights_matches_unfused	168.3	167.9
test_gmm_grad_equivalence_tokamax_v2_fp8_dynamic_ep1	136.2	27.1
test_moe_quantize_combine_bwd_method_rowwise	189.6	69.6
..._ring_context_parallel_grad_packed	38.7	4.6
test_gmm_grad_equivalence_megablox_fp8_dynamic_ep1	25.0	135.2
test_ragged_sort_loss_and_grad_no_ring_of_experts	42.1	125.8
test_gmm_grad_equivalence_tokamax_v2_fp8_static_ep4	32.6	123.5
Total tests flagged	8 (warn-only → pass)	11 (fail)

Bold = flagged (more than 20% and more than 15s over the baseline). Each scheduled main run overwrites the baseline with its own raw durations. So this PR was compared directly against the 16:39 main run, including the tests that were randomly slow in that run.

[dipannita08]

Overall, LGTM!

[dipannita08]

Left a couple small comments, LGTM otherwise!