From b30cc56dfd5a0a417f5784b88355d9fda4688a81 Mon Sep 17 00:00:00 2001 From: mridul-sahu Date: Fri, 14 Aug 2026 00:23:50 +0530 Subject: [PATCH] Detect GCSFuse mounts and default to commit-file atomicity. In AUTO mode, paths on GCSFuse mounts previously resolved to ATOMIC_RENAME because scheme-based dispatch only recognizes gs://. Directory renames through GCSFuse are per-object copy-and-delete operations, so finalization cost scaled with checkpoint size. - Add gcs_utils.is_gcsfuse_path, which resolves the deepest mount containing the path from the system mount table (read once per process) and checks for a gcsfuse filesystem type. - AUTO dispatch in atomicity_defaults now selects CommitFileTemporaryPath for GCSFuse paths, for writers and readers alike. Explicitly configured modes are unaffected. - Warn on every dispatch when ATOMIC_RENAME is explicitly requested for a GCS or GCSFuse path. Validation on GCSFuse paths consequently requires commit_success.txt, matching direct GCS. Checkpoints written by Orbax versions predating the marker convention need allow_legacy_atomic_rename=True or a one-time commit_success.txt stamp; without either they read as in-progress saves. The v1 API inherits the behavior through context.atomicity.v0(). Also modernizes typing aliases and removes unused imports in the touched files to satisfy lint. --- checkpoint/CHANGELOG.md | 15 +++ .../_src/path/atomicity_defaults.py | 48 +++++----- .../checkpoint/_src/path/atomicity_test.py | 92 ++++++++++++++++++- .../orbax/checkpoint/_src/path/gcs_utils.py | 62 +++++++++++++ .../checkpoint/_src/path/gcs_utils_test.py | 51 ++++++++++ .../checkpoint/_src/path/temporary_paths.py | 11 +-- .../_src/path/temporary_paths_test.py | 90 +++++++++++++++++- checkpoint/orbax/checkpoint/options.py | 23 +++-- 8 files changed, 343 insertions(+), 49 deletions(-) diff --git a/checkpoint/CHANGELOG.md b/checkpoint/CHANGELOG.md index b0cd928f82..65291b37b4 100644 --- a/checkpoint/CHANGELOG.md +++ b/checkpoint/CHANGELOG.md @@ -7,6 +7,21 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [Unreleased] +### Added + +- Detect GCSFuse mounts in `AUTO` atomicity mode and use the commit-file + protocol for them instead of atomic rename. + +### Changed + +- Warn whenever `ATOMIC_RENAME` is explicitly requested for a GCS or GCSFuse + path. +- On GCSFuse paths in `AUTO` mode, validation now requires + `commit_success.txt` like direct GCS. Checkpoints written by Orbax versions + that predate the marker need `allow_legacy_atomic_rename=True` or a + one-time `commit_success.txt` stamp; otherwise they are treated as + in-progress saves and may be cleaned up. + ## [0.12.4] - 2026-08-12 ### Fixed diff --git a/checkpoint/orbax/checkpoint/_src/path/atomicity_defaults.py b/checkpoint/orbax/checkpoint/_src/path/atomicity_defaults.py index 4e241c7b45..0499a1a97a 100644 --- a/checkpoint/orbax/checkpoint/_src/path/atomicity_defaults.py +++ b/checkpoint/orbax/checkpoint/_src/path/atomicity_defaults.py @@ -25,8 +25,6 @@ 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 - _ATOMICITY_MODE_TO_PATH_CLASS: dict[ options_lib.AtomicityMode, type[atomicity_types.TemporaryPath] @@ -63,6 +61,26 @@ def _resolve_temporary_path_class_from_options( raise ValueError(f'Unsupported atomicity_mode: {atomicity_options.mode}.') +def _maybe_warn_atomic_rename_on_object_store(path: epath.Path) -> None: + """Warns when ATOMIC_RENAME was explicitly requested for a GCS-backed path.""" + if gcs_utils.is_gcs_path(path): + logging.warning( + 'ATOMIC_RENAME atomicity was requested for the GCS path %s. GCS has' + ' no atomic directory rename, so finalization copies and deletes' + ' every object in the checkpoint. Prefer COMMIT_FILE (or AUTO) for' + ' GCS paths.', + path, + ) + elif gcs_utils.is_gcsfuse_path(path): + logging.warning( + 'ATOMIC_RENAME atomicity was requested for %s, which is on a GCSFuse' + ' mount. Directory renames through GCSFuse are slow copy-and-delete' + ' operations and may fail entirely depending on mount options.' + ' Prefer COMMIT_FILE (or AUTO) for GCSFuse paths.', + path, + ) + + def get_item_default_temporary_path_class( path: epath.Path, *, @@ -79,17 +97,10 @@ def get_item_default_temporary_path_class( """ path_cls = _resolve_temporary_path_class_from_options(atomicity_options) if path_cls is not None: - if ( - gcs_utils.is_gcs_path(path) - and path_cls == atomicity.AtomicRenameTemporaryPath - ): - logging.warning( - 'AtomicRenameTemporaryPath can cause major performance issues for GCS' - ' paths. ' - ) - + if path_cls == atomicity.AtomicRenameTemporaryPath: + _maybe_warn_atomic_rename_on_object_store(path) return path_cls - if gcs_utils.is_gcs_path(path): + if gcs_utils.is_gcs_path(path) or gcs_utils.is_gcsfuse_path(path): return atomicity.CommitFileTemporaryPath else: return atomicity.AtomicRenameTemporaryPath @@ -111,17 +122,10 @@ def get_default_temporary_path_class( """ path_cls = _resolve_temporary_path_class_from_options(atomicity_options) if path_cls is not None: - if ( - gcs_utils.is_gcs_path(path) - and path_cls == atomicity.AtomicRenameTemporaryPath - ): - logging.warning( - 'AtomicRenameTemporaryPath can cause major performance issues for GCS' - ' paths. ' - ) - + if path_cls == atomicity.AtomicRenameTemporaryPath: + _maybe_warn_atomic_rename_on_object_store(path) return path_cls - if gcs_utils.is_gcs_path(path): + if gcs_utils.is_gcs_path(path) or gcs_utils.is_gcsfuse_path(path): return atomicity.CommitFileTemporaryPath else: return atomicity.AtomicRenameTemporaryPath diff --git a/checkpoint/orbax/checkpoint/_src/path/atomicity_test.py b/checkpoint/orbax/checkpoint/_src/path/atomicity_test.py index d1dd43235a..31fc70380d 100644 --- a/checkpoint/orbax/checkpoint/_src/path/atomicity_test.py +++ b/checkpoint/orbax/checkpoint/_src/path/atomicity_test.py @@ -14,8 +14,9 @@ import asyncio import concurrent.futures -import stat import unittest +from unittest import mock + from absl.testing import absltest from absl.testing import parameterized from etils import epath @@ -24,11 +25,9 @@ from orbax.checkpoint._src.multihost import multihost from orbax.checkpoint._src.path import atomicity from orbax.checkpoint._src.path import atomicity_defaults -from orbax.checkpoint._src.path import atomicity_types -from orbax.checkpoint._src.path import temporary_paths +from orbax.checkpoint._src.path import gcs_utils from orbax.checkpoint._src.path.snapshot import snapshot as snapshot_lib - AtomicRenameTemporaryPath = atomicity.AtomicRenameTemporaryPath CommitFileTemporaryPath = atomicity.CommitFileTemporaryPath TMP_DIR_SUFFIX = atomicity.TMP_DIR_SUFFIX @@ -237,6 +236,91 @@ def test_set_path_twice_raises(self): dp.set_path(epath.Path('/second')) +class GetDefaultTemporaryPathClassTest(parameterized.TestCase): + + @parameterized.parameters( + (atomicity_defaults.get_default_temporary_path_class,), + (atomicity_defaults.get_item_default_temporary_path_class,), + ) + def test_gcs_path_uses_commit_file(self, get_path_cls): + self.assertIs( + get_path_cls(epath.Path('gs://bucket/ckpt')), + CommitFileTemporaryPath, + ) + + @parameterized.parameters( + (atomicity_defaults.get_default_temporary_path_class,), + (atomicity_defaults.get_item_default_temporary_path_class,), + ) + def test_local_path_uses_atomic_rename(self, get_path_cls): + with mock.patch.object(gcs_utils, 'is_gcsfuse_path', return_value=False): + self.assertIs( + get_path_cls(epath.Path('/local/ckpt')), + AtomicRenameTemporaryPath, + ) + + @parameterized.product( + get_path_cls=( + atomicity_defaults.get_default_temporary_path_class, + atomicity_defaults.get_item_default_temporary_path_class, + ), + atomicity_options=(None, options_lib.AtomicityOptions()), + ) + def test_gcsfuse_path_uses_commit_file(self, get_path_cls, atomicity_options): + with mock.patch.object(gcs_utils, 'is_gcsfuse_path', return_value=True): + self.assertIs( + get_path_cls( + epath.Path('/mnt/gcs/ckpt'), atomicity_options=atomicity_options + ), + CommitFileTemporaryPath, + ) + + def test_explicit_mode_overrides_gcsfuse_detection(self): + atomicity_options = options_lib.AtomicityOptions( + mode=options_lib.AtomicityMode.ATOMIC_RENAME + ) + with mock.patch.object(gcs_utils, 'is_gcsfuse_path', return_value=True): + self.assertIs( + atomicity_defaults.get_default_temporary_path_class( + epath.Path('/mnt/gcs/ckpt'), atomicity_options=atomicity_options + ), + AtomicRenameTemporaryPath, + ) + + @parameterized.parameters( + ('gs://bucket/ckpt', False), + ('/mnt/gcs/ckpt', True), + ) + def test_explicit_atomic_rename_on_gcs_backed_path_warns( + self, path, gcsfuse_detected + ): + atomicity_options = options_lib.AtomicityOptions( + mode=options_lib.AtomicityMode.ATOMIC_RENAME + ) + with mock.patch.object( + gcs_utils, 'is_gcsfuse_path', return_value=gcsfuse_detected + ): + with mock.patch.object( + atomicity_defaults.logging, 'warning' + ) as mock_log: + atomicity_defaults.get_default_temporary_path_class( + epath.Path(path), atomicity_options=atomicity_options + ) + mock_log.assert_called_once() + + def test_explicit_atomic_rename_on_local_path_does_not_warn(self): + atomicity_options = options_lib.AtomicityOptions( + mode=options_lib.AtomicityMode.ATOMIC_RENAME + ) + with mock.patch.object(gcs_utils, 'is_gcsfuse_path', return_value=False): + with mock.patch.object( + atomicity_defaults.logging, 'warning' + ) as mock_log: + atomicity_defaults.get_default_temporary_path_class( + epath.Path('/local/ckpt'), atomicity_options=atomicity_options + ) + mock_log.assert_not_called() + if __name__ == '__main__': absltest.main() diff --git a/checkpoint/orbax/checkpoint/_src/path/gcs_utils.py b/checkpoint/orbax/checkpoint/_src/path/gcs_utils.py index 26a8336cc5..7b35a9bd5a 100644 --- a/checkpoint/orbax/checkpoint/_src/path/gcs_utils.py +++ b/checkpoint/orbax/checkpoint/_src/path/gcs_utils.py @@ -18,16 +18,78 @@ import os import pathlib from urllib import parse + from absl import logging from etils import epath _GCS_PATH_PREFIX = ('gs://',) +_MOUNTS_FILE = '/proc/mounts' +_GCSFUSE_FSTYPES = ('fuse.gcsfuse', 'gcsfuse') +# Octal escapes used by the kernel for whitespace in mount table fields. +_MOUNT_POINT_ESCAPES = ( + ('\\040', ' '), + ('\\011', '\t'), + ('\\012', '\n'), + ('\\134', '\\'), +) + def is_gcs_path(path: pathlib.PurePosixPath) -> bool: return path.as_posix().startswith(_GCS_PATH_PREFIX) +def _unescape_mount_point(mount_point: str) -> str: + for escape, char in _MOUNT_POINT_ESCAPES: + mount_point = mount_point.replace(escape, char) + return mount_point + + +@functools.lru_cache(maxsize=1) +def _mount_table() -> tuple[tuple[str, str], ...]: + """Returns (mount_point, fstype) pairs, or () if the table is unreadable.""" + try: + with open(_MOUNTS_FILE, 'rt') as f: + lines = f.read().splitlines() + except OSError: + return () + table = [] + for line in lines: + fields = line.split() + if len(fields) >= 3: + table.append((_unescape_mount_point(fields[1]), fields[2])) + return tuple(table) + + +def is_gcsfuse_path(path: epath.PathLike) -> bool: + """Returns whether `path` resides on a GCSFuse mount. + + The system mount table is read once per process and cached, so mounts + established after the first call are not detected. + + Args: + path: A local filesystem path. URI-style paths (e.g. `gs://...`) are never + considered GCSFuse paths. + + Returns: + True if the deepest mount containing `path` is a GCSFuse filesystem. + """ + path_str = os.fspath(path) + if '://' in path_str: + return False + resolved = os.path.realpath(path_str) + best_mount_point = '' + best_fstype = '' + for mount_point, fstype in _mount_table(): + if resolved == mount_point or resolved.startswith( + mount_point.rstrip('/') + '/' + ): + if len(mount_point) > len(best_mount_point): + best_mount_point = mount_point + best_fstype = fstype + return best_fstype in _GCSFUSE_FSTYPES + + def parse_gcs_path(path: epath.PathLike) -> tuple[str, str]: parsed = parse.urlparse(str(path)) assert parsed.scheme == 'gs', f'Unsupported scheme for GCS: {parsed.scheme}' diff --git a/checkpoint/orbax/checkpoint/_src/path/gcs_utils_test.py b/checkpoint/orbax/checkpoint/_src/path/gcs_utils_test.py index cb38a39e55..0b68b7999e 100644 --- a/checkpoint/orbax/checkpoint/_src/path/gcs_utils_test.py +++ b/checkpoint/orbax/checkpoint/_src/path/gcs_utils_test.py @@ -15,11 +15,62 @@ """Tests for gcs_utils functions.""" from unittest import mock + from absl.testing import absltest from etils import epath from orbax.checkpoint._src.path import gcs_utils +class IsGcsFusePathTest(absltest.TestCase): + + def setUp(self): + super().setUp() + gcs_utils._mount_table.cache_clear() + self.addCleanup(gcs_utils._mount_table.cache_clear) + + def _patch_mounts(self, content: str) -> mock._patch: + mounts_file = self.create_tempfile('mounts', content=content) + return mock.patch.object(gcs_utils, '_MOUNTS_FILE', mounts_file.full_path) + + def test_path_on_gcsfuse_mount(self): + mounts = ( + '/dev/sda1 / ext4 rw 0 0\n' + 'my-bucket /mnt/gcs fuse.gcsfuse rw,nosuid 0 0\n' + ) + with self._patch_mounts(mounts): + self.assertTrue(gcs_utils.is_gcsfuse_path('/mnt/gcs/ckpts/step_1')) + self.assertTrue(gcs_utils.is_gcsfuse_path(epath.Path('/mnt/gcs'))) + self.assertFalse(gcs_utils.is_gcsfuse_path('/mnt/gcs2/ckpts')) + self.assertFalse(gcs_utils.is_gcsfuse_path('/home/user/ckpts')) + + def test_nested_non_gcsfuse_mount_wins(self): + mounts = ( + '/dev/sda1 / ext4 rw 0 0\n' + 'my-bucket /mnt/gcs fuse.gcsfuse rw 0 0\n' + 'tmpfs /mnt/gcs/scratch tmpfs rw 0 0\n' + ) + with self._patch_mounts(mounts): + self.assertTrue(gcs_utils.is_gcsfuse_path('/mnt/gcs/ckpts')) + self.assertFalse(gcs_utils.is_gcsfuse_path('/mnt/gcs/scratch/ckpts')) + + def test_escaped_mount_point(self): + mounts = 'my-bucket /mnt/gcs\\040dir fuse.gcsfuse rw 0 0\n' + with self._patch_mounts(mounts): + self.assertTrue(gcs_utils.is_gcsfuse_path('/mnt/gcs dir/ckpts')) + + def test_gcs_uri_is_not_gcsfuse(self): + mounts = 'my-bucket /mnt/gcs fuse.gcsfuse rw 0 0\n' + with self._patch_mounts(mounts): + self.assertFalse(gcs_utils.is_gcsfuse_path('gs://my-bucket/ckpts')) + self.assertFalse(gcs_utils.is_gcsfuse_path(epath.Path('gs://b/k'))) + + def test_missing_mount_table(self): + with mock.patch.object( + gcs_utils, '_MOUNTS_FILE', '/nonexistent/mounts-file' + ): + self.assertFalse(gcs_utils.is_gcsfuse_path('/mnt/gcs/ckpts')) + + class GcsUtilsTest(absltest.TestCase): def test_rmtree_non_gcs_path(self): diff --git a/checkpoint/orbax/checkpoint/_src/path/temporary_paths.py b/checkpoint/orbax/checkpoint/_src/path/temporary_paths.py index ea1727e319..7c196fbff2 100644 --- a/checkpoint/orbax/checkpoint/_src/path/temporary_paths.py +++ b/checkpoint/orbax/checkpoint/_src/path/temporary_paths.py @@ -19,7 +19,7 @@ """ import asyncio -from typing import Iterable, Type +from collections.abc import Iterable from absl import logging from etils import epath @@ -30,7 +30,6 @@ from orbax.checkpoint._src.path import atomicity_defaults from orbax.checkpoint._src.path import atomicity_types - ValidationError = atomicity_types.ValidationError TMP_DIR_SUFFIX = atomicity_types.TMP_DIR_SUFFIX @@ -39,7 +38,7 @@ async def is_path_temporary( path: epath.PathLike, *, - temporary_path_cls: Type[atomicity_types.TemporaryPath] | None = None, + temporary_path_cls: type[atomicity_types.TemporaryPath] | None = None, atomicity_options: options_lib.AtomicityOptions | None = None, ) -> bool: """Determines if the given path represents a temporary checkpoint. @@ -94,7 +93,7 @@ async def is_path_temporary( async def is_path_finalized( path: epath.PathLike, *, - temporary_path_cls: Type[atomicity_types.TemporaryPath] | None = None, + temporary_path_cls: type[atomicity_types.TemporaryPath] | None = None, atomicity_options: options_lib.AtomicityOptions | None = None, ) -> bool: """Determines if the given path represents a finalized checkpoint. @@ -149,7 +148,7 @@ async def is_path_finalized( async def all_temporary_paths( root_directory: epath.PathLike, *, - temporary_path_cls: Type[atomicity_types.TemporaryPath] | None = None, + temporary_path_cls: type[atomicity_types.TemporaryPath] | None = None, atomicity_options: options_lib.AtomicityOptions | None = None, ) -> Iterable[atomicity_types.TemporaryPath]: """Returns a list of tmp checkpoint dir names in `root_directory`.""" @@ -181,7 +180,7 @@ async def cleanup_temporary_paths( directory: epath.PathLike, *, multiprocessing_options: options_lib.MultiprocessingOptions | None = None, - temporary_path_cls: Type[atomicity_types.TemporaryPath] | None = None, + temporary_path_cls: type[atomicity_types.TemporaryPath] | None = None, atomicity_options: options_lib.AtomicityOptions | None = None, ): """Cleanup steps in `directory` with tmp files, as these are not finalized. diff --git a/checkpoint/orbax/checkpoint/_src/path/temporary_paths_test.py b/checkpoint/orbax/checkpoint/_src/path/temporary_paths_test.py index 123b83d9ce..f54c6c2655 100644 --- a/checkpoint/orbax/checkpoint/_src/path/temporary_paths_test.py +++ b/checkpoint/orbax/checkpoint/_src/path/temporary_paths_test.py @@ -13,14 +13,17 @@ # limitations under the License. import asyncio -from typing import Type import unittest +from unittest import mock + from absl.testing import absltest from absl.testing import parameterized from etils import epath +from orbax.checkpoint import options as options_lib from orbax.checkpoint._src.path import async_path 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 temporary_paths @@ -37,7 +40,7 @@ def setUp(self): (atomicity.CommitFileTemporaryPath,), ) async def test_temporary_path( - self, tmp_path_cls: Type[atomicity_types.TemporaryPath] + self, tmp_path_cls: type[atomicity_types.TemporaryPath] ): tmp_path = tmp_path_cls.from_final(self.directory / 'ckpt') await tmp_path.create() @@ -57,7 +60,7 @@ async def test_temporary_path( (atomicity.CommitFileTemporaryPath,), ) async def test_finalized_path( - self, tmp_path_cls: Type[atomicity_types.TemporaryPath] + self, tmp_path_cls: type[atomicity_types.TemporaryPath] ): tmp_path = tmp_path_cls.from_final(self.directory / 'ckpt') await tmp_path.create() @@ -138,7 +141,7 @@ async def test_incorrect_path_class(self): (atomicity.CommitFileTemporaryPath,), ) async def test_all_temporary_paths( - self, tmp_path_cls: Type[atomicity_types.TemporaryPath] + self, tmp_path_cls: type[atomicity_types.TemporaryPath] ): num_paths = 3 tmp_paths = [ @@ -172,7 +175,7 @@ async def test_all_temporary_paths( (atomicity.CommitFileTemporaryPath,), ) async def test_cleanup_temporary_paths( - self, tmp_path_cls: Type[atomicity_types.TemporaryPath] + self, tmp_path_cls: type[atomicity_types.TemporaryPath] ): num_paths = 3 tmp_paths = [ @@ -198,5 +201,82 @@ async def test_cleanup_temporary_paths( self.assertFalse(await async_path.exists(path.get_final())) +class LegacyAtomicRenameFallbackTest( + parameterized.TestCase, unittest.IsolatedAsyncioTestCase +): + + def setUp(self): + super().setUp() + self.directory = epath.Path(self.create_tempdir('ckpt').full_path) + + async def _create_legacy_finalized(self, name: str) -> epath.Path: + """Creates a finalized rename-style checkpoint without a commit marker.""" + tmp_path = atomicity.AtomicRenameTemporaryPath.from_final( + self.directory / name + ) + await tmp_path.create() + await tmp_path.finalize() + final_path = tmp_path.get_final() + (final_path / atomicity_types.COMMIT_SUCCESS_FILE).unlink() + return final_path + + @parameterized.parameters( + (None,), + (options_lib.AtomicityOptions(),), + ( + options_lib.AtomicityOptions( + mode=options_lib.AtomicityMode.COMMIT_FILE + ), + ), + ) + async def test_gcsfuse_does_not_read_legacy_checkpoint_without_flag( + self, atomicity_options + ): + final_path = await self._create_legacy_finalized('step_1') + with mock.patch.object(gcs_utils, 'is_gcsfuse_path', return_value=True): + self.assertFalse( + await temporary_paths.is_path_finalized( + final_path, atomicity_options=atomicity_options + ) + ) + # Without the marker the directory reads as an in-progress commit-file + # save, so cleanup passes will remove it. + self.assertTrue( + await temporary_paths.is_path_temporary( + final_path, atomicity_options=atomicity_options + ) + ) + + async def test_gcsfuse_treats_rename_tmp_dir_as_temporary(self): + tmp_path = atomicity.AtomicRenameTemporaryPath.from_final( + self.directory / 'step_1' + ) + await tmp_path.create() + with mock.patch.object(gcs_utils, 'is_gcsfuse_path', return_value=True): + self.assertTrue(await temporary_paths.is_path_temporary(tmp_path.get())) + self.assertFalse(await temporary_paths.is_path_finalized(tmp_path.get())) + + @parameterized.parameters((True,), (False,)) + async def test_explicit_flag_enables_fallback(self, gcsfuse_detected): + final_path = await self._create_legacy_finalized('step_1') + atomicity_options = options_lib.AtomicityOptions( + mode=options_lib.AtomicityMode.COMMIT_FILE, + allow_legacy_atomic_rename=True, + ) + with mock.patch.object( + gcs_utils, 'is_gcsfuse_path', return_value=gcsfuse_detected + ): + self.assertTrue( + await temporary_paths.is_path_finalized( + final_path, atomicity_options=atomicity_options + ) + ) + self.assertFalse( + await temporary_paths.is_path_temporary( + final_path, atomicity_options=atomicity_options + ) + ) + + if __name__ == '__main__': absltest.main() diff --git a/checkpoint/orbax/checkpoint/options.py b/checkpoint/orbax/checkpoint/options.py index 5b6fa78dd7..b8fd5487f3 100644 --- a/checkpoint/orbax/checkpoint/options.py +++ b/checkpoint/orbax/checkpoint/options.py @@ -14,14 +14,13 @@ """Configuration options for APIs like CheckpointManager and Checkpointer.""" +from collections.abc import Callable import dataclasses import enum -from typing import Callable, Optional, Set from orbax.checkpoint._src.multihost import multihost - @dataclasses.dataclass class AsyncOptions: """Options used to configure async behavior. @@ -32,8 +31,8 @@ class AsyncOptions: timeout_secs: int = ( 1200 # 20 minutes. Same as default in `AsyncCheckpointer`. ) - barrier_sync_fn: Optional[multihost.BarrierSyncFn] = None - post_finalization_callback: Optional[Callable[[], None]] = None + barrier_sync_fn: multihost.BarrierSyncFn | None = None + post_finalization_callback: Callable[[], None] | None = None create_directories_asynchronously: bool = True @@ -55,9 +54,9 @@ class MultiprocessingOptions: other barrier syncs if another CheckpointManager is being used concurrently. """ - primary_host: Optional[int] = 0 - active_processes: Optional[Set[int]] = None - barrier_sync_key_prefix: Optional[str] = None + primary_host: int | None = 0 + active_processes: set[int] | None = None + barrier_sync_key_prefix: str | None = None @@ -76,11 +75,11 @@ class AtomicityOptions: Attributes: mode: Specifies the atomicity mode for saving checkpoints. - AUTO: - Automatically selects based on storage backend (GCSDirect -> COMMIT_FILE, - POSIX/GCSFuse -> ATOMIC_RENAME). - COMMIT_FILE: Writes in-place and - creates commit_success.txt upon completion (ideal for GCSFuse). - - ATOMIC_RENAME: Writes to temporary directory and renames to final - directory. + Automatically selects based on storage backend (GCSDirect and detected + GCSFuse mounts -> COMMIT_FILE, other POSIX -> ATOMIC_RENAME). - + COMMIT_FILE: Writes in-place and creates commit_success.txt upon + completion (ideal for GCSFuse). - ATOMIC_RENAME: Writes to temporary + directory and renames to final directory. allow_legacy_atomic_rename: Optional. If True, permits reading legacy checkpoints saved via AtomicRename that do not contain a commit_success.txt file. Default is False.