diff --git a/checkpoint/CHANGELOG.md b/checkpoint/CHANGELOG.md index 6d83a501c..24e1cfd73 100644 --- a/checkpoint/CHANGELOG.md +++ b/checkpoint/CHANGELOG.md @@ -7,9 +7,12 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] -### Changed +## [0.12.5] - 2026-09-17 + +### Added -- Batch array persistence writes by location and device assignment in `CloudPathwaysArrayHandler.serialize` and symmetrically group reads by location and mesh in `deserialize`. +- Add `FileOptions.skip_sync_file_validations` to skip synchronized file + validation during async checkpointing. ## [0.12.4] - 2026-08-12 diff --git a/checkpoint/orbax/checkpoint/_src/checkpointers/async_checkpointer.py b/checkpoint/orbax/checkpoint/_src/checkpointers/async_checkpointer.py index 8d5848d84..91e4e3400 100644 --- a/checkpoint/orbax/checkpoint/_src/checkpointers/async_checkpointer.py +++ b/checkpoint/orbax/checkpoint/_src/checkpointers/async_checkpointer.py @@ -380,7 +380,7 @@ def __init__( *, async_options: options_lib.AsyncOptions = options_lib.AsyncOptions(), multiprocessing_options: options_lib.MultiprocessingOptions = options_lib.MultiprocessingOptions(), - file_options: options_lib.FileOptions = options_lib.FileOptions(), + file_options: Optional[options_lib.FileOptions] = None, atomicity_options: Optional[options_lib.AtomicityOptions] = None, checkpoint_metadata_store: Optional[checkpoint.MetadataStore] = None, temporary_path_class: Optional[ @@ -412,7 +412,7 @@ def __init__( else f'{multiprocessing_options.barrier_sync_key_prefix}' ) self._barrier_sync_key_prefix = barrier_sync_key_prefix - self._file_options = file_options + self._file_options = file_options or options_lib.FileOptions() self._atomicity_options = atomicity_options self._metadata_store = ( checkpoint_metadata_store @@ -507,6 +507,33 @@ def _callback() -> None: return _callback + async def _prepare_destination_async( + self, directory: epath.Path, *, force: bool + ) -> None: + """Removes existing destination if force=True, or checks for collisions.""" + skip = self._file_options.skip_sync_file_validations + + # 1. Force overwrite: only the primary host performs cleanup. + if force: + if not utils.is_primary_host(self._primary_host): + return + should_remove = skip or await async_path.exists(directory) + if should_remove: + logging.info( + '[process=%s] Specified `force`: removing existing directory.', + multihost.process_index(), + ) + await async_path.rmtree( + directory, + missing_ok=skip, + ) # Post-sync handled by create_tmp_directory. + return + + # 2. Collision validation: verify destination directory does not exist. + if not skip: + if await async_path.exists(directory): + raise ValueError(f'Destination {directory} already exists.') + async def _save( self, tmpdir: atomicity_types.TemporaryPath, @@ -515,18 +542,7 @@ async def _save( **kwargs, ): directory = tmpdir.get_final() - if await async_path.exists(directory): - if force: - if utils.is_primary_host(self._primary_host): - logging.info( - '[process=%s] Specified `force`: removing existing directory.', - multihost.process_index(), - ) - await async_path.rmtree( - directory - ) # Post-sync handled by create_tmp_directory. - else: - raise ValueError(f'Destination {directory} already exists.') + await self._prepare_destination_async(directory, force=force) commit_ops = [] if self._create_directories_asynchronously: diff --git a/checkpoint/orbax/checkpoint/_src/checkpointers/async_checkpointer_test.py b/checkpoint/orbax/checkpoint/_src/checkpointers/async_checkpointer_test.py index e24fed2ee..b57e4e14b 100644 --- a/checkpoint/orbax/checkpoint/_src/checkpointers/async_checkpointer_test.py +++ b/checkpoint/orbax/checkpoint/_src/checkpointers/async_checkpointer_test.py @@ -550,6 +550,86 @@ def test_save_metrics(self): ) checkpointer.close() + @parameterized.parameters((True,), (False,)) + def test_skip_sync_file_validations(self, skip_sync_file_validations): + file_options = options_lib.FileOptions( + skip_sync_file_validations=skip_sync_file_validations + ) + checkpointer = self.checkpointer( + PyTreeCheckpointHandler(), file_options=file_options + ) + with mock.patch.object( + async_path, 'exists', wraps=async_path.exists + ) as mock_exists: + checkpointer.save(self.directory, self.pytree) + self.wait_if_async(checkpointer) + directory_exists_calls = [ + call + for call in mock_exists.call_args_list + if call.args and call.args[0] == self.directory + ] + if skip_sync_file_validations: + self.assertEmpty(directory_exists_calls) + else: + if multihost.is_primary_host(checkpointer._primary_host): + self.assertNotEmpty(directory_exists_calls) + restored = checkpointer.restore( + self.directory, restore_args=self.pytree_restore_args + ) + test_utils.assert_tree_equal(self, self.pytree, restored) + checkpointer.close() + + @parameterized.parameters((True,), (False,)) + def test_overwrite_existing_primary_only_exists( + self, skip_sync_file_validations + ): + file_options = options_lib.FileOptions( + skip_sync_file_validations=skip_sync_file_validations + ) + checkpointer = self.checkpointer( + PyTreeCheckpointHandler(), file_options=file_options + ) + checkpointer.save(self.directory, self.pytree) + self.wait_if_async(checkpointer) + with ( + mock.patch.object( + async_path, 'exists', wraps=async_path.exists + ) as mock_exists, + mock.patch.object( + async_path, 'rmtree', wraps=async_path.rmtree + ) as mock_rmtree, + ): + checkpointer.save(self.directory, self.doubled_pytree, force=True) + self.wait_if_async(checkpointer) + directory_exists_calls = [ + call + for call in mock_exists.call_args_list + if call.args and call.args[0] == self.directory + ] + directory_rmtree_calls = [ + call + for call in mock_rmtree.call_args_list + if call.args and call.args[0] == self.directory + ] + if multihost.is_primary_host(checkpointer._primary_host): + if skip_sync_file_validations: + self.assertEmpty(directory_exists_calls) + else: + self.assertNotEmpty(directory_exists_calls) + self.assertLen(directory_rmtree_calls, 1) + self.assertEqual( + directory_rmtree_calls[0].kwargs.get('missing_ok'), + skip_sync_file_validations, + ) + else: + self.assertEmpty(directory_exists_calls) + self.assertEmpty(directory_rmtree_calls) + restored = checkpointer.restore( + self.directory, restore_args=self.pytree_restore_args + ) + test_utils.assert_tree_equal(self, self.doubled_pytree, restored) + checkpointer.close() + if __name__ == '__main__': multiprocess_test.main() diff --git a/checkpoint/orbax/checkpoint/_src/path/async_path.py b/checkpoint/orbax/checkpoint/_src/path/async_path.py index 090995190..1faba6803 100644 --- a/checkpoint/orbax/checkpoint/_src/path/async_path.py +++ b/checkpoint/orbax/checkpoint/_src/path/async_path.py @@ -91,15 +91,22 @@ async def async_stat(path: epath.Path): return await asyncio.to_thread(path.stat) -async def rmtree(path: epath.Path): +async def rmtree(path: epath.Path, *, missing_ok: bool = False) -> None: + """Removes a directory tree asynchronously. + + Args: + path: The directory path to remove. + missing_ok: Whether to ignore if the path does not exist. + """ def _rmtree(): # TODO(b/493110683): Cleanup with refactoring of HNS GCS logic into # StorageBackend. if gcs_utils.is_gcs_path(path): - gcs_utils.rmtree(path) + gcs_utils.rmtree(path, missing_ok=missing_ok) else: - path.rmtree() - return await asyncio.to_thread(_rmtree) + path.rmtree(missing_ok=missing_ok) + + await asyncio.to_thread(_rmtree) async def touch(path: epath.Path, *, exist_ok: bool = False): diff --git a/checkpoint/orbax/checkpoint/_src/path/async_path_test.py b/checkpoint/orbax/checkpoint/_src/path/async_path_test.py index 7be9f33f6..26e5bcf4f 100644 --- a/checkpoint/orbax/checkpoint/_src/path/async_path_test.py +++ b/checkpoint/orbax/checkpoint/_src/path/async_path_test.py @@ -170,6 +170,27 @@ async def _test(): asyncio.run(_test()) mock_open.assert_called_once_with(mode='r') + def test_rmtree(self): + target_dir = self.test_dir / 'sub_dir' + target_dir.mkdir() + (target_dir / 'file.txt').write_text('content') + + async def _test(): + await async_path.rmtree(target_dir) + self.assertFalse(target_dir.exists()) + + asyncio.run(_test()) + + def test_rmtree_missing(self): + missing_dir = self.test_dir / 'non_existent' + + async def _test(): + await async_path.rmtree(missing_dir, missing_ok=True) + with self.assertRaisesRegex(FileNotFoundError, r'non_existent'): + await async_path.rmtree(missing_dir, missing_ok=False) + + asyncio.run(_test()) + if __name__ == '__main__': absltest.main() diff --git a/checkpoint/orbax/checkpoint/_src/path/atomicity.py b/checkpoint/orbax/checkpoint/_src/path/atomicity.py index 10693a10a..4f007fbe7 100644 --- a/checkpoint/orbax/checkpoint/_src/path/atomicity.py +++ b/checkpoint/orbax/checkpoint/_src/path/atomicity.py @@ -752,11 +752,10 @@ def create_all_async( barrier_sync_key_prefix = multiprocessing_options.barrier_sync_key_prefix active_processes = multiprocessing_options.active_processes primary_host = multiprocessing_options.primary_host - # Sync for existence check to complete on all hosts before directory - # creation starts. + # Sync all hosts before directory creation starts. multihost.sync_global_processes( multihost.unique_barrier_key( - 'create_tmp_directory:post_existence_check', + 'create_tmp_directory:sync_before_directory_creation', prefix=barrier_sync_key_prefix, ), timeout=multihost.coordination_timeout(), diff --git a/checkpoint/orbax/checkpoint/_src/path/gcs_utils.py b/checkpoint/orbax/checkpoint/_src/path/gcs_utils.py index 26a8336cc..3438f7c66 100644 --- a/checkpoint/orbax/checkpoint/_src/path/gcs_utils.py +++ b/checkpoint/orbax/checkpoint/_src/path/gcs_utils.py @@ -100,19 +100,26 @@ def cleanup_hns_folders(path: epath.Path) -> None: ) -def rmtree(path: epath.Path) -> None: +def rmtree(path: epath.Path, *, missing_ok: bool = False) -> None: """Deletes a GCS path, performing HNS folder cleanup if necessary. Args: path: the global path to delete, must be a GCS path. + missing_ok: Whether to ignore if the path does not exist. Raises: ValueError: if path is not a GCS path. + FileNotFoundError: if path does not exist and missing_ok is False. """ if not is_gcs_path(path): raise ValueError(f'Path is not a GCS path: {path}') - path.rmtree() + try: + path.rmtree() + except FileNotFoundError: + if not missing_ok: + raise + return # For HNS, clean up the remaining empty directory structure. if is_hierarchical_namespace_enabled(path): diff --git a/checkpoint/orbax/checkpoint/_src/path/gcs_utils_test.py b/checkpoint/orbax/checkpoint/_src/path/gcs_utils_test.py index c57065a3b..03fc74cd9 100644 --- a/checkpoint/orbax/checkpoint/_src/path/gcs_utils_test.py +++ b/checkpoint/orbax/checkpoint/_src/path/gcs_utils_test.py @@ -113,6 +113,22 @@ def test_cleanup_hns_folders(self): self.assertIn('projects/_/buckets/my-bucket/folders/dir/a/', deleted_names) self.assertIn('projects/_/buckets/my-bucket/folders/dir/', deleted_names) + def test_rmtree_missing_ok(self): + gcs_path = mock.MagicMock(spec=epath.Path) + gcs_path.rmtree.side_effect = FileNotFoundError('Directory not found') + with ( + mock.patch.object(gcs_utils, 'is_gcs_path', return_value=True), + mock.patch.object( + gcs_utils, 'is_hierarchical_namespace_enabled', return_value=True + ), + mock.patch.object(gcs_utils, 'cleanup_hns_folders') as mock_cleanup, + ): + # Should not raise when missing_ok=True and should not clean up HNS. + gcs_utils.rmtree(gcs_path, missing_ok=True) + mock_cleanup.assert_not_called() + with self.assertRaisesRegex(FileNotFoundError, 'Directory not found'): + gcs_utils.rmtree(gcs_path, missing_ok=False) + if __name__ == '__main__': absltest.main() diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/context/options.py b/checkpoint/orbax/checkpoint/experimental/v1/_src/context/options.py index 297b9e4fe..b1e7fe8d1 100644 --- a/checkpoint/orbax/checkpoint/experimental/v1/_src/context/options.py +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/context/options.py @@ -223,15 +223,20 @@ class FileOptions(_ActiveContextGuard): The implementation of :py:class:`~.v1.path.Path` to use. Defaults to `etils.epath.Path`, but may be overridden to some other subclass of :py:class:`~.v1.path.Path`. + skip_sync_file_validations: + If True, bypasses synchronous filesystem existence and validation checks + prior to async saving. Default is False. """ path_permission_mode: int | None = None path_class: type[path_types.Path] = epath.Path + skip_sync_file_validations: bool = False def v0(self) -> v0_options_lib.FileOptions: """Converts this :py:class:`~.v1.options.FileOptions` to a v0 :py:class:`~orbax.checkpoint.options.FileOptions`.""" return v0_options_lib.FileOptions( path_permission_mode=self.path_permission_mode, + skip_sync_file_validations=self.skip_sync_file_validations, ) diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/context/options_test.py b/checkpoint/orbax/checkpoint/experimental/v1/_src/context/options_test.py index 8025f252f..c32164838 100644 --- a/checkpoint/orbax/checkpoint/experimental/v1/_src/context/options_test.py +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/context/options_test.py @@ -35,15 +35,18 @@ def test_v0_conversion_with_none_options(self): v0_opts = opts.v0() self.assertIsInstance(v0_opts, v0_options_lib.FileOptions) self.assertIsNone(v0_opts.path_permission_mode) + self.assertFalse(v0_opts.skip_sync_file_validations) def test_v0_conversion_with_all_options(self): opts = ocp_options.FileOptions( path_permission_mode=0o777, + skip_sync_file_validations=True, ) v0_opts = opts.v0() self.assertIsInstance(v0_opts, v0_options_lib.FileOptions) self.assertEqual(v0_opts.path_permission_mode, 0o777) + self.assertTrue(v0_opts.skip_sync_file_validations) class AtomicityOptionsTest(parameterized.TestCase): diff --git a/checkpoint/orbax/checkpoint/experimental/v1/_src/saving/path_utils.py b/checkpoint/orbax/checkpoint/experimental/v1/_src/saving/path_utils.py index cb2dd22ff..6f94577df 100644 --- a/checkpoint/orbax/checkpoint/experimental/v1/_src/saving/path_utils.py +++ b/checkpoint/orbax/checkpoint/experimental/v1/_src/saving/path_utils.py @@ -60,7 +60,7 @@ async def remove_existing_path( path: path_types.Path, *, context: context_lib.Context, -): +) -> None: """Removes the existing path. Args: @@ -72,7 +72,9 @@ async def remove_existing_path( '[process=%s] Specified `overwrite`: removing existing path.', multihost.process_index(), ) - await async_path.rmtree(path) + await async_path.rmtree( + path, missing_ok=context.file_options.skip_sync_file_validations + ) await multihost.sync_global_processes( multihost.unique_barrier_key( 'save_checkpointables_async:rmtree', @@ -99,6 +101,10 @@ async def maybe_overwrite_existing( Raises: ValueError: If the path exists and overwrite is False. """ + if context.file_options.skip_sync_file_validations: + if overwrite: + await remove_existing_path(path, context=context) + return # Sync before and after existence check, since otherwise the processes may # not agree and it is possible for one process to run ahead and create the # directory before another has checked for its existence. diff --git a/checkpoint/orbax/checkpoint/options.py b/checkpoint/orbax/checkpoint/options.py index 5b6fa78dd..d5a1902cb 100644 --- a/checkpoint/orbax/checkpoint/options.py +++ b/checkpoint/orbax/checkpoint/options.py @@ -90,6 +90,7 @@ class AtomicityOptions: allow_legacy_atomic_rename: bool = False +# pyformat: disable @dataclasses.dataclass(frozen=True) class FileOptions: """Options used to configure checkpoint directories and files. @@ -98,9 +99,16 @@ class FileOptions: path_permission_mode: Path permission mode for step directories, user metadata files. e.g. 0o750. Please check https://github.com/google/etils/blob/main/etils/epath/backend.py if your + path is supported. default=None. + skip_sync_file_validations: If True, bypasses synchronous filesystem + existence and validation checks prior to async saving in + AsyncCheckpointer. Default is False. This is experimental and should only + be used if you are certain of the effects. """ +# pyformat: enable path_permission_mode: int | None = None + skip_sync_file_validations: bool = False @dataclasses.dataclass diff --git a/checkpoint/orbax/checkpoint/version.py b/checkpoint/orbax/checkpoint/version.py index e94cafce8..fb24b9d09 100644 --- a/checkpoint/orbax/checkpoint/version.py +++ b/checkpoint/orbax/checkpoint/version.py @@ -17,7 +17,7 @@ # A new PyPI release will be pushed everytime `__version__` is increased. # Also modify version and date in CHANGELOG. # LINT.IfChange -__version__ = '0.12.4' +__version__ = '0.12.5' # LINT.ThenChange(//depot/orbax/checkpoint/CHANGELOG.md)