Skip to content

Fix configuration metadata in stitched restart output - #2521

Open
AdemVessell wants to merge 1 commit into
google-deepmind:mainfrom
AdemVessell:codex/restart-config-provenance
Open

AdemVessell wants to merge 1 commit into
google-deepmind:mainfrom
AdemVessell:codex/restart-config-provenance

Conversation

@AdemVessell

@AdemVessell AdemVessell commented Sep 27, 2026 •

Copy link
Copy Markdown

Thanks @jcitrin for taking the time to reproduce this and for the encouraging response in #2520. I appreciated the clear guidance on what would be useful here, and I'm glad to contribute this fix.

Stitched restarts currently retain the original file's config attribute, losing the resumed run's configuration. This change keeps the resumed configuration in config and saves the prior file's configuration as previous_config.

previous_config describes the immediately preceding run, not a complete history of chained restarts. It is omitted if the prior file lacks config. The change is confined to the stitching boundary; time concatenation and overlap handling are unchanged. The output documentation describes these attributes.

Validation on upstream base b53645a, macOS 27.0.1 arm64 / Python 3.12.13:

  • Five focused regression cases fail against unchanged production and pass with the fix. They cover normal and legacy prior metadata, stale and older provenance, repeated restarts, overlap handling, preservation of the actual input file, and NetCDF round trips.
  • All 35 tests in the current output and stitch test files pass with TORAX_ERRORS_ENABLED=1 and TORAX_JAXTYPING=1; all 8 output-key tests also pass.
  • The regressions now live in output_stitch_test.py because upstream refactored the original output test file. Production behavior remains the same as the earlier patch.

The earlier 12-shard Ubuntu/Python 3.12 CI pass belongs to 5f547126; it does not validate this updated revision. The full repository suite and Linux CI for the updated head have not been run locally. Code and test preparation were AI-assisted with Codex.

@AdemVessell
AdemVessell force-pushed the codex/restart-config-provenance branch from 5f54712 to 7ec22ed Compare November 15, 2026 09:26
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant