Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions .github/workflows/cicd-main.yml
Original file line number Diff line number Diff line change
Expand Up @@ -222,6 +222,7 @@ jobs:
build-args: |
MAX_JOBS=4
NEMO_RL_COMMIT=${{ github.sha }}
SKIP_SGLANG_BUILD=1

cicd-doc-tests:
strategy:
Expand Down
48 changes: 41 additions & 7 deletions docker/Dockerfile
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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
Expand Down
2 changes: 1 addition & 1 deletion docs/design-docs/dependency-management.md
Original file line number Diff line number Diff line change
Expand Up @@ -104,7 +104,7 @@ ACTOR_ENVIRONMENT_REGISTRY: dict[str, str] = {

## Container Pre-caching

When a [release container](../docker.md#release-image) is built, it pre-caches:
When a [release container](../docker.md#building-the-release-image) is built, it pre-caches:

1. **Virtual environments**: All worker virtual environments are created and stored in the container
2. **UV cache**: Python packages are pre-downloaded into the UV cache directory
Expand Down
72 changes: 43 additions & 29 deletions docs/docker.md
Original file line number Diff line number Diff line change
@@ -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.
32 changes: 30 additions & 2 deletions nemo_rl/utils/prefetch_venvs.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 = []
Expand All @@ -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)")
Expand Down Expand Up @@ -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:
Expand Down Expand Up @@ -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(
Expand All @@ -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,
)
2 changes: 1 addition & 1 deletion tests/functional/L1_Functional_Tests_GPU.sh
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ time uv run --no-sync bash ./tests/functional/grpo_megatron.sh
time uv run --no-sync bash ./tests/functional/grpo_megatron_generation.sh
time uv run --no-sync bash ./tests/functional/grpo_multiturn.sh
time uv run --no-sync bash ./tests/functional/grpo_non_colocated.sh
time uv run --no-sync bash ./tests/functional/grpo_sglang.sh
# time uv run --no-sync bash ./tests/functional/grpo_sglang.sh
time uv run --no-sync bash ./tests/functional/grpo_multiple_datasets.sh
time uv run --no-sync bash ./tests/functional/dpo_automodel_lora.sh
time uv run --no-sync bash ./tests/functional/dpo.sh
Expand Down
14 changes: 7 additions & 7 deletions tests/unit/L0_Unit_Tests_Generation.sh
Original file line number Diff line number Diff line change
Expand Up @@ -46,10 +46,10 @@ else
uv run --extra vllm bash -x ./tests/run_unit.sh unit/models/generation/ --cov=nemo_rl --cov-append --cov-report=term-missing --cov-report=json --hf-gated --vllm-only
fi

# Check and run sglang tests
exit_code=$(uv run --extra sglang pytest tests/unit/models/generation/ --collect-only --hf-gated --sglang-only -q >/dev/null 2>&1; echo $?)
if [[ $exit_code -eq 5 ]]; then
echo "No sglang tests to run"
else
uv run --extra sglang bash -x ./tests/run_unit.sh unit/models/generation/ --cov=nemo_rl --cov-append --cov-report=term-missing --cov-report=json --hf-gated --sglang-only
fi
# # Check and run sglang tests
# exit_code=$(uv run --extra sglang pytest tests/unit/models/generation/ --collect-only --hf-gated --sglang-only -q >/dev/null 2>&1; echo $?)
# if [[ $exit_code -eq 5 ]]; then
# echo "No sglang tests to run"
# else
# uv run --extra sglang bash -x ./tests/run_unit.sh unit/models/generation/ --cov=nemo_rl --cov-append --cov-report=term-missing --cov-report=json --hf-gated --sglang-only
# fi
14 changes: 7 additions & 7 deletions tests/unit/L0_Unit_Tests_Other.sh
Original file line number Diff line number Diff line change
Expand Up @@ -46,13 +46,13 @@ else
uv run --extra vllm bash -x ./tests/run_unit.sh unit/ --ignore=unit/models/generation/ --ignore=unit/models/policy/ --cov=nemo_rl --cov-append --cov-report=term-missing --cov-report=json --hf-gated --vllm-only
fi

# Check and run sglang tests
exit_code=$(uv run --extra sglang pytest tests/unit/ --ignore=tests/unit/models/generation/ --ignore=tests/unit/models/policy/ --collect-only --hf-gated --sglang-only -q >/dev/null 2>&1; echo $?)
if [[ $exit_code -eq 5 ]]; then
echo "No sglang tests to run"
else
uv run --extra sglang bash -x ./tests/run_unit.sh unit/ --ignore=unit/models/generation/ --ignore=unit/models/policy/ --cov=nemo_rl --cov-append --cov-report=term-missing --cov-report=json --hf-gated --sglang-only
fi
# # Check and run sglang tests
# exit_code=$(uv run --extra sglang pytest tests/unit/ --ignore=tests/unit/models/generation/ --ignore=tests/unit/models/policy/ --collect-only --hf-gated --sglang-only -q >/dev/null 2>&1; echo $?)
# if [[ $exit_code -eq 5 ]]; then
# echo "No sglang tests to run"
# else
# uv run --extra sglang bash -x ./tests/run_unit.sh unit/ --ignore=unit/models/generation/ --ignore=unit/models/policy/ --cov=nemo_rl --cov-append --cov-report=term-missing --cov-report=json --hf-gated --sglang-only
# fi

# Research unit tests
for i in research/*/tests/unit; do
Expand Down
14 changes: 7 additions & 7 deletions tests/unit/L0_Unit_Tests_Policy.sh
Original file line number Diff line number Diff line change
Expand Up @@ -46,10 +46,10 @@ else
uv run --extra vllm bash -x ./tests/run_unit.sh unit/models/policy/ --cov=nemo_rl --cov-append --cov-report=term-missing --cov-report=json --hf-gated --vllm-only
fi

# Check and run sglang tests
exit_code=$(uv run --extra sglang pytest tests/unit/models/policy/ --collect-only --hf-gated --sglang-only -q >/dev/null 2>&1; echo $?)
if [[ $exit_code -eq 5 ]]; then
echo "No sglang tests to run"
else
uv run --extra sglang bash -x ./tests/run_unit.sh unit/models/policy/ --cov=nemo_rl --cov-append --cov-report=term-missing --cov-report=json --hf-gated --sglang-only
fi
# # Check and run sglang tests
# exit_code=$(uv run --extra sglang pytest tests/unit/models/policy/ --collect-only --hf-gated --sglang-only -q >/dev/null 2>&1; echo $?)
# if [[ $exit_code -eq 5 ]]; then
# echo "No sglang tests to run"
# else
# uv run --extra sglang bash -x ./tests/run_unit.sh unit/models/policy/ --cov=nemo_rl --cov-append --cov-report=term-missing --cov-report=json --hf-gated --sglang-only
# fi
Loading