Skip to content

GH-47371: [Python] Fix installation of Numba-CUDA in CI environment #47372

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
gmarkall wants to merge 1 commit into
base: main
Choose a base branch
from
Open

GH-47371: [Python] Fix installation of Numba-CUDA in CI environment #47372

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .env
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -67,7 +67,7 @@ LLVM=18
MAVEN=3.8.7
NODE=20
NUMBA=latest
NUMBA_CUDA=latest
NUMBA_CUDA=0.18.1
NUMPY=latest
PANDAS=latest
PYTHON=3.9
Expand Down
3 changes: 2 additions & 1 deletion ci/docker/linux-apt-python-3.dockerfile
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -33,9 +33,10 @@ RUN python3 -m venv ${ARROW_PYTHON_VENV} && \

ARG numba
ARG numba_cuda
ARG cuda
COPY ci/scripts/install_numba.sh /arrow/ci/scripts/
RUN if [ "${numba}" != "" ]; then \
/arrow/ci/scripts/install_numba.sh ${numba} ${numba_cuda} \
/arrow/ci/scripts/install_numba.sh ${numba} ${numba_cuda} ${cuda} \
; fi

ENV ARROW_ACERO=ON \
Expand Down
21 changes: 16 additions & 5 deletions ci/scripts/install_numba.sh
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,8 @@

set -e

if [ "$#" -ne 1 ] && [ "$#" -ne 2 ]; then
echo "Usage: $0 <numba version> [numba-cuda version]"
if [ "$#" -ne 1 ] && [ "$#" -ne 2 ] && [ "$#" -ne 3 ]; then
echo "Usage: $0 <numba version> [numba-cuda version] [cuda version]"
exit 1
fi

Expand Down Expand Up @@ -48,10 +48,21 @@ fi

numba_cuda=$2

if [ "$#" -eq 3 ]; then
cuda=$3
else
# Default to CUDA 11
cuda=11
fi
Comment on lines +51 to +56
Copy link
Member

@pitrou pitrou Aug 19, 2025

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I'm curious, is it possible to query the system to get the CUDA version instead of having to pass it explicitly?


# Variants are cu11, cu12, etc. depending on CUDA version
variant="cu${cuda%%.*}"

if [ "${numba_cuda}" = "master" ]; then
pip install https://github.com/NVIDIA/numba-cuda/archive/main.tar.gz#egg=numba-cuda
pip install "numba-cuda[$variant] @ https://github.com/NVIDIA/numba-cuda/archive/main.tar.gz"
elif [ "${numba_cuda}" = "latest" ]; then
pip install numba-cuda
pip install numba-cuda[$variant]
else
pip install "numba-cuda==${numba_cuda}"
echo installing for $variant
pip install "numba-cuda[$variant]==${numba_cuda}"
fi
1 change: 1 addition & 0 deletions docker-compose.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -966,6 +966,7 @@ services:
base: ${REPO}:${ARCH}-ubuntu-${UBUNTU}-cuda-${CUDA}-cpp
numba: ${NUMBA}
numba_cuda: ${NUMBA_CUDA}
cuda: ${CUDA}
shm_size: *shm-size
environment:
<<: [*common, *ccache, *sccache]
Expand Down
Loading