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