refactor(nemo-gym): centralize actor construction by ananthsub · Pull Request #3367 · NVIDIA-NeMo/RL

ananthsub · 2026-07-27T21:40:34Z

Summary

Centralizes construction of the NemoGym Ray actor in nemo_rl/environments/nemo_gym.py.

grpo.py, distillation.py, and single_controller_utils/setup.py previously assembled the actor configuration independently. Those paths could diverge in how they handled NeMo RL integration keys, uv directories, router replay, routed-expert dtype, and use_fastokens.

The shared implementation now separates those responsibilities:

build_nemo_gym_config() removes NeMo RL integration settings from the configuration forwarded to Gym and applies model-dependent rollout settings.
spinup_nemo_gym_actor() creates the actor runtime environment, preserves existing placement behavior, starts Gym, and installs the tokenizer.
SingleController accepts deferred model-server URLs while generation starts concurrently.
GRPO and SingleController now set VIRTUAL_ENV and UV_PROJECT_ENVIRONMENT to the virtual-environment root through make_actor_runtime_env(), matching the directory values already used by distillation.
Distillation now promotes tokenizer_config and pad_dynamic_image_shapes into the top-level NemoGymConfig fields consumed by the actor. This enables the VLM distillation path, although no current in-tree distillation recipe sets is_vlm: true.

Stack

This is the first PR in the NeMo-Gym actor-sharding stack. Actor lifetime fixes follow in #3368.

copy-pr-bot · 2026-07-27T21:40:37Z

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

copy-pr-bot · 2026-08-13T21:54:28Z

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

ananthsub · 2026-08-13T21:55:15Z

/ok to test a0df64f

ananthsub · 2026-08-13T22:14:27Z

/ok to test 76447f3

ananthsub · 2026-09-01T05:17:32Z

/ok to test 9a140b1

ananthsub · 2026-09-01T06:31:44Z

/ok to test e190399

ananthsub · 2026-09-01T17:42:09Z

/ok to test 45e305f

ananthsub · 2026-09-02T16:15:40Z

/ok to test a77c04e

terrykong

Three call sites collapse into one factory, and the refactor holds up — this is real dedup, not code moved around. Distillation drops ~40 hand-rolled lines, and the routed-experts dtype resolution that was copy-pasted in grpo.py and setup.py now lives in one place. Nothing here blocks the merge.

The new tests in test_nemo_gym_utils.py are a genuine improvement. The coverage migration is complete, and it adds three cases that existed nowhere before: the num_gpu_nodes == 0 branch, dtype resolution with router replay on, and the default to int16 with it off. The ray.get ordering assertion is stricter than the one it replaces. build_nemo_gym_config also copies env_configs["nemo_gym"] rather than mutating the caller's config, and there is a test pinning that.

Two behavior changes are not mentioned in the PR description. Both look like fixes; no code change needed, but each deserves a line in the description:

VIRTUAL_ENV / UV_PROJECT_ENVIRONMENT for the grpo and SingleController Gym actors now point at the venv root instead of the python binary, because those paths moved onto make_actor_runtime_env. Distillation already worked this way, and every consumer in the tree treats these as directories, so the new value is the right one.
Distillation now gets tokenizer_config and pad_dynamic_image_shapes promoted to top-level NemoGymConfig fields for the first time, which fixes VLM distillation. To be exact: this was not broken in production — no in-tree config exercises it. All three is_vlm: true configs under examples/configs/recipes/vlm/ are GRPO recipes, and no distillation config sets it. So it fixes a path nothing currently runs.

Environment caveat, stated once: this review is read from source at the pinned SHAs, with no GPU. The exception is the unit tests, which were actually run in a hand-built CPU-only venv — 21 passed in test_nemo_gym_utils.py, 22 in test_distillation.py, 16 selected in test_grpo.py, and 68 in test_setup.py with one pre-existing failure unrelated to this PR (test_single_controller_mopd_recipe_resolves_to_runtime_contract, which dies on an unset HF_HOME). Ruff is clean: the pinned 0.9.9 hooks (check, check --select I, format --check) all pass. CI is 60/60 green.

Generated by Claude Code

ananthsub · 2026-09-03T06:13:57Z

/ok to test 8aa1cb9

ananthsub · 2026-09-04T19:08:15Z

/ok to test 8aa1cb9

The NemoGym actor was constructed in three places (grpo.setup, distillation.setup, single_controller_utils.setup) that had drifted apart, each forgetting a different field: - single_controller never passed use_fastokens or routed_experts_dtype - distillation never passed require_routed_experts - grpo and single_controller built runtime_env inline and set VIRTUAL_ENV to the *interpreter path* rather than the venv root Collapse all three onto spinup_nemo_gym_actor(), split out build_nemo_gym_config() so the config split (NeMo-RL knobs vs Gym's global config) is testable on its own, and use the shared make_actor_runtime_env() helper, which resolves VIRTUAL_ENV correctly. This makes the actor-creation path a single place to extend for multi-actor sharding. Distillation keeps require_routed_experts off: it never configures vLLM to emit routed experts, so requiring them would fail every rollout. The construction assertions that lived in the distillation tests move to the factory's own test module, where they cover all three callers at once. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

The deferred vLLM path reserves one URL slot per data-parallel rank, so the list can contain None until the servers finish loading. Match the shared NeMo-Gym factory's input type instead of narrowing it to list[str]. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

Update the bottom stack layer's mocks for tokenizer installation and the centralized factory signature so every PR remains independently testable. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

Keep the concurrent setup double aligned with the actor helper so the test exercises the intended handshake instead of failing in its worker thread. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

Make router replay and tokenizer behavior explicit at every call site, keep model dtype resolution patchable in the owning module, and document and test the config fields promoted for actor use. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

ananthsub · 2026-09-04T19:15:11Z

/ok to test c02605e

… lint after main merge The origin/main merge (f6305dc) auto-merged nemo_gym.py without conflict but dropped the token_capture kwarg from the new spinup_nemo_gym_actor wrapper introduced by #3367, while setup.py still passed it. Restore the parameter and forward it to build_nemo_gym_config; assert in the unit test that it reaches the actor config. Also fix the Lint CI job: point pyrefly.toml at the renamed rollout_reassembler*.py modules (the stale blackbox_finalizer / finalizer_actor entries made pyrefly abort), and apply ruff format/isort to the files the merge left unformatted. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@nvidia.com>

The four generation backends disagreed on what `dp_openai_server_base_urls` holds when there is no HTTP server to report. Megatron returned `[]`; vLLM, TRT-LLM and Dynamo returned a list of `None`s. The placeholder does not stand for anything -- in those branches no server exists, so the length carries no information -- but it does make the value lie to a plain truth test, since `[None]` is truthy. Callers had to know that and write `if not any(urls)` instead of `if not urls`, and a sync run printed `Reserved 1 vLLM server URLs: [None]`. Match Megatron and narrow the type to `list[str]`. `NemoGymConfig.base_urls` and `spinup_nemo_gym_actor` already declare `list[str]`, so the producers now match what the consumers ask for; NVIDIA-NeMo#3367 had to add a `cast(list[str], ...)` in `distillation.py` to bridge the gap. vLLM needed more than swapping the sentinel. A worker leaves `base_url` unset unless `expose_http_server` started a server, so an async engine without one answered `[None] * dp_size` from the workers themselves, not from the guarded branch. Both `_report_dp_openai_server_base_urls` and `_collect_reserved_urls` now report that case as empty. Emptiness is decided for the list as a whole rather than by dropping entries. `expose_http_server` is one config value for every worker, so the answer is all-or-nothing; and `GenerationFleetHealth` indexes these URLs by shard and rejects a list whose length does not match `shard_count`, so filtering entry by entry would turn a degraded shard into a startup failure with a length error that names neither cause. No caller changes needed: `_shard_base_urls` already collapses an all-None list to None and does the same with an empty one, `_maybe_start_generation_router` filters falsy entries before its own emptiness check, and NeMo-Gym is reachable only behind the `should_expose_http_server` assert, so the no-server list never gets there. Nothing reads the length in the no-server case. The redundant `any(...)` guards stay as they are. Signed-off-by: Yupeng Tang <85978465+yupengtang@users.noreply.github.com>

ananthsub force-pushed the ansubramania/gym-shards/factory branch from 1414943 to a0df64f Compare August 13, 2026 21:44

ananthsub added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Aug 13, 2026

ananthsub marked this pull request as ready for review August 13, 2026 21:54

ananthsub requested review from a team as code owners August 13, 2026 21:54

ananthsub mentioned this pull request Aug 13, 2026

fix(nemo-gym): define actor failure and teardown ownership #3368

Open

ananthsub force-pushed the ansubramania/gym-shards/factory branch from 76447f3 to 9a7060f Compare August 14, 2026 17:07

ananthsub force-pushed the ansubramania/gym-shards/factory branch from 9a7060f to 9a140b1 Compare August 31, 2026 18:02

ananthsub changed the title ~~refactor(nemo-gym): centralize the gym actor factory~~ refactor(nemo-gym): centralize actor construction Aug 31, 2026

ananthsub force-pushed the ansubramania/gym-shards/factory branch from 9a140b1 to e190399 Compare September 1, 2026 06:26

ananthsub force-pushed the ansubramania/gym-shards/factory branch from e190399 to 45e305f Compare September 1, 2026 17:42

ananthsub force-pushed the ansubramania/gym-shards/factory branch from 45e305f to a77c04e Compare September 2, 2026 05:50

terrykong mentioned this pull request Sep 3, 2026

Generation backends disagree on how to report "no OpenAI server": [None] vs [] #3977

Open

terrykong reviewed Sep 3, 2026

View reviewed changes

Comment thread nemo_rl/environments/nemo_gym.py Outdated

Comment thread tests/unit/environments/test_nemo_gym_utils.py

Comment thread nemo_rl/environments/nemo_gym.py Outdated

Comment thread nemo_rl/environments/nemo_gym.py

Comment thread nemo_rl/environments/nemo_gym.py Outdated

ananthsub force-pushed the ansubramania/gym-shards/factory branch from a77c04e to 8aa1cb9 Compare September 3, 2026 06:06

ananthsub requested a review from terrykong September 3, 2026 09:10

ananthsub added 3 commits September 4, 2026 12:14

test(nemo-gym): align factory expectations with actor setup

b9742fa

Update the bottom stack layer's mocks for tokenizer installation and the centralized factory signature so every PR remains independently testable. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

ananthsub added 2 commits September 4, 2026 12:14

test(nemo-gym): match positional environment config

60a6489

Keep the concurrent setup double aligned with the actor helper so the test exercises the intended handshake instead of failing in its worker thread. Signed-off-by: Ananth Subramaniam <ansubramania@nvidia.com>

ananthsub force-pushed the ansubramania/gym-shards/factory branch from 8aa1cb9 to c02605e Compare September 4, 2026 19:14

terrykong approved these changes Sep 4, 2026

View reviewed changes

terrykong merged commit 971d03c into main Sep 4, 2026
86 checks passed

terrykong deleted the ansubramania/gym-shards/factory branch September 4, 2026 22:35

yupengtang mentioned this pull request Sep 5, 2026

fix(generation): report no OpenAI server as an empty list #3993

Open

4 tasks

Uh oh!

Conversation

ananthsub commented Jul 27, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Stack

Uh oh!

copy-pr-bot Bot commented Jul 27, 2026

Uh oh!

copy-pr-bot Bot commented Aug 13, 2026

Uh oh!

ananthsub commented Aug 13, 2026

Uh oh!

ananthsub commented Aug 13, 2026

Uh oh!

ananthsub commented Sep 1, 2026

Uh oh!

ananthsub commented Sep 1, 2026

Uh oh!

ananthsub commented Sep 1, 2026

Uh oh!

ananthsub commented Sep 2, 2026

Uh oh!

terrykong left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

ananthsub commented Sep 3, 2026

Uh oh!

ananthsub commented Sep 4, 2026

Uh oh!

ananthsub commented Sep 4, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

ananthsub commented Jul 27, 2026 •

edited

Loading