# 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 .
#
#   Non-root runtime override (the image still defaults to root):
#     docker run --rm --user 65532:65532 \
#       --env HOME=/home/nemo-runtime \
#       --env XDG_CACHE_HOME=/opt/nemo-gym/cache/xdg \
#       <registry>/nemo-gym:latest --help
#

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

ARG ENROOT_VERSION=3.5.0
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 \
    bash \
    jq \
    curl \
    git \
    wget \
    less \
    vim \
    gnupg

apt-get purge -y \
    ffmpeg \
    libavcodec* \
    libavformat* \
    libavutil* \
    libswscale* \
    libswresample* \
    libx264* \
    libx265* \
    libfdk-aac* \
    libmp3lame* \
    2>/dev/null || true

# Gym never imports dbus or GObject; drop the base image's Python bindings.
apt-get purge -y python3-dbus python3-gi

apt-get autoremove -y

# CVE-2026-45447: PKCS#7/S-MIME use-after-free in OpenSSL's PKCS7_verify().
apt-get install -y --only-upgrade openssl libssl3t64

# Install enroot for the EnrootProvider sandbox backend.
# GitHub releases provide signed .debs for amd64 and arm64.
arch=$(dpkg --print-architecture)
curl -fSsL -o /tmp/enroot.deb \
    "https://github.com/NVIDIA/enroot/releases/download/v${ENROOT_VERSION}/enroot_${ENROOT_VERSION}-1_${arch}.deb"
apt-get install -y /tmp/enroot.deb
rm /tmp/enroot.deb

# Drop the --ldconfig arg from enroot's NVIDIA hook so nvidia-container-cli
# reuses the sandbox's existing ld.so.cache instead of regenerating it.
sed -i 's/cli_args=("--no-cgroups" "--ldconfig=@$(command -v ldconfig.real || command -v ldconfig)")/cli_args=("--no-cgroups")/' \
    /etc/enroot/hooks.d/98-nvidia.sh
grep -q '^cli_args=("--no-cgroups")$' /etc/enroot/hooks.d/98-nvidia.sh

# Install nvidia-container-toolkit so enroot's NVIDIA hook
# (/etc/enroot/hooks.d/98-nvidia.sh) can find nvidia-container-cli to inject
# GPUs into enroot sandboxes. Without this, enroot sandboxes that request a
# GPU fail at the hook step even though the outer container has GPU access.
curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey \
    | gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg
curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list \
    | sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' \
    > /etc/apt/sources.list.d/nvidia-container-toolkit.list
apt-get update
apt-get install -y nvidia-container-toolkit
rm -f /etc/apt/sources.list.d/nvidia-container-toolkit.list /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg

# Remove the CUDA base image's bundled Nsight Systems. Gym does not profile
# with nsys, and the bundled copy ships prebuilt Go binaries flagged by scans.
rm -rf /usr/local/cuda*/NsightSystems-cli-* /opt/nvidia/nsight-systems*
apt-get remove --purge -y --allow-change-held-packages 'nsight-systems*' || true
rm -f /usr/local/cuda/bin/nsys /usr/local/bin/nsys

apt-get clean
rm -rf /var/lib/apt/lists/*
EOF

ARG UV_VERSION=0.11.29
ARG PYTHON_VERSION=3.13.14
ARG TARGETARCH
ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
ENV UV_CACHE_DIR=/opt/uv/cache
RUN case "${TARGETARCH}" in \
        amd64) \
            uv_arch="x86_64"; \
# pragma: allowlist nextline secret
            uv_sha256="04f8b82f5d47f0512dcd32c67a4a6f16a0ea27c81537c338fd0ad6b23cebe829"; \
            ;; \
        arm64) \
            uv_arch="aarch64"; \
# pragma: allowlist nextline secret
            uv_sha256="94500fb064ae3c971a873cba64d94694c50677e0a4dbf78735c80509e7429919"; \
            ;; \
        *) \
            echo "Unsupported TARGETARCH for uv: ${TARGETARCH}" >&2; \
            exit 1; \
            ;; \
    esac && \
    uv_archive="uv-${uv_arch}-unknown-linux-gnu.tar.gz" && \
    curl -fLSs -o "/tmp/${uv_archive}" \
        "https://github.com/astral-sh/uv/releases/download/${UV_VERSION}/${uv_archive}" && \
    echo "${uv_sha256}  /tmp/${uv_archive}" | sha256sum -c - && \
    tar -xzf "/tmp/${uv_archive}" -C /tmp && \
    install -m 0755 "/tmp/uv-${uv_arch}-unknown-linux-gnu/uv" /usr/local/bin/uv && \
    install -m 0755 "/tmp/uv-${uv_arch}-unknown-linux-gnu/uvx" /usr/local/bin/uvx && \
    rm -rf "/tmp/${uv_archive}" "/tmp/uv-${uv_arch}-unknown-linux-gnu" && \
    /usr/local/bin/uv python install ${PYTHON_VERSION}
ENV PATH="/usr/local/bin:$PATH"

ARG RUNTIME_UID=65532
ARG RUNTIME_GID=65532
RUN if ! getent group "${RUNTIME_GID}" >/dev/null; then \
        groupadd --gid "${RUNTIME_GID}" nemo-runtime; \
    fi && \
    if ! getent passwd "${RUNTIME_UID}" >/dev/null; then \
        useradd --no-log-init --uid "${RUNTIME_UID}" --gid "${RUNTIME_GID}" \
            --create-home --home-dir /home/nemo-runtime --shell /bin/bash nemo-runtime; \
    fi && \
    install -d -o "${RUNTIME_UID}" -g "${RUNTIME_GID}" \
        /home/nemo-runtime \
        /opt/nemo-gym \
        /opt/nemo-gym/cache \
        /opt/nemo-gym/results \
        /opt/uv/cache

ENV RAY_USAGE_STATS_ENABLED=0


FROM base AS hermetic

WORKDIR /opt/nemo-gym

ARG BASE_IMAGE
ARG UV_VERSION
ARG RUNTIME_UID
ARG RUNTIME_GID

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 --chown=${RUNTIME_UID}:${RUNTIME_GID} pyproject.toml uv.lock ./
COPY --from=nemo-gym --chown=${RUNTIME_UID}:${RUNTIME_GID} \
    nemo_gym/__init__.py nemo_gym/package_info.py ./nemo_gym/

# The telemetry extra is included even though telemetry is off by default. It is a few MB
# of pure-Python packages, and without it in the image nobody can switch tracing on with
# an environment variable — they would have to rebuild. Off remains the default: the code
# is inert until NEMO_LENS_ENABLED=1.
RUN uv venv --seed && \
    uv sync --link-mode symlink --locked --extra vllm --extra telemetry --extra dev --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 while the package build introspects the preserved Git
# checkout. Exclude pyproject.toml and uv.lock since they are already present
# from the hermetic stage.
#
# --chown matches the hermetic-stage copies above and makes the source readable
# by the runtime UID regardless of the build context's file modes. Without it
# the image only works when the context happens to be world-readable: a builder
# with a restrictive umask (NeMo CI uses 0007) yields mode 0660 root-owned
# files, and the later chown pass only covers directories, so `gym --help`
# fails with PermissionError on nemo_gym/__init__.py once USER is switched.
COPY --from=nemo-gym --chown=${RUNTIME_UID}:${RUNTIME_GID} \
    --exclude=pyproject.toml --exclude=uv.lock . /opt/nemo-gym

# Install the nemo-gym project itself along with the vllm and telemetry extras. The
# worktree was created for the runtime UID in the base stage, so restore build ownership
# before setuptools invokes Git.
RUN chown root:root /opt/nemo-gym && \
    UV_LINK_MODE=symlink uv sync --locked --extra vllm --extra telemetry --extra dev

# Seed the runtime uv cache with the full dependency set (including the dev
# extra and the project itself) so that CI scripts can run offline: setup_dev.sh
# resolves `uv sync --extra dev --offline` from this cache in a fresh venv.
# The runtime cache is separate from the root-owned build cache (/opt/uv/cache)
# that backs the symlinked project environment.
ENV UV_CACHE_DIR=/opt/nemo-gym/cache/uv
RUN uv venv /opt/nemo-gym/cache/seed-venv && \
    UV_PROJECT_ENVIRONMENT=/opt/nemo-gym/cache/seed-venv \
        uv sync --locked --extra vllm --extra telemetry --extra dev && \
    rm -rf /opt/nemo-gym/cache/seed-venv && \
    chown -R "${RUNTIME_UID}:${RUNTIME_GID}" /opt/nemo-gym/cache/uv

# Optional: seed the same cache with every server module's venv so that
# scripts/ci/server_tests.sh can run offline too. Pass
# --build-arg NEMO_GYM_SEED_SERVER_ENVS=1 for an image that runs the server
# suite without a package index. Each server installs its own unlocked
# requirements.txt or pyproject.toml (including git sources) outside the root
# lock, so the seed above does not cover them, and together they are large
# (about 25 GB more cache for all 212 servers in October 2026), hence opt-in.
# `ng_test_all +setup_only=true` installs each venv exactly as the server suite
# does (same discovery, requirements, and head-server pins), outside the source
# tree, and deletes it; only the downloaded and built distributions stay.
ARG NEMO_GYM_SEED_SERVER_ENVS=""
RUN <<"EOF" bash -exu
if [[ -n "${NEMO_GYM_SEED_SERVER_ENVS:-}" ]]; then
    seed_venv_root="$(mktemp -d)"
    ng_test_all +setup_only=true +delete_venvs_after_each_test=true \
        "+uv_cache_dir=${UV_CACHE_DIR}" "+uv_venv_dir=${seed_venv_root}"
    rm -rf "${seed_venv_root}"
    chown -R "${RUNTIME_UID}:${RUNTIME_GID}" /opt/nemo-gym/cache/uv
fi
EOF

# Seed the pre-commit hook cache (PRE_COMMIT_HOME) so the image's offline lint
# path needs no network: lint.sh reuses the baked pre-commit executable, but a
# bare `pre-commit run` still clones the hook repositories declared in
# .pre-commit-config.yaml and pip-installs the local hooks' dependencies on
# first run. Resolving every hook environment here, while the build has
# network, keeps the runtime lint hermetic.
ENV PRE_COMMIT_HOME=/opt/nemo-gym/cache/pre-commit
# Local build contexts exclude .git, but pre-commit requires a repository to
# resolve the configuration. Use disposable metadata while seeding the cache.
RUN git init --quiet && \
    pre-commit install-hooks && \
    chown -R "${RUNTIME_UID}:${RUNTIME_GID}" "${PRE_COMMIT_HOME}" && \
    rm -rf .git
RUN install -d -o "${RUNTIME_UID}" -g "${RUNTIME_GID}" \
        "${UV_CACHE_DIR}" \
        /opt/nemo-gym/cache/xdg

# 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

# ray's bundled Java worker jar ships an outdated httpcore5 with a known CVE
# (GHSA-hf6x-8p5f-cgmf). Gym never uses Ray's Java cross-language worker, and
# removing the jar doesn't affect ray.init(), remote tasks, actors, or
# ray.serve (verified locally). Covers both installed venvs (site-packages)
# and uv's own cache (archive-v0): confirmed uv silently skips a missing
# cached file and installs everything else fine, even fully offline, so
# purging it from the cache too doesn't break installing a fresh venv later
# for a server not prefetched here.
find /opt -name ray_dist.jar -delete

# Keep ownership changes for generated caches and server venvs in the layer
# that creates them so later layers do not duplicate their contents.
chown -R "${RUNTIME_UID}:${RUNTIME_GID}" /opt/nemo-gym/cache
find /opt/nemo-gym -type d -name .venv -prune \
    -exec chown -R "${RUNTIME_UID}:${RUNTIME_GID}" {} +
EOF

# Own source and project-environment directories so the runtime identity can
# mutate their contents without copying prefetched artifacts into this layer.
RUN venv_site_packages="$(/opt/nemo_gym_venv/bin/python -c \
        'import site; print(site.getsitepackages()[0])')" && \
    touch "${venv_site_packages}/.nonroot-uninstall-probe" && \
    find /opt/nemo-gym \
        -path /opt/nemo-gym/cache -prune -o \
        -type d -name .venv -prune -o \
        -type d -exec chown "${RUNTIME_UID}:${RUNTIME_GID}" {} + && \
    find /opt/nemo_gym_venv -type d \
        -exec chown "${RUNTIME_UID}:${RUNTIME_GID}" {} +

USER ${RUNTIME_UID}:${RUNTIME_GID}

# Runtime callers select this UID and set HOME/cache variables explicitly; the
# published image remains root by default. Prove that the CLI, uv toolchain,
# project environment, source tree, and runtime paths support that contract.
RUN export HOME=/home/nemo-runtime \
        XDG_CACHE_HOME=/opt/nemo-gym/cache/xdg && \
    test "$(id -u)" = "${RUNTIME_UID}" && \
    test -w "${HOME}" && \
    test -w "${XDG_CACHE_HOME}" && \
    test -w "${UV_CACHE_DIR}" && \
    test -w "${PRE_COMMIT_HOME}" && \
    test -w /opt/nemo-gym/cache && \
    test -w /opt/nemo-gym/results && \
    test -x /usr/local/bin/uv && \
    test -x /usr/local/bin/uvx && \
    test -x /opt/nemo_gym_venv/bin/python && \
    venv_site_packages="$(python -c 'import site; print(site.getsitepackages()[0])')" && \
    rm "${venv_site_packages}/.nonroot-uninstall-probe" && \
    touch "${HOME}/.nonroot-write-probe" \
        "${XDG_CACHE_HOME}/.nonroot-write-probe" \
        "${UV_CACHE_DIR}/.nonroot-write-probe" \
        /opt/nemo-gym/.nonroot-write-probe \
        /opt/nemo-gym/cache/.nonroot-write-probe \
        /opt/nemo-gym/results/.nonroot-write-probe \
        "${venv_site_packages}/.nonroot-write-probe" && \
    gym --help >/dev/null && \
    uv venv /opt/nemo-gym/cache/container-smoke && \
    /opt/nemo-gym/cache/container-smoke/bin/python -c \
        "import sys; assert sys.version_info[:2] == (3, 13)" && \
    rm -r /opt/nemo-gym/cache/container-smoke \
        "${HOME}/.nonroot-write-probe" \
        "${XDG_CACHE_HOME}/.nonroot-write-probe" \
        "${UV_CACHE_DIR}/.nonroot-write-probe" \
        /opt/nemo-gym/.nonroot-write-probe \
        /opt/nemo-gym/cache/.nonroot-write-probe \
        /opt/nemo-gym/results/.nonroot-write-probe \
        "${venv_site_packages}/.nonroot-write-probe"

USER root
ENTRYPOINT ["gym"]
CMD ["--help"]
