Update CI to install numba-cuda

Author: gmarkall

.env

@@ -67,6 +67,7 @@ LLVM=18
 MAVEN=3.8.7
 NODE=20
 NUMBA=latest
+NUMBA_CUDA=latest
 NUMPY=latest
 PANDAS=latest
 PYTHON=3.9

ci/docker/linux-apt-python-3.dockerfile

@@ -32,9 +32,10 @@ RUN python3 -m venv ${ARROW_PYTHON_VENV} && \
       -r arrow/python/requirements-test.txt
 
 ARG numba
+ARG numba_cuda
 COPY ci/scripts/install_numba.sh /arrow/ci/scripts/
 RUN if [ "${numba}" != "" ]; then \
-        /arrow/ci/scripts/install_numba.sh ${numba} \
+        /arrow/ci/scripts/install_numba.sh ${numba} ${numba_cuda} \
     ; fi
 
 ENV ARROW_ACERO=ON \

ci/scripts/install_numba.sh

@@ -19,8 +19,8 @@
 
 set -e
 
-if [ "$#" -ne 1 ]; then
-  echo "Usage: $0 <numba version>"
+if [ "$#" -ne 1 ] && [ "$#" -ne 2 ]; then
+  echo "Usage: $0 <numba version> [numba-cuda version]"
   exit 1
 fi
 
@@ -41,3 +41,17 @@ elif [ "${numba}" = "latest" ]; then
 else
   pip install "numba==${numba}"
 fi
+
+if [ "$#" -eq 1 ]; then
+  exit 0
+fi
+
+numba_cuda=$2
+
+if [ "${numba_cuda}" = "master" ]; then
+  pip install https://github.com/NVIDIA/numba-cuda/archive/main.tar.gz#egg=numba-cuda
+elif [ "${numba_cuda}" = "latest" ]; then
+  pip install numba-cuda
+else
+  pip install "numba-cuda==${numba_cuda}"
+fi

docker-compose.yml

@@ -560,10 +560,11 @@ services:
 
   ubuntu-cuda-cpp:
     # Usage:
-    #   docker compose build cuda-cpp
-    #   docker compose run --rm cuda-cpp
+    #   docker compose build ubuntu-cuda-cpp
+    #   docker compose run --rm ubuntu-cuda-cpp
     # Parameters:
     #   ARCH: amd64
+    #   UBUNTU: 22.04, 24.04
     #   CUDA: <depends on your nvidia driver, should match system CUDA>
     image: ${REPO}:${ARCH}-ubuntu-${UBUNTU}-cuda-${CUDA}-cpp
     build:
@@ -946,13 +947,15 @@ services:
 
   ubuntu-cuda-python:
     # Usage:
-    #   docker compose build cuda-cpp
-    #   docker compose build cuda-python
-    #   docker compose run --rm cuda-python
+    #   docker compose build ubuntu-cuda-cpp
+    #   docker compose build ubuntu-cuda-python
+    #   docker compose run --rm ubuntu-cuda-python
     # Parameters:
     #   ARCH: amd64
     #   CUDA: <depends on your nvidia driver, should match system CUDA>
     #   UBUNTU: 22.04, 24.04
+    #   NUMBA: master, latest, <version>
+    #   NUMBA_CUDA: master, latest, <version>
     image: ${REPO}:${ARCH}-ubuntu-${UBUNTU}-cuda-${CUDA}-python-3
     build:
       context: .
@@ -962,6 +965,7 @@ services:
       args:
         base: ${REPO}:${ARCH}-ubuntu-${UBUNTU}-cuda-${CUDA}-cpp
         numba: ${NUMBA}
+        numba_cuda: ${NUMBA_CUDA}
     shm_size: *shm-size
     environment:
       <<: [*common, *ccache, *sccache]