Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

qa: cleanup MLflow #2945

Merged
merged 2 commits into from Sep 27, 2022
Merged

qa: cleanup MLflow #2945

Changes from 1 commit
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
24e44f7
qa: cleanup MLflow
aarnphm Aug 26, 2022
7d208e9
revert: lazy
aarnphm Aug 26, 2022
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
27 changes: 15 additions & 12 deletions bentoml/_internal/frameworks/mlflow.py
View file
Open in desktop
Expand Up @@ -4,7 +4,6 @@
import shutil
import typing as t
import logging
import tempfile
parano marked this conversation as resolved.
Show resolved Hide resolved
from typing import TYPE_CHECKING

import bentoml
Expand All @@ -24,6 +23,7 @@

if TYPE_CHECKING:
import mlflow
import mlflow.models
else:
mlflow = LazyLoader(
"mlflow",
Expand Down Expand Up @@ -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
Expand All @@ -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:
ssheng marked this conversation as resolved.
Show resolved Hide resolved
if TYPE_CHECKING:
# We need to trick type checker here, since local_path
# will always be set
local_path = ""
aarnphm marked this conversation as resolved.
Show resolved Hide resolved
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):
Expand Down Expand Up @@ -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
Expand Down