chore(customizer): update RL image

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 = ""

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; \

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/<env>         per-ENVIRONMENT venvs
-  ├─ built-in Gym envs   ── prefetched at build (NEMO_GYM_PREFETCH_CONFIGS)
+/opt/gym_venvs/<env>         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/<env>
-                          ├─ 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. |
 

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