From 24e44f74b23d2af9a2716ca70e5971138df57bac Mon Sep 17 00:00:00 2001
From: Aaron Pham <29749331+aarnphm@users.noreply.github.com>
Date: Fri, 26 Aug 2022 00:01:22 -0700
Subject: [PATCH 1/2] qa: cleanup MLflow

Signed-off-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>
---
 bentoml/_internal/frameworks/mlflow.py | 27 ++++++++++++++------------
 1 file changed, 15 insertions(+), 12 deletions(-)

diff --git a/bentoml/_internal/frameworks/mlflow.py b/bentoml/_internal/frameworks/mlflow.py
index f401563dbe8..4b4768c15bb 100644
--- a/bentoml/_internal/frameworks/mlflow.py
+++ b/bentoml/_internal/frameworks/mlflow.py
@@ -4,7 +4,6 @@
 import shutil
 import typing as t
 import logging
-import tempfile
 from typing import TYPE_CHECKING
 
 import bentoml
@@ -24,6 +23,7 @@
 
 if TYPE_CHECKING:
     import mlflow
+    import mlflow.models
 else:
     mlflow = LazyLoader(
         "mlflow",
@@ -185,6 +185,8 @@ def import_model(
         metadata=metadata,
         context=context,
     ) as bento_model:
+        import tempfile
+
         from mlflow.models import Model as MLflowModel
         from mlflow.pyfunc import FLAVOR_NAME as PYFUNC_FLAVOR_NAME
         from mlflow.models.model import MLMODEL_FILE_NAME
@@ -205,19 +207,20 @@ def import_model(
             # For MLflow < 1.25
             from mlflow.tracking.artifact_utils import _download_artifact_from_uri
 
-            local_path = _download_artifact_from_uri(
+            local_path: str = _download_artifact_from_uri(
                 artifact_uri=model_uri, output_path=download_dir
             )
-
-        mlflow_model_path = bento_model.path_of(MLFLOW_MODEL_FOLDER)
-        # Rename model folder from original artifact name to fixed "mlflow_model"
-        shutil.move(local_path, mlflow_model_path)
-        # If the temp dir we created still exists now, we never needed it because
-        # the provided mlflow url must have provided enough path information.  Just
-        # delete it.  If it's not here, it means we needed it but it's been renamed
-        # by now so we don't need to remove it.
-        if os.path.isdir(download_dir):
+        finally:
+            if TYPE_CHECKING:
+                # We need to trick type checker here, since local_path
+                # will always be set
+                local_path = ""
+            mlflow_model_path = bento_model.path_of(MLFLOW_MODEL_FOLDER)
+            # Rename model folder from original artifact name to fixed "mlflow_model"
+            shutil.move(local_path, mlflow_model_path)
+            # Remove the tempdir
             shutil.rmtree(download_dir)
+
         mlflow_model_file = os.path.join(mlflow_model_path, MLMODEL_FILE_NAME)
 
         if not os.path.exists(mlflow_model_file):
@@ -258,7 +261,7 @@ def __init__(self):
             input_spec=None,
             output_spec=None,
         )
-        def predict(self, input_data):
+        def predict(self, input_data: t.Any) -> t.Any:
             return self.model.predict(input_data)
 
     return MLflowPyfuncRunnable

From 7d208e90e700e1cfbbc056e1c1a10f530725a79a Mon Sep 17 00:00:00 2001
From: Aaron Pham <29749331+aarnphm@users.noreply.github.com>
Date: Fri, 26 Aug 2022 15:25:25 -0700
Subject: [PATCH 2/2] revert: lazy

Signed-off-by: Aaron Pham <29749331+aarnphm@users.noreply.github.com>
---
 bentoml/_internal/frameworks/mlflow.py | 9 ++-------
 1 file changed, 2 insertions(+), 7 deletions(-)

diff --git a/bentoml/_internal/frameworks/mlflow.py b/bentoml/_internal/frameworks/mlflow.py
index 4b4768c15bb..96e5e0ac8c8 100644
--- a/bentoml/_internal/frameworks/mlflow.py
+++ b/bentoml/_internal/frameworks/mlflow.py
@@ -4,6 +4,7 @@
 import shutil
 import typing as t
 import logging
+import tempfile
 from typing import TYPE_CHECKING
 
 import bentoml
@@ -185,8 +186,6 @@ def import_model(
         metadata=metadata,
         context=context,
     ) as bento_model:
-        import tempfile
-
         from mlflow.models import Model as MLflowModel
         from mlflow.pyfunc import FLAVOR_NAME as PYFUNC_FLAVOR_NAME
         from mlflow.models.model import MLMODEL_FILE_NAME
@@ -211,13 +210,9 @@ def import_model(
                 artifact_uri=model_uri, output_path=download_dir
             )
         finally:
-            if TYPE_CHECKING:
-                # We need to trick type checker here, since local_path
-                # will always be set
-                local_path = ""
             mlflow_model_path = bento_model.path_of(MLFLOW_MODEL_FOLDER)
             # Rename model folder from original artifact name to fixed "mlflow_model"
-            shutil.move(local_path, mlflow_model_path)
+            shutil.move(local_path, mlflow_model_path)  # type: ignore (local_path is bound)
             # Remove the tempdir
             shutil.rmtree(download_dir)