[CI] Fix CUDA 12 NVIDIA wheel cleanup (#31035)

Signed-off-by: Hank Han <hanhan7630@outlook.com>
This commit is contained in:
Hank Han
2026-07-15 10:06:11 +08:00
committed by GitHub
parent 76dc427806
commit b22f20b660
+25 -2
View File
@@ -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]}"
}