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;
14in
15buildPythonPackage rec {
16 pname = "tensorrt";
17 version = buildVersion;
18
19 src = cudaPackages.tensorrt.src;
20
21 format = "wheel";
22 # We unpack the wheel ourselves because of the odd packaging.
23 dontUseWheelUnpack = true;
24
25 nativeBuildInputs = [
26 unzip
27 autoPatchelfHook
28 autoAddDriverRunpath
29 ];
30
31 preUnpack = ''
32 mkdir -p dist
33 tar --strip-components=2 -xf "$src" --directory=dist \
34 "TensorRT-${buildVersion}/python/tensorrt-${buildVersion}-cp${pyVersion}-none-linux_x86_64.whl"
35 '';
36
37 sourceRoot = ".";
38
39 buildInputs = [
40 cudaPackages.cudnn
41 cudaPackages.tensorrt
42 ];
43
44 pythonImportsCheck = [ "tensorrt" ];
45
46 meta = with lib; {
47 description = "Python bindings for TensorRT, a high-performance deep learning interface";
48 homepage = "https://developer.nvidia.com/tensorrt";
49 license = licenses.unfree;
50 platforms = [ "x86_64-linux" ];
51 maintainers = with maintainers; [ aidalgol ];
52 broken = !(cudaPackages ? tensorrt) || !(cudaPackages ? cudnn);
53 };
54}