From 32709ae95a8a3c68623006ef2f741fe2ab867e17 Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Wed, 6 May 2026 11:12:04 +0000
Subject: [PATCH] python3Packages.sagemaker-mlflow: 0.3.0 -> 0.4.0

Diff: https://github.com/aws/sagemaker-mlflow/compare/v0.3.0...v0.4.0

Changelog: https://github.com/aws/sagemaker-mlflow/releases/tag/v0.4.0
---
 .../sagemaker-mlflow/default.nix              | 29 ++++++++++++-------
 1 file changed, 19 insertions(+), 10 deletions(-)

diff --git a/pkgs/development/python-modules/sagemaker-mlflow/default.nix b/pkgs/development/python-modules/sagemaker-mlflow/default.nix
index 3a84be94939d..06cfb393bdbe 100644
--- a/pkgs/development/python-modules/sagemaker-mlflow/default.nix
+++ b/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 ];
   };
-}
+})