import json import os import pickle import shutil import tempfile from pathlib import Path from typing import Any, NamedTuple from unittest import mock import cloudpickle import numpy as np import pandas as pd import pytest import sklearn import sklearn.linear_model as glm import sklearn.naive_bayes as nb import sklearn.neighbors as knn import skops import yaml from packaging.version import Version from sklearn import datasets from sklearn.pipeline import Pipeline as SKPipeline from sklearn.pipeline import make_pipeline from sklearn.preprocessing import FunctionTransformer as SKFunctionTransformer import mlflow.pyfunc.scoring_server as pyfunc_scoring_server import mlflow.sklearn from mlflow import pyfunc from mlflow.entities.model_registry.model_version import ModelVersion, ModelVersionStatus from mlflow.exceptions import MlflowException from mlflow.models import Model, ModelSignature from mlflow.models.utils import _read_example, load_serving_example from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE, ErrorCode from mlflow.store._unity_catalog.registry.rest_store import UcModelRegistryStore from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository from mlflow.store.artifact.s3_artifact_repo import S3ArtifactRepository from mlflow.tracking.artifact_utils import _download_artifact_from_uri from mlflow.types import DataType from mlflow.types.schema import ColSpec, Schema from mlflow.utils.environment import _mlflow_conda_env from mlflow.utils.file_utils import TempDir from mlflow.utils.model_utils import _get_flavor_configuration from tests.helper_functions import ( _assert_pip_requirements, _compare_conda_env_requirements, _compare_logged_code_paths, _is_available_on_pypi, _mlflow_major_version_string, assert_register_model_called_with_local_model_path, pyfunc_serve_and_score_model, ) from tests.store._unity_catalog.conftest import ( configure_client_for_uc, # noqa: F401 mock_databricks_uc_host_creds, # noqa: F401 ) EXTRA_PYFUNC_SERVING_TEST_ARGS = ( [] if _is_available_on_pypi("scikit-learn", module="sklearn") else ["--env-manager", "local"] ) class ModelWithData(NamedTuple): model: Any inference_data: Any @pytest.fixture(scope="module") def iris_df(): iris = datasets.load_iris() X = iris.data y = iris.target X_df = pd.DataFrame(X, columns=iris.feature_names) X_df = X_df.iloc[:, :2] # we only take the first two features. y_series = pd.Series(y) return X_df, y_series @pytest.fixture(scope="module") def iris_signature(): return ModelSignature( inputs=Schema([ ColSpec(name="sepal length (cm)", type=DataType.double), ColSpec(name="sepal width (cm)", type=DataType.double), ]), outputs=Schema([ColSpec(type=DataType.long)]), ) @pytest.fixture(scope="module") def sklearn_knn_model(iris_df): X, y = iris_df knn_model = knn.KNeighborsClassifier() knn_model.fit(X, y) return ModelWithData(model=knn_model, inference_data=X) # To load sklearn KNN model as skops format, # We need to mark these types as `skops_trusted_types` # related ticket: https://github.com/skops-dev/skops/issues/498 sklearn_knn_model_skops_trusted_types = [ "sklearn.metrics._dist_metrics.EuclideanDistance64", "sklearn.neighbors._kd_tree.KDTree", ] @pytest.fixture(scope="module") def sklearn_logreg_model(iris_df): X, y = iris_df linear_lr = glm.LogisticRegression() linear_lr.fit(X, y) return ModelWithData(model=linear_lr, inference_data=X) @pytest.fixture(scope="module") def sklearn_gaussian_model(iris_df): X, y = iris_df gaussian_nb = nb.GaussianNB() gaussian_nb.fit(X, y) return ModelWithData(model=gaussian_nb, inference_data=X) @pytest.fixture(scope="module") def sklearn_custom_transformer_model(sklearn_knn_model, iris_df): def transform(vec): return vec + 1 transformer = SKFunctionTransformer(transform, validate=True) pipeline = SKPipeline([("custom_transformer", transformer), ("knn", sklearn_knn_model.model)]) X, _ = iris_df return ModelWithData(pipeline, inference_data=X) @pytest.fixture def model_path(tmp_path): return os.path.join(tmp_path, "model") @pytest.fixture def sklearn_custom_env(tmp_path): conda_env = os.path.join(tmp_path, "conda_env.yml") _mlflow_conda_env(conda_env, additional_pip_deps=["scikit-learn", "pytest"]) return conda_env @pytest.mark.parametrize("serialization_format", mlflow.sklearn.SUPPORTED_SERIALIZATION_FORMATS) def test_model_save_load(sklearn_logreg_model, model_path, serialization_format): from mlflow.utils.requirements_utils import _parse_requirements sk_model = sklearn_logreg_model.model mlflow.sklearn.save_model( sk_model=sk_model, path=model_path, serialization_format=serialization_format ) reloaded_model = mlflow.sklearn.load_model(model_uri=model_path) reloaded_pyfunc = pyfunc.load_model(model_uri=model_path) sklearn_conf = _get_flavor_configuration( model_path=model_path, flavor_name=mlflow.sklearn.FLAVOR_NAME ) assert "serialization_format" in sklearn_conf assert sklearn_conf["serialization_format"] == serialization_format req_map = { mlflow.sklearn.SERIALIZATION_FORMAT_SKOPS: f"skops=={skops.__version__}", mlflow.sklearn.SERIALIZATION_FORMAT_CLOUDPICKLE: f"cloudpickle=={cloudpickle.__version__}", } logged_reqs = [ req.req_str for req in _parse_requirements( os.path.join(model_path, "requirements.txt"), is_constraint=False ) ] if serialization_format != mlflow.sklearn.SERIALIZATION_FORMAT_PICKLE: assert req_map[serialization_format] in logged_reqs np.testing.assert_array_equal( sk_model.predict(sklearn_logreg_model.inference_data), reloaded_model.predict(sklearn_logreg_model.inference_data), ) np.testing.assert_array_equal( reloaded_model.predict(sklearn_logreg_model.inference_data), reloaded_pyfunc.predict(sklearn_logreg_model.inference_data), ) def test_model_skops_format_trusted_type(sklearn_knn_model, model_path): sk_model = sklearn_knn_model.model with pytest.raises(MlflowException, match="The saved sklearn model references untrusted type"): mlflow.sklearn.save_model( sk_model=sk_model, path=model_path, serialization_format="skops", ) shutil.rmtree(model_path) mlflow.sklearn.save_model( sk_model=sklearn_knn_model.model, path=model_path, serialization_format="skops", skops_trusted_types=sklearn_knn_model_skops_trusted_types, ) reloaded_model = mlflow.sklearn.load_model(model_uri=model_path) reloaded_pyfunc = pyfunc.load_model(model_uri=model_path) np.testing.assert_array_equal( sk_model.predict(sklearn_knn_model.inference_data), reloaded_model.predict(sklearn_knn_model.inference_data), ) np.testing.assert_array_equal( reloaded_model.predict(sklearn_knn_model.inference_data), reloaded_pyfunc.predict(sklearn_knn_model.inference_data), ) def test_log_model_skops_no_pip_requirements_warning(sklearn_logreg_model, recwarn): with mlflow.start_run(): mlflow.sklearn.log_model( sklearn_logreg_model.model, serialization_format=mlflow.sklearn.SERIALIZATION_FORMAT_SKOPS, ) warning_messages = [str(w.message) for w in recwarn] assert not any("Fall back to return" in msg for msg in warning_messages) def test_model_save_behavior_with_preexisting_folders(sklearn_logreg_model, tmp_path): sklearn_model_path = tmp_path / "sklearn_model_empty_exists" sklearn_model_path.mkdir() mlflow.sklearn.save_model(sk_model=sklearn_logreg_model.model, path=sklearn_model_path) sklearn_model_path = tmp_path / "sklearn_model_filled_exists" sklearn_model_path.mkdir() (sklearn_model_path / "foo.txt").write_text("dummy content") with pytest.raises(MlflowException, match="already exists and is not empty"): mlflow.sklearn.save_model(sk_model=sklearn_logreg_model.model, path=sklearn_model_path) def test_signature_and_examples_are_saved_correctly(sklearn_knn_model, iris_signature): data = sklearn_knn_model.inference_data model = sklearn_knn_model.model example_ = data[:3] for signature in (None, iris_signature): for example in (None, example_): with TempDir() as tmp: path = tmp.path("model") mlflow.sklearn.save_model( model, path=path, signature=signature, input_example=example, skops_trusted_types=sklearn_knn_model_skops_trusted_types, ) mlflow_model = Model.load(path) if signature is None and example is None: assert mlflow_model.signature is None else: assert mlflow_model.signature == iris_signature if example is None: assert mlflow_model.saved_input_example_info is None else: np.testing.assert_array_equal(_read_example(mlflow_model, path), example) def test_model_load_from_remote_uri_succeeds(sklearn_logreg_model, model_path, mock_s3_bucket): mlflow.sklearn.save_model(sk_model=sklearn_logreg_model.model, path=model_path) artifact_root = f"s3://{mock_s3_bucket}" artifact_path = "model" artifact_repo = S3ArtifactRepository(artifact_root) artifact_repo.log_artifacts(model_path, artifact_path=artifact_path) model_uri = artifact_root + "/" + artifact_path reloaded_model = mlflow.sklearn.load_model(model_uri=model_uri) np.testing.assert_array_equal( sklearn_logreg_model.model.predict(sklearn_logreg_model.inference_data), reloaded_model.predict(sklearn_logreg_model.inference_data), ) def test_model_log(sklearn_logreg_model, model_path): with TempDir(chdr=True, remove_on_exit=True) as tmp: for should_start_run in [False, True]: try: if should_start_run: mlflow.start_run() artifact_path = "linear" conda_env = os.path.join(tmp.path(), "conda_env.yaml") _mlflow_conda_env(conda_env, additional_pip_deps=["scikit-learn"]) model_info = mlflow.sklearn.log_model( sklearn_logreg_model.model, name=artifact_path, conda_env=conda_env, ) reloaded_logsklearn_knn_model = mlflow.sklearn.load_model( model_uri=model_info.model_uri ) np.testing.assert_array_equal( sklearn_logreg_model.model.predict(sklearn_logreg_model.inference_data), reloaded_logsklearn_knn_model.predict(sklearn_logreg_model.inference_data), ) model_path = _download_artifact_from_uri(artifact_uri=model_info.model_uri) model_config = Model.load(os.path.join(model_path, "MLmodel")) assert pyfunc.FLAVOR_NAME in model_config.flavors assert pyfunc.ENV in model_config.flavors[pyfunc.FLAVOR_NAME] env_path = model_config.flavors[pyfunc.FLAVOR_NAME][pyfunc.ENV]["conda"] assert os.path.exists(os.path.join(model_path, env_path)) finally: mlflow.end_run() def test_log_model_calls_register_model(sklearn_logreg_model): artifact_path = "linear" register_model_patch = mock.patch("mlflow.tracking._model_registry.fluent._register_model") with mlflow.start_run(), register_model_patch, TempDir(chdr=True, remove_on_exit=True) as tmp: conda_env = os.path.join(tmp.path(), "conda_env.yaml") _mlflow_conda_env(conda_env, additional_pip_deps=["scikit-learn"]) model_info = mlflow.sklearn.log_model( sklearn_logreg_model.model, name=artifact_path, conda_env=conda_env, registered_model_name="AdsModel1", ) assert_register_model_called_with_local_model_path( register_model_mock=mlflow.tracking._model_registry.fluent._register_model, model_uri=model_info.model_uri, registered_model_name="AdsModel1", ) def test_log_model_call_register_model_to_uc(configure_client_for_uc, sklearn_logreg_model): artifact_path = "linear" mock_model_version = ModelVersion( name="AdsModel1", version=1, creation_timestamp=123, status=ModelVersionStatus.to_string(ModelVersionStatus.READY), ) with ( mock.patch.object(UcModelRegistryStore, "create_registered_model"), mock.patch.object( UcModelRegistryStore, "create_model_version", return_value=mock_model_version, autospec=True, ) as mock_create_mv, TempDir(chdr=True, remove_on_exit=True) as tmp, ): with mlflow.start_run() as run: conda_env = os.path.join(tmp.path(), "conda_env.yaml") _mlflow_conda_env(conda_env, additional_pip_deps=["scikit-learn"]) model_info = mlflow.sklearn.log_model( sklearn_logreg_model.model, name=artifact_path, conda_env=conda_env, registered_model_name="AdsModel1", ) source = model_info.artifact_path [(args, kwargs)] = mock_create_mv.call_args_list assert args[1:] == ("AdsModel1", source, run.info.run_id, [], None, None) assert kwargs["local_model_path"].startswith(tempfile.gettempdir()) def test_log_model_no_registered_model_name(sklearn_logreg_model): artifact_path = "model" register_model_patch = mock.patch("mlflow.tracking._model_registry.fluent._register_model") with mlflow.start_run(), register_model_patch, TempDir(chdr=True, remove_on_exit=True) as tmp: conda_env = os.path.join(tmp.path(), "conda_env.yaml") _mlflow_conda_env(conda_env, additional_pip_deps=["scikit-learn"]) mlflow.sklearn.log_model( sklearn_logreg_model.model, name=artifact_path, conda_env=conda_env, ) mlflow.tracking._model_registry.fluent._register_model.assert_not_called() def test_custom_transformer_can_be_saved_and_loaded_with_cloudpickle_format( sklearn_custom_transformer_model, tmp_path ): custom_transformer_model = sklearn_custom_transformer_model.model # Because the model contains a customer transformer that is not defined at the top level of the # current test module, we expect pickle to fail when attempting to serialize it. In contrast, # we expect cloudpickle to successfully locate the transformer definition and serialize the # model successfully. pickle_format_model_path = os.path.join(tmp_path, "pickle_model") with pytest.raises(AttributeError, match="Can't pickle local object"): mlflow.sklearn.save_model( sk_model=custom_transformer_model, path=pickle_format_model_path, serialization_format=mlflow.sklearn.SERIALIZATION_FORMAT_PICKLE, ) cloudpickle_format_model_path = os.path.join(tmp_path, "cloud_pickle_model") mlflow.sklearn.save_model( sk_model=custom_transformer_model, path=cloudpickle_format_model_path, serialization_format=mlflow.sklearn.SERIALIZATION_FORMAT_CLOUDPICKLE, ) reloaded_custom_transformer_model = mlflow.sklearn.load_model( model_uri=cloudpickle_format_model_path ) np.testing.assert_array_equal( custom_transformer_model.predict(sklearn_custom_transformer_model.inference_data), reloaded_custom_transformer_model.predict(sklearn_custom_transformer_model.inference_data), ) def test_model_save_persists_specified_conda_env_in_mlflow_model_directory( sklearn_logreg_model, model_path, sklearn_custom_env ): mlflow.sklearn.save_model( sk_model=sklearn_logreg_model.model, path=model_path, conda_env=sklearn_custom_env, ) pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME) saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV]["conda"]) assert os.path.exists(saved_conda_env_path) assert saved_conda_env_path != sklearn_custom_env with open(sklearn_custom_env) as f: sklearn_custom_env_parsed = yaml.safe_load(f) with open(saved_conda_env_path) as f: saved_conda_env_parsed = yaml.safe_load(f) assert saved_conda_env_parsed == sklearn_custom_env_parsed def test_model_save_persists_requirements_in_mlflow_model_directory( sklearn_knn_model, model_path, sklearn_custom_env ): mlflow.sklearn.save_model( sk_model=sklearn_knn_model.model, path=model_path, conda_env=sklearn_custom_env, skops_trusted_types=sklearn_knn_model_skops_trusted_types, ) saved_pip_req_path = os.path.join(model_path, "requirements.txt") _compare_conda_env_requirements(sklearn_custom_env, saved_pip_req_path) def test_log_model_with_pip_requirements(sklearn_knn_model, tmp_path): expected_mlflow_version = _mlflow_major_version_string() # Path to a requirements file req_file = tmp_path.joinpath("requirements.txt") req_file.write_text("a") with mlflow.start_run(): model_info = mlflow.sklearn.log_model( sklearn_knn_model.model, name="model", pip_requirements=str(req_file), skops_trusted_types=sklearn_knn_model_skops_trusted_types, ) _assert_pip_requirements(model_info.model_uri, [expected_mlflow_version, "a"], strict=True) # List of requirements with mlflow.start_run(): model_info = mlflow.sklearn.log_model( sklearn_knn_model.model, name="model", pip_requirements=[f"-r {req_file}", "b"], skops_trusted_types=sklearn_knn_model_skops_trusted_types, ) _assert_pip_requirements( model_info.model_uri, [expected_mlflow_version, "a", "b"], strict=True ) # Constraints file with mlflow.start_run(): model_info = mlflow.sklearn.log_model( sklearn_knn_model.model, name="model", pip_requirements=[f"-c {req_file}", "b"], skops_trusted_types=sklearn_knn_model_skops_trusted_types, ) _assert_pip_requirements( model_info.model_uri, [expected_mlflow_version, "b", "-c constraints.txt"], ["a"], strict=True, ) def test_log_model_with_extra_pip_requirements(sklearn_knn_model, tmp_path): expected_mlflow_version = _mlflow_major_version_string() default_reqs = mlflow.sklearn.get_default_pip_requirements(include_skops=True) # Path to a requirements file req_file = tmp_path.joinpath("requirements.txt") req_file.write_text("a") with mlflow.start_run(): model_info = mlflow.sklearn.log_model( sklearn_knn_model.model, name="model", extra_pip_requirements=str(req_file), skops_trusted_types=sklearn_knn_model_skops_trusted_types, ) _assert_pip_requirements( model_info.model_uri, [expected_mlflow_version, *default_reqs, "a"] ) # List of requirements with mlflow.start_run(): model_info = mlflow.sklearn.log_model( sklearn_knn_model.model, name="model", extra_pip_requirements=[f"-r {req_file}", "b"], skops_trusted_types=sklearn_knn_model_skops_trusted_types, ) _assert_pip_requirements( model_info.model_uri, [expected_mlflow_version, *default_reqs, "a", "b"] ) # Constraints file with mlflow.start_run(): model_info = mlflow.sklearn.log_model( sklearn_knn_model.model, name="model", extra_pip_requirements=[f"-c {req_file}", "b"], skops_trusted_types=sklearn_knn_model_skops_trusted_types, ) _assert_pip_requirements( model_info.model_uri, [expected_mlflow_version, *default_reqs, "b", "-c constraints.txt"], ["a"], ) def test_model_save_accepts_conda_env_as_dict(sklearn_knn_model, model_path): conda_env = dict(mlflow.sklearn.get_default_conda_env()) conda_env["dependencies"].append("pytest") mlflow.sklearn.save_model( sk_model=sklearn_knn_model.model, path=model_path, conda_env=conda_env, skops_trusted_types=sklearn_knn_model_skops_trusted_types, ) pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME) saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV]["conda"]) assert os.path.exists(saved_conda_env_path) with open(saved_conda_env_path) as f: saved_conda_env_parsed = yaml.safe_load(f) assert saved_conda_env_parsed == conda_env def test_model_log_persists_specified_conda_env_in_mlflow_model_directory( sklearn_knn_model, sklearn_custom_env ): artifact_path = "model" with mlflow.start_run(): model_info = mlflow.sklearn.log_model( sklearn_knn_model.model, name=artifact_path, conda_env=sklearn_custom_env, skops_trusted_types=sklearn_knn_model_skops_trusted_types, ) model_uri = model_info.model_uri model_path = _download_artifact_from_uri(artifact_uri=model_uri) pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME) saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV]["conda"]) assert os.path.exists(saved_conda_env_path) assert saved_conda_env_path != sklearn_custom_env with open(sklearn_custom_env) as f: sklearn_custom_env_parsed = yaml.safe_load(f) with open(saved_conda_env_path) as f: saved_conda_env_parsed = yaml.safe_load(f) assert saved_conda_env_parsed == sklearn_custom_env_parsed def test_model_log_persists_requirements_in_mlflow_model_directory( sklearn_knn_model, sklearn_custom_env ): with mlflow.start_run(): model_info = mlflow.sklearn.log_model( sklearn_knn_model.model, name="model", conda_env=sklearn_custom_env, skops_trusted_types=sklearn_knn_model_skops_trusted_types, ) model_path = _download_artifact_from_uri(artifact_uri=model_info.model_uri) saved_pip_req_path = os.path.join(model_path, "requirements.txt") _compare_conda_env_requirements(sklearn_custom_env, saved_pip_req_path) def test_model_save_throws_exception_if_serialization_format_is_unrecognized( sklearn_knn_model, model_path ): with pytest.raises(MlflowException, match="Unrecognized serialization format") as exc: mlflow.sklearn.save_model( sk_model=sklearn_knn_model.model, path=model_path, serialization_format="not a valid format", ) assert exc.value.error_code == ErrorCode.Name(INVALID_PARAMETER_VALUE) # The unsupported serialization format should have been detected prior to the execution of # any directory creation or state-mutating persistence logic that would prevent a second # serialization call with the same model path from succeeding assert not os.path.exists(model_path) mlflow.sklearn.save_model( sk_model=sklearn_knn_model.model, path=model_path, skops_trusted_types=sklearn_knn_model_skops_trusted_types, ) def test_model_save_without_specified_conda_env_uses_default_env_with_expected_dependencies( sklearn_logreg_model, model_path ): mlflow.sklearn.save_model(sk_model=sklearn_logreg_model.model, path=model_path) _assert_pip_requirements( model_path, mlflow.sklearn.get_default_pip_requirements(include_skops=True) ) def test_model_log_without_specified_conda_env_uses_default_env_with_expected_dependencies( sklearn_logreg_model, ): with mlflow.start_run(): model_info = mlflow.sklearn.log_model(sklearn_logreg_model.model, name="model") _assert_pip_requirements( model_info.model_uri, mlflow.sklearn.get_default_pip_requirements(include_skops=True) ) def test_model_save_uses_skops_serialization_format_by_default(sklearn_logreg_model, model_path): mlflow.sklearn.save_model(sk_model=sklearn_logreg_model.model, path=model_path) sklearn_conf = _get_flavor_configuration( model_path=model_path, flavor_name=mlflow.sklearn.FLAVOR_NAME ) assert "serialization_format" in sklearn_conf assert sklearn_conf["serialization_format"] == mlflow.sklearn.SERIALIZATION_FORMAT_SKOPS def test_model_log_uses_skops_serialization_format_by_default(sklearn_logreg_model): with mlflow.start_run(): model_info = mlflow.sklearn.log_model(sklearn_logreg_model.model, name="model") model_path = _download_artifact_from_uri(artifact_uri=model_info.model_uri) sklearn_conf = _get_flavor_configuration( model_path=model_path, flavor_name=mlflow.sklearn.FLAVOR_NAME ) assert "serialization_format" in sklearn_conf assert sklearn_conf["serialization_format"] == mlflow.sklearn.SERIALIZATION_FORMAT_SKOPS def test_model_save_with_cloudpickle_format_adds_cloudpickle_to_conda_environment( sklearn_knn_model, model_path ): mlflow.sklearn.save_model( sk_model=sklearn_knn_model.model, path=model_path, serialization_format=mlflow.sklearn.SERIALIZATION_FORMAT_CLOUDPICKLE, ) sklearn_conf = _get_flavor_configuration( model_path=model_path, flavor_name=mlflow.sklearn.FLAVOR_NAME ) assert "serialization_format" in sklearn_conf assert sklearn_conf["serialization_format"] == mlflow.sklearn.SERIALIZATION_FORMAT_CLOUDPICKLE pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME) saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV]["conda"]) assert os.path.exists(saved_conda_env_path) with open(saved_conda_env_path) as f: saved_conda_env_parsed = yaml.safe_load(f) pip_deps = [ dependency for dependency in saved_conda_env_parsed["dependencies"] if type(dependency) == dict and "pip" in dependency ] assert len(pip_deps) == 1 assert any("cloudpickle" in pip_dep for pip_dep in pip_deps[0]["pip"]) def test_model_save_without_cloudpickle_format_does_not_add_cloudpickle_to_conda_environment( sklearn_logreg_model, model_path ): non_cloudpickle_serialization_formats = list(mlflow.sklearn.SUPPORTED_SERIALIZATION_FORMATS) non_cloudpickle_serialization_formats.remove(mlflow.sklearn.SERIALIZATION_FORMAT_CLOUDPICKLE) for serialization_format in non_cloudpickle_serialization_formats: mlflow.sklearn.save_model( sk_model=sklearn_logreg_model.model, path=model_path, serialization_format=serialization_format, ) sklearn_conf = _get_flavor_configuration( model_path=model_path, flavor_name=mlflow.sklearn.FLAVOR_NAME ) assert "serialization_format" in sklearn_conf assert sklearn_conf["serialization_format"] == serialization_format pyfunc_conf = _get_flavor_configuration( model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME ) saved_conda_env_path = os.path.join(model_path, pyfunc_conf[pyfunc.ENV]["conda"]) assert os.path.exists(saved_conda_env_path) with open(saved_conda_env_path) as f: saved_conda_env_parsed = yaml.safe_load(f) assert all( "cloudpickle" not in dependency for dependency in saved_conda_env_parsed["dependencies"] ) shutil.rmtree(model_path) def test_load_pyfunc_succeeds_for_older_models_with_pyfunc_data_field( sklearn_knn_model, model_path ): """ This test verifies that scikit-learn models saved in older versions of MLflow are loaded successfully by ``mlflow.pyfunc.load_model``. These older models specify a pyfunc ``data`` field referring directly to a serialized scikit-learn model file. In contrast, newer models omit the ``data`` field. """ mlflow.sklearn.save_model( sk_model=sklearn_knn_model.model, path=model_path, serialization_format=mlflow.sklearn.SERIALIZATION_FORMAT_PICKLE, ) model_conf_path = os.path.join(model_path, "MLmodel") model_conf = Model.load(model_conf_path) pyfunc_conf = model_conf.flavors.get(pyfunc.FLAVOR_NAME) sklearn_conf = model_conf.flavors.get(mlflow.sklearn.FLAVOR_NAME) assert sklearn_conf is not None assert pyfunc_conf is not None pyfunc_conf[pyfunc.DATA] = sklearn_conf["pickled_model"] reloaded_knn_pyfunc = pyfunc.load_model(model_uri=model_path) np.testing.assert_array_equal( sklearn_knn_model.model.predict(sklearn_knn_model.inference_data), reloaded_knn_pyfunc.predict(sklearn_knn_model.inference_data), ) def test_add_pyfunc_flavor_only_when_model_defines_predict(model_path): from sklearn.cluster import AgglomerativeClustering sk_model = AgglomerativeClustering() assert not hasattr(sk_model, "predict") mlflow.sklearn.save_model( sk_model=sk_model, path=model_path, serialization_format=mlflow.sklearn.SERIALIZATION_FORMAT_PICKLE, ) model_conf_path = os.path.join(model_path, "MLmodel") model_conf = Model.load(model_conf_path) assert pyfunc.FLAVOR_NAME not in model_conf.flavors def test_pyfunc_serve_and_score(sklearn_logreg_model): model, inference_dataframe = sklearn_logreg_model artifact_path = "model" with mlflow.start_run(): model_info = mlflow.sklearn.log_model( model, name=artifact_path, input_example=inference_dataframe, serialization_format="cloudpickle", ) inference_payload = load_serving_example(model_info.model_uri) resp = pyfunc_serve_and_score_model( model_info.model_uri, data=inference_payload, content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON, extra_args=EXTRA_PYFUNC_SERVING_TEST_ARGS, ) scores = pd.DataFrame( data=json.loads(resp.content.decode("utf-8"))["predictions"] ).values.squeeze() np.testing.assert_array_almost_equal(scores, model.predict(inference_dataframe)) def test_sklearn_compatible_with_mlflow_2_4_0(sklearn_knn_model, tmp_path): model, inference_dataframe = sklearn_knn_model model_predict = model.predict(inference_dataframe) # save test model tmp_path.joinpath("MLmodel").write_text( f""" artifact_path: model flavors: python_function: env: conda: conda.yaml virtualenv: python_env.yaml loader_module: mlflow.sklearn model_path: model.pkl predict_fn: predict python_version: 3.11.15 sklearn: code: null pickled_model: model.pkl serialization_format: cloudpickle sklearn_version: {sklearn.__version__} mlflow_version: 2.4.0 model_uuid: c9833d74b1ff4013a1c9eff05d39eeef run_id: 8146a2ae86104f5b853351e600fc9d7b utc_time_created: '2023-07-04 07:19:43.561797' """ ) tmp_path.joinpath("python_env.yaml").write_text( """ python: 3.11.15 build_dependencies: - pip==25.1.1 - setuptools==80.4.0 - wheel==0.45.1 dependencies: - -r requirements.txt """ ) tmp_path.joinpath("requirements.txt").write_text( f""" mlflow==2.4.0 cloudpickle numpy psutil scikit-learn=={sklearn.__version__} scipy """ ) with open(tmp_path / "model.pkl", "wb") as out: pickle.dump(model, out, protocol=pickle.DEFAULT_PROTOCOL) assert Version(mlflow.__version__) > Version("2.4.0") model_uri = str(tmp_path) pyfunc_loaded = mlflow.pyfunc.load_model(model_uri) # predict is compatible local_predict = pyfunc_loaded.predict(inference_dataframe) np.testing.assert_array_almost_equal(local_predict, model_predict) # model serving is compatible resp = pyfunc_serve_and_score_model( model_uri, data=pd.DataFrame(inference_dataframe), content_type=pyfunc_scoring_server.CONTENT_TYPE_JSON, extra_args=EXTRA_PYFUNC_SERVING_TEST_ARGS, ) scores = pd.DataFrame( data=json.loads(resp.content.decode("utf-8"))["predictions"] ).values.squeeze() np.testing.assert_array_almost_equal(scores, model_predict) # Issues a warning if params are specified prior to MLflow support in 2.5.0 with mock.patch("mlflow.models.utils._logger.warning") as mock_warning: pyfunc_loaded.predict(inference_dataframe, params={"top_k": 2}) mock_warning.assert_called_with( "`params` can only be specified at inference time if the model signature defines a params " "schema. This model does not define a params schema. Ignoring provided params: " "['top_k']" ) def test_log_model_with_code_paths(sklearn_knn_model): artifact_path = "model" with ( mlflow.start_run(), mock.patch("mlflow.sklearn._add_code_from_conf_to_system_path") as add_mock, ): model_info = mlflow.sklearn.log_model( sklearn_knn_model.model, name=artifact_path, code_paths=[__file__], skops_trusted_types=sklearn_knn_model_skops_trusted_types, ) _compare_logged_code_paths(__file__, model_info.model_uri, mlflow.sklearn.FLAVOR_NAME) mlflow.sklearn.load_model(model_uri=model_info.model_uri) add_mock.assert_called() @pytest.mark.parametrize( "predict_fn", ["predict", "predict_proba", "predict_log_proba", "predict_joint_log_proba"] ) def test_log_model_with_custom_pyfunc_predict_fn(sklearn_gaussian_model, predict_fn): if Version(sklearn.__version__) < Version("1.2.0") and predict_fn == "predict_joint_log_proba": pytest.skip("predict_joint_log_proba is not available in scikit-learn < 1.2.0") model, inference_dataframe = sklearn_gaussian_model expected_scores = getattr(model, predict_fn)(inference_dataframe) artifact_path = "model" with mlflow.start_run(): model_info = mlflow.sklearn.log_model( model, name=artifact_path, pyfunc_predict_fn=predict_fn ) loaded_model = pyfunc.load_model(model_info.model_uri) actual_scores = loaded_model.predict(inference_dataframe) np.testing.assert_array_almost_equal(expected_scores, actual_scores) def test_virtualenv_subfield_points_to_correct_path(sklearn_logreg_model, model_path): mlflow.sklearn.save_model(sklearn_logreg_model.model, path=model_path) pyfunc_conf = _get_flavor_configuration(model_path=model_path, flavor_name=pyfunc.FLAVOR_NAME) python_env_path = Path(model_path, pyfunc_conf[pyfunc.ENV]["virtualenv"]) assert python_env_path.exists() assert python_env_path.is_file() def test_model_save_load_with_metadata(sklearn_logreg_model, model_path): mlflow.sklearn.save_model( sklearn_logreg_model.model, path=model_path, metadata={"metadata_key": "metadata_value"} ) reloaded_model = mlflow.pyfunc.load_model(model_uri=model_path) assert reloaded_model.metadata.metadata["metadata_key"] == "metadata_value" def test_model_log_with_metadata(sklearn_logreg_model): artifact_path = "model" with mlflow.start_run(): model_info = mlflow.sklearn.log_model( sklearn_logreg_model.model, name=artifact_path, metadata={"metadata_key": "metadata_value"}, ) reloaded_model = mlflow.pyfunc.load_model(model_uri=model_info.model_uri) assert reloaded_model.metadata.metadata["metadata_key"] == "metadata_value" def test_model_log_with_signature_inference(sklearn_logreg_model, iris_signature): artifact_path = "model" X = sklearn_logreg_model.inference_data example = X.iloc[[0]] with mlflow.start_run(): model_info = mlflow.sklearn.log_model( sklearn_logreg_model.model, name=artifact_path, input_example=example ) mlflow_model = Model.load(model_info.model_uri) assert mlflow_model.signature == iris_signature def test_model_size_bytes(sklearn_logreg_model, tmp_path): mlflow.sklearn.save_model(sklearn_logreg_model.model, path=tmp_path) # expected size only counts for files saved before the MLmodel file is saved model_file = tmp_path.joinpath("model.skops") with model_file.open("rb") as fp: expected_size = len(fp.read()) mlmodel = yaml.safe_load(tmp_path.joinpath("MLmodel").read_bytes()) assert mlmodel["model_size_bytes"] == expected_size def test_model_registration_metadata_handling(sklearn_logreg_model, tmp_path): artifact_path = "model" with mlflow.start_run(): mlflow.sklearn.log_model( sklearn_logreg_model.model, name=artifact_path, registered_model_name="test", ) model_uri = "models:/test/1" artifact_repository = get_artifact_repository(model_uri) dst_full = tmp_path.joinpath("full") dst_full.mkdir() artifact_repository.download_artifacts("MLmodel", dst_full) # This validates that the models artifact repo will not attempt to create a # "registered model metadata" file if the source of an artifact download is a file. assert os.listdir(dst_full) == ["MLmodel"] def test_pipeline_predict_proba(sklearn_logreg_model, model_path): logreg_model = sklearn_logreg_model.model pipeline = make_pipeline(logreg_model) mlflow.sklearn.save_model(sk_model=pipeline, path=model_path, pyfunc_predict_fn="predict_proba") reloaded_logreg_pyfunc = pyfunc.load_model(model_uri=model_path) np.testing.assert_array_equal( logreg_model.predict_proba(sklearn_logreg_model.inference_data), reloaded_logreg_pyfunc.predict(sklearn_logreg_model.inference_data), ) def test_get_raw_model(sklearn_logreg_model): with mlflow.start_run(): model_info = mlflow.sklearn.log_model( sklearn_logreg_model.model, name="model", input_example=sklearn_logreg_model.inference_data, ) pyfunc_model = pyfunc.load_model(model_info.model_uri) raw_model = pyfunc_model.get_raw_model() assert type(raw_model) == type(sklearn_logreg_model.model) np.testing.assert_array_equal( raw_model.predict(sklearn_logreg_model.inference_data), sklearn_logreg_model.model.predict(sklearn_logreg_model.inference_data), )