Skip to content

Commit ddd25d1

Browse files
authored
GH-47128: [Python] Numba-CUDA interop with NVIDIA bindings (#47150)
### Rationale for this change Testing with [Numba-CUDA](https://github.com/NVIDIA/numba-cuda) which uses the [NVIDIA CUDA Python bindings](https://github.com/NVIDIA/cuda-python) by default identified that PyArrow Numba interop has an incompatibility with Numba / Numba-CUDA using the NVIDIA bindings. See Issue #47128. ### What changes are included in this PR? The fix is to get device pointer values from their `device_pointer_value` property, which is consistent across the ctypes and NVIDIA bindings in Numba. I also attempted to update the CI config to install Numba-CUDA. I think some of the comments in `docker-compose.yml` were a bit out of sync with changes to it, so I also updated comments that appeared to be relevant to reflect what I had to run locally. I could have got the CI changes all wrong - happy to change these, as they're not really the critical part of this PR. Fixes #47128. ### Are these changes tested? Yes, by the existing `test_cuda_numba_interop.py` and the CI changes in this PR. ### Are there any user-facing changes? No. * GitHub Issue: #47128 Authored-by: Graham Markall <[email protected]> Signed-off-by: Sutou Kouhei <[email protected]>
1 parent cd8e677 commit ddd25d1

File tree

5 files changed

+30
-10
lines changed

5 files changed

+30
-10
lines changed

.env

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ LLVM=18
6767
MAVEN=3.8.7
6868
NODE=20
6969
NUMBA=latest
70+
NUMBA_CUDA=latest
7071
NUMPY=latest
7172
PANDAS=latest
7273
PYTHON=3.9

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

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -32,9 +32,10 @@ RUN python3 -m venv ${ARROW_PYTHON_VENV} && \
3232
-r arrow/python/requirements-test.txt
3333

3434
ARG numba
35+
ARG numba_cuda
3536
COPY ci/scripts/install_numba.sh /arrow/ci/scripts/
3637
RUN if [ "${numba}" != "" ]; then \
37-
/arrow/ci/scripts/install_numba.sh ${numba} \
38+
/arrow/ci/scripts/install_numba.sh ${numba} ${numba_cuda} \
3839
; fi
3940

4041
ENV ARROW_ACERO=ON \

ci/scripts/install_numba.sh

Lines changed: 16 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -19,8 +19,8 @@
1919

2020
set -e
2121

22-
if [ "$#" -ne 1 ]; then
23-
echo "Usage: $0 <numba version>"
22+
if [ "$#" -ne 1 ] && [ "$#" -ne 2 ]; then
23+
echo "Usage: $0 <numba version> [numba-cuda version]"
2424
exit 1
2525
fi
2626

@@ -41,3 +41,17 @@ elif [ "${numba}" = "latest" ]; then
4141
else
4242
pip install "numba==${numba}"
4343
fi
44+
45+
if [ "$#" -eq 1 ]; then
46+
exit 0
47+
fi
48+
49+
numba_cuda=$2
50+
51+
if [ "${numba_cuda}" = "master" ]; then
52+
pip install https://github.com/NVIDIA/numba-cuda/archive/main.tar.gz#egg=numba-cuda
53+
elif [ "${numba_cuda}" = "latest" ]; then
54+
pip install numba-cuda
55+
else
56+
pip install "numba-cuda==${numba_cuda}"
57+
fi

docker-compose.yml

Lines changed: 9 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -560,10 +560,11 @@ services:
560560

561561
ubuntu-cuda-cpp:
562562
# Usage:
563-
# docker compose build cuda-cpp
564-
# docker compose run --rm cuda-cpp
563+
# docker compose build ubuntu-cuda-cpp
564+
# docker compose run --rm ubuntu-cuda-cpp
565565
# Parameters:
566566
# ARCH: amd64
567+
# UBUNTU: 22.04, 24.04
567568
# CUDA: <depends on your nvidia driver, should match system CUDA>
568569
image: ${REPO}:${ARCH}-ubuntu-${UBUNTU}-cuda-${CUDA}-cpp
569570
build:
@@ -946,13 +947,15 @@ services:
946947

947948
ubuntu-cuda-python:
948949
# Usage:
949-
# docker compose build cuda-cpp
950-
# docker compose build cuda-python
951-
# docker compose run --rm cuda-python
950+
# docker compose build ubuntu-cuda-cpp
951+
# docker compose build ubuntu-cuda-python
952+
# docker compose run --rm ubuntu-cuda-python
952953
# Parameters:
953954
# ARCH: amd64
954955
# CUDA: <depends on your nvidia driver, should match system CUDA>
955956
# UBUNTU: 22.04, 24.04
957+
# NUMBA: master, latest, <version>
958+
# NUMBA_CUDA: master, latest, <version>
956959
image: ${REPO}:${ARCH}-ubuntu-${UBUNTU}-cuda-${CUDA}-python-3
957960
build:
958961
context: .
@@ -962,6 +965,7 @@ services:
962965
args:
963966
base: ${REPO}:${ARCH}-ubuntu-${UBUNTU}-cuda-${CUDA}-cpp
964967
numba: ${NUMBA}
968+
numba_cuda: ${NUMBA_CUDA}
965969
shm_size: *shm-size
966970
environment:
967971
<<: [*common, *ccache, *sccache]

python/pyarrow/_cuda.pyx

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -451,9 +451,9 @@ cdef class CudaBuffer(Buffer):
451451
Device buffer as a view of numba MemoryPointer.
452452
"""
453453
ctx = Context.from_numba(mem.context)
454-
if mem.device_pointer.value is None and mem.size==0:
454+
if mem.device_pointer_value is None and mem.size==0:
455455
return ctx.new_buffer(0)
456-
return ctx.foreign_buffer(mem.device_pointer.value, mem.size, base=mem)
456+
return ctx.foreign_buffer(mem.device_pointer_value, mem.size, base=mem)
457457

458458
def to_numba(self):
459459
"""Return numba memory pointer of CudaBuffer instance.

0 commit comments

Comments
 (0)