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