Internal.

PiperOrigin-RevId: 935681934

Author: JustinPan-goog

checkpoint/orbax/checkpoint/_src/path/atomicity_defaults.py

@@ -23,6 +23,7 @@
 from orbax.checkpoint._src.path import atomicity
 from orbax.checkpoint._src.path import atomicity_types
 from orbax.checkpoint._src.path import gcs_utils
+from orbax.checkpoint._src.path import utils
 
 
 def get_item_default_temporary_path_class(

checkpoint/orbax/checkpoint/_src/path/deleter.py

@@ -34,6 +34,7 @@
 from orbax.checkpoint._src.path import step as step_lib
 from orbax.checkpoint._src.path import utils as path_utils
 
+
 urlparse = parse.urlparse
 PurePosixPath = pathlib.PurePosixPath
 
@@ -403,6 +404,7 @@ def __init__(
         num_threads=num_threads,
     )
     self._delete_queue = queue.Queue()
+    self._exception = None
     # Turn on daemon=True so the thread won't block the main thread and die
     # when the program exits.
     self._delete_thread = threading.Thread(
@@ -417,10 +419,14 @@ def __init__(
   def _delete_thread_run(self) -> None:
     logging.info('Delete thread has started.')
     while True:
-      step = self._delete_queue.get(block=True)
-      if step < 0:
+      try:
+        step = self._delete_queue.get(block=True)
+        if step < 0:
+          break
+        self._standard_deleter.delete(step)
+      except Exception as e:  # pylint: disable=broad-exception-caught
+        self._exception = e
         break
-      self._standard_deleter.delete(step)
     logging.info('Delete thread exited.')
 
   def delete(self, step: int) -> None:
@@ -437,6 +443,9 @@ def close(self) -> None:
       self._delete_queue.put(-1)
       self._delete_thread.join()
 
+    if self._exception:
+      raise self._exception
+
     self._standard_deleter.close()
 
 

checkpoint/orbax/checkpoint/_src/path/deleter_test.py

@@ -39,9 +39,7 @@ def _get_save_diretory(self, step: int, directory: epath.Path) -> epath.Path:
       threaded=(False, True),
       todelete_subdir=(None, 'some_delete_dir'),
   )
-  def test_checkpoint_deleter_delete(
-      self, threaded, todelete_subdir
-  ):
+  def test_checkpoint_deleter_delete(self, threaded, todelete_subdir):
     """Test regular CheckpointDeleter."""
     deleter = deleter_lib.create_checkpoint_deleter(
         self.ckpt_dir,
@@ -158,6 +156,30 @@ def test_checkpoint_deleter_auto_threads_gcs(self):
     self.assertEqual(num_threads, 1)
 
 
+
+class ThreadedCheckpointDeleterExceptionTest(parameterized.TestCase):
+
+  def test_exception_bubbling_on_close(self):
+    """Verifies exceptions in background thread are raised on close()."""
+    mock_standard_deleter = mock.MagicMock()
+    error_msg = 'Background deletion failed'
+    mock_standard_deleter.delete.side_effect = ValueError(error_msg)
+
+    with mock.patch.object(
+        deleter_lib,
+        'StandardCheckpointDeleter',
+        return_value=mock_standard_deleter,
+    ):
+      deleter = deleter_lib.ThreadedCheckpointDeleter(
+          epath.Path('/tmp'),
+          name_format=step_lib.standard_name_format(),
+      )
+      deleter.delete(1)
+
+      with self.assertRaisesRegex(ValueError, error_msg):
+        deleter.close()
+
+
 class GcsRenameTest(unittest.TestCase):
 
   @mock.patch('orbax.checkpoint._src.path.deleter.epath.Path')
@@ -203,5 +225,6 @@ def test_gcs_rename_logic_directly(self, mock_epath_constructor):
     # Verify the Rename was actually called
     mock_step_path.rename.assert_called_with(mock_final_dest)
 
+
 if __name__ == '__main__':
   absltest.main()