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

Signed-off-by: Hank Han <hanhan7630@outlook.com>

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 < <(
+    mapfile -t INSTALLED_NVIDIA_WHEELS < <(
-        python3 -m pip list --format=freeze | sed -n 's/^\(nvidia-.*-cu12\)==.*/\1/p'
+        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]}"
 }