Internal Change.

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
 

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:

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()

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):

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()

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(),

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):

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()

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,
     )
 
 

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):

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.

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

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)