tjh.dev
1{
2 lib,
3 stdenv,
4 pythonAtLeast,
5 pythonOlder,
6 fetchFromGitHub,
7 fetchpatch2,
8 python,
9 buildPythonPackage,
10 setuptools,
11 numpy,
12 numpy_1,
13 llvmlite,
14 replaceVars,
15 writers,
16 numba,
17 pytestCheckHook,
18
19 config,
20
21 # CUDA-only dependencies:
22 addDriverRunpath,
23 autoAddDriverRunpath,
24 cudaPackages,
25
26 # CUDA flags:
27 cudaSupport ? config.cudaSupport,
28 testsWithoutSandbox ? false,
29 doFullCheck ? false,
30}:
31
32let
33 cudatoolkit = cudaPackages.cuda_nvcc;
34in
35buildPythonPackage rec {
36 version = "0.61.2";
37 pname = "numba";
38 pyproject = true;
39
40 disabled = pythonOlder "3.10" || pythonAtLeast "3.14";
41
42 src = fetchFromGitHub {
43 owner = "numba";
44 repo = "numba";
45 tag = version;
46 # Upstream uses .gitattributes to inject information about the revision
47 # hash and the refname into `numba/_version.py`, see:
48 #
49 # - https://git-scm.com/docs/gitattributes#_export_subst and
50 # - https://github.com/numba/numba/blame/5ef7c86f76a6e8cc90e9486487294e0c34024797/numba/_version.py#L25-L31
51 postFetch = ''
52 sed -i 's/git_refnames = "[^"]*"/git_refnames = " (tag: ${src.tag})"/' $out/numba/_version.py
53 '';
54 hash = "sha256-Qa2B5pOWrLb/1V3PSyiwS1x9ueXwDKRhDMDecBCAN+8=";
55 };
56
57 postPatch = ''
58 substituteInPlace numba/cuda/cudadrv/driver.py \
59 --replace-fail \
60 "dldir = [" \
61 "dldir = [ '${addDriverRunpath.driverLink}/lib', "
62
63 substituteInPlace setup.py \
64 --replace-fail 'max_numpy_run_version = "2.3"' 'max_numpy_run_version = "2.4"'
65 substituteInPlace numba/__init__.py \
66 --replace-fail "numpy_version > (2, 2)" "numpy_version > (2, 3)"
67 '';
68
69 env.NIX_CFLAGS_COMPILE = lib.optionalString stdenv.hostPlatform.isDarwin "-I${lib.getInclude stdenv.cc.libcxx}/include/c++/v1";
70
71 build-system = [
72 setuptools
73 numpy
74 ];
75
76 nativeBuildInputs = lib.optionals cudaSupport [
77 autoAddDriverRunpath
78 cudaPackages.cuda_nvcc
79 ];
80
81 buildInputs = lib.optionals cudaSupport [ cudaPackages.cuda_cudart ];
82
83 pythonRelaxDeps = [ "numpy" ];
84
85 dependencies = [
86 numpy
87 llvmlite
88 ];
89
90 patches = [
91 (fetchpatch2 {
92 url = "https://github.com/numba/numba/commit/e2c8984ba60295def17e363a926d6f75e7fa9f2d.patch";
93 includes = [ "numba/core/bytecode.py" ];
94 hash = "sha256-HIVbp3GSmnq6W7zrRIirIbhGjJsFN3PNyHSfAE8fdDw=";
95 })
96 ]
97 ++ lib.optionals cudaSupport [
98 (replaceVars ./cuda_path.patch {
99 cuda_toolkit_path = cudatoolkit;
100 cuda_toolkit_lib_path = lib.getLib cudatoolkit;
101 })
102 ];
103
104 nativeCheckInputs = [
105 pytestCheckHook
106 ];
107
108 preCheck = ''
109 export HOME="$(mktemp -d)"
110 # https://github.com/NixOS/nixpkgs/issues/255262
111 cd $out
112 '';
113
114 enabledTestPaths = lib.optionals (!doFullCheck) [
115 # These are the most basic tests. Running all tests is too expensive, and
116 # some of them fail (also differently on different platforms), so it will
117 # be too hard to maintain such a `disabledTests` list.
118 "${python.sitePackages}/numba/tests/test_usecases.py"
119 ];
120
121 disabledTests = lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64) [
122 # captured stderr: Fatal Python error: Segmentation fault
123 "test_sum1d_pyobj"
124 ];
125
126 disabledTestPaths = lib.optionals (!testsWithoutSandbox) [
127 # See NOTE near passthru.tests.withoutSandbox
128 "${python.sitePackages}/numba/cuda/tests"
129 ];
130
131 pythonImportsCheck = [ "numba" ];
132
133 passthru.testers.cuda-detect =
134 writers.writePython3Bin "numba-cuda-detect"
135 { libraries = [ (numba.override { cudaSupport = true; }) ]; }
136 ''
137 from numba import cuda
138 cuda.detect()
139 '';
140 passthru.tests = {
141 # CONTRIBUTOR NOTE: numba also contains CUDA tests, though these cannot be run in
142 # this sandbox environment. Consider building the derivation below with
143 # --no-sandbox to get a view of how many tests succeed outside the sandbox.
144 withoutSandbox = numba.override {
145 doFullCheck = true;
146 cudaSupport = true;
147 testsWithoutSandbox = true;
148 };
149 withSandbox = numba.override {
150 cudaSupport = false;
151 doFullCheck = true;
152 testsWithoutSandbox = false;
153 };
154 numpy_1 = numba.override {
155 numpy = numpy_1;
156 };
157 };
158
159 meta = with lib; {
160 changelog = "https://numba.readthedocs.io/en/stable/release/${version}-notes.html";
161 description = "Compiling Python code using LLVM";
162 homepage = "https://numba.pydata.org/";
163 license = licenses.bsd2;
164 mainProgram = "numba";
165 };
166}