fix(trainer): improve checkpoint consistency on resume

[Kaiwei LIU (KAIWEILIUCC)]

Summary

Improve checkpoint consistency when resuming trainer runs.

• cross-check runtime state, history, step records, and skill snapshots
• find the last complete committed step
• remove artifacts created after that step
• atomically write history and runtime state
• use runtime state as the final commit marker
• restore scheduler state and rebuild the current epoch's step buffer
• track slow-update and meta-skill commit phases

No new CLI flags or configuration options are required.

Problem

The previous resume logic primarily relied on runtime_state.last_completed_step, or the last history entry, without cross-validating all related step records and skill snapshots.

If a run is interrupted while checkpoint files are being written, there is a possibility that runtime state, history, skill snapshots, step directories, and learning-rate history may reflect slightly different progress.

For a mid-epoch resume, the accumulated step buffer may also be unavailable, and scheduler progress may be reconstructed from the global step number instead of its previously persisted internal state.

Changes

On startup, the trainer now finds the longest contiguous sequence of valid committed steps.

If the last valid step is N, it:

• removes step directories and skill snapshots after N
• truncates history and learning-rate history after N
• removes incomplete slow-update and meta-skill artifacts
• restores the committed scheduler state
• rebuilds the current epoch's step buffer
• resumes by rerunning step N + 1

New checkpoints include skill hashes, commit IDs, scheduler state, and the current commit phase.

runtime_state.json is written last and acts as the commit marker. History, runtime state, skill snapshots, and step records are written atomically.

Compatibility

• Existing commands and configuration continue to work.
• Uninterrupted training behavior is unchanged.
• Legacy checkpoints without hashes or commit IDs use structural validation.
• Recovery is single-process; concurrent writers are not handled.
• Rerunning a model call may still produce different output because remote LLM generation is not deterministic.

Tests

Focused resume-related tests:

74 passed

The full suite was also run on macOS. Three existing tests fail because macOS resolves /tmp to /private/tmp. The same failures occur on a clean origin/main checkout:
- test_absolute_paths_needing_normalisation_are_accepted
- test_current_directory_segments_are_normalised_not_refused
- test_one_row_per_skill_in_order
No trainer or resume-related tests failed.

[Kaiwei LIU (KAIWEILIUCC)]

@microsoft-github-policy-service agree

[Copilot]

Pull request overview

Improves trainer checkpoint recovery and resume consistency.

Changes:

• Validates committed checkpoint prefixes and removes stale artifacts.
• Restores scheduler and step-buffer state.
• Adds atomic persistence and recovery tests.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Review summary
tests/test_trainer_resume.py	Adds focused recovery, phase, buffer, and atomic-write tests.
skillopt/engine/trainer.py	Implements checkpoint recovery and resume restoration. Two unresolved moderate findings concern unguarded marker conversion and trusting best_skill.md after recovery.

Suppressed comments (7)

skillopt/engine/trainer.py:1394

• An accepted slow update can change current_score, but this amendment only persists the best fields. The amended history/step record therefore reports the pre-slow current score, making epoch summaries and history-only recovery inconsistent with the runtime state. Persist current_score here as well.

                "best_origin": best_origin,
                "best_score": best_score,
                "best_step": best_step,
                "best_skill_hash": skill_hash(best_skill),

skillopt/engine/trainer.py:625

• These checks never compare runtime_state.scheduler_state with the committed step record. If the marker has the same step/hash metadata but a stale or malformed scheduler state, the resume path below loads it and either uses the wrong learning-rate phase or fails on the next scheduler.step(). Validate it against last_record (or fall back to the record state) before accepting the marker.

            runtime.get("current_skill_hash") in (None, previous_skill_hash)
            and runtime.get("history_hash") in (None, _json_digest(committed))
            and runtime.get("step_record_hash")
            in (None, _json_digest(last_record) if last_record is not None else "")

skillopt/engine/trainer.py:1341

• This clears the hash whenever a phase checkpoint calls _persist_runtime_state without step_record (the meta-skill and complete calls do exactly that). Recovery treats None as an opt-out, so after an epoch/final checkpoint it no longer validates the last step record against the runtime marker. Preserve the hash for history[-1] when it is the record for last_completed_step, or pass that record at these call sites.

                    "step_record_hash": (
                        _json_digest(step_record) if step_record is not None else None
                    ),

skillopt/engine/trainer.py:1304

• With skill-aware reflection enabled, this conditional persists raw skill_init before the inject_empty_appendix_field call at line 1314. The first step records its input hash after injection, while recovery computes the previous hash from raw skill_v0000 and rejects step 1 on the next startup, so a baseline resume becomes unrecoverable after that step. Persist the normalized initial skill before starting/resuming.

        if not (runtime_state or history) or not os.path.exists(skill_zero_path):
            _save_skill(out_root, 0, skill_init)

skillopt/engine/trainer.py:638

• The normalization above replaces any non-list history with [], so not isinstance(history, list) is always false here. For a non-list or otherwise invalid history.json, changed is also false when committed is empty, meaning the invalid file is never rewritten and subsequent startups continue without a usable history document. Preserve a validity flag before normalization and rewrite invalid history.

    if changed or (os.path.exists(history_path) and not isinstance(history, list)):

skillopt/engine/trainer.py:1230

• For legacy runtime markers, best_skill_hash is absent, so this validation is skipped and best_skill.md is trusted. The old writer updated that convenience file before publishing its runtime marker; a crash can therefore leave it at an uncommitted step, and the resumed gate will use the wrong best skill. When the hash is absent, rebuild/validate best_skill from the versioned best_step snapshot instead of accepting this file.

            expected_best_hash = runtime_state.get("best_skill_hash")
            if expected_best_hash and skill_hash(best_skill) != expected_best_hash:

skillopt/engine/trainer.py:636

• Pre-v2 runtime markers are valid dictionaries but have no phase. Defaulting them to step makes an epoch-boundary legacy checkpoint delete that epoch's slow_update and meta_skill directories, including completed done artifacts. The next resume then reruns potentially nondeterministic LLM updates and can produce a different skill; missing legacy phase should preserve structurally complete artifacts or otherwise distinguish unknown phase from an in-progress step.

        phase=str(runtime.get("phase", "step")) if runtime else "step",

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.