From f21d23a21105bf2ef04c0a771220f881e1b4d361 Mon Sep 17 00:00:00 2001 From: Alison Shao <54658187+alisonshao@users.noreply.github.com> Date: Sat, 11 Apr 2026 23:32:51 -0500 Subject: [PATCH] ci: use local NVIDIA wheels to avoid re-downloading ~2GB every CI run (#22602) Co-authored-by: Alison Shao --- scripts/ci/cuda/cache_nvidia_wheels.sh | 20 +++++++++++++++++--- 1 file changed, 17 insertions(+), 3 deletions(-) diff --git a/scripts/ci/cuda/cache_nvidia_wheels.sh b/scripts/ci/cuda/cache_nvidia_wheels.sh index 2a0f8dbb9..9b3c49a96 100755 --- a/scripts/ci/cuda/cache_nvidia_wheels.sh +++ b/scripts/ci/cuda/cache_nvidia_wheels.sh @@ -2,15 +2,22 @@ # Cache and pre-install nvidia wheels that torch pins. # # pypi.nvidia.com returns Cache-Control: no-store, so pip re-downloads -# cudnn (~707 MB) and nvshmem (~125 MB) on every CI run. This script -# caches the wheels locally and installs them so that the subsequent -# `pip install -e "python[dev]"` sees "Requirement already satisfied". +# ~2 GB of NVIDIA wheels on every CI run. This script: +# 1. Caches cudnn + nvshmem wheels locally and pre-installs them +# 2. Points pip at a local wheel directory via PIP_FIND_LINKS so that +# all NVIDIA torch dependencies (cublas, cufft, nvrtc, etc.) are +# installed from local files instead of re-downloading. +# +# Pre-cache the wheels on the host at /opt/ci-cache/nvidia-pip-wheels/ +# (mounted as /root/.cache/nvidia-pip-wheels inside containers). +# See the 5090 ops guide post-reboot checklist for how to populate this. # # Integrity: uses `unzip -t` to detect partial/corrupt downloads. # # Usage: source scripts/ci/cuda/cache_nvidia_wheels.sh NVIDIA_WHEEL_CACHE="/root/.cache/nvidia-wheels" +NVIDIA_PIP_WHEELS="/root/.cache/nvidia-pip-wheels" mkdir -p "$NVIDIA_WHEEL_CACHE" for url in \ @@ -22,3 +29,10 @@ done pip install --no-deps "$NVIDIA_WHEEL_CACHE"/nvidia_cudnn_cu12-*.whl \ "$NVIDIA_WHEEL_CACHE"/nvidia_nvshmem_cu12-*.whl 2>/dev/null || true + +# If pre-cached NVIDIA pip wheels exist, tell pip to check there first. +# This avoids re-downloading ~2 GB of cublas/cufft/nvrtc/etc. every run +# (pypi.nvidia.com sends Cache-Control: no-store). +if [ -d "$NVIDIA_PIP_WHEELS" ] && ls "$NVIDIA_PIP_WHEELS"/*.whl &>/dev/null; then + export PIP_FIND_LINKS="${PIP_FIND_LINKS:+$PIP_FIND_LINKS }$NVIDIA_PIP_WHEELS" +fi