bifurcate pynvjitlink test scripts

.github/workflows/pr.yaml

@@ -56,7 +56,7 @@ jobs:
     uses: ./.github/workflows/conda-python-tests.yaml
     with:
       build_type: pull-request
-      script: "ci/test_conda.sh false"
+      script: "ci/test_conda.sh"
       run_codecov: false
       matrix_filter: ${{ needs.compute-matrix.outputs.TEST_MATRIX }}
   test-conda-pynvjitlink:
@@ -66,7 +66,7 @@ jobs:
     uses: ./.github/workflows/conda-python-tests.yaml
     with:
       build_type: pull-request
-      script: "ci/test_conda.sh true"
+      script: "ci/test_conda_pynvjitlink.sh"
       run_codecov: false
       matrix_filter: map(select(.ARCH == "amd64" and .CUDA_VER == "12.5.1" and .PY_VER == "3.11"))
   build-wheels:
@@ -90,7 +90,7 @@ jobs:
     uses: ./.github/workflows/wheels-test.yaml
     with:
       build_type: pull-request
-      script: "ci/test_wheel.sh true"
+      script: "ci/test_wheel_pynvjitlink.sh"
       matrix_filter: map(select(.ARCH == "amd64" and .CUDA_VER == "12.5.1" and .PY_VER == "3.12"))
   build-docs:
     needs:

ci/test_conda.sh

@@ -3,12 +3,6 @@
 
 set -euo pipefail
 
-if [ "$1" == "true" ]; then
-    USE_PYNVJITLINK=true
-else
-    USE_PYNVJITLINK=false
-fi
-
 . /opt/conda/etc/profile.d/conda.sh
 
 if [ "${CUDA_VER%.*.*}" = "11" ]; then
@@ -53,16 +47,8 @@ EXITCODE=0
 trap "EXITCODE=1" ERR
 set +e
 
-if [ "$USE_PYNVJITLINK" == true ]; then
-    rapids-logger "Install pynvjitlink"
-    set +u
-    conda install -c rapidsai pynvjitlink
-    set -u
-    sh build_tests.sh
-fi
-
 rapids-logger "Run Tests"
-ENABLE_PYNVJITLINK=1 python -m numba.runtests numba.cuda.tests -v
+python -m numba.runtests numba.cuda.tests -v
 
 popd
 

ci/test_conda_pynvjitlink.sh

@@ -0,0 +1,77 @@
+#!/bin/bash
+# Copyright (c) 2024, NVIDIA CORPORATION
+
+set -euo pipefail
+
+. /opt/conda/etc/profile.d/conda.sh
+
+if [ "${CUDA_VER%.*.*}" = "11" ]; then
+  CTK_PACKAGES="cudatoolkit"
+else
+  CTK_PACKAGES="cuda-nvcc-impl cuda-nvrtc"
+fi
+
+rapids-logger "Install testing dependencies"
+# TODO: Replace with rapids-dependency-file-generator
+rapids-mamba-retry create -n test \
+    c-compiler \
+    cxx-compiler \
+    ${CTK_PACKAGES} \
+    cuda-python \
+    cuda-version=${CUDA_VER%.*} \
+    make \
+    psutil \
+    pytest \
+    python=${RAPIDS_PY_VERSION}
+
+# Temporarily allow unbound variables for conda activation.
+set +u
+conda activate test
+set -u
+
+rapids-mamba-retry install -c `pwd`/conda-repo numba-cuda
+
+RAPIDS_TESTS_DIR=${RAPIDS_TESTS_DIR:-"${PWD}/test-results"}/
+mkdir -p "${RAPIDS_TESTS_DIR}"
+pushd "${RAPIDS_TESTS_DIR}"
+
+rapids-print-env
+
+rapids-logger "Check GPU usage"
+nvidia-smi
+
+rapids-logger "Show Numba system info"
+python -m numba --sysinfo
+
+EXITCODE=0
+trap "EXITCODE=1" ERR
+set +e
+
+
+rapids-logger "Install pynvjitlink"
+set +u
+conda install -c rapidsai pynvjitlink
+set -u
+
+rapids_logger "Build tests"
+
+PY_SCRIPT="
+import numba_cuda
+root = numba_cuda.__file__.rstrip('__init__.py')
+test_dir = root + \"numba/cuda/tests/test_binary_generation/\"
+print(test_dir)
+"
+
+TEST_DIR=$(python -c "$PY_SCRIPT")
+pushd $TEST_DIR
+make
+popd
+
+
+rapids-logger "Run Tests"
+ENABLE_PYNVJITLINK=1 python -m numba.runtests numba.cuda.tests -v
+
+popd
+
+rapids-logger "Test script exiting with value: $EXITCODE"
+exit ${EXITCODE}

ci/test_wheel.sh

@@ -3,25 +3,13 @@
 
 set -euo pipefail
 
-if [ "$1" == "true" ]; then
-    USE_PYNVJITLINK=true
-else
-    USE_PYNVJITLINK=false
-fi
-
 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-cu12
-    sh build_tests.sh
-fi
-
 rapids-logger "Install wheel"
 package=$(realpath wheel/numba_cuda*.whl)
 echo "Package path: $package"
@@ -38,6 +26,6 @@ rapids-logger "Show Numba system info"
 python -m numba --sysinfo
 
 rapids-logger "Run Tests"
-ENABLE_PYNVJITLINK=1 python -m numba.runtests numba.cuda.tests -v
+python -m numba.runtests numba.cuda.tests -v
 
 popd

ci/test_wheel_pynvjitlink.sh

@@ -0,0 +1,47 @@
+#!/bin/bash
+# Copyright (c) 2023-2024, NVIDIA CORPORATION
+
+set -euo pipefail
+
+rapids-logger "Install testing dependencies"
+# TODO: Replace with rapids-dependency-file-generator
+python -m pip install \
+    psutil \
+    cuda-python \
+    pytest
+
+rapids-logger "Install pynvjitlink"
+python -m pip install pynvjitlink-cu12
+
+rapids-logger "Build tests"
+PY_SCRIPT="
+import numba_cuda
+root = numba_cuda.__file__.rstrip('__init__.py')
+test_dir = root + \"numba/cuda/tests/test_binary_generation/\"
+print(test_dir)
+"
+
+TEST_DIR=$(python -c "$PY_SCRIPT")
+pushd $TEST_DIR
+make
+popd
+
+rapids-logger "Install wheel"
+package=$(realpath wheel/numba_cuda*.whl)
+echo "Package path: $package"
+python -m pip install $package
+
+rapids-logger "Check GPU usage"
+nvidia-smi
+
+RAPIDS_TESTS_DIR=${RAPIDS_TESTS_DIR:-"${PWD}/test-results"}/
+mkdir -p "${RAPIDS_TESTS_DIR}"
+pushd "${RAPIDS_TESTS_DIR}"
+
+rapids-logger "Show Numba system info"
+python -m numba --sysinfo
+
+rapids-logger "Run Tests"
+ENABLE_PYNVJITLINK=1 python -m numba.runtests numba.cuda.tests -v
+
+popd