feat: resume interrupted NeMo Gym episodes from durable checkpoints - #3
Draft
lauradang wants to merge 2 commits into
Draft
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Resume interrupted NeMo Gym episodes from a durable checkpoint instead of restarting them from scratch.
What this adds
nemo_gym/episode_checkpoint.py—EpisodeCheckpointschema,FilesystemCheckpointStore(atomic writes, revision guard, advisory per-episodeflocklease), and theSessionSnapshot/RestoreSessionRequestprotocol types.responses_api_agents/simple_agent/app.py— opt-inmid_episode_resumeconfig. When enabled,/runrequires_ng_episode_id(and_ng_policy_versionunderrequire_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 itsphaserather than re-running the episode. A completed episode returns its cached result. Mismatchedpolicy_versionarchives 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_internalchange 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: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 ofcount=10against a dataset row whoseinitial_countis 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_rundoes 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 inServerClient, 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