Why can't you rename a file on Google Cloud Storage?

Object stores like GCS have no real rename. A key like gs://bucket/step_42.tmp/a is an immutable object; "renaming" it to step_42/a means copy-then-delete, and a directory of thousands of shards is thousands of non-atomic copies. So a rename cannot be the commit point — a reader could catch the directory half-moved. Orbax detects the gs:// prefix and switches strategy: it writes bytes directly under the final path and commits with a single marker file.

How does Orbax make a checkpoint save atomic?

Two interchangeable strategies in atomicity.py. AtomicRenameTemporaryPath writes into a sibling directory suffixed ".orbax-checkpoint-tmp" and finalizes with one os.rename into the final name — atomic on POSIX. CommitFileTemporaryPath writes straight to the final path and finalizes by writing a tiny "commit_success.txt" as the very last step. In both, the commit is a single indivisible operation, so a reader sees either "not there / not committed" or "complete" — never a half.

What is the .orbax-checkpoint-tmp directory?

It is the staging directory for the atomic-rename strategy. Orbax builds the temporary path as final_path.name + ".orbax-checkpoint-tmp" in the same parent, writes every shard there, and at finalize renames that directory to the real step name. The suffix is load-bearing: validation refuses to treat a path ending in it as final, and refuses to treat a path without it as temporary, so a half-written directory can always be told apart from a finished one and cleaned up.

What is the commit_success.txt file in a checkpoint?

It is the commit marker for the object-store strategy. Because there is no atomic rename, CommitFileTemporaryPath writes the checkpoint bytes directly under the final path and then, as the last action of finalize, writes a small file named "commit_success.txt" containing "Checkpoint commit was successful to ". Validation of a final path requires that file to exist; its presence is the definition of "done," and its absence means the directory is a partial write to be ignored.

What happens if a checkpoint save crashes mid-write?

Nothing is ever mistaken for complete. On POSIX the bytes live under the ".orbax-checkpoint-tmp" directory and the final name simply never appears until the single rename runs, so a crash before finalize leaves only an orphan temp dir that later gets cleaned up. On an object store the bytes are under the final path but "commit_success.txt" was never written, so validation rejects the directory as uncommitted. Either way, restore sees no usable checkpoint at that step rather than a corrupt one.

What "Atomic" Means on a Filesystem vs. an Object Store

A checkpoint must appear all-or-nothing, even if the job dies with one byte left to write. We walk the real Orbax code that guarantees it two ways — an atomic rename on a POSIX filesystem, and a commit_success.txt marker on an object store that has no rename at all.

Code walk · AI / ML. The source ↗

A free, interactive, animated visual explainer of What "Atomic" Means on a Filesystem vs. an Object Store — built to be understood, not skimmed.

Questions

Why can't you rename a file on Google Cloud Storage?: Object stores like GCS have no real rename. A key like gs://bucket/step_42.tmp/a is an immutable object; "renaming" it to step_42/a means copy-then-delete, and a directory of thousands of shards is thousands of non-atomic copies. So a rename cannot be the commit point — a reader could catch the directory half-moved. Orbax detects the gs:// prefix and switches strategy: it writes bytes directly under the final path and commits with a single marker file.
How does Orbax make a checkpoint save atomic?: Two interchangeable strategies in atomicity.py. AtomicRenameTemporaryPath writes into a sibling directory suffixed ".orbax-checkpoint-tmp" and finalizes with one os.rename into the final name — atomic on POSIX. CommitFileTemporaryPath writes straight to the final path and finalizes by writing a tiny "commit_success.txt" as the very last step. In both, the commit is a single indivisible operation, so a reader sees either "not there / not committed" or "complete" — never a half.
What is the .orbax-checkpoint-tmp directory?: It is the staging directory for the atomic-rename strategy. Orbax builds the temporary path as final_path.name + ".orbax-checkpoint-tmp" in the same parent, writes every shard there, and at finalize renames that directory to the real step name. The suffix is load-bearing: validation refuses to treat a path ending in it as final, and refuses to treat a path without it as temporary, so a half-written directory can always be told apart from a finished one and cleaned up.
What is the commit_success.txt file in a checkpoint?: It is the commit marker for the object-store strategy. Because there is no atomic rename, CommitFileTemporaryPath writes the checkpoint bytes directly under the final path and then, as the last action of finalize, writes a small file named "commit_success.txt" containing "Checkpoint commit was successful to <path>". Validation of a final path requires that file to exist; its presence is the definition of "done," and its absence means the directory is a partial write to be ignored.
What happens if a checkpoint save crashes mid-write?: Nothing is ever mistaken for complete. On POSIX the bytes live under the ".orbax-checkpoint-tmp" directory and the final name simply never appears until the single rename runs, so a crash before finalize leaves only an orphan temp dir that later gets cleaned up. On an object store the bytes are under the final path but "commit_success.txt" was never written, so validation rejects the directory as uncommitted. Either way, restore sees no usable checkpoint at that step rather than a corrupt one.

What "Atomic" Means on a Filesystem vs. an Object Store

Questions

Related explainers