fix(checkpointing): re-raise fatal checkpointing errors in non-elastic mode

[RexBearIU]

Description

In standard (non-elastic) training, 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 check if elastic_utils.elastic_enabled(config): before raising exceptions.StopTraining. This ensures that StopTraining is only raised when elastic training is actually active (both Pathways backend used and config.elastic_enabled enabled).
• In non-elastic mode, unexpected checkpoint errors are now re-raised directly, ensuring the job fails with a non-zero exit code and displays the actual stack trace.
• Added comprehensive unit test coverage in tests/unit/checkpointing_test.py and tests/unit/train_state_nnx_checkpoint_test.py to verify error re-raising in non-elastic mode and StopTraining raising in elastic mode.

Tests

• Pre-commit linters and formatters:
  • codespell, pylint, pyink, mdformat, yamllint, actionlint (all passed)
• Unit test suites:
  • python3 -m pytest tests/unit/checkpointing_test.py (14 tests passed)
  • python3 -m pytest tests/unit/train_state_nnx_checkpoint_test.py (22 tests passed)
  • python3 -m pytest tests/unit/elastic_utils_test.py (43 tests passed)
  • All checkpoint suites (162 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.

[gemini-code-assist]

Using config.elastic_enabled directly can lead to unexpected behavior if the configuration flag is set to True but the job is not actually running in an elastic environment (e.g., when the Pathways backend is not used). To ensure that we only raise StopTraining when elastic training is actually active, we should use elastic_utils.elastic_enabled(config) instead. This also keeps the behavior consistent with other elastic checks in the codebase.

Suggested change

	"""Handles checkpointing errors, when not in an elastic context."""
	raise exceptions.StopTraining(f"Checkpointing failed. {str(err)}") from err
	if config.elastic_enabled:
	raise exceptions.StopTraining(f"Checkpointing failed. {str(err)}") from err
	raise err
	"""Handles checkpointing errors."""
	if elastic_utils.elastic_enabled(config):
	raise exceptions.StopTraining(f"Checkpointing failed. {str(err)}") from err
	raise err

[SurbhiJainUSC]

This is a valid suggestion. Can you please do this?

[RexBearIU]

Done! Updated _checkpoint_error_handler to use elastic_utils.elastic_enabled(config) and added comprehensive unit test coverage in tests/unit/checkpointing_test.py and tests/unit/train_state_nnx_checkpoint_test.py to verify error propagation in both elastic and non-elastic modes.

[codecov]

Codecov Report

❌ Patch coverage is 0% with 3 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
src/maxtext/common/checkpointing.py	0.00%	3 Missing ⚠️

📢 Thoughts on this report? Let us know!