chore: import upstream snapshot with attribution

tests/pyfunc/docker/conftest.py

@@ -0,0 +1,78 @@
+import logging
+import os
+import subprocess
+from functools import lru_cache
+
+import docker
+import pytest
+import requests
+from packaging.version import Version
+
+import mlflow
+
+TEST_IMAGE_NAME = "test_image"
+MLFLOW_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..", ".."))
+RESOURCE_DIR = os.path.join(MLFLOW_ROOT, "tests", "resources", "dockerfile")
+
+docker_client = docker.from_env()
+
+_logger = logging.getLogger(__name__)
+
+
+@pytest.fixture(autouse=True)
+def clean_up_docker():
+    yield
+
+    # Get all containers using the test image
+    containers = docker_client.containers.list(filters={"ancestor": TEST_IMAGE_NAME})
+    for container in containers:
+        container.remove(force=True)
+
+    # Clean up the image
+    try:
+        docker_client.images.remove(TEST_IMAGE_NAME, force=True)
+    except docker.errors.ImageNotFound:
+        pass
+
+    # Clean up the build cache and volumes
+    try:
+        subprocess.check_call(["docker", "builder", "prune", "-a", "-f"])
+    except subprocess.CalledProcessError as e:
+        _logger.warning("Failed to clean up docker system: %s", e)
+
+
+@lru_cache(maxsize=1)
+def get_released_mlflow_version():
+    url = "https://pypi.org/pypi/mlflow/json"
+    response = requests.get(url)
+    response.raise_for_status()
+    data = response.json()
+    versions = [
+        v for v in map(Version, data["releases"]) if not (v.is_devrelease or v.is_prerelease)
+    ]
+    return str(max(versions))
+
+
+def save_model_with_latest_mlflow_version(flavor, extra_pip_requirements=None, **kwargs):
+    """
+    Save a model with overriding MLflow version from dev version to the latest released version.
+    By default a model is saved with the dev version of MLflow, which is not available on PyPI.
+    Usually we can be workaround this by adding --serve-wheel flag that starts local PyPI server,
+    however, this doesn't work when installing dependencies inside Docker container. Hence, this
+    function uses `extra_pip_requirements` to save the model with the latest released MLflow.
+    """
+    latest_mlflow_version = get_released_mlflow_version()
+    if flavor == "langchain":
+        kwargs["pip_requirements"] = [
+            f"mlflow[gateway]=={latest_mlflow_version}",
+            "langchain<1.1.0",
+        ]
+    else:
+        extra_pip_requirements = extra_pip_requirements or []
+        extra_pip_requirements.append(f"mlflow=={latest_mlflow_version}")
+        if flavor == "lightgbm":
+            # Adding pyarrow < 18 to prevent pip installation resolution conflicts.
+            extra_pip_requirements.append("pyarrow<18")
+        kwargs["extra_pip_requirements"] = extra_pip_requirements
+    flavor_module = getattr(mlflow, flavor)
+    flavor_module.save_model(**kwargs)

tests/pyfunc/docker/test_docker.py

@@ -0,0 +1,157 @@
+import difflib
+import os
+import shutil
+from dataclasses import dataclass
+from pathlib import Path
+from unittest import mock
+
+import pytest
+import sklearn
+import sklearn.neighbors
+from packaging.version import Version
+
+import mlflow
+from mlflow.environment_variables import _MLFLOW_RUN_SLOW_TESTS
+from mlflow.models import Model
+from mlflow.models.docker_utils import build_image_from_context
+from mlflow.models.flavor_backend_registry import get_flavor_backend
+from mlflow.utils import PYTHON_VERSION
+from mlflow.utils.env_manager import CONDA, LOCAL, VIRTUALENV
+from mlflow.version import VERSION
+
+from tests.pyfunc.docker.conftest import RESOURCE_DIR, get_released_mlflow_version
+
+
+def _get_mlflow_install_specifier():
+    if Version(VERSION).is_devrelease:
+        return "https://github.com/mlflow/mlflow/archive/refs/heads/master.zip"
+    return f"mlflow=={VERSION}"
+
+
+def assert_dockerfiles_equal(actual_dockerfile_path: Path, expected_dockerfile_path: Path):
+    actual_dockerfile = actual_dockerfile_path.read_text()
+    expected_dockerfile = (
+        expected_dockerfile_path
+        .read_text()
+        .replace("${{ MLFLOW_INSTALL }}", _get_mlflow_install_specifier())
+        .replace("${{ PYTHON_VERSION }}", PYTHON_VERSION)
+    )
+    assert actual_dockerfile == expected_dockerfile, (
+        "Generated Dockerfile does not match expected one. Diff:\n"
+        + "\n".join(
+            difflib.unified_diff(expected_dockerfile.splitlines(), actual_dockerfile.splitlines())
+        )
+    )
+
+
+def save_model(tmp_path):
+    knn_model = sklearn.neighbors.KNeighborsClassifier()
+    model_path = os.path.join(tmp_path, "model")
+    mlflow.sklearn.save_model(
+        knn_model,
+        path=model_path,
+        pip_requirements=[
+            f"mlflow=={get_released_mlflow_version()}",
+            f"scikit-learn=={sklearn.__version__}",
+        ],  # Skip requirements inference for speed up
+    )
+    return model_path
+
+
+def add_spark_flavor_to_model(model_path):
+    model_config_path = os.path.join(model_path, "MLmodel")
+    model = Model.load(model_config_path)
+    model.add_flavor("spark", spark_version="3.5.0")
+    model.save(model_config_path)
+
+
+@dataclass
+class Param:
+    expected_dockerfile: str
+    env_manager: str | None = None
+    mlflow_home: str | None = None
+    install_mlflow: bool = False
+    # If True, image is built with --model-uri param
+    specify_model_uri: bool = True
+
+
+@pytest.mark.parametrize(
+    "params",
+    [
+        Param(expected_dockerfile="Dockerfile_default"),
+        Param(expected_dockerfile="Dockerfile_default", env_manager=LOCAL),
+        Param(expected_dockerfile="Dockerfile_java_flavor", env_manager=VIRTUALENV),
+        Param(expected_dockerfile="Dockerfile_conda", env_manager=CONDA),
+        Param(install_mlflow=True, expected_dockerfile="Dockerfile_install_mlflow"),
+        Param(mlflow_home=".", expected_dockerfile="Dockerfile_with_mlflow_home"),
+        Param(specify_model_uri=False, expected_dockerfile="Dockerfile_no_model_uri"),
+    ],
+)
+def test_build_image(tmp_path, params):
+    model_uri = save_model(tmp_path) if params.specify_model_uri else None
+
+    backend = get_flavor_backend(model_uri, docker_build=True, env_manager=params.env_manager)
+
+    # Copy the context dir to a temp dir so we can verify the generated Dockerfile
+    def _build_image_with_copy(context_dir, image_name):
+        shutil.copytree(context_dir, dst_dir)
+        # Build the image if the slow-tests flag is enabled
+        if _MLFLOW_RUN_SLOW_TESTS.get():
+            for _ in range(3):
+                try:
+                    # Docker image build is unstable on GitHub Actions, retry up to 3 times
+                    build_image_from_context(context_dir, image_name)
+                    break
+                except RuntimeError:
+                    pass
+            else:
+                raise RuntimeError("Docker image build failed.")
+
+    dst_dir = tmp_path / "context"
+    with mock.patch(
+        "mlflow.models.docker_utils.build_image_from_context",
+        side_effect=_build_image_with_copy,
+    ):
+        backend.build_image(
+            model_uri=model_uri,
+            image_name="test_image",
+            mlflow_home=params.mlflow_home,
+            install_mlflow=params.install_mlflow,
+        )
+
+    actual = dst_dir / "Dockerfile"
+    expected = Path(RESOURCE_DIR) / params.expected_dockerfile
+    assert_dockerfiles_equal(actual, expected)
+
+
+def test_generate_dockerfile_for_java_flavor(tmp_path):
+    model_path = save_model(tmp_path)
+    add_spark_flavor_to_model(model_path)
+
+    backend = get_flavor_backend(model_path, docker_build=True, env_manager=None)
+
+    backend.generate_dockerfile(
+        model_uri=model_path,
+        output_dir=tmp_path,
+    )
+
+    actual = tmp_path / "Dockerfile"
+    expected = Path(RESOURCE_DIR) / "Dockerfile_java_flavor"
+    assert_dockerfiles_equal(actual, expected)
+
+
+def test_generate_dockerfile_for_custom_image(tmp_path):
+    model_path = save_model(tmp_path)
+    add_spark_flavor_to_model(model_path)
+
+    backend = get_flavor_backend(model_path, docker_build=True, env_manager=None)
+
+    backend.generate_dockerfile(
+        base_image="quay.io/jupyter/scipy-notebook:latest",
+        model_uri=model_path,
+        output_dir=tmp_path,
+    )
+
+    actual = tmp_path / "Dockerfile"
+    expected = Path(RESOURCE_DIR) / "Dockerfile_custom_scipy"
+    assert_dockerfiles_equal(actual, expected)

tests/pyfunc/docker/test_docker_flavors.py

@@ -0,0 +1,403 @@
+import contextlib
+import os
+import shutil
+import sys
+import threading
+import time
+
+import pandas as pd
+import pytest
+import requests
+
+import mlflow
+from mlflow.environment_variables import _MLFLOW_RUN_SLOW_TESTS
+from mlflow.models.flavor_backend_registry import get_flavor_backend
+from mlflow.models.utils import load_serving_example
+
+# Only import model fixtures if when MLFLOW_RUN_SLOW_TESTS environment variable is set to true
+if _MLFLOW_RUN_SLOW_TESTS.get():
+    from tests.catboost.test_catboost_model_export import reg_model  # noqa: F401
+    from tests.h2o.test_h2o_model_export import h2o_iris_model  # noqa: F401
+    from tests.helper_functions import get_safe_port
+    from tests.langchain.test_langchain_model_export import fake_chat_model  # noqa: F401
+    from tests.lightgbm.test_lightgbm_model_export import lgb_model  # noqa: F401
+    from tests.models.test_model import iris_data, sklearn_knn_model  # noqa: F401
+    from tests.pmdarima.test_pmdarima_model_export import (  # noqa: F401
+        auto_arima_object_model,
+        test_data,
+    )
+    from tests.prophet.test_prophet_model_export import (
+        prophet_model as prophet_raw_model,  # noqa: F401
+    )
+    from tests.pyfunc.docker.conftest import (
+        MLFLOW_ROOT,
+        TEST_IMAGE_NAME,
+        docker_client,
+        save_model_with_latest_mlflow_version,
+    )
+    from tests.spacy.test_spacy_model_export import spacy_model_with_data  # noqa: F401
+    from tests.spark.test_spark_model_export import (  # noqa: F401
+        iris_df,
+        spark,
+        spark_model_iris,
+    )
+    from tests.statsmodels.model_fixtures import ols_model
+    from tests.tensorflow.test_tensorflow2_core_model_export import tf2_toy_model  # noqa: F401
+    from tests.transformers.helper import load_text_classification_pipeline
+
+
+pytestmark = pytest.mark.skipif(
+    not _MLFLOW_RUN_SLOW_TESTS.get(),
+    reason="Skip slow tests. Set MLFLOW_RUN_SLOW_TESTS environment variable to run them.",
+)
+
+
+@pytest.fixture
+def model_path(tmp_path):
+    model_path = tmp_path.joinpath("model")
+
+    yield model_path
+
+    # Pytest keeps the temporary directory created by `tmp_path` fixture for 3 recent test sessions
+    # by default. This is useful for debugging during local testing, but in CI it just wastes the
+    # disk space.
+    if os.environ.get("GITHUB_ACTIONS") == "true":
+        shutil.rmtree(model_path, ignore_errors=True)
+
+
+@contextlib.contextmanager
+def start_container(port: int):
+    container = docker_client.containers.run(
+        image=TEST_IMAGE_NAME,
+        ports={8080: port},
+        detach=True,
+    )
+
+    def stream_logs():
+        for line in container.logs(stream=True):
+            sys.stdout.write(line.decode("utf-8"))
+
+    # Start a thread to stream logs from the container
+    t = threading.Thread(name="docker-log-stream", target=stream_logs, daemon=True)
+    t.start()
+
+    try:
+        # Wait for the server to start
+        for _ in range(30):
+            try:
+                response = requests.get(url=f"http://localhost:{port}/ping")
+                if response.ok:
+                    break
+            except requests.exceptions.ConnectionError as e:
+                sys.stdout.write(f"An exception occurred when calling the server: {e}\n")
+
+            container.reload()  # update container status
+            if container.status == "exited":
+                raise Exception("Container exited unexpectedly.")
+
+            sys.stdout.write(f"Container status: {container.status}\n")
+            time.sleep(5)
+
+        else:
+            raise TimeoutError("Failed to start server.")
+
+        yield container
+    finally:
+        container.stop()
+        container.remove()
+        t.join(timeout=5)
+
+
+@pytest.mark.parametrize(
+    ("flavor"),
+    [
+        "catboost",
+        "h2o",
+        # "johnsnowlabs", # Couldn't test JohnSnowLab locally due to license issue
+        "keras",
+        "langchain",
+        "lightgbm",
+        "onnx",
+        # "openai", # OPENAI API KEY is not necessarily available for everyone
+        # "paddle",  # Disabled: https://github.com/PaddlePaddle/PaddleOCR/issues/16402
+        "pmdarima",
+        "prophet",
+        "pyfunc",
+        "pytorch",
+        "sklearn",
+        "spacy",
+        "spark",
+        "statsmodels",
+        "tensorflow",
+        "transformers_pt",  # Test with Pytorch-based model
+    ],
+)
+def test_build_image_and_serve(flavor, request):
+    model_path = str(request.getfixturevalue(f"{flavor}_model"))
+    flavor = flavor.split("_")[0]  # Remove _pt or _tf from the flavor name
+
+    # Build an image
+    backend = get_flavor_backend(model_uri=model_path, docker_build=True, env_manager=None)
+    backend.build_image(
+        model_uri=model_path,
+        image_name=TEST_IMAGE_NAME,
+        mlflow_home=MLFLOW_ROOT,  # Required to prevent installing dev version of MLflow from PyPI
+    )
+
+    # Run a container
+    port = get_safe_port()
+    with start_container(port):
+        # Make a scoring request with a saved serving input example
+        inference_payload = load_serving_example(model_path)
+
+        response = requests.post(
+            url=f"http://localhost:{port}/invocations",
+            data=inference_payload,
+            headers={"Content-Type": "application/json"},
+        )
+
+        assert response.status_code == 200, f"Response: {response.text}"
+        if flavor == "langchain":
+            # "messages" key is unified llm input, output is not wrapped into predictions
+            assert response.json() == ["Hi"]
+        else:
+            assert "predictions" in response.json(), f"Response: {response.text}"
+
+
+@pytest.fixture
+def catboost_model(model_path, reg_model):
+    save_model_with_latest_mlflow_version(
+        flavor="catboost",
+        cb_model=reg_model.model,
+        path=model_path,
+        input_example=reg_model.inference_dataframe[:1],
+    )
+    return model_path
+
+
+@pytest.fixture
+def h2o_model(model_path, h2o_iris_model):
+    save_model_with_latest_mlflow_version(
+        flavor="h2o",
+        h2o_model=h2o_iris_model.model,
+        path=model_path,
+        input_example=h2o_iris_model.inference_data.as_data_frame()[:1],
+    )
+    return model_path
+
+
+@pytest.fixture
+def keras_model(model_path, iris_data):
+    from sklearn import datasets
+    from tensorflow.keras.layers import Dense
+    from tensorflow.keras.models import Sequential
+
+    model = Sequential()
+    model.add(Dense(3, input_dim=4))
+    model.add(Dense(1))
+
+    X, y = datasets.load_iris(return_X_y=True)
+    save_model_with_latest_mlflow_version(
+        flavor="tensorflow",
+        model=model,
+        path=model_path,
+        input_example=X[:3, :],
+    )
+    return model_path
+
+
+@pytest.fixture
+def langchain_model(model_path, tmp_path):
+    # LangChain v1+ requires models-from-code
+    model_code = """
+from operator import itemgetter
+from langchain_core.runnables import RunnablePassthrough
+import mlflow
+
+mlflow.models.set_model(RunnablePassthrough() | itemgetter("messages"))
+"""
+    code_path = tmp_path / "langchain_model.py"
+    code_path.write_text(model_code)
+
+    save_model_with_latest_mlflow_version(
+        flavor="langchain",
+        lc_model=str(code_path),
+        path=model_path,
+        input_example={"messages": "Hi"},
+    )
+    return model_path
+
+
+@pytest.fixture
+def lightgbm_model(model_path, lgb_model):
+    save_model_with_latest_mlflow_version(
+        flavor="lightgbm",
+        lgb_model=lgb_model.model,
+        path=model_path,
+        input_example=lgb_model.inference_dataframe.to_numpy()[:1],
+    )
+    return model_path
+
+
+@pytest.fixture
+def onnx_model(tmp_path, model_path):
+    import numpy as np
+    import onnx
+    import torch
+    from torch import nn
+
+    model = torch.nn.Sequential(nn.Linear(4, 3), nn.ReLU(), nn.Linear(3, 1))
+    onnx_model_path = os.path.join(tmp_path, "torch_onnx")
+    torch.onnx.export(
+        model,
+        torch.randn(1, 4),
+        onnx_model_path,
+        dynamic_axes={"input": {0: "batch"}},
+        input_names=["input"],
+    )
+    onnx_model = onnx.load(onnx_model_path)
+
+    model_path = str(tmp_path / "onnx_model")
+    save_model_with_latest_mlflow_version(
+        flavor="onnx",
+        onnx_model=onnx_model,
+        path=model_path,
+        input_example=np.random.rand(1, 4).astype(np.float32),
+    )
+    return model_path
+
+
+# Paddle fixture disabled: https://github.com/PaddlePaddle/PaddleOCR/issues/16402
+# @pytest.fixture
+# def paddle_model(model_path, pd_model):
+#     save_model_with_latest_mlflow_version(
+#         flavor="paddle",
+#         pd_model=pd_model.model,
+#         path=model_path,
+#         input_example=pd_model.inference_dataframe[:1],
+#     )
+#     return model_path
+
+
+@pytest.fixture
+def pmdarima_model(model_path, auto_arima_object_model):
+    save_model_with_latest_mlflow_version(
+        flavor="pmdarima",
+        pmdarima_model=auto_arima_object_model,
+        path=model_path,
+        input_example=pd.DataFrame({"n_periods": [30]}),
+    )
+    return model_path
+
+
+@pytest.fixture
+def prophet_model(model_path, prophet_raw_model):
+    save_model_with_latest_mlflow_version(
+        flavor="prophet",
+        pr_model=prophet_raw_model.model,
+        path=model_path,
+        input_example=prophet_raw_model.data[:1],
+        # Prophet does not handle numpy 2 yet. https://github.com/facebook/prophet/issues/2595
+        extra_pip_requirements=["numpy<2"],
+    )
+    return model_path
+
+
+@pytest.fixture
+def pyfunc_model(model_path):
+    class CustomModel(mlflow.pyfunc.PythonModel):
+        def __init__(self):
+            pass
+
+        def predict(self, context, model_input):
+            return model_input
+
+    save_model_with_latest_mlflow_version(
+        flavor="pyfunc",
+        python_model=CustomModel(),
+        path=model_path,
+        input_example=[1, 2, 3],
+    )
+    return model_path
+
+
+@pytest.fixture
+def pytorch_model(model_path):
+    from torch import nn, randn
+
+    model = nn.Sequential(nn.Linear(4, 3), nn.ReLU(), nn.Linear(3, 1))
+    save_model_with_latest_mlflow_version(
+        flavor="pytorch",
+        pytorch_model=model,
+        path=model_path,
+        input_example=randn(1, 4).numpy(),
+    )
+    return model_path
+
+
+@pytest.fixture
+def sklearn_model(model_path, sklearn_knn_model, iris_data):
+    save_model_with_latest_mlflow_version(
+        flavor="sklearn",
+        sk_model=sklearn_knn_model,
+        path=model_path,
+        input_example=iris_data[0][:1],
+    )
+    return model_path
+
+
+@pytest.fixture
+def spacy_model(model_path, spacy_model_with_data):
+    save_model_with_latest_mlflow_version(
+        flavor="spacy",
+        spacy_model=spacy_model_with_data.model,
+        path=model_path,
+        input_example=spacy_model_with_data.inference_data[:1],
+    )
+    return model_path
+
+
+@pytest.fixture
+def spark_model(model_path, spark_model_iris):
+    save_model_with_latest_mlflow_version(
+        flavor="spark",
+        spark_model=spark_model_iris.model,
+        path=model_path,
+        input_example=spark_model_iris.spark_df.toPandas()[:1],
+    )
+    return model_path
+
+
+@pytest.fixture
+def statsmodels_model(model_path):
+    model = ols_model()
+    save_model_with_latest_mlflow_version(
+        flavor="statsmodels",
+        statsmodels_model=model.model,
+        path=model_path,
+        input_example=model.inference_dataframe[:1],
+    )
+    return model_path
+
+
+@pytest.fixture
+def tensorflow_model(model_path, tf2_toy_model):
+    save_model_with_latest_mlflow_version(
+        flavor="tensorflow",
+        model=tf2_toy_model.model,
+        path=model_path,
+        input_example=tf2_toy_model.inference_data[:1],
+    )
+    return model_path
+
+
+@pytest.fixture
+def transformers_pt_model(model_path):
+    pipeline = load_text_classification_pipeline()
+    save_model_with_latest_mlflow_version(
+        flavor="transformers",
+        transformers_model=pipeline,
+        path=model_path,
+        input_example="hi",
+    )
+    return model_path