python3Packages.sagemaker-mlflow: 0.3.0 -> 0.4.0 (#517275)

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

@@ -0,0 +1,73 @@
+{
+  buildPythonPackage,
+  mlflow,
+
+  # build-system
+  setuptools,
+
+  # dependencies
+  cachetools,
+  click,
+  cloudpickle,
+  databricks-sdk,
+  fastapi,
+  gitpython,
+  importlib-metadata,
+  opentelemetry-api,
+  opentelemetry-proto,
+  opentelemetry-sdk,
+  packaging,
+  protobuf,
+  pydantic,
+  python-dotenv,
+  pyyaml,
+  requests,
+  sqlparse,
+  starlette,
+  typing-extensions,
+  uvicorn,
+}:
+
+buildPythonPackage (finalAttrs: {
+  pname = "mlflow-skinny";
+  inherit (mlflow) version src;
+  pyproject = true;
+  __structuredAttrs = true;
+
+  sourceRoot = "${finalAttrs.src.name}/libs/skinny";
+
+  build-system = [ setuptools ];
+
+  dependencies = [
+    cachetools
+    click
+    cloudpickle
+    databricks-sdk
+    fastapi
+    gitpython
+    importlib-metadata
+    opentelemetry-api
+    opentelemetry-proto
+    opentelemetry-sdk
+    packaging
+    protobuf
+    pydantic
+    python-dotenv
+    pyyaml
+    requests
+    sqlparse
+    starlette
+    typing-extensions
+    uvicorn
+  ];
+
+  pythonImportsCheck = [ "mlflow" ];
+
+  # No tests in the skinny subtree.
+  doCheck = false;
+
+  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";
+  };
+})

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

@@ -0,0 +1,58 @@
+{
+  lib,
+  buildPythonPackage,
+  mlflow,
+
+  # build-system
+  setuptools,
+
+  # dependencies
+  cachetools,
+  databricks-sdk,
+  opentelemetry-api,
+  opentelemetry-proto,
+  opentelemetry-sdk,
+  packaging,
+  protobuf,
+  pydantic,
+
+  # tests
+  pytestCheckHook,
+}:
+buildPythonPackage (finalAttrs: {
+  pname = "mlflow-tracing";
+  inherit (mlflow) version;
+  pyproject = true;
+  __structuredAttrs = true;
+
+  inherit (mlflow) src;
+
+  sourceRoot = "${finalAttrs.src.name}/libs/tracing";
+
+  build-system = [
+    setuptools
+  ];
+
+  dependencies = [
+    cachetools
+    databricks-sdk
+    opentelemetry-api
+    opentelemetry-proto
+    opentelemetry-sdk
+    packaging
+    protobuf
+    pydantic
+  ];
+
+  pythonImportsCheck = [ "mlflow.tracing" ];
+
+  # No tests
+  doCheck = false;
+
+  meta = {
+    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 ];
+  };
+})

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

@@ -9,195 +9,68 @@
   # dependencies
   aiohttp,
   alembic,
-  cachetools,
-  click,
-  cloudpickle,
   cryptography,
-  databricks-sdk,
   docker,
-  fastapi,
   flask,
   flask-cors,
-  gitpython,
   graphene,
   gunicorn,
   huey,
-  importlib-metadata,
-  jinja2,
-  markdown,
   matplotlib,
+  mlflow-skinny,
+  mlflow-tracing,
   numpy,
-  opentelemetry-api,
-  opentelemetry-proto,
-  opentelemetry-sdk,
-  packaging,
   pandas,
-  protobuf,
   pyarrow,
-  python-dotenv,
-  pyyaml,
-  requests,
   scikit-learn,
   scipy,
   skops,
   sqlalchemy,
-  sqlparse,
-  uvicorn,
-
-  # tests
-  azure-core,
-  azure-storage-blob,
-  azure-storage-file,
-  boto3,
-  botocore,
-  catboost,
-  datasets,
-  google-cloud-storage,
-  httpx,
-  jwt,
-  keras,
-  langchain,
-  librosa,
-  moto,
-  opentelemetry-exporter-otlp,
-  optuna,
-  pydantic,
-  pyspark,
-  pytestCheckHook,
-  pytorch-lightning,
-  sentence-transformers,
-  shap,
-  starlette,
-  statsmodels,
-  tensorflow,
-  torch,
-  transformers,
-  xgboost,
 }:
 
-buildPythonPackage rec {
+buildPythonPackage (finalAttrs: {
   pname = "mlflow";
-  version = "3.11.1";
+  version = "3.12.0";
   pyproject = true;
 
   src = fetchFromGitHub {
     owner = "mlflow";
     repo = "mlflow";
-    tag = "v${version}";
-    hash = "sha256-Oe6nBnnOz7MvGUNCcCGhHl6ZbyDfAhQ0LlfMBF4p6Hc=";
+    tag = "v${finalAttrs.version}";
+    hash = "sha256-OxhM+KCem0sb9cwtyzrUD/MGfoiiCfgU47qipYRDaFk=";
   };
 
-  pythonRelaxDeps = [
-    "cryptography"
-    "gunicorn"
-    "importlib_metadata"
-    "packaging"
-    "protobuf"
-    "pytz"
-    "pyarrow"
-  ];
+  # ppyproject.release.toml is the one shipped in the Pypi package, so we use it too.
+  postPatch = ''
+    mv pyproject.release.toml pyproject.toml
+  '';
 
   build-system = [ setuptools ];
 
   dependencies = [
     aiohttp
     alembic
-    cachetools
-    click
-    cloudpickle
     cryptography
-    databricks-sdk
     docker
-    fastapi
     flask
     flask-cors
-    gitpython
     graphene
     gunicorn
     huey
-    importlib-metadata
-    jinja2
-    markdown
     matplotlib
+    mlflow-skinny
+    mlflow-tracing
     numpy
-    opentelemetry-api
-    opentelemetry-proto
-    opentelemetry-sdk
-    packaging
     pandas
-    protobuf
     pyarrow
-    pydantic
-    python-dotenv
-    pyyaml
-    requests
     scikit-learn
     scipy
-    shap
     skops
     sqlalchemy
-    sqlparse
-    uvicorn
   ];
 
   pythonImportsCheck = [ "mlflow" ];
 
-  nativeCheckInputs = [
-    aiohttp
-    azure-core
-    azure-storage-blob
-    azure-storage-file
-    boto3
-    botocore
-    catboost
-    datasets
-    google-cloud-storage
-    httpx
-    jwt
-    keras
-    langchain
-    librosa
-    moto
-    opentelemetry-exporter-otlp
-    optuna
-    pydantic
-    pyspark
-    pytestCheckHook
-    pytorch-lightning
-    sentence-transformers
-    starlette
-    statsmodels
-    tensorflow
-    torch
-    transformers
-    uvicorn
-    xgboost
-  ];
-
-  disabledTestPaths = [
-    # Requires unpackaged `autogen`
-    "tests/autogen/test_autogen_autolog.py"
-
-    # Requires unpackaged `diviner`
-    "tests/diviner/test_diviner_model_export.py"
-
-    # Requires unpackaged `sktime`
-    "examples/sktime/test_sktime_model_export.py"
-
-    # Requires `fastai` which would cause a circular dependency
-    "tests/fastai/test_fastai_autolog.py"
-    "tests/fastai/test_fastai_model_export.py"
-
-    # Requires `spacy` which would cause a circular dependency
-    "tests/spacy/test_spacy_model_export.py"
-
-    # Requires `tensorflow.keras` which is not included in our outdated version of `tensorflow` (2.13.0)
-    "tests/gateway/providers/test_ai21labs.py"
-    "tests/tensorflow/test_keras_model_export.py"
-    "tests/tensorflow/test_keras_pyfunc_model_works_with_all_input_types.py"
-    "tests/tensorflow/test_mlflow_callback.py"
-  ];
-
   # I (@GaetanLepage) gave up at enabling tests:
   # - They require a lot of dependencies (some unpackaged);
   # - Many errors occur at collection time;
@@ -208,8 +81,10 @@ buildPythonPackage rec {
     description = "Open source platform for the machine learning lifecycle";
     mainProgram = "mlflow";
     homepage = "https://github.com/mlflow/mlflow";
-    changelog = "https://github.com/mlflow/mlflow/blob/${src.tag}/CHANGELOG.md";
+    changelog = "https://github.com/mlflow/mlflow/blob/${finalAttrs.src.tag}/CHANGELOG.md";
     license = lib.licenses.asl20;
-    maintainers = [ ];
+    maintainers = with lib.maintainers; [
+      GaetanLepage
+    ];
   };
-}
+})

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

@@ -8,23 +8,24 @@
 
   # dependencies
   boto3,
-  mlflow,
+  mlflow-skinny,
 
   # tests
   pytestCheckHook,
   scikit-learn,
 }:
 
-buildPythonPackage rec {
+buildPythonPackage (finalAttrs: {
   pname = "sagemaker-mlflow";
-  version = "0.3.0";
+  version = "0.4.0";
   pyproject = true;
+  __structuredAttrs = true;
 
   src = fetchFromGitHub {
     owner = "aws";
     repo = "sagemaker-mlflow";
-    tag = "v${version}";
-    hash = "sha256-riCoUpao9QIrQMb7r9stJO/xTSsDfL+uNS664W5GRQc=";
+    tag = "v${finalAttrs.version}";
+    hash = "sha256-QE40ZBW7N3GPC+eJqj5uzS3L+A6Wu2/LgHOiUsEXKMw=";
   };
 
   build-system = [
@@ -33,7 +34,7 @@ buildPythonPackage rec {
 
   dependencies = [
     boto3
-    mlflow
+    mlflow-skinny
   ];
 
   pythonImportsCheck = [ "sagemaker_mlflow" ];
@@ -56,16 +57,24 @@ buildPythonPackage rec {
     "test_presigned_url"
     "test_presigned_url_with_fields"
 
+    # Exercises a `sqlite://` model-registry store, only available with the
+    # sqlalchemy backend of the full `mlflow` package (not `mlflow-skinny`).
+    "test_store_instantiation_none"
+  ];
+
+  disabledTestPaths = [
+    # `from mlflow.models import infer_signature` fails to import at collection
+    # time: `infer_signature` is only available in the full `mlflow` package,
+    # not in `mlflow-skinny`. Also see:
     # https://github.com/aws/sagemaker-mlflow/issues/16
-    # TypeError: LogisticRegression.__init__() got an unexpected keyword argument 'multi_class'
-    "test_model_registry"
+    "test/integration/tests/test_model_registry.py"
   ];
 
   meta = {
     description = "MLFlow plugin for SageMaker";
     homepage = "https://github.com/aws/sagemaker-mlflow";
-    changelog = "https://github.com/aws/sagemaker-mlflow/releases/tag/v${version}";
+    changelog = "https://github.com/aws/sagemaker-mlflow/releases/tag/${finalAttrs.src.tag}";
     license = lib.licenses.asl20;
     maintainers = with lib.maintainers; [ GaetanLepage ];
   };
-}
+})

pkgs/top-level/python-packages.nix

@@ -10165,6 +10165,10 @@ self: super: with self; {
 
   mlflow = callPackage ../development/python-modules/mlflow { };
 
+  mlflow-skinny = callPackage ../development/python-modules/mlflow-skinny { };
+
+  mlflow-tracing = callPackage ../development/python-modules/mlflow-tracing { };
+
   mlrose = callPackage ../development/python-modules/mlrose { };
 
   mlt = toPythonModule (