Detect GCSFuse mounts and default to commit-file atomicity - #3537
Open
mridul-sahu wants to merge 1 commit into
Open
Detect GCSFuse mounts and default to commit-file atomicity#3537mridul-sahu wants to merge 1 commit into
mridul-sahu wants to merge 1 commit into
Conversation
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.
mridul-sahu
force-pushed
the
feature/gcsfuse-atomicity
branch
from
August 13, 2026 19:20
8388974 to
b30cc56
Compare
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
When checkpoints are saved through a GCSFuse mount, the path looks local, so
AUTOatomicity mode selectsATOMIC_RENAME. Renames through GCSFuse are per-object copy-and-delete operations, so finalization cost scales with checkpoint size;COMMIT_FILEavoids this but is only a manual opt-in today.This PR detects GCSFuse mounts (from the system mount table) and makes
AUTOmode:CommitFileTemporaryPathfor GCSFuse paths, for writers and readers alike;ATOMIC_RENAMEis explicitly requested for a GCS or GCSFuse path.Explicitly configured modes are unchanged. The v1 API inherits the behavior through
context.atomicity.v0().Migration note: validation on GCSFuse paths now requires
commit_success.txt, matching direct GCS. Checkpoints written by Orbax versions predating the marker convention needallow_legacy_atomic_rename=Trueor a one-timecommit_success.txtstamp in each finalized step directory; without either they read as in-progress saves and may be removed by temporary-path cleanup.Type of change
Checklist
CHANGELOG.mdunder the[Unreleased]section (or the relevantcheckpoint//export/changelog).