fix: update mlflow version

Signed-off-by: s0nicboOm <[email protected]>
numaproj · Sep 15, 2023 · d20b1e3 · d20b1e3
1 parent d1789c0
commit d20b1e3
Show file tree

Hide file tree

Showing 9 changed files with 6,292 additions and 13 deletions.
diff --git a/.btoutput/app1:failure/data.csv b/.btoutput/app1:failure/data.csv
diff --git a/.btoutput/app1:failure/models.pt b/.btoutput/app1:failure/models.pt
diff --git a/.btoutput/app1:failure/output.csv b/.btoutput/app1:failure/output.csv
diff --git a/.btoutput/app1:failure/plot.png b/.btoutput/app1:failure/plot.png
diff --git a/.btoutput/failure/models.pt b/.btoutput/failure/models.pt
diff --git a/.btoutput/failure/output.csv b/.btoutput/failure/output.csv
diff --git a/poetry.lock b/poetry.lock
diff --git a/pyproject.toml b/pyproject.toml
@@ -40,7 +40,7 @@ pynumaflow = "~0.4"
 
 
 # extras
-mlflow-skinny = { version = "~2.6", optional = true }
+mlflow-skinny = { version = "^2.0", optional = true }
 redis = {extras = ["hiredis"], version = "^5.0", optional = true}
 boto3 = { version = "^1.24.64", optional = true }
 pydruid = {version = "^0.6", optional = true}

diff --git a/tests/registry/test_mlflow_registry.py b/tests/registry/test_mlflow_registry.py
@@ -124,12 +124,15 @@ def test_load_model_when_pytorch_model_exist2(self):
     @patch("mlflow.tracking.MlflowClient.transition_model_version_stage", mock_transition_stage)
     @patch("mlflow.tracking.MlflowClient.get_latest_versions", mock_get_model_version)
     @patch("mlflow.tracking.MlflowClient.search_model_versions", mock_list_of_model_version2)
-    @patch("mlflow.sklearn.load_model", Mock(return_value=StandardScaler()))
     @patch("mlflow.tracking.MlflowClient.get_run", Mock(return_value=return_empty_rundata()))
     @patch.object(
         MLflowRegistry,
         "load",
-        Mock(return_value=ArtifactData(artifact=StandardScaler(), extras=None, metadata={})),
+        Mock(
+            return_value=ArtifactData(
+                artifact=StandardScaler(), extras={"metric": ["error"]}, metadata={}
+            )
+        ),
     )
     def test_load_model_when_sklearn_model_exist(self):
         ml = MLflowRegistry(TRACKING_URI)