From 1510c916955700c3ab5819a39b5f4d2be2b455bd Mon Sep 17 00:00:00 2001
From: Graham Markall <gmarkall@nvidia.com>
Date: Fri, 2 May 2025 18:09:52 +0100
Subject: [PATCH 1/2] Fix package installation for wheels CI

Normal wheel tests should not install any CUDA pip packages because they
should use the toolkit from the container.

The test depending on the toolkit wheels should install the specific
toolkit version used for the test (as opposed to just using the latest
version, which they were set up to do before).
---
 ci/test_wheel.sh             |  2 --
 ci/test_wheel_deps_wheels.sh | 16 +++++++++++-----
 ci/test_wheel_pynvjitlink.sh |  1 -
 3 files changed, 11 insertions(+), 8 deletions(-)

diff --git a/ci/test_wheel.sh b/ci/test_wheel.sh
index bee1d2ec5..071c3ea8e 100755
--- a/ci/test_wheel.sh
+++ b/ci/test_wheel.sh
@@ -9,8 +9,6 @@ python -m pip install \
     psutil \
     cffi \
     cuda-python \
-    nvidia-cuda-cccl-cu12 \
-    nvidia-curand-cu12 \
     pytest
 
 rapids-logger "Install wheel"
diff --git a/ci/test_wheel_deps_wheels.sh b/ci/test_wheel_deps_wheels.sh
index c8c135b49..2518e7b07 100755
--- a/ci/test_wheel_deps_wheels.sh
+++ b/ci/test_wheel_deps_wheels.sh
@@ -3,16 +3,22 @@
 
 set -euo pipefail
 
+# cuRAND versions don't follow the toolkit versions - map toolkit versions to
+# appropriate cuRAND versions
+declare -A CTK_CURAND_VMAP=( ["12.8"]="10.3.9" ["12.9"]="10.3.10")
+CUDA_VER_MAJOR_MINOR=${CUDA_VER%.*}
+CURAND_VER="${CTK_CURAND_VMAP[${CUDA_VER_MAJOR_MINOR}]}"
+
 rapids-logger "Install testing dependencies"
 # TODO: Replace with rapids-dependency-file-generator
 python -m pip install \
     psutil \
     cffi \
-    cuda-python \
-    nvidia-cuda-runtime-cu12 \
-    nvidia-curand-cu12 \
-    nvidia-cuda-nvcc-cu12 \
-    nvidia-cuda-nvrtc-cu12 \
+    "cuda-python==${CUDA_VER_MAJOR_MINOR}.*" \
+    "nvidia-cuda-runtime-cu12==${CUDA_VER_MAJOR_MINOR}.*" \
+    "nvidia-curand-cu12==${CURAND_VER}.*" \
+    "nvidia-cuda-nvcc-cu12==${CUDA_VER_MAJOR_MINOR}.*" \
+    "nvidia-cuda-nvrtc-cu12==${CUDA_VER_MAJOR_MINOR}.*" \
     pynvjitlink-cu12 \
     pytest
 
diff --git a/ci/test_wheel_pynvjitlink.sh b/ci/test_wheel_pynvjitlink.sh
index 55110234c..ab586b6b1 100755
--- a/ci/test_wheel_pynvjitlink.sh
+++ b/ci/test_wheel_pynvjitlink.sh
@@ -9,7 +9,6 @@ python -m pip install \
     psutil \
     cffi \
     cuda-python \
-    nvidia-curand-cu12 \
     pytest
 
 rapids-logger "Install pynvjitlink"

From 40c26ffabdb07f04f98f16a7f06653792018e740 Mon Sep 17 00:00:00 2001
From: Graham Markall <gmarkall@nvidia.com>
Date: Sat, 3 May 2025 11:31:59 +0100
Subject: [PATCH 2/2] Add CUDA_VER to wheels test script environment

---
 .github/workflows/wheels-test.yaml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/wheels-test.yaml b/.github/workflows/wheels-test.yaml
index b5d0ea7e2..b4a2a3a25 100644
--- a/.github/workflows/wheels-test.yaml
+++ b/.github/workflows/wheels-test.yaml
@@ -142,6 +142,7 @@ jobs:
     - name: Run tests
       run: ${{ inputs.script }}
       env:
+        CUDA_VER: ${{ matrix.CUDA_VER }}
         GH_TOKEN: ${{ github.token }}
         RAPIDS_AUX_SECRET_1: ${{ secrets.RAPIDS_AUX_SECRET_1 }}