pyrox.dev
1{
2 lib,
3 stdenv,
4 buildPythonPackage,
5 autoPatchelfHook,
6 pythonRelaxDepsHook,
7 onnxruntime,
8 coloredlogs,
9 numpy,
10 packaging,
11 oneDNN,
12 re2,
13
14}:
15
16# onnxruntime requires an older protobuf.
17# Doing an override in protobuf in the python-packages set
18# can give you a functioning Python package but note not
19# all Python packages will be compatible then.
20#
21# Because protobuf is not always needed we remove it
22# as a runtime dependency from our wheel.
23#
24# We do include here the non-Python protobuf so the shared libs
25# link correctly. If you do also want to include the Python
26# protobuf, you can add it to your Python env, but be aware
27# the version likely mismatches with what is used here.
28
29buildPythonPackage {
30 inherit (onnxruntime) pname version;
31 format = "wheel";
32 src = onnxruntime.dist;
33
34 unpackPhase = ''
35 cp -r $src dist
36 chmod +w dist
37 '';
38
39 nativeBuildInputs = [ pythonRelaxDepsHook ] ++ lib.optionals stdenv.isLinux [ autoPatchelfHook ];
40
41 # This project requires fairly large dependencies such as sympy which we really don't always need.
42 pythonRemoveDeps = [
43 "flatbuffers"
44 "protobuf"
45 "sympy"
46 ];
47
48 # Libraries are not linked correctly.
49 buildInputs =
50 [
51 oneDNN
52 re2
53 onnxruntime.protobuf
54 ]
55 ++ lib.optionals onnxruntime.passthru.cudaSupport (
56 with onnxruntime.passthru.cudaPackages;
57 [
58 libcublas # libcublasLt.so.XX libcublas.so.XX
59 libcurand # libcurand.so.XX
60 libcufft # libcufft.so.XX
61 cudnn # libcudnn.soXX
62 cuda_cudart # libcudart.so.XX
63 ]
64 );
65
66 propagatedBuildInputs = [
67 coloredlogs
68 # flatbuffers
69 numpy
70 packaging
71 # protobuf
72 # sympy
73 ];
74
75 meta = onnxruntime.meta;
76}