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

src/maxtext/common/checkpointing.py

@@ -893,8 +893,10 @@ def maybe_save_checkpoint(checkpoint_manager, state, config, data_iterator, step
     return
 
   def _checkpoint_error_handler(err):
-    """Handles checkpointing errors, when not in an elastic context."""
-    raise exceptions.StopTraining(f"Checkpointing failed. {str(err)}") from err
+    """Handles checkpointing errors."""
+    if elastic_utils.elastic_enabled(config):
+      raise exceptions.StopTraining(f"Checkpointing failed. {str(err)}") from err
+    raise err
 
   with checkpoint_exception_guard(config, checkpoint_manager, _checkpoint_error_handler):
     checkpoint_saved = save_checkpoint(

tests/unit/checkpointing_test.py

@@ -33,6 +33,7 @@
 )
 from maxtext.common import checkpointing
 from maxtext.common import grain_utility
+from maxtext.utils import exceptions
 import numpy as np
 import optax
 import pytest
@@ -480,5 +481,59 @@ async def await_creation(self):
     self.assertEqual(iterators_restore_v1[1].state, 20)
 
 
+class CheckpointErrorHandlerTest(parameterized.TestCase):
+  """Tests for checkpoint error handling in maybe_save_checkpoint."""
+
+  def setUp(self):
+    super().setUp()
+    self.mock_manager = mock.MagicMock()
+    self.mock_manager.latest_step.return_value = None
+    self.mock_manager.reached_preemption.return_value = False
+    self.state = mock.Mock()
+
+  def test_error_handler_reraises_in_non_elastic_mode(self):
+    """Unexpected checkpointing errors should be re-raised when elastic mode is inactive."""
+    config = mock.Mock()
+    config.checkpoint_period = 1
+    config.pure_nnx = True
+    config.enable_diloco = False
+    config.async_checkpointing = False
+    config.enable_continuous_checkpointing = False
+    config.enable_emergency_checkpoint = False
+    config.enable_multi_tier_checkpointing = False
+    config.local_checkpoint_period = 0
+    config.enable_autocheckpoint = False
+    config.elastic_enabled = True
+
+    with (
+        mock.patch.object(checkpointing, "save_checkpoint", side_effect=RuntimeError("GCS failure")),
+        mock.patch.object(checkpointing.elastic_utils, "elastic_enabled", return_value=False),
+    ):
+      with self.assertRaisesRegex(RuntimeError, "GCS failure"):
+        checkpointing.maybe_save_checkpoint(self.mock_manager, self.state, config, data_iterator=None, step=1)
+
+  def test_error_handler_raises_stop_training_in_elastic_mode(self):
+    """Checkpointing errors should raise StopTraining when elastic mode is active."""
+    config = mock.Mock()
+    config.checkpoint_period = 1
+    config.pure_nnx = True
+    config.enable_diloco = False
+    config.async_checkpointing = False
+    config.enable_continuous_checkpointing = False
+    config.enable_emergency_checkpoint = False
+    config.enable_multi_tier_checkpointing = False
+    config.local_checkpoint_period = 0
+    config.enable_autocheckpoint = False
+    config.elastic_enabled = True
+
+    with (
+        mock.patch.object(checkpointing, "save_checkpoint", side_effect=RuntimeError("GCS failure")),
+        mock.patch.object(checkpointing.elastic_utils, "elastic_enabled", return_value=True),
+    ):
+      with self.assertRaises(exceptions.StopTraining) as cm:
+        checkpointing.maybe_save_checkpoint(self.mock_manager, self.state, config, data_iterator=None, step=1)
+      self.assertIn("Checkpointing failed. GCS failure", str(cm.exception))
+
+
 if __name__ == "__main__":
   absltest.main()

tests/unit/train_state_nnx_checkpoint_test.py

@@ -28,6 +28,7 @@
 import jax.numpy as jnp
 from maxtext.common import checkpointing
 from maxtext.common import train_state_nnx
+from maxtext.utils import exceptions
 import optax
 import orbax.checkpoint as ocp
 
@@ -624,6 +625,37 @@ def test_maybe_save_checkpoint_checks_scale_up_after_unsaved_dispatch(self):
     save_checkpoint_mock.assert_called_once()
     mock_maybe_scale_up.assert_called_once_with(config, mgr)
 
+  def test_maybe_save_checkpoint_error_handler_reraises_in_non_elastic_mode(self):
+    """Fatal checkpoint errors are re-raised when elastic mode is not active."""
+    state = mock.Mock()
+    config = self._config(checkpoint_period=1, elastic_enabled=True)
+    mgr = mock.MagicMock()
+    mgr.latest_step.return_value = None
+    mgr.reached_preemption.return_value = False
+
+    with (
+        mock.patch.object(checkpointing, "save_checkpoint", side_effect=RuntimeError("Disk I/O error")),
+        mock.patch.object(checkpointing.elastic_utils, "elastic_enabled", return_value=False),
+    ):
+      with self.assertRaisesRegex(RuntimeError, "Disk I/O error"):
+        checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=5)
+
+  def test_maybe_save_checkpoint_error_handler_raises_stop_training_in_elastic_mode(self):
+    """Checkpoint errors raise StopTraining when elastic mode is active."""
+    state = mock.Mock()
+    config = self._config(checkpoint_period=1, elastic_enabled=True)
+    mgr = mock.MagicMock()
+    mgr.latest_step.return_value = None
+    mgr.reached_preemption.return_value = False
+
+    with (
+        mock.patch.object(checkpointing, "save_checkpoint", side_effect=RuntimeError("Disk I/O error")),
+        mock.patch.object(checkpointing.elastic_utils, "elastic_enabled", return_value=True),
+    ):
+      with self.assertRaises(exceptions.StopTraining) as cm:
+        checkpointing.maybe_save_checkpoint(mgr, state, config, data_iterator=None, step=5)
+      self.assertIn("Checkpointing failed. Disk I/O error", str(cm.exception))
+
 
 class TestLinenCheckpointFormatConverters(unittest.TestCase):
   """to_linen_checkpoint_dict / from_linen_checkpoint_dict (NNX <-> Linen on-disk layout)."""