fix(checkpointing): raise RuntimeError on fatal checkpointing errors

[RexBearIU]

Description

When checkpointing encounters an unexpected runtime or I/O failure, _checkpoint_error_handler previously converted the exception into exceptions.StopTraining(...).

In train.py, all StopTraining exceptions are treated as graceful completions (_job_completed_gracefully = True), which causes the Python process to exit with status code 0 (success). This silently masks fatal checkpointing failures from upstream orchestration tools (such as Airflow DAGs, SLURM scripts, and CI runners), causing subsequent pipeline stages (like evaluation or decode) to fail cryptically with missing checkpoint errors.

Fix

• Updated _checkpoint_error_handler in src/maxtext/common/checkpointing.py to raise RuntimeError(f"Checkpointing failed. {str(err)}") from err directly on unexpected errors.
• Elasticity signals (JaxRuntimeError, ScaleUpSignalError) continue to be bubbled up upstream via checkpoint_exception_guard.
• Fatal checkpoint errors now properly cause training to fail loudly with a non-zero exit code and full stack trace attached via from err.
• Added unit test coverage in tests/unit/checkpointing_test.py to verify RuntimeError raising and exception chaining (__cause__).

Tests

• Pre-commit linters and formatters:
  • codespell, pylint, pyink, mdformat, yamllint, actionlint (all passed on diff files)
• Unit test suites:
  • tests/unit/checkpointing_test.py (13 tests passed)

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.

[gemini-code-assist]

Code Review

This pull request modifies the checkpoint error handler in checkpointing.py to only raise a StopTraining exception when elastic training is enabled, and otherwise re-raise the original error. The reviewer suggests using elastic_utils.elastic_enabled(config) instead of checking config.elastic_enabled directly to ensure consistency and avoid unexpected behavior when the job is not running in an elastic environment.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!