feat(checkpoint): add Gym participant coordination - #4264
macandro96 wants to merge 4 commits into
Conversation
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
terrykong
left a comment
terrykong
left a comment
There was a problem hiding this comment.
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):
- The Gym bump breaks two
nemo_gym-lane tests and thefinalize/terminal_selection_*metric names on every token-capture run. - The new Gym pin is not on Gym
main— re-pin to a GymmainSHA once the Gym stack lands. - 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
featureslist, so new Gym features don't invalidate old checkpoints. - A typed
GymControlRequestErrorthat carrieserror_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
| @@ -1 +1 @@ | |||
| Subproject commit 267305e2a7bf88153e8ebbca1029ecc1e0f67b25 | |||
| Subproject commit f4fcf8cf98272740fd7585aa49cee00fa2c099f0 | |||
There was a problem hiding this comment.
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:
KeyError: 'finalize/terminal_selection_heuristic_count'at #L227.- The test at #L152 expects
"empty_manifest"but getsinvalid_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.
| @@ -1 +1 @@ | |||
| Subproject commit 267305e2a7bf88153e8ebbca1029ecc1e0f67b25 | |||
| Subproject commit f4fcf8cf98272740fd7585aa49cee00fa2c099f0 | |||
There was a problem hiding this comment.
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.
| _get_node_ip_local, | ||
| ) | ||
| from nemo_rl.environments.gym_checkpoint import ( | ||
| GYM_AGENT_CONTINUATION_INDEX_FEATURE, |
There was a problem hiding this comment.
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_FEATUREmust come afterGYM_AGENT_CHECKPOINT_PREFIX,GymCheckpointResumeResultafterGymCheckpointRestoreResult, andGymCoordinatorModelStatusResponseafterGymControlCapabilities. - format:
nemo_gym.py#L1710-L1712andtest_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.
| from pydantic import BaseModel, ConfigDict, Field, FiniteFloat, model_validator | ||
|
|
||
| GYM_CHECKPOINT_SCHEMA_VERSION = 1 | ||
| GYM_CHECKPOINT_CONTROL_PREFIX = "/ng-control/v1" |
There was a problem hiding this comment.
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._checkpointis private in Gym. Importing it is still safer than copying: if Gym renames something, the actor fails with anImportErroron 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.
| 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." | ||
| ) |
There was a problem hiding this comment.
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:
- Discovery and a topology per actor, with one fingerprint over all of them.
- One cut across all actors: prepare every actor before any actor commits, and abort all of them if one fails.
- A save folder per actor inside the snapshot.
- 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).
There was a problem hiding this comment.
will add one and link it here
| # AllTaskProcessedDataset wraps the raw rows; a plain sequence is also fine. | ||
| rows = getattr(dataset, "dataset", dataset) | ||
| if rows is None: | ||
| return set() |
There was a problem hiding this comment.
not sure i caught it, but was this change intentional? do we need it for this PR?
There was a problem hiding this comment.
this wasn't intentional - removed
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
Signed-off-by: Anish Mahishi <amahishi@nvidia.com>
This branch has not been deployed
2 participants
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:
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:
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 executionsParticipant ordering is deliberate:
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
NemoGym.nemo_gym.pyin Pyrefly coverage and document persisted/wire contracts.Scope
Test plan
Targeted coverage is in:
tests/unit/environments/test_gym_checkpoint.pytests/unit/environments/test_nemo_gym_checkpoint.pytests/unit/environments/test_nemo_gym_utils.pytests/unit/single_controller/test_setup.pySuggested command:
Before review