def test_log_model_with_extra_pip_requirements(tmpdir):
ols = ols_model()
default_reqs = mlflow.statsmodels.get_default_pip_requirements()
# Path to a requirements file
req_file = tmpdir.join("requirements.txt")
req_file.write("a")
with mlflow.start_run():
mlflow.statsmodels.log_model(ols.model, "model", extra_pip_requirements=req_file.strpath)
_assert_pip_requirements(mlflow.get_artifact_uri("model"), ["mlflow", *default_reqs, "a"])
# List of requirements
with mlflow.start_run():
mlflow.statsmodels.log_model(
ols.model, "model", extra_pip_requirements=[f"-r {req_file.strpath}", "b"]
)
_assert_pip_requirements(
mlflow.get_artifact_uri("model"), ["mlflow", *default_reqs, "a", "b"]
)
# Constraints file
with mlflow.start_run():
mlflow.statsmodels.log_model(
ols.model, "model", extra_pip_requirements=[f"-c {req_file.strpath}", "b"]
)
_assert_pip_requirements(
mlflow.get_artifact_uri("model"),
["mlflow", *default_reqs, "b", "-c constraints.txt"],
["a"],
)
def test_log_model_with_pip_requirements(main_scoped_model_class, tmpdir):
python_model = main_scoped_model_class(predict_fn=None)
# Path to a requirements file
req_file = tmpdir.join("requirements.txt")
req_file.write("a")
with mlflow.start_run():
mlflow.pyfunc.log_model(
"model", python_model=python_model, pip_requirements=req_file.strpath
)
_assert_pip_requirements(mlflow.get_artifact_uri("model"), ["mlflow", "a"])
# List of requirements
with mlflow.start_run():
mlflow.pyfunc.log_model(
"model", python_model=python_model, pip_requirements=[f"-r {req_file.strpath}", "b"]
)
_assert_pip_requirements(mlflow.get_artifact_uri("model"), ["mlflow", "a", "b"])
# Constraints file
with mlflow.start_run():
mlflow.pyfunc.log_model(
"model", python_model=python_model, pip_requirements=[f"-c {req_file.strpath}", "b"]
)
_assert_pip_requirements(
mlflow.get_artifact_uri("model"), ["mlflow", "b", "-c constraints.txt"], ["a"]
)
def test_log_model_with_extra_pip_requirements(sklearn_knn_model, tmpdir):
default_reqs = mlflow.sklearn.get_default_pip_requirements(
include_cloudpickle=True)
# Path to a requirements file
req_file = tmpdir.join("requirements.txt")
req_file.write("a")
with mlflow.start_run():
mlflow.sklearn.log_model(sklearn_knn_model.model,
"model",
extra_pip_requirements=req_file.strpath)
_assert_pip_requirements(mlflow.get_artifact_uri("model"),
["mlflow", *default_reqs, "a"])
# List of requirements
with mlflow.start_run():
mlflow.sklearn.log_model(
sklearn_knn_model.model,
"model",
extra_pip_requirements=[f"-r {req_file.strpath}", "b"])
_assert_pip_requirements(mlflow.get_artifact_uri("model"),
["mlflow", *default_reqs, "a", "b"])
# Constraints file
with mlflow.start_run():
mlflow.sklearn.log_model(
sklearn_knn_model.model,
"model",
extra_pip_requirements=[f"-c {req_file.strpath}", "b"])
_assert_pip_requirements(
mlflow.get_artifact_uri("model"),
["mlflow", *default_reqs, "b", "-c constraints.txt"],
["a"],
)
def test_log_model_with_pip_requirements(sklearn_knn_model, tmpdir):
# Path to a requirements file
req_file = tmpdir.join("requirements.txt")
req_file.write("a")
with mlflow.start_run():
mlflow.sklearn.log_model(sklearn_knn_model.model,
"model",
pip_requirements=req_file.strpath)
_assert_pip_requirements(mlflow.get_artifact_uri("model"),
["mlflow", "a"],
strict=True)
# List of requirements
with mlflow.start_run():
mlflow.sklearn.log_model(
sklearn_knn_model.model,
"model",
pip_requirements=[f"-r {req_file.strpath}", "b"])
_assert_pip_requirements(mlflow.get_artifact_uri("model"),
["mlflow", "a", "b"],
strict=True)
# Constraints file
with mlflow.start_run():
mlflow.sklearn.log_model(
sklearn_knn_model.model,
"model",
pip_requirements=[f"-c {req_file.strpath}", "b"])
_assert_pip_requirements(
mlflow.get_artifact_uri("model"),
["mlflow", "b", "-c constraints.txt"],
["a"],
strict=True,
)
def test_save_model_with_extra_pip_requirements(sequential_model, tmpdir):
default_reqs = mlflow.pytorch.get_default_pip_requirements()
# Path to a requirements file
tmpdir1 = tmpdir.join("1")
req_file = tmpdir.join("requirements.txt")
req_file.write("a")
mlflow.pytorch.save_model(sequential_model,
tmpdir1.strpath,
extra_pip_requirements=req_file.strpath)
_assert_pip_requirements(tmpdir1.strpath, ["mlflow", *default_reqs, "a"])
# List of requirements
tmpdir2 = tmpdir.join("2")
mlflow.pytorch.save_model(
sequential_model,
tmpdir2.strpath,
extra_pip_requirements=[f"-r {req_file.strpath}", "b"])
_assert_pip_requirements(tmpdir2.strpath,
["mlflow", *default_reqs, "a", "b"])
# Constraints file
tmpdir3 = tmpdir.join("3")
mlflow.pytorch.save_model(
sequential_model,
tmpdir3.strpath,
extra_pip_requirements=[f"-c {req_file.strpath}", "b"])
_assert_pip_requirements(
tmpdir3.strpath, ["mlflow", *default_reqs, "b", "-c constraints.txt"],
["a"])
def test_save_model_with_pip_requirements(xgb_model, tmpdir):
# Path to a requirements file
tmpdir1 = tmpdir.join("1")
req_file = tmpdir.join("requirements.txt")
req_file.write("a")
mlflow.xgboost.save_model(xgb_model.model,
tmpdir1.strpath,
pip_requirements=req_file.strpath)
_assert_pip_requirements(tmpdir1.strpath, ["mlflow", "a"], strict=True)
# List of requirements
tmpdir2 = tmpdir.join("2")
mlflow.xgboost.save_model(xgb_model.model,
tmpdir2.strpath,
pip_requirements=[f"-r {req_file.strpath}", "b"])
_assert_pip_requirements(tmpdir2.strpath, ["mlflow", "a", "b"],
strict=True)
# Constraints file
tmpdir3 = tmpdir.join("3")
mlflow.xgboost.save_model(xgb_model.model,
tmpdir3.strpath,
pip_requirements=[f"-c {req_file.strpath}", "b"])
_assert_pip_requirements(tmpdir3.strpath,
["mlflow", "b", "-c constraints.txt"], ["a"],
strict=True)
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())
def test_log_model_with_extra_pip_requirements(sklearn_knn_model,
main_scoped_model_class,
tmpdir):
sklearn_model_path = tmpdir.join("sklearn_model").strpath
mlflow.sklearn.save_model(sk_model=sklearn_knn_model,
path=sklearn_model_path)
python_model = main_scoped_model_class(predict_fn=None)
default_reqs = mlflow.pyfunc.get_default_pip_requirements()
# Path to a requirements file
req_file = tmpdir.join("requirements.txt")
req_file.write("a")
with mlflow.start_run():
mlflow.pyfunc.log_model(
"model",
python_model=python_model,
artifacts={"sk_model": sklearn_model_path},
extra_pip_requirements=req_file.strpath,
)
_assert_pip_requirements(mlflow.get_artifact_uri("model"),
["mlflow", *default_reqs, "a"])
# List of requirements
with mlflow.start_run():
mlflow.pyfunc.log_model(
"model",
artifacts={"sk_model": sklearn_model_path},
python_model=python_model,
extra_pip_requirements=[f"-r {req_file.strpath}", "b"],
)
_assert_pip_requirements(mlflow.get_artifact_uri("model"),
["mlflow", *default_reqs, "a", "b"])
# Constraints file
with mlflow.start_run():
mlflow.pyfunc.log_model(
"model",
artifacts={"sk_model": sklearn_model_path},
python_model=python_model,
extra_pip_requirements=[f"-c {req_file.strpath}", "b"],
)
_assert_pip_requirements(
mlflow.get_artifact_uri("model"),
["mlflow", *default_reqs, "b", "-c constraints.txt"],
["a"],
)
def test_log_model_with_pip_requirements(saved_tf_iris_model, tmpdir):
# Path to a requirements file
req_file = tmpdir.join("requirements.txt")
req_file.write("a")
with mlflow.start_run():
mlflow.tensorflow.log_model(
tf_saved_model_dir=saved_tf_iris_model.path,
tf_meta_graph_tags=saved_tf_iris_model.meta_graph_tags,
tf_signature_def_key=saved_tf_iris_model.signature_def_key,
artifact_path="model",
pip_requirements=req_file.strpath,
)
_assert_pip_requirements(mlflow.get_artifact_uri("model"),
["mlflow", "a"],
strict=True)
# List of requirements
with mlflow.start_run():
mlflow.tensorflow.log_model(
tf_saved_model_dir=saved_tf_iris_model.path,
tf_meta_graph_tags=saved_tf_iris_model.meta_graph_tags,
tf_signature_def_key=saved_tf_iris_model.signature_def_key,
artifact_path="model",
pip_requirements=[f"-r {req_file.strpath}", "b"],
)
_assert_pip_requirements(mlflow.get_artifact_uri("model"),
["mlflow", "a", "b"],
strict=True)
# Constraints file
with mlflow.start_run():
mlflow.tensorflow.log_model(
tf_saved_model_dir=saved_tf_iris_model.path,
tf_meta_graph_tags=saved_tf_iris_model.meta_graph_tags,
tf_signature_def_key=saved_tf_iris_model.signature_def_key,
artifact_path="model",
pip_requirements=[f"-c {req_file.strpath}", "b"],
)
_assert_pip_requirements(
mlflow.get_artifact_uri("model"),
["mlflow", "b", "-c constraints.txt"],
["a"],
strict=True,
)
def test_diviner_log_model_with_pip_requirements(grouped_prophet, tmp_path):
req_file = tmp_path.joinpath("requirements.txt")
req_file.write_text("a")
with mlflow.start_run():
mlflow.diviner.log_model(grouped_prophet, "model", pip_requirements=str(req_file))
_assert_pip_requirements(mlflow.get_artifact_uri("model"), ["mlflow", "a"], strict=True)
# List of requirements
with mlflow.start_run():
mlflow.diviner.log_model(grouped_prophet, "model", pip_requirements=[f"-r {req_file}", "b"])
_assert_pip_requirements(
mlflow.get_artifact_uri("model"), ["mlflow", "a", "b"], strict=True
)
# Constraints file
with mlflow.start_run():
mlflow.diviner.log_model(grouped_prophet, "model", pip_requirements=[f"-c {req_file}", "b"])
_assert_pip_requirements(
mlflow.get_artifact_uri("model"),
["mlflow", "b", "-c constraints.txt"],
["a"],
strict=True,
)
def test_model_save_without_specified_conda_env_uses_default_env_with_expected_dependencies(
model_path,
):
ols = ols_model()
mlflow.statsmodels.save_model(statsmodels_model=ols.model, path=model_path)
_assert_pip_requirements(model_path, mlflow.statsmodels.get_default_pip_requirements())
def test_model_save_without_specified_conda_env_uses_default_env_with_expected_dependencies(
h2o_iris_model, model_path):
mlflow.h2o.save_model(h2o_model=h2o_iris_model.model, path=model_path)
_assert_pip_requirements(model_path,
mlflow.h2o.get_default_pip_requirements())
def test_model_save_without_specified_conda_env_uses_default_env_with_expected_dependencies(
reg_model, model_path):
mlflow.catboost.save_model(reg_model.model, model_path)
_assert_pip_requirements(model_path,
mlflow.catboost.get_default_pip_requirements())
def test_model_save_without_specified_conda_env_uses_default_env_with_expected_dependencies(
spacy_model_with_data, model_path
):
mlflow.spacy.save_model(spacy_model=spacy_model_with_data.model, path=model_path)
_assert_pip_requirements(model_path, mlflow.spacy.get_default_pip_requirements())
def test_model_save_without_specified_conda_env_uses_default_env_with_expected_dependencies(
sequential_model, model_path):
mlflow.pytorch.save_model(pytorch_model=sequential_model, path=model_path)
_assert_pip_requirements(model_path,
mlflow.pytorch.get_default_pip_requirements())
def test_model_save_without_specified_conda_env_uses_default_env_with_expected_dependencies(
lgb_model, model_path):
mlflow.lightgbm.save_model(lgb_model=lgb_model.model, path=model_path)
_assert_pip_requirements(model_path,
mlflow.lightgbm.get_default_pip_requirements())
def test_model_save_without_specified_conda_env_uses_default_env_with_expected_dependencies(
pd_model, model_path):
mlflow.paddle.save_model(pd_model=pd_model.model, path=model_path)
_assert_pip_requirements(model_path,
mlflow.onnx.get_default_pip_requirements())
def test_sparkml_model_save_without_specified_conda_env_uses_default_env_with_expected_dependencies(
spark_model_iris, model_path):
sparkm.save_model(spark_model=spark_model_iris.model, path=model_path)
_assert_pip_requirements(model_path, sparkm.get_default_pip_requirements())
def test_pmdarima_model_save_without_conda_env_uses_default_env_with_expected_dependencies(
auto_arima_model, model_path):
mlflow.pmdarima.save_model(auto_arima_model, model_path)
_assert_pip_requirements(model_path,
mlflow.pmdarima.get_default_pip_requirements())
def test_model_save_without_specified_conda_env_uses_default_env_with_expected_dependencies(
fastai_model, model_path):
mlflow.fastai.save_model(fastai_learner=fastai_model.model,
path=model_path)
_assert_pip_requirements(model_path,
mlflow.fastai.get_default_pip_requirements())
def test_diviner_model_save_without_conda_env_uses_default_env_with_expected_dependencies(
grouped_prophet, model_path
):
mlflow.diviner.save_model(grouped_prophet, model_path)
_assert_pip_requirements(model_path, mlflow.diviner.get_default_pip_requirements())
def test_model_log_without_specified_conda_env_uses_default_env_with_expected_dependencies(model):
artifact_path = "model"
with mlflow.start_run():
mlflow.keras.log_model(keras_model=model, artifact_path=artifact_path)
model_uri = mlflow.get_artifact_uri(artifact_path)
_assert_pip_requirements(model_uri, mlflow.keras.get_default_pip_requirements())
