diff --git a/checkpoint/orbax/checkpoint/_src/path/atomicity_defaults.py b/checkpoint/orbax/checkpoint/_src/path/atomicity_defaults.py index 4ae0ccfe9..3a2b065eb 100644 --- a/checkpoint/orbax/checkpoint/_src/path/atomicity_defaults.py +++ b/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( diff --git a/checkpoint/orbax/checkpoint/_src/path/deleter.py b/checkpoint/orbax/checkpoint/_src/path/deleter.py index 4db34865d..f5a94e0d1 100644 --- a/checkpoint/orbax/checkpoint/_src/path/deleter.py +++ b/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() diff --git a/checkpoint/orbax/checkpoint/_src/path/deleter_test.py b/checkpoint/orbax/checkpoint/_src/path/deleter_test.py index c0c77c64b..c234ebfd2 100644 --- a/checkpoint/orbax/checkpoint/_src/path/deleter_test.py +++ b/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() diff --git a/checkpoint/orbax/checkpoint/_src/path/step_test.py b/checkpoint/orbax/checkpoint/_src/path/step_test.py index 41afee71e..33145239a 100644 --- a/checkpoint/orbax/checkpoint/_src/path/step_test.py +++ b/checkpoint/orbax/checkpoint/_src/path/step_test.py @@ -245,6 +245,27 @@ def test_find_step( name_format.find_step(self.directory, step).commit_timestamp_nsecs ) + def test_custom_prefix_cvl(self): + """Verifies StandardNameFormat works with custom prefixes like 'v=step'.""" + name_format = step_lib.standard_name_format( + step_prefix='v=step', step_format_fixed_length=None + ) + self.assertEqual('v=step_1', name_format.build_name(1)) + + # Setup directories + (self.directory / 'v=step_1').mkdir() + (self.directory / 'v=step_2').mkdir() + + # Test find_step + metadata = name_format.find_step(self.directory, 1) + self.assertEqual(metadata.step, 1) + self.assertEqual(self.directory / 'v=step_1', metadata.path) + + # Test find_all + all_metadata = list(name_format.find_all(self.directory)) + self.assertLen(all_metadata, 2) + self.assertEqual({m.step for m in all_metadata}, {1, 2}) + @parameterized.parameters(True, False) def test_find_step_path_with_uncommitted_checkpoint(self, gcs: bool): """Tests for `step.find_step_path(include_uncommitted=True)`."""