def test_log_model_with_code_paths(grouped_pmdarima):
artifact_path = "model"
with mlflow.start_run(), mock.patch(
"mlflow.diviner._add_code_from_conf_to_system_path"
) as add_mock:
mlflow.diviner.log_model(grouped_pmdarima, artifact_path, code_paths=[__file__])
model_uri = mlflow.get_artifact_uri(artifact_path)
_compare_logged_code_paths(__file__, model_uri, mlflow.diviner.FLAVOR_NAME)
mlflow.diviner.load_model(model_uri)
add_mock.assert_called()
def test_log_model_with_code_paths(spacy_model_with_data):
artifact_path = "model"
with mlflow.start_run(), mock.patch(
"mlflow.spacy._add_code_from_conf_to_system_path"
) as add_mock:
mlflow.spacy.log_model(spacy_model_with_data.model, artifact_path, code_paths=[__file__])
model_uri = mlflow.get_artifact_uri(artifact_path)
_compare_logged_code_paths(__file__, model_uri, mlflow.spacy.FLAVOR_NAME)
mlflow.spacy.load_model(model_uri)
add_mock.assert_called()
def test_log_model_with_code_paths(saved_tf_iris_model):
artifact_path = "model"
with mlflow.start_run(), mock.patch(
"mlflow.tensorflow._add_code_from_conf_to_system_path"
) as add_mock:
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=artifact_path,
code_paths=[__file__],
)
model_uri = mlflow.get_artifact_uri(artifact_path)
_compare_logged_code_paths(__file__, model_uri, mlflow.tensorflow.FLAVOR_NAME)
mlflow.tensorflow.load_model(model_uri)
add_mock.assert_called()
