apache / arrow-nanoarrow

Helpers for Arrow C Data & Arrow C Stream interfaces
https://arrow.apache.org/nanoarrow
Apache License 2.0
149 stars 34 forks source link

feat(python): Implement CUDA build in Python bindings #547

Closed paleolimbot closed 5 days ago

paleolimbot commented 6 days ago

This PR updates the build system such that nanoarrow's CUDA capability can be exposed in Python (it doesn't quite expose any of the functionality yet). I'm not sure this is the best long-term strategy...I think we might be able to get CUDA plugged in at runtime because we don't use much of the API...but I think this would at least enable distributing a cuda-enabled build on conda-forge.

To build, one has to do something like:

export NANOARROW_PYTHON_CUDA_HOME=/usr/local/cuda
pip install .

On Windows, this would be something like:

$Env:NANOARROW_PYTHON_CUDA_HOME = 'C:/Program Files/NVIDIA GPU Computing Toolkit/CUDA/v12.5'
python -m pip install .

Then, you should be able to resolve a CUDA device:

from nanoarrow import device

device.resolve(device.DeviceType.CUDA.value, 0)
#> <nanoarrow.device.Device>
#> - device_type: CUDA <2>
#> - device_id: 0