pyrox.dev
1{ lib
2, python
3, buildPythonPackage
4, autoPatchelfHook
5, unzip
6, cudaPackages
7}:
8
9let
10 pyVersion = "${lib.versions.major python.version}${lib.versions.minor python.version}";
11in
12buildPythonPackage rec {
13 pname = "tensorrt";
14 version = cudaPackages.tensorrt.version;
15
16 src = cudaPackages.tensorrt.src;
17
18 format = "wheel";
19 # We unpack the wheel ourselves because of the odd packaging.
20 dontUseWheelUnpack = true;
21
22 nativeBuildInputs = [
23 unzip
24 autoPatchelfHook
25 cudaPackages.autoAddOpenGLRunpathHook
26 ];
27
28 preUnpack = ''
29 mkdir -p dist
30 tar --strip-components=2 -xf "$src" --directory=dist \
31 "TensorRT-${version}/python/tensorrt-${version}-cp${pyVersion}-none-linux_x86_64.whl"
32 '';
33
34 sourceRoot = ".";
35
36 buildInputs = [
37 cudaPackages.cudnn
38 cudaPackages.tensorrt
39 ];
40
41 pythonImportsCheck = [
42 "tensorrt"
43 ];
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 };
52}