add a new ci job for testing with pynvjitlink

NVIDIA · Sep 30, 2024 · c97767c · c97767c
1 parent 6327ec2
commit c97767c
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 4 deletions.
diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml
@@ -54,9 +54,20 @@ jobs:
     uses: ./.github/workflows/conda-python-tests.yaml
     with:
       build_type: pull-request
-      script: "ci/test_conda.sh"
+      script: "ci/test_conda.sh false"
       run_codecov: false
       matrix_filter: ${{ needs.compute-matrix.outputs.TEST_MATRIX }}
+  test-conda-pynvjitlink:
+    needs:
+      - build-conda
+      - compute-matrix
+    uses: ./.github/workflows/conda-python-tests.yaml
+    with:
+      build_type: pull-request
+      script: "ci/test_conda.sh true"
+      run_codecov: false
+      # This selects "ARCH=amd64 + the latest supported Python + CUDA".
+      matrix_filter: map(select(.ARCH == "amd64")) | group_by(.CUDA_VER|split(".")|map(tonumber)|.[0]) | map(max_by([(.PY_VER|split(".")|map(tonumber)), (.CUDA_VER|split(".")|map(tonumber))]))
   build-wheels:
     needs:
       - compute-matrix
@@ -72,8 +83,18 @@ jobs:
     uses: ./.github/workflows/wheels-test.yaml
     with:
       build_type: pull-request
-      script: "ci/test_wheel.sh"
+      script: "ci/test_wheel.sh false"
       matrix_filter: ${{ needs.compute-matrix.outputs.TEST_MATRIX }}
+  test-wheels-pynvjitlink:
+    needs:
+      - build-wheels
+      - compute-matrix
+    uses: ./.github/workflows/wheels-test.yaml
+    with:
+      build_type: pull-request
+      script: "ci/test_wheel.sh true"
+      # This selects "ARCH=amd64 + the latest supported Python + CUDA".
+      matrix_filter: map(select(.ARCH == "amd64")) | group_by(.CUDA_VER|split(".")|map(tonumber)|.[0]) | map(max_by([(.PY_VER|split(".")|map(tonumber)), (.CUDA_VER|split(".")|map(tonumber))]))
   build-docs:
     needs:
       - build-conda

diff --git a/ci/test_conda.sh b/ci/test_conda.sh
@@ -3,6 +3,8 @@
 
 set -euo pipefail
 
+USE_PYNVJITLINK = $1
+
 . /opt/conda/etc/profile.d/conda.sh
 
 rapids-logger "Install testing dependencies"
@@ -42,8 +44,13 @@ EXITCODE=0
 trap "EXITCODE=1" ERR
 set +e
 
+if [ "$USE_PYNVJITLINK" == true ]; then
+    rapids-logger "Install pynvjitlink"
+    conda install -c rapidsai pynvjitlink
+fi
+
 rapids-logger "Run Tests"
-python -m numba.runtests numba.cuda.tests -v
+ENABLE_PYNVJITLINK=$USE_PYNVJITLINK python -m numba.runtests numba.cuda.tests -v
 
 popd
 

diff --git a/ci/test_wheel.sh b/ci/test_wheel.sh
@@ -3,13 +3,20 @@
 
 set -euo pipefail
 
+USE_PYNVJITLINK = $1
+
 rapids-logger "Install testing dependencies"
 # TODO: Replace with rapids-dependency-file-generator
 python -m pip install \
     psutil \
     cuda-python \
     pytest
 
+if [ "$USE_PYNVJITLINK" == true ]; then
+    rapids-logger "Install pynvjitlink"
+    python -m pip install pynvjitlink
+fi
+
 rapids-logger "Install wheel"
 package=$(realpath wheel/numba_cuda*.whl)
 echo "Package path: $package"
@@ -26,6 +33,6 @@ rapids-logger "Show Numba system info"
 python -m numba --sysinfo
 
 rapids-logger "Run Tests"
-python -m numba.runtests numba.cuda.tests -v
+ENABLE_PYNVJITLINK=$USE_PYNVJITLINK python -m numba.runtests numba.cuda.tests -v
 
 popd