onnxruntime: add Python support

2 files changed, 31 insertions, 5 deletions

diff --git a/pkgs/development/libraries/onnxruntime/default.nix b/pkgs/development/libraries/onnxruntime/default.nix
index 77730cb8820..e297f835f32 100644
--- a/pkgs/development/libraries/onnxruntime/default.nix
+++ b/pkgs/development/libraries/onnxruntime/default.nix
@@ -5,7 +5,7 @@
 , fetchurl
 , pkg-config
 , cmake
-, python3
+, python3Packages
 , libpng
 , zlib
 , eigen
@@ -15,6 +15,9 @@
 , boost
 , oneDNN
 , gtest
+, pythonSupport ? true
+, nsync
+, flatbuffers
 }:
 
 let
@@ -49,9 +52,13 @@ stdenv.mkDerivation rec {
   nativeBuildInputs = [
     cmake
     pkg-config
-    python3
+    python3Packages.python
     gtest
-  ];
+  ] ++ lib.optionals pythonSupport (with python3Packages; [
+    setuptools
+    wheel
+    pip
+  ]);
 
   buildInputs = [
     libpng
@@ -61,10 +68,16 @@ stdenv.mkDerivation rec {
     nlohmann_json
     boost
     oneDNN
-  ];
+  ] ++ lib.optionals pythonSupport ([
+    flatbuffers
+    nsync
+  ] ++ (with python3Packages; [
+    numpy
+    pybind11
+  ]));
 
   # TODO: build server, and move .so's to lib output
-  outputs = [ "out" "dev" ];
+  outputs = [ "out" "dev" ] ++ lib.optionals pythonSupport [ "python" ];
 
   enableParallelBuilding = true;
 
@@ -79,6 +92,8 @@ stdenv.mkDerivation rec {
     "-Donnxruntime_USE_MPI=ON"
     "-Deigen_SOURCE_PATH=${eigen.src}"
     "-Donnxruntime_USE_DNNL=YES"
+  ] ++ lib.optionals pythonSupport [
+    "-Donnxruntime_ENABLE_PYTHON=ON"
   ];
 
   doCheck = true;
@@ -91,12 +106,18 @@ stdenv.mkDerivation rec {
       --replace '$'{prefix}/@CMAKE_INSTALL_ @CMAKE_INSTALL_
   '';
 
+  postBuild = lib.optionalString pythonSupport ''
+    ${python3Packages.python.interpreter} ../setup.py bdist_wheel
+  '';
+
   postInstall = ''
     # perform parts of `tools/ci_build/github/linux/copy_strip_binary.sh`
     install -m644 -Dt $out/include \
       ../include/onnxruntime/core/framework/provider_options.h \
       ../include/onnxruntime/core/providers/cpu/cpu_provider_factory.h \
       ../include/onnxruntime/core/session/onnxruntime_*.h
+  '' + lib.optionalString pythonSupport ''
+    pip install dist/*.whl --no-index --no-warn-script-location --prefix="$python" --no-cache --no-deps
   '';
 
   meta = with lib; {
diff --git a/pkgs/top-level/python-packages.nix b/pkgs/top-level/python-packages.nix
index 9a57da59836..f2ce500f808 100644
--- a/pkgs/top-level/python-packages.nix
+++ b/pkgs/top-level/python-packages.nix
@@ -6388,6 +6388,11 @@ in {
 
   onnx = callPackage ../development/python-modules/onnx { };
 
+  onnxruntime = (toPythonModule (pkgs.onnxruntime.override {
+    python3Packages = self;
+    pythonSupport = true;
+  })).python;
+
   onvif-zeep-async = callPackage ../development/python-modules/onvif-zeep-async { };
 
   oocsi = callPackage ../development/python-modules/oocsi { };