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