From c97767c7a0ac21341e981ef50165aeaa87dd23a7 Mon Sep 17 00:00:00 2001
From: brandon-b-miller <brmiller@nvidia.com>
Date: Mon, 30 Sep 2024 12:31:24 -0700
Subject: [PATCH] add a new ci job for testing with pynvjitlink

---
 .github/workflows/pr.yaml | 25 +++++++++++++++++++++++--
 ci/test_conda.sh          |  9 ++++++++-
 ci/test_wheel.sh          |  9 ++++++++-
 3 files changed, 39 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/pr.yaml b/.github/workflows/pr.yaml
index 7498eb1..d4c54d5 100644
--- 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
index 5861128..a932630 100755
--- 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
index 5ffd014..7121517 100755
--- a/ci/test_wheel.sh
+++ b/ci/test_wheel.sh
@@ -3,6 +3,8 @@
 
 set -euo pipefail
 
+USE_PYNVJITLINK = $1
+
 rapids-logger "Install testing dependencies"
 # TODO: Replace with rapids-dependency-file-generator
 python -m pip install \
@@ -10,6 +12,11 @@ python -m pip install \
     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