feat: add way of excluding generation backends

Signed-off-by: Terry Kong <[email protected]>

Author: terrykong

docker/Dockerfile

@@ -1,9 +1,20 @@
 # syntax=docker/dockerfile:1
 # Usage:
-# Self-contained build (default: builds from main): docker buildx build -f docker/Dockerfile --tag <registry>/nemo-rl:latest --push .
-# Self-contained build (specific git ref): docker buildx build -f docker/Dockerfile --build-arg NRL_GIT_REF=r0.3.0 --tag <registry>/nemo-rl:r0.3.0 --push .
-# Self-contained build (remote NeMo RL source; no need for a local clone of NeMo RL): docker buildx build -f docker/Dockerfile --build-arg NRL_GIT_REF=r0.3.0 --tag <registry>/nemo-rl:r0.3.0 --push https://github.com/NVIDIA-NeMo/RL.git
-# Local NeMo RL source override: docker buildx build --build-context nemo-rl=. -f docker/Dockerfile --tag <registry>/nemo-rl:latest --push .
+#   Self-contained build (default: builds from main):
+#     docker buildx build -f docker/Dockerfile --tag <registry>/nemo-rl:latest --push .
+#
+#   Self-contained build (specific git ref):
+#     docker buildx build -f docker/Dockerfile --build-arg NRL_GIT_REF=r0.3.0 --tag <registry>/nemo-rl:r0.3.0 --push .
+#
+#   Self-contained build (remote NeMo RL source; no need for a local clone of NeMo RL):
+#     docker buildx build -f docker/Dockerfile --build-arg NRL_GIT_REF=r0.3.0 --tag <registry>/nemo-rl:r0.3.0 --push https://github.com/NVIDIA-NeMo/RL.git
+#
+#   Local NeMo RL source override:
+#     docker buildx build --build-context nemo-rl=. -f docker/Dockerfile --tag <registry>/nemo-rl:latest --push .
+#
+# Optional build args to skip vLLM or SGLang dependencies:
+#   --build-arg SKIP_VLLM_BUILD=1    # Skip vLLM dependencies
+#   --build-arg SKIP_SGLANG_BUILD=1  # Skip SGLang dependencies
 
 ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:25.05-cuda12.9-devel-ubuntu24.04
 FROM scratch AS nemo-rl
@@ -84,6 +95,9 @@ ARG MAX_JOBS
 ARG NVTE_BUILD_THREADS_PER_JOB
 # Only use for custom vllm installs. Learn more at https://github.com/NVIDIA-NeMo/RL/blob/main/docs/guides/use-custom-vllm.md
 ARG BUILD_CUSTOM_VLLM
+# Skip building vLLM or SGLang dependencies (set to any non-empty value to skip)
+ARG SKIP_VLLM_BUILD
+ARG SKIP_SGLANG_BUILD
 
 ENV UV_PROJECT_ENVIRONMENT=/opt/nemo_rl_venv
 ENV UV_LINK_MODE=copy
@@ -113,8 +127,12 @@ fi
 
 # The venv is symlinked to avoid bloating the layer size
 uv sync --link-mode symlink --locked --no-install-project
-uv sync --link-mode symlink --locked --extra vllm --no-install-project
-uv sync --link-mode symlink --locked --extra sglang --no-install-project
+if [[ -z "${SKIP_VLLM_BUILD:-}" ]]; then
+    uv sync --link-mode symlink --locked --extra vllm --no-install-project
+fi
+if [[ -z "${SKIP_SGLANG_BUILD:-}" ]]; then
+    uv sync --link-mode symlink --locked --extra sglang --no-install-project
+fi
 uv sync --link-mode symlink --locked --extra mcore --no-install-project
 uv sync --link-mode symlink --locked --extra automodel --no-install-project
 uv sync --link-mode symlink --locked --all-groups --no-install-project
@@ -131,6 +149,9 @@ WORKDIR /opt/nemo-rl
 
 FROM hermetic AS release
 
+# Re-declare build args for this stage
+ARG SKIP_VLLM_BUILD
+ARG SKIP_SGLANG_BUILD
 ARG NEMO_RL_COMMIT
 ARG NVIDIA_BUILD_ID
 ARG NVIDIA_BUILD_REF
@@ -151,7 +172,20 @@ COPY --from=nemo-rl --exclude=pyproject.toml --exclude=uv.lock . /opt/nemo-rl
 # Potentially not necessary if the repo is passed in as a complete repository (w/ full git history),
 # so do a quick check before trying to unshallow.
 RUN git rev-parse --is-shallow-repository | grep -q true && git fetch --unshallow || true
-RUN UV_LINK_MODE=symlink uv run nemo_rl/utils/prefetch_venvs.py
+RUN <<"EOF" bash -exu
+NEGATIVE_FILTERS=""
+if [[ -n "${SKIP_VLLM_BUILD:-}" ]]; then
+    NEGATIVE_FILTERS="$NEGATIVE_FILTERS vllm"
+fi
+if [[ -n "${SKIP_SGLANG_BUILD:-}" ]]; then
+    NEGATIVE_FILTERS="$NEGATIVE_FILTERS sglang"
+fi
+if [[ -n "$NEGATIVE_FILTERS" ]]; then
+    UV_LINK_MODE=symlink 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
+fi
+EOF
 
 # Generate container fingerprint for frozen environment support
 # Store outside /opt/nemo-rl to avoid being overwritten by user mounts

docs/docker.md

@@ -1,50 +1,64 @@
 # Build Docker Images
 
-This guide provides two methods for building Docker images:
+This guide explains how to build the NeMo RL Docker image.
 
-* **release**: Contains everything from the hermetic image, plus the nemo-rl source code and pre-fetched virtual environments for isolated workers.
-* **hermetic**: Includes the base image plus pre-fetched NeMo RL python packages in the `uv` cache.
+The **release** image is our recommended option as it provides the most complete environment. It includes the base image with pre-fetched NeMo RL python packages in the `uv` cache, plus the nemo-rl source code and pre-fetched virtual environments for isolated workers. This is the ideal choice for production deployments.
 
-Use the:
-* **release** (recommended): if you want to pre-fetch the NeMo RL [worker virtual environments](./design-docs/uv.md#worker-configuration) and copy in the project source code.
-* **hermetic**: if you want to pre-fetch NeMo RL python packages into the `uv` cache to eliminate the initial overhead of program start.
-
-## Release Image
-
-The release image is our recommended option as it provides the most complete environment. It includes everything from the hermetic image, plus the nemo-rl source code and pre-fetched virtual environments for isolated workers. This is the ideal choice for production deployments.
+## Building the Release Image
 
 ```sh
 # Self-contained build (default: builds from main):
-docker buildx build --target release -f docker/Dockerfile --tag <registry>/nemo-rl:latest --push .
+docker buildx build -f docker/Dockerfile \
+    --tag <registry>/nemo-rl:latest \
+    --push .
 
 # Self-contained build (specific git ref):
-docker buildx build --target release -f docker/Dockerfile --build-arg NRL_GIT_REF=r0.3.0 --tag <registry>/nemo-rl:r0.3.0 --push .
+docker buildx build -f docker/Dockerfile \
+    --build-arg NRL_GIT_REF=r0.3.0 \
+    --tag <registry>/nemo-rl:r0.3.0 \
+    --push .
 
 # Self-contained build (remote NeMo RL source; no need for a local clone of NeMo RL):
-docker buildx build --target release -f docker/Dockerfile --build-arg NRL_GIT_REF=r0.3.0 --tag <registry>/nemo-rl:r0.3.0 --push https://github.com/NVIDIA-NeMo/RL.git
+docker buildx build -f docker/Dockerfile \
+    --build-arg NRL_GIT_REF=r0.3.0 \
+    --tag <registry>/nemo-rl:r0.3.0 \
+    --push https://github.com/NVIDIA-NeMo/RL.git
 
 # Local NeMo RL source override:
-docker buildx build --target release --build-context nemo-rl=. -f docker/Dockerfile --tag <registry>/nemo-rl:latest --push .
+docker buildx build --build-context nemo-rl=. -f docker/Dockerfile \
+    --tag <registry>/nemo-rl:latest \
+    --push .
 ```
 
-**Note:** The `--tag <registry>/nemo-rl:latest --push` flags are not necessary if you just want to build locally.
+> [!NOTE]
+> The `--tag <registry>/nemo-rl:latest --push` flags are not necessary if you just want to build locally.
 
-## Hermetic Image
+## Skipping vLLM or SGLang Dependencies
 
-The hermetic image includes all Python dependencies pre-downloaded in the `uv` cache, eliminating the initial overhead of downloading packages at runtime. This is useful when you need a more predictable environment or have limited network connectivity.
+If you don't need vLLM or SGLang support, you can skip building those dependencies to reduce build time and image size. Use the `SKIP_VLLM_BUILD` and/or `SKIP_SGLANG_BUILD` build arguments:
 
 ```sh
-# Self-contained build (default: builds from main):
-docker buildx build --target hermetic -f docker/Dockerfile --tag <registry>/nemo-rl:latest --push .
-
-# Self-contained build (specific git ref):
-docker buildx build --target hermetic -f docker/Dockerfile --build-arg NRL_GIT_REF=r0.3.0 --tag <registry>/nemo-rl:r0.3.0 --push .
-
-# Self-contained build (remote NeMo RL source; no need for a local clone of NeMo RL):
-docker buildx build --target hermetic -f docker/Dockerfile --build-arg NRL_GIT_REF=r0.3.0 --tag <registry>/nemo-rl:r0.3.0 --push https://github.com/NVIDIA-NeMo/RL.git
-
-# Local NeMo RL source override:
-docker buildx build --target hermetic --build-context nemo-rl=. -f docker/Dockerfile --tag <registry>/nemo-rl:latest --push .
+# Skip vLLM dependencies:
+docker buildx build -f docker/Dockerfile \
+    --build-arg SKIP_VLLM_BUILD=1 \
+    --tag <registry>/nemo-rl:latest \
+    .
+
+# Skip SGLang dependencies:
+docker buildx build -f docker/Dockerfile \
+    --build-arg SKIP_SGLANG_BUILD=1 \
+    --tag <registry>/nemo-rl:latest \
+    .
+
+# Skip both vLLM and SGLang dependencies:
+docker buildx build -f docker/Dockerfile \
+    --build-arg SKIP_VLLM_BUILD=1 \
+    --build-arg SKIP_SGLANG_BUILD=1 \
+    --tag <registry>/nemo-rl:latest \
+    .
 ```
 
-**Note:** The `--tag <registry>/nemo-rl:latest --push` flags are not necessary if you just want to build locally.
+When these build arguments are set, the corresponding `uv sync --extra` commands are skipped, and the virtual environment prefetching will exclude actors that depend on those packages.
+
+> [!NOTE]
+> If you skip vLLM or SGLang during the build but later try to use those backends at runtime, the dependencies will be fetched and built on-demand. This may add significant setup time on first use.

nemo_rl/utils/prefetch_venvs.py

@@ -22,20 +22,25 @@
 from nemo_rl.utils.venvs import create_local_venv
 
 
-def prefetch_venvs(filters=None):
+def prefetch_venvs(filters=None, negative_filters=None):
     """Prefetch all virtual environments that will be used by workers.
 
     Args:
         filters: List of strings to match against actor FQNs. If provided, only
                 actors whose FQN contains at least one of the filter strings will
                 be prefetched. If None, all venvs are prefetched.
+        negative_filters: List of strings to exclude from prefetching. Actors whose
+                FQN contains any of these strings will be skipped.
     """
     print("Prefetching virtual environments...")
     if filters:
         print(f"Filtering for: {filters}")
+    if negative_filters:
+        print(f"Excluding: {negative_filters}")
 
     # Track statistics for summary
     skipped_by_filter = []
+    skipped_by_negative_filter = []
     skipped_system_python = []
     prefetched = []
     failed = []
@@ -47,6 +52,10 @@ def prefetch_venvs(filters=None):
         if filters and not any(f in actor_fqn for f in filters):
             skipped_by_filter.append(actor_fqn)
             continue
+        # Apply negative filters if provided
+        if negative_filters and any(f in actor_fqn for f in negative_filters):
+            skipped_by_negative_filter.append(actor_fqn)
+            continue
         # Skip system python as it doesn't need a venv
         if py_executable == "python" or py_executable == sys.executable:
             print(f"Skipping {actor_fqn} (uses system Python)")
@@ -88,6 +97,10 @@ def prefetch_venvs(filters=None):
         print(f"  Skipped (filtered out): {len(skipped_by_filter)}")
         for actor_fqn in skipped_by_filter:
             print(f"    - {actor_fqn}")
+    if negative_filters:
+        print(f"  Skipped (negative filter): {len(skipped_by_negative_filter)}")
+        for actor_fqn in skipped_by_negative_filter:
+            print(f"    - {actor_fqn}")
     if failed:
         print(f"  Failed: {len(failed)}")
         for actor_fqn in failed:
@@ -202,6 +215,12 @@ def create_frozen_environment_symlinks(venv_configs):
 
   # Prefetch multiple specific venvs
   python -m nemo_rl.utils.prefetch_venvs vllm policy environment
+
+  # Prefetch all venvs except vLLM-related ones
+  python -m nemo_rl.utils.prefetch_venvs --negative-filters vllm
+
+  # Prefetch all venvs except vLLM and SGLang
+  python -m nemo_rl.utils.prefetch_venvs --negative-filters vllm sglang
         """,
     )
     parser.add_argument(
@@ -211,6 +230,15 @@ def create_frozen_environment_symlinks(venv_configs):
         "contains at least one of these strings will be prefetched. "
         "If not provided, all venvs are prefetched.",
     )
+    parser.add_argument(
+        "--negative-filters",
+        nargs="*",
+        help="Filter strings to exclude from prefetching. Actors whose FQN "
+        "contains any of these strings will be skipped.",
+    )
     args = parser.parse_args()
 
-    prefetch_venvs(filters=args.filters if args.filters else None)
+    prefetch_venvs(
+        filters=args.filters if args.filters else None,
+        negative_filters=args.negative_filters if args.negative_filters else None,
+    )