diff --git a/.github/container/Dockerfile.base b/.github/container/Dockerfile.base index e4426aa00..ba97a3638 100644 --- a/.github/container/Dockerfile.base +++ b/.github/container/Dockerfile.base @@ -1,5 +1,5 @@ # syntax=docker/dockerfile:1-labs -ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.04 +ARG BASE_IMAGE=nvcr.io/nvidia/cuda-dl-base:26.06-cuda13.3-devel-ubuntu24.04 ARG GIT_USER_NAME="JAX Toolbox" ARG GIT_USER_EMAIL=jax@nvidia.com ARG JAX_TOOLBOX_REF=main @@ -133,3 +133,7 @@ ENV MANIFEST_FILE="/opt/manifest.d/manifest.yaml" ADD manifest.yaml create-distribution.sh bump.sh /opt/manifest.d/ COPY patches/ /opt/manifest.d/patches/ + +# Temporary workaround for TransformerEngine, which added a >=1.25.0 requirement +# when no public cuda-dl-base container had that CUDNN_FRONTEND_VERSION. +ENV CUDNN_FRONTEND_VERSION=1.25.0 diff --git a/.github/container/Dockerfile.jax b/.github/container/Dockerfile.jax index 4ff927c5a..1ec8709f8 100644 --- a/.github/container/Dockerfile.jax +++ b/.github/container/Dockerfile.jax @@ -132,6 +132,7 @@ COPY --from=builder ${SRC_PATH_TRANSFORMER_ENGINE} ${SRC_PATH_TRANSFORMER_ENGINE RUN <<"EOF" bash -ex ls ${SRC_PATH_TRANSFORMER_ENGINE}/dist/*.whl echo "transformer-engine @ file://$(ls ${SRC_PATH_TRANSFORMER_ENGINE}/dist/*.whl)" > /opt/pip-tools.d/requirements-te.in +echo "nvidia-cudnn-frontend==${CUDNN_FRONTEND_VERSION}" >> /opt/pip-tools.d/requirements-te.in EOF ## nvidia-cutlass-dsl[cu13] diff --git a/.github/container/build-te.sh b/.github/container/build-te.sh index a6b78c2b5..66d3349be 100755 --- a/.github/container/build-te.sh +++ b/.github/container/build-te.sh @@ -138,13 +138,15 @@ pushd ${SRC_PATH_TE} # JAX, or the wheel-based installation of CUDA. Note that when we build TE as # part of building the JAX containers, JAX and XLA are not yet installed. python - << EOF -import subprocess, sys, tomllib +import os, subprocess, sys, tomllib with open("pyproject.toml", "rb") as ifile: data = tomllib.load(ifile) subprocess.run( [sys.executable, "-m", "pip", "install"] + [r for r in data["build-system"]["requires"] - if r.startswith("pybind11") or r.startswith("cmake") or r.startswith("ninja")]) + if r.startswith("pybind11") or r.startswith("cmake") or r.startswith("ninja")] + + [f"nvidia-cudnn-frontend=={os.environ['CUDNN_FRONTEND_VERSION']}"] +) EOF if [[ "${CCACHE}" == "1" ]]; then # Install ccache if not present (needs >= 4.1 for Redis remote storage support) diff --git a/docs/reference/container-versions.md b/docs/reference/container-versions.md index 27f9fb5fc..edef5efb3 100644 --- a/docs/reference/container-versions.md +++ b/docs/reference/container-versions.md @@ -8,6 +8,8 @@ The table below shows the first nightly that shipped with each base container. | First nightly with new base container | Base container | | ------------------------------------- | -------------- | +| 2026-07-08 | `nvcr.io/nvidia/cuda-dl-base:26.06-cuda13.3-devel-ubuntu24.04` | +| 2026-06-26 | `nvcr.io/nvidia/cuda-dl-base:26.05-cuda13.2-devel-ubuntu24.04` | | 2026-05-01 | `nvcr.io/nvidia/cuda-dl-base:26.04-cuda13.2-devel-ubuntu24.04` | | 2026-03-12 | `nvcr.io/nvidia/cuda-dl-base:26.02-cuda13.1-devel-ubuntu24.04` | | 2025-12-17 | `nvcr.io/nvidia/cuda-dl-base:25.11-cuda13.0-devel-ubuntu24.04` |