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