def test_log_model_without_specified_conda_env_uses_default_env_with_expected_dependencies(
sklearn_knn_model, main_scoped_model_class):
sklearn_artifact_path = "sk_model"
with mlflow.start_run():
mlflow.sklearn.log_model(sk_model=sklearn_knn_model,
artifact_path=sklearn_artifact_path)
sklearn_run_id = mlflow.active_run().info.run_id
pyfunc_artifact_path = "pyfunc_model"
with mlflow.start_run():
mlflow.pyfunc.log_model(
artifact_path=pyfunc_artifact_path,
artifacts={
"sk_model":
utils_get_artifact_uri(artifact_path=sklearn_artifact_path,
run_id=sklearn_run_id)
},
python_model=main_scoped_model_class(predict_fn=None))
pyfunc_model_path = _download_artifact_from_uri(
"runs:/{run_id}/{artifact_path}".format(
run_id=mlflow.active_run().info.run_id,
artifact_path=pyfunc_artifact_path))
pyfunc_conf = _get_flavor_configuration(
model_path=pyfunc_model_path, flavor_name=mlflow.pyfunc.FLAVOR_NAME)
conda_env_path = os.path.join(pyfunc_model_path,
pyfunc_conf[mlflow.pyfunc.ENV])
with open(conda_env_path, "r") as f:
conda_env = yaml.safe_load(f)
assert conda_env == mlflow.pyfunc.model.get_default_conda_env()
def test_log_model_persists_specified_conda_env_in_mlflow_model_directory(
sklearn_knn_model, main_scoped_model_class, pyfunc_custom_env):
sklearn_artifact_path = "sk_model"
with mlflow.start_run():
mlflow.sklearn.log_model(sk_model=sklearn_knn_model, artifact_path=sklearn_artifact_path)
sklearn_run_id = mlflow.active_run().info.run_id
pyfunc_artifact_path = "pyfunc_model"
with mlflow.start_run():
mlflow.pyfunc.log_model(artifact_path=pyfunc_artifact_path,
artifacts={
"sk_model": utils_get_artifact_uri(
artifact_path=sklearn_artifact_path,
run_id=sklearn_run_id)
},
python_model=main_scoped_model_class(predict_fn=None),
conda_env=pyfunc_custom_env)
pyfunc_model_path = _download_artifact_from_uri("runs:/{run_id}/{artifact_path}".format(
run_id=mlflow.active_run().info.run_id, artifact_path=pyfunc_artifact_path))
pyfunc_conf = _get_flavor_configuration(
model_path=pyfunc_model_path, flavor_name=mlflow.pyfunc.FLAVOR_NAME)
saved_conda_env_path = os.path.join(pyfunc_model_path, pyfunc_conf[mlflow.pyfunc.ENV])
assert os.path.exists(saved_conda_env_path)
assert saved_conda_env_path != pyfunc_custom_env
with open(pyfunc_custom_env, "r") as f:
pyfunc_custom_env_parsed = yaml.safe_load(f)
with open(saved_conda_env_path, "r") as f:
saved_conda_env_parsed = yaml.safe_load(f)
assert saved_conda_env_parsed == pyfunc_custom_env_parsed
def test_model_log_load(sklearn_knn_model, main_scoped_model_class, iris_data):
sklearn_artifact_path = "sk_model"
with mlflow.start_run():
mlflow.sklearn.log_model(sk_model=sklearn_knn_model,
artifact_path=sklearn_artifact_path)
sklearn_run_id = mlflow.active_run().info.run_uuid
def test_predict(sk_model, model_input):
return sk_model.predict(model_input) * 2
pyfunc_artifact_path = "pyfunc_model"
with mlflow.start_run():
mlflow.pyfunc.log_model(
artifact_path=pyfunc_artifact_path,
artifacts={
"sk_model":
utils_get_artifact_uri(artifact_path=sklearn_artifact_path,
run_id=sklearn_run_id)
},
python_model=main_scoped_model_class(test_predict))
pyfunc_run_id = mlflow.active_run().info.run_uuid
loaded_pyfunc_model = mlflow.pyfunc.load_pyfunc(path=pyfunc_artifact_path,
run_id=pyfunc_run_id)
np.testing.assert_array_equal(
loaded_pyfunc_model.predict(model_input=iris_data[0]),
test_predict(sk_model=sklearn_knn_model, model_input=iris_data[0]))
def test_model_log_persists_requirements_in_mlflow_model_directory(
sklearn_knn_model, main_scoped_model_class, pyfunc_custom_env):
sklearn_artifact_path = "sk_model"
with mlflow.start_run():
mlflow.sklearn.log_model(sk_model=sklearn_knn_model,
artifact_path=sklearn_artifact_path)
sklearn_run_id = mlflow.active_run().info.run_id
pyfunc_artifact_path = "pyfunc_model"
with mlflow.start_run():
mlflow.pyfunc.log_model(
artifact_path=pyfunc_artifact_path,
artifacts={
"sk_model":
utils_get_artifact_uri(artifact_path=sklearn_artifact_path,
run_id=sklearn_run_id)
},
python_model=main_scoped_model_class(predict_fn=None),
conda_env=pyfunc_custom_env,
)
pyfunc_model_path = _download_artifact_from_uri(
"runs:/{run_id}/{artifact_path}".format(
run_id=mlflow.active_run().info.run_id,
artifact_path=pyfunc_artifact_path))
saved_pip_req_path = os.path.join(pyfunc_model_path, "requirements.txt")
_compare_conda_env_requirements(pyfunc_custom_env, saved_pip_req_path)
def test_log_model_without_specified_conda_env_uses_default_env_with_expected_dependencies(
sklearn_knn_model, main_scoped_model_class
):
sklearn_artifact_path = "sk_model"
with mlflow.start_run():
mlflow.sklearn.log_model(sk_model=sklearn_knn_model, artifact_path=sklearn_artifact_path)
sklearn_run_id = mlflow.active_run().info.run_id
pyfunc_artifact_path = "pyfunc_model"
with mlflow.start_run():
mlflow.pyfunc.log_model(
artifact_path=pyfunc_artifact_path,
artifacts={
"sk_model": utils_get_artifact_uri(
artifact_path=sklearn_artifact_path, run_id=sklearn_run_id
)
},
python_model=main_scoped_model_class(predict_fn=None),
)
model_uri = mlflow.get_artifact_uri(pyfunc_artifact_path)
_assert_pip_requirements(model_uri, mlflow.pyfunc.get_default_pip_requirements())
