Skip to content

chore(customizer): fix setuptools CVE in the RL image + docs cleanup #1187

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
anubhutivyas merged 1 commit into from
Aug 11, 2026
Merged

chore(customizer): fix setuptools CVE in the RL image + docs cleanup #1187

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
38 changes: 0 additions & 38 deletions docker-bake.hcl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -52,10 +52,6 @@ variable "FFMPEG_VLM_WHEEL_CONTEXT" {
default = ""
}

variable "TRANSFORMER_ENGINE_WHEEL_CONTEXT" {
default = ""
}

variable "DISTROLESS_BASE" {
default = "nvcr.io/nvidia/distroless/python:3.11-v4.0.8"
}
Expand Down Expand Up @@ -171,11 +167,6 @@ function "get_ffmpeg_vlm_wheel_image" {
result = "${WHEELS_REGISTRY}/ffmpeg-vlm-wheel:${WHEELS_TAG}"
}

function "get_transformer_engine_wheel_image" {
params = []
result = "${WHEELS_REGISTRY}/transformer-engine-wheel:${WHEELS_TAG}"
}

function "get_arch_tag" {
params = []
result = BUILD_ARCH == "linux/arm64" ? "linux-arm64" : "linux-amd64"
Expand Down Expand Up @@ -213,11 +204,6 @@ function "ffmpeg_vlm_wheel_context" {
result = notequal(FFMPEG_VLM_WHEEL_CONTEXT, "") ? FFMPEG_VLM_WHEEL_CONTEXT : notequal(USE_LOCAL_WHEELS, "") ? "target:ffmpeg-vlm-wheel" : "docker-image://${get_ffmpeg_vlm_wheel_image()}"
}

function "transformer_engine_wheel_context" {
params = []
result = notequal(TRANSFORMER_ENGINE_WHEEL_CONTEXT, "") ? TRANSFORMER_ENGINE_WHEEL_CONTEXT : notequal(USE_LOCAL_WHEELS, "") ? "target:transformer-engine-wheel" : "docker-image://${get_transformer_engine_wheel_image()}"
}

function "wheel_tags" {
params = [name]
result = [
Expand Down Expand Up @@ -704,30 +690,6 @@ target "ffmpeg-vlm-wheel" {
platforms = get_platforms()
}

# transformer-engine wheel (cp313 / cu130), pinned to NeMo-RL's TE ref (release_v2.15, from RL's
# override-dependencies) and arch (90;100, matching RL).
#
# NOT consumed by any image today, and deliberately not in a build group: Transformer-Engine only
# ships in the `automodel` and `mcore` extras, and the customizer's DPO + GRPO path uses neither
# (DPO -> fsdp, GRPO generation -> vllm). Kept because it builds cleanly and is the drop-in for
# nmp-rl-base if a Megatron/automodel backend is ever adopted. Build on demand:
# docker buildx bake transformer-engine-wheel
target "transformer-engine-wheel" {
target = "transformer-engine-wheel"
context = "."
dockerfile = "docker/base/Dockerfile.python-wheels"
cache-to = maybe_registry_cache_to("transformer-engine-wheel")
cache-from = maybe_registry_cache_from("transformer-engine-wheel")
tags = wheel_tags("transformer-engine-wheel")
output = image_output()
args = {
TE_REF = "release_v2.15"
NVTE_CUDA_ARCHS = "90;100"
}
platforms = get_platforms()
}


target "safe-synthesizer-tasks-docker" {
target = "runtime"
context = "."
Expand Down
97 changes: 0 additions & 97 deletions docker/base/Dockerfile.python-wheels
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
# Builds Python wheels for:
# - causal-conv1d (CUDA extension required by mamba-ssm)
# - mamba-ssm (selective state space model)
# - transformer-engine (cp313 / cu130, for nmp-rl-base — the longest source compile)
# - av + opencv-python-headless + decord2 (FFmpeg 8.1.2 bundled wheels for nspect / VLM)
#
# causal-conv1d / mamba-ssm only ship source on PyPI and require nvcc.
Expand Down Expand Up @@ -278,102 +277,6 @@ COPY --from=mamba-ssm-23-wheel-builder-py312 /wheels /wheels
COPY --from=mamba-ssm-23-wheel-builder-py312-cu13.1.1 /wheels /wheels/cu13.1.1
COPY --from=mamba-ssm-23-wheel-builder-py312-cu13.2 /wheels /wheels/cu13.2

# =============================================================================
# transformer-engine wheel — Python 3.13 / CUDA 13.2
#
# NOT CONSUMED BY ANY IMAGE TODAY, and intentionally so. Transformer-Engine ships only in
# NeMo-RL's `automodel` and `mcore` extras, and the customizer's DPO + GRPO path uses
# neither (DPO/GRPO train on `fsdp`; GRPO generates on `vllm`). nmp-rl-base therefore does
# not build TE at all. This stage is kept ready — TE is the longest CUDA compile in that
# image, so if the Megatron/automodel backend is ever enabled, wiring this wheel in
# (`uv pip install --no-deps`, plus --no-install-package on the syncs) avoids recompiling
# it on every base rebuild.
#
# ABI must match the consuming image EXACTLY or the wheel won't import: same base image
# (cuda-dl-base 26.03-cuda13.2), cp313, and the torch + TE ref NeMo-RL pins (torch 2.11.0
# cu130; TE release_v2.15 from override-dependencies). Bumping RL's TE ref means rebuilding
# this wheel.
#
# Verified to build and produce a single self-contained wheel
# (`transformer_engine-<ver>-cp313-cp313-linux_x86_64.whl`) — the pytorch extension
# (`transformer_engine_torch*.so`) is bundled inside it, there is no separate distribution.
# Build on demand with:
# docker buildx bake transformer-engine-wheel
# =============================================================================
FROM nvcr.io/nvidia/cuda-dl-base:26.03-cuda13.2-devel-ubuntu24.04 AS te-wheel-base-py313

ARG PYTHON_VERSION=3.13.13
ARG CMAKE_VERSION=4.0.3
# Mirror NeMo-RL's pins (do not drift): torch 2.11.0 on the cu130 index; cuDNN wheel
# versions from RL's override-dependencies.
ARG TORCH_VERSION=2.11.0

ENV DEBIAN_FRONTEND=noninteractive
RUN apt-get update && apt-get install -y --no-install-recommends \
build-essential ca-certificates git curl ninja-build libopenmpi-dev && \
rm -rf /var/lib/apt/lists/*

RUN ARCH="$(uname -m)" && F="cmake-${CMAKE_VERSION}-linux-${ARCH}" && \
curl --retry 3 --retry-delay 2 -fsSL -o "${F}.tar.gz" \
"https://github.com/Kitware/CMake/releases/download/v${CMAKE_VERSION}/${F}.tar.gz" && \
tar -xzf "${F}.tar.gz" && cp -r "${F}/bin/"* /usr/local/bin/ && \
cp -r "${F}/share/"* /usr/local/share/ && rm -rf "${F}" "${F}.tar.gz"

# uv 0.11.18 (matches the nmp-rl base). Older uv (0.9.14, used by the mamba stages for
# their lock's required-version) has no python-build-standalone entry for 3.13.13, so
# `uv python install 3.13.13` fails with "No download found".
COPY --from=ghcr.io/astral-sh/uv:0.11.18 /uv /usr/local/bin/uv
ENV UV_PROJECT_ENVIRONMENT=/opt/venv UV_LINK_MODE=copy
# cuda-dl-base / Ubuntu 24.04 ships no Python 3.13 — have uv provision the managed
# interpreter before creating the venv (same as the nmp-rl base does).
RUN uv python install ${PYTHON_VERSION} && uv venv --python ${PYTHON_VERSION} /opt/venv
ENV PATH="/opt/venv/bin:$PATH" VIRTUAL_ENV=/opt/venv

# Build environment: torch is the ABI anchor (match nmp-rl-base); TE also links cuDNN
# and (for some kernels) mathdx. Versions mirror NeMo-RL's pyproject.
RUN uv pip install --torch-backend=cu130 "torch==${TORCH_VERSION}" && \
uv pip install setuptools wheel packaging pybind11 ninja \
"nvidia-cudnn-cu13==9.20.0.48" \
"nvidia-cudnn-frontend==1.23.0" \
"nvidia-mathdx==25.1.1"

ENV CUDA_HOME=/usr/local/cuda \
CPLUS_INCLUDE_PATH=/usr/local/cuda/include/cccl \
CUDNN_HOME=/opt/venv/lib/python3.13/site-packages/nvidia/cudnn \
CUDNN_PATH=/opt/venv/lib/python3.13/site-packages/nvidia/cudnn \
LD_LIBRARY_PATH="/opt/venv/lib/python3.13/site-packages/nvidia/cudnn/lib:${LD_LIBRARY_PATH}"

FROM te-wheel-base-py313 AS transformer-engine-wheel-builder

# NeMo-RL pins TE via override-dependencies (@release_v2.15 wins across all extras).
ARG TE_REF=release_v2.15
# Includes 80 (A100) on purpose. Transformer-Engine is NOT Hopper-only - upstream Automodel
# builds it for "80;90;100;120". Only deep_ep / deep_gemm have a hard SM90 floor, and neither is
# part of TE. Building 80;90;100 here means the wheel works on A100 as well as Hopper/Blackwell,
# so the Megatron (mcore) / automodel backends could run on A100 if they are ever enabled -
# without having to rediscover that TE was the blocker. Cost is a longer compile (one extra arch).
# NeMo-RL itself pins 90;100; docker-bake.hcl currently overrides this arg back to 90;100 to keep
# today's behavior byte-identical. Drop that override to get an A100-capable wheel.
ARG NVTE_CUDA_ARCHS=80;90;100
ARG NVTE_BUILD_THREADS_PER_JOB=8

RUN mkdir -p /wheels && \
git clone https://github.com/NVIDIA/TransformerEngine.git /src/TransformerEngine && \
cd /src/TransformerEngine && \
git fetch origin ${TE_REF} && git checkout FETCH_HEAD && \
git submodule update --init --recursive && \
NVTE_FRAMEWORK=pytorch NVTE_CUDA_ARCHS="${NVTE_CUDA_ARCHS}" \
NVTE_BUILD_THREADS_PER_JOB=${NVTE_BUILD_THREADS_PER_JOB} \
uv build --wheel --no-build-isolation --out-dir /wheels . && \
rm -rf /src/TransformerEngine

# The final transformer-engine-wheel image contains (cp313 / cu130):
# - transformer_engine-*-cp313-*.whl
# - transformer_engine_torch-*-cp313-*.whl (if TE emits it separately — verify)
# Consumers install with `uv pip install --no-deps /wheels/*.whl`.
FROM scratch AS transformer-engine-wheel
COPY --from=transformer-engine-wheel-builder /wheels /wheels

# =============================================================================
# FFmpeg / VLM wheels (manylinux_2_28) — av + opencv-python-headless + decord2
# =============================================================================
Expand Down
Loading
Loading