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 all commits
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
20 changes: 9 additions & 11 deletions bentoml/_internal/frameworks/mlflow.py
View file
Open in desktop
Expand Up @@ -24,6 +24,7 @@

if TYPE_CHECKING:
import mlflow
import mlflow.models
else:
mlflow = LazyLoader(
"mlflow",
Expand Down Expand Up @@ -205,19 +206,16 @@ 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
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) # type: ignore (local_path is bound)
# 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 +256,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