diff --git a/scripts/ci/cuda/ci_install_dependency.sh b/scripts/ci/cuda/ci_install_dependency.sh index 9bd9969e3..fa9645fc5 100755 --- a/scripts/ci/cuda/ci_install_dependency.sh +++ b/scripts/ci/cuda/ci_install_dependency.sh @@ -197,14 +197,27 @@ setup_pip_toolchain() { } remove_stale_cuda12_nvidia_wheels() { + local package_name spec + local -a INSTALLED_NVIDIA_WHEELS=() + local -a NVIDIA_WHEELS_TO_RESTORE=() + local -a STALE_CUDA12_NVIDIA_WHEELS=() + if [ "$CU_MAJOR" != "13" ]; then mark_step_done "${FUNCNAME[0]}" return fi - mapfile -t STALE_CUDA12_NVIDIA_WHEELS < <( - python3 -m pip list --format=freeze | sed -n 's/^\(nvidia-.*-cu12\)==.*/\1/p' + mapfile -t INSTALLED_NVIDIA_WHEELS < <( + python3 -m pip list --format=freeze | sed -n '/^nvidia-.*==/p' ) + for spec in "${INSTALLED_NVIDIA_WHEELS[@]}"; do + package_name="${spec%%==*}" + case "$package_name" in + *-cu12) STALE_CUDA12_NVIDIA_WHEELS+=("$package_name") ;; + *) NVIDIA_WHEELS_TO_RESTORE+=("$spec") ;; + esac + done + if [ ${#STALE_CUDA12_NVIDIA_WHEELS[@]} -eq 0 ]; then echo "No stale CUDA 12 NVIDIA wheels found for ${CU_VERSION} job" mark_step_done "${FUNCNAME[0]}" @@ -214,6 +227,16 @@ remove_stale_cuda12_nvidia_wheels() { echo "Removing stale CUDA 12 NVIDIA wheels from ${CU_VERSION} job: ${STALE_CUDA12_NVIDIA_WHEELS[*]}" $PIP_UNINSTALL_CMD "${STALE_CUDA12_NVIDIA_WHEELS[@]}" $PIP_UNINSTALL_SUFFIX + # CUDA 12 and CUDA 13 wheels can own the same nvidia/* paths. Uninstalling + # the stale variant deletes those shared files even though the remaining + # wheel metadata still says they are installed. Restore every remaining + # NVIDIA wheel at its already-installed version to make the transition + # atomic and avoid package-specific payload checks. + if [ ${#NVIDIA_WHEELS_TO_RESTORE[@]} -gt 0 ]; then + echo "Restoring NVIDIA wheels after CUDA 12 cleanup: ${NVIDIA_WHEELS_TO_RESTORE[*]}" + $PIP_CMD install --force-reinstall --no-deps "${NVIDIA_WHEELS_TO_RESTORE[@]}" $PIP_INSTALL_SUFFIX + fi + mark_step_done "${FUNCNAME[0]}" }