import json import os import re import shutil import subprocess import sys import warnings from dataclasses import dataclass from io import BytesIO, StringIO from pathlib import Path from unittest import mock import numpy as np import pandas as pd import pytest import sklearn import sklearn.datasets import sklearn.linear_model from click.testing import CliRunner from packaging.requirements import Requirement import mlflow import mlflow.models.cli as models_cli import mlflow.sklearn from mlflow.environment_variables import MLFLOW_DISABLE_ENV_MANAGER_CONDA_WARNING from mlflow.exceptions import MlflowException from mlflow.models.flavor_backend_registry import get_flavor_backend from mlflow.models.model import get_model_requirements_files, update_model_requirements from mlflow.models.utils import load_serving_example from mlflow.protos.databricks_pb2 import BAD_REQUEST, ErrorCode from mlflow.pyfunc.backend import PyFuncBackend from mlflow.pyfunc.scoring_server import ( CONTENT_TYPE_CSV, CONTENT_TYPE_JSON, ) from mlflow.store.artifact.models_artifact_repo import ModelsArtifactRepository from mlflow.utils import PYTHON_VERSION from mlflow.utils import env_manager as _EnvManager from mlflow.utils.conda import _get_conda_env_name from mlflow.utils.environment import ( _get_requirements_from_file, _mlflow_conda_env, ) from mlflow.utils.file_utils import TempDir from mlflow.utils.process import ShellCommandException from tests.helper_functions import ( RestEndpoint, get_safe_port, pyfunc_build_image, pyfunc_generate_dockerfile, pyfunc_serve_and_score_model, pyfunc_serve_from_docker_image, pyfunc_serve_from_docker_image_with_env_override, ) # NB: for now, windows tests do not have conda available. no_conda = ["--env-manager", "local"] if sys.platform == "win32" else [] # NB: need to install mlflow since the pip version does not have mlflow models cli. install_mlflow = ["--install-mlflow"] if not no_conda else [] extra_options = no_conda + install_mlflow def env_with_tracking_uri() -> dict[str, str]: return {**os.environ, "MLFLOW_TRACKING_URI": mlflow.get_tracking_uri()} @pytest.fixture(scope="module") def iris_data() -> tuple[np.ndarray, np.ndarray]: iris = sklearn.datasets.load_iris() x = iris.data[:, :2] y = iris.target return x, y @pytest.fixture(scope="module") def sk_model(iris_data: tuple[np.ndarray, np.ndarray]) -> sklearn.linear_model.LogisticRegression: x, y = iris_data logreg_model = sklearn.linear_model.LogisticRegression() logreg_model.fit(x, y) return logreg_model @pytest.mark.allow_infer_pip_requirements_fallback def test_mlflow_is_not_installed_unless_specified(): if no_conda: pytest.skip("This test requires conda.") with TempDir(chdr=True) as tmp: fake_model_path = tmp.path("fake_model") mlflow.pyfunc.save_model(fake_model_path, loader_module=__name__) # Overwrite the logged `conda.yaml` to remove mlflow. _mlflow_conda_env(path=os.path.join(fake_model_path, "conda.yaml"), install_mlflow=False) # The following should fail because there should be no mlflow in the env: prc = subprocess.run( [ sys.executable, "-m", "mlflow", "models", "predict", "-m", fake_model_path, "--env-manager", "conda", ], stderr=subprocess.PIPE, cwd=tmp.path(""), check=False, text=True, env=env_with_tracking_uri(), ) assert prc.returncode != 0 if PYTHON_VERSION.startswith("3"): assert "ModuleNotFoundError: No module named 'mlflow'" in prc.stderr else: assert "ImportError: No module named mlflow.pyfunc.scoring_server" in prc.stderr def test_model_with_no_deployable_flavors_fails_pollitely(): from mlflow.models import Model with TempDir(chdr=True) as tmp: m = Model( artifact_path=None, run_id=None, utc_time_created="now", flavors={"some": {}, "useless": {}, "flavors": {}}, ) os.mkdir(tmp.path("model")) m.save(tmp.path("model", "MLmodel")) # The following should fail because there should be no suitable flavor prc = subprocess.run( [sys.executable, "-m", "mlflow", "models", "predict", "-m", tmp.path("model")], stderr=subprocess.PIPE, cwd=tmp.path(""), check=False, text=True, env=env_with_tracking_uri(), ) assert "No suitable flavor backend was found for the model." in prc.stderr def test_serve_uvicorn_opts(iris_data, sk_model): if sys.platform == "win32": pytest.skip("This test requires gunicorn which is not available on windows.") with mlflow.start_run(): x, _ = iris_data model_info = mlflow.sklearn.log_model( sk_model, name="model", registered_model_name="test", input_example=pd.DataFrame(x) ) model_uris = ["models:/test/None", model_info.model_uri] for model_uri in model_uris: with TempDir() as tpm: output_file_path = tpm.path("stdout") inference_payload = load_serving_example(model_uri) with open(output_file_path, "w") as output_file: scoring_response = pyfunc_serve_and_score_model( model_uri, inference_payload, content_type=CONTENT_TYPE_JSON, stdout=output_file, extra_args=["-w", "3", "--env-manager", "local"], ) with open(output_file_path) as output_file: stdout = output_file.read() actual = pd.read_json(BytesIO(scoring_response.content), orient="records") actual = actual[actual.columns[0]].values expected = sk_model.predict(x) assert all(expected == actual) expected_command_pattern = re.compile( r"uvicorn.*--workers 3.*mlflow\.pyfunc\.scoring_server\.app:app" ) assert expected_command_pattern.search(stdout) is not None @dataclass class PredictTestData: model_uri: str model_registry_uri: str input_json_path: Path input_csv_path: Path output_json_path: Path x: np.ndarray sk_model: sklearn.base.BaseEstimator @pytest.fixture def predict_test_setup( iris_data: tuple[np.ndarray, np.ndarray], sk_model: sklearn.linear_model.LogisticRegression, tmp_path: Path, ) -> PredictTestData: with mlflow.start_run() as active_run: mlflow.sklearn.log_model(sk_model, name="model", registered_model_name="impredicting") model_uri = f"runs:/{active_run.info.run_id}/model" model_registry_uri = "models:/impredicting/None" input_json_path = tmp_path / "input.json" input_csv_path = tmp_path / "input.csv" output_json_path = tmp_path / "output.json" x, _ = iris_data with open(input_json_path, "w") as f: json.dump({"dataframe_split": pd.DataFrame(x).to_dict(orient="split")}, f) pd.DataFrame(x).to_csv(input_csv_path, index=False) return PredictTestData( model_uri=model_uri, model_registry_uri=model_registry_uri, input_json_path=input_json_path, input_csv_path=input_csv_path, output_json_path=output_json_path, x=x, sk_model=sk_model, ) def test_predict_with_model_registry_uri(predict_test_setup: PredictTestData) -> None: setup = predict_test_setup subprocess.check_call( [ sys.executable, "-m", "mlflow", "models", "predict", "-m", setup.model_registry_uri, "-i", setup.input_json_path, "-o", setup.output_json_path, "--env-manager", "local", ], env=env_with_tracking_uri(), ) actual = pd.read_json(setup.output_json_path, orient="records") actual = actual[actual.columns[0]].values expected = setup.sk_model.predict(setup.x) assert all(expected == actual) def test_predict_with_conda_and_install_mlflow(predict_test_setup: PredictTestData) -> None: setup = predict_test_setup subprocess.check_call( [ sys.executable, "-m", "mlflow", "models", "predict", "-m", setup.model_uri, "-i", setup.input_json_path, "-o", setup.output_json_path, *extra_options, ], env=env_with_tracking_uri(), ) actual = pd.read_json(setup.output_json_path, orient="records") actual = actual[actual.columns[0]].values expected = setup.sk_model.predict(setup.x) assert all(expected == actual) def test_predict_explicit_json_format_default_orient(predict_test_setup: PredictTestData) -> None: setup = predict_test_setup subprocess.check_call( [ sys.executable, "-m", "mlflow", "models", "predict", "-m", setup.model_uri, "-i", setup.input_json_path, "-o", setup.output_json_path, "-t", "json", *extra_options, ], env=env_with_tracking_uri(), ) actual = pd.read_json(setup.output_json_path, orient="records") actual = actual[actual.columns[0]].values expected = setup.sk_model.predict(setup.x) assert all(expected == actual) def test_predict_explicit_json_format_split_orient(predict_test_setup: PredictTestData) -> None: # Note: This test has the same command as the previous one but tests orient==split # The comment in original code mentions this should be split orient setup = predict_test_setup subprocess.check_call( [ sys.executable, "-m", "mlflow", "models", "predict", "-m", setup.model_uri, "-i", setup.input_json_path, "-o", setup.output_json_path, "-t", "json", *extra_options, ], env=env_with_tracking_uri(), ) actual = pd.read_json(setup.output_json_path, orient="records") actual = actual[actual.columns[0]].values expected = setup.sk_model.predict(setup.x) assert all(expected == actual) def test_predict_stdin_stdout(predict_test_setup: PredictTestData) -> None: setup = predict_test_setup stdout = subprocess.check_output( [ sys.executable, "-m", "mlflow", "models", "predict", "-m", setup.model_uri, "-t", "json", *extra_options, ], input=setup.input_json_path.read_text(), env=env_with_tracking_uri(), text=True, ) predictions = re.search(r"{\"predictions\": .*}", stdout).group(0) actual = pd.read_json(StringIO(predictions), orient="records") actual = actual[actual.columns[0]].values expected = setup.sk_model.predict(setup.x) assert all(expected == actual) # NB: We do not test orient=records here because records may loose column ordering. # orient == records is tested in other test with simpler model. def test_predict_csv_format(predict_test_setup: PredictTestData) -> None: setup = predict_test_setup subprocess.check_call( [ sys.executable, "-m", "mlflow", "models", "predict", "-m", setup.model_uri, "-i", setup.input_csv_path, "-o", setup.output_json_path, "-t", "csv", *extra_options, ], env=env_with_tracking_uri(), ) actual = pd.read_json(setup.output_json_path, orient="records") actual = actual[actual.columns[0]].values expected = setup.sk_model.predict(setup.x) assert all(expected == actual) def test_predict_check_content_type(iris_data, sk_model, tmp_path): with mlflow.start_run(): mlflow.sklearn.log_model(sk_model, name="model", registered_model_name="impredicting") model_registry_uri = "models:/impredicting/None" input_json_path = tmp_path / "input.json" input_csv_path = tmp_path / "input.csv" output_json_path = tmp_path / "output.json" x, _ = iris_data with input_json_path.open("w") as f: json.dump({"dataframe_split": pd.DataFrame(x).to_dict(orient="split")}, f) pd.DataFrame(x).to_csv(input_csv_path, index=False) # Throw errors for invalid content_type prc = subprocess.run( [ sys.executable, "-m", "mlflow", "models", "predict", "-m", model_registry_uri, "-i", input_json_path, "-o", output_json_path, "-t", "invalid", "--env-manager", "local", ], stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env_with_tracking_uri(), check=False, ) assert prc.returncode != 0 assert "Content type must be one of json or csv." in prc.stderr.decode("utf-8") def test_predict_check_input_path(iris_data, sk_model, tmp_path): with mlflow.start_run(): mlflow.sklearn.log_model(sk_model, name="model", registered_model_name="impredicting") model_registry_uri = "models:/impredicting/None" input_json_path = tmp_path / "input with space.json" input_csv_path = tmp_path / "input.csv" output_json_path = tmp_path / "output.json" x, _ = iris_data with input_json_path.open("w") as f: json.dump({"dataframe_split": pd.DataFrame(x).to_dict(orient="split")}, f) pd.DataFrame(x).to_csv(input_csv_path, index=False) # Valid input path with space prc = subprocess.run( [ sys.executable, "-m", "mlflow", "models", "predict", "-m", model_registry_uri, "-i", f"{input_json_path}", "-o", output_json_path, "--env-manager", "local", ], stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env_with_tracking_uri(), check=False, text=True, ) assert prc.returncode == 0 # Throw errors for invalid input_path prc = subprocess.run( [ sys.executable, "-m", "mlflow", "models", "predict", "-m", model_registry_uri, "-i", f'{input_json_path}"; echo ThisIsABug! "', "-o", output_json_path, "--env-manager", "local", ], stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env_with_tracking_uri(), check=False, text=True, ) assert prc.returncode != 0 assert "ThisIsABug!" not in prc.stdout assert "FileNotFoundError" in prc.stderr prc = subprocess.run( [ sys.executable, "-m", "mlflow", "models", "predict", "-m", model_registry_uri, "-i", f'{input_csv_path}"; echo ThisIsABug! "', "-o", output_json_path, "-t", "csv", "--env-manager", "local", ], stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env_with_tracking_uri(), check=False, text=True, ) assert prc.returncode != 0 assert "ThisIsABug!" not in prc.stdout assert "FileNotFoundError" in prc.stderr def test_predict_check_output_path(iris_data, sk_model, tmp_path): with mlflow.start_run(): mlflow.sklearn.log_model(sk_model, name="model", registered_model_name="impredicting") model_registry_uri = "models:/impredicting/None" input_json_path = tmp_path / "input.json" input_csv_path = tmp_path / "input.csv" output_json_path = tmp_path / "output.json" x, _ = iris_data with input_json_path.open("w") as f: json.dump({"dataframe_split": pd.DataFrame(x).to_dict(orient="split")}, f) pd.DataFrame(x).to_csv(input_csv_path, index=False) prc = subprocess.run( [ sys.executable, "-m", "mlflow", "models", "predict", "-m", model_registry_uri, "-i", input_json_path, "-o", f'{output_json_path}"; echo ThisIsABug! "', "--env-manager", "local", ], stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=env_with_tracking_uri(), check=False, text=True, ) assert prc.returncode == 0 assert "ThisIsABug!" not in prc.stdout def test_prepare_env_passes(sk_model): if no_conda: pytest.skip("This test requires conda.") with TempDir(chdr=True): with mlflow.start_run() as active_run: mlflow.sklearn.log_model(sk_model, name="model") model_uri = f"runs:/{active_run.info.run_id}/model" # With conda subprocess.run( [ sys.executable, "-m", "mlflow", "models", "prepare-env", "-m", model_uri, ], env=env_with_tracking_uri(), check=True, ) # Should be idempotent subprocess.run( [ sys.executable, "-m", "mlflow", "models", "prepare-env", "-m", model_uri, ], env=env_with_tracking_uri(), check=True, ) def test_prepare_env_fails(sk_model): if no_conda: pytest.skip("This test requires conda.") with TempDir(chdr=True): with mlflow.start_run() as active_run: mlflow.sklearn.log_model( sk_model, name="model", pip_requirements=["does-not-exist-dep==abc"] ) model_uri = f"runs:/{active_run.info.run_id}/model" # With conda - should fail due to bad conda environment. prc = subprocess.run( [ sys.executable, "-m", "mlflow", "models", "prepare-env", "-m", model_uri, ], env=env_with_tracking_uri(), check=False, ) assert prc.returncode != 0 def test_generate_dockerfile(sk_model, tmp_path): with mlflow.start_run() as active_run: mlflow.sklearn.log_model(sk_model, name="model") model_uri = f"runs:/{active_run.info.run_id}/model" output_directory = tmp_path.joinpath("output_directory") pyfunc_generate_dockerfile( output_directory, model_uri, extra_args=["--install-mlflow"], env=env_with_tracking_uri(), ) assert output_directory.is_dir() assert output_directory.joinpath("Dockerfile").exists() assert output_directory.joinpath("model_dir").is_dir() # Assert file is not empty assert output_directory.joinpath("Dockerfile").stat().st_size != 0 def test_build_docker(iris_data, sk_model): with mlflow.start_run() as active_run: mlflow.sklearn.log_model(sk_model, name="model", extra_pip_requirements=["/opt/mlflow"]) model_uri = f"runs:/{active_run.info.run_id}/model" x, _ = iris_data df = pd.DataFrame(x) image_name = pyfunc_build_image( model_uri, extra_args=["--install-mlflow"], env=env_with_tracking_uri(), ) host_port = get_safe_port() scoring_proc = pyfunc_serve_from_docker_image(image_name, host_port) _validate_with_rest_endpoint(scoring_proc, host_port, df, x, sk_model) def test_build_docker_virtualenv(iris_data, sk_model): with mlflow.start_run(): model_info = mlflow.sklearn.log_model( sk_model, name="model", extra_pip_requirements=["/opt/mlflow"] ) x, _ = iris_data df = pd.DataFrame(iris_data[0]) extra_args = ["--install-mlflow", "--env-manager", "virtualenv"] image_name = pyfunc_build_image( model_info.model_uri, extra_args=extra_args, env=env_with_tracking_uri(), ) host_port = get_safe_port() scoring_proc = pyfunc_serve_from_docker_image(image_name, host_port) _validate_with_rest_endpoint(scoring_proc, host_port, df, x, sk_model) def test_build_docker_with_env_override(iris_data, sk_model): with mlflow.start_run() as active_run: mlflow.sklearn.log_model(sk_model, name="model", extra_pip_requirements=["/opt/mlflow"]) model_uri = f"runs:/{active_run.info.run_id}/model" x, _ = iris_data df = pd.DataFrame(x) image_name = pyfunc_build_image( model_uri, extra_args=["--install-mlflow"], env=env_with_tracking_uri(), ) host_port = get_safe_port() scoring_proc = pyfunc_serve_from_docker_image_with_env_override(image_name, host_port) _validate_with_rest_endpoint(scoring_proc, host_port, df, x, sk_model) def test_build_docker_without_model_uri(iris_data, sk_model, tmp_path): model_path = tmp_path.joinpath("model") mlflow.sklearn.save_model(sk_model, model_path, extra_pip_requirements=["/opt/mlflow"]) image_name = pyfunc_build_image(model_uri=None) host_port = get_safe_port() scoring_proc = pyfunc_serve_from_docker_image_with_env_override( image_name, host_port, extra_docker_run_options=["-v", f"{model_path}:/opt/ml/model"], ) x = iris_data[0] df = pd.DataFrame(x) _validate_with_rest_endpoint(scoring_proc, host_port, df, x, sk_model) def _validate_with_rest_endpoint(scoring_proc, host_port, df, x, sk_model): with RestEndpoint(proc=scoring_proc, port=host_port, validate_version=False) as endpoint: for content_type in [CONTENT_TYPE_JSON, CONTENT_TYPE_CSV]: scoring_response = endpoint.invoke(df, content_type) assert scoring_response.status_code == 200, ( f"Failed to serve prediction, got response {scoring_response.text}" ) np.testing.assert_array_equal( np.array(json.loads(scoring_response.text)["predictions"]), sk_model.predict(x) ) # Try examples of bad input, verify we get a non-200 status code for content_type in [CONTENT_TYPE_JSON, CONTENT_TYPE_CSV, CONTENT_TYPE_JSON]: scoring_response = endpoint.invoke(data="", content_type=content_type) assert scoring_response.status_code == 400, ( "Expected server failure with error code 400, " f"got response with status code {scoring_response.status_code} " f"and body {scoring_response.text}" ) scoring_response_dict = json.loads(scoring_response.content) assert "error_code" in scoring_response_dict assert scoring_response_dict["error_code"] == ErrorCode.Name(BAD_REQUEST) assert "message" in scoring_response_dict def test_env_manager_warning_for_use_of_conda(monkeypatch): with mock.patch("mlflow.models.cli.get_flavor_backend") as mock_get_flavor_backend: with pytest.warns(UserWarning, match=r"Use of conda is discouraged"): CliRunner().invoke( models_cli.serve, ["--model-uri", "model", "--env-manager", "conda"], catch_exceptions=False, ) with warnings.catch_warnings(): warnings.simplefilter("error") monkeypatch.setenv(MLFLOW_DISABLE_ENV_MANAGER_CONDA_WARNING.name, "TRUE") CliRunner().invoke( models_cli.serve, ["--model-uri", "model", "--env-manager", "conda"], catch_exceptions=False, ) assert mock_get_flavor_backend.call_count == 2 def test_env_manager_unsupported_value(): with pytest.raises(MlflowException, match=r"Invalid value for `env_manager`"): CliRunner().invoke( models_cli.serve, ["--model-uri", "model", "--env-manager", "abc"], catch_exceptions=False, ) def test_host_invalid_value(): class MyModel(mlflow.pyfunc.PythonModel): def predict(self, context, model_input): return model_input with mlflow.start_run(): model_info = mlflow.pyfunc.log_model( name="test_model", python_model=MyModel(), registered_model_name="model" ) with mock.patch( "mlflow.models.cli.get_flavor_backend", return_value=PyFuncBackend({}, env_manager=_EnvManager.VIRTUALENV), ): with pytest.raises(ShellCommandException, match=r"Non-zero exit code: -?[1-9]\d*"): CliRunner().invoke( models_cli.serve, ["--model-uri", model_info.model_uri, "--host", "localhost & echo BUG"], catch_exceptions=False, ) def test_change_conda_env_root_location(tmp_path, sk_model): def _test_model(env_root_path, model_path, sklearn_ver): env_root_path.mkdir(exist_ok=True) mlflow.sklearn.save_model( sk_model, str(model_path), pip_requirements=[f"scikit-learn=={sklearn_ver}"] ) env = get_flavor_backend( str(model_path), env_manager=_EnvManager.CONDA, install_mlflow=False, env_root_dir=str(env_root_path), ).prepare_env(model_uri=str(model_path)) conda_env_name = _get_conda_env_name( str(model_path / "conda.yaml"), env_root_dir=env_root_path ) env_path = env_root_path / "conda_envs" / conda_env_name assert env_path.exists() python_exec_path = str(env_path / "bin" / "python") # Test execution of command under the correct activated python env. env.execute( command=f"python -c \"import sys; assert sys.executable == '{python_exec_path}'; " f"import sklearn; assert sklearn.__version__ == '{sklearn_ver}'\"", ) # Cleanup model path and Conda environment to prevent out of space failures on CI shutil.rmtree(model_path) shutil.rmtree(env_path) env_root1_path = tmp_path / "root1" env_root2_path = tmp_path / "root2" # Test with model1_path model1_path = tmp_path / "model1" _test_model(env_root1_path, model1_path, "1.4.0") _test_model(env_root2_path, model1_path, "1.4.0") # Test with model2_path model2_path = tmp_path / "model2" _test_model(env_root1_path, model2_path, "1.4.2") @pytest.mark.parametrize( ("input_schema", "output_schema", "params_schema"), [(True, False, False), (False, True, False), (False, False, True)], ) def test_signature_enforcement_with_model_serving(input_schema, output_schema, params_schema): class MyModel(mlflow.pyfunc.PythonModel): def predict(self, context, model_input, params=None): return ["test"] input_data = ["test_input"] if input_schema else None output_data = ["test_output"] if output_schema else None params = {"test": "test"} if params_schema else None signature = mlflow.models.infer_signature( model_input=input_data, model_output=output_data, params=params ) with mlflow.start_run(): model_info = mlflow.pyfunc.log_model( name="test_model", python_model=MyModel(), signature=signature ) inference_payload = json.dumps({"inputs": ["test"]}) # Serve and score the model scoring_result = pyfunc_serve_and_score_model( model_uri=model_info.model_uri, data=inference_payload, content_type=CONTENT_TYPE_JSON, extra_args=["--env-manager", "local"], ) scoring_result.raise_for_status() # Assert the prediction result assert json.loads(scoring_result.content)["predictions"] == ["test"] def assert_base_model_reqs(): """ Helper function for testing model requirements. Asserts that the contents of requirements.txt and conda.yaml are as expected, then returns their filepaths so mutations can be performed. """ import cloudpickle class MyModel(mlflow.pyfunc.PythonModel): def predict(self, context, model_input, params=None): return ["test"] with mlflow.start_run(): model_info = mlflow.pyfunc.log_model(name="model", python_model=MyModel()) resolved_uri = ModelsArtifactRepository.get_underlying_uri(model_info.model_uri) local_paths = get_model_requirements_files(resolved_uri) requirements_txt_file = local_paths.requirements conda_env_file = local_paths.conda reqs = _get_requirements_from_file(requirements_txt_file) assert Requirement(f"mlflow=={mlflow.__version__}") in reqs assert Requirement(f"cloudpickle=={cloudpickle.__version__}") in reqs reqs = _get_requirements_from_file(conda_env_file) assert Requirement(f"mlflow=={mlflow.__version__}") in reqs assert Requirement(f"cloudpickle=={cloudpickle.__version__}") in reqs return model_info.model_uri def test_update_requirements_cli_adds_reqs_successfully(): import cloudpickle model_uri = assert_base_model_reqs() CliRunner().invoke( models_cli.update_pip_requirements, ["-m", f"{model_uri}", "add", "mlflow>=2.9, !=2.9.0", "coolpackage[extra]==8.8.8"], catch_exceptions=False, ) resolved_uri = ModelsArtifactRepository.get_underlying_uri(model_uri) local_paths = get_model_requirements_files(resolved_uri) # the tool should overwrite mlflow, add coolpackage, and leave cloudpickle alone reqs = _get_requirements_from_file(local_paths.requirements) assert Requirement("mlflow!=2.9.0,>=2.9") in reqs assert Requirement("coolpackage[extra]==8.8.8") in reqs assert Requirement(f"cloudpickle=={cloudpickle.__version__}") in reqs reqs = _get_requirements_from_file(local_paths.conda) assert Requirement("mlflow!=2.9.0,>=2.9") in reqs assert Requirement("coolpackage[extra]==8.8.8") in reqs assert Requirement(f"cloudpickle=={cloudpickle.__version__}") in reqs def test_update_requirements_cli_removes_reqs_successfully(): import cloudpickle model_uri = assert_base_model_reqs() CliRunner().invoke( models_cli.update_pip_requirements, ["-m", f"{model_uri}", "remove", "mlflow"], catch_exceptions=False, ) resolved_uri = ModelsArtifactRepository.get_underlying_uri(model_uri) local_paths = get_model_requirements_files(resolved_uri) # the tool should remove mlflow and leave cloudpickle alone reqs = _get_requirements_from_file(local_paths.requirements) assert reqs == [Requirement(f"cloudpickle=={cloudpickle.__version__}")] reqs = _get_requirements_from_file(local_paths.conda) assert reqs == [Requirement(f"cloudpickle=={cloudpickle.__version__}")] def test_update_requirements_cli_throws_on_incompatible_input(): model_uri = assert_base_model_reqs() with pytest.raises( MlflowException, match="The specified requirements versions are incompatible" ): CliRunner().invoke( models_cli.update_pip_requirements, ["-m", f"{model_uri}", "add", "mlflow<2.6", "mlflow>2.7"], catch_exceptions=False, ) def test_update_model_requirements_add(): import cloudpickle model_uri = assert_base_model_reqs() update_model_requirements( model_uri, "add", ["mlflow>=2.9, !=2.9.0", "coolpackage[extra]==8.8.8"] ) resolved_uri = ModelsArtifactRepository.get_underlying_uri(model_uri) local_paths = get_model_requirements_files(resolved_uri) # the tool should overwrite mlflow, add coolpackage, and leave cloudpickle alone reqs = _get_requirements_from_file(local_paths.requirements) assert Requirement("mlflow!=2.9.0,>=2.9") in reqs assert Requirement("coolpackage[extra]==8.8.8") in reqs assert Requirement(f"cloudpickle=={cloudpickle.__version__}") in reqs reqs = _get_requirements_from_file(local_paths.conda) assert Requirement("mlflow!=2.9.0,>=2.9") in reqs assert Requirement("coolpackage[extra]==8.8.8") in reqs assert Requirement(f"cloudpickle=={cloudpickle.__version__}") in reqs def test_update_model_requirements_remove(): import cloudpickle model_uri = assert_base_model_reqs() update_model_requirements(model_uri, "remove", ["mlflow"]) resolved_uri = ModelsArtifactRepository.get_underlying_uri(model_uri) local_paths = get_model_requirements_files(resolved_uri) # the tool should remove mlflow and leave cloudpickle alone reqs = _get_requirements_from_file(local_paths.requirements) assert reqs == [Requirement(f"cloudpickle=={cloudpickle.__version__}")] reqs = _get_requirements_from_file(local_paths.conda) assert reqs == [Requirement(f"cloudpickle=={cloudpickle.__version__}")]