[Python] test_cuda_numba_interop fails with numba-cuda

[pitrou]

Describe the bug, including details regarding any error messages, version, and platform.

Integrated CUDA support in Numba is deprecated in favor of the external Numba-CUDA module, maintained by NVidia.

However, when running our Numba CUDA interop tests with Numba-CUDA installed, I get the following failures on all tests:

  File "/home/antoine/arrow/dev/python/pyarrow/tests/test_cuda_numba_interop.py", line 41, in setup_module
    ctx2 = cuda.Context.from_numba(nb_ctx2)
  File "pyarrow/_cuda.pyx", line 84, in pyarrow._cuda.Context.from_numba
    handle=context.handle.value)
AttributeError: 'cuda.bindings.driver.CUcontext' object has no attribute 'value'

This probably means that our Numba interoperation code doesn't work with Numba-CUDA, only with the deprecated CUDA support in Numba itself.

Component(s)

GPU, Python

[pitrou]

@gmarkall Perhaps this is something that you or one of your colleagues might be interested in submitting a fix for?

[pitrou]

Also, ideally this should also be exercised on our CI. We do have a CUDA Python CI job, but it doesn't install Numba-CUDA, so the deprecated CUDA support is used instead.

[gmarkall]

I think this might be a bug in Numba-CUDA that is fixed in main but not yet in a release - let me check and get back to you.

[gmarkall]

Most of the test failures are resolved by NVIDIA/numba-cuda#328 and NVIDIA/numba-cuda#296.

There is one remaining issue which is not determined by whether Numba or Numba-CUDA is in use, but instead whether the Numba ctypes driver bindings or the NVIDIA CUDA Python driver bindings are in use - ctypes is the default in Numba, and the NVIDIA bindings are the default in Numba-CUDA. I'll look into a fix for this remaining issue tomorrow.

[gmarkall]

Oops, I forgot to write up the remaining issue - it is:

______________________________________ test_numba_memalloc[float32-numba.cuda] _______________________________________

c = 1, dtype = dtype('float32')

    @pytest.mark.parametrize("c", range(len(context_choice_ids)),
                             ids=context_choice_ids)
    @pytest.mark.parametrize("dtype", dtypes, ids=dtypes)
    def test_numba_memalloc(c, dtype):
        ctx, nb_ctx = context_choices[c]
        dtype = np.dtype(dtype)
        # Allocate memory using numba context
        # Warning: this will not be reflected in pyarrow context manager
        # (e.g bytes_allocated does not change)
        size = 10
        mem = nb_ctx.memalloc(size * dtype.itemsize)
        darr = DeviceNDArray((size,), (dtype.itemsize,), dtype, gpu_data=mem)
        darr[:5] = 99
        darr[5:] = 88
        np.testing.assert_equal(darr.copy_to_host()[:5], 99)
        np.testing.assert_equal(darr.copy_to_host()[5:], 88)
    
        # wrap numba allocated memory with CudaBuffer
>       cbuf = cuda.CudaBuffer.from_numba(mem)

pyarrow/tests/test_cuda_numba_interop.py:178: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

>   if mem.device_pointer.value is None and mem.size==0:
E   AttributeError: 'cuda.bindings.driver.CUdeviceptr' object has no attribute 'value'

memalloc returns c_void_p objects with the ctypes bindings, and CUdeviceptr objects with the NVIDIA bindings. This is obviously an unpleasant inconsistency, but I've not yet thought of a straightforward fix because other libraries had adapted to this inconsistency (e.g. https://github.com/rapidsai/rmm), so just straightforwardly making it consistent will break them.

[gmarkall]

With this change in Arrow:

diff --git a/python/pyarrow/_cuda.pyx b/python/pyarrow/_cuda.pyx
index 340f7e35c7..96e5c13949 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.

and the Numba-CUDA PR NVIDIA/numba-cuda#328 on top of Numba-CUDA main, all tests pass:

0 gmarkall@housel 11:32:46 (main) ~/numbadev/arrow/python
$ python -m pytest pyarrow/tests/test_cuda_numba_interop.py
============================================ test session starts =============================================
platform linux -- Python 3.10.18, pytest-8.4.1, pluggy-1.6.0
rootdir: /home/gmarkall/numbadev/arrow/python
configfile: setup.cfg
plugins: hypothesis-6.135.32
collected 50 items                                                                                                   

pyarrow/tests/test_cuda_numba_interop.py ..................................................            [100%]

====================================== 50 passed, 18 warnings in 2.48s =======================================

0 gmarkall@housel 11:32:50 (main) ~/numbadev/arrow/python
$ pip list
...
numba-cuda          0.16.0      /home/gmarkall/numbadev/numba-cuda
...

Next I'll aim to:

• Get a Numba-CUDA release out with the relevant fixes, and
• Make a PR to Arrow with the above change.

The change above swaps the use of device_pointer.value with device_pointer_value because device_pointer_value always returns None or an integer regardless of whether the device_pointer is a c_void_p or CUdeviceptr. Feel free to take this change and merge it before I get chance to make a PR, if you would prefer :-)

[pitrou]

No, I'll let you make a PR. This is not urgent since we have just released 21.0.0 and the next release will be in 3 months.

It would be nice if the code remains compatible with the old CUDA support shipped with Numba. I assume that shouldn't be too hard?

[gmarkall]

It would be nice if the code remains compatible with the old CUDA support shipped with Numba. I assume that shouldn't be too hard?

Yes - I consider it a hard requirement that Numba and Numba-CUDA are compatible for the CUDA target (notwithstanding new additions to Numba-CUDA that aren't present in upstream Numba). Any discrepancy or breaking of existing code is a bug, not a feature 🙂

[gmarkall]

For anyone using the existing Arrow release with Numba-CUDA, setting NUMBA_CUDA_USE_NVIDIA_BINDING=0 will override its default setting to use the NVIDIA CUDA Python bindings, and therefore allow it to work without any of the PyArrow or unreleased Numba-CUDA changes.