pkgs/development/python-modules/tensorrt/default.nix

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