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()