From 96be9f68710d2a356817318bc0a8d34074498d80 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Mon, 31 Aug 2026 19:22:36 -0500 Subject: [PATCH 01/21] perf(docker): hardlink worker venvs into layers Replaces the symlink scheme with hardlinks. Increased size of image by ~3%. Reduces number of symlinks from ~1.5M to ~3k. Improves SQSH conversion time from ~3 hours to ~25 minutes. Signed-off-by: Teodor-Dumitru Ene --- docker/Dockerfile | 82 ++++++++++++++--- docker/venv_prefetch_manifest.tsv | 22 +++++ nemo_rl/utils/venv_prefetch_manifest.py | 90 +++++++++++++++++++ nemo_rl/utils/venvs.py | 9 +- pyrefly.toml | 1 + .../unit/utils/test_venv_prefetch_manifest.py | 24 +++++ 6 files changed, 213 insertions(+), 15 deletions(-) create mode 100644 docker/venv_prefetch_manifest.tsv create mode 100644 nemo_rl/utils/venv_prefetch_manifest.py create mode 100644 tests/unit/utils/test_venv_prefetch_manifest.py diff --git a/docker/Dockerfile b/docker/Dockerfile index b4e8c8b969b..bf1d8a70681 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -195,6 +195,8 @@ ENV LD_LIBRARY_PATH="/opt/nemo_rl_venv/lib/python3.13/site-packages/z3/lib:/opt/ # First copy only the dependency files COPY --from=nemo-rl pyproject.toml uv.lock ./ +# Prefetch worker-venv manifest to avoid the need for symlinks. +COPY --from=nemo-rl docker/venv_prefetch_manifest.tsv ./docker/venv_prefetch_manifest.tsv # Copy in the top level __init__.py/package_info.py since build-custom-vllm.sh needs the nemo_rl package to exist. COPY --from=nemo-rl nemo_rl/__init__.py nemo_rl/package_info.py ./nemo_rl/ COPY --from=nemo-rl tools/build-custom-vllm.sh ./tools/build-custom-vllm.sh @@ -234,7 +236,7 @@ fi # to warm the uv cache, then at the end just sync the default dependencies. # Do everything in one layer to prevent large layers. -# The venv is symlinked to avoid bloating the layer size +# Everything is hardlinked against the uv cache; this layer holds exactly one real copy of every wheel. UV_LINK_MODE=hardlink uv sync --frozen --no-install-project if [[ -z "${SKIP_VLLM_BUILD:-}" ]]; then UV_LINK_MODE=hardlink uv sync --frozen --extra vllm --no-install-project @@ -242,10 +244,41 @@ fi if [[ -z "${SKIP_SGLANG_BUILD:-}" ]]; then UV_LINK_MODE=hardlink uv sync --frozen --extra sglang --no-install-project fi -uv sync --link-mode symlink --frozen --extra mcore --no-install-project -uv sync --link-mode symlink --frozen --extra automodel --no-install-project -uv sync --link-mode symlink --frozen --extra modelopt --no-install-project -uv sync --link-mode symlink --frozen --all-groups --no-install-project +UV_LINK_MODE=hardlink uv sync --frozen --extra mcore --no-install-project +UV_LINK_MODE=hardlink uv sync --frozen --extra automodel --no-install-project +UV_LINK_MODE=hardlink uv sync --frozen --extra modelopt --no-install-project +UV_LINK_MODE=hardlink uv sync --frozen --all-groups --no-install-project + +# Worker-venv prefetch, phase 1: +# Materialize each registry venv's third-party packages in THIS layer, hardlinked against the above cache, +# so the N worker venvs cost directory entries instead of N wheel copies +PREFETCH_NEGATIVE_FILTERS="" +if [[ -n "${SKIP_VLLM_BUILD:-}" ]]; then + PREFETCH_NEGATIVE_FILTERS="$PREFETCH_NEGATIVE_FILTERS vllm" +fi +if [[ -n "${SKIP_SGLANG_BUILD:-}" ]]; then + PREFETCH_NEGATIVE_FILTERS="$PREFETCH_NEGATIVE_FILTERS sglang" +fi +while IFS=$'\t' read -r venv_name stage extras; do + if [[ -z "$venv_name" || "$venv_name" == \#* ]]; then + continue + fi + skip="" + for f in $PREFETCH_NEGATIVE_FILTERS; do + if [[ "$venv_name" == *"$f"* ]]; then + skip=1 + fi + done + if [[ -n "$skip" ]]; then + continue + fi + if [[ "$stage" == "trtllm" ]]; then + extras="" + fi + venv_path="${NEMO_RL_VENV_DIR}/${venv_name}" + uv venv --allow-existing "$venv_path" + UV_PROJECT_ENVIRONMENT="$venv_path" UV_LINK_MODE=hardlink uv sync --frozen $extras --no-install-project +done < docker/venv_prefetch_manifest.tsv # Remove the aiohttp in this uv cache dir to fully address CVE GHSA-mqqc-3gqh-h2x8 # The ray install will include the older aiohttp version in its cache @@ -298,9 +331,14 @@ du -sh /root/.cache/uv /root/.cache/trtllm-wheels TRTLLM_SYNC_LOG=$(mktemp /tmp/trtllm-sync.XXXXXX.log) set +e +# Build TRT-LLM into a throwaway venv instead of the main one: +# the main venv is fully hardlinked from the dependency layer, +# and churning it here would copy files up into this layer. +# The throwaway venv is created and deleted within this RUN, so it leaves no bytes in the layer. UV_CACHE_DIR=/root/.cache/uv \ TRTLLM_WHEEL_CACHE_DIR=/root/.cache/trtllm-wheels \ - uv sync --verbose --link-mode symlink --locked --extra trtllm --no-install-project \ + UV_PROJECT_ENVIRONMENT=/tmp/trtllm-build-venv \ + uv sync --verbose --link-mode hardlink --locked --extra trtllm --no-install-project \ 2>&1 | tee "$TRTLLM_SYNC_LOG" \ | awk '/\[TRTLLM_CCACHE\]|Ninja progress:/ { print; fflush() }' TRTLLM_SYNC_STATUS=$? @@ -317,11 +355,25 @@ rm -f "$TRTLLM_SYNC_LOG" echo "TRT-LLM cache state after build:" du -sh /root/.cache/uv /root/.cache/trtllm-wheels -# Restore the intended default environment in the same layer so the transient -# TRT-LLM installation does not add a large intermediate venv layer. -uv sync --link-mode symlink --locked --all-groups --no-install-project - -# The final sync can repopulate the shared uv cache, so repeat the security +# Worker-venv prefetch, phase 1 for the "trtllm"-stage manifest entries: +# the dependency layer gave them a base-only warm; now that the tensorrt_llm wheel exists, +# top off the venv with the full extra set so that the added files hardlink in-layer. +while IFS=$'\t' read -r venv_name stage extras; do + if [[ -z "$venv_name" || "$venv_name" == \#* || "$stage" != "trtllm" ]]; then + continue + fi + venv_path="${NEMO_RL_VENV_DIR}/${venv_name}" + uv venv --allow-existing "$venv_path" + UV_PROJECT_ENVIRONMENT="$venv_path" UV_LINK_MODE=hardlink \ + TRTLLM_WHEEL_CACHE_DIR=/root/.cache/trtllm-wheels \ + uv sync --frozen $extras --no-install-project +done < docker/venv_prefetch_manifest.tsv + +# The main venv was never touched (the build ran in the throwaway venv), so +# the previous "restore the default environment" sync is no longer needed. +rm -rf /tmp/trtllm-build-venv + +# The syncs above can repopulate the shared uv cache, so repeat the security # cleanup performed in the preceding dependency layer. find /root/.cache/uv -type d -path "*ray/_private/runtime_env/agent/thirdparty_files/aiohttp*" -exec rm -rf {} + EOF @@ -417,16 +469,20 @@ fi export TRTLLM_WHEEL_CACHE_DIR=/root/.cache/trtllm-wheels export TRTLLM_WHEEL_CACHE_MIRROR_DIR=/opt/trtllm_wheels export TRTLLM_REQUIRE_CACHED_WHEEL=1 +# Worker-venv prefetch, phase 2: +# all venvs have by now been materialized in hardlink-mode; +# add the editable project install per venv. if [[ -n "$NEGATIVE_FILTERS" ]]; then - UV_LINK_MODE=symlink uv run nemo_rl/utils/prefetch_venvs.py --negative-filters $NEGATIVE_FILTERS + UV_LINK_MODE=hardlink uv run nemo_rl/utils/prefetch_venvs.py --negative-filters $NEGATIVE_FILTERS else - UV_LINK_MODE=symlink uv run nemo_rl/utils/prefetch_venvs.py + UV_LINK_MODE=hardlink uv run nemo_rl/utils/prefetch_venvs.py fi EOF # Prefetch NeMo Gym internal venvs (gym servers like code_gen, math_with_judge, etc.) # into the image. Gated on NEMO_GYM_PREFETCH_CONFIGS (no-op by default). Runs after the # source COPY above so examples/nemo_gym/prefetch_venvs.py is present. +# Gym server venvs use symlinks, not the hardlink scheme used for worker venvs. # Gym pins vllm==0.24.0; override it to vLLM's default prebuilt CUDA wheel so Gym's venvs share # RL's torch/nvidia stack. UV_TORCH_BACKEND keeps torch on the same CUDA backend. ARG NEMO_GYM_CUDA=cu130 diff --git a/docker/venv_prefetch_manifest.tsv b/docker/venv_prefetch_manifest.tsv new file mode 100644 index 00000000000..e6d901a80fa --- /dev/null +++ b/docker/venv_prefetch_manifest.tsv @@ -0,0 +1,22 @@ +# DO NOT EDIT BY HAND +# Generated by `uv run python -m nemo_rl.utils.venv_prefetch_manifest`. +# tests/unit/utils/test_venv_prefetch_manifest.py ensures this file is not stale. +# Columns: +nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector deps --extra vllm +nemo_rl.algorithms.async_utils.ReplayBuffer deps --extra vllm +nemo_rl.environments.nemo_gym.NemoGym deps --extra nemo_gym +nemo_rl.experience.sync_rollout_actor.SyncRolloutActor deps --extra vllm +nemo_rl.modelopt.models.generation.vllm_quant_worker.VllmQuantAsyncGenerationWorker deps --extra modelopt --extra vllm +nemo_rl.modelopt.models.generation.vllm_quant_worker.VllmQuantGenerationWorker deps --extra modelopt --extra vllm +nemo_rl.modelopt.models.policy.workers.dtensor_quant_policy_worker.DTensorQuantPolicyWorker deps --extra modelopt --extra automodel +nemo_rl.modelopt.models.policy.workers.dtensor_quant_policy_worker_v2.DTensorQuantPolicyWorkerV2 deps --extra modelopt --extra automodel +nemo_rl.modelopt.models.policy.workers.megatron_quant_policy_worker.MegatronQuantPolicyWorker deps --extra modelopt --extra mcore +nemo_rl.models.generation.sglang.sglang_worker.SGLangGenerationWorker deps --extra sglang +nemo_rl.models.generation.trtllm.trtllm_worker_async.TrtllmAsyncGenerationWorker trtllm --extra trtllm +nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker deps --extra vllm +nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker deps --extra vllm +nemo_rl.models.policy.workers.dtensor_policy_worker.DTensorPolicyWorker deps --extra fsdp +nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2 deps --extra automodel +nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker deps --extra mcore +nemo_rl.models.value.workers.dtensor_value_worker_v2.DTensorValueWorkerV2 deps --extra automodel +nemo_rl.models.value.workers.megatron_value_worker.MegatronValueWorker deps --extra mcore diff --git a/nemo_rl/utils/venv_prefetch_manifest.py b/nemo_rl/utils/venv_prefetch_manifest.py new file mode 100644 index 00000000000..92b653d3955 --- /dev/null +++ b/nemo_rl/utils/venv_prefetch_manifest.py @@ -0,0 +1,90 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import os +import shlex +from pathlib import Path + +MANIFEST_PATH = Path(__file__).parents[2] / "docker" / "venv_prefetch_manifest.tsv" + +_HEADER = """\ +# DO NOT EDIT BY HAND +# Generated by `uv run python -m nemo_rl.utils.venv_prefetch_manifest`. +# tests/unit/utils/test_venv_prefetch_manifest.py ensures this file is not stale. +# Columns: \t\t +""" + + +def parse_sync_extras(py_executable: str) -> list[str]: + """Extract the ``--extra `` pairs from a registry py_executable.""" + tokens = shlex.split(py_executable) + if tokens[:2] != ["uv", "run"]: + raise ValueError(f"not a uv py_executable: {py_executable!r}") + extras: list[str] = [] + i = 2 + while i < len(tokens): + token = tokens[i] + if token == "--extra": + extras.extend(tokens[i : i + 2]) + i += 2 + elif token == "--directory": + i += 2 + elif token == "--locked": + i += 1 + else: + raise ValueError( + f"unhandled py_executable flag {token!r} in {py_executable!r}; " + "teach nemo_rl/utils/venv_prefetch_manifest.py how it maps to `uv sync` flags" + ) + return extras + + +def build_manifest_rows() -> list[tuple[str, str, str]]: + """Build (venv_name, stage, extras) rows for every uv-managed actor.""" + if os.environ.get("NEMO_RL_PY_EXECUTABLES_SYSTEM", "0") == "1": + raise RuntimeError( + "NEMO_RL_PY_EXECUTABLES_SYSTEM=1 collapses the registry to system " + "python; unset it to generate the manifest" + ) + # Local import — the registry transitively pulls heavy deps (ray, modelopt). + from nemo_rl.distributed.ray_actor_environment_registry import ( + ACTOR_ENVIRONMENT_REGISTRY, + ) + + rows: list[tuple[str, str, str]] = [] + for actor_fqn, py_executable in sorted(ACTOR_ENVIRONMENT_REGISTRY.items()): + if not py_executable.startswith("uv"): + continue + extras = parse_sync_extras(py_executable) + stage = "trtllm" if "trtllm" in extras else "deps" + rows.append((actor_fqn, stage, " ".join(extras))) + return rows + + +def render_manifest() -> str: + """Render the manifest file content.""" + lines = [_HEADER] + for actor_fqn, stage, extras in build_manifest_rows(): + lines.append(f"{actor_fqn}\t{stage}\t{extras}\n") + return "".join(lines) + + +def main() -> None: + """Regenerate docker/venv_prefetch_manifest.tsv in place.""" + MANIFEST_PATH.write_text(render_manifest()) + print(f"wrote {MANIFEST_PATH}") + + +if __name__ == "__main__": + main() diff --git a/nemo_rl/utils/venvs.py b/nemo_rl/utils/venvs.py index 6e632b8d9fb..61e44444e4c 100644 --- a/nemo_rl/utils/venvs.py +++ b/nemo_rl/utils/venvs.py @@ -123,8 +123,13 @@ def create_local_venv( # Command doesn't matter, since `uv` syncs the environment no matter the command. exec_cmd.extend(["echo", f"Finished creating venv {venv_path}"]) - # Always run uv sync first to ensure the build requirements are set (for --no-build-isolation packages) - subprocess.run(["uv", "sync", "--directory", git_root], env=env, check=True) + # Run uv sync first to ensure build requirements are set (for --no-build-isolation packages). + # --inexact: this base-set sync must not prune extras out of a venv that was + # pre-materialized in the image; pruning and re-adding hardlinked packages would copy + # them up into the image's final layer. + subprocess.run( + ["uv", "sync", "--inexact", "--directory", git_root], env=env, check=True + ) subprocess.run(exec_cmd, env=env, check=True) # Return the path to the python executable in the virtual environment diff --git a/pyrefly.toml b/pyrefly.toml index 1b54fdaab50..5832f9f61c9 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -286,6 +286,7 @@ project-includes = [ "nemo_rl/utils/r3_trace.py", "nemo_rl/utils/routed_experts_codec.py", "nemo_rl/utils/timer.py", + "nemo_rl/utils/venv_prefetch_manifest.py", "nemo_rl/utils/venvs.py", "nemo_rl/utils/weight_transfer_http.py", "nemo_rl/utils/weight_transfer_sparse_codec.py", diff --git a/tests/unit/utils/test_venv_prefetch_manifest.py b/tests/unit/utils/test_venv_prefetch_manifest.py new file mode 100644 index 00000000000..e0524475b8f --- /dev/null +++ b/tests/unit/utils/test_venv_prefetch_manifest.py @@ -0,0 +1,24 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from nemo_rl.utils.venv_prefetch_manifest import MANIFEST_PATH, render_manifest + + +def test_manifest_matches_actor_registry(): + """docker/venv_prefetch_manifest.tsv must stay in lockstep with the registry.""" + assert MANIFEST_PATH.read_text() == render_manifest(), ( + "docker/venv_prefetch_manifest.tsv is stale relative to " + "ACTOR_ENVIRONMENT_REGISTRY; regenerate it with " + "`uv run python -m nemo_rl.utils.venv_prefetch_manifest`" + ) From 206a7f2ff8942ed512bc439d7dc689b0a9b5fe7a Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Tue, 1 Sep 2026 21:34:55 -0500 Subject: [PATCH 02/21] Fix TRTLLM layer Signed-off-by: Teodor-Dumitru Ene --- docker/Dockerfile | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index bf1d8a70681..7f4ffa5ef94 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -259,6 +259,9 @@ fi if [[ -n "${SKIP_SGLANG_BUILD:-}" ]]; then PREFETCH_NEGATIVE_FILTERS="$PREFETCH_NEGATIVE_FILTERS sglang" fi +if [[ -n "${SKIP_TRTLLM_BUILD:-}" ]]; then + PREFETCH_NEGATIVE_FILTERS="$PREFETCH_NEGATIVE_FILTERS trtllm" +fi while IFS=$'\t' read -r venv_name stage extras; do if [[ -z "$venv_name" || "$venv_name" == \#* ]]; then continue @@ -334,11 +337,12 @@ set +e # Build TRT-LLM into a throwaway venv instead of the main one: # the main venv is fully hardlinked from the dependency layer, # and churning it here would copy files up into this layer. -# The throwaway venv is created and deleted within this RUN, so it leaves no bytes in the layer. +# The throwaway venv must be symlink-mode: hardlinks to files in the dependency layer's uv +# cache make overlayfs copy them up into this layer, and the copies outlive the rm -rf below. UV_CACHE_DIR=/root/.cache/uv \ TRTLLM_WHEEL_CACHE_DIR=/root/.cache/trtllm-wheels \ UV_PROJECT_ENVIRONMENT=/tmp/trtllm-build-venv \ - uv sync --verbose --link-mode hardlink --locked --extra trtllm --no-install-project \ + uv sync --verbose --link-mode symlink --locked --extra trtllm --no-install-project \ 2>&1 | tee "$TRTLLM_SYNC_LOG" \ | awk '/\[TRTLLM_CCACHE\]|Ninja progress:/ { print; fflush() }' TRTLLM_SYNC_STATUS=$? @@ -358,6 +362,8 @@ du -sh /root/.cache/uv /root/.cache/trtllm-wheels # Worker-venv prefetch, phase 1 for the "trtllm"-stage manifest entries: # the dependency layer gave them a base-only warm; now that the tensorrt_llm wheel exists, # top off the venv with the full extra set so that the added files hardlink in-layer. +# tensorrt_llm itself is new in this layer; trtllm-extra deps already present in the +# dependency layer's cache (e.g. tilelang) get copied up here, a small accepted cost. while IFS=$'\t' read -r venv_name stage extras; do if [[ -z "$venv_name" || "$venv_name" == \#* || "$stage" != "trtllm" ]]; then continue From d09d3eb7def4a11dcda027d85ece520eba419621 Mon Sep 17 00:00:00 2001 From: Teodor-Dumitru Ene Date: Wed, 2 Sep 2026 01:14:52 -0500 Subject: [PATCH 03/21] Move manifest to pyproject.toml Signed-off-by: Teodor-Dumitru Ene --- docker/Dockerfile | 43 ++++++--- docker/venv_prefetch_manifest.tsv | 22 ----- docs/design-docs/dependency-management.md | 17 ++-- docs/design-docs/uv.md | 4 +- .../ray_actor_environment_registry.py | 80 ++++++++--------- nemo_rl/distributed/virtual_cluster.py | 8 +- nemo_rl/utils/venv_prefetch_manifest.py | 90 ------------------- pyproject.toml | 74 +++++++++++++++ pyrefly.toml | 1 - .../unit/utils/test_venv_prefetch_manifest.py | 24 ----- 10 files changed, 162 insertions(+), 201 deletions(-) delete mode 100644 docker/venv_prefetch_manifest.tsv delete mode 100644 nemo_rl/utils/venv_prefetch_manifest.py delete mode 100644 tests/unit/utils/test_venv_prefetch_manifest.py diff --git a/docker/Dockerfile b/docker/Dockerfile index 7f4ffa5ef94..f0e8a8626d5 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -195,8 +195,6 @@ ENV LD_LIBRARY_PATH="/opt/nemo_rl_venv/lib/python3.13/site-packages/z3/lib:/opt/ # First copy only the dependency files COPY --from=nemo-rl pyproject.toml uv.lock ./ -# Prefetch worker-venv manifest to avoid the need for symlinks. -COPY --from=nemo-rl docker/venv_prefetch_manifest.tsv ./docker/venv_prefetch_manifest.tsv # Copy in the top level __init__.py/package_info.py since build-custom-vllm.sh needs the nemo_rl package to exist. COPY --from=nemo-rl nemo_rl/__init__.py nemo_rl/package_info.py ./nemo_rl/ COPY --from=nemo-rl tools/build-custom-vllm.sh ./tools/build-custom-vllm.sh @@ -250,8 +248,22 @@ UV_LINK_MODE=hardlink uv sync --frozen --extra modelopt --no-install-project UV_LINK_MODE=hardlink uv sync --frozen --all-groups --no-install-project # Worker-venv prefetch, phase 1: -# Materialize each registry venv's third-party packages in THIS layer, hardlinked against the above cache, +# Materialize each actor venv's third-party packages in THIS layer, hardlinked against the above cache, # so the N worker venvs cost directory entries instead of N wheel copies +# Emit "\t\t" for every uv-managed actor in +# pyproject.toml's [tool.nemo_rl.actor_environments]; the runtime registry reads the same table. +list_actor_venvs() { + "${UV_PROJECT_ENVIRONMENT}/bin/python" - <<'PY' +import tomllib +with open("pyproject.toml", "rb") as f: + table = tomllib.load(f)["tool"]["nemo_rl"]["actor_environments"] +for actor_fqn, extras in sorted(table.items()): + if extras == "system": + continue + stage = "trtllm" if "trtllm" in extras else "deps" + print(actor_fqn, stage, " ".join(f"--extra {e}" for e in extras), sep="\t") +PY +} PREFETCH_NEGATIVE_FILTERS="" if [[ -n "${SKIP_VLLM_BUILD:-}" ]]; then PREFETCH_NEGATIVE_FILTERS="$PREFETCH_NEGATIVE_FILTERS vllm" @@ -263,9 +275,6 @@ if [[ -n "${SKIP_TRTLLM_BUILD:-}" ]]; then PREFETCH_NEGATIVE_FILTERS="$PREFETCH_NEGATIVE_FILTERS trtllm" fi while IFS=$'\t' read -r venv_name stage extras; do - if [[ -z "$venv_name" || "$venv_name" == \#* ]]; then - continue - fi skip="" for f in $PREFETCH_NEGATIVE_FILTERS; do if [[ "$venv_name" == *"$f"* ]]; then @@ -281,7 +290,7 @@ while IFS=$'\t' read -r venv_name stage extras; do venv_path="${NEMO_RL_VENV_DIR}/${venv_name}" uv venv --allow-existing "$venv_path" UV_PROJECT_ENVIRONMENT="$venv_path" UV_LINK_MODE=hardlink uv sync --frozen $extras --no-install-project -done < docker/venv_prefetch_manifest.tsv +done < <(list_actor_venvs) # Remove the aiohttp in this uv cache dir to fully address CVE GHSA-mqqc-3gqh-h2x8 # The ray install will include the older aiohttp version in its cache @@ -359,13 +368,27 @@ rm -f "$TRTLLM_SYNC_LOG" echo "TRT-LLM cache state after build:" du -sh /root/.cache/uv /root/.cache/trtllm-wheels -# Worker-venv prefetch, phase 1 for the "trtllm"-stage manifest entries: +# Worker-venv prefetch, phase 1 for the "trtllm"-stage actors: # the dependency layer gave them a base-only warm; now that the tensorrt_llm wheel exists, # top off the venv with the full extra set so that the added files hardlink in-layer. # tensorrt_llm itself is new in this layer; trtllm-extra deps already present in the # dependency layer's cache (e.g. tilelang) get copied up here, a small accepted cost. +# Emit "\t\t" for every uv-managed actor in +# pyproject.toml's [tool.nemo_rl.actor_environments]; the runtime registry reads the same table. +list_actor_venvs() { + "${UV_PROJECT_ENVIRONMENT}/bin/python" - <<'PY' +import tomllib +with open("pyproject.toml", "rb") as f: + table = tomllib.load(f)["tool"]["nemo_rl"]["actor_environments"] +for actor_fqn, extras in sorted(table.items()): + if extras == "system": + continue + stage = "trtllm" if "trtllm" in extras else "deps" + print(actor_fqn, stage, " ".join(f"--extra {e}" for e in extras), sep="\t") +PY +} while IFS=$'\t' read -r venv_name stage extras; do - if [[ -z "$venv_name" || "$venv_name" == \#* || "$stage" != "trtllm" ]]; then + if [[ "$stage" != "trtllm" ]]; then continue fi venv_path="${NEMO_RL_VENV_DIR}/${venv_name}" @@ -373,7 +396,7 @@ while IFS=$'\t' read -r venv_name stage extras; do UV_PROJECT_ENVIRONMENT="$venv_path" UV_LINK_MODE=hardlink \ TRTLLM_WHEEL_CACHE_DIR=/root/.cache/trtllm-wheels \ uv sync --frozen $extras --no-install-project -done < docker/venv_prefetch_manifest.tsv +done < <(list_actor_venvs) # The main venv was never touched (the build ran in the throwaway venv), so # the previous "restore the default environment" sync is no longer needed. diff --git a/docker/venv_prefetch_manifest.tsv b/docker/venv_prefetch_manifest.tsv deleted file mode 100644 index e6d901a80fa..00000000000 --- a/docker/venv_prefetch_manifest.tsv +++ /dev/null @@ -1,22 +0,0 @@ -# DO NOT EDIT BY HAND -# Generated by `uv run python -m nemo_rl.utils.venv_prefetch_manifest`. -# tests/unit/utils/test_venv_prefetch_manifest.py ensures this file is not stale. -# Columns: -nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector deps --extra vllm -nemo_rl.algorithms.async_utils.ReplayBuffer deps --extra vllm -nemo_rl.environments.nemo_gym.NemoGym deps --extra nemo_gym -nemo_rl.experience.sync_rollout_actor.SyncRolloutActor deps --extra vllm -nemo_rl.modelopt.models.generation.vllm_quant_worker.VllmQuantAsyncGenerationWorker deps --extra modelopt --extra vllm -nemo_rl.modelopt.models.generation.vllm_quant_worker.VllmQuantGenerationWorker deps --extra modelopt --extra vllm -nemo_rl.modelopt.models.policy.workers.dtensor_quant_policy_worker.DTensorQuantPolicyWorker deps --extra modelopt --extra automodel -nemo_rl.modelopt.models.policy.workers.dtensor_quant_policy_worker_v2.DTensorQuantPolicyWorkerV2 deps --extra modelopt --extra automodel -nemo_rl.modelopt.models.policy.workers.megatron_quant_policy_worker.MegatronQuantPolicyWorker deps --extra modelopt --extra mcore -nemo_rl.models.generation.sglang.sglang_worker.SGLangGenerationWorker deps --extra sglang -nemo_rl.models.generation.trtllm.trtllm_worker_async.TrtllmAsyncGenerationWorker trtllm --extra trtllm -nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker deps --extra vllm -nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker deps --extra vllm -nemo_rl.models.policy.workers.dtensor_policy_worker.DTensorPolicyWorker deps --extra fsdp -nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2 deps --extra automodel -nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker deps --extra mcore -nemo_rl.models.value.workers.dtensor_value_worker_v2.DTensorValueWorkerV2 deps --extra automodel -nemo_rl.models.value.workers.megatron_value_worker.MegatronValueWorker deps --extra mcore diff --git a/docs/design-docs/dependency-management.md b/docs/design-docs/dependency-management.md index 9cf3d9c3be5..bf4bcd304dd 100644 --- a/docs/design-docs/dependency-management.md +++ b/docs/design-docs/dependency-management.md @@ -94,15 +94,14 @@ Within the driver script, NeMo RL starts multiple [`RayWorkerGroup`](https://git - **Generation workers** (e.g., vLLM): Require `vllm` dependencies - **Environment workers** (e.g., math evaluation): Use system/base dependencies -Each worker type is mapped to a specific Python executable configuration in the [`ACTOR_ENVIRONMENT_REGISTRY`](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/distributed/ray_actor_environment_registry.py#L17-L55). This registry defines which virtual environment should be used for each actor type: - -```python -ACTOR_ENVIRONMENT_REGISTRY: dict[str, str] = { - "nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker": PY_EXECUTABLES.VLLM, - "nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker": PY_EXECUTABLES.MCORE, - "nemo_rl.environments.math_environment.MathEnvironment": PY_EXECUTABLES.SYSTEM, - # ... more mappings -} +Each worker type is mapped to the uv extras its virtual environment needs in the `[tool.nemo_rl.actor_environments]` table of `pyproject.toml`. [`ACTOR_ENVIRONMENT_REGISTRY`](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/distributed/ray_actor_environment_registry.py) is built from that table at import time, and the release container reads the same table to prefetch one virtual environment per worker type: + +```toml +[tool.nemo_rl.actor_environments] +"nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker" = ["vllm"] +"nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker" = ["mcore"] +"nemo_rl.environments.math_environment.MathEnvironment" = "system" +# ... more mappings ``` > [!NOTE] diff --git a/docs/design-docs/uv.md b/docs/design-docs/uv.md index 62b2bcb6973..17e15090646 100644 --- a/docs/design-docs/uv.md +++ b/docs/design-docs/uv.md @@ -40,7 +40,7 @@ This section outlines how workers define their required executables, details the ### Worker Configuration -In our codebase, workers (classes decorated with `@ray.remote`, e.g., `PolicyWorker`) are associated with a `PY_EXECUTABLE` which specifies what dependencies the worker needs. These are set in a global registry in [`ACTOR_ENVIRONMENT_REGISTRY`](../../nemo_rl/distributed/ray_actor_environment_registry.py). This allows different parts of our application to have their own tailored environments. +In our codebase, workers (classes decorated with `@ray.remote`, e.g., `PolicyWorker`) are associated with a `PY_EXECUTABLE` which specifies what dependencies the worker needs. These are declared in the `[tool.nemo_rl.actor_environments]` table of `pyproject.toml`, from which the global registry [`ACTOR_ENVIRONMENT_REGISTRY`](../../nemo_rl/distributed/ray_actor_environment_registry.py) is built. This allows different parts of our application to have their own tailored environments. ### Supported Python Executables @@ -72,7 +72,7 @@ When a NeMo RL job is started: 1. The driver script creates several {py:class}`RayWorkerGroup `s. 2. Each worker group will create their workers which are wrapped in a {py:class}`RayWorkerBuilder ` where the fully qualified name (FQN) of the worker class is passed as a string. 3. {py:class}`RayWorkerBuilder ` launches the worker under {py:class}`RayWorkerBuilder ` which allows us to initialize the class without importing packages not available in the base environment. -4. Before the worker class is instantiated by the `RayWorkerBuilder`, the FQN is used to lookup -- in a [global registry](../../nemo_rl/distributed/ray_actor_environment_registry.py))) -- to determine which member of `PY_EXECUTABLES` should be used to launch that set of workers. If the chosen `PY_EXECUTABLES.*` starts with `uv`; a `venv` is created with all the dependencies it needs and the `runtime_env["py_executable"]` is replaced with the `venv`'s python interpreter. +4. Before the worker class is instantiated by the `RayWorkerBuilder`, the FQN is used to lookup -- in a [global registry](../../nemo_rl/distributed/ray_actor_environment_registry.py) built from `pyproject.toml`'s `[tool.nemo_rl.actor_environments]` -- to determine which member of `PY_EXECUTABLES` should be used to launch that set of workers. If the chosen `PY_EXECUTABLES.*` starts with `uv`; a `venv` is created with all the dependencies it needs and the `runtime_env["py_executable"]` is replaced with the `venv`'s python interpreter. This approach allows a fast start-up and maintains dependency isolation. It also has the added benefit of having all the virtual environments local under `./venvs`. diff --git a/nemo_rl/distributed/ray_actor_environment_registry.py b/nemo_rl/distributed/ray_actor_environment_registry.py index dbad277abba..b02a11b3e3f 100644 --- a/nemo_rl/distributed/ray_actor_environment_registry.py +++ b/nemo_rl/distributed/ray_actor_environment_registry.py @@ -12,47 +12,44 @@ # See the License for the specific language governing permissions and # limitations under the License. -from nemo_rl.distributed.virtual_cluster import PY_EXECUTABLES +import os +import tomllib +from pathlib import Path -ACTOR_ENVIRONMENT_REGISTRY: dict[str, str] = { - "nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker": PY_EXECUTABLES.VLLM_GYM, - "nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker": PY_EXECUTABLES.VLLM_GYM, - "nemo_rl.models.generation.sglang.sglang_worker.SGLangGenerationWorker": PY_EXECUTABLES.SGLANG, - "nemo_rl.models.generation.trtllm.trtllm_worker_async.TrtllmAsyncGenerationWorker": PY_EXECUTABLES.TRTLLM, - "nemo_rl.models.generation.dynamo.dynamo_worker.DynamoVllmWorker": PY_EXECUTABLES.SYSTEM, - "nemo_rl.models.policy.workers.dtensor_policy_worker.DTensorPolicyWorker": PY_EXECUTABLES.FSDP, - "nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2": PY_EXECUTABLES.AUTOMODEL, - "nemo_rl.models.value.workers.dtensor_value_worker_v2.DTensorValueWorkerV2": PY_EXECUTABLES.AUTOMODEL, - "nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker": PY_EXECUTABLES.MCORE, - "nemo_rl.models.value.workers.megatron_value_worker.MegatronValueWorker": PY_EXECUTABLES.MCORE, - "nemo_rl.data.energon.sft_worker.SFTMegatronPolicyWorker": PY_EXECUTABLES.MCORE, - "nemo_rl.environments.math_environment.MathEnvironment": PY_EXECUTABLES.SYSTEM, - "nemo_rl.environments.math_environment.MathMultiRewardEnvironment": PY_EXECUTABLES.SYSTEM, - "nemo_rl.environments.vlm_environment.VLMEnvironment": PY_EXECUTABLES.SYSTEM, - "nemo_rl.environments.code_environment.CodeEnvironment": PY_EXECUTABLES.SYSTEM, - "nemo_rl.environments.reward_model_environment.RewardModelEnvironment": PY_EXECUTABLES.SYSTEM, - "nemo_rl.environments.code_jaccard_environment.CodeJaccardEnvironment": PY_EXECUTABLES.SYSTEM, - "nemo_rl.environments.games.sliding_puzzle.SlidingPuzzleEnv": PY_EXECUTABLES.SYSTEM, - # AsyncTrajectoryCollector needs vLLM environment to handle exceptions from VllmGenerationWorker - "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector": PY_EXECUTABLES.VLLM, - # ReplayBuffer needs vLLM environment to handle trajectory data from VllmGenerationWorker - "nemo_rl.algorithms.async_utils.ReplayBuffer": PY_EXECUTABLES.VLLM, - # SyncRolloutActor doesn't import vllm directly — policy_generation is a - # Ray actor handle. The VLLM env is needed because (1) transfer_queue is - # bundled into the VLLM venv (and the policy training venvs), and the - # actor writes flattened tensors to TQ via dp_client.put_samples; - # (2) same-node colocation with VllmGenerationWorker avoids duplicate - # venv caches. - "nemo_rl.experience.sync_rollout_actor.SyncRolloutActor": PY_EXECUTABLES.VLLM, - "nemo_rl.environments.tools.retriever.RAGEnvironment": PY_EXECUTABLES.SYSTEM, - "nemo_rl.environments.nemo_gym.NemoGym": PY_EXECUTABLES.NEMO_GYM, - # ModelOpt actors need the modelopt extra on top of their backend extra. - "nemo_rl.modelopt.models.generation.vllm_quant_worker.VllmQuantGenerationWorker": PY_EXECUTABLES.MODELOPT_VLLM, - "nemo_rl.modelopt.models.generation.vllm_quant_worker.VllmQuantAsyncGenerationWorker": PY_EXECUTABLES.MODELOPT_VLLM, - "nemo_rl.modelopt.models.policy.workers.dtensor_quant_policy_worker.DTensorQuantPolicyWorker": PY_EXECUTABLES.MODELOPT_AUTOMODEL, - "nemo_rl.modelopt.models.policy.workers.dtensor_quant_policy_worker_v2.DTensorQuantPolicyWorkerV2": PY_EXECUTABLES.MODELOPT_AUTOMODEL, - "nemo_rl.modelopt.models.policy.workers.megatron_quant_policy_worker.MegatronQuantPolicyWorker": PY_EXECUTABLES.MODELOPT_MCORE, -} +from nemo_rl.distributed.virtual_cluster import ( + PY_EXECUTABLES, + git_root, + uv_py_executable, +) + +# NEMO_RL_PY_EXECUTABLES_SYSTEM=1 (single-environment images such as Dockerfile.ngc_pytorch) +# runs every actor on the driver's interpreter instead of a per-actor uv venv. +USE_SYSTEM_EXECUTABLE = os.environ.get("NEMO_RL_PY_EXECUTABLES_SYSTEM", "0") == "1" + + +def _load_actor_environments() -> dict[str, str]: + """Build actor FQN -> py_executable from pyproject.toml's [tool.nemo_rl.actor_environments].""" + with open(Path(git_root) / "pyproject.toml", "rb") as f: + pyproject = tomllib.load(f) + declared_extras = set(pyproject["project"]["optional-dependencies"]) + registry: dict[str, str] = {} + for actor_fqn, extras in pyproject["tool"]["nemo_rl"]["actor_environments"].items(): + if extras == "system": + registry[actor_fqn] = PY_EXECUTABLES.SYSTEM + continue + unknown = set(extras) - declared_extras + if unknown: + raise ValueError( + f"[tool.nemo_rl.actor_environments] {actor_fqn!r} names extras " + f"{sorted(unknown)} that are not in [project.optional-dependencies]" + ) + registry[actor_fqn] = ( + PY_EXECUTABLES.SYSTEM if USE_SYSTEM_EXECUTABLE else uv_py_executable(extras) + ) + return registry + + +ACTOR_ENVIRONMENT_REGISTRY: dict[str, str] = _load_actor_environments() def get_actor_python_env(actor_class_fqn: str) -> str: @@ -63,8 +60,7 @@ def get_actor_python_env(actor_class_fqn: str) -> str: f"No actor environment registered for {actor_class_fqn}. " f"You're attempting to create an actor ({actor_class_fqn}) " "without specifying a python environment for it. Please either" - "specify a python environment in the registry " - "(nemo_rl.distributed.ray_actor_environment_registry.ACTOR_ENVIRONMENT_REGISTRY) " + "add the actor to the [tool.nemo_rl.actor_environments] table in pyproject.toml " "or pass a py_executable to the RayWorkerBuilder. If you're unsure about which " "environment to use, a good default is PY_EXECUTABLES.SYSTEM for ray actors that " "don't have special dependencies. If you do have special dependencies (say, you're " diff --git a/nemo_rl/distributed/virtual_cluster.py b/nemo_rl/distributed/virtual_cluster.py index 75f9ea1051f..b910b7feec1 100644 --- a/nemo_rl/distributed/virtual_cluster.py +++ b/nemo_rl/distributed/virtual_cluster.py @@ -16,7 +16,7 @@ import socket import sys import time -from typing import NamedTuple, NotRequired, Optional, TypedDict +from typing import NamedTuple, NotRequired, Optional, Sequence, TypedDict import ray from ray.util.placement_group import ( @@ -118,6 +118,12 @@ def _resolve_system_overrides(cls) -> None: PY_EXECUTABLES._resolve_system_overrides() +def uv_py_executable(extras: Sequence[str]) -> str: + """py_executable of a uv-managed venv with the given extras (same shape as PY_EXECUTABLES.*).""" + extra_flags = "".join(f"--extra {extra} " for extra in extras) + return f"uv run --locked {extra_flags}--directory {git_root}" + + # Default port ranges — kept below the OS ephemeral range. On some DGX/GB200 # nodes the ephemeral floor is as low as 9000 (32768 on stock Linux), so every # service port is pinned below 9000 to avoid TOCTOU collisions. See ray.sub for diff --git a/nemo_rl/utils/venv_prefetch_manifest.py b/nemo_rl/utils/venv_prefetch_manifest.py deleted file mode 100644 index 92b653d3955..00000000000 --- a/nemo_rl/utils/venv_prefetch_manifest.py +++ /dev/null @@ -1,90 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import os -import shlex -from pathlib import Path - -MANIFEST_PATH = Path(__file__).parents[2] / "docker" / "venv_prefetch_manifest.tsv" - -_HEADER = """\ -# DO NOT EDIT BY HAND -# Generated by `uv run python -m nemo_rl.utils.venv_prefetch_manifest`. -# tests/unit/utils/test_venv_prefetch_manifest.py ensures this file is not stale. -# Columns: \t\t -""" - - -def parse_sync_extras(py_executable: str) -> list[str]: - """Extract the ``--extra `` pairs from a registry py_executable.""" - tokens = shlex.split(py_executable) - if tokens[:2] != ["uv", "run"]: - raise ValueError(f"not a uv py_executable: {py_executable!r}") - extras: list[str] = [] - i = 2 - while i < len(tokens): - token = tokens[i] - if token == "--extra": - extras.extend(tokens[i : i + 2]) - i += 2 - elif token == "--directory": - i += 2 - elif token == "--locked": - i += 1 - else: - raise ValueError( - f"unhandled py_executable flag {token!r} in {py_executable!r}; " - "teach nemo_rl/utils/venv_prefetch_manifest.py how it maps to `uv sync` flags" - ) - return extras - - -def build_manifest_rows() -> list[tuple[str, str, str]]: - """Build (venv_name, stage, extras) rows for every uv-managed actor.""" - if os.environ.get("NEMO_RL_PY_EXECUTABLES_SYSTEM", "0") == "1": - raise RuntimeError( - "NEMO_RL_PY_EXECUTABLES_SYSTEM=1 collapses the registry to system " - "python; unset it to generate the manifest" - ) - # Local import — the registry transitively pulls heavy deps (ray, modelopt). - from nemo_rl.distributed.ray_actor_environment_registry import ( - ACTOR_ENVIRONMENT_REGISTRY, - ) - - rows: list[tuple[str, str, str]] = [] - for actor_fqn, py_executable in sorted(ACTOR_ENVIRONMENT_REGISTRY.items()): - if not py_executable.startswith("uv"): - continue - extras = parse_sync_extras(py_executable) - stage = "trtllm" if "trtllm" in extras else "deps" - rows.append((actor_fqn, stage, " ".join(extras))) - return rows - - -def render_manifest() -> str: - """Render the manifest file content.""" - lines = [_HEADER] - for actor_fqn, stage, extras in build_manifest_rows(): - lines.append(f"{actor_fqn}\t{stage}\t{extras}\n") - return "".join(lines) - - -def main() -> None: - """Regenerate docker/venv_prefetch_manifest.tsv in place.""" - MANIFEST_PATH.write_text(render_manifest()) - print(f"wrote {MANIFEST_PATH}") - - -if __name__ == "__main__": - main() diff --git a/pyproject.toml b/pyproject.toml index 72d335a838d..c99dcd56dfc 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -655,6 +655,80 @@ name = "drain3" version = "0.9.11" requires-dist = ["jsonpickle", "cachetools>=4.2.1"] +# Ray actor -> uv extras of the venv it runs in; "system" = the driver's interpreter. +# nemo_rl.distributed.ray_actor_environment_registry builds ACTOR_ENVIRONMENT_REGISTRY from +# this table, and docker/Dockerfile reads it to prefetch one venv per actor into the image. +# Extras must exist in [project.optional-dependencies]. +[tool.nemo_rl.actor_environments] +"nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker" = ["vllm"] +"nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker" = [ + "vllm", +] +"nemo_rl.models.generation.sglang.sglang_worker.SGLangGenerationWorker" = [ + "sglang", +] +"nemo_rl.models.generation.dynamo.dynamo_worker.DynamoVllmWorker" = "system" +"nemo_rl.models.policy.workers.dtensor_policy_worker.DTensorPolicyWorker" = [ + "fsdp", +] +"nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2" = [ + "automodel", +] +"nemo_rl.models.value.workers.dtensor_value_worker_v2.DTensorValueWorkerV2" = [ + "automodel", +] +"nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker" = [ + "mcore", +] +"nemo_rl.models.value.workers.megatron_value_worker.MegatronValueWorker" = [ + "mcore", +] +"nemo_rl.data.energon.sft_worker.SFTMegatronPolicyWorker" = [ + "mcore", +] +"nemo_rl.models.generation.trtllm.trtllm_worker_async.TrtllmAsyncGenerationWorker" = [ + "trtllm", +] +"nemo_rl.environments.math_environment.MathEnvironment" = "system" +"nemo_rl.environments.math_environment.MathMultiRewardEnvironment" = "system" +"nemo_rl.environments.vlm_environment.VLMEnvironment" = "system" +"nemo_rl.environments.code_environment.CodeEnvironment" = "system" +"nemo_rl.environments.reward_model_environment.RewardModelEnvironment" = "system" +"nemo_rl.environments.code_jaccard_environment.CodeJaccardEnvironment" = "system" +"nemo_rl.environments.games.sliding_puzzle.SlidingPuzzleEnv" = "system" +# AsyncTrajectoryCollector needs the vLLM environment to handle exceptions from VllmGenerationWorker +"nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector" = ["vllm"] +# ReplayBuffer needs the vLLM environment to handle trajectory data from VllmGenerationWorker +"nemo_rl.algorithms.async_utils.ReplayBuffer" = ["vllm"] +# SyncRolloutActor doesn't import vllm directly -- policy_generation is a Ray actor handle. +# The vLLM env is needed because (1) transfer_queue is bundled into the vLLM venv (and the +# policy training venvs), and the actor writes flattened tensors to TQ via dp_client.put_samples; +# (2) same-node colocation with VllmGenerationWorker avoids duplicate venv caches. +"nemo_rl.experience.sync_rollout_actor.SyncRolloutActor" = ["vllm"] +"nemo_rl.environments.tools.retriever.RAGEnvironment" = "system" +"nemo_rl.environments.nemo_gym.NemoGym" = ["nemo_gym"] +# ModelOpt quantization-aware workers +"nemo_rl.modelopt.models.generation.vllm_quant_worker.VllmQuantGenerationWorker" = [ + "modelopt", + "vllm", +] +"nemo_rl.modelopt.models.generation.vllm_quant_worker.VllmQuantAsyncGenerationWorker" = [ + "modelopt", + "vllm", +] +"nemo_rl.modelopt.models.policy.workers.dtensor_quant_policy_worker.DTensorQuantPolicyWorker" = [ + "modelopt", + "automodel", +] +"nemo_rl.modelopt.models.policy.workers.dtensor_quant_policy_worker_v2.DTensorQuantPolicyWorkerV2" = [ + "modelopt", + "automodel", +] +"nemo_rl.modelopt.models.policy.workers.megatron_quant_policy_worker.MegatronQuantPolicyWorker" = [ + "modelopt", + "mcore", +] + [tool.black] line-length = 120 include = '\.pyi?$' diff --git a/pyrefly.toml b/pyrefly.toml index 5832f9f61c9..1b54fdaab50 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -286,7 +286,6 @@ project-includes = [ "nemo_rl/utils/r3_trace.py", "nemo_rl/utils/routed_experts_codec.py", "nemo_rl/utils/timer.py", - "nemo_rl/utils/venv_prefetch_manifest.py", "nemo_rl/utils/venvs.py", "nemo_rl/utils/weight_transfer_http.py", "nemo_rl/utils/weight_transfer_sparse_codec.py", diff --git a/tests/unit/utils/test_venv_prefetch_manifest.py b/tests/unit/utils/test_venv_prefetch_manifest.py deleted file mode 100644 index e0524475b8f..00000000000 --- a/tests/unit/utils/test_venv_prefetch_manifest.py +++ /dev/null @@ -1,24 +0,0 @@ -# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -from nemo_rl.utils.venv_prefetch_manifest import MANIFEST_PATH, render_manifest - - -def test_manifest_matches_actor_registry(): - """docker/venv_prefetch_manifest.tsv must stay in lockstep with the registry.""" - assert MANIFEST_PATH.read_text() == render_manifest(), ( - "docker/venv_prefetch_manifest.tsv is stale relative to " - "ACTOR_ENVIRONMENT_REGISTRY; regenerate it with " - "`uv run python -m nemo_rl.utils.venv_prefetch_manifest`" - ) From 9d8f52bd65e88c8683dabab5e19619730a891c26 Mon Sep 17 00:00:00 2001 From: Terry Kong Date: Fri, 4 Sep 2026 14:49:14 -0700 Subject: [PATCH 04/21] refactor(docker): read actor venv list from a Python leaf module Keeps tdene's three-layer hardlink prefetch exactly as-is, but moves the actor -> uv extras mapping out of pyproject.toml and back into Python, so there is one list with one owner. - nemo_rl/distributed/actor_environments.py holds ACTOR_ENVIRONMENTS. It is stdlib-only and dependency-free, so docker/Dockerfile can run it as a script from the dependency layer, where the source tree does not exist yet. Running it as a script (not importing) also keeps nemo_rl/__init__.py from executing there. - ray_actor_environment_registry.py imports that dict instead of parsing pyproject.toml at import time. No file I/O, no tomllib. - The Dockerfile's two byte-identical tomllib heredocs are gone. The list is written once to /opt/actor_venvs.tsv with a plain redirect, so set -e catches a parse failure instead of silently prefetching nothing. - SKIP_*_BUILD now filters on the declared extras rather than a substring of the actor name, so SKIP_VLLM_BUILD also skips AsyncTrajectoryCollector, ReplayBuffer and SyncRolloutActor. - Phase 2 sets TRTLLM_REQUIRE_CACHED_WHEEL=1, matching the release stage, so a cache miss fails fast instead of starting a source build. - Adds tests: extras are declared, actor modules exist, generated py_executables match PY_EXECUTABLES, script output matches the registry, and actor_environments.py imports only stdlib. Signed-off-by: Terry Kong --- docker/Dockerfile | 67 +++----- docs/design-docs/dependency-management.md | 21 ++- docs/design-docs/uv.md | 4 +- nemo_rl/distributed/actor_environments.py | 142 ++++++++++++++++ .../ray_actor_environment_registry.py | 48 ++---- pyproject.toml | 74 -------- pyrefly.toml | 1 + .../distributed/test_actor_environments.py | 160 ++++++++++++++++++ 8 files changed, 356 insertions(+), 161 deletions(-) create mode 100644 nemo_rl/distributed/actor_environments.py create mode 100644 tests/unit/distributed/test_actor_environments.py diff --git a/docker/Dockerfile b/docker/Dockerfile index f0e8a8626d5..7051ac2431e 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -197,6 +197,10 @@ ENV LD_LIBRARY_PATH="/opt/nemo_rl_venv/lib/python3.13/site-packages/z3/lib:/opt/ COPY --from=nemo-rl pyproject.toml uv.lock ./ # Copy in the top level __init__.py/package_info.py since build-custom-vllm.sh needs the nemo_rl package to exist. COPY --from=nemo-rl nemo_rl/__init__.py nemo_rl/package_info.py ./nemo_rl/ +# The single source of truth for which extras each Ray actor's venv needs. Run as a +# script by the prefetch below; imported by nemo_rl.distributed.ray_actor_environment_registry +# at runtime. Kept dependency-free so it can be copied in without the rest of the source. +COPY --from=nemo-rl nemo_rl/distributed/actor_environments.py ./nemo_rl/distributed/ COPY --from=nemo-rl tools/build-custom-vllm.sh ./tools/build-custom-vllm.sh COPY --from=nemo-rl tools/build-custom-flashinfer.sh ./tools/build-custom-flashinfer.sh COPY --from=nemo-rl --link research/ ./research/ @@ -249,48 +253,34 @@ UV_LINK_MODE=hardlink uv sync --frozen --all-groups --no-install-project # Worker-venv prefetch, phase 1: # Materialize each actor venv's third-party packages in THIS layer, hardlinked against the above cache, -# so the N worker venvs cost directory entries instead of N wheel copies -# Emit "\t\t" for every uv-managed actor in -# pyproject.toml's [tool.nemo_rl.actor_environments]; the runtime registry reads the same table. -list_actor_venvs() { - "${UV_PROJECT_ENVIRONMENT}/bin/python" - <<'PY' -import tomllib -with open("pyproject.toml", "rb") as f: - table = tomllib.load(f)["tool"]["nemo_rl"]["actor_environments"] -for actor_fqn, extras in sorted(table.items()): - if extras == "system": - continue - stage = "trtllm" if "trtllm" in extras else "deps" - print(actor_fqn, stage, " ".join(f"--extra {e}" for e in extras), sep="\t") -PY -} -PREFETCH_NEGATIVE_FILTERS="" +# so the N worker venvs cost directory entries instead of N wheel copies. +# nemo_rl/distributed/actor_environments.py is the single source of truth for which extras each +# actor needs; the runtime registry imports the same dict. It is run as a script (not imported) +# because nemo_rl/__init__.py cannot execute in this layer -- the source tree is not here yet. +SKIP_EXTRAS="" if [[ -n "${SKIP_VLLM_BUILD:-}" ]]; then - PREFETCH_NEGATIVE_FILTERS="$PREFETCH_NEGATIVE_FILTERS vllm" + SKIP_EXTRAS="$SKIP_EXTRAS vllm" fi if [[ -n "${SKIP_SGLANG_BUILD:-}" ]]; then - PREFETCH_NEGATIVE_FILTERS="$PREFETCH_NEGATIVE_FILTERS sglang" + SKIP_EXTRAS="$SKIP_EXTRAS sglang" fi if [[ -n "${SKIP_TRTLLM_BUILD:-}" ]]; then - PREFETCH_NEGATIVE_FILTERS="$PREFETCH_NEGATIVE_FILTERS trtllm" + SKIP_EXTRAS="$SKIP_EXTRAS trtllm" fi +# Write the list once, in a plain redirect so `set -e` catches a failure here +# instead of silently prefetching nothing. +"${UV_PROJECT_ENVIRONMENT}/bin/python" nemo_rl/distributed/actor_environments.py all $SKIP_EXTRAS \ + > /opt/actor_venvs.tsv +test -s /opt/actor_venvs.tsv while IFS=$'\t' read -r venv_name stage extras; do - skip="" - for f in $PREFETCH_NEGATIVE_FILTERS; do - if [[ "$venv_name" == *"$f"* ]]; then - skip=1 - fi - done - if [[ -n "$skip" ]]; then - continue - fi if [[ "$stage" == "trtllm" ]]; then + # tensorrt_llm does not exist yet; warm the base and finish in the TRT-LLM layer. extras="" fi venv_path="${NEMO_RL_VENV_DIR}/${venv_name}" uv venv --allow-existing "$venv_path" UV_PROJECT_ENVIRONMENT="$venv_path" UV_LINK_MODE=hardlink uv sync --frozen $extras --no-install-project -done < <(list_actor_venvs) +done < /opt/actor_venvs.tsv # Remove the aiohttp in this uv cache dir to fully address CVE GHSA-mqqc-3gqh-h2x8 # The ray install will include the older aiohttp version in its cache @@ -373,20 +363,8 @@ du -sh /root/.cache/uv /root/.cache/trtllm-wheels # top off the venv with the full extra set so that the added files hardlink in-layer. # tensorrt_llm itself is new in this layer; trtllm-extra deps already present in the # dependency layer's cache (e.g. tilelang) get copied up here, a small accepted cost. -# Emit "\t\t" for every uv-managed actor in -# pyproject.toml's [tool.nemo_rl.actor_environments]; the runtime registry reads the same table. -list_actor_venvs() { - "${UV_PROJECT_ENVIRONMENT}/bin/python" - <<'PY' -import tomllib -with open("pyproject.toml", "rb") as f: - table = tomllib.load(f)["tool"]["nemo_rl"]["actor_environments"] -for actor_fqn, extras in sorted(table.items()): - if extras == "system": - continue - stage = "trtllm" if "trtllm" in extras else "deps" - print(actor_fqn, stage, " ".join(f"--extra {e}" for e in extras), sep="\t") -PY -} +# The dependency layer already wrote the actor list to /opt/actor_venvs.tsv, so this layer +# reuses it rather than re-deriving it -- one reader, one place it can go wrong. while IFS=$'\t' read -r venv_name stage extras; do if [[ "$stage" != "trtllm" ]]; then continue @@ -395,8 +373,9 @@ while IFS=$'\t' read -r venv_name stage extras; do uv venv --allow-existing "$venv_path" UV_PROJECT_ENVIRONMENT="$venv_path" UV_LINK_MODE=hardlink \ TRTLLM_WHEEL_CACHE_DIR=/root/.cache/trtllm-wheels \ + TRTLLM_REQUIRE_CACHED_WHEEL=1 \ uv sync --frozen $extras --no-install-project -done < <(list_actor_venvs) +done < /opt/actor_venvs.tsv # The main venv was never touched (the build ran in the throwaway venv), so # the previous "restore the default environment" sync is no longer needed. diff --git a/docs/design-docs/dependency-management.md b/docs/design-docs/dependency-management.md index bf4bcd304dd..61fa04707f8 100644 --- a/docs/design-docs/dependency-management.md +++ b/docs/design-docs/dependency-management.md @@ -94,16 +94,21 @@ Within the driver script, NeMo RL starts multiple [`RayWorkerGroup`](https://git - **Generation workers** (e.g., vLLM): Require `vllm` dependencies - **Environment workers** (e.g., math evaluation): Use system/base dependencies -Each worker type is mapped to the uv extras its virtual environment needs in the `[tool.nemo_rl.actor_environments]` table of `pyproject.toml`. [`ACTOR_ENVIRONMENT_REGISTRY`](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/distributed/ray_actor_environment_registry.py) is built from that table at import time, and the release container reads the same table to prefetch one virtual environment per worker type: - -```toml -[tool.nemo_rl.actor_environments] -"nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker" = ["vllm"] -"nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker" = ["mcore"] -"nemo_rl.environments.math_environment.MathEnvironment" = "system" -# ... more mappings +Each worker type is mapped to the uv extras its virtual environment needs in `ACTOR_ENVIRONMENTS` in [`nemo_rl/distributed/actor_environments.py`](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/distributed/actor_environments.py). [`ACTOR_ENVIRONMENT_REGISTRY`](https://github.com/NVIDIA-NeMo/RL/blob/main/nemo_rl/distributed/ray_actor_environment_registry.py) is built from it at import time, and `docker/Dockerfile` runs the same module as a script to prefetch one virtual environment per worker type into the image: + +```python +# nemo_rl/distributed/actor_environments.py -- None means the driver's interpreter +ACTOR_ENVIRONMENTS: dict[str, list[str] | None] = { + "nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker": ["vllm"], + "nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker": ["mcore"], + "nemo_rl.environments.math_environment.MathEnvironment": None, + # ... more mappings +} ``` +This module is deliberately dependency-free: `docker/Dockerfile` runs it as a script +from the dependency layer, where the rest of the source tree does not exist yet. + > [!NOTE] > For more details on how workers define and use their Python executables, see the [UV Documentation](uv.md#worker-configuration). diff --git a/docs/design-docs/uv.md b/docs/design-docs/uv.md index 17e15090646..31baa3889ac 100644 --- a/docs/design-docs/uv.md +++ b/docs/design-docs/uv.md @@ -40,7 +40,7 @@ This section outlines how workers define their required executables, details the ### Worker Configuration -In our codebase, workers (classes decorated with `@ray.remote`, e.g., `PolicyWorker`) are associated with a `PY_EXECUTABLE` which specifies what dependencies the worker needs. These are declared in the `[tool.nemo_rl.actor_environments]` table of `pyproject.toml`, from which the global registry [`ACTOR_ENVIRONMENT_REGISTRY`](../../nemo_rl/distributed/ray_actor_environment_registry.py) is built. This allows different parts of our application to have their own tailored environments. +In our codebase, workers (classes decorated with `@ray.remote`, e.g., `PolicyWorker`) are associated with a `PY_EXECUTABLE` which specifies what dependencies the worker needs. These are declared in `ACTOR_ENVIRONMENTS` in [`nemo_rl/distributed/actor_environments.py`](../../nemo_rl/distributed/actor_environments.py), from which the global registry [`ACTOR_ENVIRONMENT_REGISTRY`](../../nemo_rl/distributed/ray_actor_environment_registry.py) is built. Workers defined outside this repo register themselves by assigning into `ACTOR_ENVIRONMENT_REGISTRY` at runtime -- see `research/template_project`. This allows different parts of our application to have their own tailored environments. ### Supported Python Executables @@ -72,7 +72,7 @@ When a NeMo RL job is started: 1. The driver script creates several {py:class}`RayWorkerGroup `s. 2. Each worker group will create their workers which are wrapped in a {py:class}`RayWorkerBuilder ` where the fully qualified name (FQN) of the worker class is passed as a string. 3. {py:class}`RayWorkerBuilder ` launches the worker under {py:class}`RayWorkerBuilder ` which allows us to initialize the class without importing packages not available in the base environment. -4. Before the worker class is instantiated by the `RayWorkerBuilder`, the FQN is used to lookup -- in a [global registry](../../nemo_rl/distributed/ray_actor_environment_registry.py) built from `pyproject.toml`'s `[tool.nemo_rl.actor_environments]` -- to determine which member of `PY_EXECUTABLES` should be used to launch that set of workers. If the chosen `PY_EXECUTABLES.*` starts with `uv`; a `venv` is created with all the dependencies it needs and the `runtime_env["py_executable"]` is replaced with the `venv`'s python interpreter. +4. Before the worker class is instantiated by the `RayWorkerBuilder`, the FQN is used to lookup -- in a [global registry](../../nemo_rl/distributed/ray_actor_environment_registry.py) built from [`ACTOR_ENVIRONMENTS`](../../nemo_rl/distributed/actor_environments.py) -- to determine which member of `PY_EXECUTABLES` should be used to launch that set of workers. If the chosen `PY_EXECUTABLES.*` starts with `uv`; a `venv` is created with all the dependencies it needs and the `runtime_env["py_executable"]` is replaced with the `venv`'s python interpreter. This approach allows a fast start-up and maintains dependency isolation. It also has the added benefit of having all the virtual environments local under `./venvs`. diff --git a/nemo_rl/distributed/actor_environments.py b/nemo_rl/distributed/actor_environments.py new file mode 100644 index 00000000000..210d23567ba --- /dev/null +++ b/nemo_rl/distributed/actor_environments.py @@ -0,0 +1,142 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Which uv extras each Ray actor needs. The single source of truth. + +Two readers: + +* ``nemo_rl.distributed.ray_actor_environment_registry`` imports + ``ACTOR_ENVIRONMENTS`` and turns it into ``py_executable`` strings at runtime. +* ``docker/Dockerfile`` runs this file **as a script** to list the venvs it must + pre-build, so the image ships one venv per actor. + +DO NOT IMPORT ANYTHING FROM ``nemo_rl`` HERE, AND KEEP IT STDLIB-ONLY. +The Dockerfile runs this from the dependency layer, where only this file and +``pyproject.toml``/``uv.lock`` exist -- the rest of the source tree has not been +copied in yet. Running it as a script (rather than importing it) is also what +keeps ``nemo_rl/__init__.py`` from executing there. An import added here breaks +the image build in its most expensive layer. +``tests/unit/distributed/test_actor_environments.py`` enforces this. +""" + +# Keeps the annotations below from being evaluated at runtime, so this module also +# runs under interpreters older than the one the image pins. +from __future__ import annotations + +# Actor fully-qualified name -> the uv extras its virtual environment needs. +# ``None`` means the actor runs on the driver's interpreter and gets no venv. +# Every extra must exist in ``[project.optional-dependencies]`` of pyproject.toml. +ACTOR_ENVIRONMENTS: dict[str, list[str] | None] = { + "nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker": ["vllm"], + "nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker": [ + "vllm" + ], + "nemo_rl.models.generation.sglang.sglang_worker.SGLangGenerationWorker": ["sglang"], + "nemo_rl.models.generation.dynamo.dynamo_worker.DynamoVllmWorker": None, + "nemo_rl.models.policy.workers.dtensor_policy_worker.DTensorPolicyWorker": ["fsdp"], + "nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2": [ + "automodel" + ], + "nemo_rl.models.value.workers.dtensor_value_worker_v2.DTensorValueWorkerV2": [ + "automodel" + ], + "nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker": [ + "mcore" + ], + "nemo_rl.models.value.workers.megatron_value_worker.MegatronValueWorker": ["mcore"], + "nemo_rl.models.generation.trtllm.trtllm_worker_async.TrtllmAsyncGenerationWorker": [ + "trtllm" + ], + "nemo_rl.environments.math_environment.MathEnvironment": None, + "nemo_rl.environments.math_environment.MathMultiRewardEnvironment": None, + "nemo_rl.environments.vlm_environment.VLMEnvironment": None, + "nemo_rl.environments.code_environment.CodeEnvironment": None, + "nemo_rl.environments.reward_model_environment.RewardModelEnvironment": None, + "nemo_rl.environments.code_jaccard_environment.CodeJaccardEnvironment": None, + "nemo_rl.environments.games.sliding_puzzle.SlidingPuzzleEnv": None, + # AsyncTrajectoryCollector needs the vLLM environment to handle exceptions + # from VllmGenerationWorker. + "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector": ["vllm"], + # ReplayBuffer needs the vLLM environment to handle trajectory data from + # VllmGenerationWorker. + "nemo_rl.algorithms.async_utils.ReplayBuffer": ["vllm"], + # SyncRolloutActor doesn't import vllm directly -- policy_generation is a Ray + # actor handle. The vLLM env is needed because (1) transfer_queue is bundled + # into the vLLM venv (and the policy training venvs), and the actor writes + # flattened tensors to TQ via dp_client.put_samples; (2) same-node colocation + # with VllmGenerationWorker avoids duplicate venv caches. + "nemo_rl.experience.sync_rollout_actor.SyncRolloutActor": ["vllm"], + "nemo_rl.environments.tools.retriever.RAGEnvironment": None, + "nemo_rl.environments.nemo_gym.NemoGym": ["nemo_gym"], + # ModelOpt quantization-aware workers + "nemo_rl.modelopt.models.generation.vllm_quant_worker.VllmQuantGenerationWorker": [ + "modelopt", + "vllm", + ], + "nemo_rl.modelopt.models.generation.vllm_quant_worker.VllmQuantAsyncGenerationWorker": [ + "modelopt", + "vllm", + ], + "nemo_rl.modelopt.models.policy.workers.dtensor_quant_policy_worker.DTensorQuantPolicyWorker": [ + "modelopt", + "automodel", + ], + "nemo_rl.modelopt.models.policy.workers.dtensor_quant_policy_worker_v2.DTensorQuantPolicyWorkerV2": [ + "modelopt", + "automodel", + ], + "nemo_rl.modelopt.models.policy.workers.megatron_quant_policy_worker.MegatronQuantPolicyWorker": [ + "modelopt", + "mcore", + ], +} + + +def _build_stage(extras: list[str]) -> str: + """Which image layer can finish this venv. + + The tensorrt_llm wheel is only built in the TRT-LLM layer, so those venvs get + their third-party packages in two steps; everything else finishes in the + dependency layer. + """ + return "trtllm" if "trtllm" in extras else "deps" + + +def main(argv: list[str]) -> int: + r"""Print "\t\t" for uv-managed actors. + + Usage: actor_environments.py [] [ ...] + + ```` is "deps" or "trtllm" (omit for all). Any extras listed after it + are skipped, which is how the Dockerfile honors SKIP_VLLM_BUILD and friends. + Filtering on the declared extras -- rather than on a substring of the actor + name -- is what makes SKIP_VLLM_BUILD also skip actors like + AsyncTrajectoryCollector, whose name contains no "vllm". + """ + stage = argv[1] if len(argv) > 1 and argv[1] != "all" else None + skip = set(argv[2:]) + for actor_fqn, extras in sorted(ACTOR_ENVIRONMENTS.items()): + if extras is None or skip & set(extras): + continue + actor_stage = _build_stage(extras) + if stage is not None and actor_stage != stage: + continue + flags = " ".join(f"--extra {extra}" for extra in extras) + print(actor_fqn, actor_stage, flags, sep="\t") + return 0 + + +if __name__ == "__main__": + import sys + + raise SystemExit(main(sys.argv)) diff --git a/nemo_rl/distributed/ray_actor_environment_registry.py b/nemo_rl/distributed/ray_actor_environment_registry.py index b02a11b3e3f..f78c126b58e 100644 --- a/nemo_rl/distributed/ray_actor_environment_registry.py +++ b/nemo_rl/distributed/ray_actor_environment_registry.py @@ -13,43 +13,23 @@ # limitations under the License. import os -import tomllib -from pathlib import Path -from nemo_rl.distributed.virtual_cluster import ( - PY_EXECUTABLES, - git_root, - uv_py_executable, -) +from nemo_rl.distributed.actor_environments import ACTOR_ENVIRONMENTS +from nemo_rl.distributed.virtual_cluster import PY_EXECUTABLES, uv_py_executable # NEMO_RL_PY_EXECUTABLES_SYSTEM=1 (single-environment images such as Dockerfile.ngc_pytorch) # runs every actor on the driver's interpreter instead of a per-actor uv venv. USE_SYSTEM_EXECUTABLE = os.environ.get("NEMO_RL_PY_EXECUTABLES_SYSTEM", "0") == "1" - -def _load_actor_environments() -> dict[str, str]: - """Build actor FQN -> py_executable from pyproject.toml's [tool.nemo_rl.actor_environments].""" - with open(Path(git_root) / "pyproject.toml", "rb") as f: - pyproject = tomllib.load(f) - declared_extras = set(pyproject["project"]["optional-dependencies"]) - registry: dict[str, str] = {} - for actor_fqn, extras in pyproject["tool"]["nemo_rl"]["actor_environments"].items(): - if extras == "system": - registry[actor_fqn] = PY_EXECUTABLES.SYSTEM - continue - unknown = set(extras) - declared_extras - if unknown: - raise ValueError( - f"[tool.nemo_rl.actor_environments] {actor_fqn!r} names extras " - f"{sorted(unknown)} that are not in [project.optional-dependencies]" - ) - registry[actor_fqn] = ( - PY_EXECUTABLES.SYSTEM if USE_SYSTEM_EXECUTABLE else uv_py_executable(extras) - ) - return registry - - -ACTOR_ENVIRONMENT_REGISTRY: dict[str, str] = _load_actor_environments() +# Actor FQN -> the py_executable its workers launch under. The extras come from +# nemo_rl.distributed.actor_environments, which docker/Dockerfile also reads to +# pre-build one venv per actor into the image. +ACTOR_ENVIRONMENT_REGISTRY: dict[str, str] = { + actor_fqn: PY_EXECUTABLES.SYSTEM + if extras is None or USE_SYSTEM_EXECUTABLE + else uv_py_executable(extras) + for actor_fqn, extras in ACTOR_ENVIRONMENTS.items() +} def get_actor_python_env(actor_class_fqn: str) -> str: @@ -59,8 +39,10 @@ def get_actor_python_env(actor_class_fqn: str) -> str: raise ValueError( f"No actor environment registered for {actor_class_fqn}. " f"You're attempting to create an actor ({actor_class_fqn}) " - "without specifying a python environment for it. Please either" - "add the actor to the [tool.nemo_rl.actor_environments] table in pyproject.toml " + "without specifying a python environment for it. Please either " + "add the actor to ACTOR_ENVIRONMENTS in nemo_rl/distributed/actor_environments.py, " + "register it at runtime with ACTOR_ENVIRONMENT_REGISTRY[fqn] = " + "(the path for workers defined outside this repo), " "or pass a py_executable to the RayWorkerBuilder. If you're unsure about which " "environment to use, a good default is PY_EXECUTABLES.SYSTEM for ray actors that " "don't have special dependencies. If you do have special dependencies (say, you're " diff --git a/pyproject.toml b/pyproject.toml index c99dcd56dfc..72d335a838d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -655,80 +655,6 @@ name = "drain3" version = "0.9.11" requires-dist = ["jsonpickle", "cachetools>=4.2.1"] -# Ray actor -> uv extras of the venv it runs in; "system" = the driver's interpreter. -# nemo_rl.distributed.ray_actor_environment_registry builds ACTOR_ENVIRONMENT_REGISTRY from -# this table, and docker/Dockerfile reads it to prefetch one venv per actor into the image. -# Extras must exist in [project.optional-dependencies]. -[tool.nemo_rl.actor_environments] -"nemo_rl.models.generation.vllm.vllm_worker.VllmGenerationWorker" = ["vllm"] -"nemo_rl.models.generation.vllm.vllm_worker_async.VllmAsyncGenerationWorker" = [ - "vllm", -] -"nemo_rl.models.generation.sglang.sglang_worker.SGLangGenerationWorker" = [ - "sglang", -] -"nemo_rl.models.generation.dynamo.dynamo_worker.DynamoVllmWorker" = "system" -"nemo_rl.models.policy.workers.dtensor_policy_worker.DTensorPolicyWorker" = [ - "fsdp", -] -"nemo_rl.models.policy.workers.dtensor_policy_worker_v2.DTensorPolicyWorkerV2" = [ - "automodel", -] -"nemo_rl.models.value.workers.dtensor_value_worker_v2.DTensorValueWorkerV2" = [ - "automodel", -] -"nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker" = [ - "mcore", -] -"nemo_rl.models.value.workers.megatron_value_worker.MegatronValueWorker" = [ - "mcore", -] -"nemo_rl.data.energon.sft_worker.SFTMegatronPolicyWorker" = [ - "mcore", -] -"nemo_rl.models.generation.trtllm.trtllm_worker_async.TrtllmAsyncGenerationWorker" = [ - "trtllm", -] -"nemo_rl.environments.math_environment.MathEnvironment" = "system" -"nemo_rl.environments.math_environment.MathMultiRewardEnvironment" = "system" -"nemo_rl.environments.vlm_environment.VLMEnvironment" = "system" -"nemo_rl.environments.code_environment.CodeEnvironment" = "system" -"nemo_rl.environments.reward_model_environment.RewardModelEnvironment" = "system" -"nemo_rl.environments.code_jaccard_environment.CodeJaccardEnvironment" = "system" -"nemo_rl.environments.games.sliding_puzzle.SlidingPuzzleEnv" = "system" -# AsyncTrajectoryCollector needs the vLLM environment to handle exceptions from VllmGenerationWorker -"nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector" = ["vllm"] -# ReplayBuffer needs the vLLM environment to handle trajectory data from VllmGenerationWorker -"nemo_rl.algorithms.async_utils.ReplayBuffer" = ["vllm"] -# SyncRolloutActor doesn't import vllm directly -- policy_generation is a Ray actor handle. -# The vLLM env is needed because (1) transfer_queue is bundled into the vLLM venv (and the -# policy training venvs), and the actor writes flattened tensors to TQ via dp_client.put_samples; -# (2) same-node colocation with VllmGenerationWorker avoids duplicate venv caches. -"nemo_rl.experience.sync_rollout_actor.SyncRolloutActor" = ["vllm"] -"nemo_rl.environments.tools.retriever.RAGEnvironment" = "system" -"nemo_rl.environments.nemo_gym.NemoGym" = ["nemo_gym"] -# ModelOpt quantization-aware workers -"nemo_rl.modelopt.models.generation.vllm_quant_worker.VllmQuantGenerationWorker" = [ - "modelopt", - "vllm", -] -"nemo_rl.modelopt.models.generation.vllm_quant_worker.VllmQuantAsyncGenerationWorker" = [ - "modelopt", - "vllm", -] -"nemo_rl.modelopt.models.policy.workers.dtensor_quant_policy_worker.DTensorQuantPolicyWorker" = [ - "modelopt", - "automodel", -] -"nemo_rl.modelopt.models.policy.workers.dtensor_quant_policy_worker_v2.DTensorQuantPolicyWorkerV2" = [ - "modelopt", - "automodel", -] -"nemo_rl.modelopt.models.policy.workers.megatron_quant_policy_worker.MegatronQuantPolicyWorker" = [ - "modelopt", - "mcore", -] - [tool.black] line-length = 120 include = '\.pyi?$' diff --git a/pyrefly.toml b/pyrefly.toml index 1b54fdaab50..52d64fb3ee9 100644 --- a/pyrefly.toml +++ b/pyrefly.toml @@ -150,6 +150,7 @@ project-includes = [ "nemo_rl/data_plane/tq_token_sink.py", "nemo_rl/data_plane/worker_mixin.py", "nemo_rl/distributed/__init__.py", + "nemo_rl/distributed/actor_environments.py", "nemo_rl/distributed/collectives.py", "nemo_rl/distributed/held_port.py", "nemo_rl/distributed/named_sharding.py", diff --git a/tests/unit/distributed/test_actor_environments.py b/tests/unit/distributed/test_actor_environments.py new file mode 100644 index 00000000000..af8e55d6aeb --- /dev/null +++ b/tests/unit/distributed/test_actor_environments.py @@ -0,0 +1,160 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Guards on ACTOR_ENVIRONMENTS, the actor -> uv extras table. + +docker/Dockerfile runs nemo_rl/distributed/actor_environments.py as a script from +the dependency layer to decide which venvs to pre-build, and the runtime registry +imports the same dict. These tests keep the two readers honest. +""" + +import ast +import subprocess +import sys +import tomllib +from pathlib import Path + +import pytest + +from nemo_rl.distributed.actor_environments import ACTOR_ENVIRONMENTS +from nemo_rl.distributed.ray_actor_environment_registry import ( + ACTOR_ENVIRONMENT_REGISTRY, + USE_SYSTEM_EXECUTABLE, +) +from nemo_rl.distributed.virtual_cluster import ( + PY_EXECUTABLES, + git_root, + uv_py_executable, +) + +MODULE_PATH = Path(git_root) / "nemo_rl" / "distributed" / "actor_environments.py" + +with open(Path(git_root) / "pyproject.toml", "rb") as _f: + DECLARED_EXTRAS = set(tomllib.load(_f)["project"]["optional-dependencies"]) + + +@pytest.mark.parametrize("actor_fqn", sorted(ACTOR_ENVIRONMENTS)) +def test_actor_extras_are_declared(actor_fqn): + """Every extra names a real [project.optional-dependencies] entry.""" + extras = ACTOR_ENVIRONMENTS[actor_fqn] + if extras is None: + return + assert isinstance(extras, list) and all(isinstance(e, str) for e in extras), ( + f"{actor_fqn}: value must be None or a list of extras, got {extras!r}" + ) + undeclared = set(extras) - DECLARED_EXTRAS + assert not undeclared, ( + f"{actor_fqn} names extras {sorted(undeclared)} that are not in " + "[project.optional-dependencies] of pyproject.toml" + ) + + +@pytest.mark.parametrize("actor_fqn", sorted(ACTOR_ENVIRONMENTS)) +def test_actor_module_exists(actor_fqn): + """The FQN still points at a module that exists. + + Parsed, not imported: most of these pull in vllm or megatron. + """ + module_name, _, class_name = actor_fqn.rpartition(".") + path = Path(git_root) / (module_name.replace(".", "/") + ".py") + if not path.exists(): + path = Path(git_root) / module_name.replace(".", "/") / "__init__.py" + assert path.exists(), f"{actor_fqn}: no module file for {module_name}" + + defined = set() + for node in ast.walk(ast.parse(path.read_text())): + if isinstance(node, (ast.ClassDef, ast.FunctionDef, ast.AsyncFunctionDef)): + defined.add(node.name) + elif isinstance(node, (ast.Import, ast.ImportFrom)): + defined.update(a.asname or a.name.split(".")[-1] for a in node.names) + assert class_name in defined, f"{actor_fqn}: {class_name} not defined in {path}" + + +@pytest.mark.skipif( + USE_SYSTEM_EXECUTABLE, + reason="NEMO_RL_PY_EXECUTABLES_SYSTEM=1 puts every actor on the driver interpreter", +) +def test_registry_matches_py_executables(): + """The generated py_executable is the string the worker actually needs.""" + expected = { + ("vllm",): PY_EXECUTABLES.VLLM, + ("sglang",): PY_EXECUTABLES.SGLANG, + ("fsdp",): PY_EXECUTABLES.FSDP, + ("automodel",): PY_EXECUTABLES.AUTOMODEL, + ("mcore",): PY_EXECUTABLES.MCORE, + ("trtllm",): PY_EXECUTABLES.TRTLLM, + ("nemo_gym",): PY_EXECUTABLES.NEMO_GYM, + } + for actor_fqn, extras in ACTOR_ENVIRONMENTS.items(): + got = ACTOR_ENVIRONMENT_REGISTRY[actor_fqn] + if extras is None: + assert got == PY_EXECUTABLES.SYSTEM, actor_fqn + else: + assert got == expected.get(tuple(extras), uv_py_executable(extras)), ( + actor_fqn + ) + + +def test_actor_environments_module_is_stdlib_only(): + """docker/Dockerfile runs this module from the dependency layer. + + Only pyproject.toml, uv.lock and a couple of nemo_rl files exist there, so an + import of anything else -- especially anything from nemo_rl -- breaks the image + build in its most expensive layer. + """ + stdlib = set(sys.stdlib_module_names) | {"__future__"} + tree = ast.parse(MODULE_PATH.read_text()) + for node in ast.walk(tree): + roots = [] + if isinstance(node, ast.Import): + roots = [a.name.split(".")[0] for a in node.names] + elif isinstance(node, ast.ImportFrom): + roots = [(node.module or "").split(".")[0]] + for root in roots: + assert root in stdlib, ( + f"{MODULE_PATH.name} imports {root!r}, which is not in the standard " + "library. This module must stay dependency-free -- docker/Dockerfile " + "runs it from a layer where only pyproject.toml and uv.lock exist." + ) + + +def test_script_output_matches_the_registry(): + """Running the module as a script lists exactly the venvs the runtime expects.""" + proc = subprocess.run( + [sys.executable, str(MODULE_PATH), "all"], + capture_output=True, + text=True, + check=True, + cwd=git_root, + ) + listed = {line.split("\t")[0] for line in proc.stdout.splitlines() if line.strip()} + expected = {fqn for fqn, extras in ACTOR_ENVIRONMENTS.items() if extras is not None} + assert listed == expected + + +def test_script_skips_by_extra_not_by_name(): + """SKIP_VLLM_BUILD must drop actors that need vllm even without 'vllm' in the name.""" + proc = subprocess.run( + [sys.executable, str(MODULE_PATH), "all", "vllm"], + capture_output=True, + text=True, + check=True, + cwd=git_root, + ) + listed = {line.split("\t")[0] for line in proc.stdout.splitlines() if line.strip()} + assert "nemo_rl.algorithms.async_utils.AsyncTrajectoryCollector" not in listed + assert "nemo_rl.experience.sync_rollout_actor.SyncRolloutActor" not in listed + assert ( + "nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker" + in listed + ) From 667fde794c3d894dee79ff0dc5113ae7c1c1b235 Mon Sep 17 00:00:00 2001 From: Terry Kong Date: Fri, 4 Sep 2026 15:13:56 -0700 Subject: [PATCH 05/21] fix: hash the actor environment table in the container fingerprint Venvs under NEMO_RL_VENV_DIR are reused rather than rebuilt, and after #3947 nothing prunes them -- the base sync passes --inexact and 'uv run' is inexact by default. So changing an actor's extras leaves the old extra's packages installed in a venv that is silently reused. _check_container_fingerprint() is what catches that and points the user at NRL_FORCE_REBUILD_VENVS=true, but generate_fingerprint.py hashes only pyproject.toml, uv.lock and the submodule SHAs. While the table lived in pyproject.toml that was covered by accident; moving it to a .py file dropped the coverage. Hash the table explicitly so the check still fires. Verified: swapping MegatronValueWorker from ['mcore'] to ['automodel'] changes the fingerprint hash. Adds a test so the coupling cannot break silently. Signed-off-by: Terry Kong --- .../distributed/test_actor_environments.py | 25 +++++++++++++++++++ tools/generate_fingerprint.py | 10 ++++++++ 2 files changed, 35 insertions(+) diff --git a/tests/unit/distributed/test_actor_environments.py b/tests/unit/distributed/test_actor_environments.py index af8e55d6aeb..3d92a4e6132 100644 --- a/tests/unit/distributed/test_actor_environments.py +++ b/tests/unit/distributed/test_actor_environments.py @@ -19,6 +19,7 @@ """ import ast +import importlib.util import subprocess import sys import tomllib @@ -142,6 +143,30 @@ def test_script_output_matches_the_registry(): assert listed == expected +def test_fingerprint_covers_the_actor_table(): + """Editing an actor's extras must invalidate the container fingerprint. + + Venvs at NEMO_RL_VENV_DIR are reused rather than rebuilt, and nothing prunes + them (the base sync runs --inexact and `uv run` is inexact by default). So a + changed extras list has to trip _check_container_fingerprint(), which is what + tells the user to set NRL_FORCE_REBUILD_VENVS=true. + """ + spec = importlib.util.spec_from_file_location( + "_gen_fingerprint", Path(git_root) / "tools" / "generate_fingerprint.py" + ) + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + fingerprint = module.generate_fingerprint() + assert "nemo_rl/distributed/actor_environments.py" in fingerprint, ( + "tools/generate_fingerprint.py must hash the actor -> extras table, or a " + "changed actor environment leaves a stale venv with no warning" + ) + assert fingerprint["nemo_rl/distributed/actor_environments.py"] not in ( + "", + "missing", + ) + + def test_script_skips_by_extra_not_by_name(): """SKIP_VLLM_BUILD must drop actors that need vllm even without 'vllm' in the name.""" proc = subprocess.run( diff --git a/tools/generate_fingerprint.py b/tools/generate_fingerprint.py index 2a235788044..4cabfb0db3d 100755 --- a/tools/generate_fingerprint.py +++ b/tools/generate_fingerprint.py @@ -109,6 +109,7 @@ def generate_fingerprint() -> dict[str, str]: Dictionary mapping component names to their hashes/commits: - "pyproject.toml": MD5 hash of pyproject.toml - "uv.lock": MD5 hash of uv.lock + - "nemo_rl/distributed/actor_environments.py": MD5 hash of the actor -> extras table - "submodules/": Commit SHA for each submodule """ repo_root = get_repo_root() @@ -121,6 +122,15 @@ def generate_fingerprint() -> dict[str, str]: # Hash uv.lock fingerprint["uv.lock"] = compute_file_hash(repo_root / "uv.lock") + # Hash the actor -> uv extras table. Changing an actor's extras changes what its + # pre-built venv should contain, and venvs are reused rather than rebuilt, so this + # has to invalidate the fingerprint the same way a dependency change does. + # Without it a stale venv keeps packages from the extra the actor no longer declares + # and nothing tells the user to set NRL_FORCE_REBUILD_VENVS=true. + fingerprint["nemo_rl/distributed/actor_environments.py"] = compute_file_hash( + repo_root / "nemo_rl" / "distributed" / "actor_environments.py" + ) + # Get submodule SHAs (sorted by path for consistency) submodules = get_submodule_shas(repo_root) for path, sha in sorted(submodules.items()): From 0516535aae7cece1f1dcf96b92352e0f718bd43f Mon Sep 17 00:00:00 2001 From: Terry Kong Date: Fri, 4 Sep 2026 15:18:21 -0700 Subject: [PATCH 06/21] chore(docker): set NEMO_RL_VENV_DIR once in the base stage It was declared three times with the same value: base, hermetic, and release. hermetic and release both derive from base, so the later two were dead. Worse, the hermetic one came *after* the worker-venv prefetch that reads it, so changing it there would have looked effective while the prefetch kept using base's value. Keep the base declaration, note that the later stages inherit it. Signed-off-by: Terry Kong --- docker/Dockerfile | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/docker/Dockerfile b/docker/Dockerfile index 7051ac2431e..401bdb04011 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -129,6 +129,9 @@ ENV RAY_USAGE_STATS_ENABLED=0 # There is severe contention and performance issues with this enabled considering our dependencies are so large and occasionally # need to be compiled, so NeMo RL has an implementation in nemo_rl/utils/venv.py that does it once per node as opposed to once per task. ENV RAY_ENABLE_UV_RUN_RUNTIME_ENV=0 +# Set once here: `hermetic` and `release` both derive from `base`, so they inherit it. +# The worker-venv prefetch in the hermetic stage reads it -- do not re-declare it in a +# later stage, or the value the prefetch uses and the value the image ships can drift. ENV NEMO_RL_VENV_DIR=/opt/ray_venvs ENV NEMO_GYM_VENV_DIR=/opt/gym_venvs @@ -387,7 +390,6 @@ find /root/.cache/uv -type d -path "*ray/_private/runtime_env/agent/thirdparty_f EOF ENV PATH="/opt/nemo_rl_venv/bin:$PATH" -ENV NEMO_RL_VENV_DIR=/opt/ray_venvs # Custom setup layer (override with: --build-context custom-setup= --build-arg CUSTOM_SETUP_FNAME=