Detect GCSFuse mounts and default to commit-file atomicity

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

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

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

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}'

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