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

   1 { lib
   2 , stdenv
   3 , buildPythonPackage
   4 , autoPatchelfHook
   5 , pythonRelaxDepsHook
   6 , onnxruntime
   7 , coloredlogs
   8 , numpy
   9 , packaging
  10 , oneDNN
  11 , re2
  12
  13 }:
  14
  15 # onnxruntime requires an older protobuf.
  16 # Doing an override in protobuf in the python-packages set
  17 # can give you a functioning Python package but note not
  18 # all Python packages will be compatible then.
  19 #
  20 # Because protobuf is not always needed we remove it
  21 # as a runtime dependency from our wheel.
  22 #
  23 # We do include here the non-Python protobuf so the shared libs
  24 # link correctly. If you do also want to include the Python
  25 # protobuf, you can add it to your Python env, but be aware
  26 # the version likely mismatches with what is used here.
  27
  28 buildPythonPackage {
  29   inherit (onnxruntime) pname version;
  30   format = "wheel";
  31   src = onnxruntime.dist;
  32
  33   unpackPhase = ''
  34     cp -r $src dist
  35     chmod +w dist
  36   '';
  37
  38   nativeBuildInputs = [
  39     pythonRelaxDepsHook
  40   ] ++ lib.optionals stdenv.isLinux [
  41     autoPatchelfHook
  42   ];
  43
  44   # This project requires fairly large dependencies such as sympy which we really don't always need.
  45   pythonRemoveDeps = [
  46     "flatbuffers"
  47     "protobuf"
  48     "sympy"
  49   ];
  50
  51   # Libraries are not linked correctly.
  52   buildInputs = [
  53     oneDNN
  54     re2
  55     onnxruntime.protobuf
  56   ] ++ lib.optionals onnxruntime.passthru.cudaSupport (with onnxruntime.passthru.cudaPackages; [
  57     libcublas # libcublasLt.so.XX libcublas.so.XX
  58     libcurand # libcurand.so.XX
  59     libcufft # libcufft.so.XX
  60     cudnn # libcudnn.soXX
  61     cuda_cudart # libcudart.so.XX
  62   ]);
  63
  64   propagatedBuildInputs = [
  65     coloredlogs
  66     # flatbuffers
  67     numpy
  68     packaging
  69     # protobuf
  70     # sympy
  71   ];
  72
  73   meta = onnxruntime.meta // { maintainers = with lib.maintainers; [ fridh ]; };
  74 }