cp: feat: add standalone Gym Docker container (2161) into r0.5.0

.github/workflows/build-container.yml

@@ -0,0 +1,47 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# 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.
+name: Build container
+on:
+  pull_request:
+    paths:
+      - docker/Dockerfile
+      - pyproject.toml
+      - uv.lock
+  push:
+    branches: [main]
+    paths:
+      - docker/Dockerfile
+      - pyproject.toml
+      - uv.lock
+  workflow_dispatch:
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout
+        uses: actions/checkout@df4cb1c069e1874edd31b4311f1884172cec0e10 # v6
+
+      - name: Set up Docker Buildx
+        uses: docker/setup-buildx-action@8d2750c68a42422c14e847fe6c8ac0403b4cbd6f # v3
+
+      - name: Build
+        uses: docker/build-push-action@ca052bb54ab0790a636c9b5f226502c73d547a25 # v5
+        with:
+          context: .
+          file: docker/Dockerfile
+          build-contexts: nemo-gym=.
+          target: release
+          push: false

docker/Dockerfile

@@ -0,0 +1,106 @@
+# syntax=docker/dockerfile:1
+# Usage:
+#   Self-contained build (default: builds from main):
+#     docker buildx build -f docker/Dockerfile --tag <registry>/nemo-gym:latest --push .
+#
+#   Self-contained build (specific git ref):
+#     docker buildx build -f docker/Dockerfile --build-arg GYM_GIT_REF=v0.4.0 \
+#       --tag <registry>/nemo-gym:v0.4.0 --push .
+#
+#   Local source override:
+#     docker buildx build --build-context nemo-gym=. -f docker/Dockerfile \
+#       --tag <registry>/nemo-gym:latest --push .
+#
+
+ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.03-cuda13.2-devel-ubuntu24.04
+
+# Source stage: clones Gym from GitHub. Override with --build-context nemo-gym=.
+# to inject a local checkout instead.
+FROM scratch AS nemo-gym
+ARG GYM_GIT_REF=main
+ADD --keep-git-dir=true https://github.com/NVIDIA-NeMo/Gym.git#${GYM_GIT_REF} /
+
+FROM ${BASE_IMAGE} AS base
+ENV NEMO_GYM_CONTAINER=1
+USER root
+
+RUN <<"EOF" bash -exu -o pipefail
+export DEBIAN_FRONTEND=noninteractive
+export TZ=America/Los_Angeles
+
+apt-get update
+apt-get install -y --no-install-recommends \
+    jq \
+    curl \
+    git \
+    rsync \
+    wget \
+    less \
+    vim
+
+apt-get clean
+rm -rf /var/lib/apt/lists/*
+EOF
+
+ARG UV_VERSION=0.11.6
+ARG PYTHON_VERSION=3.12.11
+ENV PATH="/root/.local/bin:$PATH"
+RUN curl -LsSf https://astral.sh/uv/${UV_VERSION}/install.sh | sh && \
+    uv python install ${PYTHON_VERSION}
+
+ENV RAY_USAGE_STATS_ENABLED=0
+
+
+FROM base AS hermetic
+
+WORKDIR /opt/nemo-gym
+
+ARG BASE_IMAGE
+ARG UV_VERSION
+
+ENV UV_PROJECT_ENVIRONMENT=/opt/nemo_gym_venv
+ENV UV_LINK_MODE=copy
+
+# Copy only dependency metadata first for layer caching.
+COPY --from=nemo-gym pyproject.toml uv.lock ./
+COPY --from=nemo-gym nemo_gym/__init__.py nemo_gym/package_info.py ./nemo_gym/
+
+RUN uv venv --seed && \
+    uv sync --link-mode symlink --locked --extra vllm --no-install-project
+
+ENV PATH="/opt/nemo_gym_venv/bin:$PATH"
+
+
+FROM hermetic AS release
+
+ARG NEMO_GYM_PREFETCH_CONFIGS=""
+ARG NEMO_GYM_COMMIT
+ARG NVIDIA_BUILD_ID
+ARG NVIDIA_BUILD_REF
+ARG RC_DATE=00.00
+ARG TARGETARCH
+ENV NEMO_GYM_COMMIT=${NEMO_GYM_COMMIT:-<unknown>}
+ENV NVIDIA_BUILD_ID=${NVIDIA_BUILD_ID:-<unknown>}
+ENV NVIDIA_BUILD_REF=${NVIDIA_BUILD_REF:-<unknown>}
+LABEL com.nvidia.build.id="${NVIDIA_BUILD_ID}"
+LABEL com.nvidia.build.ref="${NVIDIA_BUILD_REF}"
+
+# Copy full source. Exclude pyproject.toml and uv.lock since they are
+# already present from the hermetic stage.
+COPY --from=nemo-gym --exclude=pyproject.toml --exclude=uv.lock . /opt/nemo-gym
+
+# Install the nemo-gym project itself along with the vllm extra.
+RUN UV_LINK_MODE=symlink uv sync --locked --extra vllm
+
+# Optional: pre-warm per-server venvs at build time so they are ready at runtime.
+# Pass --build-arg NEMO_GYM_PREFETCH_CONFIGS="path/to/config1.yaml,path/to/config2.yaml"
+# Uses 'gym env prefetch' which accepts the same config format as 'gym env start',
+# installs each server's venv serially, and exits without starting any server process.
+RUN <<"EOF" bash -exu
+if [[ -n "${NEMO_GYM_PREFETCH_CONFIGS:-}" ]]; then
+    gym env prefetch "+config_paths=[${NEMO_GYM_PREFETCH_CONFIGS}]"
+fi
+EOF
+
+ENTRYPOINT ["gym"]
+CMD ["--help"]

nemo_gym/cli/env.py

@@ -427,6 +427,74 @@ def check_http_server_statuses(self, successful_servers: List[str]) -> List[Tupl
         return statuses
 
 
+@exit_cleanly_on_config_error
+def prefetch(
+    global_config_dict_parser_config: Optional[GlobalConfigDictParserConfig] = None,
+):  # pragma: no cover
+    """
+    Pre-warm per-server venvs without starting any servers.
+
+    Accepts the same config format as 'gym env start'. For each server in the config,
+    creates its isolated venv and installs dependencies serially — similar to how
+    RL's prefetch_venvs.py installs actor venvs at container build time.
+
+    No Ray is initialised and no server processes are started. Intended for use in
+    Dockerfile builds so venvs are ready at runtime with no network access needed.
+
+    Configuration Parameter:
+        config_paths (List[str]): Paths to YAML configuration files. Specify via Hydra: `+config_paths="[file1.yaml,file2.yaml]"`
+
+    Examples:
+
+    ```bash
+    # Pre-warm venvs for a specific config
+    config_paths="responses_api_models/local_vllm_model/configs/local_vllm_model.yaml,\\
+    resources_servers/math/configs/math.yaml"
+    gym env prefetch "+config_paths=[${config_paths}]"
+    ```
+    """
+    global_config_dict = get_global_config_dict(global_config_dict_parser_config=global_config_dict_parser_config)
+    GlobalConfigDictParser().raise_on_no_server_instances(global_config_dict)
+
+    top_level_paths = [k for k in global_config_dict.keys() if k not in NEMO_GYM_RESERVED_TOP_LEVEL_KEYS]
+
+    seen_dirs: set = set()
+    for top_level_path in top_level_paths:
+        server_config_dict = global_config_dict[top_level_path]
+        if not isinstance(server_config_dict, DictConfig):
+            continue
+
+        first_key = list(server_config_dict)[0]
+        server_config_dict = server_config_dict[first_key]
+        if not isinstance(server_config_dict, DictConfig):
+            continue
+        second_key = list(server_config_dict)[0]
+        server_config_dict = server_config_dict[second_key]
+        if not isinstance(server_config_dict, DictConfig):
+            continue
+
+        if "entrypoint" not in server_config_dict:
+            continue
+
+        dir_path = _resolve_server_dir(Path(first_key, second_key))
+        if dir_path in seen_dirs:
+            print(f"Skipping duplicate: {dir_path}")
+            continue
+        seen_dirs.add(dir_path)
+        print(f"Pre-warming venv for: {dir_path}")
+
+        # Run only the setup (venv creation + dep install) — no server start.
+        # Serial execution reduces uv cache size (same reasoning as dry_run in RunHelper.start).
+        setup_cmd = setup_env_command(dir_path, global_config_dict, top_level_path)
+        process = run_command(setup_cmd, dir_path, server_name=top_level_path)
+        returncode = process.wait()
+        if returncode != 0:
+            print(f"ERROR: prefetch failed for {dir_path} (exit {returncode})")
+            raise SystemExit(returncode)
+
+    print("Prefetch complete.")
+
+
 @exit_cleanly_on_config_error
 def run(
     global_config_dict_parser_config: Optional[GlobalConfigDictParserConfig] = None,

nemo_gym/cli/main.py

@@ -530,6 +530,18 @@ def _dataset_download(args: argparse.Namespace, overrides: list[str]) -> None:
             MODEL_API_KEY,
         ),
     ),
+    "env prefetch": Command(
+        target="nemo_gym.cli.env:prefetch",
+        summary="Pre-warm per-server venvs without starting servers.",
+        flags=(
+            CONFIG,
+            BENCHMARK,
+            ENVIRONMENT,
+            RESOURCES_SERVER_CONFIG,
+            MODEL_TYPE,
+            SEARCH_DIR,
+        ),
+    ),
     "env status": Command(target="nemo_gym.cli.env:status", summary="Print the server status.", flags=(JSON,)),
     "eval prepare": Command(
         target="nemo_gym.cli.eval:prepare_benchmark",

pyproject.toml

@@ -228,6 +228,15 @@ all = [
     "nemo-gym[dev,sandbox]",
 ]
 
+vllm = [
+    # Pin to the same version used by local_vllm_model / genrm_model server venvs.
+    "vllm==0.20.0",
+    # Pin flashinfer to the exact version vllm 0.20.0 specifies, ensuring pre-compiled
+    # CUDA kernels are used (avoids JIT compilation on first generation == slow step time).
+    # flashinfer>=0.2 bundles compiled kernels in the python package directly.
+    "flashinfer-python==0.6.8.post1",
+]
+
 sandbox = [
     # Tenacity: Retry helpers used by sandbox providers.
     # Updated: Sat May 09, 2026 with tenacity==9.1.4

responses_api_models/genrm_model/setup.py

@@ -23,7 +23,7 @@
 
     # We specifically pin the vllm dependency because we have tested on this version.
     # Updated Tue Jun 23, 2026 with vllm==0.20.0
-    # License: Apache 2.0 https://github.com/vllm-project/vllm/blob/88d34c6409e9fb3c7b8ca0c04756f061d2099eb1/LICENSE
+    # License: Apache 2.0 https://github.com/vllm-project/vllm/blob/main/LICENSE
     # "vllm==0.20.0",
     # VLLM is resolved below since installation on Macs requires special workarounds.
 
@@ -48,6 +48,11 @@
     dependencies.append("vllm==0.11.0")
 else:
     dependencies.append("vllm==0.20.0")
+    # Pin flashinfer to the exact version vllm 0.20.0 requires — pre-compiled CUDA kernels,
+    # avoids JIT compilation on first generation. Must stay in sync with pyproject.toml [vllm].
+    # Updated Tue Jun 23, 2026 with flashinfer-python==0.6.8.post1
+    # License: Apache 2.0 https://github.com/flashinfer-ai/flashinfer/blob/main/LICENSE
+    dependencies.append("flashinfer-python==0.6.8.post1")
 
 
 setuptools.setup(install_requires=dependencies)

responses_api_models/local_vllm_model/setup.py

@@ -20,7 +20,7 @@
 
     # We specifically pin the vllm dependency because we have tested on this version.
     # Updated Tue Jun 23, 2026 with vllm==0.20.0
-    # License: Apache 2.0 https://github.com/vllm-project/vllm/blob/88d34c6409e9fb3c7b8ca0c04756f061d2099eb1/LICENSE
+    # License: Apache 2.0 https://github.com/vllm-project/vllm/blob/main/LICENSE
     # "vllm==0.20.0",
     # VLLM is resolved below since installation on Macs requires special workarounds.
 
@@ -45,6 +45,11 @@
     dependencies.append("vllm==0.11.0")
 else:
     dependencies.append("vllm==0.20.0")
+    # Pin flashinfer to the exact version vllm 0.20.0 requires — pre-compiled CUDA kernels,
+    # avoids JIT compilation on first generation. Must stay in sync with pyproject.toml [vllm].
+    # Updated Tue Jun 23, 2026 with flashinfer-python==0.6.8.post1
+    # License: Apache 2.0 https://github.com/flashinfer-ai/flashinfer/blob/main/LICENSE
+    dependencies.append("flashinfer-python==0.6.8.post1")
 
 
 setuptools.setup(install_requires=dependencies)