diff --git a/.github/workflows/cuda_extra.yml b/.github/workflows/cuda_extra.yml index fedf672fb1b..0254314fe06 100644 --- a/.github/workflows/cuda_extra.yml +++ b/.github/workflows/cuda_extra.yml @@ -90,10 +90,6 @@ jobs: ubuntu: 24.04 image: ubuntu-cuda-python title: AMD64 Ubuntu 24 CUDA 12.9.0 Python - - cuda: 11.7.1 - ubuntu: 22.04 - image: ubuntu-cuda-python - title: AMD64 Ubuntu 22 CUDA 11.7.1 Python env: ARCHERY_DEBUG: 1 ARROW_ENABLE_TIMING_TESTS: OFF diff --git a/ci/scripts/install_numba.sh b/ci/scripts/install_numba.sh index 91c36ca5713..4569b323d9b 100755 --- a/ci/scripts/install_numba.sh +++ b/ci/scripts/install_numba.sh @@ -34,6 +34,12 @@ if [ -n "${ARROW_PYTHON_VENV:-}" ]; then . "${ARROW_PYTHON_VENV}/bin/activate" fi +# Numba-CUDA 0.30.4 uses np.row_stack, which was removed in NumPy 2.5. +# TODO: Remove after https://github.com/NVIDIA/numba-cuda/pull/916 is released. +if [ "$#" -ge 2 ]; then + pip install "numpy<2.5" +fi + if [ "${numba}" = "master" ]; then pip install https://github.com/numba/numba/archive/main.tar.gz#egg=numba elif [ "${numba}" = "latest" ]; then diff --git a/python/pyarrow/_cuda.pyx b/python/pyarrow/_cuda.pyx index a6d65e517c0..7bebb819428 100644 --- a/python/pyarrow/_cuda.pyx +++ b/python/pyarrow/_cuda.pyx @@ -16,6 +16,10 @@ # under the License. +import ctypes + +from pyarrow.vendored.version import Version + from pyarrow.lib cimport * from pyarrow.includes.libarrow_cuda cimport * from pyarrow.lib import allocate_buffer, as_buffer, ArrowTypeError @@ -23,6 +27,32 @@ from pyarrow.util import get_contiguous_span cimport cpython as cp +_NUMBA_CUDA_NATIVE_CONTEXT_VERSION = Version("0.28") + + +def _numba_context_handle_value(handle): + """Return the integer value of a legacy or native Numba context handle.""" + if hasattr(handle, "value"): + return handle.value + return int(handle) + + +def _make_numba_context_handle(uintptr_t handle): + """Create the context handle representation expected by Numba.""" + import numba.cuda + + # numba-cuda 0.28 replaced ctypes context handles with CUContext type. + if getattr(numba.cuda, "implementation", None) == "NVIDIA": + import numba_cuda + + version = Version(numba_cuda.__version__) + if version >= _NUMBA_CUDA_NATIVE_CONTEXT_VERSION: + from cuda.bindings.driver import CUcontext + return CUcontext(handle) + + return ctypes.c_void_p(handle) + + cdef class Context(_Weakrefable): """ CUDA driver context. @@ -81,7 +111,7 @@ cdef class Context(_Weakrefable): import numba.cuda context = numba.cuda.current_context() return Context(device_number=context.device.id, - handle=context.handle.value) + handle=_numba_context_handle_value(context.handle)) def to_numba(self): """ @@ -92,10 +122,9 @@ cdef class Context(_Weakrefable): context : numba.cuda.cudadrv.driver.Context Numba CUDA context instance. """ - import ctypes import numba.cuda device = numba.cuda.gpus[self.device_number] - handle = ctypes.c_void_p(self.handle) + handle = _make_numba_context_handle(self.handle) context = numba.cuda.cudadrv.driver.Context(device, handle) class DummyPendingDeallocs(object): diff --git a/python/pyarrow/tests/test_cuda_numba_interop.py b/python/pyarrow/tests/test_cuda_numba_interop.py index fa91cd0b490..f34e39b500e 100644 --- a/python/pyarrow/tests/test_cuda_numba_interop.py +++ b/python/pyarrow/tests/test_cuda_numba_interop.py @@ -15,6 +15,8 @@ # specific language governing permissions and limitations # under the License. +import ctypes + import pytest import pyarrow as pa try: @@ -50,8 +52,14 @@ def teardown_module(module): ids=context_choice_ids) def test_context(c): ctx, nb_ctx = context_choices[c] - assert ctx.handle == nb_ctx.handle.value - assert ctx.handle == ctx.to_numba().handle.value + converted = ctx.to_numba() + assert type(converted.handle) is type(nb_ctx.handle) + if isinstance(converted.handle, ctypes.c_void_p): + assert converted.handle.value == ctx.handle + else: + from cuda.bindings import driver + assert type(converted.handle) is driver.CUcontext + assert int(converted.handle) == ctx.handle ctx2 = cuda.Context.from_numba(nb_ctx) assert ctx.handle == ctx2.handle size = 10 @@ -203,7 +211,7 @@ def test_numba_context(c, dtype): with nb_cuda.gpus[0]: arr, cbuf = make_random_buffer(size, target='device', dtype=dtype, ctx=ctx) - assert cbuf.context.handle == nb_ctx.handle.value + assert cbuf.context.handle == cuda.Context.from_numba(nb_ctx).handle mem = cbuf.to_numba() darr = DeviceNDArray(arr.shape, arr.strides, arr.dtype, gpu_data=mem) np.testing.assert_equal(darr.copy_to_host(), arr)