diff --git a/ci/docker/linux-apt-python-3.dockerfile b/ci/docker/linux-apt-python-3.dockerfile index d68bed26288..d5e7bf00886 100644 --- a/ci/docker/linux-apt-python-3.dockerfile +++ b/ci/docker/linux-apt-python-3.dockerfile @@ -33,9 +33,10 @@ RUN python3 -m venv ${ARROW_PYTHON_VENV} && \ ARG numba ARG numba_cuda +ARG cuda COPY ci/scripts/install_numba.sh /arrow/ci/scripts/ RUN if [ "${numba}" != "" ]; then \ - /arrow/ci/scripts/install_numba.sh ${numba} ${numba_cuda} \ + /arrow/ci/scripts/install_numba.sh ${numba} ${numba_cuda} ${cuda} \ ; fi ENV ARROW_ACERO=ON \ diff --git a/ci/scripts/install_numba.sh b/ci/scripts/install_numba.sh index 053579bd6da..91c36ca5713 100755 --- a/ci/scripts/install_numba.sh +++ b/ci/scripts/install_numba.sh @@ -19,8 +19,8 @@ set -e -if [ "$#" -ne 1 ] && [ "$#" -ne 2 ]; then - echo "Usage: $0 [numba-cuda version]" +if [ "$#" -ne 1 ] && [ "$#" -ne 2 ] && [ "$#" -ne 3 ]; then + echo "Usage: $0 [numba-cuda version] [CUDA version]" exit 1 fi @@ -34,8 +34,6 @@ if [ -n "${ARROW_PYTHON_VENV:-}" ]; then . "${ARROW_PYTHON_VENV}/bin/activate" fi -# TODO: GH-47371 (install Python CUDA bindings explicitly) - if [ "${numba}" = "master" ]; then pip install https://github.com/numba/numba/archive/main.tar.gz#egg=numba elif [ "${numba}" = "latest" ]; then @@ -50,10 +48,19 @@ fi numba_cuda=$2 +DEFAULT_CUDA_VERSION="12" + +if [ "$#" -eq 3 ]; then + # Extract the CUDA major version only + cuda_version="${3%%.*}" +else + cuda_version="${DEFAULT_CUDA_VERSION}" +fi + if [ "${numba_cuda}" = "master" ]; then - pip install https://github.com/NVIDIA/numba-cuda/archive/main.tar.gz#egg=numba-cuda + pip install "numba-cuda[cu${cuda_version}] @ https://github.com/NVIDIA/numba-cuda/archive/main.tar.gz" elif [ "${numba_cuda}" = "latest" ]; then - pip install numba-cuda + pip install "numba-cuda[cu${cuda_version}]" else - pip install "numba-cuda==${numba_cuda}" + pip install "numba-cuda[cu${cuda_version}]==${numba_cuda}" fi diff --git a/compose.yaml b/compose.yaml index e83f1446cd4..8abc9ad6604 100644 --- a/compose.yaml +++ b/compose.yaml @@ -999,6 +999,7 @@ services: base: ${REPO}:${ARCH}-ubuntu-${UBUNTU}-cuda-${CUDA}-cpp numba: ${NUMBA} numba_cuda: ${NUMBA_CUDA} + cuda: ${CUDA} shm_size: *shm-size environment: <<: [*common, *ccache, *sccache] diff --git a/python/pyarrow/tests/test_cuda_numba_interop.py b/python/pyarrow/tests/test_cuda_numba_interop.py index 0ce25bb6082..876f3c7f761 100644 --- a/python/pyarrow/tests/test_cuda_numba_interop.py +++ b/python/pyarrow/tests/test_cuda_numba_interop.py @@ -22,9 +22,6 @@ except ImportError: pytestmark = pytest.mark.numpy -pytest.skip("Numba integration tests broken by Numba API changes, see GH-48265", - allow_module_level=True) - dtypes = ['uint8', 'int16', 'float32'] cuda = pytest.importorskip("pyarrow.cuda") nb_cuda = pytest.importorskip("numba.cuda")