Support consistent application checkpoint recovery and safe sampling replay

[shuangwu]

Summary

Support consistent resumption of application-owned checkpoints through an instance-owned controller adapter. This is recovery, not exact replay: unfinished work may be regenerated and some previously consumed prompts may repeat; an advanced fetched cursor must not silently skip uncommitted work.

• Typed, versioned metadata separates training commands, optimizer updates, remaining completions, sampling epoch and immutable checkpoint identity.
• Restore the effective sampler or batch sampler after initial set_epoch, before constructing/consuming its iterator. Explicit missing/corrupt checkpoints fail; automatic discovery misses bootstrap; step zero is valid.
• SamplingReplayLedger advances a safe sampling boundary only across a contiguous settled prefix. Out-of-order updates preserve an earlier replay boundary and its remaining-work budget. Snapshots contain sampling state, not rollout payloads.
• Optional CheckpointManifest publishes local-file checkpoints only after the required artifacts are durable. Atomic no-replace publication, artifact hashes, expected shard sets and compatibility checks reject incomplete, mixed or incompatible saves.
• Adapter-enabled controllers create a fresh execution ID. Network workers report that pinned ID; stale rollout/completion/end reports are rejected before changing counters, queues or transport ownership. Legacy runs without an adapter retain their existing behavior.
• Preserve worker/controller resume agreement, now including shared checkpoint identity. General fail-fast handling remains independent in Fail fast on checkpoint resume metadata disagreement #756.

Ownership and scope

Cosmos owns restore ordering, initialization of controller progress and old-attempt rollout rejection. Applications retain trainer serialization, checkpoint discovery, sampler/shard semantics, and the connection between issued work, completed updates and saved state. The replay ledger and POSIX manifest are reusable opt-in helpers, not automatic conversion of native checkpoints or a new mandatory checkpoint format.

Applications publish only after all required trainer shards have saved the same completed-update boundary. Their trainer validates the same checkpoint and restores model, optimizer, scheduler, RNG and any additional algorithm state before reporting agreement. The adapter restores the corresponding safe sampling state. Configured hard training-step limits still apply to replayed work.

No persisted rollout queues, partial-generation restoration, mandatory rollout drain, deterministic async scheduling, or identical future numerical trajectory is required. Resume starts fresh worker/transport processes; this is not live recovery of an existing worker cohort.

Validation conclusions

• 138 tests passed across adapter/replay/manifest/fencing, existing resume indices, completion admission, ranked rollout end, terminal checkpoint/drain and network utility suites.
• 65 tests passed with Support consistent application checkpoint recovery and safe sampling replay #749 and the independent Fail fast on checkpoint resume metadata disagreement #756 composed in a separate checkout; the changes apply cleanly without adding a PR dependency.
• Fresh two-rank CPU/Gloo save and restart completed successfully.
• Fresh two-rank GPU/NCCL save and restart completed successfully. The fixture fetched three batches, trained the first and third, and saved while the second remained outstanding. Both restarted ranks restored model, optimizer momentum, scheduler and RNG; regenerated the second batch; repeated the third; and completed two new updates.
• Negative tests cover interrupted publication, missing/mixed artifacts, incompatible dataset/topology, different checkpoint identities with equal counters, partial completion groups, duplicate/old settlements and stale network reports.
• Ruff and diff whitespace checks passed.

Reproduction and limits

The portable save/restart fixture is included as tests/controller_resume_replay_canary.py, with CPU/GPU launch instructions. Negative/failure-injection cases are in tests/test_checkpoint_replay.py and tests/test_resume_execution_fence.py. The earlier exact-parity fixture remains available as tests/controller_resume_gpu_canary.py; its parity is not a general replay guarantee.

The distributed fixture directly exercises the public adapter and helpers with application-owned trainer serialization/shard cursors. It is not a full application-launcher or large-model sharded-checkpoint validation. Custom formats/stores must implement the documented artifact validation and publication obligations. See docs/controller_resume_adapter.md.

Independence

Based directly on upstream main. Does not include the watchdog, W&B, requested-stop, or metadata fail-fast PRs. No downstream runtime/wheel patch is added.