[macandro96]

What does this PR do?

Saves a rollout snapshot at the checkpointing.checkpoint_must_save_by deadline, so a turn-level run keeps its unfinished turns across a walltime stop, and lets a run save environment state only at that deadline.

Stacked on #5049 (batch age).

Why

When the deadline fired after a train step, the controller saved a trainer checkpoint and stopped at once. A trainer checkpoint holds no environment state, and stopping cancelled the snapshot pump before it could save one, so every unfinished turn was redone after the restore. There was also no way to save environment state only at the end: turn-level recovery required snapshot_attempt_interval_s.

Config

rollout_checkpointing:
  snapshot_attempt_interval_s: null     # existing: extra snapshots on a timer
  snapshot_on_trainer_checkpoint: true  # new

The recovery config decides what a snapshot preserves; these two fields decide when, beyond the deadline save, which is automatic with turn-level recovery.

Saves environment state	snapshot_on_trainer_checkpoint	snapshot_attempt_interval_s
After every trainer checkpoint and at the deadline (default)	true	null
Only at the deadline	false	null
After every trainer checkpoint, at the deadline, and every 300 s	true	300

How it works

• Snapshots on: when the interval is set or rollout_recovery.target_level is turn (rollout_snapshots_enabled). This replaces the eight checks of the interval in validation, data-plane support, bootstrap identity, restore and the pump.
• Deadline: after the trainer checkpoint the deadline triggers, the train pump publishes any deferred save and calls RolloutSnapshotter.save_at_deadline(), a forced snapshot, before it stops. A failure that leaves checkpoint state consistent only warns, since the trainer checkpoint is durable; an invariant failure still raises.
• Pump: with no interval it waits only for a trainer checkpoint's wakeup, and before the first step also for the deadline. A wait timed for that deadline that the first step overtook saves nothing.
• Deadline estimate: each completed environment snapshot records its duration in TimeoutChecker.save_reserve_seconds, and the deadline fires that much earlier, so the trainer checkpoint and the snapshot both fit.

Notes for review

• Validation: turn level no longer requires an interval. Setup rejects turn level with snapshot_on_trainer_checkpoint: false, no interval and no checkpoint_must_save_by, which would never save.
• Deferred finalize: still rejected with an interval or with snapshot_on_trainer_checkpoint at turn level. Deadline-only is allowed: the deadline publishes step_N before its snapshot, and the run then stops.
• Sibling level: unchanged. The trainer checkpoint already holds that rollout state, so no snapshot follows it.
• Backward compatible: existing turn-level configs set an interval, which keeps working; they now also get the deadline snapshot.

Testing

• Unit:
  • the pump with no interval;
  • the trigger after a trainer checkpoint, on and off;
  • the pre-step deadline;
  • ignoring a stale deadline wait;
  • each deadline-snapshot outcome (saved, skipped, I/O error, invariant failure);
  • the train pump snapshotting before it stops;
  • deferred-finalize rules;
  • the never-saves rejection;
  • the timer countdown and reserve.
• 563 passed (checkpointing, setup, timer and trace tests), Gym lane 7 passed. ruff and pyrefly clean.
• Not yet run: a GPU functional run of the deadline path.

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.