pyrox.dev
1{
2 lib,
3 python,
4 autoAddDriverRunpath,
5 buildPythonPackage,
6 autoPatchelfHook,
7 unzip,
8 cudaPackages,
9}:
10
11let
12 pyVersion = "${lib.versions.major python.version}${lib.versions.minor python.version}";
13 buildVersion = lib.optionalString (cudaPackages ? tensorrt) cudaPackages.tensorrt.version;
14 wheelVersion = lib.optionalString (cudaPackages ? tensorrt) (
15 if
16 (builtins.elem buildVersion [
17 "8.6.1.6"
18 "10.3.0.26"
19 ])
20 then
21 builtins.concatStringsSep "." (lib.take 3 (builtins.splitVersion buildVersion))
22 else
23 buildVersion
24 );
25in
26buildPythonPackage rec {
27 pname = "tensorrt";
28 version = wheelVersion;
29
30 src = cudaPackages.tensorrt.src;
31
32 format = "wheel";
33 # We unpack the wheel ourselves because of the odd packaging.
34 dontUseWheelUnpack = true;
35
36 nativeBuildInputs = [
37 unzip
38 autoPatchelfHook
39 autoAddDriverRunpath
40 ];
41
42 preUnpack = ''
43 mkdir -p dist
44 tar --strip-components=2 -xf "$src" --directory=dist \
45 "TensorRT-${buildVersion}/python/tensorrt-${wheelVersion}-cp${pyVersion}-none-linux_x86_64.whl"
46 '';
47
48 sourceRoot = ".";
49
50 buildInputs = [
51 cudaPackages.cudnn
52 cudaPackages.tensorrt
53 ];
54
55 pythonImportsCheck = [ "tensorrt" ];
56
57 meta = with lib; {
58 description = "Python bindings for TensorRT, a high-performance deep learning interface";
59 homepage = "https://developer.nvidia.com/tensorrt";
60 license = licenses.unfree;
61 platforms = [ "x86_64-linux" ];
62 maintainers = with maintainers; [ aidalgol ];
63 broken = !(cudaPackages ? tensorrt) || !(cudaPackages ? cudnn);
64 };
65}