How RL gets every Gym service to stop at the same point before a checkpoint

NVIDIA-NeMo/RL #4264, 1/3 of a stack split from #4117 (then #4265, #4266). All permalinks pin head ad3c9cb.

Before
RL could snapshot its own rollout state, but it had no way to ask the Gym services (policy model server, agent servers, resource servers) to stop and save theirs. Gym work still running at save time could not be captured.
This PR
Adds the RL side of Gym's checkpoint protocol on the NemoGym actor: find out who can take part, then prepare → commit → resume (or restore, or abort) in a fixed order. Off by default.
Status
New, opt-in, not yet driven. Only discovery runs in this PR (at setup). The Single Controller calls the phases in #4266.

Nine pages on this stack — this one: #4264: what the new protocol does.

Background, if you have not used Gym. NeMo Gym runs RL's rollouts as a few small HTTP servers, all started by the NemoGym actor. There are three kinds:
• policy model server: passes each model call from an agent through to RL's own inference engine (vLLM);
• agent server: runs the multi-turn loop for each rollout (call the model, call a tool, repeat);
• resource server: holds the tool and environment state for each rollout, and scores the result.
All three hold state for rollouts that are still running, so a checkpoint that can resume them needs every one of them to stop and save. Dotted words have a short explanation: hover or tap them.

At setup, when rollout_checkpointing.gym.capability_discovery_enabled is true (config.py:746, default false), RL asks every Gym server what it can do in a checkpoint and keeps a list of the ones that can take part (discover_checkpoint_capabilities, called from setup.py:1798). Each later phase calls those servers one by one, in an order fixed per phase by a small table (nemo_gym.py:410-431). Policy models always pause and resume. Everything else takes part only if it can export its state (_participates_in_checkpoint_phase). One deadline, an absolute clock time, bounds every call.

The four phases, on one small setup

PhaseWhat it doesWhat exists afterwards
prepareStop. The policy model refuses new calls and waits for running ones to finish; agents park each rollout at a turn end; resources freeze. Nothing is written.All of Gym paused at one point. Nothing on disk yet, so backing out is free.
commitEach Gym server writes its own part into a temporary folder RL picked. Gym stays paused.No usable checkpoint yet: only Gym's half, unpublished. The Single Controller (#4266) still adds RL's half (TQ, loader place, recovery ledger, manifest) and then renames the folder into place. That rename is the real commit; a crash before it leaves the previous snapshot in use.
resumeReopen: resources, then the model, then agents. Called after the snapshot is published.A complete published snapshot, and Gym running again.
restoreAfter a restart, each Gym server loads its own part from a published snapshot. Gym stays paused until RL has loaded its half and calls resume.Gym back at the saved turns, paused.
rollback (abort)Undo a prepare or an unpublished commit by calling the resume routes, so nothing half-done is kept.Gym running again; the previous snapshot is still the latest.

Gym's "commit" means "write my part", not "make it permanent". The step-by-step order, with RL's own save, is on the sequence page.

Who calls what. In #4266 the Single Controller calls the actor's methods (prepare_checkpoint, commit_checkpoint, …); the NemoGym actor then makes each HTTP call below, one server at a time. The full sequence, including RL's own save and the restore path, is on the sequence page.

Example setup: one policy model policy_model, one agent simple_agent that can export state, one stateless agent math_agent, one resource server tools.

The actor's HTTP calls, per phase (left → right, one at a time) policy model agent resource server long-poll wait prepare policy_modelAdmission closed: new model calls get HTTP 409 and retry later; calls already running finish.pause: now draining simple_agentThe agent stops each rollout at the end of a turn and holds it there.park at a safe turn toolsThe resource server stops accepting changes to tool and environment state.freeze sessions policy_model statusRL asks the model server for its status; the server answers once nothing is in flight, or the deadline passes.wait until 3 in flight → 0 Why: close the model first so agents can't start a new turn and park fast; freeze tools once agents stop calling them. commit simple_agentA file listing, for each parked rollout, the saved turn it can continue from.write continuation index policy_modelThe model server's record of the model calls behind each saved turn.write model ledger toolsTool and environment state for each parked rollout.write session state agent's index is passed into the model commit (continuation_indexes) one line of continuations.jsonl (one line per parked rollout): {"rollout_id": "p7_g2", "attempt_index": 0, "capture_key": "p7_g2", "last_committed_model_call_id": "call-2", "resource_state_revisions": {"tools": 3}} Why: agents first, because the model commit needs their continuation indexes (above) to know which calls to save. restore policy_modelload ledger, stay paused simple_agentload index, stay paused toolsload sessions Why (not written down in the PR): likely model first, because restored agent turns point at calls in the model ledger. resume toolsaccept calls again policy_modelopen admission simple_agentParked rollouts continue with their next turn.release parked turns Why: open what agents call (tools, then the model) before releasing the agents, so their first call is accepted. rollback prepare fails or misses the deadline at any step → resume every server RL already contacted (in resume order), even one whose reply was lost, then raise. The previous checkpoint stays valid. abort_checkpoint uses the same resume calls. Not called in any phase: math_agent (stateless) and any model marked instance_role: auxiliary (for example a judge).

Where it joins what RL already had

RL already talked to one Gym server over a control API: the token-capture ledger on the policy model. The new code reuses that same helper, _control (nemo_gym.py:815), which sends every call through Gym's ServerClient.request. From there the two differ in only three ways:

Token-capture control (existing)Checkpoint control (new)
Which serverspolicy_model onlyevery discovered participant, by name
Password (bearer)token-capture tokenNEMO_GYM_CHECKPOINT_CONTROL_TOKEN, or the token-capture one if unset (picked by path, :824-828)
Time limit per callfixed control_timeout_s (60 s)time left until the shared deadline

Two setup checks keep the first version small: exactly one Gym actor (config.py:1280, sole_nemo_gym_checkpoint_actor), and multi-worker servers that Gym cannot coordinate are refused at discovery.

What to watch in review. Off by default, so nothing changes for existing runs from the new code. Three things matter anyway. (1) The Gym submodule now points at a commit on a Gym draft branch, not Gym main. (2) That same bump breaks four existing finalize metrics on every token-capture run (next page). (3) Gym marks every model server as a policy model unless told otherwise, so a separate judge server would be paused during prepare, and prepare would never finish on that recipe.

🔴 review comment — Gym submodule pin review comment — Gym bump breaks metrics review comment — nemo_gym.py:935 (judge servers)