mlflow{,-server}: fix server startup and missing UI (#519345)

maintainers/maintainer-list.nix

@@ -10141,6 +10141,12 @@
     githubId = 6893840;
     name = "Yacine Hmito";
   };
+  gquetel = {
+    email = "gregor.quetel@telecom-paris.fr";
+    github = "gquetel";
+    githubId = 48437427;
+    name = "Grégor Quetel";
+  };
   gracicot = {
     email = "dev@gracicot.com";
     matrix = "@gracicot-59e8f173d73408ce4f7ac803:gitter.im";

pkgs/by-name/ml/mlflow-server/package.nix

@@ -1,37 +1,10 @@
-{ python3Packages, writers }:
+{ python3Packages }:
 
-let
-  py = python3Packages;
-
-  gunicornScript = writers.writePython3 "gunicornMlflow" { } ''
-    import re
-    import sys
-    from gunicorn.app.wsgiapp import run
-    if __name__ == '__main__':
-        sys.argv[0] = re.sub(r'(-script\.pyw|\.exe)?$', ''', sys.argv[0])
-        sys.exit(run())
-  '';
-in
-py.toPythonApplication (
-  py.mlflow.overridePythonAttrs (old: {
-
-    propagatedBuildInputs = old.dependencies ++ [
-      py.boto3
-      py.mysqlclient
+python3Packages.toPythonApplication (
+  python3Packages.mlflow.overridePythonAttrs (old: {
+    propagatedBuildInputs = (old.propagatedBuildInputs or [ ]) ++ [
+      python3Packages.boto3
+      python3Packages.mysqlclient
     ];
-
-    postPatch = (old.postPatch or "") + ''
-      cat mlflow/utils/process.py
-
-      substituteInPlace mlflow/utils/process.py --replace-fail \
-        "process = subprocess.Popen(" \
-        "cmd[0]='${gunicornScript}'; process = subprocess.Popen("
-    '';
-
-    postInstall = ''
-      gpath=$out/bin/gunicornMlflow
-      cp ${gunicornScript} $gpath
-      chmod 555 $gpath
-    '';
   })
 )

pkgs/development/python-modules/mlflow-skinny/default.nix

@@ -1,5 +1,7 @@
 {
+  lib,
   buildPythonPackage,
+  fetchFromGitHub,
   mlflow,
 
   # build-system
@@ -30,10 +32,17 @@
 
 buildPythonPackage (finalAttrs: {
   pname = "mlflow-skinny";
-  inherit (mlflow) version src;
+  inherit (mlflow) version;
   pyproject = true;
   __structuredAttrs = true;
 
+  src = fetchFromGitHub {
+    owner = "mlflow";
+    repo = "mlflow";
+    rev = "v${finalAttrs.version}";
+    hash = "sha256-OxhM+KCem0sb9cwtyzrUD/MGfoiiCfgU47qipYRDaFk=";
+  };
+
   sourceRoot = "${finalAttrs.src.name}/libs/skinny";
 
   build-system = [ setuptools ];
@@ -69,5 +78,6 @@ buildPythonPackage (finalAttrs: {
   meta = mlflow.meta // {
     description = "Lightweight version of MLflow that is designed to minimize package size";
     homepage = "https://github.com/mlflow/mlflow/tree/master/libs/skinny";
+    sourceProvenance = with lib.sourceTypes; [ fromSource ];
   };
 })

pkgs/development/python-modules/mlflow-tracing/default.nix

@@ -1,6 +1,7 @@
 {
   lib,
   buildPythonPackage,
+  fetchFromGitHub,
   mlflow,
 
   # build-system
@@ -15,9 +16,6 @@
   packaging,
   protobuf,
   pydantic,
-
-  # tests
-  pytestCheckHook,
 }:
 buildPythonPackage (finalAttrs: {
   pname = "mlflow-tracing";
@@ -25,13 +23,16 @@ buildPythonPackage (finalAttrs: {
   pyproject = true;
   __structuredAttrs = true;
 
-  inherit (mlflow) src;
+  src = fetchFromGitHub {
+    owner = "mlflow";
+    repo = "mlflow";
+    rev = "v${finalAttrs.version}";
+    hash = "sha256-OxhM+KCem0sb9cwtyzrUD/MGfoiiCfgU47qipYRDaFk=";
+  };
 
   sourceRoot = "${finalAttrs.src.name}/libs/tracing";
 
-  build-system = [
-    setuptools
-  ];
+  build-system = [ setuptools ];
 
   dependencies = [
     cachetools
@@ -53,6 +54,9 @@ buildPythonPackage (finalAttrs: {
     description = "Open-Source SDK for observability and monitoring GenAI applications";
     homepage = "https://github.com/mlflow/mlflow/tree/master/libs/tracing";
     inherit (mlflow.meta) license;
-    maintainers = with lib.maintainers; [ GaetanLepage ];
+    maintainers = with lib.maintainers; [
+      GaetanLepage
+      gquetel
+    ];
   };
 })

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

@@ -1,10 +1,7 @@
 {
   lib,
   buildPythonPackage,
-  fetchFromGitHub,
-
-  # build-system
-  setuptools,
+  fetchPypi,
 
   # dependencies
   aiohttp,
@@ -31,26 +28,32 @@
 buildPythonPackage (finalAttrs: {
   pname = "mlflow";
   version = "3.12.0";
-  pyproject = true;
+  format = "wheel";
   __structuredAttrs = true;
 
-  src = fetchFromGitHub {
-    owner = "mlflow";
-    repo = "mlflow";
-    tag = "v${finalAttrs.version}";
-    hash = "sha256-OxhM+KCem0sb9cwtyzrUD/MGfoiiCfgU47qipYRDaFk=";
+  # We build from the PyPI wheel rather than fetchFromGitHub, because the mlflow-server
+  # JS UI is absent from GitHub but provided in the wheel.
+  src = fetchPypi {
+    pname = "mlflow";
+    inherit (finalAttrs) version;
+    format = "wheel";
+    dist = "py3";
+    python = "py3";
+    hash = "sha256-4cKO1MSFV8xSx2bxfxylgmdT3fJB1D8w+ZxF9+prPOA=";
   };
 
-  # ppyproject.release.toml is the one shipped in the Pypi package, so we use it too.
-  postPatch = ''
-    mv pyproject.release.toml pyproject.toml
+  # Nix-wrapped python populates sys.path via NIX_PYTHONPATH/site hooks,
+  # but PYTHONPATH stays unset in os.environ. mlflow spawns the server
+  # in a subprocess with a curated env, so without this patch the child
+  # interpreter cannot import uvicorn / mlflow itself.
+  postInstall = ''
+    patch -p1 -d "$out/lib/python"*/site-packages < ${./subprocess-pythonpath.patch}
   '';
 
-  build-system = [ setuptools ];
-
   pythonRelaxDeps = [
     "cryptography"
   ];
+
   dependencies = [
     aiohttp
     alembic
@@ -85,10 +88,15 @@ buildPythonPackage (finalAttrs: {
     description = "Open source platform for the machine learning lifecycle";
     mainProgram = "mlflow";
     homepage = "https://github.com/mlflow/mlflow";
-    changelog = "https://github.com/mlflow/mlflow/blob/${finalAttrs.src.tag}/CHANGELOG.md";
+    changelog = "https://github.com/mlflow/mlflow/blob/v${finalAttrs.version}/CHANGELOG.md";
     license = lib.licenses.asl20;
+    # Build from wheel which contains pure Python and pre-built JS bundle.
+    sourceProvenance = with lib.sourceTypes; [
+      binaryBytecode
+    ];
     maintainers = with lib.maintainers; [
       GaetanLepage
+      gquetel
     ];
   };
 })

pkgs/development/python-modules/mlflow/subprocess-pythonpath.patch

@@ -0,0 +1,14 @@
+--- a/mlflow/server/__init__.py
++++ b/mlflow/server/__init__.py
+@@ -471,6 +471,11 @@
+             else tempfile.mkdtemp()
+         )
++    # In Nix-packaged environments sys.path is populated by wrappers but
++    # PYTHONPATH is never set in os.environ, so subprocesses (uvicorn,
++    # gunicorn) cannot find packages. Propagate it when not already set.
++    if "PYTHONPATH" not in os.environ:
++        env_map.setdefault("PYTHONPATH", os.pathsep.join(p for p in sys.path if p))
+
+     server_proc = _exec_cmd(
+         full_command,
+         extra_env=env_map,