|
| 1 | +# ============================================================================= |
| 2 | +# NVFlow NeMo-RL Container |
| 3 | +# ============================================================================= |
| 4 | +# Extends the NeMo-RL nightly container with NeMo-Skills and the NeMo-Gym |
| 5 | +# finance agent for NVFlow workflows. The base image ships with frozen |
| 6 | +# environments and pre-built Ray venvs; this Dockerfile adds the NeMo-Skills |
| 7 | +# package, replaces the Gym submodule with a feature branch that includes the |
| 8 | +# finance-SEC-search resource server and finance agent, pre-builds all Gym |
| 9 | +# component venvs, and relocates paths for Slurm/enroot compatibility. |
| 10 | +# |
| 11 | +# Build: |
| 12 | +# docker build -f dockerfiles/Dockerfile.nemo-rl -t nvflow-nemo-rl:latest . |
| 13 | +# ============================================================================= |
| 14 | + |
| 15 | +ARG BASE_IMAGE=nvcr.io/nvidia/nemo-rl:v0.6.0 |
| 16 | +FROM ${BASE_IMAGE} |
| 17 | + |
| 18 | +# --- Symlink for NeMo-Skills code that references /opt/NeMo-RL (wrong case) -- |
| 19 | +RUN ln -sf /opt/nemo-rl /opt/NeMo-RL |
| 20 | + |
| 21 | +# --- Upgrade uv (nemo-gym may require newer features than what the nightly ships) |
| 22 | +RUN curl -LsSf https://astral.sh/uv/install.sh | sh |
| 23 | + |
| 24 | +# --- Pre-cache Python interpreter for uv (air-gapped safety net) ------------- |
| 25 | +# Gym's cli_setup_command runs `uv venv --python <version>`. If |
| 26 | +# skip_venv_if_present ever misses, uv still needs a local interpreter. |
| 27 | +RUN /root/.local/bin/uv python install 3.12 |
| 28 | + |
| 29 | +# --- Install NeMo-Skills into the frozen venv -------------------------------- |
| 30 | +ARG NEMO_SKILLS_COMMIT=022904023ad7a83a87662a313cf72e7df5891d55 |
| 31 | +RUN git clone https://github.com/NVIDIA-NeMo/Skills.git /opt/NeMo-Skills && \ |
| 32 | + cd /opt/NeMo-Skills && git checkout ${NEMO_SKILLS_COMMIT} && \ |
| 33 | + /root/.local/bin/uv pip install --python /opt/nemo_rl_venv/bin/python . |
| 34 | + |
| 35 | +# --- Replace NeMo-Gym submodule with feature branch ------------------------- |
| 36 | +# The feature branch includes the finance-SEC-search resource server and |
| 37 | +# finance agent that are not yet on main. |
| 38 | +ARG NEMO_GYM_BRANCH=ude/finance-sec-search-v2 |
| 39 | +RUN rm -rf /opt/nemo-rl/3rdparty/Gym-workspace/Gym && \ |
| 40 | + git clone --branch ${NEMO_GYM_BRANCH} \ |
| 41 | + https://github.com/NVIDIA-NeMo/Gym.git \ |
| 42 | + /opt/nemo-rl/3rdparty/Gym-workspace/Gym |
| 43 | + |
| 44 | +# --- Pre-build Gym venv ------------------------------------------------------ |
| 45 | +WORKDIR /opt/nemo-rl/3rdparty/Gym-workspace/Gym |
| 46 | +RUN /root/.local/bin/uv venv .venv --python 3.12 && \ |
| 47 | + . .venv/bin/activate && \ |
| 48 | + /root/.local/bin/uv sync --active --extra dev |
| 49 | + |
| 50 | +# Install finance-specific dependencies into Gym venv |
| 51 | +# uvicorn>=0.37.0 is required for timeout_worker_healthcheck support; |
| 52 | +# uv sync resolves from the parent nemo-rl workspace lock (0.35.0) instead |
| 53 | +# of the Gym lock, so we force the correct version here. |
| 54 | +RUN . .venv/bin/activate && \ |
| 55 | + /root/.local/bin/uv pip install aiohttp beautifulsoup4 "tavily==1.1.0" tenacity "uvicorn>=0.37.0" |
| 56 | + |
| 57 | +# --- Symlink component venvs to the main Gym venv --------------------------- |
| 58 | +# Each NeMo-Gym component expects its own .venv/; symlinking avoids multi-GB |
| 59 | +# duplication and guarantees every component runs with the same packages. |
| 60 | +RUN for component in \ |
| 61 | + resources_servers/equivalence_llm_judge \ |
| 62 | + resources_servers/finance_sec_search \ |
| 63 | + responses_api_agents/simple_agent \ |
| 64 | + responses_api_agents/finance_agent \ |
| 65 | + responses_api_models/openai_model \ |
| 66 | + responses_api_models/vllm_model; do \ |
| 67 | + dir="/opt/nemo-rl/3rdparty/Gym-workspace/Gym/$component"; \ |
| 68 | + [ -d "$dir" ] && ln -sf /opt/nemo-rl/3rdparty/Gym-workspace/Gym/.venv "$dir/.venv"; \ |
| 69 | + done |
| 70 | + |
| 71 | +WORKDIR / |
| 72 | + |
| 73 | +# --- Install Gym into the NemoGym Ray venv ------------------------------------ |
| 74 | +# The pre-built Ray venv from the base image is stale (built from the old Gym |
| 75 | +# submodule). Install the new Gym branch editable + all deps so the Ray actor |
| 76 | +# can import nemo_gym without missing modules (e.g. gprof2dot, pydot). |
| 77 | +RUN /root/.local/bin/uv pip install \ |
| 78 | + --python /opt/ray_venvs/nemo_rl.environments.nemo_gym.NemoGym/bin/python \ |
| 79 | + -e /opt/nemo-rl/3rdparty/Gym-workspace/Gym |
| 80 | + |
| 81 | +# --- Align numpy across all Ray venvs to match the main venv ---------------- |
| 82 | +# NeMo-Skills may upgrade numpy; mismatched versions cause pickle failures |
| 83 | +# when Ray serializes data between the main process and worker processes. |
| 84 | +RUN MAIN_NP=$(/opt/nemo_rl_venv/bin/python -c "import numpy; print(numpy.__version__)") && \ |
| 85 | + for venv in /opt/ray_venvs/*/; do \ |
| 86 | + "$venv/bin/pip" install --no-cache-dir "numpy==$MAIN_NP" 2>/dev/null || true; \ |
| 87 | + done |
| 88 | + |
| 89 | +# --- Relocate /root/.local/ → /opt/ ----------------------------------------- |
| 90 | +# enroot/pyxis on Slurm mounts the user's home directory over /root at runtime, |
| 91 | +# which shadows everything uv installed there during the Docker build. |
| 92 | +# NOTE: Do NOT move /root/.cache/uv — base-image venvs symlink into it. |
| 93 | +RUN REAL_PYTHON=$(readlink /opt/nemo_rl_venv/bin/python) && \ |
| 94 | + mv /root/.local/share/uv/python /opt/uv-python && \ |
| 95 | + find /opt/uv-python -maxdepth 1 -type l | while read link; do \ |
| 96 | + target=$(readlink "$link") && \ |
| 97 | + new_target=$(echo "$target" | sed "s|/root/.local/share/uv/python|/opt/uv-python|") && \ |
| 98 | + ln -sf "$new_target" "$link"; \ |
| 99 | + done && \ |
| 100 | + NEW_PYTHON=$(echo "$REAL_PYTHON" | sed "s|/root/.local/share/uv/python|/opt/uv-python|") && \ |
| 101 | + ln -sf "$NEW_PYTHON" /opt/nemo_rl_venv/bin/python && \ |
| 102 | + sed -i "s|/root/.local/share/uv/python|/opt/uv-python|g" /opt/nemo_rl_venv/pyvenv.cfg && \ |
| 103 | + mv /root/.local/bin /opt/uv-bin |
| 104 | + |
| 105 | +# --- Fix pre-built Ray venvs (same /root/ relocation) ----------------------- |
| 106 | +RUN for cfg in /opt/ray_venvs/*/pyvenv.cfg; do \ |
| 107 | + sed -i "s|/root/.local/share/uv/python|/opt/uv-python|g" "$cfg"; \ |
| 108 | + done && \ |
| 109 | + find /opt/ray_venvs/ -type l | while read link; do \ |
| 110 | + target=$(readlink "$link") && \ |
| 111 | + case "$target" in */root/.local/share/uv/python*) \ |
| 112 | + new_target=$(echo "$target" | sed "s|/root/.local/share/uv/python|/opt/uv-python|") && \ |
| 113 | + ln -sf "$new_target" "$link" ;; \ |
| 114 | + esac; \ |
| 115 | + done |
| 116 | + |
| 117 | +# --- Fix Gym venv (same /root/ relocation) ---------------------------------- |
| 118 | +RUN GYM_VENV=/opt/nemo-rl/3rdparty/Gym-workspace/Gym/.venv && \ |
| 119 | + sed -i "s|/root/.local/share/uv/python|/opt/uv-python|g" "$GYM_VENV/pyvenv.cfg" && \ |
| 120 | + find "$GYM_VENV" -type l | while read link; do \ |
| 121 | + target=$(readlink "$link") && \ |
| 122 | + case "$target" in */root/.local/share/uv/python*) \ |
| 123 | + new_target=$(echo "$target" | sed "s|/root/.local/share/uv/python|/opt/uv-python|") && \ |
| 124 | + ln -sf "$new_target" "$link" ;; \ |
| 125 | + esac; \ |
| 126 | + done |
| 127 | + |
| 128 | +# --- Runtime environment ----------------------------------------------------- |
| 129 | +ENV VIRTUAL_ENV=/opt/nemo_rl_venv |
| 130 | +ENV PATH=/opt/uv-bin:/opt/nemo_rl_venv/bin:$PATH |
| 131 | +ENV UV_PYTHON_INSTALL_DIR=/opt/uv-python |
0 commit comments