pyrox.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.62.0";
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-y/mvmzMwTHc/tWg4WFqFJOThbFiIF71OHLvtztkT+hE=";
55 };
56
57 postPatch = ''
58 substituteInPlace numba/cuda/cudadrv/driver.py \
59 --replace-fail \
60 "dldir = [" \
61 "dldir = [ '${addDriverRunpath.driverLink}/lib', "
62 '';
63
64 build-system = [
65 setuptools
66 numpy
67 ];
68
69 nativeBuildInputs = lib.optionals cudaSupport [
70 autoAddDriverRunpath
71 cudaPackages.cuda_nvcc
72 ];
73
74 buildInputs = lib.optionals cudaSupport [ cudaPackages.cuda_cudart ];
75
76 pythonRelaxDeps = [
77 "numpy"
78 ];
79
80 dependencies = [
81 numpy
82 llvmlite
83 ];
84
85 patches = lib.optionals cudaSupport [
86 (replaceVars ./cuda_path.patch {
87 cuda_toolkit_path = cudatoolkit;
88 cuda_toolkit_lib_path = lib.getLib cudatoolkit;
89 })
90 ];
91
92 nativeCheckInputs = [
93 pytestCheckHook
94 ];
95
96 preCheck = ''
97 export HOME="$(mktemp -d)"
98 # https://github.com/NixOS/nixpkgs/issues/255262
99 cd $out
100 '';
101
102 enabledTestPaths =
103 if doFullCheck then
104 null
105 else
106 [
107 # These are the most basic tests. Running all tests is too expensive, and
108 # some of them fail (also differently on different platforms), so it will
109 # be too hard to maintain such a `disabledTests` list.
110 "${python.sitePackages}/numba/tests/test_usecases.py"
111 ];
112
113 disabledTests = lib.optionals (stdenv.hostPlatform.isLinux && stdenv.hostPlatform.isAarch64) [
114 # captured stderr: Fatal Python error: Segmentation fault
115 "test_sum1d_pyobj"
116 ];
117
118 disabledTestPaths = lib.optionals (!testsWithoutSandbox) [
119 # See NOTE near passthru.tests.withoutSandbox
120 "${python.sitePackages}/numba/cuda/tests"
121 ];
122
123 pythonImportsCheck = [ "numba" ];
124
125 passthru.testers.cuda-detect =
126 writers.writePython3Bin "numba-cuda-detect"
127 { libraries = [ (numba.override { cudaSupport = true; }) ]; }
128 ''
129 from numba import cuda
130 cuda.detect()
131 '';
132 passthru.tests = {
133 # CONTRIBUTOR NOTE: numba also contains CUDA tests, though these cannot be run in
134 # this sandbox environment. Consider building the derivation below with
135 # --no-sandbox to get a view of how many tests succeed outside the sandbox.
136 withoutSandbox = numba.override {
137 doFullCheck = true;
138 cudaSupport = true;
139 testsWithoutSandbox = true;
140 };
141 withSandbox = numba.override {
142 cudaSupport = false;
143 doFullCheck = true;
144 testsWithoutSandbox = false;
145 };
146 numpy_1 = numba.override {
147 numpy = numpy_1;
148 };
149 };
150
151 meta = with lib; {
152 changelog = "https://numba.readthedocs.io/en/stable/release/${version}-notes.html";
153 description = "Compiling Python code using LLVM";
154 homepage = "https://numba.pydata.org/";
155 license = licenses.bsd2;
156 mainProgram = "numba";
157 };
158}