import json import os import pathlib import time import uuid from datetime import date from unittest import mock import numpy as np import pandas as pd import pydantic import pytest import sklearn.datasets import sklearn.linear_model from packaging.version import Version from scipy.sparse import csc_matrix import mlflow from mlflow.exceptions import MlflowException from mlflow.models import Model, ModelSignature, infer_signature, set_model, validate_schema from mlflow.models.model import METADATA_FILES, SET_MODEL_ERROR from mlflow.models.resources import DatabricksServingEndpoint, DatabricksVectorSearchIndex from mlflow.models.utils import _read_example, _save_example from mlflow.store.artifact.s3_artifact_repo import S3ArtifactRepository from mlflow.tracking.artifact_utils import _download_artifact_from_uri from mlflow.types.schema import ColSpec, DataType, ParamSchema, ParamSpec, Schema, TensorSpec from mlflow.utils.databricks_utils import DatabricksRuntimeVersion from mlflow.utils.file_utils import TempDir from mlflow.utils.model_utils import _validate_and_prepare_target_save_path from mlflow.utils.proto_json_utils import dataframe_from_raw_json @pytest.fixture(scope="module") def iris_data(): iris = sklearn.datasets.load_iris() x = iris.data[:, :2] y = iris.target return x, y @pytest.fixture(scope="module") def sklearn_knn_model(iris_data): x, y = iris_data logreg_model = sklearn.linear_model.LogisticRegression() logreg_model.fit(x, y) return logreg_model def test_model_save_load(): m = Model( artifact_path="model", run_id="123", flavors={"flavor1": {"a": 1, "b": 2}, "flavor2": {"x": 1, "y": 2}}, signature=ModelSignature( inputs=Schema([ColSpec("integer", "x"), ColSpec("integer", "y")]), outputs=Schema([ColSpec(name=None, type="double")]), ), saved_input_example_info={"x": 1, "y": 2}, ) assert m.get_input_schema() == m.signature.inputs assert m.get_output_schema() == m.signature.outputs x = Model(artifact_path="some/other/path", run_id="1234") assert x.get_input_schema() is None assert x.get_output_schema() is None n = Model( artifact_path="model", run_id="123", flavors={"flavor1": {"a": 1, "b": 2}, "flavor2": {"x": 1, "y": 2}}, signature=ModelSignature( inputs=Schema([ColSpec("integer", "x"), ColSpec("integer", "y")]), outputs=Schema([ColSpec(name=None, type="double")]), ), saved_input_example_info={"x": 1, "y": 2}, ) n.utc_time_created = m.utc_time_created n.model_uuid = m.model_uuid assert m == n n.signature = None assert m != n with TempDir() as tmp: m.save(tmp.path("MLmodel")) o = Model.load(tmp.path("MLmodel")) assert m == o assert m.to_json() == o.to_json() assert m.to_yaml() == o.to_yaml() def test_model_load_remote(tmp_path, mock_s3_bucket): model = Model( artifact_path="model", run_id="123", flavors={"flavor1": {"a": 1, "b": 2}, "flavor2": {"x": 1, "y": 2}}, signature=ModelSignature( inputs=Schema([ColSpec("integer", "x"), ColSpec("integer", "y")]), outputs=Schema([ColSpec(name=None, type="double")]), ), saved_input_example_info={"x": 1, "y": 2}, ) model_path = tmp_path / "MLmodel" model.save(model_path) artifact_root = f"s3://{mock_s3_bucket}" artifact_repo = S3ArtifactRepository(artifact_root) artifact_repo.log_artifact(str(model_path)) model_reloaded_1 = Model.load(f"{artifact_root}/MLmodel") assert model_reloaded_1 == model model_reloaded_2 = Model.load(artifact_root) assert model_reloaded_2 == model class TestFlavor: @classmethod def save_model(cls, path, mlflow_model, signature=None, input_example=None): mlflow_model.flavors["flavor1"] = {"a": 1, "b": 2} mlflow_model.flavors["flavor2"] = {"x": 1, "y": 2} _validate_and_prepare_target_save_path(path) if signature is not None: mlflow_model.signature = signature if input_example is not None: _save_example(mlflow_model, input_example, path) mlflow_model.save(os.path.join(path, "MLmodel")) def _log_model_with_signature_and_example( tmp_path, sig, input_example, metadata=None, resources=None ): experiment_id = mlflow.create_experiment("test") with mlflow.start_run(experiment_id=experiment_id) as run: model = Model.log( "model", TestFlavor, signature=sig, input_example=input_example, metadata=metadata, resources=resources, ) # TODO: remove this after replacing all `with TempDir(chdr=True) as tmp` # with tmp_path fixture output_path = tmp_path if isinstance(tmp_path, pathlib.PosixPath) else tmp_path.path("") local_path = _download_artifact_from_uri(model.model_uri, output_path=output_path) return local_path, run def test_model_log(): with TempDir(chdr=True) as tmp: sig = ModelSignature( inputs=Schema([ColSpec("integer", "x"), ColSpec("integer", "y")]), outputs=Schema([ColSpec(name=None, type="double")]), ) input_example = {"x": 1, "y": 2} local_path, r = _log_model_with_signature_and_example(tmp, sig, input_example) loaded_model = Model.load(os.path.join(local_path, "MLmodel")) assert loaded_model.run_id == r.info.run_id assert loaded_model.flavors == { "flavor1": {"a": 1, "b": 2}, "flavor2": {"x": 1, "y": 2}, } assert loaded_model.signature == sig x = _read_example( Model(saved_input_example_info=loaded_model.saved_input_example_info), local_path ) assert x == input_example assert not hasattr(loaded_model, "databricks_runtime") loaded_example = loaded_model.load_input_example(local_path) assert loaded_example == input_example assert Version(loaded_model.mlflow_version) == Version(mlflow.version.VERSION) def test_model_log_without_run(tmp_path): model_info = Model.log("model", TestFlavor) assert model_info.run_id is None def test_model_log_with_active_run(tmp_path): with mlflow.start_run() as run: model_info = Model.log("model", TestFlavor) assert model_info.run_id == run.info.run_id def test_model_log_inactive_run_id(tmp_path): experiment_id = mlflow.create_experiment("test", artifact_location=str(tmp_path)) run = mlflow.MlflowClient().create_run(experiment_id=experiment_id) model_info = Model.log("model", TestFlavor, run_id=run.info.run_id) assert model_info.run_id == run.info.run_id def test_model_log_calls_maybe_render_agent_eval_recipe(tmp_path): sig = ModelSignature( inputs=Schema([ColSpec("integer", "x"), ColSpec("integer", "y")]), outputs=Schema([ColSpec(name=None, type="double")]), ) input_example = {"x": 1, "y": 2} with mock.patch("mlflow.models.display_utils.maybe_render_agent_eval_recipe") as render_mock: _log_model_with_signature_and_example(tmp_path, sig, input_example) render_mock.assert_called_once() def test_model_info(): with TempDir(chdr=True) as tmp: sig = ModelSignature( inputs=Schema([ColSpec("integer", "x"), ColSpec("integer", "y")]), outputs=Schema([ColSpec(name=None, type="double")]), ) input_example = {"x": 1, "y": 2} experiment_id = mlflow.create_experiment("test") with mlflow.start_run(experiment_id=experiment_id) as run: model_info = Model.log("model", TestFlavor, signature=sig, input_example=input_example) model_uri = f"models:/{model_info.model_id}" model_info_fetched = mlflow.models.get_model_info(model_uri) local_path = _download_artifact_from_uri(model_uri, output_path=tmp.path("")) assert model_info.run_id == run.info.run_id assert model_info_fetched.run_id == run.info.run_id assert model_info.model_uri == model_uri assert model_info_fetched.model_uri == model_uri loaded_model = Model.load(os.path.join(local_path, "MLmodel")) assert model_info.utc_time_created == loaded_model.utc_time_created assert model_info_fetched.utc_time_created == loaded_model.utc_time_created assert model_info.model_uuid == loaded_model.model_uuid assert model_info_fetched.model_uuid == loaded_model.model_uuid assert model_info.flavors == { "flavor1": {"a": 1, "b": 2}, "flavor2": {"x": 1, "y": 2}, } x = _read_example( Model(saved_input_example_info=model_info.saved_input_example_info), local_path ) assert x == input_example model_signature = model_info_fetched.signature assert model_signature.to_dict() == sig.to_dict() assert model_info.mlflow_version == loaded_model.mlflow_version assert model_info_fetched.mlflow_version == loaded_model.mlflow_version def test_model_info_with_model_version(tmp_path): experiment_id = mlflow.create_experiment("test", artifact_location=str(tmp_path)) with mlflow.start_run(experiment_id=experiment_id): model_info = Model.log("model", TestFlavor, registered_model_name="model_abc") assert model_info.registered_model_version == 1 model_info = Model.log("model", TestFlavor, registered_model_name="model_abc") assert model_info.registered_model_version == 2 model_info = Model.log("model", TestFlavor) assert model_info.registered_model_version is None def test_model_log_tags_propagated_to_registered_model_version(tmp_path): experiment_id = mlflow.create_experiment("test_tags", artifact_location=str(tmp_path)) tags = {"stage": "training", "framework": "test"} with mlflow.start_run(experiment_id=experiment_id): model_info = Model.log( "model", TestFlavor, registered_model_name="model_with_tags", tags=tags, ) client = mlflow.MlflowClient() model_version = client.get_model_version("model_with_tags", model_info.registered_model_version) assert model_version.tags == tags def test_model_metadata(): with TempDir(chdr=True) as tmp: metadata = {"metadata_key": "metadata_value"} local_path, _ = _log_model_with_signature_and_example(tmp, None, None, metadata) loaded_model = Model.load(os.path.join(local_path, "MLmodel")) assert loaded_model.metadata["metadata_key"] == "metadata_value" def test_load_model_without_mlflow_version(): with TempDir(chdr=True) as tmp: model = Model(artifact_path="model", run_id="1234", mlflow_version=None) path = tmp.path("MLmodel") with open(path, "w") as out: model.to_yaml(out) loaded_model = Model.load(path) assert loaded_model.mlflow_version is None def test_model_log_with_databricks_runtime(): dbr_version = "8.3.x" with mlflow.start_run(): with mock.patch( "mlflow.models.model.get_databricks_runtime_version", return_value=dbr_version ) as mock_get_dbr_version: model = Model.log("path", TestFlavor, signature=None, input_example=None) mock_get_dbr_version.assert_called() loaded_model = Model.load(model.model_uri) assert loaded_model.databricks_runtime == dbr_version def test_model_log_with_databricks_runtime_gpu(): dbr_version = "client.8.1-gpu" with mlflow.start_run(): with mock.patch( "mlflow.models.model.get_databricks_runtime_version", return_value=dbr_version ) as mock_get_dbr_version: model = Model.log("path", TestFlavor, signature=None, input_example=None) mock_get_dbr_version.assert_called() # Verify the GPU suffix is preserved in the MLmodel file loaded_model = Model.load(model.model_uri) assert loaded_model.databricks_runtime == dbr_version # Verify that the version can be parsed correctly and is_gpu_image is True parsed_version = DatabricksRuntimeVersion.parse(loaded_model.databricks_runtime) assert parsed_version.is_client_image is True assert parsed_version.major == 8 assert parsed_version.minor == 1 assert parsed_version.is_gpu_image is True def test_model_log_with_input_example_succeeds(): with TempDir(chdr=True) as tmp: sig = ModelSignature( inputs=Schema([ ColSpec("integer", "a"), ColSpec("string", "b"), ColSpec("boolean", "c"), ColSpec("string", "d"), ColSpec("datetime", "e"), ]), outputs=Schema([ColSpec(name=None, type="double")]), ) input_example = pd.DataFrame( { "a": np.int32(1), "b": "test string", "c": True, "d": date.today(), "e": np.datetime64("2020-01-01T00:00:00"), }, index=[0], ) local_path, _ = _log_model_with_signature_and_example(tmp, sig, input_example) loaded_model = Model.load(os.path.join(local_path, "MLmodel")) path = os.path.join(local_path, loaded_model.saved_input_example_info["artifact_path"]) x = dataframe_from_raw_json(path, schema=sig.inputs) # date column will get deserialized into string input_example["d"] = input_example["d"].apply(lambda x: x.isoformat()) # datetime Datatype numpy type is [ns] input_example["e"] = input_example["e"].astype(np.dtype("datetime64[ns]")) pd.testing.assert_frame_equal(x, input_example) loaded_example = loaded_model.load_input_example(local_path) assert isinstance(loaded_example, pd.DataFrame) pd.testing.assert_frame_equal(loaded_example, input_example) def test_model_input_example_with_params_log_load_succeeds(tmp_path): pdf = pd.DataFrame( { "a": np.int32(1), "b": "test string", "c": True, "d": date.today(), "e": np.datetime64("2020-01-01T00:00:00"), }, index=[0], ) input_example = (pdf, {"a": 1, "b": "string"}) sig = ModelSignature( inputs=Schema([ ColSpec("integer", "a"), ColSpec("string", "b"), ColSpec("boolean", "c"), ColSpec("string", "d"), ColSpec("datetime", "e"), ]), outputs=Schema([ColSpec(name=None, type="double")]), params=ParamSchema([ ParamSpec("a", DataType.long, 1), ParamSpec("b", DataType.string, "string"), ]), ) local_path, _ = _log_model_with_signature_and_example(tmp_path, sig, input_example) loaded_model = Model.load(os.path.join(local_path, "MLmodel")) # date column will get deserialized into string pdf["d"] = pdf["d"].apply(lambda x: x.isoformat()) loaded_example = loaded_model.load_input_example(local_path) assert isinstance(loaded_example, pd.DataFrame) # datetime Datatype numpy type is [ns] pdf["e"] = pdf["e"].astype(np.dtype("datetime64[ns]")) pd.testing.assert_frame_equal(loaded_example, pdf) params = loaded_model.load_input_example_params(local_path) assert params == input_example[1] def test_model_load_input_example_numpy(): with TempDir(chdr=True) as tmp: input_example = np.array([[3, 4, 5]], dtype=np.int32) sig = ModelSignature( inputs=Schema([TensorSpec(type=input_example.dtype, shape=input_example.shape)]), outputs=Schema([ColSpec(name=None, type="double")]), ) local_path, _ = _log_model_with_signature_and_example(tmp, sig, input_example) loaded_model = Model.load(os.path.join(local_path, "MLmodel")) loaded_example = loaded_model.load_input_example(local_path) assert isinstance(loaded_example, np.ndarray) np.testing.assert_array_equal(input_example, loaded_example) def test_model_load_input_example_scipy(): with TempDir(chdr=True) as tmp: input_example = csc_matrix(np.arange(0, 12, 0.5).reshape(3, 8)) sig = ModelSignature( inputs=Schema([TensorSpec(type=input_example.data.dtype, shape=input_example.shape)]), outputs=Schema([ColSpec(name=None, type="double")]), ) local_path, _ = _log_model_with_signature_and_example(tmp, sig, input_example) loaded_model = Model.load(os.path.join(local_path, "MLmodel")) loaded_example = loaded_model.load_input_example(local_path) assert isinstance(loaded_example, csc_matrix) np.testing.assert_array_equal(input_example.data, loaded_example.data) def test_model_load_input_example_failures(): with TempDir(chdr=True) as tmp: input_example = np.array([[3, 4, 5]], dtype=np.int32) sig = ModelSignature( inputs=Schema([TensorSpec(type=input_example.dtype, shape=input_example.shape)]), outputs=Schema([ColSpec(name=None, type="double")]), ) local_path, _ = _log_model_with_signature_and_example(tmp, sig, input_example) loaded_model = Model.load(os.path.join(local_path, "MLmodel")) loaded_example = loaded_model.load_input_example(local_path) assert loaded_example is not None with pytest.raises(MlflowException, match="No such artifact"): loaded_model.load_input_example(os.path.join(local_path, "folder_which_does_not_exist")) path = os.path.join(local_path, loaded_model.saved_input_example_info["artifact_path"]) os.remove(path) with pytest.raises(MlflowException, match="No such artifact"): loaded_model.load_input_example(local_path) def test_model_load_input_example_no_signature(): with TempDir(chdr=True) as tmp: input_example = np.array([[3, 4, 5]], dtype=np.int32) sig = ModelSignature( inputs=Schema([TensorSpec(type=input_example.dtype, shape=input_example.shape)]), outputs=Schema([ColSpec(name=None, type="double")]), ) local_path, _ = _log_model_with_signature_and_example(tmp, sig, input_example=None) loaded_model = Model.load(os.path.join(local_path, "MLmodel")) loaded_example = loaded_model.load_input_example(local_path) assert loaded_example is None def _is_valid_uuid(val): try: uuid.UUID(str(val)) return True except ValueError: return False def test_model_uuid(): m = Model() assert m.model_uuid is not None assert _is_valid_uuid(m.model_uuid) m2 = Model() assert m.model_uuid != m2.model_uuid m_dict = m.to_dict() assert m_dict["model_uuid"] == m.model_uuid m3 = Model.from_dict(m_dict) assert m3.model_uuid == m.model_uuid m_dict.pop("model_uuid") m4 = Model.from_dict(m_dict) assert m4.model_uuid is None def test_validate_schema(sklearn_knn_model, iris_data, tmp_path): sk_model_path = os.path.join(tmp_path, "sk_model") X, y = iris_data signature = infer_signature(X, y) mlflow.sklearn.save_model( sklearn_knn_model, sk_model_path, signature=signature, ) validate_schema(X, signature.inputs) prediction = sklearn_knn_model.predict(X) reloaded_model = mlflow.sklearn.load_model(sk_model_path) np.testing.assert_array_equal(prediction, reloaded_model.predict(X)) validate_schema(prediction, signature.outputs) def test_save_load_input_example_without_conversion(tmp_path): class MyModel(mlflow.pyfunc.PythonModel): def predict(self, context, model_input, params=None): return model_input input_example = { "messages": [ {"role": "user", "content": "Hello!"}, ] } with mlflow.start_run() as run: mlflow.pyfunc.log_model( name="test_model", python_model=MyModel(), input_example=input_example, ) local_path = _download_artifact_from_uri( f"runs:/{run.info.run_id}/test_model", output_path=tmp_path ) loaded_model = Model.load(os.path.join(local_path, "MLmodel")) assert loaded_model.saved_input_example_info["type"] == "json_object" loaded_example = loaded_model.load_input_example(local_path) assert loaded_example == input_example def test_save_load_input_example_with_pydantic_model(tmp_path): class Message(pydantic.BaseModel): role: str content: str class MyModel(mlflow.pyfunc.PythonModel): def predict(self, context, model_input: list[Message], params=None): return model_input with mlflow.start_run(): model_info = mlflow.pyfunc.log_model( name="test_model", python_model=MyModel(), input_example=[Message(role="user", content="Hello!")], ) local_path = _download_artifact_from_uri(model_info.model_uri, output_path=tmp_path) loaded_model = Model.load(os.path.join(local_path, "MLmodel")) assert loaded_model.saved_input_example_info["type"] == "json_object" loaded_example = loaded_model.load_input_example(local_path) assert loaded_example == [{"role": "user", "content": "Hello!"}] def test_model_saved_by_save_model_can_be_loaded(tmp_path, sklearn_knn_model): mlflow.sklearn.save_model(sklearn_knn_model, tmp_path) info = Model.load(tmp_path).get_model_info() assert info.run_id is None assert info.artifact_path is None def test_copy_metadata(mock_is_in_databricks, sklearn_knn_model): with mlflow.start_run(): model_info = mlflow.sklearn.log_model(sklearn_knn_model, name="model") artifact_path = mlflow.artifacts.download_artifacts(model_info.model_uri) metadata_path = os.path.join(artifact_path, "metadata") # Metadata should be copied only in Databricks if mock_is_in_databricks.return_value: assert set(os.listdir(metadata_path)) == set(METADATA_FILES) else: assert not os.path.exists(metadata_path) mock_is_in_databricks.assert_called_once() class LegacyTestFlavor: @classmethod def save_model(cls, path, mlflow_model): mlflow_model.flavors["flavor1"] = {"a": 1, "b": 2} mlflow_model.flavors["flavor2"] = {"x": 1, "y": 2} _validate_and_prepare_target_save_path(path) mlflow_model.save(os.path.join(path, "MLmodel")) def test_legacy_flavor(mock_is_in_databricks): with mlflow.start_run(): model_info = Model.log("model", LegacyTestFlavor) artifact_path = _download_artifact_from_uri(model_info.model_uri) metadata_path = os.path.join(artifact_path, "metadata") # Metadata should be copied only in Databricks if mock_is_in_databricks.return_value: assert set(os.listdir(metadata_path)) == {"MLmodel"} else: assert not os.path.exists(metadata_path) mock_is_in_databricks.assert_called_once() def test_pyfunc_set_model(): class MyModel(mlflow.pyfunc.PythonModel): def predict(self, context, model_input): return model_input set_model(MyModel()) assert isinstance(mlflow.models.model.__mlflow_model__, mlflow.pyfunc.PythonModel) def test_langchain_set_model(): from langchain_core.runnables import RunnableLambda def create_runnable(): def my_runnable(input): return f"Input was: {input}" runnable = RunnableLambda(my_runnable) set_model(runnable) create_runnable() assert isinstance(mlflow.models.model.__mlflow_model__, RunnableLambda) def test_error_set_model(sklearn_knn_model): with pytest.raises(mlflow.MlflowException, match=SET_MODEL_ERROR): set_model(sklearn_knn_model) def test_model_resources(): expected_resources = { "api_version": "1", "databricks": { "serving_endpoint": [ {"name": "databricks-mixtral-8x7b-instruct"}, {"name": "databricks-bge-large-en"}, {"name": "azure-eastus-model-serving-2_vs_endpoint"}, ], "vector_search_index": [{"name": "rag.studio_bugbash.databricks_docs_index"}], }, } with TempDir(chdr=True) as tmp: resources = [ DatabricksServingEndpoint(endpoint_name="databricks-mixtral-8x7b-instruct"), DatabricksServingEndpoint(endpoint_name="databricks-bge-large-en"), DatabricksServingEndpoint(endpoint_name="azure-eastus-model-serving-2_vs_endpoint"), DatabricksVectorSearchIndex(index_name="rag.studio_bugbash.databricks_docs_index"), ] local_path, _ = _log_model_with_signature_and_example(tmp, None, None, resources=resources) loaded_model = Model.load(os.path.join(local_path, "MLmodel")) assert loaded_model.resources == expected_resources def test_save_load_model_with_run_uri(): class MyModel(mlflow.pyfunc.PythonModel): def predict(self, context, model_input: list[str], params=None): return model_input with mlflow.start_run() as run: mlflow.pyfunc.log_model( name="test_model", python_model=MyModel(), input_example=["a", "b", "c"], ) mlflow_model = Model.load(f"runs:/{run.info.run_id}/test_model/MLmodel") assert mlflow_model.load_input_example() == ["a", "b", "c"] model = Model.load(f"runs:/{run.info.run_id}/test_model") assert model == mlflow_model model = Model.load(f"runs:/{run.info.run_id}/test_model/") assert model == mlflow_model def test_save_model_with_prompts(): prompt_1 = mlflow.register_prompt("prompt-1", "Hello, {{title}} {{name}}!") time.sleep(0.001) # To avoid timestamp precision issue in Windows prompt_2 = mlflow.register_prompt("prompt-2", "Hello, {{title}} {{name}}!") class MyModel(mlflow.pyfunc.PythonModel): def predict(self, model_input: list[str]): return model_input with mlflow.start_run(): model_info = mlflow.pyfunc.log_model( name="test_model", python_model=MyModel(), # The 'prompts' parameter should accept both prompt object and URI prompts=[prompt_1, prompt_2.uri], ) assert model_info.prompts == [prompt_1.uri, prompt_2.uri] # Prompts should be recorded in the yaml file model = Model.load(model_info.model_uri) assert model.prompts == [prompt_1.uri, prompt_2.uri] # Check that prompts were linked to the run via the linkedPrompts tag from mlflow.tracing.constant import TraceTagKey run = mlflow.MlflowClient().get_run(model_info.run_id) linked_prompts_tag = run.data.tags.get(TraceTagKey.LINKED_PROMPTS) assert linked_prompts_tag is not None linked_prompts = json.loads(linked_prompts_tag) assert len(linked_prompts) == 2 assert {p["name"] for p in linked_prompts} == {prompt_1.name, prompt_2.name} def test_logged_model_status(): def predict_fn(model_input: list[str]): return model_input model_info = mlflow.pyfunc.log_model( name="test_model", python_model=predict_fn, input_example=["a", "b", "c"], ) logged_model = mlflow.get_logged_model(model_info.model_id) assert logged_model.status == "READY" with pytest.raises(Exception, match=r"mock exception"): with mock.patch( "mlflow.pyfunc.model._save_model_with_class_artifacts_params", side_effect=Exception("mock exception"), ): mlflow.pyfunc.log_model( name="test_model", python_model=predict_fn, input_example=["a", "b", "c"], ) logged_model = mlflow.last_logged_model() assert logged_model.status == "FAILED" def test_model_log_links_prompts_to_logged_model(): client = mlflow.MlflowClient() # Create actual prompts in the registry client.create_prompt(name="test_prompt_1") prompt_1 = client.create_prompt_version(name="test_prompt_1", template="Hello {{name}}") client.create_prompt(name="test_prompt_2") prompt_2 = client.create_prompt_version(name="test_prompt_2", template="Goodbye {{name}}") with mlflow.start_run() as run: model_info = Model.log("model", TestFlavor, prompts=[prompt_1, prompt_2]) # Verify prompts were linked to the run run_data = client.get_run(run.info.run_id) linked_prompts_tag = run_data.data.tags.get("mlflow.linkedPrompts") assert linked_prompts_tag is not None linked_prompts = json.loads(linked_prompts_tag) assert len(linked_prompts) == 2 assert {p["name"] for p in linked_prompts} == {"test_prompt_1", "test_prompt_2"} # Verify prompts were linked to the LoggedModel logged_model = client.get_logged_model(model_info.model_id) model_linked_prompts_tag = logged_model.tags.get("mlflow.linkedPrompts") assert model_linked_prompts_tag is not None model_linked_prompts = json.loads(model_linked_prompts_tag) assert len(model_linked_prompts) == 2 assert {p["name"] for p in model_linked_prompts} == {"test_prompt_1", "test_prompt_2"} def test_get_model_info_with_logged_model(): def model(model_input: list[str]) -> list[str]: return model_input model_info_log_model = mlflow.pyfunc.log_model( name="test_model", python_model=model, input_example=["a", "b", "c"] ) model_info_get_model_info = mlflow.models.get_model_info(model_info_log_model.model_uri) assert model_info_log_model.model_id == model_info_get_model_info.model_id assert model_info_log_model.name == model_info_get_model_info.name