Attempt to add test of patch to CI

.github/workflows/pr.yaml

@@ -19,6 +19,7 @@ jobs:
  - test-conda
  - build-wheels
  - test-wheels
+ - test-patch
  secrets: inherit
  uses: rapidsai/shared-workflows/.github/workflows/pr-builder.yaml@branch-24.02
  checks:
@@ -53,6 +54,16 @@ jobs:
  build_type: pull-request
  test_script: "ci/test_conda.sh"
  matrix_filter: ${{ needs.compute-matrix.outputs.TEST_MATRIX }}
+ test-patch:
+ needs:
+ - build-conda
+ - compute-matrix
+ uses: rapidsai/shared-workflows/.github/workflows/conda-python-tests.yaml@branch-24.02
+ with:
+ build_type: pull-request
+ test_script: "ci/test_patch.sh"
+ run_codecov: false
+ matrix_filter: ${{ needs.compute-matrix.outputs.TEST_MATRIX }}
  build-wheels:
  needs:
  - compute-matrix

ci/test_patch.py

@@ -0,0 +1,11 @@
+#!/usr/bin/env python
+# Copyright (c) 2024, NVIDIA CORPORATION
+
+from pynvjitlink import patch
+patch.patch_numba_linker()
+
+if __name__ == '__main__':
+ from numba.testing._runtests import _main
+ import sys
+
+ sys.exit(0 if _main(sys.argv) else 1)

ci/test_patch.sh

@@ -0,0 +1,41 @@
+#!/bin/bash
+# Copyright (c) 2024, NVIDIA CORPORATION
+
+set -euo pipefail
+
+. /opt/conda/etc/profile.d/conda.sh
+
+rapids-logger "Generate testing dependencies"
+# TODO: Replace with rapids-dependency-file-generator
+rapids-mamba-retry create -n test \
+ cuda-nvcc-impl \
+ cuda-nvrtc \
+ cuda-version=${RAPIDS_CUDA_VERSION%.*} \
+ "numba>=0.58" \
+ python=${RAPIDS_PY_VERSION}
+
+# Temporarily allow unbound variables for conda activation.
+set +u
+conda activate test
+set -u
+
+PYTHON_CHANNEL=$(rapids-download-conda-from-s3 python)
+
+rapids-print-env
+
+rapids-mamba-retry install \
+ --channel "${PYTHON_CHANNEL}" \
+ pynvjitlink
+
+rapids-logger "Check GPU usage"
+nvidia-smi
+
+EXITCODE=0
+trap "EXITCODE=1" ERR
+set +e
+
+rapids-logger "Test Numba with patch"
+python ci/test_patch.py numba.cuda.tests -v -m
+
+rapids-logger "Test script exiting with value: $EXITCODE"
+exit ${EXITCODE}