From fc3f3b5b7b75d2163ac7cf294806052c3b5f3121 Mon Sep 17 00:00:00 2001 From: Graham Markall Date: Mon, 21 Jul 2025 11:19:20 +0100 Subject: [PATCH 1/2] Numba interop: fix device pointer value with NVIDIA bindings Numba and Numba-CUDA return a different type from `Context.memalloc()` depending on whether their built-in ctypes bindings or the NVIDIA CUDA Python bindings are in use - either e `ctypes.c_void_p` or a `cuda.bindings.driver.CUdeviceptr`. Whilst this inconsistency is unfortunate, it's hard to change as existing code in the wild depends on it. The issue in Arrow is that the value of the pointer cannot be obtained by `device_pointer.value` for a `CUdeviceptr`. Numba and Numba-CUDA do provide another property, `device_pointer_value`, that provides the device pointer as an `int` regardless of the kind of binding in use, so we can switch to use this for consistency between the two kinds of bindings. Fixes #47128. --- python/pyarrow/_cuda.pyx | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/python/pyarrow/_cuda.pyx b/python/pyarrow/_cuda.pyx index 340f7e35c73..96e5c139492 100644 --- a/python/pyarrow/_cuda.pyx +++ b/python/pyarrow/_cuda.pyx @@ -451,9 +451,9 @@ cdef class CudaBuffer(Buffer): Device buffer as a view of numba MemoryPointer. """ ctx = Context.from_numba(mem.context) - if mem.device_pointer.value is None and mem.size==0: + if mem.device_pointer_value is None and mem.size==0: return ctx.new_buffer(0) - return ctx.foreign_buffer(mem.device_pointer.value, mem.size, base=mem) + return ctx.foreign_buffer(mem.device_pointer_value, mem.size, base=mem) def to_numba(self): """Return numba memory pointer of CudaBuffer instance. From 232c66ebc3861c166a624b6387f317646a38955c Mon Sep 17 00:00:00 2001 From: Graham Markall Date: Mon, 21 Jul 2025 12:36:10 +0100 Subject: [PATCH 2/2] Update CI to install numba-cuda --- .env | 1 + ci/docker/linux-apt-python-3.dockerfile | 3 ++- ci/scripts/install_numba.sh | 18 ++++++++++++++++-- docker-compose.yml | 14 +++++++++----- 4 files changed, 28 insertions(+), 8 deletions(-) diff --git a/.env b/.env index 81638c67efd..6b9bea19f94 100644 --- a/.env +++ b/.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 diff --git a/ci/docker/linux-apt-python-3.dockerfile b/ci/docker/linux-apt-python-3.dockerfile index e215976d448..d68bed26288 100644 --- a/ci/docker/linux-apt-python-3.dockerfile +++ b/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 \ diff --git a/ci/scripts/install_numba.sh b/ci/scripts/install_numba.sh index 87c256e4b8a..22e0df2b3c6 100755 --- a/ci/scripts/install_numba.sh +++ b/ci/scripts/install_numba.sh @@ -19,8 +19,8 @@ set -e -if [ "$#" -ne 1 ]; then - echo "Usage: $0 " +if [ "$#" -ne 1 ] && [ "$#" -ne 2 ]; then + echo "Usage: $0 [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 diff --git a/docker-compose.yml b/docker-compose.yml index 528e4cc8fce..7bdd82c5757 100644 --- a/docker-compose.yml +++ b/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: 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: # UBUNTU: 22.04, 24.04 + # NUMBA: master, latest, + # NUMBA_CUDA: master, latest, 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]