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