diff --git a/docker/Dockerfile.fw_final b/docker/Dockerfile.fw_final index f029622fcc..635dcca7a2 100644 --- a/docker/Dockerfile.fw_final +++ b/docker/Dockerfile.fw_final @@ -57,6 +57,11 @@ RUN pushd /opt/Export-Deploy && \ echo "Container built with NeMo-Run commit hash: $NEMO_RUN_COMMIT" && \ popd +# cuDNN Frontend is replaced after the sync below with this ref of NVIDIA/cudnn-frontend +# (feat/1.29.0-megatron). Docker-only: outside fw_pyproject.toml and uv.lock. Set empty to keep the +# PyPI wheel. +ARG CUDNN_FRONTEND_GIT_REF=8924d4b41fd678ba1b98cd7f98bb78b222637573 + # Sync environment COPY --chmod=644 docker/common/fw_pyproject.toml /opt/NeMo-FW/pyproject.toml COPY --chmod=644 docker/common/_version.py docker/common/print_sha.sh /opt/NeMo-FW/ @@ -69,6 +74,10 @@ RUN \ BASE_RUNTIME_UV_ARGS="--no-install-package nvidia-cudnn-frontend --no-install-package transformer-engine --no-install-package transformer-engine-torch"; \ fi && \ uv sync --no-cache-dir --all-groups --inexact ${BASE_RUNTIME_UV_ARGS} && \ + if [ -n "$CUDNN_FRONTEND_GIT_REF" ] && [ "${MBRIDGE_PRESERVE_BASE_RUNTIME:-False}" != "True" ]; then \ + uv pip install --no-cache-dir --no-deps --reinstall \ + "nvidia-cudnn-frontend @ git+https://github.com/NVIDIA/cudnn-frontend.git@${CUDNN_FRONTEND_GIT_REF}"; \ + fi && \ python /opt/preserve_base_runtime.py check /opt/base-runtime-packages.json && \ uv pip list && \ chmod +x print_sha.sh && \ diff --git a/docker/README.md b/docker/README.md index 3db8399c61..f69f61bbb1 100644 --- a/docker/README.md +++ b/docker/README.md @@ -250,6 +250,7 @@ docker build \ | `NEMO_EXPORT_DEPLOY_COMMIT` | NeMo Export-Deploy git commit SHA | | `NEMO_EVALUATOR_COMMIT` | NeMo Evaluator git commit SHA | | `NEMO_RUN_COMMIT` | NeMo Run git commit SHA | +| `CUDNN_FRONTEND_GIT_REF` | Commit of [NVIDIA/cudnn-frontend](https://github.com/NVIDIA/cudnn-frontend) (`feat/1.29.0-megatron`) installed over the PyPI cuDNN frontend after the environment sync; empty keeps the PyPI wheel. Docker-only, outside `fw_pyproject.toml`; skipped when `PRESERVE_BASE_RUNTIME=True`. | ---