Skip to content

feat: resume interrupted NeMo Gym episodes from durable checkpoints - #3

Draft
lauradang wants to merge 2 commits into
jalbericiola:hsg-validated-gym-20260728from
lauradang:laurad/mid-episode-resume
Draft

feat: resume interrupted NeMo Gym episodes from durable checkpoints#3
lauradang wants to merge 2 commits into
jalbericiola:hsg-validated-gym-20260728from
lauradang:laurad/mid-episode-resume

Conversation

@lauradang

Copy link
Copy Markdown

Resume interrupted NeMo Gym episodes from a durable checkpoint instead of restarting them from scratch.

What this adds

  • nemo_gym/episode_checkpoint.pyEpisodeCheckpoint schema, FilesystemCheckpointStore (atomic writes, revision guard, advisory per-episode flock lease), and the SessionSnapshot / RestoreSessionRequest protocol types.
  • responses_api_agents/simple_agent/app.py — opt-in mid_episode_resume config. When enabled, /run requires _ng_episode_id (and _ng_policy_version under require_policy_version), and checkpoints at each durable boundary: after seed, after each model step, and after each tool call. A retried request loads the checkpoint and continues from its phase rather than re-running the episode. A completed episode returns its cached result. Mismatched policy_version archives the checkpoint and starts fresh; a reused episode ID with a different request fingerprint is a 409. Lease contention returns 409 + Retry-After.
  • resources_servers/example_session_state_mgmt/app.py — implements /session_checkpointing, /checkpoint_session, /restore_session, /discard_session, so the stateful-counter example can opt in.

Disabled by default everywhere. A Resources server that doesn't implement session snapshots is rejected up front by the capability probe.

Verification

End-to-end on the hsg2 cluster against the counter env, GRPO / group-submission + group-consumption, generation lag 5, 8 nodes, driven by the companion mrl_internal change that stamps the episode identity.

Killed a live job mid-episode with 3013 episodes in flight (all at policy_version: iteration:0, no training checkpoint saved), then resubmitted:

resume_count histogram : {0: 3480, 1: 1221}   (2598 resumed by a later pass)
restart_count histogram: {0: 4701}
policy_mismatch archives: 0

277 of those were genuine mid-turn resumes — step=1, phase=ready_for_tool, revision 3–4, 5–6 saved model outputs, and a restored counter snapshot of count=10 against a dataset row whose initial_count is 9, i.e. mutated in-episode state recovered from disk rather than reseeded.

Known limitation

A crash of the gym process currently cannot reach this logic. ng_run does not restart a dead child — it raises and the supervisor restarts the whole stack, which comes up on new random ports, and the training client keeps dialing the old port. That retry loop lives in ServerClient, below the caller's retry handling, so the connection is never re-established and the checkpoint is never consulted. Fixing that needs pinned ports or URL re-resolution and is out of scope here; what this PR covers is retry-within-a-live-rollout and recovery across a full process restart.

🤖 Generated with Claude Code

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant