onnnxruntime, python3Packages.onnxruntime: improve packaging

The Python bindings to onnxruntime were added by me in #193188.

Adding Python support this way is not a good way. Here a wheel was
created (which is fine) and installed in a python output. The
propagated-build-inputs file is put in the dev output. That's fine,
except that stdenv.mkDerivation does not automatically add dev when
python is included, because other outputs are only added when no
output is explicitly selected. This means that when you want to use
these bindings in another Python package, pip will complain it cannot
find the dependencies of the bindings.

In this PR, the onnxruntime derivation outputs a wheel in the dist output.
Then, in python-packages.nix we have a separate onnxruntime package
which installs the bindings.

The Python bindings have quite some dependencies which, depending on
your use case, are not required. Thus the dependency relax hook is
used to remove some of these dependencies.

Note there is also an issue with protobuf versions. The onnxruntime
bindings require an older protobuf and Python protobuf which we
cannot offer. Thus protobuf is also removed as Python dependency.

pkgs/development/libraries/onnxruntime/default.nix

@@ -15,11 +15,18 @@
 , boost
 , oneDNN
 , gtest
-, pythonSupport ? true
+, pythonSupport ? false
 , nsync
 , flatbuffers
 }:
 
+# Python Support
+#
+# When enabling Python support a wheel is made and stored in a `dist` output.
+# This wheel is then installed in a separate derivation.
+
+assert pythonSupport -> lib.versionOlder protobuf.version "3.20";
+
 let
   # prefetch abseil
   # Note: keep URL in sync with `cmake/external/abseil-cpp.cmake`
@@ -58,26 +65,26 @@ stdenv.mkDerivation rec {
     setuptools
     wheel
     pip
+    pythonOutputDistHook
   ]);
 
   buildInputs = [
     libpng
     zlib
-    protobuf
     howard-hinnant-date
     nlohmann_json
     boost
     oneDNN
-  ] ++ lib.optionals pythonSupport ([
-    flatbuffers
+    protobuf
+  ] ++ lib.optionals pythonSupport [
     nsync
-  ] ++ (with python3Packages; [
-    numpy
-    pybind11
-  ]));
+    python3Packages.numpy
+    python3Packages.pybind11
+  ];
 
   # TODO: build server, and move .so's to lib output
-  outputs = [ "out" "dev" ] ++ lib.optionals pythonSupport [ "python" ];
+  # Python's wheel is stored in a separate dist output
+  outputs = [ "out" "dev" ] ++ lib.optionals pythonSupport [ "dist" ];
 
   enableParallelBuilding = true;
 
@@ -116,10 +123,15 @@ stdenv.mkDerivation rec {
       ../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
   '';
 
+  passthru = {
+    inherit protobuf;
+    tests = lib.optionalAttrs pythonSupport {
+      python = python3Packages.onnxruntime;
+    };
+  };
+
   meta = with lib; {
     description = "Cross-platform, high performance scoring engine for ML models";
     longDescription = ''

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

@@ -0,0 +1,64 @@
+{ lib
+, buildPythonPackage
+, autoPatchelfHook
+, pythonRelaxDepsHook
+, onnxruntime
+, coloredlogs
+, numpy
+, packaging
+, oneDNN
+
+}:
+
+# onnxruntime requires an older protobuf.
+# Doing an override in protobuf in the python-packages set
+# can give you a functioning Python package but note not
+# all Python packages will be compatible then.
+#
+# Because protobuf is not always needed we remove it
+# as a runtime dependency from our wheel.
+#
+# We do include here the non-Python protobuf so the shared libs
+# link correctly. If you do also want to include the Python
+# protobuf, you can add it to your Python env, but be aware
+# the version likely mismatches with what is used here.
+
+buildPythonPackage {
+  inherit (onnxruntime) pname version;
+  format = "wheel";
+  src = onnxruntime.dist;
+
+  unpackPhase = ''
+    cp -r $src dist
+    chmod +w dist
+  '';
+
+  nativeBuildInputs = [
+    autoPatchelfHook
+    pythonRelaxDepsHook
+  ];
+
+  # This project requires fairly large dependencies such as sympy which we really don't always need.
+  pythonRemoveDeps = [
+    "flatbuffers"
+    "protobuf"
+    "sympy"
+  ];
+
+  # Libraries are not linked correctly.
+  buildInputs = [
+    oneDNN
+    onnxruntime.protobuf
+  ];
+
+  propagatedBuildInputs = [
+    coloredlogs
+    # flatbuffers
+    numpy
+    packaging
+    # protobuf
+    # sympy
+  ];
+
+  meta = onnxruntime.meta // { maintainers = with lib.maintainers; [ fridh ]; };
+}

pkgs/top-level/python-packages.nix

@@ -6436,10 +6436,12 @@ in {
 
   onnxconverter-common = callPackage ../development/python-modules/onnxconverter-common { };
 
-  onnxruntime = (toPythonModule (pkgs.onnxruntime.override {
-    python3Packages = self;
-    pythonSupport = true;
-  })).python;
+  onnxruntime = callPackage ../development/python-modules/onnxruntime {
+    onnxruntime = pkgs.onnxruntime.override {
+      python3Packages = self;
+      pythonSupport = true;
+    };
+  };
 
   onvif-zeep-async = callPackage ../development/python-modules/onvif-zeep-async { };