fix(train): crash-safe checkpoints — completeness marker gating resume

[jessiewei7]

fix(train): crash-safe checkpoints — completeness marker gating resume

Summary

save_checkpoint() writes model.safetensors / optimizer_state_dict.pt /
etc. directly to their final paths, and _get_previous_epoch() (which resume
uses to find the latest checkpoint) trusts any directory whose name parses
as an integer — nothing distinguishes a fully written checkpoint from one
truncated by a crash mid-save.

Failure mode:

1. save_checkpoint() for epoch N is killed mid-save_pretrained() (OOM,
   preemption, node failure) — directory N/ now exists with a truncated
   model.safetensors.
2. The job restarts; _get_previous_epoch() sees N, returns it.
3. load_model_state_dict() hits the truncated file: best case safetensors
   throws a parse error and the run is stuck until someone deletes N/ by
   hand; worst case (optimizer_state_dict.pt via torch.load) a partial
   read surfaces as a confusing unpickling error.

training_state.json does not guard this: resume only reads it to pick
mid-epoch vs end-of-epoch, and a missing/corrupt one silently falls back to
"advance to next epoch" — still loading model/optimizer from the broken dir.

Fix

SpecForge had the identical bug class; the fix follows the same design as
sgl-project/SpecForge#659: a completeness marker written last, gating resume.

• BaseCheckpointer.mark_checkpoint_complete(epoch): fsync every file in the
  epoch dir, then write a checkpoint_complete marker (fsync it and the
  directory). The fsyncs matter: after a hard host crash, files can survive
  at full length with truncated content, so the marker must not claim
  completeness before the data is durable.
• Trainer calls it once all files of an epoch are in place (model,
  optimizer, scheduler, training_state.json, val metrics — the last writes
  happen in the trainer, which is why the marker call lives there).
• _get_previous_epoch() skips integer dirs without the marker, logging a
  warning with explicit instructions (touch .../checkpoint_complete) —
  the same warn-and-instruct pattern already used for interrupted/ dirs.
  This also covers checkpoints from before this change.

Test plan

tests/unit/train/test_checkpoint.py (extends the existing suite, CPU-only):

• a dir with files but no marker is skipped; resume falls back to the newest
  complete checkpoint (the crash scenario);
• mark_checkpoint_complete makes a checkpoint visible to a fresh
  checkpointer;
• maybe_save_checkpoint end-to-end writes the marker.

Full tests/unit/train/ suite passes (113 passed; 1 pre-existing failure is
a hub-connectivity issue in test_setup_model.py, unrelated).

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 7e87864b-b6d1-4747-81ea-ccea2127eae4

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

This PR adds a checkpoint completeness marker mechanism to BaseCheckpointer, gating resume/previous-epoch detection on the presence of a checkpoint_complete marker file written via fsync, wires the new mark_checkpoint_complete() call into Trainer's save and best-update paths, and adds corresponding unit tests.

Changes

Checkpoint Completeness Marker

Layer / File(s)	Summary
Completeness marker constant and helpers src/speculators/train/checkpointer.py	Adds os import, COMPLETE_MARKER_FILENAME constant, complete_marker_path(), and mark_checkpoint_complete(), which fsyncs epoch files, writes/fsyncs the marker, then fsyncs the directory to persist completion.
Resume detection gated by completeness marker src/speculators/train/checkpointer.py	_get_previous_epoch() now skips epoch directories missing the completeness marker, warning when saves were interrupted or predate marker behavior.
Trainer wiring of completion marking src/speculators/train/trainer.py	maybe_save_checkpoint and maybe_update_best now call checkpointer.mark_checkpoint_complete(epoch) on rank 0 or non-distributed runs after saving/symlinking and after persisting new best state, respectively.
Tests for marker gating and trainer integration tests/unit/train/test_checkpoint.py	Adds a _mark_complete helper, tests that unmarked epoch directories are ignored during resume, a test that mark_checkpoint_complete updates previous_epoch, and a test that maybe_save_checkpoint produces the marker file.

Possibly related PRs

• vllm-project/speculators#501: Both PRs modify checkpoint save/resume handling around epoch directories in checkpointer.py/trainer.py.
• vllm-project/speculators#603: Both PRs modify BaseCheckpointer._get_previous_epoch to change how epoch subdirectories are considered valid for resume.

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly summarizes the main change: making training checkpoints crash-safe by gating resume on a completeness marker.
Description check	✅ Passed	The description directly matches the implemented checkpoint completeness-marker workflow and the accompanying tests.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands.}

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@src/speculators/train/checkpointer.py`:
- Around line 123-132: Clear stale complete markers before rewriting an existing
checkpoint epoch so resume logic in the checkpoint scanner does not trust a
partially rewritten directory. Update the checkpoint save flow in the save
method that writes to self.path / str(epoch) to either remove or replace
COMPLETE_MARKER_FILENAME before writing model/optimizer/scheduler/metrics, or
write into a temporary epoch directory and atomically promote it after all
writes succeed. Ensure the marker is only recreated after the full save
completes, and keep the existing scan logic that checks COMPLETE_MARKER_FILENAME
aligned with this new incomplete-then-complete lifecycle.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: f2d3e154-7ca3-42ed-a96f-f1da3d77b2ce

📥 Commits

Reviewing files that changed from the base of the PR and between 2d97b17 and 9e018a7.

📒 Files selected for processing (3)

• src/speculators/train/checkpointer.py
• src/speculators/train/trainer.py
• tests/unit/train/test_checkpoint.py

[orestis-z]

LGTM — reviewed design, fsync ordering, distributed rank-gating, and test coverage. The completeness-marker pattern is sound, the stale-marker edge case is handled, and the regression test covers the core crash scenario.

[jessiewei7]

Rebased onto latest main. CI needs maintainer approval to run whenever you have a moment — thanks!