fix(checkpoint): survive an interrupted save instead of hanging all ranks (#3261)

* fix(checkpoint): survive an interrupted save instead of hanging all ranks

A save cut short by a wall-clock limit, preemption, or OOM leaves a partial
epoch_<E>_step_<N> directory behind. Because checkpoint directory names are
derived from the step, the resumed run recomputed step N and collided with
that leftover. save_checkpoint raised FileExistsError on rank 0 only, so
rank 0 died while every other rank blocked in the next collective — an
indefinite NCCL hang that idled 128 GPUs until the cluster reaper killed the
job, repeating in every subsequent resume window for zero net progress.

Mark each checkpoint directory as in progress while it is being written and
clear the mark only at publication, so an interrupted save stays observable
on disk. Replace the rank-0 guard with a cleanup whose outcome is reduced
across ranks, so a failure aborts the whole job rather than stranding it.

Fixes #3254

Changelog:
- Add mark_/clear_/is_checkpoint_incomplete to the checkpoint component and
  write the marker for the whole save window.
- Replace the rank-0 FileExistsError with BaseRecipe._reserve_checkpoint_dir:
  rank 0 removes any leftover directory, and the result is all-reduced so
  every rank raises together instead of hanging on a rank-0 death.
- Add BaseRecipe._publish_checkpoint as the single point where a checkpoint
  becomes resumable (clear marker, then advance LATEST), covering both the
  sync path and deferred async publication.
- Skip marked directories in find_latest_checkpoint's step-scan fallback, so
  a run whose first save was interrupted cannot resume from a partial
  checkpoint before any LATEST pointer exists.
- Exclude marked directories from the max_recent_checkpoints retention
  window so a leftover neither evicts a good checkpoint nor survives as an
  orphan.
- Route the four speculative recipes (DFlash, EAGLE-1, EAGLE-3, DSpark),
  which each carried a copy of the same guard, through the shared methods.
- Document the behavior in the checkpointing guide, including that
  replacement also applies to a complete checkpoint when resuming from an
  older one.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

* fix(checkpoint): close the review gaps in interrupted-save handling

Addresses three review findings on the interrupted-save fix.

Pointer targets bypassed the marker filter. find_latest_checkpoint and
resolve_restore_from_to_checkpoint_dir returned a LATEST or named-pointer
target before the incomplete check ran, so re-saving a step that LATEST
already pointed at and then being interrupted resumed from the partial
directory. Both now skip an incomplete pointer target: LATEST falls back to
the most recent complete checkpoint, and any other pointer fails loudly
instead of restoring partial state.

Rank-0 bookkeeping could still strand peers. os.makedirs, the marker write,
and the marker clear ran on rank 0 after the reduction, so an I/O error there
killed rank 0 while its peers entered the next barrier. Directory creation
and the marker write now sit inside the same guarded block as the removal and
share one reduction; the marker clear is non-fatal and leaves LATEST where it
is, so the checkpoint stays flagged incomplete and resume falls back.

The lifecycle silently did nothing for msc:// roots. Path, os.path.exists,
shutil.rmtree, and os.makedirs treated a remote path as a local one, writing
the marker into a bogus local msc:/... directory while the real checkpoint
went unmarked. Remote roots are now detected and skipped explicitly, with a
warning and a documented caveat, matching the existing local-only scope of
checkpoint discovery and retention.

Changelog:
- Skip incomplete pointer targets in find_latest_checkpoint and
  resolve_restore_from_to_checkpoint_dir, falling back to the most recent
  complete checkpoint for LATEST.
- Cover directory creation and marker creation by the same cross-rank
  reduction as the leftover removal, extracted as BaseRecipe._any_rank_failed.
- Make the marker clear in _publish_checkpoint non-fatal and leave LATEST
  unchanged when it fails.
- Move is_cloud_path into the checkpoint utils module, make the marker
  helpers no-ops for msc:// paths, and skip the reserve step for remote
  checkpoint roots with a warning.
- Document the remote-storage limitation in the checkpointing guide.
- Add seven tests: interrupted pointer target with and without a fallback,
  explicit pointer restore, makedirs and marker-write failures aborting every
  rank, non-fatal marker clear, and the remote-path no-op.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

* test: stop leaking a process group across the unit test session

tests/unit_tests/recipes/test_vlm_kd_tp_cp_correctness.py initialized a
single-rank gloo process group in a module-scoped fixture and never destroyed
it, so torch.distributed.is_initialized() stayed true for the rest of the
pytest session.

That silently corrupts later tests. get_world_size_safe() returns the live
group's size instead of reading WORLD_SIZE from the environment, so
test_save_time_warns_for_large_inline_consolidation asserted world_size=256
and observed world_size=1. Tests that skip when a group already exists, such
as the DDP no_sync and DSpark max-steps reducer tests, stop running at all.

Whether it bites depends on collection order. Running the checkpoint suite
before the recipes suite hides it, which is why CI stayed green while
`pytest tests/unit_tests/recipes tests/unit_tests/checkpoint` failed.

Changelog:
- Convert trivial_pg to a yield fixture that destroys the group it created in
  a finally block, leaving a pre-existing group alone.
- Add a session-scoped autouse guard in tests/unit_tests/conftest.py that
  fails the session when a process group outlives it, so the next leak is
  reported directly instead of surfacing as an unrelated assertion.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

* fix(checkpoint): repair remote saves and pointer handling for interrupted saves

Addresses three review findings on e27fa9ed.

msc:// saves broke. _reserve_checkpoint_dir returned early for remote roots
without creating the directory, but the rest of the save path writes rank-0
files under it with plain open and torch.save. losses.json raised
FileNotFoundError on rank 0 before the first barrier, so the peers blocked --
the exact hang this branch set out to remove. Directory creation now runs for
remote roots too; only the leftover cleanup and the marker, which are local
filesystem operations, are skipped.

Incomplete named pointers still restored partial state. Skipping an incomplete
LOWEST_VAL target fell through to <checkpoint_dir>/LOWEST_VAL, which is the
pointer itself and resolves straight back to the marked checkpoint. A
non-LATEST pointer at an incomplete checkpoint now raises; LATEST keeps its
fallback to the most recent complete checkpoint.

A failed marker clear could still publish LOWEST_VAL. _publish_checkpoint
returned silently while its callers went on to _update_best_symlink and
_prune_old_checkpoints, leaving LOWEST_VAL aimed at a directory whose marker
was still present. Publication now reports success and gates every dependent
pointer and retention update, in the sync path and in deferred async
publication.

Changelog:
- Keep directory creation in _reserve_checkpoint_dir for msc:// roots and skip
  only the local-filesystem cleanup and marker.
- Add save_losses to the checkpoint component, mirroring save_config's storage
  handling, and route losses.json through it from all five recipes; failures
  warn instead of raising on rank 0.
- Raise from resolve_restore_from_to_checkpoint_dir when a named pointer other
  than LATEST targets an incomplete checkpoint.
- Return publication success from _publish_checkpoint and gate the best-pointer
  and retention updates on it in base_recipe and the four speculative recipes.
- Add tests over the complete remote save path, the loss-metadata helper, the
  named-pointer rejection, and both sync and async publication gating.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

* fix(checkpoint): reject remote checkpoint roots and unpin interrupted checkpoints

Addresses the review of 6b699414 by fixing the cause rather than the three
symptoms: the branch had been accumulating msc:// special cases around a save
path that cannot write to object storage at all.

msc:// training checkpoints have never worked. Only the DCP model and optimizer
shards and config.yaml are MSC-aware; RNG, dataloader, and recipe metadata go
through torch.save, and LATEST and LOWEST_VAL are published with os.symlink.
Verified against origin/main: save_on_dp_ranks raises "Parent directory
msc://.../rng does not exist" and os.makedirs leaves a local msc: shadow tree,
both before this branch existed. Rather than keep patching, checkpoint_dir now
rejects remote roots at config construction. Every rank builds the config
identically, so the failure is collective and lands at startup instead of
partway through a save. This deletes the msc:// branches this branch had added
to _reserve_checkpoint_dir and to the marker helpers.

Retention pinned unusable checkpoints. _prune_old_checkpoints excluded the whole
pointer-protected set from deletion, including checkpoints carrying the
in-progress marker, so a LOWEST_VAL pointer at a checkpoint that was overwritten
and then interrupted kept it forever and broke the documented retention window.
Incomplete checkpoints no longer receive pointer protection and their dangling
pointer is dropped with them.

save_losses could still deadlock. _ensure_msc_available raises ImportError,
which the handler did not catch, so rank 0 could exit before a barrier. The
handler now covers it too, in addition to the config-level rejection.

Changelog:
- Reject msc:// checkpoint.checkpoint_dir in CheckpointingConfig, replacing the
  two narrower msc:// rejections whose guidance no longer holds.
- Drop pointer protection for checkpoints left by an interrupted save so
  retention deletes them and removes the dangling pointer.
- Catch ImportError in save_losses.
- Remove the msc:// special casing from _reserve_checkpoint_dir, the marker
  helpers, and checkpoint_lifecycle_is_supported, now unreachable.
- Document that checkpoint_dir must be a local or shared filesystem path, and
  that interrupted checkpoints are never pointer-protected.
- Extend tests: config rejection across retention and consolidation settings,
  retention with pointers at complete and incomplete targets, save_losses
  without multistorageclient, and a sync/async by restore-source lifecycle
  matrix over an interrupted save.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

* fix(checkpoint): keep interrupted checkpoints out of best-metric and pointer state

Two follow-ups from the review of 4d40fd8e.

An incomplete LOWEST_VAL target poisoned best-metric tracking.
_update_best_symlink runs before pruning and seeds _best_val_loss from the
pointer target's losses.json, including a target left behind by an interrupted
save. That metric belongs to a checkpoint the run can never restore, so it held
the baseline below every later checkpoint: with 0.1 recorded on the marked
target, a 0.2 checkpoint did not replace it, pruning then removed the target and
its pointer, and a later 0.15 checkpoint still could not recreate LOWEST_VAL
because the cached value survived. The baseline now skips marked targets.

Custom pointers were left dangling. Dropping pointer protection for marked
checkpoints lets retention delete a directory that any top-level pointer may
target, but cleanup only handled LATEST and LOWEST_VAL, so a user's PINNED
pointer survived its deleted target. Cleanup now covers every pointer the
protection scan already discovers, matching what the retention documentation
claims.

Staleness is now tested by target non-existence rather than "is not a
directory". Two existing tests caught that the stricter check would delete
pointers aimed at a file inside a live checkpoint, which are legitimate.

Changelog:
- Skip incomplete LOWEST_VAL targets in _initialize_best_val_loss_from_pointer.
- Add find_checkpoint_pointers to the checkpoint utils, returning every
  top-level pointer and its target, and build find_pointer_protected_checkpoints
  on it.
- Replace _remove_stale_checkpoint_pointer with _remove_stale_checkpoint_pointers,
  which removes any pointer whose target no longer exists and leaves pointers
  with live targets alone.
- Document that dangling pointers are cleaned up and live ones are never touched.
- Add regression tests for the multi-save best-metric sequence and for custom
  pointer cleanup alongside a preserved pointer.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

* fix(checkpoint): narrow the interrupted-save fix to marking and ignoring

The review of b671b442 was right that the three remaining blockers were not
independent: each one only existed because this branch had grown past "mark a
checkpoint while it is being written and ignore marked checkpoints". Narrowing
the semantics removes all three by construction rather than by another targeted
patch, and shrinks the change.

Only unpublished checkpoints are replaced. A directory carrying the marker was
never published and cannot be restored, so reclaiming it loses nothing. A
directory without the marker holds a published checkpoint and is no longer
overwritten; the save aborts on every rank instead. An async checkpoint whose
files finished but was never published still carries the marker, so the case the
issue reported stays covered. Because a published checkpoint is never rewritten,
its recorded metric can no longer go stale behind LOWEST_VAL, which was the first
blocker, and no pointer can come to target a marked checkpoint, which is why the
pointer-protection and custom-pointer cleanup added earlier are both reverted --
including the dictionary that collapsed the symlink and .txt forms, the second
blocker.

Every rank-0 filesystem mutation in the lifecycle now runs through one helper,
_run_rank_0_checkpoint_step, which reduces the outcome so either every rank
continues or every rank raises. That covers the pointer unlink of the third
blocker, and also the pre-existing rank-0 torch.save of tracked state, which
could strand peers the same way. Publication no longer reports success through a
return value, because failures are collective instead of silently degrading.

Changelog:
- Replace only marked checkpoint directories; raise FileExistsError on every rank
  when the target holds a published checkpoint.
- Add BaseRecipe._run_rank_0_checkpoint_step as the single guard for rank-0
  checkpoint filesystem work, and route reservation, publication, best-pointer
  updates, retention, and the tracked-state torch.save through it.
- Revert find_checkpoint_pointers, the pointer-protection exclusion, and the
  all-pointer stale cleanup; restore the original two-pointer cleanup, now
  covered by the guard.
- Keep the retention window computed over complete checkpoints only, so a marked
  leftover neither consumes a slot nor survives as an orphan.
- Restructure the tests as a lifecycle state machine (published, interrupted,
  ignored, replaced, published) plus failure injection at each rank-0 filesystem
  operation, including the pointer unlink, and drop the nine tests that asserted
  the reverted behavior.
- Document that only marked directories are replaced and that a published
  checkpoint is never overwritten.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

* fix(checkpoint): make reservation failure-atomic and guard every callback failure

Two guard failures from the review of 906d5babc, both in code the previous
rework introduced.

Reservation was not failure-atomic. It created the final directory and then
wrote the marker, so a marker write that failed -- or a process that stopped
between the two calls -- left the directory unmarked. Since the absence of the
marker is exactly what makes a directory look published, find_latest_checkpoint
then resumed from an empty directory and the next save refused to replace it.
Reproduced both outcomes. Reservation now builds the directory complete with its
marker under a `.reserving-<uuid>` staging name and renames it into place, which
is atomic within a filesystem, so the final path is never observable without its
marker. A reservation that dies partway leaves only the staging directory, which
neither checkpoint listing matches because the name does not end in step_<N>, and
which the next reservation of that step clears.

The rank-0 guard caught only OSError. Its callbacks now reach state_dict(),
torch.save, pointer publication, and retention, which raise RuntimeError and
pickling errors too. A RuntimeError injected into torch.save escaped before the
reduction, so rank 0 died while its peers waited -- the original hang. The guard
now catches Exception, logs it on rank 0, reports it through the reduction, and
chains the original as the cause of the collective RuntimeError.

Auditing the same invariant across the recipes found the equivalent unguarded
rank-0 metadata write in all four speculative recipes, which is now routed
through the guard as well.

Changelog:
- Reserve checkpoint directories through a marked staging directory renamed into
  place, and clear stale staging directories for the step being reserved.
- Catch Exception rather than OSError in _run_rank_0_checkpoint_step, preserving
  the original exception as the cause.
- Route the rank-0 recipe-metadata write in the DFlash, EAGLE-1, EAGLE-3, and
  DSpark recipes through the guard.
- Add regression tests for a failed reservation leaving nothing resumable, for a
  stale staging directory being ignored and cleaned up, and for a non-OSError
  rank-0 failure aborting every rank with its cause preserved.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

* fix(checkpoint): make publication failure-atomic

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

* docs(checkpoint): clarify async publication lifecycle

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

* fix(checkpoint): preserve MSC DCP support

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

* refactor(checkpoint): extract publication lifecycle

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

* docs(checkpoint): apply review feedback

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

* test(recipe): update checkpointer finalization fixture

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>

---------

Signed-off-by: Yuhe Zhang <yuhez@nvidia.com>
Signed-off-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>
Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
Co-authored-by: Alexandros Koumparoulis <153118171+akoumpa@users.noreply.github.com>
Co-authored-by: Alexandros Koumparoulis <akoumparouli@nvidia.com>