Skip to content

Detect GCSFuse mounts and default to commit-file atomicity - #3537

Open
mridul-sahu wants to merge 1 commit into
google:mainfrom
mridul-sahu:feature/gcsfuse-atomicity
Open

Detect GCSFuse mounts and default to commit-file atomicity#3537
mridul-sahu wants to merge 1 commit into
google:mainfrom
mridul-sahu:feature/gcsfuse-atomicity

Conversation

@mridul-sahu

@mridul-sahu mridul-sahu commented Aug 13, 2026

Copy link
Copy Markdown
Collaborator

Description

When checkpoints are saved through a GCSFuse mount, the path looks local, so AUTO atomicity mode selects ATOMIC_RENAME. Renames through GCSFuse are per-object copy-and-delete operations, so finalization cost scales with checkpoint size; COMMIT_FILE avoids this but is only a manual opt-in today.

This PR detects GCSFuse mounts (from the system mount table) and makes AUTO mode:

  • select CommitFileTemporaryPath for GCSFuse paths, for writers and readers alike;
  • warn whenever ATOMIC_RENAME is 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 need allow_legacy_atomic_rename=True or a one-time commit_success.txt stamp in each finalized step directory; without either they read as in-progress saves and may be removed by temporary-path cleanup.

Type of change

  • Bug fix (non-breaking change that fixes an issue)
  • New feature (non-breaking change that adds functionality)
  • Breaking change (alters existing behavior or a public API)
  • Docs / tests / internal tooling (no user-facing behavior change)

Checklist

  • I have read the contribution guidelines.
  • Tests cover this change and pass locally.
  • Public APIs and non-trivial functions are documented.
  • If this change is user-facing, I have updated CHANGELOG.md under the [Unreleased] section (or the relevant checkpoint/ / export/ changelog).

@github-actions github-actions Bot added the pull ready Ready to be pulled from GitHub into Google label Aug 13, 2026
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
mridul-sahu force-pushed the feature/gcsfuse-atomicity branch from 8388974 to b30cc56 Compare August 13, 2026 19:20
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

pull ready Ready to be pulled from GitHub into Google

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant