diff --git a/docker-bake.hcl b/docker-bake.hcl index 78450452257..7839541fd69 100644 --- a/docker-bake.hcl +++ b/docker-bake.hcl @@ -100,7 +100,7 @@ variable "NEMO_RL_REPO" { # RL pins Gym as a git submodule (-> soluwalana/Gym over https), so Gym rides in with the RL git ADD # - no separate Gym pin needed. variable "NEMO_RL_REF" { - default = "14c0e020aa0c2c8ab55f4c57e7cea645ae6eb373" # soluwalana/RL nmp/customizer + default = "706382658e8b3b6292d0b96e83102a610b5f5491" # soluwalana/RL nmp/customizer } variable "RL_BASE_CONTEXT" { default = "" diff --git a/docker/rl/Dockerfile.nmp-rl-base b/docker/rl/Dockerfile.nmp-rl-base index 3ec7491478c..abbd8e6056b 100644 --- a/docker/rl/Dockerfile.nmp-rl-base +++ b/docker/rl/Dockerfile.nmp-rl-base @@ -278,16 +278,29 @@ RUN uv pip install --python /opt/nemo_rl_venv/bin/python --no-deps -e /opt/nemo- # deep_ep / mamba-ssm / causal-conv1d on every job, with a cold uv cache. Prefetching bakes them in # (same approach as NeMo-RL's own release stage). # -# Prefetched (filters are SUBSTRING matches against ACTOR FQNs, not extra names). The four filters -# below resolve to FIVE venvs - exactly what DPO + GRPO + Gym need: +# Prefetched (filters are SUBSTRING matches against ACTOR FQNs, not extra names). One venv is built +# PER ACTOR, not per extra - prefetch_venvs.py passes the actor FQN as the venv name - so actors +# sharing an extra still need their own filter. The six filters below resolve to SEVEN venvs, +# exactly what DPO + GRPO + sandboxed Gym need: # dtensor_policy_worker.DTensorPolicyWorker -> --extra fsdp (DPO/GRPO policy training) # vllm.vllm_worker -> --extra vllm (GRPO generation) - matches BOTH # VllmGenerationWorker and VllmAsyncGenerationWorker, which is wanted: NeMo-Gym forces async # rollouts, so the async engine is on the GRPO+Gym path. # sync_rollout_actor.SyncRolloutActor -> --extra vllm (GRPO rollout driver, sync path) -# nemo_gym.NemoGym -> --extra nemo_gym (Gym environment actor) +# nemo_gym.NemoGym -> --extra nemo_gym (Gym environment actor, mode A) +# nemo_gym_actor.SandboxedGymActor -> --extra nemo_gym (trusted proxy actor, mode B - the +# Gym stack itself runs in the job sandbox; this actor never imports user environment code) +# broker_actor.SandboxEpisodeBrokerActor -> --extra nemo_gym (trusted episode broker: creates +# per-episode sandboxes on the agent's behalf so the untrusted job sandbox never holds the +# OpenSandbox credential) # The filters are intentionally specific: a bare "vllm" would also match nemo_rl.modelopt's -# vllm_quant_worker and drag in the modelopt+vllm combination. +# vllm_quant_worker and drag in the modelopt+vllm combination, and a bare "nemo_gym" would not +# distinguish the three nemo_gym-extra actors (it would match all three, which is fine, but naming +# them keeps this list readable against ray_actor_environment_registry.py). +# +# The opensandbox / tenacity SDKs the sandbox path needs arrive through the extra itself: RL +# declares `nemo_gym = ["nemo_gym[sandbox]"]`, and Gym's `sandbox` extra pins opensandbox>=0.1.9 and +# tenacity>=9.1.4. # # NOT prefetched: automodel / mcore (unused by our DPO+GRPO path - see README), sglang, trtllm, # modelopt-quant, and the async-GRPO bookkeeping actors AsyncTrajectoryCollector / ReplayBuffer @@ -299,40 +312,35 @@ uv run nemo_rl/utils/prefetch_venvs.py \ dtensor_policy_worker.DTensorPolicyWorker \ vllm.vllm_worker \ sync_rollout_actor.SyncRolloutActor \ - nemo_gym.NemoGym - -# Gym sandbox CLIENT SDK, into the NemoGym ACTOR venv. NeMo-RL's lock carries opensandbox only -# behind nemo_gym's own `sandbox` extra, so `uv run --extra nemo_gym` (the executable this actor -# launches with) does NOT install it. NeMo-Gym imports it lazily in -# nemo_gym/sandbox/providers/opensandbox/provider.py, so a missing package would surface only on -# first sandbox use rather than at startup. Installing into the base venv alone is not enough: -# the actor runs from /opt/ray_venvs, not /opt/nemo_rl_venv. -GYM_VENV=/opt/ray_venvs/nemo_rl.environments.nemo_gym.NemoGym -if [ -x "${GYM_VENV}/bin/python" ]; then - uv pip install --python "${GYM_VENV}/bin/python" "opensandbox>=0.1.9" "tenacity>=9.1.4" -else - echo "ERROR: expected prefetched NemoGym venv at ${GYM_VENV}" >&2 - exit 1 -fi + nemo_gym.NemoGym \ + nemo_gym_actor.SandboxedGymActor \ + broker_actor.SandboxEpisodeBrokerActor + +# Fail loudly if the registry moved an actor out from under a filter above - otherwise the venv +# would silently be built on the node at first use instead, turning a build error into a slow job. +for actor in \ + nemo_rl.environments.nemo_gym.NemoGym \ + nemo_rl.environments.sandbox.nemo_gym_actor.SandboxedGymActor \ + nemo_rl.environments.sandbox.broker_actor.SandboxEpisodeBrokerActor +do + if [ ! -x "/opt/ray_venvs/${actor}/bin/python" ]; then + echo "ERROR: expected prefetched venv at /opt/ray_venvs/${actor}" >&2 + echo " check the filters above against ray_actor_environment_registry.py" >&2 + exit 1 + fi +done EOF -# ---- Bundled NeMo-Gym environment venvs (GRPO only) ---- +# ---- Bundled NeMo-Gym environment venvs (GRPO only; OFF by default, opt in below) ---- # SECOND, separate venv layer from /opt/ray_venvs above: NeMo-Gym gives every environment server its # own venv under NEMO_GYM_VENV_DIR (/opt/gym_venvs), built from that server's own pyproject.toml / # requirements.txt. The isolation is required - env deps (rdkit, spacy, numpy 2.x) would otherwise # clobber RL's pinned stack. # -# Baking the SHIPPED environments in removes their install cost from job startup. The default config -# is NeMo-RL's curated union of the environments used across the RLVR / SWE / RLHF stages; it is -# purpose-built for this (dry-run NemoGym._spinup, dummy interpolation values) and is kept in step -# with the Gym pin. Override with a space-separated list of config paths, or set empty to skip. -# -# NOT covered (by design): user-supplied environment FileSets, which are downloaded per job and so -# cannot exist at build time - those install into /opt/gym_venvs at first use. A `wheels-v1` FileSet -# vendors its wheels and installs without PyPI egress; `native-v1` installs from source. -# -# DPO does not use NeMo-Gym at all (no rollout/environment stage), so this layer is GRPO-only. -ARG NEMO_GYM_PREFETCH_CONFIGS=examples/nemo_gym/prefetch_super_all_envs.yaml +# Nothing is prefetched: every environment installs into /opt/gym_venvs at first use on the node. +# To bake a set in, pass a space-separated list of Gym config paths - e.g. NeMo-RL's curated +# examples/nemo_gym/prefetch_super_all_envs.yaml. +ARG NEMO_GYM_PREFETCH_CONFIGS="" # Gym's model servers pin a plain `vllm==0.20.0`, which resolves to the CUDA 12 wheel on PyPI - # it would fail to load CUDA libraries inside this CUDA 13 image. Override it to vLLM's prebuilt # cu130 GitHub wheel (the same one RL's policy pins) and keep torch on the matching backend, so @@ -343,7 +351,7 @@ ARG TARGETARCH RUN <<"EOF" bash -exu export UV_LINK_MODE=symlink mkdir -p /opt/gym_venvs -if [ -n "${NEMO_GYM_PREFETCH_CONFIGS}" ]; then +if [ -n "${NEMO_GYM_PREFETCH_CONFIGS:-}" ]; then cd /opt/nemo-rl case "${TARGETARCH:-amd64}" in arm64) WHEEL_ARCH=aarch64 ;; *) WHEEL_ARCH=x86_64 ;; esac GYM_VLLM_OVERRIDE=/opt/nemo-rl/.gym-vllm-override.txt @@ -395,9 +403,10 @@ FROM builder AS nmp-rl-base # toolkit / nvcc (deep_ep JIT-compiles kernels at runtime), gcc/build-essential + libibverbs-dev # (JIT + native-v1 env-dep compiles), git + curl (git+https env deps, downloads). # -# /opt/uv_cache is NOT pruned and must never be: the venvs under /opt/ray_venvs and /opt/gym_venvs -# are symlinked into it (UV_LINK_MODE=symlink), so deleting it breaks every one of them. It also -# doubles as a warm cache for environments installed at runtime (user Gym FileSets, non-prefetched +# /opt/uv_cache is NOT pruned and must never be: the per-worker venvs under /opt/ray_venvs are +# symlinked into it (UV_LINK_MODE=symlink), so deleting it breaks every one of them - as it would +# any /opt/gym_venvs environment baked in via NEMO_GYM_PREFETCH_CONFIGS. It also doubles as a warm +# cache for everything installed at runtime (Gym environments, user FileSets, non-prefetched # actors), making those installs faster. RUN apt-get purge -y ccache vim vim-common less >/dev/null 2>&1 || true; \ apt-get autoremove -y >/dev/null 2>&1 || true; \ diff --git a/docker/rl/README.md b/docker/rl/README.md index 9e08155b648..48ed864e9f8 100644 --- a/docker/rl/README.md +++ b/docker/rl/README.md @@ -182,15 +182,15 @@ Two sources of environments, with different timing: | Environment source | Venv built | Cost | |---|---|---| -| **Built into the Gym repo** (math_with_judge, code_gen, swe_agents, …) | At **image build** — `NEMO_GYM_PREFETCH_CONFIGS` defaults to NeMo-RL's curated `examples/nemo_gym/prefetch_super_all_envs.yaml` | Paid once at build; zero at job start | +| **Built into the Gym repo** (math_with_judge, code_gen, swe_agents, …) | At **runtime**, on first spin-up — prefetch is off by default (`NEMO_GYM_PREFETCH_CONFIGS` is empty) | Paid per job/node, until you opt back in | | **User-supplied env FileSet** (downloaded per job) | At **runtime**, on first spin-up | Paid per job/node — unavoidable, the env isn't known at build time | -The default prefetch config is the union of the environments NeMo-RL uses across its RLVR / -SWE / RLHF stages, and is maintained in step with the Gym pin. Set -`NEMO_GYM_PREFETCH_CONFIGS` to a different space-separated list of config paths to bake a -different set, or to empty to skip the layer entirely (every environment then installs at -runtime). Baking these environments is a deliberate build-time / image-size cost traded for -job startup latency. +**Prefetch is disabled by default.** +To bake a set back in, set `NEMO_GYM_PREFETCH_CONFIGS` to a space-separated list of config +paths — NeMo-RL's `examples/nemo_gym/prefetch_super_all_envs.yaml` is the curated union of the +environments it uses across its RLVR / SWE / RLHF stages, purpose-built for prefetching and +maintained in step with the Gym pin. That trades build time and image size for job startup +latency. User environments therefore *do* add startup time, and cannot be prebaked. Two things bound it: @@ -222,8 +222,9 @@ IMAGE (built once) RUNTIME ├─ …SyncRolloutActor [vllm] ────────> rollout driver (GRPO only) └─ …NemoGym [nemo_gym] ────────> Gym actor (GRPO only) -/opt/gym_venvs/ per-ENVIRONMENT venvs - ├─ built-in Gym envs ── prefetched at build (NEMO_GYM_PREFETCH_CONFIGS) +/opt/gym_venvs/ per-ENVIRONMENT venvs — empty in the shipped image + ├─ built-in Gym envs ── created at RUNTIME (prefetch off; opt in via + │ NEMO_GYM_PREFETCH_CONFIGS to bake them at build) └─ user FileSet envs ── created at RUNTIME from the env's pyproject/requirements ``` @@ -244,9 +245,13 @@ driver (base venv) ├─ Generation ─ VllmGenerationWorker ×N → /opt/ray_venvs/…VllmGenerationWorker [vllm] │ └─ deep_ep / deep_gemm → Hopper+ only ├─ Rollout ─ SyncRolloutActor → /opt/ray_venvs/…SyncRolloutActor [vllm] - └─ Env ─ NemoGym actor → /opt/ray_venvs/…NemoGym [nemo_gym] + └─ Env ─ mode A: NemoGym → /opt/ray_venvs/…NemoGym [nemo_gym] + mode B: SandboxedGymActor + SandboxEpisodeBrokerActor + → /opt/ray_venvs/…SandboxedGymActor [nemo_gym] + /opt/ray_venvs/…BrokerActor [nemo_gym] + └─ Gym stack + user FileSet run in an isolated OpenSandbox pod └─ per-environment venv → /opt/gym_venvs/ - ├─ shipped env → prebuilt in image (no startup cost) + ├─ shipped env → built on first use (prefetch off by default) └─ user FileSet → built on first use (startup cost; wheels-v1 avoids PyPI, native-v1 needs egress) ``` @@ -392,12 +397,19 @@ Prefetched (the filters match **actor FQNs**, not extra names): | `dtensor_policy_worker.DTensorPolicyWorker` | `fsdp` | DPO + GRPO policy training | | `vllm.vllm_worker` | `vllm` | GRPO generation — matches **both** `VllmGenerationWorker` and `VllmAsyncGenerationWorker` (NeMo-Gym forces async rollouts, so both are on the path) | | `sync_rollout_actor.SyncRolloutActor` | `vllm` | GRPO rollout driver (sync path) | -| `nemo_gym.NemoGym` | `nemo_gym` | Gym environment actor | +| `nemo_gym.NemoGym` | `nemo_gym` | Gym environment actor (mode A, colocated) | +| `nemo_gym_actor.SandboxedGymActor` | `nemo_gym` | Sandboxed Gym (mode B) — the trusted proxy actor in the training pod | +| `broker_actor.SandboxEpisodeBrokerActor` | `nemo_gym` | Trusted episode broker — creates per-episode sandboxes so the job sandbox never holds the OpenSandbox credential | -Filters are **substring matches on actor FQNs**, so four filters yield five venvs. They are +Filters are **substring matches on actor FQNs**, so six filters yield seven venvs. They are deliberately specific — a bare `vllm` would also match `nemo_rl.modelopt`'s `vllm_quant_worker` and pull in the modelopt+vllm combination. +One venv is built **per actor, not per extra** — `prefetch_venvs.py` passes the actor FQN as +the venv name — so the three `nemo_gym`-extra actors above each get their own directory and +each needs its own filter. `opensandbox` / `tenacity` come in through the extra itself, since +RL declares `nemo_gym = ["nemo_gym[sandbox]"]`. + Without this, each venv is built **on the node at first run**, re-resolving and recompiling `deep_ep` / `mamba-ssm` / `causal-conv1d` against a cold uv cache on every job. Not prefetched (they build on the node if a config selects them): `automodel`, @@ -469,12 +481,12 @@ every import. ### Levers, if size becomes a problem -- **`NEMO_GYM_PREFETCH_CONFIGS`** — the environment venvs are 4.2 GB, but the build step that - produces them costs far more than that, because resolving 25 environments pulls their - dependencies into the shared cache and materializes the SWE agent setups (~5.9 GB). - Narrowing the list, or setting it empty, moves that cost from image size to job-start - latency. It is kept broad on purpose so GRPO runs do not install environments at startup, - and `swe_agents` in particular is required by the sandboxed-Gym work. +- **`NEMO_GYM_PREFETCH_CONFIGS`** — already empty by default, which is the single largest + saving available: the environment venvs are 4.2 GB, and the build step that produces them + costs far more than that, because resolving 25 environments pulls their dependencies into the + shared cache and materializes the SWE agent setups (~5.9 GB). Setting it back to a non-empty + list moves that cost from job-start latency to image size and build time. Note `swe_agents` is + the set the sandboxed-Gym work uses, and it is also the least reliable to build. - **`--extra modelopt`** can be dropped from the warmup sync if quantization is not needed. - The **CUDA devel toolkit** cannot be removed: `deep_ep` JIT-compiles kernels at runtime and `native-v1` environments compile dependencies on the node. @@ -488,5 +500,5 @@ every import. | `TORCH_CUDA_ARCH_LIST` | `"9.0 10.0"` | Archs for the torch-based source extensions (deep_ep, deep_gemm, mamba, causal-conv1d). | | `NVTE_CUDA_ARCHS` | `90;100` | Archs for Transformer-Engine. Inert today (TE is only in the unused `automodel`/`mcore` extras); kept for when either is enabled. | | `UV_SYNC_MODE` | `--frozen` | Reproducible sync. Set to empty to relock if a bumped RL commit's lock has drifted. | -| `NEMO_GYM_PREFETCH_CONFIGS` | `examples/nemo_gym/prefetch_super_all_envs.yaml` | Space-separated Gym config paths whose environment venvs are baked into `/opt/gym_venvs`. Set empty to skip (every environment then installs at runtime). | +| `NEMO_GYM_PREFETCH_CONFIGS` | *(empty — prefetch off)* | Space-separated Gym config paths whose environment venvs are baked into `/opt/gym_venvs`. Empty means every environment installs at runtime on first use. Set to `examples/nemo_gym/prefetch_super_all_envs.yaml` to bake NeMo-RL's curated set back in. | diff --git a/tests/smoke_gpu/test_rl_training.py b/tests/smoke_gpu/test_rl_training.py index c997c0d435c..3d607751106 100644 --- a/tests/smoke_gpu/test_rl_training.py +++ b/tests/smoke_gpu/test_rl_training.py @@ -34,9 +34,9 @@ # Actor FQN -> packages that must import inside that actor's venv. # -# MUST list every venv the build prefetches (the four filters in docker/rl/Dockerfile.nmp-rl-base -# resolve to FIVE actors, because `vllm.vllm_worker` matches the sync and async workers alike). -# Listing all five is what makes a broken prefetch filter fail here instead of silently shipping an +# MUST list every venv the build prefetches (the six filters in docker/rl/Dockerfile.nmp-rl-base +# resolve to SEVEN actors, because `vllm.vllm_worker` matches the sync and async workers alike). +# Listing all seven is what makes a broken prefetch filter fail here instead of silently shipping an # image whose workers rebuild their venv on the node at job start. WORKER_VENV_IMPORTS = { # DPO + GRPO policy training (--extra fsdp) @@ -66,11 +66,27 @@ "torch", "vllm", ], - # NeMo-Gym environment actor (--extra nemo_gym); opensandbox is the sandbox client SDK + # NeMo-Gym environment actor (--extra nemo_gym). opensandbox is the sandbox client SDK; it + # reaches this venv only through RL's `nemo_gym = ["nemo_gym[sandbox]"]`, so this assertion is + # what proves the extra really carries it. "nemo_rl.environments.nemo_gym.NemoGym": [ "nemo_gym", "opensandbox", ], + # Sandboxed-Gym mode B: the trusted proxy actor in the training pod. Shares the nemo_gym extra + # with NemoGym, but venvs are named per ACTOR, so it gets its own and needs its own filter. + "nemo_rl.environments.sandbox.nemo_gym_actor.SandboxedGymActor": [ + "nemo_gym", + "opensandbox", + "nemo_rl.environments.sandbox.nemo_gym_actor", + ], + # Trusted episode broker: creates per-episode sandboxes so the untrusted job sandbox never + # holds the OpenSandbox credential. + "nemo_rl.environments.sandbox.broker_actor.SandboxEpisodeBrokerActor": [ + "nemo_gym", + "opensandbox", + "nemo_rl.environments.sandbox.broker_actor", + ], } # Packages that link libcuda.so.1 — the NVIDIA *driver*, which the container runtime injects on a