From ae55d117806227628a8dedcccebd113d83fff6f4 Mon Sep 17 00:00:00 2001 From: Gaetan Lepage Date: Tue, 11 Aug 2026 18:00:56 +0000 Subject: [PATCH] python3Packages.numba: cleanup --- .../python-modules/numba/cuda_path.patch | 79 ------------------- .../python-modules/numba/default.nix | 42 +++++++--- .../numba/nvidia-libs-paths.patch | 43 ++++++++++ 3 files changed, 74 insertions(+), 90 deletions(-) delete mode 100644 pkgs/development/python-modules/numba/cuda_path.patch create mode 100644 pkgs/development/python-modules/numba/nvidia-libs-paths.patch diff --git a/pkgs/development/python-modules/numba/cuda_path.patch b/pkgs/development/python-modules/numba/cuda_path.patch deleted file mode 100644 index 73f38e283d8c..000000000000 --- a/pkgs/development/python-modules/numba/cuda_path.patch +++ /dev/null @@ -1,79 +0,0 @@ -diff --git a/numba/cuda/cuda_paths.py b/numba/cuda/cuda_paths.py -index d195bbc29..0699b365b 100644 ---- a/numba/cuda/cuda_paths.py -+++ b/numba/cuda/cuda_paths.py -@@ -24,11 +24,7 @@ def _find_valid_path(options): - - def _get_libdevice_path_decision(): - options = [ -- ('Conda environment', get_conda_ctk()), -- ('Conda environment (NVIDIA package)', get_nvidia_libdevice_ctk()), -- ('CUDA_HOME', get_cuda_home('nvvm', 'libdevice')), -- ('System', get_system_ctk('nvvm', 'libdevice')), -- ('Debian package', get_debian_pkg_libdevice()), -+ ('Nix store', get_nix_ctk('nvvm', 'libdevice')), - ] - by, libdir = _find_valid_path(options) - return by, libdir -@@ -36,17 +32,14 @@ def _get_libdevice_path_decision(): - - def _nvvm_lib_dir(): - if IS_WIN32: -- return 'nvvm', 'bin' -+ return 'bin', - else: -- return 'nvvm', 'lib64' -+ return 'lib64', - - - def _get_nvvm_path_decision(): - options = [ -- ('Conda environment', get_conda_ctk()), -- ('Conda environment (NVIDIA package)', get_nvidia_nvvm_ctk()), -- ('CUDA_HOME', get_cuda_home(*_nvvm_lib_dir())), -- ('System', get_system_ctk(*_nvvm_lib_dir())), -+ ('Nix store', get_nix_ctk(*_nvvm_lib_dir())), - ] - by, path = _find_valid_path(options) - return by, path -@@ -66,7 +59,7 @@ def _cudalib_path(): - if IS_WIN32: - return 'bin' - else: -- return 'lib64' -+ return 'lib' - - - def _cuda_home_static_cudalib_path(): -@@ -78,10 +71,7 @@ def _cuda_home_static_cudalib_path(): - - def _get_cudalib_dir_path_decision(): - options = [ -- ('Conda environment', get_conda_ctk()), -- ('Conda environment (NVIDIA package)', get_nvidia_cudalib_ctk()), -- ('CUDA_HOME', get_cuda_home(_cudalib_path())), -- ('System', get_system_ctk(_cudalib_path())), -+ ('Nix store', get_nix_lib_ctk(_cudalib_path())), - ] - by, libdir = _find_valid_path(options) - return by, libdir -@@ -239,3 +229,19 @@ def get_debian_pkg_libdevice(): - if not os.path.exists(pkg_libdevice_location): - return None - return pkg_libdevice_location -+ -+ -+def get_nix_ctk(*subdirs): -+ """Return path to nix store cudatoolkit; or, None if it doesn't exist. -+ """ -+ base = '@cuda_toolkit_path@' -+ if os.path.exists(base): -+ return os.path.join(base, *subdirs) -+ -+ -+def get_nix_lib_ctk(*subdirs): -+ """Return path to nix store cudatoolkit-lib; or, None if it doesn't exist. -+ """ -+ base = '@cuda_toolkit_lib_path@' -+ if os.path.exists(base): -+ return os.path.join(base, *subdirs) diff --git a/pkgs/development/python-modules/numba/default.nix b/pkgs/development/python-modules/numba/default.nix index dbf3255e8e66..f9e436a384c4 100644 --- a/pkgs/development/python-modules/numba/default.nix +++ b/pkgs/development/python-modules/numba/default.nix @@ -39,12 +39,27 @@ }: let - cudatoolkit = cudaPackages.cuda_nvcc; + # `libnvvm.so` and `libdevice.10.bc` are bundled in `cuda_nvcc` up to CUDA 12.x, and shipped as a + # dedicated redistributable from CUDA 13.0 onwards + nvvmRoot = + if cudaPackages.cudaOlder "13.0" then + "${lib.getLib cudaPackages.cuda_nvcc}/nvvm" + else + "${lib.getLib cudaPackages.libnvvm}"; + + libCudaPath = + # Use cuda_compat to provide libcuda.so on pre-Thor Jetsons + if (cudaPackages.cuda_compat.meta.available or false) then + cudaPackages.cuda_compat + # Else, use the host CUDA driver library + else + addDriverRunpath.driverLink; in buildPythonPackage (finalAttrs: { version = "0.66.0"; pname = "numba"; pyproject = true; + __structuredAttrs = true; src = fetchFromGitHub { owner = "numba"; @@ -65,22 +80,27 @@ buildPythonPackage (finalAttrs: { ./numpy2.5.patch ] ++ lib.optionals cudaSupport [ - (replaceVars ./cuda_path.patch { - cuda_toolkit_path = cudatoolkit; - cuda_toolkit_lib_path = lib.getLib cudatoolkit; + # Hardcode the paths of the NVIDIA libraries which numba looks up at runtime, instead of + # relying on its discovery heuristics (conda environment, `CUDA_HOME`, `/usr/local/cuda`, ...) + (replaceVars ./nvidia-libs-paths.patch { + libcuda = libCudaPath; + libcudart = lib.getLib cudaPackages.cuda_cudart; + libcudart_static = lib.getOutput "static" cudaPackages.cuda_cudart; + libnvrtc = lib.getLib cudaPackages.cuda_nvrtc; + libnvvm = nvvmRoot; }) ]; postPatch = '' - substituteInPlace numba/cuda/cudadrv/driver.py \ - --replace-fail \ - "dldir = [" \ - "dldir = [ '${addDriverRunpath.driverLink}/lib', " - substituteInPlace setup.py \ - --replace-fail 'max_numpy_run_version = "2.5"' 'max_numpy_run_version = "2.6"' + --replace-fail \ + 'max_numpy_run_version = "2.5"' \ + 'max_numpy_run_version = "2.6"' + substituteInPlace numba/__init__.py \ - --replace-fail "(2, 4)" "(2, 6)" + --replace-fail \ + "(2, 4)" \ + "(2, 6)" ''; build-system = [ diff --git a/pkgs/development/python-modules/numba/nvidia-libs-paths.patch b/pkgs/development/python-modules/numba/nvidia-libs-paths.patch new file mode 100644 index 000000000000..475adba60153 --- /dev/null +++ b/pkgs/development/python-modules/numba/nvidia-libs-paths.patch @@ -0,0 +1,43 @@ +diff --git a/numba/cuda/cuda_paths.py b/numba/cuda/cuda_paths.py +index 8ff96640d..7eb262a22 100644 +--- a/numba/cuda/cuda_paths.py ++++ b/numba/cuda/cuda_paths.py +@@ -235,12 +235,20 @@ def get_cuda_paths(): + if hasattr(get_cuda_paths, '_cached_result'): + return get_cuda_paths._cached_result + else: +- # Not in cache ++ # Not in cache. ++ # The CUDA libraries are installed at fixed, known locations in the Nix ++ # store, so the discovery heuristics above (conda environment, ++ # CUDA_HOME, /usr/local/cuda, ...) are bypassed entirely. Note that ++ # `find_file` also accepts a list of directories, which is needed since ++ # each library lives in its own store path. + d = { +- 'nvvm': _get_nvvm_path(), +- 'libdevice': _get_libdevice_paths(), +- 'cudalib_dir': _get_cudalib_dir(), +- 'static_cudalib_dir': _get_static_cudalib_dir(), ++ 'nvvm': _env_path_tuple('Nix store', '@libnvvm@/lib/libnvvm.so'), ++ 'libdevice': _env_path_tuple( ++ 'Nix store', '@libnvvm@/libdevice/libdevice.10.bc'), ++ 'cudalib_dir': _env_path_tuple( ++ 'Nix store', ['@libcudart@/lib', '@libnvrtc@/lib']), ++ 'static_cudalib_dir': _env_path_tuple( ++ 'Nix store', '@libcudart_static@/lib'), + } + # Cache result + get_cuda_paths._cached_result = d +diff --git a/numba/cuda/cudadrv/driver.py b/numba/cuda/cudadrv/driver.py +index bea271fce..79312f1ae 100644 +--- a/numba/cuda/cudadrv/driver.py ++++ b/numba/cuda/cudadrv/driver.py +@@ -118,7 +118,7 @@ def locate_driver_and_loader(): + else: + # Assume to be *nix like + dlloader = ctypes.CDLL +- dldir = ['/usr/lib', '/usr/lib64'] ++ dldir = ['@libcuda@/lib'] + dlnames = ['libcuda.so', 'libcuda.so.1'] + + if envpath: