feat(checkpoint): add Gym participant coordination

[macandro96]

Summary

Adds the NeMo-RL side of the Gym checkpoint participant protocol: capability discovery, topology validation, typed control-plane messages, and dependency-safe prepare/commit/restore/resume/abort fan-out.

This is 1/3 in a stacked decomposition of #4117:

Stack	PR	Scope
1/3	#4264 (this PR)	Gym participant protocol
2/3	#4265	Durable rollout state and validation
3/3	#4266	Single Controller scheduling and recovery

This layer does not schedule rollout snapshots from Single Controller. It establishes and tests the participant transaction that the later PRs consume.

Why

Trainer checkpoints cannot safely capture active Gym work unless every relevant service agrees on one checkpoint cut. NeMo-RL needs to know:

• which models, agents, and resource servers support checkpointing;
• which services must pause or export state in each phase;
• when policy admission is closed and accepted model calls have drained;
• the safe dependency order for commit, restore, and resume;
• whether every participant handled the same checkpoint ID.

Failing these checks at setup or at the control boundary is safer than publishing a partial snapshot.

Participant flow

sequenceDiagram
    participant SC as NeMo-RL
    participant Gym as Gym actor
    participant Policy as Policy model
    participant Agent as Agent server
    participant Resource as Resource server

    SC->>Gym: discover_checkpoint_capabilities
    Gym->>Policy: GET capabilities
    Gym->>Agent: GET capabilities
    Gym->>Resource: GET capabilities
    Gym-->>SC: validated topology and fingerprint

    SC->>Gym: prepare(checkpoint_id, deadline)
    Gym->>Policy: close admission
    Policy-->>Gym: paused or draining
    Gym->>Agent: park or freeze executions
    Gym->>Resource: freeze resource state
    Gym->>Policy: poll until inflight_total is zero
    Gym-->>SC: all participants ready

    SC->>Gym: commit(checkpoint_id, temp_dir)
    Gym->>Agent: export continuation index
    Gym->>Policy: export model ledger and references
    Gym->>Resource: export resource state
    Gym-->>SC: participant manifests

    SC->>Gym: resume(checkpoint_id)
    Gym->>Resource: resume
    Gym->>Policy: reopen admission
    Gym->>Agent: release parked executions

Loading

Participant ordering is deliberate:

• Prepare: policy models → agents → resources.
• Commit: agents → policy models → resources.
• Restore: policy models → agents → resources.
• Resume: resources → policy models → agents.

The policy pause closes admission immediately but may initially report draining. Commit is not permitted until accepted model calls finish and every configured worker acknowledges the pause. A shared deadline bounds the transaction; failed prepares resume the participants that may have observed the request.

Main changes

• Add strict wire models for Gym checkpoint capabilities and phase responses.
• Discover, normalize, fingerprint, and validate the participant topology during setup.
• Add phase-specific participation rules and dependency-safe ordering.
• Add prepare, commit, restore, resume, and abort control methods to NemoGym.
• Validate checkpoint IDs, deadlines, participant manifests, and worker acknowledgements.
• Fail setup for unsupported multi-Gym-actor checkpoint topologies.
• Use Gym's authoritative resolved server configuration for control clients.
• Include nemo_gym.py in Pyrefly coverage and document persisted/wire contracts.
• Update the Gym submodule to the matching participant-protocol implementation.

Scope

• Opt-in and disabled by default.
• Exactly one NeMo Gym actor is supported by this stack.
• Turn-level policy generations drain during prepare; token-prefix cuts are a later layer.
• No completed-result ACK outbox or rollout redispatch is introduced in this PR.
• No Single Controller snapshot pump is enabled in this PR.

Test plan

Targeted coverage is in:

• tests/unit/environments/test_gym_checkpoint.py
• tests/unit/environments/test_nemo_gym_checkpoint.py
• tests/unit/environments/test_nemo_gym_utils.py
• tests/unit/single_controller/test_setup.py

Suggested command:

uv run pytest -q \
  tests/unit/environments/test_gym_checkpoint.py \
  tests/unit/environments/test_nemo_gym_checkpoint.py \
  tests/unit/environments/test_nemo_gym_utils.py \
  tests/unit/single_controller/test_setup.py

Before review

• Contributor conventions followed.
• Unit coverage added for participant discovery and phase coordination.
• Full target environment test run recorded in CI.

[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.

[terrykong]

Thanks for splitting #4117 into a stack — it made this much easier to review.

First, please rebase: the PR conflicts with main.

Must fix before merge (details in the inline comments on the Gym submodule line and on nemo_gym.py#L59):

1. The Gym bump breaks two nemo_gym-lane tests and the finalize/terminal_selection_* metric names on every token-capture run.
2. The new Gym pin is not on Gym main — re-pin to a Gym main SHA once the Gym stack lands.
3. Lint fails (isort + ruff format).

Worth keeping:

• A separate wire-format module, with strict models for what is saved and lenient ones for live responses (gym_checkpoint.py#L60-L69).
• Phase order written as a data table instead of spread through control flow.
• The topology fingerprint ignores the additive features list, so new Gym features don't invalidate old checkpoints.
• A typed GymControlRequestError that carries error_code.
• Failing loud on more than one Gym actor, both at config time and at setup.
• Off by default, behind a nested config block.

For Gym, not an ask on this PR: the coordinator's model-admission resume leaves out IDLE (coordinator.py#L636-L647), while the single-worker route allows it (model_admission.py#L180-L184). After a failed coordinator pause the fence is IDLE, so RL's rollback resume would get a 409 and _active_gym_checkpoint_id would never clear. No server reports coordinator mode at this pin, so nothing hits it today; worth fixing in Gym #3349 before #3563 turns coordinator mode on.

Stack notes:

• Field docs and the exemplar YAML land in #4266 — fine. Please also document NEMO_GYM_CHECKPOINT_CONTROL_TOKEN (nemo_gym.py#L391) and its fallback to the token-capture bearer token in #4266's docs section; today it only appears in #4266's test scripts.
• On its own, this PR can't finish an agent prepare (the acks arrive in #4265). Fine as long as the stack lands together.

Visual explainer: https://terrykong.github.io/gh-pages-poc/terryk/pr-4264-gym-checkpoint-participants.html
Visual explainer: https://terrykong.github.io/gh-pages-poc/terryk/pr-4264-gym-bump-metrics.html

Generated by Claude Code

[terrykong]

1 action item — please fix in this PR.

TL;DR — The Gym bump breaks two existing nemo_gym-lane tests and renames the per-method terminal_selection metrics into garbage on every token-capture run, not just checkpoint runs. Caused by the bump, not by the new checkpoint code.

Gym changed RolloutReceipt.terminal_selection from a plain Literal (old pin) to Literal[...] | None = None (new pin). rollout_reassembler.py builds metric names from get_args(...) on that annotation, which now returns (Literal[...], NoneType). So the run publishes finalize/terminal_selection_typing.Literal[...]_count and finalize/terminal_selection_<class 'NoneType'>_count, and the four real per-method metrics disappear.

How it shows up — test_rollout_reassembler.py passes 13/13 on main and on this PR with the old Gym pin, and fails 2 with the new pin:

1. KeyError: 'finalize/terminal_selection_heuristic_count' at #L227.
2. The test at #L152 expects "empty_manifest" but gets invalid_receipt: ... an unpoisoned receipt must name terminal_model_call_id and terminal_selection, from Gym's new receipt validator (records.py#L421-L441).

Neither #4265 nor #4266 fixes this, and no CI has run on this PR yet, so nothing has caught it.

AI-1

In rollout_reassembler.py, drop the None from the annotation before reading the method names:

annotation = RolloutReceipt.model_fields["terminal_selection"].annotation
terminal_selection_methods = tuple(
    method
    for arg in get_args(annotation)
    if arg is not type(None)
    for method in (get_args(arg) or (arg,))
)

Then update the receipt used by the test at #L152 so it passes Gym's new validator and still reaches the empty_manifest check. Not a suggestion block: the file is outside this PR's diff.

[macandro96]

Fixed in fdd1d64

[terrykong]

1 action item — please fix before merge.

TL;DR — The new Gym pin f4fcf8c is not on Gym main; it is the head of an open draft stack. Once that stack is rebased or squash-merged, the commit can become unfetchable and git submodule update --init breaks for every fresh clone and container of RL main. Introduced by this PR's bump.

f4fcf8c is the head of Gym draft PR #3349, which sits on 8 more open drafts (#2939–#2946). Compared with Gym main it has diverged: 69 commits ahead, 18 behind. Every earlier RL Gym pin was a Gym main commit, and .gitmodules tracks branch = main with shallow = true.

CI will not catch this. _submodule_check.yml only checks that the new SHA descends from main's current pin, and it does.

AI-1

Land the Gym stack on Gym main first, then re-pin 3rdparty/Gym-workspace/Gym to that Gym main SHA before this PR merges. Not a suggestion block: it needs the Gym stack merged first, then a new SHA.

[terrykong]

1 action item — please fix in this PR.

TL;DR — Any judge or reward model server that doesn't set instance_role is treated as a policy model, so RL closes its admission during prepare. On a recipe with a separate judge, every prepare hits its deadline and rolls back, so no checkpoint can ever be taken. Introduced by this PR.

How it shows up:

• Gym defaults every model server to instance_role="policy" (base_responses_api_model.py#L211).
• RL makes every policy-role model join prepare and resume (_participates_in_checkpoint_phase), and this line only checks that at least one exists.
• stage3_rlhf.yaml has a separate judge, genrm_model, with no instance_role. RL pauses the judge; policy rollouts waiting on the judge never drain; prepare times out and rolls back, every time.

Gym's own docstring says judges must not pause (model_admission.py#L17-L22).

"Exactly one policy model" is the wrong check: the same recipe has policy_model_reasoning_off: _copy: policy_model (#L314-L321), a real second policy server.

AI-1

In discover_checkpoint_capabilities, raise for any responses_api_models participant with instance_role == "policy" that is not served by RL's own policy endpoint (the policy_base_url RL sets at nemo_gym.py#L653). The error should tell the user to set instance_role: auxiliary on that server. Shape of the check (how to read a server's base URL from the resolved Gym config is up to you):

for item in participants:
    caps = item.capabilities
    if (
        caps.component == "responses_api_models"
        and caps.instance_role == "policy"
        and not _served_by_rl_policy(item.participant.server_name)
    ):
        raise RuntimeError(
            f"Gym model server {item.participant.server_name!r} reports "
            "instance_role='policy' but is not served by the RL policy. "
            "Set instance_role: auxiliary on it (e.g. a judge or reward model), "
            "or RL will pause it during every checkpoint prepare."
        )

Not a suggestion block: it adds a new loop and a helper, not a change to this line.

[macandro96]

good point! fixed in fdd1d64

[terrykong]

1 action item — please fix in this PR.

TL;DR — Lint fails on this PR (main passes): the new gym_checkpoint import block is out of isort order, and ruff format --check fails in two spots. Introduced by this PR.

With the pinned ruff v0.9.9:

• isort: GYM_AGENT_CONTINUATION_INDEX_FEATURE must come after GYM_AGENT_CHECKPOINT_PREFIX, GymCheckpointResumeResult after GymCheckpointRestoreResult, and GymCoordinatorModelStatusResponse after GymControlCapabilities.
• format: nemo_gym.py#L1710-L1712 and test_nemo_gym_checkpoint.py#L93-L95.

AI-1

Run pre-commit run --all-files and commit the result. Not a suggestion block: the fixes span several spots in two files.

[macandro96]

fixed in fdd1d64

[terrykong]

1 action item.

TL;DR — Nothing checks RL's copies of Gym's models against Gym, and nothing checks RL's models against the replies Gym actually sends, so a Gym bump can break checkpointing with every unit test green. Two response models have already drifted: the first model server with a generation-cut backend will fail every commit and restore. Introduced by this PR. Low severity.

Keeping these copies is fine — the module docstring explains why. Everything that can be imported instead is covered in the comment on line 36. What has to stay copied, with the Gym source:

RL copy	Gym source
gym_capture_key	capture_key_for
_IDENTITY_PATTERN	ROLLOUT_ID_PATTERN
GymCompletionReceipt	AgentAcknowledgeRequest
capability values in GymCheckpointParticipantContract (checkpoint_mode, instance_role, multi_process.mode, admission states)	ControlCapabilities
GymModelCommitResponse / GymModelRestoreResponse	CaptureLedgerCommitResult / CaptureLedgerRestoreResult

The drift that exists today: the two RL response models are extra="forbid", but Gym's result models also declare optional generation_cut_receipt / generation_cut_proof. Gym only sends them when set, and nothing sets them at this pin, so it is hidden for now.

AI-1

Add a new test file, test_gym_checkpoint_contract.py in tests/unit/environments/, with pytestmark = pytest.mark.nemo_gym and plain top-level Gym imports (no importorskip — that lane errors, not skips, when nemo_gym is missing: conftest.py#L205-L215). It needs two kinds of checks:

• The copies in the table above: for each row, assert the RL copy matches the Gym source. The commit/restore check fails today, so add the two optional fields to GymModelCommitResponse and GymModelRestoreResponse in the same change.
• The replies Gym builds as plain dicts (prepare, status, resume, and the agent and resources restore replies — Gym has no model to compare these against): start Gym's real control routes in-process with FastAPI's TestClient (install_model_admission, install_model_checkpoint, install_agent_checkpoint, install_resources_checkpoint), drive one prepare → commit → resume and one restore → resume, and parse every reply with the RL model. Gym's own test_checkpoint_e2e.py builds the same apps.

Both run on CPU, so any Gym bump that changes something RL depends on fails here instead of in the GPU functional tests or at the first checkpoint of a real run. Not a suggestion block: it adds a new test file.

[macandro96]

fixed in fdd1d64

[terrykong]

1 action item (nit).

TL;DR — The client HTTP timeout and the server long-poll timeout are both remaining, so the client usually gives up first. A slow drain then reports "control plane unreachable or stalled" instead of the real cause, "remained 'draining' with N in-flight". Introduced by this PR.

Gym waits min(timeout_s, deadline.remaining()) (model_admission.py#L148), the same budget the client's timeout_s=remaining on #L1222 allows, so the client's wait_for usually fires first and the user sees the timeout message at #L842-L846.

AI-1

Give the server a little less time than the client, and keep the reason next to the code:

Suggested change

	"timeout_s": remaining,
	# End the server's wait 1 s before ours, so a slow drain comes back as a
	# "still draining" reply instead of our own "unreachable" timeout.
	"timeout_s": max(0.0, remaining - 1.0),

While there, include status.inflight (the stuck rollout IDs) in the TimeoutError at #L1249-L1259 — today it is dropped even when the server does reply.

[macandro96]

fixed in fdd1d64

[terrykong]

1 action item.

TL;DR — checkpoint_id and source_checkpoint_id are two different IDs, and the docstring does not say so: the first names this restore operation, the second names the earlier save being loaded. A caller cannot tell which to pass where, or that a failed restore leaves Gym paused until they call abort or resume. Introduced by this PR.

How the only caller uses them (#4266): the save was taken under an ID like rollout-step-40-snapshot-3-<nonce> (single_controller.py:4635-4638); the restore runs under a fresh ID restore-<uuid> (setup.py:1972) and passes the save's ID as source_checkpoint_id (setup.py:1975-1982). The source check is at nemo_gym.py:1446-1458; on failure the caller must use abort_checkpoint or resume_checkpoint.

AI-1

Say it in the docstring, with the example values:

Suggested change

	"""Restore every stateful participant but leave admission paused."""
	"""Load each participant's saved Gym state from a snapshot and leave Gym paused.
	Two different IDs are involved:
	- ``checkpoint_id`` names this restore operation, for example
	``"restore-3f2a..."``. Each Gym server refuses control calls with any
	other ID until ``resume_checkpoint`` or ``abort_checkpoint`` is called
	with this one. It must be new, because Gym refuses an ID it has
	already finished.
	- ``source_checkpoint_id`` names the earlier save being loaded, for
	example ``"rollout-step-40-snapshot-3-9c1e..."``: the ``checkpoint_id``
	that was passed to ``commit_checkpoint`` when this snapshot was
	written. Each server reports which save it loaded; if any reports a
	different one, this raises ``RuntimeError``. ``None`` skips that check.
	Only participants whose ``checkpoint_mode`` is ``"export_restore"`` are
	called; auxiliary models never are. If any call or the source check
	fails, the participants already restored stay paused, so the caller
	must call ``abort_checkpoint`` or ``resume_checkpoint`` with
	``checkpoint_id``.
	Args:
	checkpoint_id: ID of this restore operation (see above).
	deadline_ts: Absolute Unix time that bounds every control call.
	checkpoint_dir: Snapshot folder that holds the saved Gym files.
	source_checkpoint_id: ID of the save being loaded (see above).
	Returns:
	``GymCheckpointRestoreResult`` as JSON.
	"""

[macandro96]

fixed in fdd1d64

[terrykong]

1 action item.

TL;DR — About 80 of this file's 681 lines (12%) are copies of Gym code that can simply be imported: every place that uses them runs inside the NemoGym actor, which has Gym installed. Importing them there removes the copies and any chance of them drifting from Gym. Introduced by this PR.

What can be imported instead of copied (Gym links are at the pinned commit):

RL copy (lines)	Gym source
route prefixes and the capabilities path, #L36-L43 (8)	CONTROL_URL_PREFIX, MODEL_ADMISSION_URL_PREFIX, MODEL_CHECKPOINT_URL_PREFIX, AGENT_CHECKPOINT_URL_PREFIX, RESOURCES_CHECKPOINT_URL_PREFIX
feature strings, #L44-L45 (2)	artifacts.py#L30-L32
request models GymCheckpointControlRequest, GymCheckpointDirectoryRequest, GymAgentCheckpointDirectoryRequest, GymModelCheckpointCommitRequest, GymModelCheckpointRestoreRequest (22)	CheckpointControlRequest, AgentCommitRequest / AgentRestoreRequest, ResourcesCommitRequest / ResourcesRestoreRequest, ModelCheckpointCommitRequest / ModelCheckpointRestoreRequest — same fields
GymControlCapabilities (38)	ControlCapabilities
_CHECKPOINT_CONTROL_ENV, nemo_gym.py#L391	CHECKPOINT_CONTROL_TOKEN_ENV
folder and file names in _participant_checkpoint_path	MODEL_LEDGER_SUBDIR, RESOURCES_STATE_SUBDIR, AGENT_STATE_SUBDIR / AGENT_MANIFEST_NAME; the instance-<sha256> folder comes from the private _agent_checkpoint_directory

Gym is installed in that actor: actor_environments.py#L96.

AI-1

Delete these copies and import the Gym names inside the NemoGym methods that use them, the same deferred way _control_client already imports ServerClient. For example, inside prepare_checkpoint:

from nemo_gym._checkpoint import CheckpointControlRequest, MODEL_ADMISSION_URL_PREFIX

request = CheckpointControlRequest(
    checkpoint_id=checkpoint_id, deadline_ts=deadline_ts
).model_dump(mode="json")

For the capabilities reply, drop the three fields the route adds from ControlFence.snapshot() (phase, active_checkpoint_id, deadline_ts) before calling ControlCapabilities.model_validate, because Gym's model rejects unknown fields. GymDiscoveredParticipant can then hold Gym's model; only the actor builds it. Not a suggestion block: the change spans both files.

Context — no action.

• nemo_gym._checkpoint is private in Gym. Importing it is still safer than copying: if Gym renames something, the actor fails with an ImportError on the first checkpoint call instead of silently drifting. A small public Gym module would be the long-term fix.
• What is left in this file is used on the driver or the Single Controller, where Gym is not installed, so it cannot be imported. The comment on line 35 covers testing those copies.

[macandro96]

fixed in fdd1d64

[terrykong]

1 action item.

TL;DR — Any sharded or replicated Gym stack (Ananth's shard stack, #3369–#3374, already on main) is refused by this feature, and the error only says "currently", so a user cannot tell whether support is planned. Please file a tracking issue for the multi-actor case and put its link in this error message, so the message documents itself. Deliberate scope limit of this PR (PR description: "Exactly one NeMo Gym actor is supported by this stack").

This raise and the ShardSetupError in sole_nemo_gym_checkpoint_actor are the two places that refuse more than one actor. The error message is the better home for the link than a code comment: it is what the user actually reads when they hit the limit.

AI-1

File the tracking issue, then add its URL to both messages. For this one:

            raise NotImplementedError(
                "Gym participant checkpointing currently supports exactly one "
                f"NeMo-Gym actor, but env.nemo_gym.shards configures "
                f"{gym_actor_count}. Configure one shard with replicas=1, or "
                "disable rollout_checkpointing.gym.capability_discovery_enabled. "
                "Support for sharded or replicated Gym stacks is tracked in "
                "https://github.com/NVIDIA-NeMo/RL/issues/<N>."
            )

Not a suggestion block: the issue number does not exist yet, and the same edit goes in a second file.

What the issue could list: the parts multi-actor support needs

Each Gym actor runs its own Gym servers, so this is more than a loop over actors:

1. Discovery and a topology per actor, with one fingerprint over all of them.
2. One cut across all actors: prepare every actor before any actor commits, and abort all of them if one fails.
3. A save folder per actor inside the snapshot.
4. On restore, send each resumed rollout back to the actor that holds its saved turns (the replica-routing work in fix(nemo-gym): preserve replica identity during routing #3374 is the starting point).

[macandro96]

will add one and link it here

[terrykong]

not sure i caught it, but was this change intentional? do we need it for this PR?

[macandro96]

this wasn't intentional - removed