star

venv/lib/python3.9/site-packages/mlflow/models/__init__.py

@@ -0,0 +1,99 @@
+"""
+The ``mlflow.models`` module provides an API for saving machine learning models in
+"flavors" that can be understood by different downstream tools.
+
+The built-in flavors are:
+
+- :py:mod:`mlflow.catboost`
+- :py:mod:`mlflow.diviner`
+- :py:mod:`mlflow.dspy`
+- :py:mod:`mlflow.fastai`
+- :py:mod:`mlflow.h2o`
+- :py:mod:`mlflow.langchain`
+- :py:mod:`mlflow.lightgbm`
+- :py:mod:`mlflow.llama_index`
+- :py:mod:`mlflow.mleap`
+- :py:mod:`mlflow.onnx`
+- :py:mod:`mlflow.openai`
+- :py:mod:`mlflow.paddle`
+- :py:mod:`mlflow.pmdarima`
+- :py:mod:`mlflow.prophet`
+- :py:mod:`mlflow.pyfunc`
+- :py:mod:`mlflow.pyspark.ml`
+- :py:mod:`mlflow.pytorch`
+- :py:mod:`mlflow.sklearn`
+- :py:mod:`mlflow.spacy`
+- :py:mod:`mlflow.spark`
+- :py:mod:`mlflow.statsmodels`
+- :py:mod:`mlflow.tensorflow`
+- :py:mod:`mlflow.transformers`
+- :py:mod:`mlflow.xgboost`
+
+For details, see `MLflow Models <../models.html>`_.
+"""
+
+from mlflow.models.dependencies_schemas import set_retriever_schema
+from mlflow.models.evaluation import (
+    EvaluationArtifact,
+    EvaluationMetric,
+    EvaluationResult,
+    MetricThreshold,
+    evaluate,
+    list_evaluators,
+    make_metric,
+)
+from mlflow.models.flavor_backend import FlavorBackend
+from mlflow.models.model import Model, get_model_info, set_model, update_model_requirements
+from mlflow.models.model_config import ModelConfig
+from mlflow.models.python_api import build_docker
+from mlflow.models.resources import Resource, ResourceType
+from mlflow.utils.environment import infer_pip_requirements
+
+__all__ = [
+    "Model",
+    "FlavorBackend",
+    "infer_pip_requirements",
+    "evaluate",
+    "make_metric",
+    "EvaluationMetric",
+    "EvaluationArtifact",
+    "EvaluationResult",
+    "get_model_info",
+    "set_model",
+    "set_retriever_schema",
+    "list_evaluators",
+    "MetricThreshold",
+    "build_docker",
+    "Resource",
+    "ResourceType",
+    "ModelConfig",
+    "update_model_requirements",
+]
+
+
+# Under skinny-mlflow requirements, the following packages cannot be imported
+# because of lack of numpy/pandas library, so wrap them with try...except block
+try:
+    from mlflow.models.python_api import predict
+    from mlflow.models.signature import ModelSignature, infer_signature, set_signature
+    from mlflow.models.utils import (
+        ModelInputExample,
+        add_libraries_to_model,
+        convert_input_example_to_serving_input,
+        validate_schema,
+        validate_serving_input,
+    )
+
+    __all__ += [
+        "ModelSignature",
+        "ModelInputExample",
+        "infer_signature",
+        "validate_schema",
+        "add_libraries_to_model",
+        "convert_input_example_to_serving_input",
+        "set_signature",
+        "predict",
+        "validate_serving_input",
+    ]
+except ImportError:
+    pass

venv/lib/python3.9/site-packages/mlflow/models/auth_policy.py

@@ -0,0 +1,83 @@
+from typing import Optional
+
+from mlflow.models.resources import Resource, _ResourceBuilder
+from mlflow.utils.annotations import experimental
+
+
+@experimental
+class UserAuthPolicy:
+    """
+    A minimal list of scopes that the user should have access to
+    in order to invoke this model
+
+    Note: This is only compatible with Databricks Environment currently.
+    TODO: Add Databricks Documentation for User Auth Policy
+
+    Args:
+        api_scopes: A list of scopes. Example: "vectorsearch.vector-search-indexes", "sql"
+    """
+
+    def __init__(self, api_scopes: list[str]):
+        self._api_scopes = api_scopes
+
+    @property
+    def api_scopes(self) -> list[str]:
+        return self._api_scopes
+
+    @api_scopes.setter
+    def api_scopes(self, value: list[str]):
+        self._api_scopes = value
+
+    def to_dict(self):
+        return {"api_scopes": self.api_scopes}
+
+
+class SystemAuthPolicy:
+    """
+    System Auth Policy, which defines a list of resources required to
+    serve this model
+    """
+
+    def __init__(self, resources: list[Resource]):
+        self._resources = resources
+
+    @property
+    def resources(self) -> list[Resource]:
+        return self._resources
+
+    @resources.setter
+    def resources(self, value: list[Resource]):
+        self._resources = value
+
+    def to_dict(self):
+        serialized_resources = _ResourceBuilder.from_resources(self.resources)
+        return {"resources": serialized_resources}
+
+
+class AuthPolicy:
+    """
+    Specifies the authentication policy for the model, which includes two key
+    components.
+        System Auth Policy: A list of resources required to serve this model
+        User Auth Policy: A minimal list of scopes that the user should
+                          have access to, in order to invoke this model
+    """
+
+    def __init__(
+        self,
+        user_auth_policy: Optional[UserAuthPolicy] = None,
+        system_auth_policy: Optional[SystemAuthPolicy] = None,
+    ):
+        self.user_auth_policy = user_auth_policy
+        self.system_auth_policy = system_auth_policy
+
+    def to_dict(self):
+        """
+        Serialize Auth Policy to a dictionary
+        """
+        return {
+            "system_auth_policy": self.system_auth_policy.to_dict()
+            if self.system_auth_policy
+            else {},
+            "user_auth_policy": self.user_auth_policy.to_dict() if self.user_auth_policy else {},
+        }

venv/lib/python3.9/site-packages/mlflow/models/cli.py

@@ -0,0 +1,352 @@
+import logging
+
+import click
+
+from mlflow.models import python_api
+from mlflow.models.flavor_backend_registry import get_flavor_backend
+from mlflow.models.model import update_model_requirements
+from mlflow.utils import cli_args
+from mlflow.utils import env_manager as _EnvManager
+
+_logger = logging.getLogger(__name__)
+
+
+@click.group("models")
+def commands():
+    """
+    Deploy MLflow models locally.
+
+    To deploy a model associated with a run on a tracking server, set the MLFLOW_TRACKING_URI
+    environment variable to the URL of the desired server.
+    """
+
+
+@commands.command("serve")
+@cli_args.MODEL_URI
+@cli_args.PORT
+@cli_args.HOST
+@cli_args.TIMEOUT
+@cli_args.MODELS_WORKERS
+@cli_args.ENV_MANAGER
+@cli_args.NO_CONDA
+@cli_args.INSTALL_MLFLOW
+@cli_args.ENABLE_MLSERVER
+def serve(
+    model_uri,
+    port,
+    host,
+    timeout,
+    workers,
+    env_manager=None,
+    no_conda=False,
+    install_mlflow=False,
+    enable_mlserver=False,
+):
+    """
+    Serve a model saved with MLflow by launching a webserver on the specified host and port.
+    The command supports models with the ``python_function`` or ``crate`` (R Function) flavor.
+    For information about the input data formats accepted by the webserver, see the following
+    documentation: https://www.mlflow.org/docs/latest/models.html#built-in-deployment-tools.
+
+    .. warning::
+
+        Models built using MLflow 1.x will require adjustments to the endpoint request payload
+        if executed in an environment that has MLflow 2.x installed. In 1.x, a request payload
+        was in the format: ``{'columns': [str], 'data': [[...]]}``. 2.x models require
+        payloads that are defined by the structural-defining keys of either ``dataframe_split``,
+        ``instances``, ``inputs`` or ``dataframe_records``. See the examples below for
+        demonstrations of the changes to the invocation API endpoint in 2.0.
+
+    .. note::
+
+        Requests made in pandas DataFrame structures can be made in either `split` or `records`
+        oriented formats.
+        See https://pandas.pydata.org/docs/reference/api/pandas.DataFrame.to_json.html for
+        detailed information on orientation formats for converting a pandas DataFrame to json.
+
+    Example:
+
+    .. code-block:: bash
+
+        $ mlflow models serve -m runs:/my-run-id/model-path &
+
+        # records orientation input format for serializing a pandas DataFrame
+        $ curl http://127.0.0.1:5000/invocations -H 'Content-Type: application/json' -d '{
+            "dataframe_records": [{"a":1, "b":2}, {"a":3, "b":4}, {"a":5, "b":6}]
+        }'
+
+        # split orientation input format for serializing a pandas DataFrame
+        $ curl http://127.0.0.1:5000/invocations -H 'Content-Type: application/json' -d '{
+            "dataframe_split": {"columns": ["a", "b"],
+                                "index": [0, 1, 2],
+                                "data": [[1, 2], [3, 4], [5, 6]]}
+        }'
+
+        # inputs format for List submission of array, tensor, or DataFrame data
+        $ curl http://127.0.0.1:5000/invocations -H 'Content-Type: application/json' -d '{
+            "inputs": [[1, 2], [3, 4], [5, 6]]
+        }'
+
+        # instances format for submission of Tensor data
+        curl http://127.0.0.1:5000/invocations -H 'Content-Type: application/json' -d '{
+            "instances": [
+                {"a": "t1", "b": [1, 2, 3]},
+                {"a": "t2", "b": [4, 5, 6]},
+                {"a": "t3", "b": [7, 8, 9]}
+            ]
+        }'
+
+    """
+    env_manager = _EnvManager.LOCAL if no_conda else env_manager
+
+    return get_flavor_backend(
+        model_uri, env_manager=env_manager, workers=workers, install_mlflow=install_mlflow
+    ).serve(
+        model_uri=model_uri, port=port, host=host, timeout=timeout, enable_mlserver=enable_mlserver
+    )
+
+
+class KeyValueType(click.ParamType):
+    name = "key=value"
+
+    def convert(self, value, param, ctx):
+        if "=" not in value:
+            self.fail(f"{value!r} is not a valid key value pair, expecting `key=value`", param, ctx)
+        return value.split("=", 1)
+
+
+@commands.command("predict")
+@cli_args.MODEL_URI
+@click.option(
+    "--input-path", "-i", default=None, help="CSV containing pandas DataFrame to predict against."
+)
+@click.option(
+    "--output-path",
+    "-o",
+    default=None,
+    help="File to output results to as json file. If not provided, output to stdout.",
+)
+@click.option(
+    "--content-type",
+    "-t",
+    default="json",
+    help="Content type of the input file. Can be one of {'json', 'csv'}.",
+)
+@cli_args.ENV_MANAGER
+@cli_args.INSTALL_MLFLOW
+@click.option(
+    "--pip-requirements-override",
+    "-r",
+    default=None,
+    help="Specify packages and versions to override the dependencies defined "
+    "in the model. Must be a comma-separated string like x==y,z==a.",
+)
+@click.option(
+    "--env",
+    default=None,
+    type=KeyValueType(),
+    multiple=True,
+    help="Extra environment variables to set when running the model. Must be "
+    "key value pairs, e.g. `--env key=value`.",
+)
+def predict(
+    model_uri,
+    input_data=None,
+    input_path=None,
+    content_type=python_api._CONTENT_TYPE_JSON,
+    output_path=None,
+    env_manager=_EnvManager.VIRTUALENV,
+    install_mlflow=False,
+    pip_requirements_override=None,
+    env=None,
+):
+    """
+    Generate predictions in json format using a saved MLflow model. For information about the input
+    data formats accepted by this function, see the following documentation:
+    https://www.mlflow.org/docs/latest/models.html#built-in-deployment-tools.
+    """
+    return python_api.predict(
+        model_uri=model_uri,
+        input_data=input_data,
+        input_path=input_path,
+        content_type=content_type,
+        output_path=output_path,
+        env_manager=env_manager,
+        install_mlflow=install_mlflow,
+        pip_requirements_override=pip_requirements_override,
+        extra_envs=dict(env),
+    )
+
+
+@commands.command("prepare-env")
+@cli_args.MODEL_URI
+@cli_args.ENV_MANAGER
+@cli_args.INSTALL_MLFLOW
+def prepare_env(
+    model_uri,
+    env_manager,
+    install_mlflow,
+):
+    """
+    Performs any preparation necessary to predict or serve the model, for example
+    downloading dependencies or initializing a conda environment. After preparation,
+    calling predict or serve should be fast.
+    """
+    return get_flavor_backend(
+        model_uri, env_manager=env_manager, install_mlflow=install_mlflow
+    ).prepare_env(model_uri=model_uri)
+
+
+@commands.command("generate-dockerfile")
+@cli_args.MODEL_URI_BUILD_DOCKER
+@click.option(
+    "--output-directory",
+    "-d",
+    default="mlflow-dockerfile",
+    help="Output directory where the generated Dockerfile is stored.",
+)
+@cli_args.ENV_MANAGER_DOCKERFILE
+@cli_args.MLFLOW_HOME
+@cli_args.INSTALL_JAVA
+@cli_args.INSTALL_MLFLOW
+@cli_args.ENABLE_MLSERVER
+def generate_dockerfile(
+    model_uri,
+    output_directory,
+    env_manager,
+    mlflow_home,
+    install_java,
+    install_mlflow,
+    enable_mlserver,
+):
+    """
+    Generates a directory with Dockerfile whose default entrypoint serves an MLflow model at port
+    8080 using the python_function flavor. The generated Dockerfile is written to the specified
+    output directory, along with the model (if specified). This Dockerfile defines an image that
+    is equivalent to the one produced by ``mlflow models build-docker``.
+    """
+    if model_uri:
+        _logger.info("Generating Dockerfile for model %s", model_uri)
+    else:
+        _logger.info("Generating Dockerfile")
+    backend = get_flavor_backend(model_uri, docker_build=True, env_manager=env_manager)
+    if backend.can_build_image():
+        backend.generate_dockerfile(
+            model_uri,
+            output_directory,
+            mlflow_home=mlflow_home,
+            install_java=install_java,
+            install_mlflow=install_mlflow,
+            enable_mlserver=enable_mlserver,
+        )
+        _logger.info("Generated Dockerfile in directory %s", output_directory)
+    else:
+        _logger.error(
+            "Cannot build docker image for selected backend",
+            extra={"backend": backend.__class__.__name__},
+        )
+        raise NotImplementedError("Cannot build docker image for selected backend")
+
+
+@commands.command("build-docker")
+@cli_args.MODEL_URI_BUILD_DOCKER
+@click.option("--name", "-n", default="mlflow-pyfunc-servable", help="Name to use for built image")
+@cli_args.ENV_MANAGER
+@cli_args.MLFLOW_HOME
+@cli_args.INSTALL_JAVA
+@cli_args.INSTALL_MLFLOW
+@cli_args.ENABLE_MLSERVER
+def build_docker(**kwargs):
+    """
+    Builds a Docker image whose default entrypoint serves an MLflow model at port 8080, using the
+    python_function flavor. The container serves the model referenced by ``--model-uri``, if
+    specified when ``build-docker`` is called. If ``--model-uri`` is not specified when build_docker
+    is called, an MLflow Model directory must be mounted as a volume into the /opt/ml/model
+    directory in the container.
+
+    Building a Docker image with ``--model-uri``:
+
+    .. code:: bash
+
+        # Build a Docker image named 'my-image-name' that serves the model from run 'some-run-uuid'
+        # at run-relative artifact path 'my-model'
+        mlflow models build-docker --model-uri "runs:/some-run-uuid/my-model" --name "my-image-name"
+        # Serve the model
+        docker run -p 5001:8080 "my-image-name"
+
+    Building a Docker image without ``--model-uri``:
+
+    .. code:: bash
+
+        # Build a generic Docker image named 'my-image-name'
+        mlflow models build-docker --name "my-image-name"
+        # Mount the model stored in '/local/path/to/artifacts/model' and serve it
+        docker run --rm -p 5001:8080 -v /local/path/to/artifacts/model:/opt/ml/model "my-image-name"
+
+    .. important::
+
+        Since MLflow 2.10.1, the Docker image built with ``--model-uri`` does **not install Java**
+        for improved performance, unless the model flavor is one of ``["johnsnowlabs", "h2o",
+        "mleap", "spark"]``. If you need to install Java for other flavors, e.g. custom Python model
+        that uses SparkML, please specify the ``--install-java`` flag to enforce Java installation.
+
+    .. warning::
+
+        The image built without ``--model-uri`` doesn't support serving models with RFunc / Java
+        MLeap model server.
+
+    NB: by default, the container will start nginx and gunicorn processes. If you don't need the
+    nginx process to be started (for instance if you deploy your container to Google Cloud Run),
+    you can disable it via the DISABLE_NGINX environment variable:
+
+    .. code:: bash
+
+        docker run -p 5001:8080 -e DISABLE_NGINX=true "my-image-name"
+
+    See https://www.mlflow.org/docs/latest/python_api/mlflow.pyfunc.html for more information on the
+    'python_function' flavor.
+    """
+    python_api.build_docker(**kwargs)
+
+
+@commands.command("update-pip-requirements")
+@cli_args.MODEL_URI
+@click.argument("operation", type=click.Choice(["add", "remove"]))
+@click.argument("requirement_strings", type=str, nargs=-1)
+def update_pip_requirements(model_uri, operation, requirement_strings):
+    """
+    Add or remove requirements from a model's conda.yaml and requirements.txt files.
+    If using a remote tracking server, please make sure to set the MLFLOW_TRACKING_URI
+    environment variable to the URL of the desired server.
+
+    REQUIREMENT_STRINGS is a list of pip requirements specifiers.
+    See below for examples.
+
+    Sample usage:
+
+    .. code::
+
+        # Add requirements using the model's "runs:/" URI
+
+        mlflow models update-pip-requirements -m runs:/<run_id>/<model_path> \\
+            add "pandas==1.0.0" "scikit-learn" "mlflow >= 2.8, != 2.9.0"
+
+        # Remove requirements from a local model
+
+        mlflow models update-pip-requirements -m /path/to/local/model \\
+            remove "torchvision" "pydantic"
+
+    Note that model registry URIs (i.e. URIs in the form ``models:/``) are not
+    supported, as artifacts in the model registry are intended to be read-only.
+    Editing requirements is read-only artifact repositories is also not supported.
+
+    If adding requirements, the function will overwrite any existing requirements
+    that overlap, or else append the new requirements to the existing list.
+
+    If removing requirements, the function will ignore any version specifiers,
+    and remove all the specified package names. Any requirements that are not
+    found in the existing files will be ignored.
+    """
+    update_model_requirements(model_uri, operation, requirement_strings)
+
+    _logger.info(f"Successfully updated the requirements for the model at {model_uri}!")

venv/lib/python3.9/site-packages/mlflow/models/container/__init__.py

@@ -0,0 +1,317 @@
+"""
+Initialize the environment and start model serving in a Docker container.
+
+To be executed only during the model deployment.
+
+"""
+
+import logging
+import multiprocessing
+import os
+import shutil
+import signal
+import sys
+from pathlib import Path
+from subprocess import Popen, check_call
+
+import mlflow
+import mlflow.version
+from mlflow import mleap, pyfunc
+from mlflow.environment_variables import MLFLOW_DEPLOYMENT_FLAVOR_NAME, MLFLOW_DISABLE_ENV_CREATION
+from mlflow.models import Model
+from mlflow.models.model import MLMODEL_FILE_NAME
+from mlflow.pyfunc import _extract_conda_env, mlserver, scoring_server
+from mlflow.store.artifact.models_artifact_repo import REGISTERED_MODEL_META_FILE_NAME
+from mlflow.utils import env_manager as em
+from mlflow.utils.environment import _PythonEnv
+from mlflow.utils.file_utils import read_yaml
+from mlflow.utils.virtualenv import _get_or_create_virtualenv
+from mlflow.version import VERSION as MLFLOW_VERSION
+
+MODEL_PATH = "/opt/ml/model"
+
+
+DEFAULT_SAGEMAKER_SERVER_PORT = 8080
+DEFAULT_INFERENCE_SERVER_PORT = 8000
+DEFAULT_NGINX_SERVER_PORT = 8080
+DEFAULT_MLSERVER_PORT = 8080
+
+SUPPORTED_FLAVORS = [pyfunc.FLAVOR_NAME, mleap.FLAVOR_NAME]
+
+DISABLE_NGINX = "DISABLE_NGINX"
+ENABLE_MLSERVER = "ENABLE_MLSERVER"
+
+SERVING_ENVIRONMENT = "SERVING_ENVIRONMENT"
+
+
+_logger = logging.getLogger(__name__)
+
+
+def _init(cmd, env_manager):  # noqa: D417
+    """
+    Initialize the container and execute command.
+
+    Args:
+        cmd: Command param passed by Sagemaker. Can be "serve" or "train" (unimplemented).
+    """
+    if cmd == "serve":
+        _serve(env_manager)
+    elif cmd == "train":
+        _train()
+    else:
+        raise Exception(f"Unrecognized command {cmd}, full args = {sys.argv}")
+
+
+def _serve(env_manager):
+    """
+    Serve the model.
+
+    Read the MLmodel config, initialize the Conda environment if needed and start python server.
+    """
+    model_config_path = os.path.join(MODEL_PATH, MLMODEL_FILE_NAME)
+    m = Model.load(model_config_path)
+
+    # Older versions of mlflow may not specify a deployment configuration
+    serving_flavor = MLFLOW_DEPLOYMENT_FLAVOR_NAME.get() or pyfunc.FLAVOR_NAME
+
+    if serving_flavor == mleap.FLAVOR_NAME:
+        _serve_mleap()
+    elif pyfunc.FLAVOR_NAME in m.flavors:
+        _serve_pyfunc(m, env_manager)
+    else:
+        raise Exception("This container only supports models with the MLeap or PyFunc flavors.")
+
+
+def _install_pyfunc_deps(
+    model_path=None, install_mlflow=False, enable_mlserver=False, env_manager=em.VIRTUALENV
+):
+    """
+    Creates a conda env for serving the model at the specified path and installs almost all serving
+    dependencies into the environment - MLflow is not installed as it's not available via conda.
+    """
+    activate_cmd = _install_model_dependencies_to_env(model_path, env_manager) if model_path else []
+
+    # NB: install gunicorn[gevent] from pip rather than from conda because gunicorn is already
+    # dependency of mlflow on pip and we expect mlflow to be part of the environment.
+    server_deps = ["gunicorn[gevent]"]
+    if enable_mlserver:
+        server_deps = [
+            "'mlserver>=1.2.0,!=1.3.1,<1.4.0'",
+            "'mlserver-mlflow>=1.2.0,!=1.3.1,<1.4.0'",
+        ]
+
+    install_server_deps = [f"pip install {' '.join(server_deps)}"]
+    if Popen(["bash", "-c", " && ".join(activate_cmd + install_server_deps)]).wait() != 0:
+        raise Exception("Failed to install serving dependencies into the model environment.")
+
+    # NB: If we don't use virtualenv or conda env, we don't need to install mlflow here as
+    # it's already installed in the container.
+    if len(activate_cmd):
+        if _container_includes_mlflow_source():
+            # If the MLflow source code is copied to the container,
+            # we always need to run `pip install /opt/mlflow` otherwise
+            # the MLflow dependencies are not installed.
+            install_mlflow_cmd = ["pip install /opt/mlflow/."]
+        elif install_mlflow:
+            install_mlflow_cmd = [f"pip install mlflow=={MLFLOW_VERSION}"]
+        else:
+            install_mlflow_cmd = []
+
+        if install_mlflow_cmd:
+            if Popen(["bash", "-c", " && ".join(activate_cmd + install_mlflow_cmd)]).wait() != 0:
+                raise Exception("Failed to install mlflow into the model environment.")
+    return activate_cmd
+
+
+def _install_model_dependencies_to_env(model_path, env_manager) -> list[str]:
+    """:
+    Installs model dependencies to the specified environment, which can be either a local
+    environment, a conda environment, or a virtualenv.
+
+    Returns:
+        Empty list if local environment, otherwise a list of bash commands to activate the
+        virtualenv or conda environment.
+    """
+    model_config_path = os.path.join(model_path, MLMODEL_FILE_NAME)
+    model = Model.load(model_config_path)
+
+    conf = model.flavors.get(pyfunc.FLAVOR_NAME, {})
+    if pyfunc.ENV not in conf:
+        return []
+    env_conf = conf[mlflow.pyfunc.ENV]
+
+    if env_manager == em.LOCAL:
+        # Install pip dependencies directly into the local environment
+        python_env_config_path = os.path.join(model_path, env_conf[em.VIRTUALENV])
+        python_env = _PythonEnv.from_yaml(python_env_config_path)
+        deps = " ".join(python_env.build_dependencies + python_env.dependencies)
+        deps = deps.replace("requirements.txt", os.path.join(model_path, "requirements.txt"))
+        if Popen(["bash", "-c", f"python -m pip install {deps}"]).wait() != 0:
+            raise Exception("Failed to install model dependencies.")
+        return []
+
+    _logger.info("creating and activating custom environment")
+
+    env = _extract_conda_env(env_conf)
+    env_path_dst = os.path.join("/opt/mlflow/", env)
+    env_path_dst_dir = os.path.dirname(env_path_dst)
+    if not os.path.exists(env_path_dst_dir):
+        os.makedirs(env_path_dst_dir)
+    shutil.copy2(os.path.join(MODEL_PATH, env), env_path_dst)
+
+    if env_manager == em.CONDA:
+        conda_create_model_env = f"conda env create -n custom_env -f {env_path_dst}"
+        if Popen(["bash", "-c", conda_create_model_env]).wait() != 0:
+            raise Exception("Failed to create model environment.")
+        activate_cmd = ["source /miniconda/bin/activate custom_env"]
+
+    elif env_manager == em.VIRTUALENV:
+        env_activate_cmd = _get_or_create_virtualenv(model_path, env_manager=env_manager)
+        path = env_activate_cmd.split(" ")[-1]
+        os.symlink(path, "/opt/activate")
+        activate_cmd = [env_activate_cmd]
+
+    return activate_cmd
+
+
+def _serve_pyfunc(model, env_manager):
+    # option to disable manually nginx. The default behavior is to enable nginx.
+    disable_nginx = os.getenv(DISABLE_NGINX, "false").lower() == "true"
+    enable_mlserver = os.getenv(ENABLE_MLSERVER, "false").lower() == "true"
+    disable_env_creation = MLFLOW_DISABLE_ENV_CREATION.get()
+
+    conf = model.flavors[pyfunc.FLAVOR_NAME]
+    bash_cmds = []
+    if pyfunc.ENV in conf:
+        # NB: MLFLOW_DISABLE_ENV_CREATION is False only for SageMaker deployment, where the model
+        # files are loaded into the container at runtime rather than build time. In this case,
+        # we need to create a virtual environment and install the model dependencies into it when
+        # starting the container.
+        if not disable_env_creation:
+            _install_pyfunc_deps(
+                MODEL_PATH,
+                install_mlflow=True,
+                enable_mlserver=enable_mlserver,
+                env_manager=env_manager,
+            )
+        if env_manager == em.CONDA:
+            bash_cmds.append("source /miniconda/bin/activate custom_env")
+        elif env_manager == em.VIRTUALENV:
+            bash_cmds.append("source /opt/activate")
+    procs = []
+
+    start_nginx = True
+    if disable_nginx or enable_mlserver:
+        start_nginx = False
+
+    if start_nginx:
+        nginx_conf = Path(mlflow.models.__file__).parent.joinpath(
+            "container", "scoring_server", "nginx.conf"
+        )
+
+        nginx = Popen(["nginx", "-c", nginx_conf]) if start_nginx else None
+
+        # link the log streams to stdout/err so they will be logged to the container logs.
+        # Default behavior is to do the redirection unless explicitly specified
+        # by environment variable.
+        check_call(["ln", "-sf", "/dev/stdout", "/var/log/nginx/access.log"])
+        check_call(["ln", "-sf", "/dev/stderr", "/var/log/nginx/error.log"])
+
+        procs.append(nginx)
+
+    cpu_count = multiprocessing.cpu_count()
+    inference_server_kwargs = {}
+    if enable_mlserver:
+        inference_server = mlserver
+        # Allows users to choose the number of workers using MLServer var env settings.
+        # Default to cpu count
+        nworkers = int(os.getenv("MLSERVER_INFER_WORKERS", cpu_count))
+        # Since MLServer will run without NGINX, expose the server in the `8080`
+        # port, which is the assumed "public" port.
+        port = DEFAULT_MLSERVER_PORT
+
+        model_meta = _read_registered_model_meta(MODEL_PATH)
+        model_dict = model.to_dict()
+        inference_server_kwargs = {
+            "model_name": model_meta.get("model_name"),
+            "model_version": model_meta.get(
+                "model_version", model_dict.get("run_id", model_dict.get("model_uuid"))
+            ),
+        }
+    else:
+        inference_server = scoring_server
+        nworkers = cpu_count
+        port = DEFAULT_INFERENCE_SERVER_PORT
+
+    cmd, cmd_env = inference_server.get_cmd(
+        model_uri=MODEL_PATH, nworkers=nworkers, port=port, **inference_server_kwargs
+    )
+
+    bash_cmds.append(cmd)
+    inference_server_process = Popen(["/bin/bash", "-c", " && ".join(bash_cmds)], env=cmd_env)
+    procs.append(inference_server_process)
+
+    signal.signal(signal.SIGTERM, lambda a, b: _sigterm_handler(pids=[p.pid for p in procs]))
+    # If either subprocess exits, so do we.
+    awaited_pids = _await_subprocess_exit_any(procs=procs)
+    _sigterm_handler(awaited_pids)
+
+
+def _read_registered_model_meta(model_path):
+    model_meta = {}
+    if os.path.isfile(os.path.join(model_path, REGISTERED_MODEL_META_FILE_NAME)):
+        model_meta = read_yaml(model_path, REGISTERED_MODEL_META_FILE_NAME)
+
+    return model_meta
+
+
+def _serve_mleap():
+    serve_cmd = [
+        "java",
+        "-cp",
+        '"/opt/java/jars/*"',
+        "org.mlflow.sagemaker.ScoringServer",
+        MODEL_PATH,
+        str(DEFAULT_SAGEMAKER_SERVER_PORT),
+    ]
+    # Invoke `Popen` with a single string command in the shell to support wildcard usage
+    # with the mlflow jar version.
+    serve_cmd = " ".join(serve_cmd)
+    mleap = Popen(serve_cmd, shell=True)
+    signal.signal(signal.SIGTERM, lambda a, b: _sigterm_handler(pids=[mleap.pid]))
+    awaited_pids = _await_subprocess_exit_any(procs=[mleap])
+    _sigterm_handler(awaited_pids)
+
+
+def _container_includes_mlflow_source():
+    return os.path.exists("/opt/mlflow/pyproject.toml")
+
+
+def _train():
+    raise Exception("Train is not implemented.")
+
+
+def _await_subprocess_exit_any(procs):
+    pids = [proc.pid for proc in procs]
+    while True:
+        pid, _ = os.wait()
+        if pid in pids:
+            break
+    return pids
+
+
+def _sigterm_handler(pids):
+    """
+    Cleanup when terminating.
+
+    Attempt to kill all launched processes and exit.
+
+    """
+    _logger.info("Got sigterm signal, exiting.")
+    for pid in pids:
+        try:
+            os.kill(pid, signal.SIGTERM)
+        except OSError:
+            pass
+
+    sys.exit(0)

venv/lib/python3.9/site-packages/mlflow/models/container/scoring_server/nginx.conf

@@ -0,0 +1,39 @@
+worker_processes 1;
+daemon off; # Prevent forking
+
+
+pid /tmp/nginx.pid;
+error_log /var/log/nginx/error.log;
+
+events {
+  # defaults
+}
+
+http {
+  include /etc/nginx/mime.types;
+  default_type application/octet-stream;
+  access_log /var/log/nginx/access.log combined;
+  
+  upstream uvicorn {
+    server 127.0.0.1:8000;
+  }
+
+  server {
+    listen 8080 deferred;
+    client_max_body_size 5m;
+
+    keepalive_timeout 75;
+
+    location ~ ^/(ping|invocations) {
+      proxy_set_header X-Forwarded-For $proxy_add_x_forwarded_for;
+      proxy_set_header Host $http_host;
+      proxy_redirect off;
+      proxy_pass http://uvicorn;
+      client_max_body_size 100m;
+    }
+
+    location / {
+      return 404 "{}";
+    }
+  }
+}

venv/lib/python3.9/site-packages/mlflow/models/dependencies_schemas.py

@@ -0,0 +1,287 @@
+import json
+import logging
+from abc import ABC, abstractmethod
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from enum import Enum
+from typing import TYPE_CHECKING, Optional
+
+from mlflow.utils.annotations import experimental
+
+if TYPE_CHECKING:
+    from mlflow.models.model import Model
+
+_logger = logging.getLogger(__name__)
+
+
+class DependenciesSchemasType(Enum):
+    """
+    Enum to define the different types of dependencies schemas for the model.
+    """
+
+    RETRIEVERS = "retrievers"
+
+
+@experimental
+def set_retriever_schema(
+    *,
+    primary_key: str,
+    text_column: str,
+    doc_uri: Optional[str] = None,
+    other_columns: Optional[list[str]] = None,
+    name: Optional[str] = "retriever",
+):
+    """
+    Specify the return schema of a retriever span within your agent or generative AI app code.
+
+    **Note**: MLflow recommends that your retriever return the default MLflow retriever output
+    schema described in https://mlflow.org/docs/latest/tracing/tracing-schema#retriever-spans,
+    in which case you do not need to call `set_retriever_schema`. APIs that read MLflow traces
+    and look for retriever spans, such as MLflow evaluation, will automatically detect retriever
+    spans that match MLflow's default retriever schema.
+
+    If your retriever does not return the default MLflow retriever output schema, call this API to
+    specify which fields in each retrieved document correspond to the page content, document
+    URI, document ID, etc. This enables downstream features like MLflow evaluation to properly
+    identify these fields. Note that `set_retriever_schema` assumes that your retriever span
+    returns a list of objects.
+
+
+    Args:
+        primary_key: The primary key of the retriever or vector index.
+        text_column: The name of the text column to use for the embeddings.
+        doc_uri: The name of the column that contains the document URI.
+        other_columns: A list of other columns that are part of the vector index
+                          that need to be retrieved during trace logging.
+        name: The name of the retriever tool or vector store index.
+
+    .. code-block:: Python
+            :caption: Example
+
+            from mlflow.models import set_retriever_schema
+
+            # The following call sets the schema for a custom retriever that retrieves content from
+            # MLflow documentation, with an output schema like:
+            # [
+            #     {
+            #         'document_id': '9a8292da3a9d4005a988bf0bfdd0024c',
+            #         'chunk_text': 'MLflow is an open-source platform, purpose-built to assist...',
+            #         'doc_uri': 'https://mlflow.org/docs/latest/index.html',
+            #         'title': 'MLflow: A Tool for Managing the Machine Learning Lifecycle'
+            #     },
+            #     {
+            #         'document_id': '7537fe93c97f4fdb9867412e9c1f9e5b',
+            #         'chunk_text': 'A great way to get started with MLflow is...',
+            #         'doc_uri': 'https://mlflow.org/docs/latest/getting-started/',
+            #         'title': 'Getting Started with MLflow'
+            #     },
+            # ...
+            # ]
+            set_retriever_schema(
+                primary_key="chunk_id",
+                text_column="chunk_text",
+                doc_uri="doc_uri",
+                other_columns=["title"],
+                name="my_custom_retriever",
+            )
+    """
+    retriever_schemas = globals().get(DependenciesSchemasType.RETRIEVERS.value, [])
+
+    # Check if a retriever schema with the same name already exists
+    existing_schema = next((schema for schema in retriever_schemas if schema["name"] == name), None)
+
+    if existing_schema is not None:
+        # Compare all relevant fields
+        if (
+            existing_schema["primary_key"] == primary_key
+            and existing_schema["text_column"] == text_column
+            and existing_schema["doc_uri"] == doc_uri
+            and existing_schema["other_columns"] == (other_columns or [])
+        ):
+            # No difference, no need to warn or update
+            return
+        else:
+            # Differences found, issue a warning
+            _logger.warning(
+                f"A retriever schema with the name '{name}' already exists. "
+                "Overriding the existing schema."
+            )
+            # Override the fields of the existing schema
+            existing_schema["primary_key"] = primary_key
+            existing_schema["text_column"] = text_column
+            existing_schema["doc_uri"] = doc_uri
+            existing_schema["other_columns"] = other_columns or []
+    else:
+        retriever_schemas.append(
+            {
+                "primary_key": primary_key,
+                "text_column": text_column,
+                "doc_uri": doc_uri,
+                "other_columns": other_columns or [],
+                "name": name,
+            }
+        )
+
+    globals()[DependenciesSchemasType.RETRIEVERS.value] = retriever_schemas
+
+
+def _get_retriever_schema():
+    """
+    Get the vector search schema defined by the user.
+
+    Returns:
+        VectorSearchIndex: The vector search index schema.
+    """
+    retriever_schemas = globals().get(DependenciesSchemasType.RETRIEVERS.value, [])
+    if not retriever_schemas:
+        return []
+
+    return [
+        RetrieverSchema(
+            name=retriever.get("name"),
+            primary_key=retriever.get("primary_key"),
+            text_column=retriever.get("text_column"),
+            doc_uri=retriever.get("doc_uri"),
+            other_columns=retriever.get("other_columns"),
+        )
+        for retriever in retriever_schemas
+    ]
+
+
+def _clear_retriever_schema():
+    """
+    Clear the vector search schema defined by the user.
+    """
+    globals().pop(DependenciesSchemasType.RETRIEVERS.value, None)
+
+
+def _clear_dependencies_schemas():
+    """
+    Clear all the dependencies schema defined by the user.
+    """
+    # Clear the vector search schema
+    _clear_retriever_schema()
+
+
+@contextmanager
+def _get_dependencies_schemas():
+    dependencies_schemas = DependenciesSchemas(retriever_schemas=_get_retriever_schema())
+    try:
+        yield dependencies_schemas
+    finally:
+        _clear_dependencies_schemas()
+
+
+def _get_dependencies_schema_from_model(model: "Model") -> Optional[dict]:
+    """
+    Get the dependencies schema from the logged model metadata.
+
+    `dependencies_schemas` is a dictionary that defines the dependencies schemas, such as
+    the retriever schemas. This code is now only useful for Databricks integration.
+    """
+    if model.metadata and "dependencies_schemas" in model.metadata:
+        dependencies_schemas = model.metadata["dependencies_schemas"]
+        return {
+            "dependencies_schemas": {
+                dependency: json.dumps(schema)
+                for dependency, schema in dependencies_schemas.items()
+            }
+        }
+    return None
+
+
+@dataclass
+class Schema(ABC):
+    """
+    Base class for defining the resources needed to serve a model.
+
+    Args:
+        type (ResourceType): The type of the schema.
+    """
+
+    type: DependenciesSchemasType
+
+    @abstractmethod
+    def to_dict(self):
+        """
+        Convert the resource to a dictionary.
+        Subclasses must implement this method.
+        """
+
+    @classmethod
+    @abstractmethod
+    def from_dict(cls, data: dict[str, str]):
+        """
+        Convert the dictionary to a Resource.
+        Subclasses must implement this method.
+        """
+
+
+@dataclass
+class RetrieverSchema(Schema):
+    """
+    Define vector search index resource to serve a model.
+
+    Args:
+        name (str): The name of the vector search index schema.
+        primary_key (str): The primary key for the index.
+        text_column (str): The main text column for the index.
+        doc_uri (Optional[str]): The document URI for the index.
+        other_columns (Optional[List[str]]): Additional columns in the index.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        primary_key: str,
+        text_column: str,
+        doc_uri: Optional[str] = None,
+        other_columns: Optional[list[str]] = None,
+    ):
+        super().__init__(type=DependenciesSchemasType.RETRIEVERS)
+        self.name = name
+        self.primary_key = primary_key
+        self.text_column = text_column
+        self.doc_uri = doc_uri
+        self.other_columns = other_columns or []
+
+    def to_dict(self):
+        return {
+            self.type.value: [
+                {
+                    "name": self.name,
+                    "primary_key": self.primary_key,
+                    "text_column": self.text_column,
+                    "doc_uri": self.doc_uri,
+                    "other_columns": self.other_columns,
+                }
+            ]
+        }
+
+    @classmethod
+    def from_dict(cls, data: dict[str, str]):
+        return cls(
+            name=data["name"],
+            primary_key=data["primary_key"],
+            text_column=data["text_column"],
+            doc_uri=data.get("doc_uri"),
+            other_columns=data.get("other_columns", []),
+        )
+
+
+@dataclass
+class DependenciesSchemas:
+    retriever_schemas: list[RetrieverSchema] = field(default_factory=list)
+
+    def to_dict(self) -> dict[str, dict[DependenciesSchemasType, list[dict]]]:
+        if not self.retriever_schemas:
+            return None
+
+        return {
+            "dependencies_schemas": {
+                DependenciesSchemasType.RETRIEVERS.value: [
+                    index.to_dict()[DependenciesSchemasType.RETRIEVERS.value][0]
+                    for index in self.retriever_schemas
+                ],
+            }
+        }

venv/lib/python3.9/site-packages/mlflow/models/display_utils.py

@@ -0,0 +1,158 @@
+import html
+from pathlib import Path
+
+from mlflow.models.model import ModelInfo
+from mlflow.models.signature import ModelSignature
+from mlflow.types import schema
+from mlflow.utils import databricks_utils
+
+
+def _is_input_string(inputs: schema.Schema) -> bool:
+    return (
+        not inputs.has_input_names()
+        and len(inputs.input_types()) == 1
+        and inputs.input_types()[0] == schema.DataType.string
+    )
+
+
+def _is_input_agent_compatible(inputs: schema.Schema) -> bool:
+    if _is_input_string(inputs):
+        return True
+    if not inputs.has_input_names():
+        return False
+    messages = inputs.input_dict().get("messages")
+    if not messages:
+        return False
+    if not isinstance(messages.type, schema.Array):
+        return False
+    items = messages.type.dtype
+    if not isinstance(items, schema.Object):
+        return False
+    properties = items.properties
+    content = next(filter(lambda prop: prop.name == "content", properties), None)
+    role = next(filter(lambda prop: prop.name == "role", properties), None)
+    return (
+        content
+        and content.dtype == schema.DataType.string
+        and role
+        and role.dtype == schema.DataType.string
+    )
+
+
+def _is_output_string_response(outputs: schema.Schema) -> bool:
+    if not outputs.has_input_names():
+        return False
+    content = outputs.input_dict().get("content")
+    if not content:
+        return False
+    return content.type == schema.DataType.string
+
+
+def _is_output_string(outputs: schema.Schema) -> bool:
+    return (
+        not outputs.has_input_names()
+        and len(outputs.input_types()) == 1
+        and outputs.input_types()[0] == schema.DataType.string
+    )
+
+
+def _is_output_chat_completion_response(outputs: schema.Schema) -> bool:
+    if not outputs.has_input_names():
+        return False
+    choices = outputs.input_dict().get("choices")
+    if not choices:
+        return False
+    if not isinstance(choices.type, schema.Array):
+        return False
+    items = choices.type.dtype
+    if not isinstance(items, schema.Object):
+        return False
+    properties = items.properties
+    message = next(filter(lambda prop: prop.name == "message", properties), None)
+    if not message:
+        return False
+    if not isinstance(message.dtype, schema.Object):
+        return False
+    message_properties = message.dtype.properties
+    content = next(filter(lambda prop: prop.name == "content", message_properties), None)
+    role = next(filter(lambda prop: prop.name == "role", message_properties), None)
+    return (
+        content
+        and content.dtype == schema.DataType.string
+        and role
+        and role.dtype == schema.DataType.string
+    )
+
+
+def _is_output_agent_compatible(outputs: schema.Schema) -> bool:
+    return (
+        _is_output_string_response(outputs)
+        or _is_output_string(outputs)
+        or _is_output_chat_completion_response(outputs)
+    )
+
+
+def _is_signature_agent_compatible(signature: ModelSignature) -> bool:
+    """Determines whether the given signature is compatible with the agent eval schema.
+
+    See https://docs.databricks.com/en/generative-ai/agent-evaluation/evaluation-schema.html.
+    The schema accepts the OpenAI spec, as well as simpler formats such as vanilla string response
+    and `StringResponse`.
+    """
+    return _is_input_agent_compatible(signature.inputs) and _is_output_agent_compatible(
+        signature.outputs
+    )
+
+
+def _should_render_agent_eval_template(signature: ModelSignature) -> bool:
+    if not databricks_utils.is_in_databricks_runtime():
+        return False
+    from IPython import get_ipython
+
+    if get_ipython() is None:
+        return False
+    return _is_signature_agent_compatible(signature)
+
+
+def _generate_agent_eval_recipe(model_uri: str) -> str:
+    resources_dir = Path(__file__).parent / "notebook_resources"
+    pip_install_command = """%pip install -U databricks-agents
+dbutils.library.restartPython()
+## Run the above in a separate cell ##"""
+    eval_with_synthetic_code = (
+        (resources_dir / "eval_with_synthetic_example.py")
+        .read_text()
+        .replace("{{pipInstall}}", pip_install_command)
+        .replace("{{modelUri}}", model_uri)
+    )
+    eval_with_dataset_code = (
+        (resources_dir / "eval_with_dataset_example.py")
+        .read_text()
+        .replace("{{pipInstall}}", pip_install_command)
+        .replace("{{modelUri}}", model_uri)
+    )
+
+    # Remove the ruff noqa comments.
+    ruff_line = "# ruff: noqa: F821, I001\n"
+    eval_with_synthetic_code = eval_with_synthetic_code.replace(ruff_line, "")
+    eval_with_dataset_code = eval_with_dataset_code.replace(ruff_line, "")
+
+    return (
+        (resources_dir / "agent_evaluation_template.html")
+        .read_text()
+        .replace("{{eval_with_synthetic_code}}", html.escape(eval_with_synthetic_code))
+        .replace("{{eval_with_dataset_code}}", html.escape(eval_with_dataset_code))
+    )
+
+
+def maybe_render_agent_eval_recipe(model_info: ModelInfo) -> None:
+    # For safety, we wrap in try/catch to make sure we don't break `mlflow.*.log_model`.
+    try:
+        if not _should_render_agent_eval_template(model_info.signature):
+            return
+
+        from IPython.display import HTML, display
+
+        display(HTML(_generate_agent_eval_recipe(model_info.model_uri)))
+    except Exception:
+        pass

venv/lib/python3.9/site-packages/mlflow/models/docker_utils.py

@@ -0,0 +1,238 @@
+import logging
+import os
+from subprocess import Popen
+from typing import Optional, Union
+from urllib.parse import urlparse
+
+from mlflow.environment_variables import MLFLOW_DOCKER_OPENJDK_VERSION
+from mlflow.utils import env_manager as em
+from mlflow.utils.file_utils import _copy_project
+from mlflow.version import VERSION
+
+_logger = logging.getLogger(__name__)
+
+UBUNTU_BASE_IMAGE = "ubuntu:20.04"
+PYTHON_SLIM_BASE_IMAGE = "python:{version}-slim"
+
+
+SETUP_PYENV_AND_VIRTUALENV = r"""# Setup pyenv
+RUN DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get -y install tzdata \
+    libssl-dev zlib1g-dev libbz2-dev libreadline-dev libsqlite3-dev wget curl llvm \
+    libncursesw5-dev xz-utils tk-dev libxml2-dev libxmlsec1-dev libffi-dev liblzma-dev
+RUN git clone \
+    --depth 1 \
+    --branch $(git ls-remote --tags --sort=v:refname https://github.com/pyenv/pyenv.git | grep -o -E 'v[1-9]+(\.[1-9]+)+$' | tail -1) \
+    https://github.com/pyenv/pyenv.git /root/.pyenv
+ENV PYENV_ROOT="/root/.pyenv"
+ENV PATH="$PYENV_ROOT/bin:$PATH"
+RUN apt install -y python3.9 python3.9-distutils \
+    && ln -s -f $(which python3.9) /usr/bin/python \
+    && wget https://bootstrap.pypa.io/get-pip.py -O /tmp/get-pip.py \
+    && python /tmp/get-pip.py
+RUN pip install virtualenv
+"""  # noqa: E501
+
+_DOCKERFILE_TEMPLATE = """# Build an image that can serve mlflow models.
+FROM {base_image}
+
+{setup_python_venv}
+
+{setup_java}
+
+WORKDIR /opt/mlflow
+
+{install_mlflow}
+
+{install_model_and_deps}
+
+ENV MLFLOW_DISABLE_ENV_CREATION={disable_env_creation}
+ENV ENABLE_MLSERVER={enable_mlserver}
+
+# granting read/write access and conditional execution authority to all child directories
+# and files to allow for deployment to AWS Sagemaker Serverless Endpoints
+# (see https://docs.aws.amazon.com/sagemaker/latest/dg/serverless-endpoints.html)
+RUN chmod o+rwX /opt/mlflow/
+
+# clean up apt cache to reduce image size
+RUN rm -rf /var/lib/apt/lists/*
+
+ENTRYPOINT ["python", "-c", "{entrypoint}"]
+"""
+
+
+SETUP_MINICONDA = """# Setup miniconda
+RUN curl --fail -L https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh > miniconda.sh
+RUN bash ./miniconda.sh -b -p /miniconda && rm ./miniconda.sh
+ENV PATH="/miniconda/bin:$PATH"
+"""  # noqa: E501
+
+
+def generate_dockerfile(
+    output_dir: str,
+    base_image: str,
+    model_install_steps: Optional[str],
+    entrypoint: str,
+    env_manager: Union[em.CONDA, em.LOCAL, em.VIRTUALENV],
+    mlflow_home: Optional[str] = None,
+    enable_mlserver: bool = False,
+    disable_env_creation_at_runtime: bool = True,
+    install_java: Optional[bool] = None,
+):
+    """
+    Generates a Dockerfile that can be used to build a docker image, that serves ML model
+    stored and tracked in MLflow.
+    """
+
+    setup_java_steps = ""
+    setup_python_venv_steps = ""
+    install_mlflow_steps = _pip_mlflow_install_step(output_dir, mlflow_home)
+
+    if base_image.startswith("python:"):
+        if install_java:
+            _logger.warning(
+                "`install_java` option is not supported when using python base image, "
+                "switch to UBUNTU_BASE_IMAGE to enable java installation."
+            )
+        setup_python_venv_steps = (
+            "RUN apt-get -y update && apt-get install -y --no-install-recommends nginx"
+        )
+
+    elif base_image == UBUNTU_BASE_IMAGE:
+        setup_python_venv_steps = (
+            "RUN apt-get -y update && DEBIAN_FRONTEND=noninteractive TZ=Etc/UTC apt-get install -y "
+            "--no-install-recommends wget curl nginx ca-certificates bzip2 build-essential cmake "
+            "git-core\n\n"
+        )
+        setup_python_venv_steps += (
+            SETUP_MINICONDA if env_manager == em.CONDA else SETUP_PYENV_AND_VIRTUALENV
+        )
+        if install_java is not False:
+            jdk_ver = MLFLOW_DOCKER_OPENJDK_VERSION.get()
+            setup_java_steps = (
+                "# Setup Java\n"
+                f"RUN apt-get install -y --no-install-recommends openjdk-{jdk_ver}-jdk maven\n"
+                f"ENV JAVA_HOME=/usr/lib/jvm/java-{jdk_ver}-openjdk-amd64"
+            )
+
+            install_mlflow_steps += "\n\n" + _java_mlflow_install_step(mlflow_home)
+
+    with open(os.path.join(output_dir, "Dockerfile"), "w") as f:
+        f.write(
+            _DOCKERFILE_TEMPLATE.format(
+                base_image=base_image,
+                setup_python_venv=setup_python_venv_steps,
+                setup_java=setup_java_steps,
+                install_mlflow=install_mlflow_steps,
+                install_model_and_deps=model_install_steps,
+                entrypoint=entrypoint,
+                enable_mlserver=enable_mlserver,
+                disable_env_creation=disable_env_creation_at_runtime,
+            )
+        )
+
+
+def _java_mlflow_install_step(mlflow_home):
+    maven_proxy = _get_maven_proxy()
+    if mlflow_home:
+        return (
+            "# Install Java mlflow-scoring from local source\n"
+            "RUN cd /opt/mlflow/mlflow/java/scoring && "
+            f"mvn --batch-mode package -DskipTests {maven_proxy} && "
+            "mkdir -p /opt/java/jars && "
+            "mv /opt/mlflow/mlflow/java/scoring/target/"
+            "mlflow-scoring-*-with-dependencies.jar /opt/java/jars\n"
+        )
+    else:
+        return (
+            "# Install Java mlflow-scoring from Maven Central\n"
+            "RUN mvn"
+            " --batch-mode dependency:copy"
+            f" -Dartifact=org.mlflow:mlflow-scoring:{VERSION}:pom"
+            f" -DoutputDirectory=/opt/java {maven_proxy}\n"
+            "RUN mvn"
+            " --batch-mode dependency:copy"
+            f" -Dartifact=org.mlflow:mlflow-scoring:{VERSION}:jar"
+            f" -DoutputDirectory=/opt/java/jars {maven_proxy}\n"
+            f"RUN cp /opt/java/mlflow-scoring-{VERSION}.pom /opt/java/pom.xml\n"
+            "RUN cd /opt/java && mvn "
+            "--batch-mode dependency:copy-dependencies "
+            f"-DoutputDirectory=/opt/java/jars {maven_proxy}\n"
+        )
+
+
+def _get_maven_proxy():
+    http_proxy = os.getenv("http_proxy")
+    https_proxy = os.getenv("https_proxy")
+    if not http_proxy or not https_proxy:
+        return ""
+
+    # Expects proxies as either PROTOCOL://{USER}:{PASSWORD}@HOSTNAME:PORT
+    # or PROTOCOL://HOSTNAME:PORT
+    parsed_http_proxy = urlparse(http_proxy)
+    assert parsed_http_proxy.hostname is not None, "Invalid `http_proxy` hostname."
+    assert parsed_http_proxy.port is not None, f"Invalid proxy port: {parsed_http_proxy.port}"
+
+    parsed_https_proxy = urlparse(https_proxy)
+    assert parsed_https_proxy.hostname is not None, "Invalid `https_proxy` hostname."
+    assert parsed_https_proxy.port is not None, f"Invalid proxy port: {parsed_https_proxy.port}"
+
+    maven_proxy_options = (
+        "-DproxySet=true",
+        f"-Dhttp.proxyHost={parsed_http_proxy.hostname}",
+        f"-Dhttp.proxyPort={parsed_http_proxy.port}",
+        f"-Dhttps.proxyHost={parsed_https_proxy.hostname}",
+        f"-Dhttps.proxyPort={parsed_https_proxy.port}",
+        "-Dhttps.nonProxyHosts=repo.maven.apache.org",
+    )
+
+    if parsed_http_proxy.username is None or parsed_http_proxy.password is None:
+        return " ".join(maven_proxy_options)
+
+    return " ".join(
+        (
+            *maven_proxy_options,
+            f"-Dhttp.proxyUser={parsed_http_proxy.username}",
+            f"-Dhttp.proxyPassword={parsed_http_proxy.password}",
+        )
+    )
+
+
+def _pip_mlflow_install_step(dockerfile_context_dir, mlflow_home):
+    """
+    Get docker build commands for installing MLflow given a Docker context dir and optional source
+    directory
+    """
+    if mlflow_home:
+        mlflow_dir = _copy_project(
+            src_path=os.path.abspath(mlflow_home), dst_path=dockerfile_context_dir
+        )
+        return (
+            "# Install MLflow from local source\n"
+            f"COPY {mlflow_dir} /opt/mlflow\n"
+            "RUN pip install /opt/mlflow"
+        )
+    else:
+        return f"# Install MLflow\nRUN pip install mlflow=={VERSION}"
+
+
+def build_image_from_context(context_dir: str, image_name: str):
+    import docker
+
+    client = docker.from_env()
+    # In Docker < 19, `docker build` doesn't support the `--platform` option
+    is_platform_supported = int(client.version()["Version"].split(".")[0]) >= 19
+    # Enforcing the AMD64 architecture build for Apple M1 users
+    platform_option = ["--platform", "linux/amd64"] if is_platform_supported else []
+    commands = [
+        "docker",
+        "build",
+        "-t",
+        image_name,
+        "-f",
+        "Dockerfile",
+        *platform_option,
+        ".",
+    ]
+    proc = Popen(commands, cwd=context_dir)
+    if proc.wait():
+        raise RuntimeError("Docker build failed.")

venv/lib/python3.9/site-packages/mlflow/models/evaluation/__init__.py

@@ -0,0 +1,23 @@
+from mlflow.data.evaluation_dataset import EvaluationDataset
+from mlflow.models.evaluation.base import (
+    EvaluationArtifact,
+    EvaluationMetric,
+    EvaluationResult,
+    ModelEvaluator,
+    evaluate,
+    list_evaluators,
+    make_metric,
+)
+from mlflow.models.evaluation.validation import MetricThreshold
+
+__all__ = [
+    "ModelEvaluator",
+    "EvaluationDataset",
+    "EvaluationResult",
+    "EvaluationMetric",
+    "EvaluationArtifact",
+    "make_metric",
+    "evaluate",
+    "list_evaluators",
+    "MetricThreshold",
+]

venv/lib/python3.9/site-packages/mlflow/models/evaluation/_shap_patch.py

@@ -0,0 +1,64 @@
+import pickle
+
+import numpy as np
+import shap
+from shap._serializable import Deserializer, Serializable, Serializer
+
+
+class _PatchedKernelExplainer(shap.KernelExplainer):
+    @staticmethod
+    def not_equal(i, j):
+        # `shap.KernelExplainer.not_equal` method fails on some special types such as
+        # timestamp, this breaks the kernel explainer routine.
+        # `PatchedKernelExplainer` fixes this issue.
+        # See https://github.com/slundberg/shap/pull/2586
+        number_types = (int, float, np.number)
+        if isinstance(i, number_types) and isinstance(j, number_types):
+            return 0 if np.isclose(i, j, equal_nan=True) else 1
+        else:
+            return 0 if i == j else 1
+
+    def save(self, out_file, model_saver=None, masker_saver=None):
+        """
+        This patched `save` method fix `KernelExplainer.save`.
+        Issues in original `KernelExplainer.save`:
+         - It saves model by calling model.save, but shap.utils._legacy.Model has no save method
+         - It tries to save "masker", but there's no "masker" in KernelExplainer
+         - It does not save "KernelExplainer.data" attribute, the attribute is required when
+           loading back
+        Note: `model_saver` and `masker_saver` are meaningless argument for `KernelExplainer.save`,
+        the model in "KernelExplainer" is an instance of `shap.utils._legacy.Model`
+        (it wraps the predict function), we can only use pickle to dump it.
+        and no `masker` for KernelExplainer so `masker_saver` is meaningless.
+        but I preserve the 2 argument for overridden API compatibility.
+        """
+        pickle.dump(type(self), out_file)
+        with Serializer(out_file, "shap.Explainer", version=0) as s:
+            s.save("model", self.model)
+            s.save("link", self.link)
+            s.save("data", self.data)
+
+    @classmethod
+    def load(cls, in_file, model_loader=None, masker_loader=None, instantiate=True):
+        """
+        This patched `load` method fix `KernelExplainer.load`.
+        Issues in original KernelExplainer.load:
+         - Use mismatched model loader to load model
+         - Try to load non-existent "masker" attribute
+         - Does not load "data" attribute and then cause calling " KernelExplainer"
+           constructor lack of "data" argument.
+        Note: `model_loader` and `masker_loader` are meaningless argument for
+        `KernelExplainer.save`, because the `model` object is saved by pickle dump,
+        we must use pickle load to load it.
+        and no `masker` for KernelExplainer so `masker_loader` is meaningless.
+        but I preserve the 2 argument for overridden API compatibility.
+        """
+        if instantiate:
+            return cls._instantiated_load(in_file, model_loader=None, masker_loader=None)
+
+        kwargs = Serializable.load(in_file, instantiate=False)
+        with Deserializer(in_file, "shap.Explainer", min_version=0, max_version=0) as s:
+            kwargs["model"] = s.load("model")
+            kwargs["link"] = s.load("link")
+            kwargs["data"] = s.load("data")
+        return kwargs

venv/lib/python3.9/site-packages/mlflow/models/evaluation/artifacts.py

@@ -0,0 +1,194 @@
+import json
+import pathlib
+import pickle
+from collections import namedtuple
+from json import JSONDecodeError
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+
+from mlflow.exceptions import MlflowException
+from mlflow.models.evaluation.base import EvaluationArtifact
+from mlflow.utils.annotations import developer_stable
+from mlflow.utils.proto_json_utils import NumpyEncoder
+
+
+@developer_stable
+class ImageEvaluationArtifact(EvaluationArtifact):
+    def _save(self, output_artifact_path):
+        self._content.save(output_artifact_path)
+
+    def _load_content_from_file(self, local_artifact_path):
+        from PIL.Image import open as open_image
+
+        self._content = open_image(local_artifact_path)
+        self._content.load()  # Load image and close the file descriptor.
+        return self._content
+
+
+@developer_stable
+class CsvEvaluationArtifact(EvaluationArtifact):
+    def _save(self, output_artifact_path):
+        self._content.to_csv(output_artifact_path, index=False)
+
+    def _load_content_from_file(self, local_artifact_path):
+        self._content = pd.read_csv(local_artifact_path)
+        return self._content
+
+
+@developer_stable
+class ParquetEvaluationArtifact(EvaluationArtifact):
+    def _save(self, output_artifact_path):
+        self._content.to_parquet(output_artifact_path, compression="brotli")
+
+    def _load_content_from_file(self, local_artifact_path):
+        self._content = pd.read_parquet(local_artifact_path)
+        return self._content
+
+
+@developer_stable
+class NumpyEvaluationArtifact(EvaluationArtifact):
+    def _save(self, output_artifact_path):
+        np.save(output_artifact_path, self._content, allow_pickle=False)
+
+    def _load_content_from_file(self, local_artifact_path):
+        self._content = np.load(local_artifact_path, allow_pickle=False)
+        return self._content
+
+
+@developer_stable
+class JsonEvaluationArtifact(EvaluationArtifact):
+    def _save(self, output_artifact_path):
+        with open(output_artifact_path, "w") as f:
+            json.dump(self._content, f)
+
+    def _load_content_from_file(self, local_artifact_path):
+        with open(local_artifact_path) as f:
+            self._content = json.load(f)
+        return self._content
+
+
+@developer_stable
+class TextEvaluationArtifact(EvaluationArtifact):
+    def _save(self, output_artifact_path):
+        with open(output_artifact_path, "w") as f:
+            f.write(self._content)
+
+    def _load_content_from_file(self, local_artifact_path):
+        with open(local_artifact_path) as f:
+            self._content = f.read()
+        return self._content
+
+
+@developer_stable
+class PickleEvaluationArtifact(EvaluationArtifact):
+    def _save(self, output_artifact_path):
+        with open(output_artifact_path, "wb") as f:
+            pickle.dump(self._content, f)
+
+    def _load_content_from_file(self, local_artifact_path):
+        with open(local_artifact_path, "rb") as f:
+            self._content = pickle.load(f)
+        return self._content
+
+
+_EXT_TO_ARTIFACT_MAP = {
+    ".png": ImageEvaluationArtifact,
+    ".jpg": ImageEvaluationArtifact,
+    ".jpeg": ImageEvaluationArtifact,
+    ".json": JsonEvaluationArtifact,
+    ".npy": NumpyEvaluationArtifact,
+    ".csv": CsvEvaluationArtifact,
+    ".parquet": ParquetEvaluationArtifact,
+    ".txt": TextEvaluationArtifact,
+}
+
+_TYPE_TO_EXT_MAP = {
+    pd.DataFrame: ".csv",
+    np.ndarray: ".npy",
+    plt.Figure: ".png",
+}
+
+_TYPE_TO_ARTIFACT_MAP = {
+    pd.DataFrame: CsvEvaluationArtifact,
+    np.ndarray: NumpyEvaluationArtifact,
+    plt.Figure: ImageEvaluationArtifact,
+}
+
+_InferredArtifactProperties = namedtuple(
+    "_InferredArtifactProperties", ["from_path", "type", "ext"]
+)
+
+
+def _infer_artifact_type_and_ext(artifact_name, raw_artifact, custom_metric_tuple):
+    """
+    This function performs type and file extension inference on the provided artifact
+
+    Args:
+        artifact_name: The name of the provided artifact
+        raw_artifact: The artifact object
+        custom_metric_tuple: Containing a user provided function and its index in the
+            ``custom_metrics`` parameter of ``mlflow.evaluate``
+
+    Returns:
+        InferredArtifactProperties namedtuple
+    """
+
+    exception_header = (
+        f"Custom metric function '{custom_metric_tuple.name}' at index "
+        f"{custom_metric_tuple.index} in the `custom_metrics` parameter produced an "
+        f"artifact '{artifact_name}'"
+    )
+
+    # Given a string, first see if it is a path. Otherwise, check if it is a JsonEvaluationArtifact
+    if isinstance(raw_artifact, str):
+        potential_path = pathlib.Path(raw_artifact)
+        if potential_path.exists():
+            raw_artifact = potential_path
+        else:
+            try:
+                json.loads(raw_artifact)
+                return _InferredArtifactProperties(
+                    from_path=False, type=JsonEvaluationArtifact, ext=".json"
+                )
+            except JSONDecodeError:
+                raise MlflowException(
+                    f"{exception_header} with string representation '{raw_artifact}' that is "
+                    f"neither a valid path to a file nor a JSON string."
+                )
+
+    # Type inference based on the file extension
+    if isinstance(raw_artifact, pathlib.Path):
+        if not raw_artifact.exists():
+            raise MlflowException(f"{exception_header} with path '{raw_artifact}' does not exist.")
+        if not raw_artifact.is_file():
+            raise MlflowException(f"{exception_header} with path '{raw_artifact}' is not a file.")
+        if raw_artifact.suffix not in _EXT_TO_ARTIFACT_MAP:
+            raise MlflowException(
+                f"{exception_header} with path '{raw_artifact}' does not match any of the supported"
+                f" file extensions: {', '.join(_EXT_TO_ARTIFACT_MAP.keys())}."
+            )
+        return _InferredArtifactProperties(
+            from_path=True, type=_EXT_TO_ARTIFACT_MAP[raw_artifact.suffix], ext=raw_artifact.suffix
+        )
+
+    # Type inference based on object type
+    if type(raw_artifact) in _TYPE_TO_ARTIFACT_MAP:
+        return _InferredArtifactProperties(
+            from_path=False,
+            type=_TYPE_TO_ARTIFACT_MAP[type(raw_artifact)],
+            ext=_TYPE_TO_EXT_MAP[type(raw_artifact)],
+        )
+
+    # Given as other python object, we first attempt to infer as JsonEvaluationArtifact. If that
+    # fails, we store it as PickleEvaluationArtifact
+    try:
+        json.dumps(raw_artifact, cls=NumpyEncoder)
+        return _InferredArtifactProperties(
+            from_path=False, type=JsonEvaluationArtifact, ext=".json"
+        )
+    except TypeError:
+        return _InferredArtifactProperties(
+            from_path=False, type=PickleEvaluationArtifact, ext=".pickle"
+        )

venv/lib/python3.9/site-packages/mlflow/models/evaluation/evaluator_registry.py

@@ -0,0 +1,80 @@
+import warnings
+
+from mlflow.exceptions import MlflowException
+from mlflow.utils.import_hooks import register_post_import_hook
+from mlflow.utils.plugins import get_entry_points
+
+
+class ModelEvaluatorRegistry:
+    """
+    Scheme-based registry for model evaluator implementations
+    """
+
+    def __init__(self):
+        self._registry = {}
+        self._builtin_evaluators = {}
+
+    def register(self, scheme, evaluator):
+        """Register model evaluator provided by other packages"""
+        self._registry[scheme] = evaluator
+
+    def register_builtin(self, scheme, evaluator):
+        """Register built-in model evaluator"""
+        self._registry[scheme] = evaluator
+        self._builtin_evaluators[scheme] = evaluator
+
+    def register_entrypoints(self):
+        # Register ModelEvaluator implementation provided by other packages
+        for entrypoint in get_entry_points("mlflow.model_evaluator"):
+            try:
+                self.register(entrypoint.name, entrypoint.load())
+            except (AttributeError, ImportError) as exc:
+                warnings.warn(
+                    'Failure attempting to register model evaluator for scheme "{}": {}'.format(
+                        entrypoint.name, str(exc)
+                    ),
+                    stacklevel=2,
+                )
+
+    def get_evaluator(self, evaluator_name):
+        """
+        Get an evaluator instance from the registry based on the name of evaluator
+        """
+        evaluator_cls = self._registry.get(evaluator_name)
+        if evaluator_cls is None:
+            raise MlflowException(
+                f"Could not find a registered model evaluator for: {evaluator_name}. "
+                f"Currently registered evaluator names are: {list(self._registry.keys())}"
+            )
+        return evaluator_cls()
+
+    def is_builtin(self, name):
+        return name in self._builtin_evaluators
+
+    def is_registered(self, name):
+        return name in self._registry
+
+
+_model_evaluation_registry = ModelEvaluatorRegistry()
+
+
+def register_evaluators(module):
+    from mlflow.models.evaluation.evaluators.classifier import ClassifierEvaluator
+    from mlflow.models.evaluation.evaluators.default import DefaultEvaluator
+    from mlflow.models.evaluation.evaluators.regressor import RegressorEvaluator
+    from mlflow.models.evaluation.evaluators.shap import ShapEvaluator
+
+    # Built-in evaluators
+    module._model_evaluation_registry.register_builtin(DefaultEvaluator.name, DefaultEvaluator)
+    module._model_evaluation_registry.register_builtin(
+        ClassifierEvaluator.name, ClassifierEvaluator
+    )
+    module._model_evaluation_registry.register_builtin(RegressorEvaluator.name, RegressorEvaluator)
+    module._model_evaluation_registry.register_builtin(ShapEvaluator.name, ShapEvaluator)
+
+    # Plugin evaluators
+    module._model_evaluation_registry.register_entrypoints()
+
+
+# Put it in post-importing hook to avoid circuit importing
+register_post_import_hook(register_evaluators, __name__, overwrite=True)

venv/lib/python3.9/site-packages/mlflow/models/evaluation/evaluators/classifier.py

@@ -0,0 +1,680 @@
+import logging
+import math
+from collections import namedtuple
+from contextlib import contextmanager
+from typing import Optional
+
+import numpy as np
+import pandas as pd
+from sklearn import metrics as sk_metrics
+
+import mlflow
+from mlflow import MlflowException
+from mlflow.environment_variables import _MLFLOW_EVALUATE_SUPPRESS_CLASSIFICATION_ERRORS
+from mlflow.models.evaluation.artifacts import CsvEvaluationArtifact
+from mlflow.models.evaluation.base import EvaluationMetric, EvaluationResult, _ModelType
+from mlflow.models.evaluation.default_evaluator import (
+    BuiltInEvaluator,
+    _extract_raw_model,
+    _get_aggregate_metrics_values,
+)
+from mlflow.models.utils import plot_lines
+
+_logger = logging.getLogger(__name__)
+
+
+_Curve = namedtuple("_Curve", ["plot_fn", "plot_fn_args", "auc"])
+
+
+class ClassifierEvaluator(BuiltInEvaluator):
+    """
+    A built-in evaluator for classifier models.
+    """
+
+    name = "classifier"
+
+    @classmethod
+    def can_evaluate(cls, *, model_type, evaluator_config, **kwargs):
+        # TODO: Also the model needs to be pyfunc model, not function or endpoint URI
+        return model_type == _ModelType.CLASSIFIER
+
+    def _evaluate(
+        self,
+        model: Optional["mlflow.pyfunc.PyFuncModel"],
+        extra_metrics: list[EvaluationMetric],
+        custom_artifacts=None,
+        **kwargs,
+    ) -> Optional[EvaluationResult]:
+        # Get classification config
+        self.y_true = self.dataset.labels_data
+        self.label_list = self.evaluator_config.get("label_list")
+        self.pos_label = self.evaluator_config.get("pos_label")
+        self.sample_weights = self.evaluator_config.get("sample_weights")
+        if self.pos_label and self.label_list and self.pos_label not in self.label_list:
+            raise MlflowException.invalid_parameter_value(
+                f"'pos_label' {self.pos_label} must exist in 'label_list' {self.label_list}."
+            )
+
+        # Check if the model_type is consistent with ground truth labels
+        inferred_model_type = _infer_model_type_by_labels(self.y_true)
+        if _ModelType.CLASSIFIER != inferred_model_type:
+            _logger.warning(
+                f"According to the evaluation dataset label values, the model type looks like "
+                f"{inferred_model_type}, but you specified model type 'classifier'. Please "
+                f"verify that you set the `model_type` and `dataset` arguments correctly."
+            )
+
+        # Run model prediction
+        input_df = self.X.copy_to_avoid_mutation()
+        self.y_pred, self.y_probs = self._generate_model_predictions(model, input_df)
+
+        self._validate_label_list()
+
+        self._compute_builtin_metrics(model)
+        self.evaluate_metrics(extra_metrics, prediction=self.y_pred, target=self.y_true)
+        self.evaluate_and_log_custom_artifacts(
+            custom_artifacts, prediction=self.y_pred, target=self.y_true
+        )
+
+        # Log metrics and artifacts
+        self.log_metrics()
+        self.log_eval_table(self.y_pred)
+
+        if len(self.label_list) == 2:
+            self._log_binary_classifier_artifacts()
+        else:
+            self._log_multiclass_classifier_artifacts()
+        self._log_confusion_matrix()
+
+        return EvaluationResult(
+            metrics=self.aggregate_metrics, artifacts=self.artifacts, run_id=self.run_id
+        )
+
+    def _generate_model_predictions(self, model, input_df):
+        predict_fn, predict_proba_fn = _extract_predict_fn_and_prodict_proba_fn(model)
+        # Classifier model is guaranteed to output single column of predictions
+        y_pred = self.dataset.predictions_data if model is None else predict_fn(input_df)
+
+        # Predict class probabilities if the model supports it
+        y_probs = predict_proba_fn(input_df) if predict_proba_fn is not None else None
+        return y_pred, y_probs
+
+    def _validate_label_list(self):
+        if self.label_list is None:
+            # If label list is not specified, infer label list from model output
+            self.label_list = np.unique(np.concatenate([self.y_true, self.y_pred]))
+        else:
+            # np.where only works for numpy array, not list
+            self.label_list = np.array(self.label_list)
+
+        # sort label_list ASC, for binary classification it makes sure the last one is pos label
+        self.label_list.sort()
+
+        is_binomial = len(self.label_list) <= 2
+        if is_binomial:
+            if self.pos_label is None:
+                self.pos_label = self.label_list[-1]
+            else:
+                if self.pos_label in self.label_list:
+                    self.label_list = np.delete(
+                        self.label_list, np.where(self.label_list == self.pos_label)
+                    )
+                self.label_list = np.append(self.label_list, self.pos_label)
+            if len(self.label_list) < 2:
+                raise MlflowException(
+                    "Evaluation dataset for classification must contain at least two unique "
+                    f"labels, but only {len(self.label_list)} unique labels were found.",
+                )
+            with _suppress_class_imbalance_errors(IndexError, log_warning=False):
+                _logger.info(
+                    "The evaluation dataset is inferred as binary dataset, positive label is "
+                    f"{self.label_list[1]}, negative label is {self.label_list[0]}."
+                )
+        else:
+            _logger.info(
+                "The evaluation dataset is inferred as multiclass dataset, number of classes "
+                f"is inferred as {len(self.label_list)}. If this is incorrect, please specify the "
+                "`label_list` parameter in `evaluator_config`."
+            )
+
+    def _compute_builtin_metrics(self, model):
+        self._evaluate_sklearn_model_score_if_scorable(model, self.y_true, self.sample_weights)
+
+        if len(self.label_list) <= 2:
+            metrics = _get_binary_classifier_metrics(
+                y_true=self.y_true,
+                y_pred=self.y_pred,
+                y_proba=self.y_probs,
+                labels=self.label_list,
+                pos_label=self.pos_label,
+                sample_weights=self.sample_weights,
+            )
+            if metrics:
+                self.metrics_values.update(_get_aggregate_metrics_values(metrics))
+                self._compute_roc_and_pr_curve()
+        else:
+            average = self.evaluator_config.get("average", "weighted")
+            metrics = _get_multiclass_classifier_metrics(
+                y_true=self.y_true,
+                y_pred=self.y_pred,
+                y_proba=self.y_probs,
+                labels=self.label_list,
+                average=average,
+                sample_weights=self.sample_weights,
+            )
+            if metrics:
+                self.metrics_values.update(_get_aggregate_metrics_values(metrics))
+
+    def _compute_roc_and_pr_curve(self):
+        if self.y_probs is not None:
+            with _suppress_class_imbalance_errors(ValueError, log_warning=False):
+                self.roc_curve = _gen_classifier_curve(
+                    is_binomial=True,
+                    y=self.y_true,
+                    y_probs=self.y_probs[:, 1],
+                    labels=self.label_list,
+                    pos_label=self.pos_label,
+                    curve_type="roc",
+                    sample_weights=self.sample_weights,
+                )
+
+                self.metrics_values.update(
+                    _get_aggregate_metrics_values({"roc_auc": self.roc_curve.auc})
+                )
+            with _suppress_class_imbalance_errors(ValueError, log_warning=False):
+                self.pr_curve = _gen_classifier_curve(
+                    is_binomial=True,
+                    y=self.y_true,
+                    y_probs=self.y_probs[:, 1],
+                    labels=self.label_list,
+                    pos_label=self.pos_label,
+                    curve_type="pr",
+                    sample_weights=self.sample_weights,
+                )
+
+                self.metrics_values.update(
+                    _get_aggregate_metrics_values({"precision_recall_auc": self.pr_curve.auc})
+                )
+
+    def _log_pandas_df_artifact(self, pandas_df, artifact_name):
+        artifact_file_name = f"{artifact_name}.csv"
+        artifact_file_local_path = self.temp_dir.path(artifact_file_name)
+        pandas_df.to_csv(artifact_file_local_path, index=False)
+        mlflow.log_artifact(artifact_file_local_path)
+        artifact = CsvEvaluationArtifact(
+            uri=mlflow.get_artifact_uri(artifact_file_name),
+            content=pandas_df,
+        )
+        artifact._load(artifact_file_local_path)
+        self.artifacts[artifact_name] = artifact
+
+    def _log_multiclass_classifier_artifacts(self):
+        per_class_metrics_collection_df = _get_classifier_per_class_metrics_collection_df(
+            y=self.y_true,
+            y_pred=self.y_pred,
+            labels=self.label_list,
+            sample_weights=self.sample_weights,
+        )
+
+        log_roc_pr_curve = False
+        if self.y_probs is not None:
+            max_classes_for_multiclass_roc_pr = self.evaluator_config.get(
+                "max_classes_for_multiclass_roc_pr", 10
+            )
+            if len(self.label_list) <= max_classes_for_multiclass_roc_pr:
+                log_roc_pr_curve = True
+            else:
+                _logger.warning(
+                    f"The classifier num_classes > {max_classes_for_multiclass_roc_pr}, skip "
+                    f"logging ROC curve and Precision-Recall curve. You can add evaluator config "
+                    f"'max_classes_for_multiclass_roc_pr' to increase the threshold."
+                )
+
+        if log_roc_pr_curve:
+            roc_curve = _gen_classifier_curve(
+                is_binomial=False,
+                y=self.y_true,
+                y_probs=self.y_probs,
+                labels=self.label_list,
+                pos_label=self.pos_label,
+                curve_type="roc",
+                sample_weights=self.sample_weights,
+            )
+
+            def plot_roc_curve():
+                roc_curve.plot_fn(**roc_curve.plot_fn_args)
+
+            self._log_image_artifact(plot_roc_curve, "roc_curve_plot")
+            per_class_metrics_collection_df["roc_auc"] = roc_curve.auc
+
+            pr_curve = _gen_classifier_curve(
+                is_binomial=False,
+                y=self.y_true,
+                y_probs=self.y_probs,
+                labels=self.label_list,
+                pos_label=self.pos_label,
+                curve_type="pr",
+                sample_weights=self.sample_weights,
+            )
+
+            def plot_pr_curve():
+                pr_curve.plot_fn(**pr_curve.plot_fn_args)
+
+            self._log_image_artifact(plot_pr_curve, "precision_recall_curve_plot")
+            per_class_metrics_collection_df["precision_recall_auc"] = pr_curve.auc
+
+        self._log_pandas_df_artifact(per_class_metrics_collection_df, "per_class_metrics")
+
+    def _log_roc_curve(self):
+        def _plot_roc_curve():
+            self.roc_curve.plot_fn(**self.roc_curve.plot_fn_args)
+
+        self._log_image_artifact(_plot_roc_curve, "roc_curve_plot")
+
+    def _log_precision_recall_curve(self):
+        def _plot_pr_curve():
+            self.pr_curve.plot_fn(**self.pr_curve.plot_fn_args)
+
+        self._log_image_artifact(_plot_pr_curve, "precision_recall_curve_plot")
+
+    def _log_lift_curve(self):
+        from mlflow.models.evaluation.lift_curve import plot_lift_curve
+
+        def _plot_lift_curve():
+            return plot_lift_curve(self.y_true, self.y_probs, pos_label=self.pos_label)
+
+        self._log_image_artifact(_plot_lift_curve, "lift_curve_plot")
+
+    def _log_binary_classifier_artifacts(self):
+        if self.y_probs is not None:
+            with _suppress_class_imbalance_errors(log_warning=False):
+                self._log_roc_curve()
+            with _suppress_class_imbalance_errors(log_warning=False):
+                self._log_precision_recall_curve()
+            with _suppress_class_imbalance_errors(ValueError, log_warning=False):
+                self._log_lift_curve()
+
+    def _log_confusion_matrix(self):
+        """
+        Helper method for logging confusion matrix
+        """
+        # normalize the confusion matrix, keep consistent with sklearn autologging.
+        confusion_matrix = sk_metrics.confusion_matrix(
+            self.y_true,
+            self.y_pred,
+            labels=self.label_list,
+            normalize="true",
+            sample_weight=self.sample_weights,
+        )
+
+        def plot_confusion_matrix():
+            import matplotlib
+            import matplotlib.pyplot as plt
+
+            with matplotlib.rc_context(
+                {
+                    "font.size": min(8, math.ceil(50.0 / len(self.label_list))),
+                    "axes.labelsize": 8,
+                }
+            ):
+                _, ax = plt.subplots(1, 1, figsize=(6.0, 4.0), dpi=175)
+                disp = sk_metrics.ConfusionMatrixDisplay(
+                    confusion_matrix=confusion_matrix,
+                    display_labels=self.label_list,
+                ).plot(cmap="Blues", ax=ax)
+                disp.ax_.set_title("Normalized confusion matrix")
+
+        if hasattr(sk_metrics, "ConfusionMatrixDisplay"):
+            self._log_image_artifact(
+                plot_confusion_matrix,
+                "confusion_matrix",
+            )
+        return
+
+
+def _is_categorical(values):
+    """
+    Infer whether input values are categorical on best effort.
+    Return True represent they are categorical, return False represent we cannot determine result.
+    """
+    dtype_name = pd.Series(values).convert_dtypes().dtype.name.lower()
+    return dtype_name in ["category", "string", "boolean"]
+
+
+def _is_continuous(values):
+    """
+    Infer whether input values is continuous on best effort.
+    Return True represent they are continuous, return False represent we cannot determine result.
+    """
+    dtype_name = pd.Series(values).convert_dtypes().dtype.name.lower()
+    return dtype_name.startswith("float")
+
+
+def _infer_model_type_by_labels(labels):
+    """
+    Infer model type by target values.
+    """
+    if _is_categorical(labels):
+        return _ModelType.CLASSIFIER
+    elif _is_continuous(labels):
+        return _ModelType.REGRESSOR
+    else:
+        return None  # Unknown
+
+
+def _extract_predict_fn_and_prodict_proba_fn(model):
+    predict_fn = None
+    predict_proba_fn = None
+
+    _, raw_model = _extract_raw_model(model)
+
+    if raw_model is not None:
+        predict_fn = raw_model.predict
+        predict_proba_fn = getattr(raw_model, "predict_proba", None)
+        try:
+            from mlflow.xgboost import (
+                _wrapped_xgboost_model_predict_fn,
+                _wrapped_xgboost_model_predict_proba_fn,
+            )
+
+            # Because shap evaluation will pass evaluation data in ndarray format
+            # (without feature names), if set validate_features=True it will raise error.
+            predict_fn = _wrapped_xgboost_model_predict_fn(raw_model, validate_features=False)
+            predict_proba_fn = _wrapped_xgboost_model_predict_proba_fn(
+                raw_model, validate_features=False
+            )
+        except ImportError:
+            pass
+    elif model is not None:
+        predict_fn = model.predict
+
+    return predict_fn, predict_proba_fn
+
+
+@contextmanager
+def _suppress_class_imbalance_errors(exception_type=Exception, log_warning=True):
+    """
+    Exception handler context manager to suppress Exceptions if the private environment
+    variable `_MLFLOW_EVALUATE_SUPPRESS_CLASSIFICATION_ERRORS` is set to `True`.
+    The purpose of this handler is to prevent an evaluation call for a binary or multiclass
+    classification automl run from aborting due to an extreme minority class imbalance
+    encountered during iterative training cycles due to the non deterministic sampling
+    behavior of Spark's DataFrame.sample() API.
+    The Exceptions caught in the usage of this are broad and are designed purely to not
+    interrupt the iterative hyperparameter tuning process. Final evaluations are done
+    in a more deterministic (but expensive) fashion.
+    """
+    try:
+        yield
+    except exception_type as e:
+        if _MLFLOW_EVALUATE_SUPPRESS_CLASSIFICATION_ERRORS.get():
+            if log_warning:
+                _logger.warning(
+                    "Failed to calculate metrics due to class imbalance. "
+                    "This is expected when the dataset is imbalanced."
+                )
+        else:
+            raise e
+
+
+def _get_binary_sum_up_label_pred_prob(positive_class_index, positive_class, y, y_pred, y_probs):
+    y = np.array(y)
+    y_bin = np.where(y == positive_class, 1, 0)
+    y_pred_bin = None
+    y_prob_bin = None
+    if y_pred is not None:
+        y_pred = np.array(y_pred)
+        y_pred_bin = np.where(y_pred == positive_class, 1, 0)
+
+    if y_probs is not None:
+        y_probs = np.array(y_probs)
+        y_prob_bin = y_probs[:, positive_class_index]
+
+    return y_bin, y_pred_bin, y_prob_bin
+
+
+def _get_common_classifier_metrics(
+    *, y_true, y_pred, y_proba, labels, average, pos_label, sample_weights
+):
+    metrics = {
+        "example_count": len(y_true),
+        "accuracy_score": sk_metrics.accuracy_score(y_true, y_pred, sample_weight=sample_weights),
+        "recall_score": sk_metrics.recall_score(
+            y_true,
+            y_pred,
+            average=average,
+            pos_label=pos_label,
+            sample_weight=sample_weights,
+        ),
+        "precision_score": sk_metrics.precision_score(
+            y_true,
+            y_pred,
+            average=average,
+            pos_label=pos_label,
+            sample_weight=sample_weights,
+        ),
+        "f1_score": sk_metrics.f1_score(
+            y_true,
+            y_pred,
+            average=average,
+            pos_label=pos_label,
+            sample_weight=sample_weights,
+        ),
+    }
+
+    if y_proba is not None:
+        with _suppress_class_imbalance_errors(ValueError):
+            metrics["log_loss"] = sk_metrics.log_loss(
+                y_true, y_proba, labels=labels, sample_weight=sample_weights
+            )
+    return metrics
+
+
+def _get_binary_classifier_metrics(
+    *, y_true, y_pred, y_proba=None, labels=None, pos_label=1, sample_weights=None
+):
+    with _suppress_class_imbalance_errors(ValueError):
+        tn, fp, fn, tp = sk_metrics.confusion_matrix(y_true, y_pred).ravel()
+        return {
+            "true_negatives": tn,
+            "false_positives": fp,
+            "false_negatives": fn,
+            "true_positives": tp,
+            **_get_common_classifier_metrics(
+                y_true=y_true,
+                y_pred=y_pred,
+                y_proba=y_proba,
+                labels=labels,
+                average="binary",
+                pos_label=pos_label,
+                sample_weights=sample_weights,
+            ),
+        }
+
+
+def _get_multiclass_classifier_metrics(
+    *,
+    y_true,
+    y_pred,
+    y_proba=None,
+    labels=None,
+    average="weighted",
+    sample_weights=None,
+):
+    metrics = _get_common_classifier_metrics(
+        y_true=y_true,
+        y_pred=y_pred,
+        y_proba=y_proba,
+        labels=labels,
+        average=average,
+        pos_label=None,
+        sample_weights=sample_weights,
+    )
+    if average in ("macro", "weighted") and y_proba is not None:
+        metrics.update(
+            roc_auc=sk_metrics.roc_auc_score(
+                y_true=y_true,
+                y_score=y_proba,
+                sample_weight=sample_weights,
+                average=average,
+                multi_class="ovr",
+            )
+        )
+    return metrics
+
+
+def _get_classifier_per_class_metrics_collection_df(y, y_pred, labels, sample_weights):
+    per_class_metrics_list = []
+    for positive_class_index, positive_class in enumerate(labels):
+        (
+            y_bin,
+            y_pred_bin,
+            _,
+        ) = _get_binary_sum_up_label_pred_prob(
+            positive_class_index, positive_class, y, y_pred, None
+        )
+        per_class_metrics = {"positive_class": positive_class}
+        binary_classifier_metrics = _get_binary_classifier_metrics(
+            y_true=y_bin,
+            y_pred=y_pred_bin,
+            pos_label=1,
+            sample_weights=sample_weights,
+        )
+        if binary_classifier_metrics:
+            per_class_metrics.update(binary_classifier_metrics)
+        per_class_metrics_list.append(per_class_metrics)
+
+    return pd.DataFrame(per_class_metrics_list)
+
+
+_Curve = namedtuple("_Curve", ["plot_fn", "plot_fn_args", "auc"])
+
+
+def _gen_classifier_curve(
+    is_binomial,
+    y,
+    y_probs,
+    labels,
+    pos_label,
+    curve_type,
+    sample_weights,
+):
+    """
+    Generate precision-recall curve or ROC curve for classifier.
+
+    Args:
+        is_binomial: True if it is binary classifier otherwise False
+        y: True label values
+        y_probs: if binary classifier, the predicted probability for positive class.
+                  if multiclass classifier, the predicted probabilities for all classes.
+        labels: The set of labels.
+        pos_label: The label of the positive class.
+        curve_type: "pr" or "roc"
+        sample_weights: Optional sample weights.
+
+    Returns:
+        An instance of "_Curve" which includes attributes "plot_fn", "plot_fn_args", "auc".
+    """
+    if curve_type == "roc":
+
+        def gen_line_x_y_label_auc(_y, _y_prob, _pos_label):
+            fpr, tpr, _ = sk_metrics.roc_curve(
+                _y,
+                _y_prob,
+                sample_weight=sample_weights,
+                # For multiclass classification where a one-vs-rest ROC curve is produced for each
+                # class, the positive label is binarized and should not be included in the plot
+                # legend
+                pos_label=_pos_label if _pos_label == pos_label else None,
+            )
+
+            auc = sk_metrics.roc_auc_score(y_true=_y, y_score=_y_prob, sample_weight=sample_weights)
+            return fpr, tpr, f"AUC={auc:.3f}", auc
+
+        xlabel = "False Positive Rate"
+        ylabel = "True Positive Rate"
+        title = "ROC curve"
+        if pos_label:
+            xlabel = f"False Positive Rate (Positive label: {pos_label})"
+            ylabel = f"True Positive Rate (Positive label: {pos_label})"
+    elif curve_type == "pr":
+
+        def gen_line_x_y_label_auc(_y, _y_prob, _pos_label):
+            precision, recall, _ = sk_metrics.precision_recall_curve(
+                _y,
+                _y_prob,
+                sample_weight=sample_weights,
+                # For multiclass classification where a one-vs-rest precision-recall curve is
+                # produced for each class, the positive label is binarized and should not be
+                # included in the plot legend
+                pos_label=_pos_label if _pos_label == pos_label else None,
+            )
+            # NB: We return average precision score (AP) instead of AUC because AP is more
+            # appropriate for summarizing a precision-recall curve
+            ap = sk_metrics.average_precision_score(
+                y_true=_y, y_score=_y_prob, pos_label=_pos_label, sample_weight=sample_weights
+            )
+            return recall, precision, f"AP={ap:.3f}", ap
+
+        xlabel = "Recall"
+        ylabel = "Precision"
+        title = "Precision recall curve"
+        if pos_label:
+            xlabel = f"Recall (Positive label: {pos_label})"
+            ylabel = f"Precision (Positive label: {pos_label})"
+    else:
+        assert False, "illegal curve type"
+
+    if is_binomial:
+        x_data, y_data, line_label, auc = gen_line_x_y_label_auc(y, y_probs, pos_label)
+        data_series = [(line_label, x_data, y_data)]
+    else:
+        curve_list = []
+        for positive_class_index, positive_class in enumerate(labels):
+            y_bin, _, y_prob_bin = _get_binary_sum_up_label_pred_prob(
+                positive_class_index, positive_class, y, labels, y_probs
+            )
+
+            x_data, y_data, line_label, auc = gen_line_x_y_label_auc(
+                y_bin, y_prob_bin, _pos_label=1
+            )
+            curve_list.append((positive_class, x_data, y_data, line_label, auc))
+
+        data_series = [
+            (f"label={positive_class},{line_label}", x_data, y_data)
+            for positive_class, x_data, y_data, line_label, _ in curve_list
+        ]
+        auc = [auc for _, _, _, _, auc in curve_list]
+
+    def _do_plot(**kwargs):
+        from matplotlib import pyplot
+
+        _, ax = plot_lines(**kwargs)
+        dash_line_args = {
+            "color": "gray",
+            "alpha": 0.3,
+            "drawstyle": "default",
+            "linestyle": "dashed",
+        }
+        if curve_type == "pr":
+            ax.plot([0, 1], [1, 0], **dash_line_args)
+        elif curve_type == "roc":
+            ax.plot([0, 1], [0, 1], **dash_line_args)
+
+        if is_binomial:
+            ax.legend(loc="best")
+        else:
+            ax.legend(loc="center left", bbox_to_anchor=(1, 0.5))
+            pyplot.subplots_adjust(right=0.6, bottom=0.25)
+
+    return _Curve(
+        plot_fn=_do_plot,
+        plot_fn_args={
+            "data_series": data_series,
+            "xlabel": xlabel,
+            "ylabel": ylabel,
+            "line_kwargs": {"drawstyle": "steps-post", "linewidth": 1},
+            "title": title,
+        },
+        auc=auc,
+    )

venv/lib/python3.9/site-packages/mlflow/models/evaluation/evaluators/default.py

@@ -0,0 +1,233 @@
+import logging
+import os
+import time
+from typing import Optional
+
+import numpy as np
+import pandas as pd
+
+import mlflow
+from mlflow.entities.metric import Metric
+from mlflow.exceptions import MlflowException
+from mlflow.metrics import (
+    MetricValue,
+    ari_grade_level,
+    exact_match,
+    flesch_kincaid_grade_level,
+    ndcg_at_k,
+    precision_at_k,
+    recall_at_k,
+    rouge1,
+    rouge2,
+    rougeL,
+    rougeLsum,
+    token_count,
+    toxicity,
+)
+from mlflow.metrics.genai.genai_metric import _GENAI_CUSTOM_METRICS_FILE_NAME
+from mlflow.models.evaluation.artifacts import JsonEvaluationArtifact
+from mlflow.models.evaluation.base import EvaluationMetric, EvaluationResult, _ModelType
+from mlflow.models.evaluation.default_evaluator import (
+    _LATENCY_METRIC_NAME,
+    BuiltInEvaluator,
+    _extract_output_and_other_columns,
+    _extract_predict_fn,
+)
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+
+_logger = logging.getLogger(__name__)
+
+
+class DefaultEvaluator(BuiltInEvaluator):
+    """
+    The default built-in evaluator for any models that cannot be evaluated
+    by other built-in evaluators, such as question-answering.
+    """
+
+    name = "default"
+
+    @classmethod
+    def can_evaluate(cls, *, model_type, evaluator_config, **kwargs):
+        return model_type in _ModelType.values() or model_type is None
+
+    def _evaluate(
+        self,
+        model: Optional["mlflow.pyfunc.PyFuncModel"],
+        extra_metrics: list[EvaluationMetric],
+        custom_artifacts=None,
+        **kwargs,
+    ) -> Optional[EvaluationResult]:
+        compute_latency = False
+        for extra_metric in extra_metrics:
+            # If latency metric is specified, we will compute latency for the model
+            # during prediction, and we will remove the metric from the list of extra
+            # metrics to be computed after prediction.
+            if extra_metric.name == _LATENCY_METRIC_NAME:
+                compute_latency = True
+                extra_metrics.remove(extra_metric)
+        self._log_genai_custom_metrics(extra_metrics)
+
+        # Generate model predictions and evaluate metrics
+        y_pred, other_model_outputs, self.predictions = self._generate_model_predictions(
+            model, input_df=self.X.copy_to_avoid_mutation(), compute_latency=compute_latency
+        )
+        y_true = self.dataset.labels_data
+
+        metrics = self._builtin_metrics() + extra_metrics
+        self.evaluate_metrics(
+            metrics,
+            prediction=y_pred,
+            target=self.dataset.labels_data,
+            other_output_df=other_model_outputs,
+        )
+        self.evaluate_and_log_custom_artifacts(custom_artifacts, prediction=y_pred, target=y_true)
+
+        # Log metrics and artifacts
+        self.log_metrics()
+        self.log_eval_table(y_pred, other_model_outputs)
+        return EvaluationResult(
+            metrics=self.aggregate_metrics, artifacts=self.artifacts, run_id=self.run_id
+        )
+
+    def _builtin_metrics(self) -> list[Metric]:
+        """
+        Get a list of builtin metrics for the model type.
+        """
+        text_metrics = [
+            token_count(),
+            toxicity(),
+            flesch_kincaid_grade_level(),
+            ari_grade_level(),
+        ]
+        builtin_metrics = []
+
+        # NB: Classifier and Regressor are handled by dedicated built-in evaluators,
+        if self.model_type == _ModelType.QUESTION_ANSWERING:
+            builtin_metrics = [*text_metrics, exact_match()]
+        elif self.model_type == _ModelType.TEXT_SUMMARIZATION:
+            builtin_metrics = [
+                *text_metrics,
+                rouge1(),
+                rouge2(),
+                rougeL(),
+                rougeLsum(),
+            ]
+        elif self.model_type == _ModelType.TEXT:
+            builtin_metrics = text_metrics
+        elif self.model_type == _ModelType.RETRIEVER:
+            # default k to 3 if not specified
+            retriever_k = self.evaluator_config.pop("retriever_k", 3)
+            builtin_metrics = [
+                precision_at_k(retriever_k),
+                recall_at_k(retriever_k),
+                ndcg_at_k(retriever_k),
+            ]
+
+        return builtin_metrics
+
+    def _generate_model_predictions(
+        self,
+        model: Optional["mlflow.pyfunc.PyFuncModel"],
+        input_df: pd.DataFrame,
+        compute_latency=False,
+    ):
+        """
+        Helper method for generating model predictions
+        """
+        predict_fn = _extract_predict_fn(model)
+
+        def predict_with_latency(X_copy):
+            y_pred_list = []
+            pred_latencies = []
+            if len(X_copy) == 0:
+                raise ValueError("Empty input data")
+
+            is_dataframe = isinstance(X_copy, pd.DataFrame)
+
+            for row in X_copy.iterrows() if is_dataframe else enumerate(X_copy):
+                i, row_data = row
+                single_input = row_data.to_frame().T if is_dataframe else row_data
+                start_time = time.time()
+                y_pred = predict_fn(single_input)
+                end_time = time.time()
+                pred_latencies.append(end_time - start_time)
+                y_pred_list.append(y_pred)
+
+            # Update latency metric
+            self.metrics_values.update({_LATENCY_METRIC_NAME: MetricValue(scores=pred_latencies)})
+
+            # Aggregate all predictions into model_predictions
+            sample_pred = y_pred_list[0]
+            if isinstance(sample_pred, pd.DataFrame):
+                return pd.concat(y_pred_list)
+            elif isinstance(sample_pred, np.ndarray):
+                return np.concatenate(y_pred_list, axis=0)
+            elif isinstance(sample_pred, list):
+                return sum(y_pred_list, [])
+            elif isinstance(sample_pred, pd.Series):
+                return pd.concat(y_pred_list, ignore_index=True)
+            elif isinstance(sample_pred, str):
+                return y_pred_list
+            else:
+                raise MlflowException(
+                    message=f"Unsupported prediction type {type(sample_pred)} for model type "
+                    f"{self.model_type}.",
+                    error_code=INVALID_PARAMETER_VALUE,
+                )
+
+        if model is not None:
+            _logger.info("Computing model predictions.")
+
+            if compute_latency:
+                model_predictions = predict_with_latency(input_df)
+            else:
+                model_predictions = predict_fn(input_df)
+        else:
+            if compute_latency:
+                _logger.warning(
+                    "Setting the latency to 0 for all entries because the model is not provided."
+                )
+                self.metrics_values.update(
+                    {_LATENCY_METRIC_NAME: MetricValue(scores=[0.0] * len(input_df))}
+                )
+            model_predictions = self.dataset.predictions_data
+
+        output_column_name = self.predictions
+        (
+            y_pred,
+            other_output_df,
+            predictions_column_name,
+        ) = _extract_output_and_other_columns(model_predictions, output_column_name)
+
+        return y_pred, other_output_df, predictions_column_name
+
+    def _log_genai_custom_metrics(self, extra_metrics: list[EvaluationMetric]):
+        genai_custom_metrics = [
+            extra_metric.genai_metric_args
+            for extra_metric in extra_metrics
+            # When the field is present, the metric is created from either make_genai_metric
+            # or make_genai_metric_from_prompt. We will log the metric definition.
+            if extra_metric.genai_metric_args is not None
+        ]
+
+        if len(genai_custom_metrics) == 0:
+            return
+
+        names = []
+        versions = []
+        metric_args_list = []
+
+        for metric_args in genai_custom_metrics:
+            names.append(metric_args["name"])
+            # Custom metrics created from make_genai_metric_from_prompt don't have version
+            versions.append(metric_args.get("version", ""))
+            metric_args_list.append(metric_args)
+
+        data = {"name": names, "version": versions, "metric_args": metric_args_list}
+
+        mlflow.log_table(data, artifact_file=_GENAI_CUSTOM_METRICS_FILE_NAME)
+
+        artifact_name = os.path.splitext(_GENAI_CUSTOM_METRICS_FILE_NAME)[0]
+        self.artifacts[artifact_name] = JsonEvaluationArtifact(
+            uri=mlflow.get_artifact_uri(_GENAI_CUSTOM_METRICS_FILE_NAME)
+        )

venv/lib/python3.9/site-packages/mlflow/models/evaluation/evaluators/regressor.py

@@ -0,0 +1,96 @@
+from typing import Optional
+
+import numpy as np
+from sklearn import metrics as sk_metrics
+
+import mlflow
+from mlflow.models.evaluation.base import EvaluationMetric, EvaluationResult, _ModelType
+from mlflow.models.evaluation.default_evaluator import (
+    BuiltInEvaluator,
+    _extract_output_and_other_columns,
+    _extract_predict_fn,
+    _get_aggregate_metrics_values,
+)
+
+
+class RegressorEvaluator(BuiltInEvaluator):
+    """
+    A built-in evaluator for regressor models.
+    """
+
+    name = "regressor"
+
+    @classmethod
+    def can_evaluate(cls, *, model_type, evaluator_config, **kwargs):
+        return model_type == _ModelType.REGRESSOR
+
+    def _evaluate(
+        self,
+        model: Optional["mlflow.pyfunc.PyFuncModel"],
+        extra_metrics: list[EvaluationMetric],
+        custom_artifacts=None,
+        **kwargs,
+    ) -> Optional[EvaluationResult]:
+        self.y_true = self.dataset.labels_data
+        self.sample_weights = self.evaluator_config.get("sample_weights", None)
+
+        input_df = self.X.copy_to_avoid_mutation()
+        self.y_pred = self._generate_model_predictions(model, input_df)
+        self._compute_buildin_metrics(model)
+
+        self.evaluate_metrics(extra_metrics, prediction=self.y_pred, target=self.y_true)
+        self.evaluate_and_log_custom_artifacts(
+            custom_artifacts, prediction=self.y_pred, target=self.y_true
+        )
+
+        self.log_metrics()
+        self.log_eval_table(self.y_pred)
+
+        return EvaluationResult(
+            metrics=self.aggregate_metrics, artifacts=self.artifacts, run_id=self.run_id
+        )
+
+    def _generate_model_predictions(self, model, input_df):
+        if predict_fn := _extract_predict_fn(model):
+            preds = predict_fn(input_df)
+            y_pred, _, _ = _extract_output_and_other_columns(preds, self.predictions)
+            return y_pred
+        else:
+            return self.dataset.predictions_data
+
+    def _compute_buildin_metrics(self, model):
+        self._evaluate_sklearn_model_score_if_scorable(model, self.y_true, self.sample_weights)
+        self.metrics_values.update(
+            _get_aggregate_metrics_values(
+                _get_regressor_metrics(self.y_true, self.y_pred, self.sample_weights)
+            )
+        )
+
+
+def _get_regressor_metrics(y, y_pred, sample_weights):
+    from mlflow.metrics.metric_definitions import _root_mean_squared_error
+
+    sum_on_target = (
+        (np.array(y) * np.array(sample_weights)).sum() if sample_weights is not None else sum(y)
+    )
+    return {
+        "example_count": len(y),
+        "mean_absolute_error": sk_metrics.mean_absolute_error(
+            y, y_pred, sample_weight=sample_weights
+        ),
+        "mean_squared_error": sk_metrics.mean_squared_error(
+            y, y_pred, sample_weight=sample_weights
+        ),
+        "root_mean_squared_error": _root_mean_squared_error(
+            y_true=y,
+            y_pred=y_pred,
+            sample_weight=sample_weights,
+        ),
+        "sum_on_target": sum_on_target,
+        "mean_on_target": sum_on_target / len(y),
+        "r2_score": sk_metrics.r2_score(y, y_pred, sample_weight=sample_weights),
+        "max_error": sk_metrics.max_error(y, y_pred),
+        "mean_absolute_percentage_error": sk_metrics.mean_absolute_percentage_error(
+            y, y_pred, sample_weight=sample_weights
+        ),
+    }

venv/lib/python3.9/site-packages/mlflow/models/evaluation/evaluators/shap.py

@@ -0,0 +1,293 @@
+import functools
+import logging
+from typing import Optional
+
+import numpy as np
+from packaging.version import Version
+from sklearn.pipeline import Pipeline as sk_Pipeline
+
+import mlflow
+from mlflow import MlflowException
+from mlflow.models.evaluation.base import EvaluationMetric, EvaluationResult, _ModelType
+from mlflow.models.evaluation.default_evaluator import (
+    BuiltInEvaluator,
+    _extract_predict_fn,
+    _extract_raw_model,
+    _get_dataframe_with_renamed_columns,
+)
+from mlflow.models.evaluation.evaluators.classifier import (
+    _is_continuous,
+    _suppress_class_imbalance_errors,
+)
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.pyfunc import _ServedPyFuncModel
+
+_logger = logging.getLogger(__name__)
+
+
+_SUPPORTED_SHAP_ALGORITHMS = ("exact", "permutation", "partition", "kernel")
+_DEFAULT_SAMPLE_ROWS_FOR_SHAP = 2000
+
+
+def _shap_predict_fn(x, predict_fn, feature_names):
+    return predict_fn(_get_dataframe_with_renamed_columns(x, feature_names))
+
+
+class ShapEvaluator(BuiltInEvaluator):
+    """
+    A built-in evaluator to get SHAP explainability insights for classifier and regressor models.
+
+    This evaluator often run with the main evaluator for the model like ClassifierEvaluator.
+    """
+
+    name = "shap"
+
+    @classmethod
+    def can_evaluate(cls, *, model_type, evaluator_config, **kwargs):
+        return model_type in (_ModelType.CLASSIFIER, _ModelType.REGRESSOR) and evaluator_config.get(
+            "log_model_explainability", True
+        )
+
+    def _evaluate(
+        self,
+        model: Optional["mlflow.pyfunc.PyFuncModel"],
+        extra_metrics: list[EvaluationMetric],
+        custom_artifacts=None,
+        **kwargs,
+    ) -> Optional[EvaluationResult]:
+        if isinstance(model, _ServedPyFuncModel):
+            _logger.warning(
+                "Skipping model explainability because a model server is used for environment "
+                "restoration."
+            )
+            return
+
+        model_loader_module, raw_model = _extract_raw_model(model)
+        if model_loader_module == "mlflow.spark":
+            # TODO: Shap explainer need to manipulate on each feature values,
+            #  but spark model input dataframe contains Vector type feature column
+            #  which shap explainer does not support.
+            #  To support this, we need expand the Vector type feature column into
+            #  multiple scalar feature columns and pass it to shap explainer.
+            _logger.warning(
+                "Logging model explainability insights is not currently supported for PySpark "
+                "models."
+            )
+            return
+
+        self.y_true = self.dataset.labels_data
+        self.label_list = self.evaluator_config.get("label_list")
+        self.pos_label = self.evaluator_config.get("pos_label")
+
+        if not (np.issubdtype(self.y_true.dtype, np.number) or self.y_true.dtype == np.bool_):
+            # Note: python bool type inherits number type but np.bool_ does not inherit np.number.
+            _logger.warning(
+                "Skip logging model explainability insights because it requires all label "
+                "values to be numeric or boolean."
+            )
+            return
+
+        algorithm = self.evaluator_config.get("explainability_algorithm", None)
+        if algorithm is not None and algorithm not in _SUPPORTED_SHAP_ALGORITHMS:
+            raise MlflowException(
+                message=f"Specified explainer algorithm {algorithm} is unsupported. Currently only "
+                f"support {','.join(_SUPPORTED_SHAP_ALGORITHMS)} algorithms.",
+                error_code=INVALID_PARAMETER_VALUE,
+            )
+
+        if algorithm != "kernel":
+            feature_dtypes = list(self.X.get_original().dtypes)
+            for feature_dtype in feature_dtypes:
+                if not np.issubdtype(feature_dtype, np.number):
+                    _logger.warning(
+                        "Skip logging model explainability insights because the shap explainer "
+                        f"{algorithm} requires all feature values to be numeric, and each feature "
+                        "column must only contain scalar values."
+                    )
+                    return
+
+        try:
+            import shap
+            from matplotlib import pyplot
+        except ImportError:
+            _logger.warning(
+                "SHAP or matplotlib package is not installed, so model explainability insights "
+                "will not be logged."
+            )
+            return
+
+        if Version(shap.__version__) < Version("0.40"):
+            _logger.warning(
+                "Shap package version is lower than 0.40, Skip log model explainability."
+            )
+            return
+
+        sample_rows = self.evaluator_config.get(
+            "explainability_nsamples", _DEFAULT_SAMPLE_ROWS_FOR_SHAP
+        )
+
+        X_df = self.X.copy_to_avoid_mutation()
+
+        sampled_X = shap.sample(X_df, sample_rows, random_state=0)
+
+        mode_or_mean_dict = _compute_df_mode_or_mean(X_df)
+        sampled_X = sampled_X.fillna(mode_or_mean_dict)
+
+        # shap explainer might call provided `predict_fn` with a `numpy.ndarray` type
+        # argument, this might break some model inference, so convert the argument into
+        # a pandas dataframe.
+        # The `shap_predict_fn` calls model's predict function, we need to restore the input
+        # dataframe with original column names, because some model prediction routine uses
+        # the column name.
+
+        predict_fn = _extract_predict_fn(model)
+        shap_predict_fn = functools.partial(
+            _shap_predict_fn, predict_fn=predict_fn, feature_names=self.dataset.feature_names
+        )
+
+        if self.label_list is None:
+            # If label list is not specified, infer label list from model output.
+            # We need to copy the input data as the model might mutate the input data.
+            y_pred = predict_fn(X_df.copy()) if predict_fn else self.dataset.predictions_data
+            self.label_list = np.unique(np.concatenate([self.y_true, y_pred]))
+
+        try:
+            if algorithm:
+                if algorithm == "kernel":
+                    # We need to lazily import shap, so lazily import `_PatchedKernelExplainer`
+                    from mlflow.models.evaluation._shap_patch import _PatchedKernelExplainer
+
+                    kernel_link = self.evaluator_config.get(
+                        "explainability_kernel_link", "identity"
+                    )
+                    if kernel_link not in ["identity", "logit"]:
+                        raise ValueError(
+                            "explainability_kernel_link config can only be set to 'identity' or "
+                            f"'logit', but got '{kernel_link}'."
+                        )
+                    background_X = shap.sample(X_df, sample_rows, random_state=3)
+                    background_X = background_X.fillna(mode_or_mean_dict)
+
+                    explainer = _PatchedKernelExplainer(
+                        shap_predict_fn, background_X, link=kernel_link
+                    )
+                else:
+                    explainer = shap.Explainer(
+                        shap_predict_fn,
+                        sampled_X,
+                        feature_names=self.dataset.feature_names,
+                        algorithm=algorithm,
+                    )
+            else:
+                if (
+                    raw_model
+                    and not len(self.label_list) > 2
+                    and not isinstance(raw_model, sk_Pipeline)
+                ):
+                    # For mulitnomial classifier, shap.Explainer may choose Tree/Linear explainer
+                    # for raw model, this case shap plot doesn't support it well, so exclude the
+                    # multinomial_classifier case here.
+                    explainer = shap.Explainer(
+                        raw_model, sampled_X, feature_names=self.dataset.feature_names
+                    )
+                else:
+                    # fallback to default explainer
+                    explainer = shap.Explainer(
+                        shap_predict_fn, sampled_X, feature_names=self.dataset.feature_names
+                    )
+
+            _logger.info(f"Shap explainer {explainer.__class__.__name__} is used.")
+
+            if algorithm == "kernel":
+                shap_values = shap.Explanation(
+                    explainer.shap_values(sampled_X), feature_names=self.dataset.feature_names
+                )
+            else:
+                shap_values = explainer(sampled_X)
+        except Exception as e:
+            # Shap evaluation might fail on some edge cases, e.g., unsupported input data values
+            # or unsupported model on specific shap explainer. Catch exception to prevent it
+            # breaking the whole `evaluate` function.
+
+            if not self.evaluator_config.get("ignore_exceptions", True):
+                raise e
+
+            _logger.warning(
+                f"Shap evaluation failed. Reason: {e!r}. "
+                "Set logging level to DEBUG to see the full traceback."
+            )
+            _logger.debug("", exc_info=True)
+            return
+        try:
+            mlflow.shap.log_explainer(explainer, artifact_path="explainer")
+        except Exception as e:
+            # TODO: The explainer saver is buggy, if `get_underlying_model_flavor` return "unknown",
+            #   then fallback to shap explainer saver, and shap explainer will call `model.save`
+            #   for sklearn model, there is no `.save` method, so error will happen.
+            _logger.warning(
+                f"Logging explainer failed. Reason: {e!r}. "
+                "Set logging level to DEBUG to see the full traceback."
+            )
+            _logger.debug("", exc_info=True)
+
+        def _adjust_color_bar():
+            pyplot.gcf().axes[-1].set_aspect("auto")
+            pyplot.gcf().axes[-1].set_box_aspect(50)
+
+        def _adjust_axis_tick():
+            pyplot.xticks(fontsize=10)
+            pyplot.yticks(fontsize=10)
+
+        def plot_beeswarm():
+            shap.plots.beeswarm(shap_values, show=False, color_bar=True)
+            _adjust_color_bar()
+            _adjust_axis_tick()
+
+        with _suppress_class_imbalance_errors(ValueError, log_warning=False):
+            self._log_image_artifact(
+                plot_beeswarm,
+                "shap_beeswarm_plot",
+            )
+
+        def plot_summary():
+            shap.summary_plot(shap_values, show=False, color_bar=True)
+            _adjust_color_bar()
+            _adjust_axis_tick()
+
+        with _suppress_class_imbalance_errors(TypeError, log_warning=False):
+            self._log_image_artifact(
+                plot_summary,
+                "shap_summary_plot",
+            )
+
+        def plot_feature_importance():
+            shap.plots.bar(shap_values, show=False)
+            _adjust_axis_tick()
+
+        with _suppress_class_imbalance_errors(IndexError, log_warning=False):
+            self._log_image_artifact(
+                plot_feature_importance,
+                "shap_feature_importance_plot",
+            )
+
+        return EvaluationResult(
+            metrics=self.aggregate_metrics,
+            artifacts=self.artifacts,
+            run_id=self.run_id,
+        )
+
+
+def _compute_df_mode_or_mean(df):
+    """
+    Compute mean (for continuous columns) and compute mode (for other columns) for the
+    input dataframe, return a dict, key is column name, value is the corresponding mode or
+    mean value, this function calls `_is_continuous` to determine whether the
+    column is continuous column.
+    """
+    continuous_cols = [c for c in df.columns if _is_continuous(df[c])]
+    df_cont = df[continuous_cols]
+    df_non_cont = df.drop(continuous_cols, axis=1)
+
+    means = {} if df_cont.empty else df_cont.mean().to_dict()
+    modes = {} if df_non_cont.empty else df_non_cont.mode().loc[0].to_dict()
+    return {**means, **modes}

venv/lib/python3.9/site-packages/mlflow/models/evaluation/lift_curve.py

@@ -0,0 +1,177 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def _cumulative_gain_curve(y_true, y_score, pos_label=None):
+    """
+    This method is copied from scikit-plot package.
+    See https://github.com/reiinakano/scikit-plot/blob/2dd3e6a76df77edcbd724c4db25575f70abb57cb/scikitplot/helpers.py#L157
+
+    This function generates the points necessary to plot the Cumulative Gain
+
+    Note: This implementation is restricted to the binary classification task.
+
+    Args:
+        y_true (array-like, shape (n_samples)): True labels of the data.
+
+        y_score (array-like, shape (n_samples)): Target scores, can either be
+            probability estimates of the positive class, confidence values, or
+            non-thresholded measure of decisions (as returned by
+            decision_function on some classifiers).
+
+        pos_label (int or str, default=None): Label considered as positive and
+            others are considered negative
+
+    Returns:
+        percentages (numpy.ndarray): An array containing the X-axis values for
+            plotting the Cumulative Gains chart.
+
+        gains (numpy.ndarray): An array containing the Y-axis values for one
+            curve of the Cumulative Gains chart.
+
+    Raises:
+        ValueError: If `y_true` is not composed of 2 classes. The Cumulative
+            Gain Chart is only relevant in binary classification.
+    """
+    y_true, y_score = np.asarray(y_true), np.asarray(y_score)
+
+    # ensure binary classification if pos_label is not specified
+    classes = np.unique(y_true)
+    if pos_label is None and not (
+        np.array_equal(classes, [0, 1])
+        or np.array_equal(classes, [-1, 1])
+        or np.array_equal(classes, [0])
+        or np.array_equal(classes, [-1])
+        or np.array_equal(classes, [1])
+    ):
+        raise ValueError("Data is not binary and pos_label is not specified")
+    elif pos_label is None:
+        pos_label = 1.0
+
+    # make y_true a boolean vector
+    y_true = y_true == pos_label
+
+    sorted_indices = np.argsort(y_score)[::-1]
+    y_true = y_true[sorted_indices]
+    gains = np.cumsum(y_true)
+
+    percentages = np.arange(start=1, stop=len(y_true) + 1)
+
+    gains = gains / float(np.sum(y_true))
+    percentages = percentages / float(len(y_true))
+
+    gains = np.insert(gains, 0, [0])
+    percentages = np.insert(percentages, 0, [0])
+
+    return percentages, gains
+
+
+def plot_lift_curve(
+    y_true,
+    y_probas,
+    title="Lift Curve",
+    ax=None,
+    figsize=None,
+    title_fontsize="large",
+    text_fontsize="medium",
+    pos_label=None,
+):
+    """
+    This method is copied from scikit-plot package.
+    See https://github.com/reiinakano/scikit-plot/blob/2dd3e6a76df77edcbd724c4db25575f70abb57cb/scikitplot/metrics.py#L1133
+
+    Generates the Lift Curve from labels and scores/probabilities
+
+    The lift curve is used to determine the effectiveness of a
+    binary classifier. A detailed explanation can be found at
+    http://www2.cs.uregina.ca/~dbd/cs831/notes/lift_chart/lift_chart.html.
+    The implementation here works only for binary classification.
+
+    Args:
+        y_true (array-like, shape (n_samples)):
+            Ground truth (correct) target values.
+
+        y_probas (array-like, shape (n_samples, n_classes)):
+            Prediction probabilities for each class returned by a classifier.
+
+        title (string, optional): Title of the generated plot. Defaults to
+            "Lift Curve".
+
+        ax (:class:`matplotlib.axes.Axes`, optional): The axes upon which to
+            plot the learning curve. If None, the plot is drawn on a new set of
+            axes.
+
+        figsize (2-tuple, optional): Tuple denoting figure size of the plot
+            e.g. (6, 6). Defaults to ``None``.
+
+        title_fontsize (string or int, optional): Matplotlib-style fontsizes.
+            Use e.g. "small", "medium", "large" or integer-values. Defaults to
+            "large".
+
+        text_fontsize (string or int, optional): Matplotlib-style fontsizes.
+            Use e.g. "small", "medium", "large" or integer-values. Defaults to
+            "medium".
+
+        pos_label (optional): Label for the positive class.
+
+    Returns:
+        ax (:class:`matplotlib.axes.Axes`): The axes on which the plot was
+            drawn.
+
+    Example:
+        >>> lr = LogisticRegression()
+        >>> lr = lr.fit(X_train, y_train)
+        >>> y_probas = lr.predict_proba(X_test)
+        >>> plot_lift_curve(y_test, y_probas)
+        <matplotlib.axes._subplots.AxesSubplot object at 0x7fe967d64490>
+        >>> plt.show()
+
+        .. image:: _static/examples/plot_lift_curve.png
+           :align: center
+           :alt: Lift Curve
+    """
+    y_true = np.array(y_true)
+    y_probas = np.array(y_probas)
+
+    classes = np.unique(y_true)
+    if len(classes) != 2:
+        raise ValueError(f"Cannot calculate Lift Curve for data with {len(classes)} category/ies")
+
+    # Compute Cumulative Gain Curves
+    percentages, gains1 = _cumulative_gain_curve(y_true, y_probas[:, 0], classes[0])
+    percentages, gains2 = _cumulative_gain_curve(y_true, y_probas[:, 1], classes[1])
+
+    percentages = percentages[1:]
+    gains1 = gains1[1:]
+    gains2 = gains2[1:]
+
+    gains1 = gains1 / percentages
+    gains2 = gains2 / percentages
+
+    if ax is None:
+        _, ax = plt.subplots(1, 1, figsize=figsize)
+
+    ax.set_title(title, fontsize=title_fontsize)
+
+    label0 = f"Class {classes[0]}"
+    label1 = f"Class {classes[1]}"
+    # show (positive) next to the positive class in the legend
+    if pos_label:
+        if pos_label == classes[0]:
+            label0 = f"Class {classes[0]} (positive)"
+        elif pos_label == classes[1]:
+            label1 = f"Class {classes[1]} (positive)"
+        # do not mark positive class if pos_label is not in classes
+
+    ax.plot(percentages, gains1, lw=3, label=label0)
+    ax.plot(percentages, gains2, lw=3, label=label1)
+
+    ax.plot([0, 1], [1, 1], "k--", lw=2, label="Baseline")
+
+    ax.set_xlabel("Percentage of sample", fontsize=text_fontsize)
+    ax.set_ylabel("Lift", fontsize=text_fontsize)
+    ax.tick_params(labelsize=text_fontsize)
+    ax.grid("on")
+    ax.legend(loc="best", fontsize=text_fontsize)
+
+    return ax

venv/lib/python3.9/site-packages/mlflow/models/evaluation/utils/metric.py

@@ -0,0 +1,123 @@
+import logging
+from dataclasses import dataclass
+from typing import Any, Callable, Optional
+
+import numpy as np
+
+from mlflow.metrics.base import MetricValue
+from mlflow.models.evaluation.base import EvaluationMetric
+
+_logger = logging.getLogger(__name__)
+
+
+@dataclass
+class MetricDefinition:
+    """
+    A namedtuple representing a metric function and its properties.
+
+    function : the metric function
+    name : the name of the metric function
+    index : the index of the function in the ``extra_metrics`` argument of mlflow.evaluate
+    """
+
+    function: Callable[..., Any]
+    name: str
+    index: int
+    version: Optional[str] = None
+    genai_metric_args: Optional[dict[str, Any]] = None
+
+    @classmethod
+    def from_index_and_metric(cls, index: int, metric: EvaluationMetric):
+        return cls(
+            function=metric.eval_fn,
+            index=index,
+            name=metric.name,
+            version=metric.version,
+            genai_metric_args=metric.genai_metric_args,
+        )
+
+    def evaluate(self, eval_fn_args) -> Optional[MetricValue]:
+        """
+        This function calls the metric function and performs validations on the returned
+        result to ensure that they are in the expected format. It will warn and will not log metrics
+        that are in the wrong format.
+
+        Args:
+            eval_fn_args: A dictionary of args needed to compute the eval metrics.
+
+        Returns:
+            MetricValue
+        """
+        if self.index < 0:
+            exception_header = f"Did not log builtin metric '{self.name}' because it"
+        else:
+            exception_header = (
+                f"Did not log metric '{self.name}' at index "
+                f"{self.index} in the `extra_metrics` parameter because it"
+            )
+
+        metric: MetricValue = self.function(*eval_fn_args)
+
+        def _is_numeric(value):
+            return isinstance(value, (int, float, np.number))
+
+        def _is_string(value):
+            return isinstance(value, str)
+
+        if metric is None:
+            _logger.warning(f"{exception_header} returned None.")
+            return
+
+        if _is_numeric(metric):
+            return MetricValue(aggregate_results={self.name: metric})
+
+        if not isinstance(metric, MetricValue):
+            _logger.warning(f"{exception_header} did not return a MetricValue.")
+            return
+
+        scores = metric.scores
+        justifications = metric.justifications
+        aggregates = metric.aggregate_results
+
+        if scores is not None:
+            if not isinstance(scores, list):
+                _logger.warning(
+                    f"{exception_header} must return MetricValue with scores as a list."
+                )
+                return
+            if any(not (_is_numeric(s) or _is_string(s) or s is None) for s in scores):
+                _logger.warning(
+                    f"{exception_header} must return MetricValue with numeric or string scores."
+                )
+                return
+
+        if justifications is not None:
+            if not isinstance(justifications, list):
+                _logger.warning(
+                    f"{exception_header} must return MetricValue with justifications as a list."
+                )
+                return
+            if any(not (_is_string(just) or just is None) for just in justifications):
+                _logger.warning(
+                    f"{exception_header} must return MetricValue with string justifications."
+                )
+                return
+
+        if aggregates is not None:
+            if not isinstance(aggregates, dict):
+                _logger.warning(
+                    f"{exception_header} must return MetricValue with aggregate_results as a dict."
+                )
+                return
+
+            if any(
+                not (isinstance(k, str) and (_is_numeric(v) or v is None))
+                for k, v in aggregates.items()
+            ):
+                _logger.warning(
+                    f"{exception_header} must return MetricValue with aggregate_results with "
+                    "str keys and numeric values."
+                )
+                return
+
+        return metric

venv/lib/python3.9/site-packages/mlflow/models/evaluation/utils/trace.py

@@ -0,0 +1,177 @@
+import contextlib
+import inspect
+import logging
+from typing import Any, Callable
+
+from mlflow.ml_package_versions import FLAVOR_TO_MODULE_NAME
+from mlflow.utils.autologging_utils import (
+    AUTOLOGGING_INTEGRATIONS,
+    autologging_conf_lock,
+    get_autolog_function,
+    is_autolog_supported,
+)
+from mlflow.utils.autologging_utils.safety import revert_patches
+from mlflow.utils.import_hooks import (
+    _post_import_hooks,
+    get_post_import_hooks,
+    register_post_import_hook,
+)
+
+_logger = logging.getLogger(__name__)
+
+
+# This flag is used to display the message only once when tracing is enabled during the evaluation.
+_SHOWN_TRACE_MESSAGE_BEFORE = False
+
+
+@contextlib.contextmanager
+@autologging_conf_lock
+def configure_autologging_for_evaluation(enable_tracing: bool = True):
+    """
+    Temporarily override the autologging configuration for all flavors during the model evaluation.
+    For example, model auto-logging must be disabled during the evaluation. After the evaluation
+    is done, the original autologging configurations are restored.
+
+    Args:
+        enable_tracing (bool): Whether to enable tracing for the supported flavors during eval.
+    """
+    original_import_hooks = {}
+    new_import_hooks = {}
+
+    # AUTOLOGGING_INTEGRATIONS can change during we iterate over flavors and enable/disable
+    # autologging, therefore, we snapshot the current configuration to restore it later.
+    global_config_snapshot = AUTOLOGGING_INTEGRATIONS.copy()
+
+    for flavor in FLAVOR_TO_MODULE_NAME:
+        if not is_autolog_supported(flavor):
+            continue
+
+        original_config = global_config_snapshot.get(flavor, {}).copy()
+
+        # If autologging is explicitly disabled, do nothing.
+        if original_config.get("disable", False):
+            continue
+
+        # NB: Using post-import hook to configure the autologging lazily when the target
+        # flavor's module is imported, rather than configuring it immediately. This is
+        # because the evaluation code usually only uses a subset of the supported flavors,
+        # hence we want to avoid unnecessary overhead of configuring all flavors.
+        @autologging_conf_lock
+        def _setup_autolog(module):
+            try:
+                autolog = get_autolog_function(flavor)
+
+                # If tracing is supported and not explicitly disabled, enable it.
+                if enable_tracing and _should_enable_tracing(flavor, global_config_snapshot):
+                    new_config = {
+                        k: False if k.startswith("log_") else v for k, v in original_config.items()
+                    }
+                    new_config |= {"log_traces": True, "silent": True}
+                    _kwargs_safe_invoke(autolog, new_config)
+
+                    global _SHOWN_TRACE_MESSAGE_BEFORE
+                    if not _SHOWN_TRACE_MESSAGE_BEFORE:
+                        _logger.info(
+                            "Auto tracing is temporarily enabled during the model evaluation "
+                            "for computing some metrics and debugging. To disable tracing, call "
+                            "`mlflow.autolog(disable=True)`."
+                        )
+                        _SHOWN_TRACE_MESSAGE_BEFORE = True
+                else:
+                    autolog(disable=True)
+
+            except Exception:
+                _logger.debug(f"Failed to update autologging config for {flavor}.", exc_info=True)
+
+        module = FLAVOR_TO_MODULE_NAME[flavor]
+        try:
+            original_import_hooks[module] = get_post_import_hooks(module)
+            new_import_hooks[module] = _setup_autolog
+            register_post_import_hook(_setup_autolog, module, overwrite=True)
+        except Exception:
+            _logger.debug(f"Failed to register post-import hook for {flavor}.", exc_info=True)
+
+    try:
+        yield
+    finally:
+        # Remove post-import hooks and patches the are registered during the evaluation.
+        for module, hooks in new_import_hooks.items():
+            # Restore original post-import hooks if any. Note that we don't use
+            # register_post_import_hook method to bypass some pre-checks and just
+            # restore the original state.
+            if hooks is None:
+                _post_import_hooks.pop(module, None)
+            else:
+                _post_import_hooks[module] = original_import_hooks[module]
+
+        # If any autologging configuration is updated, restore original autologging configurations.
+        for flavor, new_config in AUTOLOGGING_INTEGRATIONS.copy().items():
+            original_config = global_config_snapshot.get(flavor)
+            if original_config != new_config:
+                try:
+                    autolog = get_autolog_function(flavor)
+                    if original_config:
+                        _kwargs_safe_invoke(autolog, original_config)
+                        AUTOLOGGING_INTEGRATIONS[flavor] = original_config
+                    else:
+                        # If the original configuration is empty, autologging was not enabled before
+                        autolog(disable=True)
+                        # Remove all safe_patch applied by autologging
+                        revert_patches(flavor)
+                        # We also need to remove the config entry from AUTOLOGGING_INTEGRATIONS,
+                        # so as not to confuse with the case user explicitly disabled autologging.
+                        AUTOLOGGING_INTEGRATIONS.pop(flavor, None)
+                except ImportError:
+                    pass
+                except Exception as e:
+                    if original_config is None or (
+                        not original_config.get("disable", False)
+                        and not original_config.get("silent", False)
+                    ):
+                        _logger.warning(
+                            f"Exception raised while calling autologging for {flavor}: {e}"
+                        )
+
+
+def _should_enable_tracing(flavor: str, autologging_config: dict[str, Any]) -> bool:
+    """
+    Check if tracing should be enabled for the given flavor during the model evaluation.
+    """
+    # 1. Check if the autologging or tracing is globally disabled
+    # TODO: This check should not take precedence over the flavor-specific configuration
+    # set by the explicit mlflow.<flavor>.autolog() call by users.
+    # However, in Databricks, sometimes mlflow.<flavor>.autolog() is automatically
+    # called in the kernel startup, which is confused with the user's action. In
+    # such cases, even when user disables autologging globally, the flavor-specific
+    # autologging remains enabled. We are going to fix the Databricks side issue,
+    # and after that, we should move this check down after the flavor-specific check.
+    global_config = autologging_config.get("mlflow", {})
+    if global_config.get("disable", False) or (not global_config.get("log_traces", True)):
+        return False
+
+    if not _is_trace_autologging_supported(flavor):
+        return False
+
+    # 3. Check if tracing is explicitly disabled for the flavor
+    flavor_config = autologging_config.get(flavor, {})
+    return flavor_config.get("log_traces", True)
+
+
+def _kwargs_safe_invoke(func: Callable[..., Any], kwargs: dict[str, Any]):
+    """
+    Invoke the function with the given dictionary as keyword arguments, but only include the
+    arguments that are present in the function's signature.
+
+    This is particularly used for calling autolog() function with the configuration dictionary
+    stored in AUTOLOGGING_INTEGRATIONS. While the config keys mostly align with the autolog()'s
+    signature by design, some keys are not present in autolog(), such as "globally_configured".
+    """
+    sig = inspect.signature(func)
+    return func(**{k: v for k, v in kwargs.items() if k in sig.parameters})
+
+
+def _is_trace_autologging_supported(flavor_name: str) -> bool:
+    """Check if the given flavor supports trace autologging."""
+    if autolog_func := get_autolog_function(flavor_name):
+        return "log_traces" in inspect.signature(autolog_func).parameters
+    return False

venv/lib/python3.9/site-packages/mlflow/models/evaluation/validation.py

@@ -0,0 +1,456 @@
+import logging
+import operator
+import os
+from decimal import Decimal
+from typing import Optional
+
+from mlflow.exceptions import MlflowException
+from mlflow.models.evaluation import EvaluationResult
+from mlflow.protos.databricks_pb2 import BAD_REQUEST, INVALID_PARAMETER_VALUE
+from mlflow.utils.annotations import deprecated
+
+_logger = logging.getLogger(__name__)
+
+
+class MetricThreshold:
+    """
+    This class allows you to define metric thresholds for model validation.
+    Allowed thresholds are: threshold, min_absolute_change, min_relative_change.
+
+    Args:
+        threshold: (Optional) A number representing the value threshold for the metric.
+
+            - If higher is better for the metric, the metric value has to be
+              >= threshold to pass validation.
+            - Otherwise, the metric value has to be <= threshold to pass the validation.
+
+        min_absolute_change: (Optional) A positive number representing the minimum absolute
+            change required for candidate model to pass validation with
+            the baseline model.
+
+            - If higher is better for the metric, metric value has to be
+              >= baseline model metric value + min_absolute_change to pass the validation.
+            - Otherwise, metric value has to be <= baseline model metric value - min_absolute_change
+              to pass the validation.
+
+        min_relative_change: (Optional) A floating point number between 0 and 1 representing
+            the minimum relative change (in percentage of
+            baseline model metric value) for candidate model
+            to pass the comparison with the baseline model.
+
+            - If higher is better for the metric, metric value has to be
+              >= baseline model metric value * (1 + min_relative_change)
+            - Otherwise, metric value has to be
+              <= baseline model metric value * (1 - min_relative_change)
+            - Note that if the baseline model metric value is equal to 0, the
+              threshold falls back performing a simple verification that the
+              candidate metric value is better than the baseline metric value,
+              i.e. metric value >= baseline model metric value + 1e-10 if higher
+              is better; metric value <= baseline model metric value - 1e-10 if
+              lower is better.
+
+        greater_is_better: A required boolean representing whether higher value is
+            better for the metric.
+
+        higher_is_better:
+            .. deprecated:: 2.3.0
+                Use ``greater_is_better`` instead.
+
+            A required boolean representing whether higher value is better for the metric.
+    """
+
+    def __init__(
+        self,
+        threshold=None,
+        min_absolute_change=None,
+        min_relative_change=None,
+        greater_is_better=None,
+        higher_is_better=None,
+    ):
+        if threshold is not None and type(threshold) not in {int, float}:
+            raise MetricThresholdClassException("`threshold` parameter must be a number.")
+        if min_absolute_change is not None and (
+            type(min_absolute_change) not in {int, float} or min_absolute_change <= 0
+        ):
+            raise MetricThresholdClassException(
+                "`min_absolute_change` parameter must be a positive number."
+            )
+        if min_relative_change is not None:
+            if not isinstance(min_relative_change, float):
+                raise MetricThresholdClassException(
+                    "`min_relative_change` parameter must be a floating point number."
+                )
+            if min_relative_change < 0 or min_relative_change > 1:
+                raise MetricThresholdClassException(
+                    "`min_relative_change` parameter must be between 0 and 1."
+                )
+        if higher_is_better is None and greater_is_better is None:
+            raise MetricThresholdClassException("`greater_is_better` parameter must be defined.")
+        if higher_is_better is not None and greater_is_better is not None:
+            raise MetricThresholdClassException(
+                "`higher_is_better` parameter must be None when `greater_is_better` is defined."
+            )
+        if greater_is_better is None:
+            greater_is_better = higher_is_better
+        if not isinstance(greater_is_better, bool):
+            raise MetricThresholdClassException("`greater_is_better` parameter must be a boolean.")
+        if threshold is None and min_absolute_change is None and min_relative_change is None:
+            raise MetricThresholdClassException("no threshold was specified.")
+        self._threshold = threshold
+        self._min_absolute_change = min_absolute_change
+        self._min_relative_change = min_relative_change
+        self._greater_is_better = greater_is_better
+
+    @property
+    def threshold(self):
+        """
+        Value of the threshold.
+        """
+        return self._threshold
+
+    @property
+    def min_absolute_change(self):
+        """
+        Value of the minimum absolute change required to pass model comparison with baseline model.
+        """
+        return self._min_absolute_change
+
+    @property
+    def min_relative_change(self):
+        """
+        Float value of the minimum relative change required to pass model comparison with
+        baseline model.
+        """
+        return self._min_relative_change
+
+    @property
+    @deprecated("The attribute `higher_is_better` is deprecated. Use `greater_is_better` instead.")
+    def higher_is_better(self):
+        """
+        Boolean value representing whether higher value is better for the metric.
+        """
+        return self._greater_is_better
+
+    @property
+    def greater_is_better(self):
+        """
+        Boolean value representing whether higher value is better for the metric.
+        """
+        return self._greater_is_better
+
+    def __str__(self):
+        """
+        Returns a human-readable string consisting of all specified thresholds.
+        """
+        threshold_strs = []
+        if self._threshold is not None:
+            threshold_strs.append(f"Threshold: {self._threshold}.")
+        if self._min_absolute_change is not None:
+            threshold_strs.append(f"Minimum Absolute Change: {self._min_absolute_change}.")
+        if self._min_relative_change is not None:
+            threshold_strs.append(f"Minimum Relative Change: {self._min_relative_change}.")
+        if self._greater_is_better is not None:
+            if self._greater_is_better:
+                threshold_strs.append("Higher value is better.")
+            else:
+                threshold_strs.append("Lower value is better.")
+        return " ".join(threshold_strs)
+
+
+class MetricThresholdClassException(MlflowException):
+    def __init__(self, _message, **kwargs):
+        message = "Could not instantiate MetricThreshold class: " + _message
+        super().__init__(message, error_code=INVALID_PARAMETER_VALUE, **kwargs)
+
+
+class _MetricValidationResult:
+    """
+    Internal class for representing validation result per metric.
+    Not user facing, used for organizing metric failures and generating failure message
+    more conveniently.
+
+    Args:
+        metric_name: String representing the metric name
+        candidate_metric_value: value of metric for candidate model
+        metric_threshold: :py:class: `MetricThreshold<mlflow.models.validation.MetricThreshold>`
+            The MetricThreshold for the metric.
+        baseline_metric_value: value of metric for baseline model
+    """
+
+    missing_candidate = False
+    missing_baseline = False
+    threshold_failed = False
+    min_absolute_change_failed = False
+    min_relative_change_failed = False
+
+    def __init__(
+        self,
+        metric_name,
+        candidate_metric_value,
+        metric_threshold,
+        baseline_metric_value=None,
+    ):
+        self.metric_name = metric_name
+        self.candidate_metric_value = candidate_metric_value
+        self.baseline_metric_value = baseline_metric_value
+        self.metric_threshold = metric_threshold
+
+    def __str__(self):
+        """
+        Returns a human-readable string representing the validation result for the metric.
+        """
+        if self.is_success():
+            return f"Metric {self.metric_name} passed the validation."
+
+        if self.missing_candidate:
+            return (
+                f"Metric validation failed: metric {self.metric_name} was missing from the "
+                f"evaluation result of the candidate model."
+            )
+
+        result_strs = []
+        if self.threshold_failed:
+            result_strs.append(
+                f"Metric {self.metric_name} value threshold check failed: "
+                f"candidate model {self.metric_name} = {self.candidate_metric_value}, "
+                f"{self.metric_name} threshold = {self.metric_threshold.threshold}."
+            )
+        if self.missing_baseline:
+            result_strs.append(
+                f"Model comparison failed: metric {self.metric_name} was missing from "
+                f"the evaluation result of the baseline model."
+            )
+        else:
+            if self.min_absolute_change_failed:
+                result_strs.append(
+                    f"Metric {self.metric_name} minimum absolute change check failed: "
+                    f"candidate model {self.metric_name} = {self.candidate_metric_value}, "
+                    f"baseline model {self.metric_name} = {self.baseline_metric_value}, "
+                    f"{self.metric_name} minimum absolute change threshold = "
+                    f"{self.metric_threshold.min_absolute_change}."
+                )
+            if self.min_relative_change_failed:
+                result_strs.append(
+                    f"Metric {self.metric_name} minimum relative change check failed: "
+                    f"candidate model {self.metric_name} = {self.candidate_metric_value}, "
+                    f"baseline model {self.metric_name} = {self.baseline_metric_value}, "
+                    f"{self.metric_name} minimum relative change threshold = "
+                    f"{self.metric_threshold.min_relative_change}."
+                )
+        return " ".join(result_strs)
+
+    def is_success(self):
+        return (
+            not self.missing_candidate
+            and not self.missing_baseline
+            and not self.threshold_failed
+            and not self.min_absolute_change_failed
+            and not self.min_relative_change_failed
+        )
+
+
+class ModelValidationFailedException(MlflowException):
+    def __init__(self, message, **kwargs):
+        super().__init__(message, error_code=BAD_REQUEST, **kwargs)
+
+
+def validate_evaluation_results(
+    validation_thresholds: dict[str, MetricThreshold],
+    candidate_result: EvaluationResult,
+    baseline_result: Optional[EvaluationResult] = None,
+):
+    """
+    Validate the evaluation result from one model (candidate) against another
+    model (baseline). If the candidate results do not meet the validation
+    thresholds, an ModelValidationFailedException will be raised.
+
+    .. note::
+
+        This API is a replacement for the deprecated model validation
+        functionality in the :py:func:`mlflow.evaluate` API.
+
+    Args:
+        validation_thresholds: A dictionary of metric name to
+            :py:class:`mlflow.models.MetricThreshold` used for model validation.
+            Each metric name must either be the name of a builtin metric or the
+            name of a metric defined in the ``extra_metrics`` parameter.
+        candidate_result: The evaluation result of the candidate model.
+            Returned by the :py:func:`mlflow.evaluate` API.
+        baseline_result: The evaluation result of the baseline model.
+            Returned by the :py:func:`mlflow.evaluate` API.
+            If set to None, the candidate model result will be
+            compared against the threshold values directly.
+
+    Code Example:
+
+        .. code-block:: python
+            :caption: Example of Model Validation
+
+            import mlflow
+            from mlflow.models import MetricThreshold
+
+            thresholds = {
+                "accuracy_score": MetricThreshold(
+                    # accuracy should be >=0.8
+                    threshold=0.8,
+                    # accuracy should be at least 5 percent greater than baseline model accuracy
+                    min_absolute_change=0.05,
+                    # accuracy should be at least 0.05 greater than baseline model accuracy
+                    min_relative_change=0.05,
+                    greater_is_better=True,
+                ),
+            }
+
+            # Get evaluation results for the candidate model
+            candidate_result = mlflow.evaluate(
+                model="<YOUR_CANDIDATE_MODEL_URI>",
+                data=eval_dataset,
+                targets="ground_truth",
+                model_type="classifier",
+            )
+
+            # Get evaluation results for the baseline model
+            baseline_result = mlflow.evaluate(
+                model="<YOUR_BASELINE_MODEL_URI>",
+                data=eval_dataset,
+                targets="ground_truth",
+                model_type="classifier",
+            )
+
+            # Validate the results
+            mlflow.validate_evaluation_results(
+                thresholds,
+                candidate_result,
+                baseline_result,
+            )
+
+        See `the Model Validation documentation
+        <../../models/index.html#performing-model-validation>`_ for more details.
+    """
+    try:
+        assert type(validation_thresholds) is dict
+        for key in validation_thresholds.keys():
+            assert type(key) is str
+        for threshold in validation_thresholds.values():
+            assert isinstance(threshold, MetricThreshold)
+    except AssertionError:
+        raise MlflowException(
+            message="The validation thresholds argument must be a dictionary that maps strings "
+            "to MetricThreshold objects.",
+            error_code=INVALID_PARAMETER_VALUE,
+        )
+
+    _logger.info("Validating candidate model metrics against baseline")
+    _validate(
+        validation_thresholds,
+        candidate_result.metrics,
+        baseline_result.metrics if baseline_result else {},
+    )
+    _logger.info("Model validation passed!")
+
+
+def _validate(
+    validation_thresholds: dict[str, MetricThreshold],
+    candidate_metrics: dict[str, float],
+    baseline_metrics: dict[str, float],
+):
+    """
+    Validate the model based on validation_thresholds by metrics value and
+    metrics comparison between candidate model's metrics (candidate_metrics) and
+    baseline model's metrics (baseline_metrics).
+
+    Args:
+        validation_thresholds: A dictionary from metric_name to MetricThreshold.
+        candidate_metrics: The metric evaluation result of the candidate model.
+        baseline_metrics: The metric evaluation result of the baseline model.
+
+    Raises:
+        If the validation does not pass, raise an MlflowException with detail failure message.
+    """
+    validation_results = {
+        metric_name: _MetricValidationResult(
+            metric_name,
+            candidate_metrics.get(metric_name),
+            threshold,
+            baseline_metrics.get(metric_name),
+        )
+        for (metric_name, threshold) in validation_thresholds.items()
+    }
+
+    for metric_name, metric_threshold in validation_thresholds.items():
+        validation_result = validation_results[metric_name]
+
+        if metric_name not in candidate_metrics:
+            validation_result.missing_candidate = True
+            continue
+
+        candidate_metric_value = candidate_metrics[metric_name]
+        baseline_metric_value = baseline_metrics[metric_name] if baseline_metrics else None
+
+        # If metric is higher is better, >= is used, otherwise <= is used
+        # for thresholding metric value and model comparison
+        comparator_fn = operator.__ge__ if metric_threshold.greater_is_better else operator.__le__
+        operator_fn = operator.add if metric_threshold.greater_is_better else operator.sub
+
+        if metric_threshold.threshold is not None:
+            # metric threshold fails
+            # - if not (metric_value >= threshold) for higher is better
+            # - if not (metric_value <= threshold) for lower is better
+            validation_result.threshold_failed = not comparator_fn(
+                candidate_metric_value, metric_threshold.threshold
+            )
+
+        if (
+            metric_threshold.min_relative_change or metric_threshold.min_absolute_change
+        ) and metric_name not in baseline_metrics:
+            validation_result.missing_baseline = True
+            continue
+
+        if metric_threshold.min_absolute_change is not None:
+            # metric comparison absolute change fails
+            # - if not (metric_value >= baseline + min_absolute_change) for higher is better
+            # - if not (metric_value <= baseline - min_absolute_change) for lower is better
+            validation_result.min_absolute_change_failed = not comparator_fn(
+                Decimal(candidate_metric_value),
+                Decimal(operator_fn(baseline_metric_value, metric_threshold.min_absolute_change)),
+            )
+
+        if metric_threshold.min_relative_change is not None:
+            # If baseline metric value equals 0, fallback to simple comparison check
+            if baseline_metric_value == 0:
+                _logger.warning(
+                    f"Cannot perform relative model comparison for metric {metric_name} as "
+                    "baseline metric value is 0. Falling back to simple comparison: verifying "
+                    "that candidate metric value is better than the baseline metric value."
+                )
+                validation_result.min_relative_change_failed = not comparator_fn(
+                    Decimal(candidate_metric_value),
+                    Decimal(operator_fn(baseline_metric_value, 1e-10)),
+                )
+                continue
+            # metric comparison relative change fails
+            # - if (metric_value - baseline) / baseline < min_relative_change for higher is better
+            # - if (baseline - metric_value) / baseline < min_relative_change for lower is better
+            if metric_threshold.greater_is_better:
+                relative_change = (
+                    candidate_metric_value - baseline_metric_value
+                ) / baseline_metric_value
+            else:
+                relative_change = (
+                    baseline_metric_value - candidate_metric_value
+                ) / baseline_metric_value
+            validation_result.min_relative_change_failed = (
+                relative_change < metric_threshold.min_relative_change
+            )
+
+    failure_messages = []
+
+    for metric_validation_result in validation_results.values():
+        if metric_validation_result.is_success():
+            continue
+        failure_messages.append(str(metric_validation_result))
+
+    if not failure_messages:
+        return
+
+    raise ModelValidationFailedException(message=os.linesep.join(failure_messages))

venv/lib/python3.9/site-packages/mlflow/models/flavor_backend.py

@@ -0,0 +1,93 @@
+from abc import ABCMeta, abstractmethod
+
+from mlflow.utils.annotations import developer_stable
+
+
+@developer_stable
+class FlavorBackend:
+    """
+    Abstract class for Flavor Backend.
+    This class defines the API interface for local model deployment of MLflow model flavors.
+    """
+
+    __metaclass__ = ABCMeta
+
+    def __init__(self, config, **kwargs):
+        self._config = config
+
+    @abstractmethod
+    def predict(self, model_uri, input_path, output_path, content_type):
+        """
+        Generate predictions using a saved MLflow model referenced by the given URI.
+        Input and output are read from and written to a file or stdin / stdout.
+
+        Args:
+            model_uri: URI pointing to the MLflow model to be used for scoring.
+            input_path: Path to the file with input data. If not specified, data is read from
+                        stdin.
+            output_path: Path to the file with output predictions. If not specified, data is
+                         written to stdout.
+            content_type: Specifies the input format. Can be one of {``json``, ``csv``}
+        """
+
+    @abstractmethod
+    def serve(
+        self,
+        model_uri,
+        port,
+        host,
+        timeout,
+        enable_mlserver,
+        synchronous=True,
+        stdout=None,
+        stderr=None,
+    ):
+        """
+        Serve the specified MLflow model locally.
+
+        Args:
+            model_uri: URI pointing to the MLflow model to be used for scoring.
+            port: Port to use for the model deployment.
+            host: Host to use for the model deployment. Defaults to ``localhost``.
+            timeout: Timeout in seconds to serve a request. Defaults to 60.
+            enable_mlserver: Whether to use MLServer or the local scoring server.
+            synchronous: If True, wait until server process exit and return 0, if process exit
+                with non-zero return code, raise exception.
+                If False, return the server process `Popen` instance immediately.
+            stdout: Redirect server stdout
+            stderr: Redirect server stderr
+        """
+
+    def prepare_env(self, model_uri, capture_output=False):
+        """
+        Performs any preparation necessary to predict or serve the model, for example
+        downloading dependencies or initializing a conda environment. After preparation,
+        calling predict or serve should be fast.
+        """
+
+    @abstractmethod
+    def build_image(
+        self, model_uri, image_name, install_mlflow, mlflow_home, enable_mlserver, base_image=None
+    ): ...
+
+    @abstractmethod
+    def generate_dockerfile(
+        self, model_uri, output_path, install_mlflow, mlflow_home, enable_mlserver, base_image=None
+    ): ...
+
+    @abstractmethod
+    def can_score_model(self):
+        """
+        Check whether this flavor backend can be deployed in the current environment.
+
+        Returns:
+            True if this flavor backend can be applied in the current environment.
+        """
+
+    def can_build_image(self):
+        """
+        Returns:
+            True if this flavor has a `build_image` method defined for building a docker
+            container capable of serving the model, False otherwise.
+        """
+        return callable(getattr(self.__class__, "build_image", None))

venv/lib/python3.9/site-packages/mlflow/models/flavor_backend_registry.py

@@ -0,0 +1,53 @@
+"""
+Registry of supported flavor backends. Contains a mapping of flavors to flavor backends. This
+mapping is used to select suitable flavor when deploying generic MLflow models.
+
+Flavor backend can deploy particular flavor locally to generate predictions, deploy as a local
+REST api endpoint, or build a docker image for serving the model locally or remotely.
+Not all flavors have a flavor backend.
+"""
+
+import logging
+
+from mlflow.models.model import Model
+from mlflow.store.artifact.artifact_repository_registry import get_artifact_repository
+from mlflow.tracking.artifact_utils import (
+    _get_root_uri_and_artifact_path,
+)
+from mlflow.utils.file_utils import TempDir
+
+_logger = logging.getLogger(__name__)
+
+
+def _get_flavor_backend_for_local_model(model=None, build_docker=True, **kwargs):
+    from mlflow import pyfunc, rfunc
+    from mlflow.pyfunc.backend import PyFuncBackend
+    from mlflow.rfunc.backend import RFuncBackend
+
+    if not model:
+        return pyfunc.FLAVOR_NAME, PyFuncBackend({}, **kwargs)
+
+    backends = {pyfunc.FLAVOR_NAME: PyFuncBackend, rfunc.FLAVOR_NAME: RFuncBackend}
+    for flavor, Backend in backends.items():
+        if flavor in model.flavors:
+            backend = Backend(model.flavors[flavor], **kwargs)
+            if (build_docker and backend.can_build_image()) or backend.can_score_model():
+                return flavor, backend
+
+    return None, None
+
+
+def get_flavor_backend(model_uri, **kwargs):
+    if model_uri:
+        with TempDir() as tmp:
+            root_uri, artifact_path = _get_root_uri_and_artifact_path(model_uri)
+            artifact_repo = get_artifact_repository(root_uri)
+            local_path = artifact_repo.download_artifacts(artifact_path, dst_path=tmp.path())
+            model = Model.load(local_path)
+    else:
+        model = None
+    flavor_name, flavor_backend = _get_flavor_backend_for_local_model(model, **kwargs)
+    if flavor_backend is None:
+        raise Exception("No suitable flavor backend was found for the model.")
+    _logger.info("Selected backend for flavor '%s'", flavor_name)
+    return flavor_backend

venv/lib/python3.9/site-packages/mlflow/models/model_config.py

@@ -0,0 +1,150 @@
+import os
+from typing import Any, Optional, Union
+
+import yaml
+
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+
+__mlflow_model_config__ = None
+
+
+class ModelConfig:
+    """
+    ModelConfig used in code to read a YAML configuration file or a dictionary.
+
+    Args:
+        development_config: Path to the YAML configuration file or a dictionary containing the
+                        configuration. If the configuration is not provided, an error is raised
+
+    .. code-block:: python
+        :caption: Example usage in model code
+
+        from mlflow.models import ModelConfig
+
+        # Load the configuration from a dictionary
+        config = ModelConfig(development_config={"key1": "value1"})
+        print(config.get("key1"))
+
+
+    .. code-block:: yaml
+        :caption: yaml file for model configuration
+
+        key1: value1
+        another_key:
+            - value2
+            - value3
+
+    .. code-block:: python
+        :caption: Example yaml usage in model code
+
+        from mlflow.models import ModelConfig
+
+        # Load the configuration from a file
+        config = ModelConfig(development_config="config.yaml")
+        print(config.get("key1"))
+
+
+    When invoking the ModelConfig locally in a model file, development_config can be passed in
+    which would be used as configuration for the model.
+
+
+    .. code-block:: python
+        :caption: Example to use ModelConfig when logging model as code: agent.py
+
+        import mlflow
+        from mlflow.models import ModelConfig
+
+        config = ModelConfig(development_config={"key1": "value1"})
+
+
+        class TestModel(mlflow.pyfunc.PythonModel):
+            def predict(self, context, model_input, params=None):
+                return config.get("key1")
+
+
+        mlflow.models.set_model(TestModel())
+
+
+    But this development_config configuration file will be overridden when logging a model.
+    When no model_config is passed in while logging the model, an error will be raised when
+    trying to load the model using ModelConfig.
+    Note: development_config is not used when logging the model.
+
+
+    .. code-block:: python
+        :caption: Example to use agent.py to log the model: deploy.py
+
+        model_config = {"key1": "value2"}
+        with mlflow.start_run():
+            model_info = mlflow.pyfunc.log_model(
+                artifact_path="model", python_model="agent.py", model_config=model_config
+            )
+
+        loaded_model = mlflow.pyfunc.load_model(model_info.model_uri)
+
+        # This will print "value2" as the model_config passed in while logging the model
+        print(loaded_model.predict(None))
+    """
+
+    def __init__(self, *, development_config: Optional[Union[str, dict[str, Any]]] = None):
+        config = globals().get("__mlflow_model_config__", None)
+        # Here mlflow_model_config have 3 states:
+        # 1. None, this means if the mlflow_model_config is None, use development_config if
+        # available
+        # 2. "", Empty string, this means the users explicitly didn't set the model config
+        # while logging the model so if ModelConfig is used, it should throw an error
+        # 3. A valid path, this means the users have set the model config while logging the
+        # model so use that path
+        if config is not None:
+            self.config = config
+        else:
+            self.config = development_config
+
+        if not self.config:
+            raise FileNotFoundError(
+                "Config file is not provided which is needed to load the model. "
+                "Please provide a valid path."
+            )
+
+        if not isinstance(self.config, dict) and not os.path.isfile(self.config):
+            raise FileNotFoundError(f"Config file '{self.config}' not found.")
+
+    def _read_config(self):
+        """Reads the YAML configuration file and returns its contents.
+
+        Raises:
+            FileNotFoundError: If the configuration file does not exist.
+            yaml.YAMLError: If there is an error parsing the YAML content.
+
+        Returns:
+            dict or None: The content of the YAML file as a dictionary, or None if the
+            config path is not set.
+        """
+        if isinstance(self.config, dict):
+            return self.config
+
+        with open(self.config) as file:
+            try:
+                return yaml.safe_load(file)
+            except yaml.YAMLError as e:
+                raise MlflowException(
+                    f"Error parsing YAML file: {e}", error_code=INVALID_PARAMETER_VALUE
+                )
+
+    def to_dict(self):
+        """Returns the configuration as a dictionary."""
+        return self._read_config()
+
+    def get(self, key):
+        """Gets the value of a top-level parameter in the configuration."""
+        config_data = self._read_config()
+
+        if config_data and key in config_data:
+            return config_data[key]
+        else:
+            raise KeyError(f"Key '{key}' not found in configuration: {config_data}.")
+
+
+def _set_model_config(model_config):
+    globals()["__mlflow_model_config__"] = model_config

venv/lib/python3.9/site-packages/mlflow/models/notebook_resources/agent_evaluation_template.html

@@ -0,0 +1,235 @@
+<head>
+  <link
+    rel="stylesheet"
+    href="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/styles/xcode.min.css"
+  />
+  <script src="https://cdnjs.cloudflare.com/ajax/libs/highlight.js/11.9.0/highlight.min.js"></script>
+  <script>
+    hljs.highlightAll();
+  </script>
+  <style>
+    body {
+      margin: 0;
+      font-family: -apple-system, BlinkMacSystemFont, Segoe UI, Roboto,
+        Helvetica Neue, Arial, Noto Sans, sans-serif, Apple Color Emoji,
+        Segoe UI Emoji, Segoe UI Symbol, Noto Color Emoji;
+      -webkit-tap-highlight-color: rgba(0, 0, 0, 0);
+      margin: 0;
+      font-weight: 400;
+      font-size: 13px;
+      line-height: 18px;
+      color: rgb(17, 23, 28);
+    }
+    code {
+      line-height: 18px;
+      font-size: 11px;
+      background: rgb(250, 250, 250) !important;
+    }
+    pre {
+      background: rgb(250, 250, 250);
+      margin: 0;
+      display: none;
+    }
+    pre.active {
+      display: unset;
+    }
+    button {
+      white-space: nowrap;
+      text-align: center;
+      position: relative;
+      cursor: pointer;
+      background: rgba(34, 114, 180, 0) !important;
+      color: rgb(34, 114, 180) !important;
+      border-color: rgba(34, 114, 180, 0) !important;
+      padding: 4px 6px !important;
+      text-decoration: none !important;
+      line-height: 20px !important;
+      box-shadow: none !important;
+      height: 32px !important;
+      display: inline-flex !important;
+      -webkit-box-align: center !important;
+      align-items: center !important;
+      -webkit-box-pack: center !important;
+      justify-content: center !important;
+      vertical-align: middle !important;
+    }
+    p {
+      margin: 0;
+      padding: 0;
+    }
+    button:hover {
+      background: rgba(34, 114, 180, 0.08) !important;
+      color: rgb(14, 83, 139) !important;
+    }
+    button:active {
+      background: rgba(34, 114, 180, 0.16) !important;
+      color: rgb(4, 53, 93) !important;
+    }
+    h1 {
+      margin-top: 4px;
+      font-size: 22px;
+    }
+    .info {
+      font-size: 12px;
+      font-weight: 500;
+      line-height: 16px;
+      color: rgb(95, 114, 129);
+    }
+    .tabs {
+      margin-top: 10px;
+      border-bottom: 1px solid rgb(209, 217, 225) !important;
+      display: flex;
+      line-height: 24px;
+    }
+    .tab {
+      font-size: 13px;
+      font-weight: 600 !important;
+      cursor: pointer;
+      margin: 0 24px 0 2px;
+      padding-left: 2px;
+    }
+    .tab:hover {
+      color: rgb(14, 83, 139) !important;
+    }
+    .tab.active {
+      border-bottom: 3px solid rgb(34, 114, 180) !important;
+    }
+    .link {
+      margin-left: 12px;
+      display: inline-block;
+      text-decoration: none;
+      color: rgb(34, 114, 180) !important;
+      font-size: 13px;
+      font-weight: 400;
+    }
+    .link:hover {
+      color: rgb(14, 83, 139) !important;
+    }
+    .link-content {
+      display: flex;
+      gap: 6px;
+      align-items: center;
+    }
+    .caret-up {
+      transform: rotate(180deg);
+    }
+  </style>
+</head>
+<body>
+  <div style="display: flex; align-items: center">
+    The logged model is compatible with the Mosaic AI Agent Framework.
+    <button onclick="toggleCode()">
+      See how to evaluate the model&nbsp;
+      <span
+        role="img"
+        id="caret"
+        aria-hidden="true"
+        class="anticon css-6xix1i"
+        style="font-size: 14px"
+        ><svg
+          xmlns="http://www.w3.org/2000/svg"
+          width="1em"
+          height="1em"
+          fill="none"
+          viewBox="0 0 16 16"
+          aria-hidden="true"
+          focusable="false"
+          class=""
+        >
+          <path
+            fill="currentColor"
+            fill-rule="evenodd"
+            d="M8 8.917 10.947 6 12 7.042 8 11 4 7.042 5.053 6z"
+            clip-rule="evenodd"
+          ></path>
+        </svg>
+      </span>
+    </button>
+  </div>
+  <div id="code" style="display: none">
+    <h1>
+      Agent evaluation
+      <a
+        class="link"
+        href="https://docs.databricks.com/en/generative-ai/agent-evaluation/synthesize-evaluation-set.html?utm_source=mlflow.log_model&utm_medium=notebook"
+        target="_blank"
+      >
+        <span class="link-content">
+          Learn more
+          <span role="img" aria-hidden="true" class="anticon css-6xix1i"
+            ><svg
+              xmlns="http://www.w3.org/2000/svg"
+              width="1em"
+              height="1em"
+              fill="none"
+              viewBox="0 0 16 16"
+              aria-hidden="true"
+              focusable="false"
+              class=""
+            >
+              <path
+                fill="currentColor"
+                d="M10 1h5v5h-1.5V3.56L8.53 8.53 7.47 7.47l4.97-4.97H10z"
+              ></path>
+              <path
+                fill="currentColor"
+                d="M1 2.75A.75.75 0 0 1 1.75 2H8v1.5H2.5v10h10V8H14v6.25a.75.75 0 0 1-.75.75H1.75a.75.75 0 0 1-.75-.75z"
+              ></path></svg></span></span
+      ></a>
+    </h1>
+    <p class="info">
+      Copy the following code snippet in a notebook cell (right click → copy)
+    </p>
+    <div class="tabs">
+      <div class="tab active" onclick="tabClicked(0)">Using synthetic data</div>
+      <div class="tab" onclick="tabClicked(1)">Using your own dataset</div>
+    </div>
+    <div style="height: 472px">
+      <pre
+        class="active"
+      ><code class="language-python">{{eval_with_synthetic_code}}</code></pre>
+
+      <pre><code class="language-python">{{eval_with_dataset_code}}</code></pre>
+    </div>
+  </div>
+  <script>
+    var codeShown = false;
+    function clip(el) {
+      var range = document.createRange();
+      range.selectNodeContents(el);
+      var sel = window.getSelection();
+      sel.removeAllRanges();
+      sel.addRange(range);
+    }
+
+    function toggleCode() {
+      if (codeShown) {
+        document.getElementById("code").style.display = "none";
+        codeShown = false;
+      } else {
+        document.getElementById("code").style.display = "block";
+        clip(document.querySelector("pre.active"));
+        codeShown = true;
+      }
+      document.getElementById("caret").classList.toggle("caret-up");
+    }
+
+    function tabClicked(tabIndex) {
+      document.querySelectorAll(".tab").forEach((tab, index) => {
+        if (index === tabIndex) {
+          tab.classList.add("active");
+        } else {
+          tab.classList.remove("active");
+        }
+      });
+      document.querySelectorAll("pre").forEach((pre, index) => {
+        if (index === tabIndex) {
+          pre.classList.add("active");
+        } else {
+          pre.classList.remove("active");
+        }
+      });
+      clip(document.querySelector("pre.active"));
+    }
+  </script>
+</body>

venv/lib/python3.9/site-packages/mlflow/models/notebook_resources/eval_with_dataset_example.py

@@ -0,0 +1,22 @@
+# ruff: noqa: F821, I001
+{{pipInstall}}
+
+import pandas as pd
+import mlflow
+
+evals = [
+    {
+        "request": {
+            "messages": [
+                {"role": "user", "content": "How do I convert a Spark DataFrame to Pandas?"}
+            ],
+        },
+        # Optional, needed for judging correctness.
+        "expected_facts": [
+            "To convert a Spark DataFrame to Pandas, you can use the toPandas() method."
+        ],
+    }
+]
+eval_result = mlflow.evaluate(
+    data=pd.DataFrame.from_records(evals), model="{{modelUri}}", model_type="databricks-agent"
+)

venv/lib/python3.9/site-packages/mlflow/models/notebook_resources/eval_with_synthetic_example.py

@@ -0,0 +1,22 @@
+# ruff: noqa: F821, I001
+{{pipInstall}}
+
+from databricks.agents.evals import generate_evals_df
+import mlflow
+
+agent_description = "A chatbot that answers questions about Databricks."
+question_guidelines = """
+# User personas
+- A developer new to the Databricks platform
+# Example questions
+- What API lets me parallelize operations over rows of a delta table?
+"""
+# TODO: Spark/Pandas DataFrame with "content" and "doc_uri" columns.
+docs = spark.table("catalog.schema.my_table_of_docs")
+evals = generate_evals_df(
+    docs=docs,
+    num_evals=25,
+    agent_description=agent_description,
+    question_guidelines=question_guidelines,
+)
+eval_result = mlflow.evaluate(data=evals, model="{{modelUri}}", model_type="databricks-agent")

venv/lib/python3.9/site-packages/mlflow/models/python_api.py

@@ -0,0 +1,361 @@
+import logging
+import os
+import shutil
+from io import StringIO
+from typing import ForwardRef, get_args, get_origin
+
+from mlflow.exceptions import MlflowException
+from mlflow.models.flavor_backend_registry import get_flavor_backend
+from mlflow.utils import env_manager as _EnvManager
+from mlflow.utils.annotations import experimental
+from mlflow.utils.databricks_utils import is_databricks_connect
+from mlflow.utils.file_utils import TempDir
+
+_logger = logging.getLogger(__name__)
+UV_INSTALLATION_INSTRUCTIONS = (
+    "Run `pip install uv` to install uv. See "
+    "https://docs.astral.sh/uv/getting-started/installation for other installation methods."
+)
+
+
+def build_docker(
+    model_uri=None,
+    name="mlflow-pyfunc",
+    env_manager=_EnvManager.VIRTUALENV,
+    mlflow_home=None,
+    install_java=False,
+    install_mlflow=False,
+    enable_mlserver=False,
+    base_image=None,
+):
+    """
+    Builds a Docker image whose default entrypoint serves an MLflow model at port 8080, using the
+    python_function flavor. The container serves the model referenced by ``model_uri``, if
+    specified. If ``model_uri`` is not specified, an MLflow Model directory must be mounted as a
+    volume into the /opt/ml/model directory in the container.
+
+    .. important::
+
+        Since MLflow 2.10.1, the Docker image built with ``--model-uri`` does **not install Java**
+        for improved performance, unless the model flavor is one of ``["johnsnowlabs", "h2o",
+        "mleap", "spark"]``. If you need to install Java for other flavors, e.g. custom Python model
+        that uses SparkML, please specify ``install-java=True`` to enforce Java installation.
+        For earlier versions, Java is always installed to the image.
+
+
+    .. warning::
+
+        If ``model_uri`` is unspecified, the resulting image doesn't support serving models with
+        the RFunc or Java MLeap model servers.
+
+    NB: by default, the container will start nginx and gunicorn processes. If you don't need the
+    nginx process to be started (for instance if you deploy your container to Google Cloud Run),
+    you can disable it via the DISABLE_NGINX environment variable:
+
+    .. code:: bash
+
+        docker run -p 5001:8080 -e DISABLE_NGINX=true "my-image-name"
+
+    See https://www.mlflow.org/docs/latest/python_api/mlflow.pyfunc.html for more information on the
+    'python_function' flavor.
+
+    Args:
+        model_uri: URI to the model. A local path, a 'runs:/' URI, or a remote storage URI (e.g.,
+            an 's3://' URI). For more information about supported remote URIs for model artifacts,
+            see https://mlflow.org/docs/latest/tracking.html#artifact-stores
+        name: Name of the Docker image to build. Defaults to 'mlflow-pyfunc'.
+        env_manager: If specified, create an environment for MLmodel using the specified environment
+            manager. The following values are supported: (1) virtualenv (default): use virtualenv
+            and pyenv for Python version management (2) conda: use conda (3) local: use the local
+            environment without creating a new one.
+        mlflow_home: Path to local clone of MLflow project. Use for development only.
+        install_java: If specified, install Java in the image. Default is False in order to
+            reduce both the image size and the build time. Model flavors requiring Java will enable
+            this setting automatically, such as the Spark flavor. (This argument is only available
+            in MLflow 2.10.1 and later. In earlier versions, Java is always installed to the image.)
+        install_mlflow: If specified and there is a conda or virtualenv environment to be activated
+            mlflow will be installed into the environment after it has been activated.
+            The version of installed mlflow will be the same as the one used to invoke this command.
+        enable_mlserver: If specified, the image will be built with the Seldon MLserver as backend.
+        base_image: Base image for the Docker image. If not specified, the default image is either
+            UBUNTU_BASE_IMAGE = "ubuntu:20.04" or PYTHON_SLIM_BASE_IMAGE = "python:{version}-slim"
+            Note: If custom image is used, there are no guarantees that the image will work. You
+            may find greater compatibility by building your image on top of the ubuntu images. In
+            addition, you must install Java and virtualenv to have the image work properly.
+    """
+    get_flavor_backend(model_uri, docker_build=True, env_manager=env_manager).build_image(
+        model_uri,
+        name,
+        mlflow_home=mlflow_home,
+        install_java=install_java,
+        install_mlflow=install_mlflow,
+        enable_mlserver=enable_mlserver,
+        base_image=base_image,
+    )
+
+
+_CONTENT_TYPE_CSV = "csv"
+_CONTENT_TYPE_JSON = "json"
+
+
+@experimental
+def predict(
+    model_uri,
+    input_data=None,
+    input_path=None,
+    content_type=_CONTENT_TYPE_JSON,
+    output_path=None,
+    env_manager=_EnvManager.VIRTUALENV,
+    install_mlflow=False,
+    pip_requirements_override=None,
+    extra_envs=None,
+    # TODO: add an option to force recreating the env
+):
+    """
+    Generate predictions in json format using a saved MLflow model. For information about the input
+    data formats accepted by this function, see the following documentation:
+    https://www.mlflow.org/docs/latest/models.html#built-in-deployment-tools.
+
+    .. note::
+
+        To increase verbosity for debugging purposes (in order to inspect the full dependency
+        resolver operations when processing transient dependencies), consider setting the following
+        environment variables:
+
+        .. code-block:: bash
+
+            # For virtualenv
+            export PIP_VERBOSE=1
+
+            # For uv
+            export RUST_LOG=uv=debug
+
+        See also:
+
+        - https://pip.pypa.io/en/stable/topics/configuration/#environment-variables
+        - https://docs.astral.sh/uv/configuration/environment
+
+    Args:
+        model_uri: URI to the model. A local path, a local or remote URI e.g. runs:/, s3://.
+        input_data: Input data for prediction. Must be valid input for the PyFunc model. Refer
+            to the :py:func:`mlflow.pyfunc.PyFuncModel.predict()` for the supported input types.
+
+            .. note::
+                If this API fails due to errors in input_data, use
+                `mlflow.models.convert_input_example_to_serving_input` to manually validate
+                your input data.
+        input_path: Path to a file containing input data. If provided, 'input_data' must be None.
+        content_type: Content type of the input data. Can be one of {‘json’, ‘csv’}.
+        output_path: File to output results to as json. If not provided, output to stdout.
+        env_manager: Specify a way to create an environment for MLmodel inference:
+
+            - "virtualenv" (default): use virtualenv (and pyenv for Python version management)
+            - "uv": use uv
+            - "local": use the local environment
+            - "conda": use conda
+
+        install_mlflow: If specified and there is a conda or virtualenv environment to be activated
+            mlflow will be installed into the environment after it has been activated. The version
+            of installed mlflow will be the same as the one used to invoke this command.
+        pip_requirements_override: If specified, install the specified python dependencies to the
+            model inference environment. This is particularly useful when you want to add extra
+            dependencies or try different versions of the dependencies defined in the logged model.
+
+            .. tip::
+                After validating the pip requirements override works as expected, you can update
+                the logged model's dependency using `mlflow.models.update_model_requirements` API
+                without re-logging it. Note that a registered model is immutable, so you need to
+                register a new model version with the updated model.
+        extra_envs: If specified, a dictionary of extra environment variables will be passed to the
+            model inference environment. This is useful for testing what environment variables are
+            needed for the model to run correctly. By default, environment variables existing in the
+            current os.environ are passed, and this parameter can be used to override them.
+
+            .. note::
+                This parameter is only supported when `env_manager` is set to "virtualenv",
+                "conda" or "uv".
+
+    Code example:
+
+    .. code-block:: python
+
+        import mlflow
+
+        run_id = "..."
+
+        mlflow.models.predict(
+            model_uri=f"runs:/{run_id}/model",
+            input_data={"x": 1, "y": 2},
+            content_type="json",
+        )
+
+        # Run prediction with "uv" as the environment manager
+        mlflow.models.predict(
+            model_uri=f"runs:/{run_id}/model",
+            input_data={"x": 1, "y": 2},
+            env_manager="uv",
+        )
+
+        # Run prediction with additional pip dependencies and extra environment variables
+        mlflow.models.predict(
+            model_uri=f"runs:/{run_id}/model",
+            input_data={"x": 1, "y": 2},
+            content_type="json",
+            pip_requirements_override=["scikit-learn==0.23.2"],
+            extra_envs={"OPENAI_API_KEY": "some_value"},
+        )
+
+    """
+    # to avoid circular imports
+    from mlflow.pyfunc import _PREBUILD_ENV_ROOT_LOCATION
+
+    if content_type not in [_CONTENT_TYPE_JSON, _CONTENT_TYPE_CSV]:
+        raise MlflowException.invalid_parameter_value(
+            f"Content type must be one of {_CONTENT_TYPE_JSON} or {_CONTENT_TYPE_CSV}."
+        )
+    if extra_envs and env_manager not in (
+        _EnvManager.VIRTUALENV,
+        _EnvManager.CONDA,
+        _EnvManager.UV,
+    ):
+        raise MlflowException.invalid_parameter_value(
+            "Extra environment variables are only supported when env_manager is "
+            f"set to '{_EnvManager.VIRTUALENV}', '{_EnvManager.CONDA}' or '{_EnvManager.UV}'."
+        )
+    if env_manager == _EnvManager.UV:
+        if not shutil.which("uv"):
+            raise MlflowException(
+                f"Found '{env_manager}' as env_manager, but the 'uv' command is not found in the "
+                f"PATH. {UV_INSTALLATION_INSTRUCTIONS} Alternatively, you can use 'virtualenv' or "
+                "'conda' as the environment manager, but note their performances are not "
+                "as good as 'uv'."
+            )
+    else:
+        _logger.info(
+            f"It is highly recommended to use `{_EnvManager.UV}` as the environment manager for "
+            "predicting with MLflow models as its performance is significantly better than other "
+            f"environment managers. {UV_INSTALLATION_INSTRUCTIONS}"
+        )
+
+    is_dbconnect_mode = is_databricks_connect()
+    if is_dbconnect_mode:
+        if env_manager not in (_EnvManager.VIRTUALENV, _EnvManager.UV):
+            raise MlflowException(
+                f"Databricks Connect only supports '{_EnvManager.VIRTUALENV}' or '{_EnvManager.UV}'"
+                f" as the environment manager. Got {env_manager}."
+            )
+        pyfunc_backend_env_root_config = {
+            "create_env_root_dir": False,
+            "env_root_dir": _PREBUILD_ENV_ROOT_LOCATION,
+        }
+    else:
+        pyfunc_backend_env_root_config = {"create_env_root_dir": True}
+
+    def _predict(_input_path: str):
+        return get_flavor_backend(
+            model_uri,
+            env_manager=env_manager,
+            install_mlflow=install_mlflow,
+            **pyfunc_backend_env_root_config,
+        ).predict(
+            model_uri=model_uri,
+            input_path=_input_path,
+            output_path=output_path,
+            content_type=content_type,
+            pip_requirements_override=pip_requirements_override,
+            extra_envs=extra_envs,
+        )
+
+    if input_data is not None and input_path is not None:
+        raise MlflowException.invalid_parameter_value(
+            "Both input_data and input_path are provided. Only one of them should be specified."
+        )
+    elif input_data is not None:
+        input_data = _serialize_input_data(input_data, content_type)
+
+        # Write input data to a temporary file
+        with TempDir() as tmp:
+            input_path = os.path.join(tmp.path(), f"input.{content_type}")
+            with open(input_path, "w") as f:
+                f.write(input_data)
+
+            _predict(input_path)
+    else:
+        _predict(input_path)
+
+
+def _get_pyfunc_supported_input_types():
+    # Importing here as the util module depends on optional packages not available in mlflow-skinny
+    import mlflow.models.utils as base_module
+
+    supported_input_types = []
+    for input_type in get_args(base_module.PyFuncInput):
+        if isinstance(input_type, type):
+            supported_input_types.append(input_type)
+        elif isinstance(input_type, ForwardRef):
+            name = input_type.__forward_arg__
+            if hasattr(base_module, name):
+                cls = getattr(base_module, name)
+                supported_input_types.append(cls)
+        else:
+            # typing instances like List, Dict, Tuple, etc.
+            supported_input_types.append(get_origin(input_type))
+    return tuple(supported_input_types)
+
+
+def _serialize_input_data(input_data, content_type):
+    # build-docker command is available in mlflow-skinny (which doesn't contain pandas)
+    # so we shouldn't import pandas at the top level
+    import pandas as pd
+
+    # this introduces numpy as dependency, we shouldn't import it at the top level
+    # as it is not available in mlflow-skinny
+    from mlflow.models.utils import convert_input_example_to_serving_input
+
+    valid_input_types = {
+        _CONTENT_TYPE_CSV: (str, list, dict, pd.DataFrame),
+        _CONTENT_TYPE_JSON: _get_pyfunc_supported_input_types(),
+    }.get(content_type)
+
+    if not isinstance(input_data, valid_input_types):
+        raise MlflowException.invalid_parameter_value(
+            f"Input data must be one of {valid_input_types} when content type is '{content_type}', "
+            f"but got {type(input_data)}."
+        )
+
+    if content_type == _CONTENT_TYPE_CSV:
+        if isinstance(input_data, str):
+            _validate_csv_string(input_data)
+            return input_data
+        else:
+            try:
+                return pd.DataFrame(input_data).to_csv(index=False)
+            except Exception as e:
+                raise MlflowException.invalid_parameter_value(
+                    "Failed to serialize input data to CSV format."
+                ) from e
+
+    try:
+        # rely on convert_input_example_to_serving_input to validate
+        # the input_data is valid type for the loaded pyfunc model
+        return convert_input_example_to_serving_input(input_data)
+    except Exception as e:
+        raise MlflowException.invalid_parameter_value(
+            "Invalid input data, please make sure the data is acceptable by the "
+            "loaded pyfunc model. Use `mlflow.models.convert_input_example_to_serving_input` "
+            "to manually validate your input data."
+        ) from e
+
+
+def _validate_csv_string(input_data: str):
+    """
+    Validate the string must be the path to a CSV file.
+    """
+    try:
+        import pandas as pd
+
+        pd.read_csv(StringIO(input_data))
+    except Exception as e:
+        raise MlflowException.invalid_parameter_value(
+            message="Failed to deserialize input string data to Pandas DataFrame."
+        ) from e

venv/lib/python3.9/site-packages/mlflow/models/rag_signatures.py

@@ -0,0 +1,128 @@
+from dataclasses import dataclass, field
+from typing import Optional
+
+from mlflow.models import ModelSignature
+from mlflow.types.schema import (
+    Array,
+    ColSpec,
+    DataType,
+    Object,
+    Property,
+    Schema,
+)
+from mlflow.utils.annotations import deprecated
+
+
+@deprecated("mlflow.types.llm.ChatMessage")
+@dataclass
+class Message:
+    role: str = "user"  # "system", "user", or "assistant"
+    content: str = "What is mlflow?"
+
+
+@deprecated("mlflow.types.llm.ChatCompletionRequest")
+@dataclass
+class ChatCompletionRequest:
+    messages: list[Message] = field(default_factory=lambda: [Message()])
+
+
+@deprecated("mlflow.types.llm.ChatCompletionRequest")
+@dataclass
+class SplitChatMessagesRequest:
+    query: str = "What is mlflow?"
+    history: Optional[list[Message]] = field(default_factory=list)
+
+
+@deprecated("mlflow.types.llm.ChatCompletionRequest")
+@dataclass
+class MultiturnChatRequest:
+    query: str = "What is mlflow?"
+    history: Optional[list[Message]] = field(default_factory=list)
+
+
+@deprecated("mlflow.types.llm.ChatChoice")
+@dataclass
+class ChainCompletionChoice:
+    index: int = 0
+    message: Message = field(
+        default_factory=lambda: Message(
+            role="assistant",
+            content="MLflow is an open source platform for the machine learning lifecycle.",
+        )
+    )
+    finish_reason: str = "stop"
+
+
+@deprecated("mlflow.types.llm.ChatCompletionChunk")
+@dataclass
+class ChainCompletionChunk:
+    index: int = 0
+    delta: Message = field(
+        default_factory=lambda: Message(
+            role="assistant",
+            content="MLflow is an open source platform for the machine learning lifecycle.",
+        )
+    )
+    finish_reason: str = "stop"
+
+
+@deprecated("mlflow.types.llm.ChatCompletionResponse")
+@dataclass
+class ChatCompletionResponse:
+    choices: list[ChainCompletionChoice] = field(default_factory=lambda: [ChainCompletionChoice()])
+    object: str = "chat.completion"
+    # TODO: support ChainCompletionChunk in the future
+
+
+@deprecated("mlflow.types.llm.ChatCompletionResponse")
+@dataclass
+class StringResponse:
+    content: str = "MLflow is an open source platform for the machine learning lifecycle."
+
+
+CHAT_COMPLETION_REQUEST_SCHEMA = Schema(
+    [
+        ColSpec(
+            name="messages",
+            type=Array(
+                Object(
+                    [
+                        Property("role", DataType.string),
+                        Property("content", DataType.string),
+                    ]
+                )
+            ),
+        ),
+    ]
+)
+
+CHAT_COMPLETION_RESPONSE_SCHEMA = Schema(
+    [
+        ColSpec(
+            name="choices",
+            type=Array(
+                Object(
+                    [
+                        Property("index", DataType.long),
+                        Property(
+                            "message",
+                            Object(
+                                [
+                                    Property("role", DataType.string),
+                                    Property("content", DataType.string),
+                                ]
+                            ),
+                        ),
+                        Property("finish_reason", DataType.string),
+                    ]
+                )
+            ),
+        ),
+    ]
+)
+
+SIGNATURE_FOR_LLM_INFERENCE_TASK = {
+    "llm/v1/chat": ModelSignature(
+        inputs=CHAT_COMPLETION_REQUEST_SCHEMA, outputs=CHAT_COMPLETION_RESPONSE_SCHEMA
+    ),
+}

venv/lib/python3.9/site-packages/mlflow/models/resources.py

@@ -0,0 +1,298 @@
+import os
+from abc import ABC, abstractmethod
+from enum import Enum
+from typing import Any, Optional
+
+import yaml
+
+DEFAULT_API_VERSION = "1"
+
+
+class ResourceType(Enum):
+    """
+    Enum to define the different types of resources needed to serve a model.
+    """
+
+    UC_CONNECTION = "uc_connection"
+    VECTOR_SEARCH_INDEX = "vector_search_index"
+    SERVING_ENDPOINT = "serving_endpoint"
+    SQL_WAREHOUSE = "sql_warehouse"
+    FUNCTION = "function"
+    GENIE_SPACE = "genie_space"
+    TABLE = "table"
+
+
+class Resource(ABC):
+    """
+    Base class for defining the resources needed to serve a model.
+
+    Args:
+        type (ResourceType): The resource type.
+        target_uri (str): The target URI where these resources are hosted.
+    """
+
+    @property
+    @abstractmethod
+    def type(self) -> ResourceType:
+        """
+        The resource type (must be defined by subclasses).
+        """
+
+    @property
+    @abstractmethod
+    def target_uri(self) -> str:
+        """
+        The target URI where the resource is hosted (must be defined by subclasses).
+        """
+
+    @abstractmethod
+    def to_dict(self):
+        """
+        Convert the resource to a dictionary.
+        Subclasses must implement this method.
+        """
+
+    @classmethod
+    @abstractmethod
+    def from_dict(cls, data: dict[str, str]):
+        """
+        Convert the dictionary to a Resource.
+        Subclasses must implement this method.
+        """
+
+    def __eq__(self, other: Any):
+        if not isinstance(other, Resource):
+            return False
+        return self.to_dict() == other.to_dict()
+
+
+class DatabricksResource(Resource, ABC):
+    """
+    Base class to define all the Databricks resources to serve a model.
+
+    Example usage: https://docs.databricks.com/en/generative-ai/log-agent.html#specify-resources-for-pyfunc-or-langchain-agent
+    """
+
+    @property
+    def target_uri(self) -> str:
+        return "databricks"
+
+    @property
+    def type(self) -> ResourceType:
+        raise NotImplementedError("Subclasses must implement the 'type' property.")
+
+    def __init__(self, name: str, on_behalf_of_user: Optional[bool] = None):
+        self.name = name
+        self.on_behalf_of_user = on_behalf_of_user
+
+    def to_dict(self):
+        result = {self.type.value: [{"name": self.name}]}
+        if self.on_behalf_of_user is not None:
+            result[self.type.value][0]["on_behalf_of_user"] = self.on_behalf_of_user
+        return result
+
+    @classmethod
+    def from_dict(cls, data: dict[str, str]):
+        return cls(data["name"], data.get("on_behalf_of_user"))
+
+
+class DatabricksUCConnection(DatabricksResource):
+    """
+    Define a Databricks UC Connection used to serve a model.
+
+    Args:
+        connection_name (str): The name of the databricks UC connection
+        used to create the tool which was used to build the model.
+        on_behalf_of_user (Optional[bool]): If True, the resource is accessed with
+        with the permission of the invoker of the model in the serving endpoint. If set to
+        None or False, the resources is accesssed with the permissions of the creator
+    """
+
+    @property
+    def type(self) -> ResourceType:
+        return ResourceType.UC_CONNECTION
+
+    def __init__(self, connection_name: str, on_behalf_of_user: Optional[bool] = None):
+        super().__init__(connection_name, on_behalf_of_user)
+
+
+class DatabricksServingEndpoint(DatabricksResource):
+    """
+    Define Databricks LLM endpoint resource to serve a model.
+
+    Args:
+        endpoint_name (str): The name of all the databricks endpoints used by the model.
+        on_behalf_of_user (Optional[bool]): If True, the resource is accessed with
+        with the permission of the invoker of the model in the serving endpoint. If set to
+        None or False, the resources is accesssed with the permissions of the creator
+    """
+
+    @property
+    def type(self) -> ResourceType:
+        return ResourceType.SERVING_ENDPOINT
+
+    def __init__(self, endpoint_name: str, on_behalf_of_user: Optional[bool] = None):
+        super().__init__(endpoint_name, on_behalf_of_user)
+
+
+class DatabricksVectorSearchIndex(DatabricksResource):
+    """
+    Define Databricks vector search index name resource to serve a model.
+
+    Args:
+        index_name (str): The name of the databricks vector search index
+        used by the model
+        on_behalf_of_user (Optional[bool]): If True, the resource is accessed with
+        with the permission of the invoker of the model in the serving endpoint. If set to
+        None or False, the resources is accesssed with the permissions of the creator
+    """
+
+    @property
+    def type(self) -> ResourceType:
+        return ResourceType.VECTOR_SEARCH_INDEX
+
+    def __init__(self, index_name: str, on_behalf_of_user: Optional[bool] = None):
+        super().__init__(index_name, on_behalf_of_user)
+
+
+class DatabricksSQLWarehouse(DatabricksResource):
+    """
+    Define Databricks sql warehouse resource to serve a model.
+
+    Args:
+        warehouse_id (str): The id of the sql warehouse used by the model
+        on_behalf_of_user (Optional[bool]): If True, the resource is accessed with
+        with the permission of the invoker of the model in the serving endpoint. If set to
+        None or False, the resources is accesssed with the permissions of the creator
+    """
+
+    @property
+    def type(self) -> ResourceType:
+        return ResourceType.SQL_WAREHOUSE
+
+    def __init__(self, warehouse_id: str, on_behalf_of_user: Optional[bool] = None):
+        super().__init__(warehouse_id, on_behalf_of_user)
+
+
+class DatabricksFunction(DatabricksResource):
+    """
+    Define Databricks UC Function to serve a model.
+
+    Args:
+        function_name (str): The name of the function used by the model
+        on_behalf_of_user (Optional[bool]): If True, the resource is accessed with
+        with the permission of the invoker of the model in the serving endpoint. If set to
+        None or False, the resources is accesssed with the permissions of the creator
+    """
+
+    @property
+    def type(self) -> ResourceType:
+        return ResourceType.FUNCTION
+
+    def __init__(self, function_name: str, on_behalf_of_user: Optional[bool] = None):
+        super().__init__(function_name, on_behalf_of_user)
+
+
+class DatabricksGenieSpace(DatabricksResource):
+    """
+    Define a Databricks Genie Space to serve a model.
+
+    Args:
+        genie_space_id (str): The genie space id
+        on_behalf_of_user (Optional[bool]): If True, the resource is accessed with
+        with the permission of the invoker of the model in the serving endpoint. If set to
+        None or False, the resources is accesssed with the permissions of the creator
+    """
+
+    @property
+    def type(self) -> ResourceType:
+        return ResourceType.GENIE_SPACE
+
+    def __init__(self, genie_space_id: str, on_behalf_of_user: Optional[bool] = None):
+        super().__init__(genie_space_id, on_behalf_of_user)
+
+
+class DatabricksTable(DatabricksResource):
+    """
+    Defines a Databricks Unity Catalog (UC) Table, which establishes table dependencies
+    for Model Serving. This table will be referenced in Agent Model Serving endpoints,
+    where an agent queries a SQL table via either Genie or UC Functions.
+
+     Args:
+         table_name (str): The name of the table used by the model
+         on_behalf_of_user (Optional[bool]): If True, the resource is accessed with
+        with the permission of the invoker of the model in the serving endpoint. If set to
+        None or False, the resources is accesssed with the permissions of the creator
+    """
+
+    @property
+    def type(self) -> ResourceType:
+        return ResourceType.TABLE
+
+    def __init__(self, table_name: str, on_behalf_of_user: Optional[bool] = None):
+        super().__init__(table_name, on_behalf_of_user)
+
+
+def _get_resource_class_by_type(target_uri: str, resource_type: ResourceType):
+    resource_classes = {
+        "databricks": {
+            ResourceType.UC_CONNECTION.value: DatabricksUCConnection,
+            ResourceType.SERVING_ENDPOINT.value: DatabricksServingEndpoint,
+            ResourceType.VECTOR_SEARCH_INDEX.value: DatabricksVectorSearchIndex,
+            ResourceType.SQL_WAREHOUSE.value: DatabricksSQLWarehouse,
+            ResourceType.FUNCTION.value: DatabricksFunction,
+            ResourceType.GENIE_SPACE.value: DatabricksGenieSpace,
+            ResourceType.TABLE.value: DatabricksTable,
+        }
+    }
+    resource = resource_classes.get(target_uri)
+    if resource is None:
+        raise ValueError(f"Unsupported target URI: {target_uri}")
+    return resource.get(resource_type)
+
+
+class _ResourceBuilder:
+    """
+    Private builder class to build the resources dictionary.
+    """
+
+    @staticmethod
+    def from_resources(
+        resources: list[Resource], api_version: str = DEFAULT_API_VERSION
+    ) -> dict[str, dict[ResourceType, list[dict]]]:
+        resource_dict = {}
+        for resource in resources:
+            resource_data = resource.to_dict()
+            for resource_type, values in resource_data.items():
+                target_dict = resource_dict.setdefault(resource.target_uri, {})
+                target_list = target_dict.setdefault(resource_type, [])
+                target_list.extend(values)
+
+        resource_dict["api_version"] = api_version
+        return resource_dict
+
+    @staticmethod
+    def from_dict(data) -> dict[str, dict[ResourceType, list[dict]]]:
+        resources = []
+        api_version = data.pop("api_version")
+        if api_version == "1":
+            for target_uri, config in data.items():
+                for resource_type, values in config.items():
+                    resource_class = _get_resource_class_by_type(target_uri, resource_type)
+                    if resource_class:
+                        resources.extend(resource_class.from_dict(value) for value in values)
+                    else:
+                        raise ValueError(f"Unsupported resource type: {resource_type}")
+        else:
+            raise ValueError(f"Unsupported API version: {api_version}")
+
+        return _ResourceBuilder.from_resources(resources, api_version)
+
+    @staticmethod
+    def from_yaml_file(path: str) -> dict[str, dict[ResourceType, list[dict]]]:
+        if not os.path.exists(path):
+            raise OSError(f"No such file or directory: '{path}'")
+        path = os.path.abspath(path)
+        with open(path) as file:
+            data = yaml.safe_load(file)
+            return _ResourceBuilder.from_dict(data)

venv/lib/python3.9/site-packages/mlflow/models/signature.py

@@ -0,0 +1,658 @@
+"""
+The :py:mod:`mlflow.models.signature` module provides an API for specification of model signature.
+
+Model signature defines schema of model input and output. See :py:class:`mlflow.types.schema.Schema`
+for more details on Schema and data types.
+"""
+
+import inspect
+import logging
+import re
+import warnings
+from copy import deepcopy
+from dataclasses import dataclass, is_dataclass
+from typing import TYPE_CHECKING, Any, Optional, Union, get_type_hints
+
+import numpy as np
+import pandas as pd
+
+from mlflow.environment_variables import _MLFLOW_TESTING
+from mlflow.exceptions import MlflowException
+from mlflow.models import Model
+from mlflow.models.model import MLMODEL_FILE_NAME
+from mlflow.models.utils import _contains_params, _Example
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE, RESOURCE_DOES_NOT_EXIST
+from mlflow.store.artifact.models_artifact_repo import ModelsArtifactRepository
+from mlflow.store.artifact.runs_artifact_repo import RunsArtifactRepository
+from mlflow.tracking.artifact_utils import _download_artifact_from_uri, _upload_artifact_to_uri
+from mlflow.types.schema import AnyType, ColSpec, ParamSchema, Schema, convert_dataclass_to_schema
+from mlflow.types.type_hints import (
+    InvalidTypeHintException,
+    _get_data_validation_result,
+    _infer_schema_from_list_type_hint,
+    _infer_schema_from_type_hint,
+    _is_list_type_hint,
+)
+from mlflow.types.utils import (
+    InvalidDataForSignatureInferenceError,
+    _infer_param_schema,
+    _infer_schema,
+)
+from mlflow.utils.annotations import filter_user_warnings_once
+from mlflow.utils.uri import append_to_uri_path
+
+# At runtime, we don't need  `pyspark.sql.dataframe`
+if TYPE_CHECKING:
+    try:
+        import pyspark.sql.dataframe
+
+        MlflowInferableDataset = Union[
+            pd.DataFrame, np.ndarray, dict[str, np.ndarray], pyspark.sql.dataframe.DataFrame
+        ]
+    except ImportError:
+        MlflowInferableDataset = Union[pd.DataFrame, np.ndarray, dict[str, np.ndarray]]
+
+_logger = logging.getLogger(__name__)
+
+_LOG_MODEL_INFER_SIGNATURE_WARNING_TEMPLATE = (
+    "Failed to infer the model signature from the input example. Reason: %s. To see the full "
+    "traceback, set the logging level to DEBUG via "
+    '`logging.getLogger("mlflow").setLevel(logging.DEBUG)`.'
+)
+
+
+class ModelSignature:
+    """
+    ModelSignature specifies schema of model's inputs, outputs and params.
+
+    ModelSignature can be :py:func:`inferred <mlflow.models.infer_signature>` from training
+    dataset, model predictions using and params for inference, or constructed by hand by
+    passing an input and output :py:class:`Schema <mlflow.types.Schema>`, and params
+    :py:class:`ParamSchema <mlflow.types.ParamSchema>`.
+    """
+
+    def __init__(
+        self,
+        inputs: Union[Schema, dataclass] = None,
+        outputs: Union[Schema, dataclass] = None,
+        params: ParamSchema = None,
+    ):
+        if inputs and not isinstance(inputs, Schema) and not is_dataclass(inputs):
+            raise TypeError(
+                "inputs must be either None, mlflow.models.signature.Schema, or a dataclass,"
+                f"got '{type(inputs).__name__}'"
+            )
+        if outputs and not isinstance(outputs, Schema) and not is_dataclass(outputs):
+            raise TypeError(
+                "outputs must be either None, mlflow.models.signature.Schema, or a dataclass,"
+                f"got '{type(outputs).__name__}'"
+            )
+        if params and not isinstance(params, ParamSchema):
+            raise TypeError(
+                "If params are provided, they must by of type mlflow.models.signature.ParamSchema, "
+                f"got '{type(params).__name__}'"
+            )
+        if all(x is None for x in [inputs, outputs, params]):
+            raise ValueError("At least one of inputs, outputs or params must be provided")
+        if is_dataclass(inputs):
+            self.inputs = convert_dataclass_to_schema(inputs)
+        else:
+            self.inputs = inputs
+        if is_dataclass(outputs):
+            self.outputs = convert_dataclass_to_schema(outputs)
+        else:
+            self.outputs = outputs
+        self.params = params
+        self.__is_signature_from_type_hint = False
+        self.__is_type_hint_from_example = False
+
+    @property
+    def _is_signature_from_type_hint(self):
+        return self.__is_signature_from_type_hint
+
+    @_is_signature_from_type_hint.setter
+    def _is_signature_from_type_hint(self, value):
+        self.__is_signature_from_type_hint = value
+
+    @property
+    def _is_type_hint_from_example(self):
+        return self.__is_type_hint_from_example
+
+    @_is_type_hint_from_example.setter
+    def _is_type_hint_from_example(self, value):
+        self.__is_type_hint_from_example = value
+
+    def to_dict(self) -> dict[str, Any]:
+        """
+        Serialize into a 'jsonable' dictionary.
+
+        Input and output schema are represented as json strings. This is so that the
+        representation is compact when embedded in an MLmodel yaml file.
+
+        Returns:
+            dictionary representation with input and output schema represented as json strings.
+        """
+
+        return {
+            "inputs": self.inputs.to_json() if self.inputs else None,
+            "outputs": self.outputs.to_json() if self.outputs else None,
+            "params": self.params.to_json() if self.params else None,
+        }
+
+    @classmethod
+    def from_dict(cls, signature_dict: dict[str, Any]):
+        """
+        Deserialize from dictionary representation.
+
+        Args:
+            signature_dict: Dictionary representation of model signature.
+                Expected dictionary format:
+                `{'inputs': <json string>,
+                'outputs': <json string>,
+                'params': <json string>" }`
+
+        Returns:
+            ModelSignature populated with the data form the dictionary.
+        """
+        inputs = Schema.from_json(x) if (x := signature_dict.get("inputs")) else None
+        outputs = Schema.from_json(x) if (x := signature_dict.get("outputs")) else None
+        params = ParamSchema.from_json(x) if (x := signature_dict.get("params")) else None
+        return cls(inputs, outputs, params)
+
+    def __eq__(self, other) -> bool:
+        return (
+            isinstance(other, ModelSignature)
+            and self.inputs == other.inputs
+            and self.outputs == other.outputs
+            and self.params == other.params
+        )
+
+    def __repr__(self) -> str:
+        return (
+            "inputs: \n"
+            f"  {self.inputs!r}\n"
+            "outputs: \n"
+            f"  {self.outputs!r}\n"
+            "params: \n"
+            f"  {self.params!r}\n"
+        )
+
+
+def infer_signature(
+    model_input: Any = None,
+    model_output: "MlflowInferableDataset" = None,
+    params: Optional[dict[str, Any]] = None,
+) -> ModelSignature:
+    """
+    Infer an MLflow model signature from the training data (input), model predictions (output)
+    and parameters (for inference).
+
+    The signature represents model input and output as data frames with (optionally) named columns
+    and data type specified as one of types defined in :py:class:`mlflow.types.DataType`. It also
+    includes parameters schema for inference, .
+    This method will raise an exception if the user data contains incompatible types or is not
+    passed in one of the supported formats listed below.
+
+    The input should be one of these:
+      - pandas.DataFrame
+      - pandas.Series
+      - dictionary of { name -> numpy.ndarray}
+      - numpy.ndarray
+      - pyspark.sql.DataFrame
+      - scipy.sparse.csr_matrix
+      - scipy.sparse.csc_matrix
+      - dictionary / list of dictionaries of JSON-convertible types
+
+    The element types should be mappable to one of :py:class:`mlflow.types.DataType`.
+
+    For pyspark.sql.DataFrame inputs, columns of type DateType and TimestampType are both inferred
+    as type :py:data:`datetime <mlflow.types.DataType.datetime>`, which is coerced to
+    TimestampType at inference.
+
+    Args:
+      model_input: Valid input to the model. E.g. (a subset of) the training dataset.
+      model_output: Valid model output. E.g. Model predictions for the (subset of) training
+                    dataset.
+      params: Valid parameters for inference. It should be a dictionary of parameters
+              that can be set on the model during inference by passing `params` to pyfunc
+              `predict` method.
+
+              An example of valid parameters:
+
+              .. code-block:: python
+
+                    from mlflow.models import infer_signature
+                    from mlflow.transformers import generate_signature_output
+
+                    # Define parameters for inference
+                    params = {
+                        "num_beams": 5,
+                        "max_length": 30,
+                        "do_sample": True,
+                        "remove_invalid_values": True,
+                    }
+
+                    # Infer the signature including parameters
+                    signature = infer_signature(
+                        data,
+                        generate_signature_output(model, data),
+                        params=params,
+                    )
+
+                    # Saving model with model signature
+                    mlflow.transformers.save_model(
+                        model,
+                        path=model_path,
+                        signature=signature,
+                    )
+
+                    pyfunc_loaded = mlflow.pyfunc.load_model(model_path)
+
+                    # Passing params to `predict` function directly
+                    result = pyfunc_loaded.predict(data, params=params)
+
+    Returns:
+      ModelSignature
+    """
+    schemas = {"inputs": model_input, "outputs": model_output}
+    for key, data in schemas.items():
+        if data is not None:
+            try:
+                schemas[key] = (
+                    convert_dataclass_to_schema(data) if is_dataclass(data) else _infer_schema(data)
+                )
+            except InvalidDataForSignatureInferenceError:
+                raise
+            except Exception:
+                extra_msg = (
+                    ("Note that MLflow doesn't validate data types during inference for AnyType. ")
+                    if key == "inputs"
+                    else ""
+                )
+                _logger.warning(
+                    f"Failed to infer schema for {key}. "
+                    f"Setting schema to `Schema([ColSpec(type=AnyType())]` as default. {extra_msg}"
+                    "To see the full traceback, set logging level to DEBUG.",
+                    exc_info=_logger.isEnabledFor(logging.DEBUG),
+                )
+                schemas[key] = Schema([ColSpec(type=AnyType())])
+    schemas["params"] = _infer_param_schema(params) if params else None
+    return ModelSignature(**schemas)
+
+
+# `t\w*\.` matches the `typing` module or its alias
+_LIST_OF_STRINGS_PATTERN = re.compile(r"^(t\w*\.)?list\[str\]$", re.IGNORECASE)
+
+
+def _is_list_str(hint_str):
+    return _LIST_OF_STRINGS_PATTERN.match(hint_str.replace(" ", "")) is not None
+
+
+_LIST_OF_STR_DICT_PATTERN = re.compile(
+    r"^(t\w*\.)?list\[(t\w*\.)?dict\[str,str\]\]$", re.IGNORECASE
+)
+
+
+def _is_list_of_string_dict(hint_str):
+    return _LIST_OF_STR_DICT_PATTERN.match(hint_str.replace(" ", "")) is not None
+
+
+def _infer_hint_from_str(hint_str):
+    if _is_list_str(hint_str):
+        return list[str]
+    elif _is_list_of_string_dict(hint_str):
+        return list[dict[str, str]]
+    else:
+        return None
+
+
+def _get_arg_names(f):
+    return list(inspect.signature(f).parameters.keys())
+
+
+class _TypeHints:
+    def __init__(self, input_=None, output=None):
+        self.input = input_
+        self.output = output
+
+    def __repr__(self):
+        return f"<input: {self.input}, output: {self.output}>"
+
+
+def _extract_type_hints(f, input_arg_index):
+    """
+    Extract type hints from a function.
+
+    Args:
+        f: Function to extract type hints from.
+        input_arg_index: Index of the function argument that corresponds to the model input.
+
+    Returns:
+        A `_TypeHints` object containing the input and output type hints.
+    """
+    if not hasattr(f, "__annotations__") and hasattr(f, "__call__"):
+        return _extract_type_hints(f.__call__, input_arg_index)
+
+    if f.__annotations__ == {}:
+        return _TypeHints()
+
+    arg_names = list(filter(lambda x: x != "self", _get_arg_names(f)))
+    if len(arg_names) - 1 < input_arg_index:
+        raise MlflowException.invalid_parameter_value(
+            f"The specified input argument index ({input_arg_index}) is out of range for the "
+            "function signature: {}".format(input_arg_index, arg_names)
+        )
+    arg_name = arg_names[input_arg_index]
+    try:
+        hints = get_type_hints(f)
+    except (
+        TypeError,
+        NameError,  # To handle this issue: https://github.com/python/typing/issues/797
+    ):
+        # ---
+        # from __future__ import annotations # postpones evaluation of 'list[str]'
+        #
+        # def f(x: list[str]) -> list[str]:
+        #          ^^^^^^^^^ Evaluating this expression ('list[str]') results in a TypeError in
+        #                    Python < 3.9 because the built-in list type is not subscriptable.
+        #     return x
+        # ---
+        # Best effort to infer type hints from strings
+        hints = {}
+        for arg in [arg_name, "return"]:
+            if hint_str := f.__annotations__.get(arg, None):
+                if hint := _infer_hint_from_str(hint_str):
+                    hints[arg] = hint
+                else:
+                    _logger.info("Unsupported type hint: %s, skipping schema inference", hint_str)
+    except Exception as e:
+        _logger.warning("Failed to extract type hints from function %s: %s", f.__name__, repr(e))
+        return _TypeHints()
+
+    return _TypeHints(hints.get(arg_name), hints.get("return"))
+
+
+def _is_context_in_predict_function_signature(*, func=None, parameters=None):
+    if parameters is None:
+        if func is None:
+            raise ValueError("Either `func` or `parameters` must be provided.")
+        parameters = inspect.signature(func).parameters
+    return (
+        # predict(self, context, model_input, ...)
+        "context" in parameters
+        # predict(self, ctx, model_input, ...) ctx can be any parameter name
+        or len([param for param in parameters if param not in ("self", "params")]) == 2
+    )
+
+
+@filter_user_warnings_once
+def _infer_signature_from_type_hints(
+    func, type_hints: _TypeHints, input_example=None
+) -> Optional[ModelSignature]:
+    """
+    Infer the signature from type hints.
+    """
+    if type_hints.input is None:
+        return None
+
+    params = None
+    params_key = "params"
+    if _contains_params(input_example):
+        input_example, params = input_example
+
+    _logger.info("Inferring model signature from type hints")
+    try:
+        input_schema = _infer_schema_from_list_type_hint(type_hints.input)
+    except InvalidTypeHintException:
+        raise MlflowException.invalid_parameter_value(
+            "The `predict` function has unsupported type hints for the model input "
+            "arguments. Update it to one of supported type hints, or remove type hints "
+            "to bypass this check. Error: {e}"
+        )
+    except Exception as e:
+        warnings.warn(f"Failed to infer signature from type hint: {e.message}", stacklevel=3)
+        return None
+
+    # only warn if the pyfunc decorator is not used and schema can
+    # be inferred from the input type hint
+    _pyfunc_decorator_used = getattr(func, "_is_pyfunc", False)
+    if not _pyfunc_decorator_used:
+        # stacklevel is 3 because we have a decorator
+        warnings.warn(
+            "Decorate your function with `@mlflow.pyfunc.utils.pyfunc` to enable auto "
+            "data validation against model input type hints.",
+            stacklevel=3,
+        )
+
+    default_output_schema = Schema([ColSpec(type=AnyType())])
+    is_output_type_hint_valid = False
+    output_schema = None
+    if type_hints.output:
+        try:
+            # output type hint doesn't need to be a list
+            # but if it's a list, we infer the schema from the list type hint
+            # to be consistent with input schema inference
+            output_schema = (
+                _infer_schema_from_list_type_hint(type_hints.output)
+                if _is_list_type_hint(type_hints.output)
+                else _infer_schema_from_type_hint(type_hints.output)
+            )
+            is_output_type_hint_valid = True
+        except Exception as e:
+            _logger.info(
+                f"Failed to infer output type hint, setting output schema to AnyType. {e}",
+                stacklevel=2,
+            )
+            output_schema = default_output_schema
+    else:
+        output_schema = default_output_schema
+    params_schema = _infer_param_schema(params) if params else None
+
+    if input_example is not None:
+        # only validate input example here if pyfunc decorator is not used
+        # because when the decorator is used, the input is validated in the predict function
+        if not _pyfunc_decorator_used and (
+            msg := _get_data_validation_result(
+                data=input_example, type_hint=type_hints.input
+            ).error_message
+        ):
+            _logger.warning(
+                "Input example is not compatible with the type hint of the `predict` function. "
+                f"Error: {msg}"
+            )
+        else:
+            kwargs = (
+                {params_key: params}
+                if params and params_key in inspect.signature(func).parameters
+                else {}
+            )
+            # This is for PythonModel's predict function
+            if _is_context_in_predict_function_signature(func=func):
+                inputs = [None, input_example]
+            else:
+                inputs = [input_example]
+            _logger.info("Running the predict function to generate output based on input example")
+            try:
+                output_example = func(*inputs, **kwargs)
+            except Exception:
+                _logger.warning(
+                    "Failed to run the predict function on input example. To see the full "
+                    "traceback, set logging level to DEBUG.",
+                    exc_info=_logger.isEnabledFor(logging.DEBUG),
+                )
+            else:
+                if is_output_type_hint_valid and (
+                    msg := _get_data_validation_result(
+                        data=output_example, type_hint=type_hints.output
+                    ).error_message
+                ):
+                    _logger.warning(
+                        f"Failed to validate output `{output_example}` against type hint "
+                        f"`{type_hints.output}`, setting output schema to AnyType. "
+                        f"Error: {msg}"
+                    )
+                    output_schema = default_output_schema
+    if not any([input_schema, output_schema, params_schema]):
+        return None
+    signature = ModelSignature(inputs=input_schema, outputs=output_schema, params=params_schema)
+    signature._is_signature_from_type_hint = True
+    return signature
+
+
+def _infer_signature_from_input_example(
+    input_example: Optional[_Example], wrapped_model
+) -> Optional[ModelSignature]:
+    """
+    Infer the signature from an example input and a PyFunc wrapped model. Catches all exceptions.
+
+    Args:
+        input_example: Saved _Example object that contains input example instance.
+        wrapped_model: A PyFunc wrapped model which has a `predict` method.
+
+    Returns:
+        A `ModelSignature` object containing the inferred schema of both the model's inputs
+        based on the `input_example` and the model's outputs based on the prediction from the
+        `wrapped_model`.
+    """
+    from mlflow.pyfunc import _validate_prediction_input
+
+    if input_example is None:
+        return None
+
+    try:
+        # Copy the input example so that it is not mutated by predict()
+        input_data = deepcopy(input_example.inference_data)
+        params = input_example.inference_params
+
+        input_schema = _infer_schema(input_data)
+        params_schema = _infer_param_schema(params) if params else None
+        # do the same validation as pyfunc predict to make sure the signature is correctly
+        # applied to the model
+        input_data, params = _validate_prediction_input(
+            input_data, params, input_schema, params_schema
+        )
+        prediction = wrapped_model.predict(input_data, params=params)
+        # For column-based inputs, 1D numpy arrays likely signify row-based predictions. Thus, we
+        # convert them to a Pandas series for inferring as a single ColSpec Schema.
+        if (
+            not input_schema.is_tensor_spec()
+            and isinstance(prediction, np.ndarray)
+            and prediction.ndim == 1
+        ):
+            prediction = pd.Series(prediction)
+
+        output_schema = None
+        try:
+            output_schema = _infer_schema(prediction)
+        except Exception:
+            # try assign output schema if failing to infer it from prediction for langchain models
+            try:
+                from mlflow.langchain import _LangChainModelWrapper
+                from mlflow.langchain.utils.chat import _ChatResponse
+            except ImportError:
+                pass
+            else:
+                if isinstance(wrapped_model, _LangChainModelWrapper) and isinstance(
+                    prediction, _ChatResponse
+                ):
+                    output_schema = prediction.get_schema()
+            if output_schema is None:
+                _logger.warning(
+                    "Failed to infer model output schema from prediction result, setting "
+                    "output schema to AnyType. For full traceback, set logging level to debug.",
+                    exc_info=_logger.isEnabledFor(logging.DEBUG),
+                )
+                output_schema = Schema([ColSpec(type=AnyType())])
+
+        return ModelSignature(input_schema, output_schema, params_schema)
+    except Exception as e:
+        if _MLFLOW_TESTING.get():
+            raise
+        _logger.warning(
+            _LOG_MODEL_INFER_SIGNATURE_WARNING_TEMPLATE,
+            repr(e),
+            exc_info=_logger.isEnabledFor(logging.DEBUG),
+        )
+
+
+def set_signature(
+    model_uri: str,
+    signature: ModelSignature,
+):
+    """
+    Sets the model signature for specified model artifacts.
+
+    The process involves downloading the MLmodel file in the model artifacts (if it's non-local),
+    updating its model signature, and then overwriting the existing MLmodel file. Should the
+    artifact repository associated with the model artifacts disallow overwriting, this function will
+    fail.
+
+    Furthermore, as model registry artifacts are read-only, model artifacts located in the
+    model registry and represented by ``models:/`` URI schemes are not compatible with this API.
+    To set a signature on a model version, first set the signature on the source model artifacts.
+    Following this, generate a new model version using the updated model artifacts. For more
+    information about setting signatures on model versions, see
+    `this doc section <https://www.mlflow.org/docs/latest/models.html#set-signature-on-mv>`_.
+
+    Args:
+        model_uri: The location, in URI format, of the MLflow model. For example:
+
+            - ``/Users/me/path/to/local/model``
+            - ``relative/path/to/local/model``
+            - ``s3://my_bucket/path/to/model``
+            - ``runs:/<mlflow_run_id>/run-relative/path/to/model``
+            - ``mlflow-artifacts:/path/to/model``
+
+            For more information about supported URI schemes, see
+            `Referencing Artifacts <https://www.mlflow.org/docs/latest/concepts.html#
+            artifact-locations>`_.
+
+            Please note that model URIs with the ``models:/`` scheme are not supported.
+
+        signature: ModelSignature to set on the model.
+
+    .. code-block:: python
+        :caption: Example
+
+        import mlflow
+        from mlflow.models import set_signature, infer_signature
+
+        # load model from run artifacts
+        run_id = "96771d893a5e46159d9f3b49bf9013e2"
+        artifact_path = "models"
+        model_uri = f"runs:/{run_id}/{artifact_path}"
+        model = mlflow.pyfunc.load_model(model_uri)
+
+        # determine model signature
+        test_df = ...
+        predictions = model.predict(test_df)
+        signature = infer_signature(test_df, predictions)
+
+        # set the signature for the logged model
+        set_signature(model_uri, signature)
+    """
+    assert isinstance(signature, ModelSignature), (
+        "The signature argument must be a ModelSignature object"
+    )
+    if ModelsArtifactRepository.is_models_uri(model_uri):
+        raise MlflowException(
+            f'Failed to set signature on "{model_uri}". '
+            + "Model URIs with the `models:/` scheme are not supported.",
+            INVALID_PARAMETER_VALUE,
+        )
+    try:
+        resolved_uri = model_uri
+        if RunsArtifactRepository.is_runs_uri(model_uri):
+            resolved_uri = RunsArtifactRepository.get_underlying_uri(model_uri)
+        ml_model_file = _download_artifact_from_uri(
+            artifact_uri=append_to_uri_path(resolved_uri, MLMODEL_FILE_NAME)
+        )
+    except Exception as ex:
+        raise MlflowException(
+            f'Failed to download an "{MLMODEL_FILE_NAME}" model file from "{model_uri}"',
+            RESOURCE_DOES_NOT_EXIST,
+        ) from ex
+    model_meta = Model.load(ml_model_file)
+    model_meta.signature = signature
+    model_meta.save(ml_model_file)
+    _upload_artifact_to_uri(ml_model_file, resolved_uri)

venv/lib/python3.9/site-packages/mlflow/models/wheeled_model.py

@@ -0,0 +1,282 @@
+import os
+import platform
+import shutil
+import subprocess
+import sys
+
+import yaml
+
+import mlflow
+from mlflow import MlflowClient
+from mlflow.environment_variables import MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import BAD_REQUEST
+from mlflow.pyfunc.model import MLMODEL_FILE_NAME, Model
+from mlflow.store.artifact.utils.models import _parse_model_uri, get_model_name_and_version
+from mlflow.tracking.artifact_utils import _download_artifact_from_uri
+from mlflow.utils.annotations import experimental
+from mlflow.utils.environment import (
+    _REQUIREMENTS_FILE_NAME,
+    _get_pip_deps,
+    _mlflow_additional_pip_env,
+    _overwrite_pip_deps,
+)
+from mlflow.utils.model_utils import _validate_and_prepare_target_save_path
+from mlflow.utils.uri import get_databricks_profile_uri_from_artifact_uri
+
+_WHEELS_FOLDER_NAME = "wheels"
+_ORIGINAL_REQ_FILE_NAME = "original_requirements.txt"
+_PLATFORM = "platform"
+
+
+@experimental
+class WheeledModel:
+    """
+    Helper class to create a model with added dependency wheels from an existing registered model.
+    The `wheeled` model contains all the model dependencies as wheels stored as model artifacts.
+    .. note::
+        This utility only operates on a model that has been registered to the Model Registry.
+    """
+
+    def __init__(self, model_uri):
+        self._model_uri = model_uri
+        databricks_profile_uri = (
+            get_databricks_profile_uri_from_artifact_uri(model_uri) or mlflow.get_registry_uri()
+        )
+        client = MlflowClient(registry_uri=databricks_profile_uri)
+        self._model_name, _ = get_model_name_and_version(client, model_uri)
+
+    @classmethod
+    def log_model(cls, model_uri, registered_model_name=None):
+        """
+        Logs a registered model as an MLflow artifact for the current run. This only operates on
+        a model which has been registered to the Model Registry. Given a registered model_uri (
+        e.g. models:/<model_name>/<model_version>), this utility re-logs the model along with all
+        the required model libraries back to the Model Registry. The required model libraries are
+        stored along with the model as model artifacts. In addition, supporting files to the
+        model (e.g. conda.yaml, requirements.txt) are modified to use the added libraries.
+
+        By default, this utility creates a new model version under the same registered model
+        specified by ``model_uri``. This behavior can be overridden by specifying the
+        ``registered_model_name`` argument.
+
+        Args:
+            model_uri: A registered model uri in the Model Registry of the form
+                       models:/<model_name>/<model_version/stage/latest>
+            registered_model_name: The new model version (model with its libraries) is
+                                   registered under the inputted registered_model_name. If None,
+                                   a new version is logged to the existing model in the Model
+                                   Registry.
+
+        .. code-block:: python
+            :caption: Example
+
+            # Given a model uri, log the wheeled model
+            with mlflow.start_run():
+                WheeledModel.log_model(model_uri)
+        """
+        parsed_uri = _parse_model_uri(model_uri)
+        return Model.log(
+            artifact_path=None,
+            flavor=WheeledModel(model_uri),
+            registered_model_name=registered_model_name or parsed_uri.name,
+        )
+
+    def save_model(self, path, mlflow_model=None):
+        """
+        Given an existing registered model, saves the model along with it's dependencies stored as
+        wheels to a path on the local file system.
+
+        This does not modify existing model behavior or existing model flavors. It simply downloads
+        the model dependencies as wheels and modifies the requirements.txt and conda.yaml file to
+        point to the downloaded wheels.
+
+        The download_command defaults to downloading only binary packages using the
+        `--only-binary=:all:` option. This behavior can be overridden using an environment
+        variable `MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS`, which will allows setting
+        different options such as `--prefer-binary`, `--no-binary`, etc.
+
+        Args:
+            path: Local path where the model is to be saved.
+            mlflow_model: The new :py:mod:`mlflow.models.Model` metadata file to store the
+                updated model metadata.
+        """
+        from mlflow.pyfunc import ENV, FLAVOR_NAME, _extract_conda_env
+
+        path = os.path.abspath(path)
+        _validate_and_prepare_target_save_path(path)
+
+        local_model_path = _download_artifact_from_uri(self._model_uri, output_path=path)
+
+        wheels_dir = os.path.join(local_model_path, _WHEELS_FOLDER_NAME)
+        pip_requirements_path = os.path.join(local_model_path, _REQUIREMENTS_FILE_NAME)
+        model_metadata_path = os.path.join(local_model_path, MLMODEL_FILE_NAME)
+
+        model_metadata = Model.load(model_metadata_path)
+
+        # Check if the model file has `wheels` set to True
+        if model_metadata.__dict__.get(_WHEELS_FOLDER_NAME, None) is not None:
+            raise MlflowException("Model libraries are already added", BAD_REQUEST)
+
+        conda_env = _extract_conda_env(model_metadata.flavors.get(FLAVOR_NAME, {}).get(ENV, None))
+        conda_env_path = os.path.join(local_model_path, conda_env)
+        if conda_env is None and not os.path.isfile(pip_requirements_path):
+            raise MlflowException(
+                "Cannot add libraries for model with no logged dependencies.", BAD_REQUEST
+            )
+
+        if not os.path.isfile(pip_requirements_path):
+            self._create_pip_requirement(conda_env_path, pip_requirements_path)
+
+        WheeledModel._download_wheels(
+            pip_requirements_path=pip_requirements_path, dst_path=wheels_dir
+        )
+
+        # Keep a copy of the original requirement.txt
+        shutil.copy2(pip_requirements_path, os.path.join(local_model_path, _ORIGINAL_REQ_FILE_NAME))
+
+        # Update requirements.txt with wheels
+        pip_deps = self._overwrite_pip_requirements_with_wheels(
+            pip_requirements_path=pip_requirements_path, wheels_dir=wheels_dir
+        )
+
+        # Update conda.yaml with wheels
+        self._update_conda_env(pip_deps, conda_env_path)
+
+        # Update MLModel File
+        mlflow_model = self._update_mlflow_model(
+            original_model_metadata=model_metadata, mlflow_model=mlflow_model
+        )
+        mlflow_model.save(model_metadata_path)
+        return mlflow_model
+
+    def _update_conda_env(self, new_pip_deps, conda_env_path):
+        """
+        Updates the list pip packages in the conda.yaml file to the list of wheels in the wheels
+        directory.
+        {
+            "name": "env",
+            "channels": [...],
+            "dependencies": [
+                ...,
+                "pip",
+                {"pip": [...]},  <- Overwrite this with list of wheels
+            ],
+        }
+
+        Args:
+            new_pip_deps: List of pip dependencies as wheels
+            conda_env_path: Path to conda.yaml file in the model directory
+        """
+        with open(conda_env_path) as f:
+            conda_env = yaml.safe_load(f)
+
+        new_conda_env = _overwrite_pip_deps(conda_env, new_pip_deps)
+
+        with open(conda_env_path, "w") as out:
+            yaml.safe_dump(new_conda_env, stream=out, default_flow_style=False)
+
+    def _update_mlflow_model(self, original_model_metadata, mlflow_model):
+        """
+        Modifies the MLModel file to reflect updated information such as the run_id,
+        utc_time_created. Additionally, this also adds `wheels` to the MLModel file to indicate that
+        this is a `wheeled` model.
+
+        Args:
+            original_model_metadata: The model metadata stored in the original MLmodel file.
+            mlflow_model: :py:mod:`mlflow.models.Model` configuration of the newly created
+                          wheeled model
+        """
+
+        run_id = mlflow.tracking.fluent._get_or_start_run().info.run_id
+        if mlflow_model is None:
+            mlflow_model = Model(run_id=run_id)
+
+        original_model_metadata.__dict__.update(
+            {k: v for k, v in mlflow_model.__dict__.items() if v}
+        )
+        mlflow_model.__dict__.update(original_model_metadata.__dict__)
+        mlflow_model.artifact_path = WheeledModel.get_wheel_artifact_path(
+            mlflow_model.artifact_path
+        )
+
+        mlflow_model.wheels = {_PLATFORM: platform.platform()}
+        return mlflow_model
+
+    @classmethod
+    def _download_wheels(cls, pip_requirements_path, dst_path):
+        """
+        Downloads all the wheels of the dependencies specified in the requirements.txt file.
+        The pip wheel download_command defaults to downloading only binary packages using
+        the `--only-binary=:all:` option. This behavior can be overridden using an
+        environment variable `MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS`, which will allows
+        setting different options such as `--prefer-binary`, `--no-binary`, etc.
+
+        Args:
+            pip_requirements_path: Path to requirements.txt in the model directory
+            dst_path: Path to the directory where the wheels are to be downloaded
+        """
+        if not os.path.exists(dst_path):
+            os.makedirs(dst_path)
+
+        pip_wheel_options = MLFLOW_WHEELED_MODEL_PIP_DOWNLOAD_OPTIONS.get()
+
+        try:
+            subprocess.run(
+                [
+                    sys.executable,
+                    "-m",
+                    "pip",
+                    "wheel",
+                    pip_wheel_options,
+                    "--wheel-dir",
+                    dst_path,
+                    "-r",
+                    pip_requirements_path,
+                    "--no-cache-dir",
+                    "--progress-bar=off",
+                ],
+                check=True,
+                stdout=subprocess.PIPE,
+                stderr=subprocess.STDOUT,
+            )
+        except subprocess.CalledProcessError as e:
+            raise MlflowException(
+                f"An error occurred while downloading the dependency wheels: {e.stdout}"
+            )
+
+    def _overwrite_pip_requirements_with_wheels(self, pip_requirements_path, wheels_dir):
+        """
+        Overwrites the requirements.txt with the wheels of the required dependencies.
+
+        Args:
+            pip_requirements_path: Path to requirements.txt in the model directory.
+            wheels_dir: Path to directory where wheels are stored.
+        """
+        wheels = []
+        with open(pip_requirements_path, "w") as wheels_requirements:
+            for wheel_file in os.listdir(wheels_dir):
+                if wheel_file.endswith(".whl"):
+                    complete_wheel_file = os.path.join(_WHEELS_FOLDER_NAME, wheel_file)
+                    wheels.append(complete_wheel_file)
+                    wheels_requirements.write(complete_wheel_file + "\n")
+        return wheels
+
+    def _create_pip_requirement(self, conda_env_path, pip_requirements_path):
+        """
+        This method creates a requirements.txt file for the model dependencies if the file does not
+        already exist. It uses the pip dependencies found in the conda.yaml env file.
+
+        Args:
+            conda_env_path: Path to conda.yaml env file which contains the required pip
+                dependencies
+            pip_requirements_path: Path where the new requirements.txt will be created.
+        """
+        with open(conda_env_path) as f:
+            conda_env = yaml.safe_load(f)
+        pip_deps = _get_pip_deps(conda_env)
+        _mlflow_additional_pip_env(pip_deps, pip_requirements_path)
+
+    @classmethod
+    def get_wheel_artifact_path(cls, original_artifact_path):
+        return original_artifact_path + "_" + _WHEELS_FOLDER_NAME