star

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

@@ -0,0 +1,119 @@
+"""
+Exposes functionality for deploying MLflow models to custom serving tools.
+
+Note: model deployment to AWS Sagemaker can currently be performed via the
+:py:mod:`mlflow.sagemaker` module. Model deployment to Azure can be performed by using the
+`azureml library <https://pypi.org/project/azureml-mlflow/>`_.
+
+MLflow does not currently provide built-in support for any other deployment targets, but support
+for custom targets can be installed via third-party plugins. See a list of known plugins
+`here <https://mlflow.org/docs/latest/plugins.html#deployment-plugins>`_.
+
+This page largely focuses on the user-facing deployment APIs. For instructions on implementing
+your own plugin for deployment to a custom serving tool, see
+`plugin docs <http://mlflow.org/docs/latest/plugins.html#writing-your-own-mlflow-plugins>`_.
+"""
+
+import contextlib
+import json
+
+from mlflow.deployments.base import BaseDeploymentClient
+from mlflow.deployments.databricks import DatabricksDeploymentClient, DatabricksEndpoint
+from mlflow.deployments.interface import get_deploy_client, run_local
+from mlflow.deployments.openai import OpenAIDeploymentClient
+from mlflow.deployments.utils import get_deployments_target, set_deployments_target
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+
+with contextlib.suppress(Exception):
+    # MlflowDeploymentClient depends on optional dependencies and can't be imported
+    # if they are not installed.
+    from mlflow.deployments.mlflow import MlflowDeploymentClient
+
+
+class PredictionsResponse(dict):
+    """
+    Represents the predictions and metadata returned in response to a scoring request, such as a
+    REST API request sent to the ``/invocations`` endpoint of an MLflow Model Server.
+    """
+
+    def get_predictions(self, predictions_format="dataframe", dtype=None):
+        """Get the predictions returned from the MLflow Model Server in the specified format.
+
+        Args:
+            predictions_format: The format in which to return the predictions. Either
+                ``"dataframe"`` or ``"ndarray"``.
+            dtype: The NumPy datatype to which to coerce the predictions. Only used when
+                the "ndarray" predictions_format is specified.
+
+        Raises:
+            Exception: If the predictions cannot be represented in the specified format.
+
+        Returns:
+            The predictions, represented in the specified format.
+
+        """
+        import numpy as np
+        import pandas as pd
+        from pandas.core.dtypes.common import is_list_like
+
+        if predictions_format == "dataframe":
+            predictions = self["predictions"]
+            if isinstance(predictions, str):
+                return pd.DataFrame(data=[predictions])
+            if isinstance(predictions, dict) and not any(
+                is_list_like(p) and getattr(p, "ndim", 1) == 1 for p in predictions.values()
+            ):
+                return pd.DataFrame(data=predictions, index=[0])
+            return pd.DataFrame(data=predictions)
+        elif predictions_format == "ndarray":
+            return np.array(self["predictions"], dtype)
+        else:
+            raise MlflowException(
+                f"Unrecognized predictions format: '{predictions_format}'",
+                INVALID_PARAMETER_VALUE,
+            )
+
+    def to_json(self, path=None):
+        """Get the JSON representation of the MLflow Predictions Response.
+
+        Args:
+            path: If specified, the JSON representation is written to this file path.
+
+        Returns:
+            If ``path`` is unspecified, the JSON representation of the MLflow Predictions
+            Response. Else, None.
+
+        """
+        if path is not None:
+            with open(path, "w") as f:
+                json.dump(dict(self), f)
+        else:
+            return json.dumps(dict(self))
+
+    @classmethod
+    def from_json(cls, json_str):
+        try:
+            parsed_response = json.loads(json_str)
+        except Exception as e:
+            raise MlflowException("Predictions response contents are not valid JSON") from e
+        if not isinstance(parsed_response, dict) or "predictions" not in parsed_response:
+            raise MlflowException(
+                f"Invalid response. Predictions response contents must be a dictionary"
+                f" containing a 'predictions' field. Instead, received: {parsed_response}"
+            )
+        return PredictionsResponse(parsed_response)
+
+
+__all__ = [
+    "get_deploy_client",
+    "run_local",
+    "BaseDeploymentClient",
+    "DatabricksDeploymentClient",
+    "OpenAIDeploymentClient",
+    "DatabricksEndpoint",
+    "MlflowDeploymentClient",
+    "PredictionsResponse",
+    "get_deployments_target",
+    "set_deployments_target",
+]

venv/lib/python3.9/site-packages/mlflow/deployments/base.py

@@ -0,0 +1,358 @@
+"""
+This module contains the base interface implemented by MLflow model deployment plugins.
+In particular, a valid deployment plugin module must implement:
+
+1. Exactly one client class subclassed from :py:class:`BaseDeploymentClient`, exposing the primary
+   user-facing APIs used to manage deployments.
+2. :py:func:`run_local`, for testing deployment by deploying a model locally
+3. :py:func:`target_help`, which returns a help message describing target-specific URI format
+   and deployment config
+"""
+
+import abc
+
+from mlflow.exceptions import MlflowException
+from mlflow.utils.annotations import developer_stable
+
+
+def run_local(target, name, model_uri, flavor=None, config=None):
+    """Deploys the specified model locally, for testing. This function should be defined
+    within the plugin module. Also note that this function has a signature which is very
+    similar to :py:meth:`BaseDeploymentClient.create_deployment` since both does logically
+    similar operation.
+
+    .. Note::
+        This function is kept here only for documentation purpose and not implementing the
+        actual feature. It should be implemented in the plugin's top level namescope and should
+        be callable with ``plugin_module.run_local``
+
+    Args:
+        target: Which target to use. This information is used to call the appropriate plugin.
+        name: Unique name to use for deployment. If another deployment exists with the same
+            name, create_deployment will raise a
+            :py:class:`mlflow.exceptions.MlflowException`.
+        model_uri: URI of model to deploy.
+        flavor: (optional) Model flavor to deploy. If unspecified, default flavor is chosen.
+        config: (optional) Dict containing updated target-specific config for the deployment.
+
+    Returns:
+        None
+    """
+    raise NotImplementedError(
+        "This function should be implemented in the deployment plugin. It is "
+        "kept here only for documentation purpose and shouldn't be used in "
+        "your application"
+    )
+
+
+def target_help():
+    """
+    .. Note::
+        This function is kept here only for documentation purpose and not implementing the
+        actual feature. It should be implemented in the plugin's top level namescope and should
+        be callable with ``plugin_module.target_help``
+
+    Return a string containing detailed documentation on the current deployment target, to be
+    displayed when users invoke the ``mlflow deployments help -t <target-name>`` CLI. This
+    method should be defined within the module specified by the plugin author.
+    The string should contain:
+
+    * An explanation of target-specific fields in the ``config`` passed to ``create_deployment``,
+      ``update_deployment``
+    * How to specify a ``target_uri`` (e.g. for AWS SageMaker, ``target_uri`` have a scheme of
+      "sagemaker:/<aws-cli-profile-name>", where aws-cli-profile-name is the name of an AWS
+      CLI profile https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-profiles.html)
+    * Any other target-specific details.
+
+    """
+    raise NotImplementedError(
+        "This function should be implemented in the deployment plugin. It is "
+        "kept here only for documentation purpose and shouldn't be used in "
+        "your application"
+    )
+
+
+@developer_stable
+class BaseDeploymentClient(abc.ABC):
+    """
+    Base class exposing Python model deployment APIs.
+
+    Plugin implementors should define target-specific deployment logic via a subclass of
+    ``BaseDeploymentClient`` within the plugin module, and customize method docstrings with
+    target-specific information.
+
+    .. Note::
+        Subclasses should raise :py:class:`mlflow.exceptions.MlflowException` in error cases (e.g.
+        on failure to deploy a model).
+    """
+
+    def __init__(self, target_uri):
+        self.target_uri = target_uri
+
+    @abc.abstractmethod
+    def create_deployment(self, name, model_uri, flavor=None, config=None, endpoint=None):
+        """
+        Deploy a model to the specified target. By default, this method should block until
+        deployment completes (i.e. until it's possible to perform inference with the deployment).
+        In the case of conflicts (e.g. if it's not possible to create the specified deployment
+        without due to conflict with an existing deployment), raises a
+        :py:class:`mlflow.exceptions.MlflowException` or an `HTTPError` for remote
+        deployments. See target-specific plugin documentation
+        for additional detail on support for asynchronous deployment and other configuration.
+
+        Args:
+            name: Unique name to use for deployment. If another deployment exists with the same
+                name, raises a :py:class:`mlflow.exceptions.MlflowException`
+            model_uri: URI of model to deploy
+            flavor: (optional) Model flavor to deploy. If unspecified, a default flavor
+                will be chosen.
+            config: (optional) Dict containing updated target-specific configuration for the
+                deployment
+            endpoint: (optional) Endpoint to create the deployment under. May not be supported
+                by all targets
+
+        Returns:
+            Dict corresponding to created deployment, which must contain the 'name' key.
+
+        """
+
+    @abc.abstractmethod
+    def update_deployment(self, name, model_uri=None, flavor=None, config=None, endpoint=None):
+        """
+        Update the deployment with the specified name. You can update the URI of the model, the
+        flavor of the deployed model (in which case the model URI must also be specified), and/or
+        any target-specific attributes of the deployment (via `config`). By default, this method
+        should block until deployment completes (i.e. until it's possible to perform inference
+        with the updated deployment). See target-specific plugin documentation for additional
+        detail on support for asynchronous deployment and other configuration.
+
+        Args:
+            name: Unique name of deployment to update.
+            model_uri: URI of a new model to deploy.
+            flavor: (optional) new model flavor to use for deployment. If provided,
+                ``model_uri`` must also be specified. If ``flavor`` is unspecified but
+                ``model_uri`` is specified, a default flavor will be chosen and the
+                deployment will be updated using that flavor.
+            config: (optional) dict containing updated target-specific configuration for the
+                deployment.
+            endpoint: (optional) Endpoint containing the deployment to update. May not be
+                supported by all targets.
+
+        Returns:
+            None
+
+        """
+
+    @abc.abstractmethod
+    def delete_deployment(self, name, config=None, endpoint=None):
+        """Delete the deployment with name ``name`` from the specified target.
+
+        Deletion should be idempotent (i.e. deletion should not fail if retried on a non-existent
+        deployment).
+
+        Args:
+            name: Name of deployment to delete
+            config: (optional) dict containing updated target-specific configuration for the
+                deployment
+            endpoint: (optional) Endpoint containing the deployment to delete. May not be
+                supported by all targets
+
+        Returns:
+            None
+        """
+
+    @abc.abstractmethod
+    def list_deployments(self, endpoint=None):
+        """List deployments.
+
+        This method is expected to return an unpaginated list of all
+        deployments (an alternative would be to return a dict with a 'deployments' field
+        containing the actual deployments, with plugins able to specify other fields, e.g.
+        a next_page_token field, in the returned dictionary for pagination, and to accept
+        a `pagination_args` argument to this method for passing pagination-related args).
+
+        Args:
+            endpoint: (optional) List deployments in the specified endpoint. May not be
+                supported by all targets
+
+        Returns:
+            A list of dicts corresponding to deployments. Each dict is guaranteed to
+            contain a 'name' key containing the deployment name. The other fields of
+            the returned dictionary and their types may vary across deployment targets.
+        """
+
+    @abc.abstractmethod
+    def get_deployment(self, name, endpoint=None):
+        """
+        Returns a dictionary describing the specified deployment, throwing either a
+        :py:class:`mlflow.exceptions.MlflowException` or an `HTTPError` for remote
+        deployments if no deployment exists with the provided ID.
+        The dict is guaranteed to contain an 'name' key containing the deployment name.
+        The other fields of the returned dictionary and their types may vary across
+        deployment targets.
+
+        Args:
+            name: ID of deployment to fetch.
+            endpoint: (optional) Endpoint containing the deployment to get. May not be
+                supported by all targets.
+
+        Returns:
+            A dict corresponding to the retrieved deployment. The dict is guaranteed to
+            contain a 'name' key corresponding to the deployment name. The other fields of
+            the returned dictionary and their types may vary across targets.
+        """
+
+    @abc.abstractmethod
+    def predict(self, deployment_name=None, inputs=None, endpoint=None):
+        """Compute predictions on inputs using the specified deployment or model endpoint.
+
+        Note that the input/output types of this method match those of `mlflow pyfunc predict`.
+
+        Args:
+            deployment_name: Name of deployment to predict against.
+            inputs: Input data (or arguments) to pass to the deployment or model endpoint for
+                inference.
+            endpoint: Endpoint to predict against. May not be supported by all targets.
+
+        Returns:
+            A :py:class:`mlflow.deployments.PredictionsResponse` instance representing the
+            predictions and associated Model Server response metadata.
+
+        """
+
+    def predict_stream(self, deployment_name=None, inputs=None, endpoint=None):
+        """
+        Submit a query to a configured provider endpoint, and get streaming response
+
+        Args:
+            deployment_name: Name of deployment to predict against.
+            inputs: The inputs to the query, as a dictionary.
+            endpoint: The name of the endpoint to query.
+
+        Returns:
+            An iterator of dictionary containing the response from the endpoint.
+        """
+        raise NotImplementedError()
+
+    def explain(self, deployment_name=None, df=None, endpoint=None):
+        """
+        Generate explanations of model predictions on the specified input pandas Dataframe
+        ``df`` for the deployed model. Explanation output formats vary by deployment target,
+        and can include details like feature importance for understanding/debugging predictions.
+
+        Args:
+            deployment_name: Name of deployment to predict against
+            df: Pandas DataFrame to use for explaining feature importance in model prediction
+            endpoint: Endpoint to predict against. May not be supported by all targets
+
+        Returns:
+            A JSON-able object (pandas dataframe, numpy array, dictionary), or
+            an exception if the implementation is not available in deployment target's class
+        """
+        raise MlflowException(
+            "Computing model explanations is not yet supported for this deployment target"
+        )
+
+    def create_endpoint(self, name, config=None):
+        """
+        Create an endpoint with the specified target. By default, this method should block until
+        creation completes (i.e. until it's possible to create a deployment within the endpoint).
+        In the case of conflicts (e.g. if it's not possible to create the specified endpoint
+        due to conflict with an existing endpoint), raises a
+        :py:class:`mlflow.exceptions.MlflowException` or an `HTTPError` for remote
+        deployments. See target-specific plugin documentation
+        for additional detail on support for asynchronous creation and other configuration.
+
+        Args:
+            name: Unique name to use for endpoint. If another endpoint exists with the same
+                name, raises a :py:class:`mlflow.exceptions.MlflowException`.
+            config: (optional) Dict containing target-specific configuration for the
+                endpoint.
+
+        Returns:
+            Dict corresponding to created endpoint, which must contain the 'name' key.
+
+        """
+        raise MlflowException(
+            "Method is unimplemented in base client. Implementation should be "
+            "provided by specific target plugins."
+        )
+
+    def update_endpoint(self, endpoint, config=None):
+        """
+        Update the endpoint with the specified name. You can update any target-specific attributes
+        of the endpoint (via `config`). By default, this method should block until the update
+        completes (i.e. until it's possible to create a deployment within the endpoint). See
+        target-specific plugin documentation for additional detail on support for asynchronous
+        update and other configuration.
+
+        Args:
+            endpoint: Unique name of endpoint to update
+            config: (optional) dict containing target-specific configuration for the
+                endpoint
+
+        Returns:
+            None
+
+        """
+        raise MlflowException(
+            "Method is unimplemented in base client. Implementation should be "
+            "provided by specific target plugins."
+        )
+
+    def delete_endpoint(self, endpoint):
+        """
+        Delete the endpoint from the specified target. Deletion should be idempotent (i.e. deletion
+        should not fail if retried on a non-existent deployment).
+
+        Args:
+            endpoint: Name of endpoint to delete
+
+        Returns:
+            None
+        """
+        raise MlflowException(
+            "Method is unimplemented in base client. Implementation should be "
+            "provided by specific target plugins."
+        )
+
+    def list_endpoints(self):
+        """
+        List endpoints in the specified target. This method is expected to return an
+        unpaginated list of all endpoints (an alternative would be to return a dict with
+        an 'endpoints' field containing the actual endpoints, with plugins able to specify
+        other fields, e.g. a next_page_token field, in the returned dictionary for pagination,
+        and to accept a `pagination_args` argument to this method for passing
+        pagination-related args).
+
+        Returns:
+            A list of dicts corresponding to endpoints. Each dict is guaranteed to
+            contain a 'name' key containing the endpoint name. The other fields of
+            the returned dictionary and their types may vary across targets.
+        """
+        raise MlflowException(
+            "Method is unimplemented in base client. Implementation should be "
+            "provided by specific target plugins."
+        )
+
+    def get_endpoint(self, endpoint):
+        """
+        Returns a dictionary describing the specified endpoint, throwing a
+        py:class:`mlflow.exception.MlflowException` or an `HTTPError` for remote
+        deployments if no endpoint exists with the provided
+        name.
+        The dict is guaranteed to contain an 'name' key containing the endpoint name.
+        The other fields of the returned dictionary and their types may vary across targets.
+
+        Args:
+            endpoint: Name of endpoint to fetch
+
+        Returns:
+            A dict corresponding to the retrieved endpoint. The dict is guaranteed to
+            contain a 'name' key corresponding to the endpoint name. The other fields of
+            the returned dictionary and their types may vary across targets.
+        """
+        raise MlflowException(
+            "Method is unimplemented in base client. Implementation should be "
+            "provided by specific target plugins."
+        )

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

@@ -0,0 +1,510 @@
+import json
+import sys
+import warnings
+from inspect import signature
+
+import click
+
+from mlflow.deployments import interface
+from mlflow.environment_variables import MLFLOW_DEPLOYMENTS_CONFIG
+from mlflow.utils import cli_args
+from mlflow.utils.annotations import experimental
+from mlflow.utils.os import is_windows
+from mlflow.utils.proto_json_utils import NumpyEncoder, _get_jsonable_obj
+
+
+def _user_args_to_dict(user_list):
+    # Similar function in mlflow.cli is throwing exception on import
+    user_dict = {}
+    for s in user_list:
+        try:
+            # Some configs may contain '=' in the value
+            name, value = s.split("=", 1)
+        except ValueError as exc:
+            # not enough values to unpack
+            raise click.BadOptionUsage(
+                "config",
+                "Config options must be a pair and should be "
+                "provided as ``-C key=value`` or "
+                "``--config key=value``",
+            ) from exc
+        if name in user_dict:
+            raise click.ClickException(f"Repeated parameter: '{name}'")
+        user_dict[name] = value
+    return user_dict
+
+
+installed_targets = list(interface.plugin_store.registry)
+if len(installed_targets) > 0:
+    supported_targets_msg = "Support is currently installed for deployment to: {targets}".format(
+        targets=", ".join(installed_targets)
+    )
+else:
+    supported_targets_msg = (
+        "NOTE: you currently do not have support installed for any deployment targets."
+    )
+
+target_details = click.option(
+    "--target",
+    "-t",
+    required=True,
+    help=f"""
+                                   Deployment target URI. Run
+                                   `mlflow deployments help --target-name <target-name>` for
+                                   more details on the supported URI format and config options
+                                   for a given target.
+                                   {supported_targets_msg}
+
+                                   See all supported deployment targets and installation
+                                   instructions at
+                                   https://mlflow.org/docs/latest/plugins.html#community-plugins
+                                   """,
+)
+deployment_name = click.option("--name", "name", required=True, help="Name of the deployment")
+optional_deployment_name = click.option("--name", "name", help="Name of the deployment")
+parse_custom_arguments = click.option(
+    "--config",
+    "-C",
+    metavar="NAME=VALUE",
+    multiple=True,
+    help="Extra target-specific config for the model "
+    "deployment, of the form -C name=value. See "
+    "documentation/help for your deployment target for a "
+    "list of supported config options.",
+)
+
+parse_input = click.option(
+    "--input-path",
+    "-I",
+    required=True,
+    help="Path to input prediction payload file. The file can"
+    "be a JSON (Python Dict) or CSV (pandas DataFrame). If the file is a CSV, the user must specify"
+    "the --content-type csv option.",
+)
+
+parse_output = click.option(
+    "--output-path",
+    "-O",
+    help="File to output results to as a JSON file. If not provided, prints output to stdout.",
+)
+
+required_endpoint_param = click.option("--endpoint", required=True, help="Name of the endpoint")
+optional_endpoint_param = click.option("--endpoint", help="Name of the endpoint")
+
+
+@click.group(
+    "deployments",
+    help=f"""
+    Deploy MLflow models to custom targets.
+    Run `mlflow deployments help --target-name <target-name>` for
+    more details on the supported URI format and config options for a given target.
+    {supported_targets_msg}
+
+    See all supported deployment targets and installation instructions in
+    https://mlflow.org/docs/latest/plugins.html#community-plugins
+
+    You can also write your own plugin for deployment to a custom target. For instructions on
+    writing and distributing a plugin, see
+    https://mlflow.org/docs/latest/plugins.html#writing-your-own-mlflow-plugins.
+""",
+)
+def commands():
+    """
+    Deploy MLflow models to custom targets. Support is currently installed for
+    the following targets: {targets}. Run `mlflow deployments help --target-name <target-name>` for
+    more details on the supported URI format and config options for a given target.
+
+    To deploy to other targets, you must first install an
+    appropriate third-party Python plugin. See the list of known community-maintained plugins
+    at https://mlflow.org/docs/latest/plugins.html#community-plugins.
+
+    You can also write your own plugin for deployment to a custom target. For instructions on
+    writing and distributing a plugin, see
+    https://mlflow.org/docs/latest/plugins.html#writing-your-own-mlflow-plugins.
+    """
+
+
+@commands.command("create")
+@optional_endpoint_param
+@parse_custom_arguments
+@deployment_name
+@target_details
+@cli_args.MODEL_URI
+@click.option(
+    "--flavor",
+    "-f",
+    help="Which flavor to be deployed. This will be auto inferred if it's not given",
+)
+def create_deployment(flavor, model_uri, target, name, config, endpoint):
+    """
+    Deploy the model at ``model_uri`` to the specified target.
+
+    Additional plugin-specific arguments may also be passed to this command, via `-C key=value`
+    """
+    config_dict = _user_args_to_dict(config)
+    client = interface.get_deploy_client(target)
+
+    sig = signature(client.create_deployment)
+    if "endpoint" in sig.parameters:
+        deployment = client.create_deployment(
+            name, model_uri, flavor, config=config_dict, endpoint=endpoint
+        )
+    else:
+        deployment = client.create_deployment(name, model_uri, flavor, config=config_dict)
+    click.echo("\n{} deployment {} is created".format(deployment["flavor"], deployment["name"]))
+
+
+@commands.command("update")
+@optional_endpoint_param
+@parse_custom_arguments
+@deployment_name
+@target_details
+@click.option(
+    "--model-uri",
+    "-m",
+    default=None,
+    metavar="URI",
+    help="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",
+)
+@click.option(
+    "--flavor",
+    "-f",
+    help="Which flavor to be deployed. This will be auto inferred if it's not given",
+)
+def update_deployment(flavor, model_uri, target, name, config, endpoint):
+    """
+    Update the deployment with ID `deployment_id` in the specified target.
+    You can update the URI of the model and/or the flavor of the deployed model (in which case the
+    model URI must also be specified).
+
+    Additional plugin-specific arguments may also be passed to this command, via `-C key=value`.
+    """
+    config_dict = _user_args_to_dict(config)
+    client = interface.get_deploy_client(target)
+
+    sig = signature(client.update_deployment)
+    if "endpoint" in sig.parameters:
+        ret = client.update_deployment(
+            name, model_uri=model_uri, flavor=flavor, config=config_dict, endpoint=endpoint
+        )
+    else:
+        ret = client.update_deployment(name, model_uri=model_uri, flavor=flavor, config=config_dict)
+    click.echo("Deployment {} is updated (with flavor {})".format(name, ret["flavor"]))
+
+
+@commands.command("delete")
+@optional_endpoint_param
+@parse_custom_arguments
+@deployment_name
+@target_details
+def delete_deployment(target, name, config, endpoint):
+    """
+    Delete the deployment with name given at `--name` from the specified target.
+    """
+    client = interface.get_deploy_client(target)
+
+    sig = signature(client.delete_deployment)
+    if "config" in sig.parameters:
+        config_dict = _user_args_to_dict(config)
+        if "endpoint" in sig.parameters:
+            client.delete_deployment(name, config=config_dict, endpoint=endpoint)
+        else:
+            client.delete_deployment(name, config=config_dict)
+    else:
+        if "endpoint" in sig.parameters:
+            client.delete_deployment(name, endpoint=endpoint)
+        else:
+            client.delete_deployment(name)
+
+    click.echo(f"Deployment {name} is deleted")
+
+
+@commands.command("list")
+@optional_endpoint_param
+@target_details
+def list_deployment(target, endpoint):
+    """
+    List the names of all model deployments in the specified target. These names can be used with
+    the `delete`, `update`, and `get` commands.
+    """
+    client = interface.get_deploy_client(target)
+
+    sig = signature(client.list_deployments)
+    if "endpoint" in sig.parameters:
+        ids = client.list_deployments(endpoint=endpoint)
+    else:
+        ids = client.list_deployments()
+    click.echo(f"List of all deployments:\n{ids}")
+
+
+@commands.command("get")
+@optional_endpoint_param
+@deployment_name
+@target_details
+def get_deployment(target, name, endpoint):
+    """
+    Print a detailed description of the deployment with name given at ``--name`` in the specified
+    target.
+    """
+    client = interface.get_deploy_client(target)
+
+    sig = signature(client.get_deployment)
+    if "endpoint" in sig.parameters:
+        desc = client.get_deployment(name, endpoint=endpoint)
+    else:
+        desc = client.get_deployment(name)
+    for key, val in desc.items():
+        click.echo(f"{key}: {val}")
+    click.echo("\n")
+
+
+@commands.command("help")
+@target_details
+def target_help(target):
+    """
+    Display additional help for a specific deployment target, e.g. info on target-specific config
+    options and the target's URI format.
+    """
+    click.echo(interface._target_help(target))
+
+
+@commands.command("run-local")
+@parse_custom_arguments
+@deployment_name
+@target_details
+@cli_args.MODEL_URI
+@click.option(
+    "--flavor",
+    "-f",
+    help="Which flavor to be deployed. This will be auto inferred if it's not given",
+)
+def run_local(flavor, model_uri, target, name, config):
+    """
+    Deploy the model locally. This has very similar signature to ``create`` API
+    """
+    config_dict = _user_args_to_dict(config)
+    interface.run_local(target, name, model_uri, flavor, config_dict)
+
+
+def predictions_to_json(raw_predictions, output):
+    predictions = _get_jsonable_obj(raw_predictions, pandas_orient="records")
+    json.dump(predictions, output, cls=NumpyEncoder)
+
+
+@commands.command("predict")
+@click.option(
+    "--name",
+    "name",
+    help="Name of the deployment. Exactly one of --name or --endpoint must be specified.",
+)
+@click.option(
+    "--endpoint",
+    help="Name of the endpoint. Exactly one of --name or --endpoint must be specified.",
+)
+@target_details
+@parse_input
+@parse_output
+def predict(target, name, input_path, output_path, endpoint):
+    """
+    Predict the results for the deployed model for the given input(s)
+    """
+    import pandas as pd
+
+    if (name, endpoint).count(None) != 1:
+        raise click.UsageError("Must specify exactly one of --name or --endpoint.")
+
+    df = pd.read_json(input_path)
+    client = interface.get_deploy_client(target)
+
+    sig = signature(client.predict)
+    if "endpoint" in sig.parameters:
+        result = client.predict(name, df, endpoint=endpoint)
+    else:
+        result = client.predict(name, df)
+    if output_path is not None:
+        result.to_json(output_path)
+    else:
+        click.echo(result.to_json())
+
+
+@commands.command("explain")
+@click.option(
+    "--name",
+    "name",
+    help="Name of the deployment. Exactly one of --name or --endpoint must be specified.",
+)
+@click.option(
+    "--endpoint",
+    help="Name of the endpoint. Exactly one of --name or --endpoint must be specified.",
+)
+@target_details
+@parse_input
+@parse_output
+def explain(target, name, input_path, output_path, endpoint):
+    """
+    Generate explanations of model predictions on the specified input for
+    the deployed model for the given input(s). Explanation output formats vary
+    by deployment target, and can include details like feature importance for
+    understanding/debugging predictions. Run `mlflow deployments help` or
+    consult the documentation for your plugin for details on explanation format.
+    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
+    """
+    import pandas as pd
+
+    if (name, endpoint).count(None) != 1:
+        raise click.UsageError("Must specify exactly one of --name or --endpoint.")
+
+    df = pd.read_json(input_path)
+    client = interface.get_deploy_client(target)
+
+    sig = signature(client.explain)
+    if "endpoint" in sig.parameters:
+        result = client.explain(name, df, endpoint=endpoint)
+    else:
+        result = client.explain(name, df)
+    if output_path:
+        with open(output_path, "w") as fp:
+            predictions_to_json(result, fp)
+    else:
+        predictions_to_json(result, sys.stdout)
+
+
+@commands.command("create-endpoint")
+@click.option(
+    "--config",
+    "-C",
+    metavar="NAME=VALUE",
+    multiple=True,
+    help="Extra target-specific config for the endpoint, "
+    "of the form -C name=value. See "
+    "documentation/help for your deployment target for a "
+    "list of supported config options.",
+)
+@required_endpoint_param
+@target_details
+def create_endpoint(target, name, config):
+    """
+    Create an endpoint with the specified name at the specified target.
+
+    Additional plugin-specific arguments may also be passed to this command, via `-C key=value`
+    """
+    config_dict = _user_args_to_dict(config)
+    client = interface.get_deploy_client(target)
+    endpoint = client.create_endpoint(name, config=config_dict)
+    click.echo("\nEndpoint {} is created".format(endpoint["name"]))
+
+
+@commands.command("update-endpoint")
+@click.option(
+    "--config",
+    "-C",
+    metavar="NAME=VALUE",
+    multiple=True,
+    help="Extra target-specific config for the endpoint, "
+    "of the form -C name=value. See "
+    "documentation/help for your deployment target for a "
+    "list of supported config options.",
+)
+@required_endpoint_param
+@target_details
+def update_endpoint(target, endpoint, config):
+    """
+    Update the specified endpoint at the specified target.
+
+    Additional plugin-specific arguments may also be passed to this command, via `-C key=value`
+    """
+    config_dict = _user_args_to_dict(config)
+    client = interface.get_deploy_client(target)
+    client.update_endpoint(endpoint, config=config_dict)
+    click.echo(f"\nEndpoint {endpoint} is updated")
+
+
+@commands.command("delete-endpoint")
+@required_endpoint_param
+@target_details
+def delete_endpoint(target, endpoint):
+    """
+    Delete the specified endpoint at the specified target
+    """
+    client = interface.get_deploy_client(target)
+    client.delete_endpoint(endpoint)
+    click.echo(f"\nEndpoint {endpoint} is deleted")
+
+
+@commands.command("list-endpoints")
+@target_details
+def list_endpoints(target):
+    """
+    List all endpoints at the specified target
+    """
+    client = interface.get_deploy_client(target)
+    ids = client.list_endpoints()
+    click.echo(f"List of all endpoints:\n{ids}")
+
+
+@commands.command("get-endpoint")
+@required_endpoint_param
+@target_details
+def get_endpoint(target, endpoint):
+    """
+    Get details for the specified endpoint at the specified target
+    """
+    client = interface.get_deploy_client(target)
+    desc = client.get_endpoint(endpoint)
+    for key, val in desc.items():
+        click.echo(f"{key}: {val}")
+    click.echo("\n")
+
+
+def validate_config_path(_ctx, _param, value):
+    from mlflow.gateway.config import _validate_config
+
+    try:
+        _validate_config(value)
+        return value
+    except Exception as e:
+        raise click.BadParameter(str(e))
+
+
+@experimental
+@commands.command("start-server", help="Start MLflow AI Gateway")
+@click.option(
+    "--config-path",
+    envvar=MLFLOW_DEPLOYMENTS_CONFIG.name,
+    callback=validate_config_path,
+    required=True,
+    help="The path to the deployments configuration file.",
+)
+@click.option(
+    "--host",
+    default="127.0.0.1",
+    help="The network address to listen on (default: 127.0.0.1).",
+)
+@click.option(
+    "--port",
+    default=5000,
+    help="The port to listen on (default: 5000).",
+)
+@click.option(
+    "--workers",
+    default=2,
+    help="The number of workers.",
+)
+def start_server(config_path: str, host: str, port: str, workers: int):
+    warnings.warn(
+        "`mlflow deployments start-server` is deprecated and will be removed in a future release. "
+        "Use `mlflow gateway start` instead.",
+        FutureWarning,
+    )
+    if is_windows():
+        raise click.ClickException("MLflow AI Gateway does not support Windows.")
+
+    from mlflow.gateway.runner import run_app
+
+    run_app(config_path=config_path, host=host, port=port, workers=workers)

venv/lib/python3.9/site-packages/mlflow/deployments/constants.py

@@ -0,0 +1,13 @@
+# Abridged retryable error codes for deployments clients.
+# These are modified from the standard MLflow Tracking server retry codes for the MLflowClient to
+# remove timeouts from the list of the retryable conditions. A long-running timeout with
+# retries for the proxied providers generally indicates an issue with the underlying query or
+# the model being served having issues responding to the query due to parameter configuration.
+MLFLOW_DEPLOYMENT_CLIENT_REQUEST_RETRY_CODES = frozenset(
+    [
+        429,  # Too many requests
+        500,  # Server Error
+        502,  # Bad Gateway
+        503,  # Service Unavailable
+    ]
+)

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

@@ -0,0 +1,835 @@
+import json
+import posixpath
+import warnings
+from typing import Any, Iterator, Optional
+
+from mlflow.deployments import BaseDeploymentClient
+from mlflow.deployments.constants import (
+    MLFLOW_DEPLOYMENT_CLIENT_REQUEST_RETRY_CODES,
+)
+from mlflow.environment_variables import (
+    MLFLOW_DEPLOYMENT_PREDICT_TIMEOUT,
+    MLFLOW_HTTP_REQUEST_TIMEOUT,
+)
+from mlflow.exceptions import MlflowException
+from mlflow.utils import AttrDict
+from mlflow.utils.annotations import deprecated, experimental
+from mlflow.utils.databricks_utils import get_databricks_host_creds
+from mlflow.utils.rest_utils import augmented_raise_for_status, http_request
+
+
+class DatabricksEndpoint(AttrDict):
+    """
+    A dictionary-like object representing a Databricks serving endpoint.
+
+    .. code-block:: python
+
+        endpoint = DatabricksEndpoint(
+            {
+                "name": "chat",
+                "creator": "alice@company.com",
+                "creation_timestamp": 0,
+                "last_updated_timestamp": 0,
+                "state": {...},
+                "config": {...},
+                "tags": [...],
+                "id": "88fd3f75a0d24b0380ddc40484d7a31b",
+            }
+        )
+        assert endpoint.name == "chat"
+    """
+
+
+@experimental
+class DatabricksDeploymentClient(BaseDeploymentClient):
+    """
+    Client for interacting with Databricks serving endpoints.
+
+    Example:
+
+    First, set up credentials for authentication:
+
+    .. code-block:: bash
+
+        export DATABRICKS_HOST=...
+        export DATABRICKS_TOKEN=...
+
+    .. seealso::
+
+        See https://docs.databricks.com/en/dev-tools/auth.html for other authentication methods.
+
+    Then, create a deployment client and use it to interact with Databricks serving endpoints:
+
+    .. code-block:: python
+
+        from mlflow.deployments import get_deploy_client
+
+        client = get_deploy_client("databricks")
+        endpoints = client.list_endpoints()
+        assert endpoints == [
+            {
+                "name": "chat",
+                "creator": "alice@company.com",
+                "creation_timestamp": 0,
+                "last_updated_timestamp": 0,
+                "state": {...},
+                "config": {...},
+                "tags": [...],
+                "id": "88fd3f75a0d24b0380ddc40484d7a31b",
+            },
+        ]
+    """
+
+    def create_deployment(self, name, model_uri, flavor=None, config=None, endpoint=None):
+        """
+        .. warning::
+
+            This method is not implemented for `DatabricksDeploymentClient`.
+        """
+        raise NotImplementedError
+
+    def update_deployment(self, name, model_uri=None, flavor=None, config=None, endpoint=None):
+        """
+        .. warning::
+
+            This method is not implemented for `DatabricksDeploymentClient`.
+        """
+        raise NotImplementedError
+
+    def delete_deployment(self, name, config=None, endpoint=None):
+        """
+        .. warning::
+
+            This method is not implemented for `DatabricksDeploymentClient`.
+        """
+        raise NotImplementedError
+
+    def list_deployments(self, endpoint=None):
+        """
+        .. warning::
+
+            This method is not implemented for `DatabricksDeploymentClient`.
+        """
+        raise NotImplementedError
+
+    def get_deployment(self, name, endpoint=None):
+        """
+        .. warning::
+
+            This method is not implemented for `DatabricksDeploymentClient`.
+        """
+        raise NotImplementedError
+
+    def _call_endpoint(
+        self,
+        *,
+        method: str,
+        prefix: str = "/api/2.0",
+        route: Optional[str] = None,
+        json_body: Optional[dict[str, Any]] = None,
+        timeout: Optional[int] = None,
+    ):
+        call_kwargs = {}
+        if method.lower() == "get":
+            call_kwargs["params"] = json_body
+        else:
+            call_kwargs["json"] = json_body
+
+        response = http_request(
+            host_creds=get_databricks_host_creds(self.target_uri),
+            endpoint=posixpath.join(prefix, "serving-endpoints", route or ""),
+            method=method,
+            timeout=MLFLOW_HTTP_REQUEST_TIMEOUT.get() if timeout is None else timeout,
+            raise_on_status=False,
+            retry_codes=MLFLOW_DEPLOYMENT_CLIENT_REQUEST_RETRY_CODES,
+            extra_headers={"X-Databricks-Endpoints-API-Client": "Databricks Deployment Client"},
+            **call_kwargs,
+        )
+        augmented_raise_for_status(response)
+        return DatabricksEndpoint(response.json())
+
+    def _call_endpoint_stream(
+        self,
+        *,
+        method: str,
+        prefix: str = "/api/2.0",
+        route: Optional[str] = None,
+        json_body: Optional[dict[str, Any]] = None,
+        timeout: Optional[int] = None,
+    ) -> Iterator[str]:
+        call_kwargs = {}
+        if method.lower() == "get":
+            call_kwargs["params"] = json_body
+        else:
+            call_kwargs["json"] = json_body
+
+        response = http_request(
+            host_creds=get_databricks_host_creds(self.target_uri),
+            endpoint=posixpath.join(prefix, "serving-endpoints", route or ""),
+            method=method,
+            timeout=MLFLOW_HTTP_REQUEST_TIMEOUT.get() if timeout is None else timeout,
+            raise_on_status=False,
+            retry_codes=MLFLOW_DEPLOYMENT_CLIENT_REQUEST_RETRY_CODES,
+            extra_headers={"X-Databricks-Endpoints-API-Client": "Databricks Deployment Client"},
+            stream=True,  # Receive response content in streaming way.
+            **call_kwargs,
+        )
+        augmented_raise_for_status(response)
+
+        # Streaming response content are composed of multiple lines.
+        # Each line format depends on specific endpoint
+        return (
+            line.strip()
+            for line in response.iter_lines(decode_unicode=True)
+            if line.strip()  # filter out keep-alive new lines
+        )
+
+    @experimental
+    def predict(self, deployment_name=None, inputs=None, endpoint=None):
+        """
+        Query a serving endpoint with the provided model inputs.
+        See https://docs.databricks.com/api/workspace/servingendpoints/query for request/response
+        schema.
+
+        Args:
+            deployment_name: Unused.
+            inputs: A dictionary containing the model inputs to query.
+            endpoint: The name of the serving endpoint to query.
+
+        Returns:
+            A :py:class:`DatabricksEndpoint` object containing the query response.
+
+        Example:
+
+        .. code-block:: python
+
+            from mlflow.deployments import get_deploy_client
+
+            client = get_deploy_client("databricks")
+            response = client.predict(
+                endpoint="chat",
+                inputs={
+                    "messages": [
+                        {"role": "user", "content": "Hello!"},
+                    ],
+                },
+            )
+            assert response == {
+                "id": "chatcmpl-8OLm5kfqBAJD8CpsMANESWKpLSLXY",
+                "object": "chat.completion",
+                "created": 1700814265,
+                "model": "gpt-4-0613",
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": "Hello! How can I assist you today?",
+                        },
+                        "finish_reason": "stop",
+                    }
+                ],
+                "usage": {
+                    "prompt_tokens": 9,
+                    "completion_tokens": 9,
+                    "total_tokens": 18,
+                },
+            }
+        """
+        return self._call_endpoint(
+            method="POST",
+            prefix="/",
+            route=posixpath.join(endpoint, "invocations"),
+            json_body=inputs,
+            timeout=MLFLOW_DEPLOYMENT_PREDICT_TIMEOUT.get(),
+        )
+
+    @experimental
+    def predict_stream(
+        self, deployment_name=None, inputs=None, endpoint=None
+    ) -> Iterator[dict[str, Any]]:
+        """
+        Submit a query to a configured provider endpoint, and get streaming response
+
+        Args:
+            deployment_name: Unused.
+            inputs: The inputs to the query, as a dictionary.
+            endpoint: The name of the endpoint to query.
+
+        Returns:
+            An iterator of dictionary containing the response from the endpoint.
+
+        Example:
+
+        .. code-block:: python
+
+            from mlflow.deployments import get_deploy_client
+
+            client = get_deploy_client("databricks")
+            chunk_iter = client.predict_stream(
+                endpoint="databricks-llama-2-70b-chat",
+                inputs={
+                    "messages": [{"role": "user", "content": "Hello!"}],
+                    "temperature": 0.0,
+                    "n": 1,
+                    "max_tokens": 500,
+                },
+            )
+            for chunk in chunk_iter:
+                print(chunk)
+                # Example:
+                # {
+                #     "id": "82a834f5-089d-4fc0-ad6c-db5c7d6a6129",
+                #     "object": "chat.completion.chunk",
+                #     "created": 1712133837,
+                #     "model": "llama-2-70b-chat-030424",
+                #     "choices": [
+                #         {
+                #             "index": 0, "delta": {"role": "assistant", "content": "Hello"},
+                #             "finish_reason": None,
+                #         }
+                #     ],
+                #     "usage": {"prompt_tokens": 11, "completion_tokens": 1, "total_tokens": 12},
+                # }
+        """
+        inputs = inputs or {}
+
+        # Add stream=True param in request body to get streaming response
+        # See https://docs.databricks.com/api/workspace/servingendpoints/query#stream
+        chunk_line_iter = self._call_endpoint_stream(
+            method="POST",
+            prefix="/",
+            route=posixpath.join(endpoint, "invocations"),
+            json_body={**inputs, "stream": True},
+            timeout=MLFLOW_DEPLOYMENT_PREDICT_TIMEOUT.get(),
+        )
+
+        for line in chunk_line_iter:
+            splits = line.split(":", 1)
+            if len(splits) < 2:
+                raise MlflowException(
+                    f"Unknown response format: '{line}', "
+                    "expected 'data: <value>' for streaming response."
+                )
+            key, value = splits
+            if key != "data":
+                raise MlflowException(
+                    f"Unknown response format with key '{key}'. "
+                    f"Expected 'data: <value>' for streaming response, got '{line}'."
+                )
+
+            value = value.strip()
+            if value == "[DONE]":
+                # Databricks endpoint streaming response ends with
+                # a line of "data: [DONE]"
+                return
+
+            yield json.loads(value)
+
+    @experimental
+    def create_endpoint(self, name=None, config=None, route_optimized=False):
+        """
+        Create a new serving endpoint with the provided name and configuration.
+
+        See https://docs.databricks.com/api/workspace/servingendpoints/create for request/response
+        schema.
+
+        Args:
+            name: The name of the serving endpoint to create.
+
+                .. warning::
+                    Deprecated. Include `name` in `config` instead.
+
+            config: A dictionary containing either the full API request payload
+                or the configuration of the serving endpoint to create.
+            route_optimized: A boolean which defines whether databricks serving endpoint
+                is optimized for routing traffic. Only used in the deprecated approach.
+
+                .. warning::
+                    Deprecated. Include `route_optimized` in `config` instead.
+
+        Returns:
+            A :py:class:`DatabricksEndpoint` object containing the request response.
+
+        Example:
+
+        .. code-block:: python
+
+            from mlflow.deployments import get_deploy_client
+
+            client = get_deploy_client("databricks")
+            endpoint = client.create_endpoint(
+                config={
+                    "name": "test",
+                    "config": {
+                        "served_entities": [
+                            {
+                                "external_model": {
+                                    "name": "gpt-4",
+                                    "provider": "openai",
+                                    "task": "llm/v1/chat",
+                                    "openai_config": {
+                                        "openai_api_key": "{{secrets/scope/key}}",
+                                    },
+                                },
+                            }
+                        ],
+                        "route_optimized": True,
+                    },
+                },
+            )
+            assert endpoint == {
+                "name": "test",
+                "creator": "alice@company.com",
+                "creation_timestamp": 0,
+                "last_updated_timestamp": 0,
+                "state": {...},
+                "config": {...},
+                "tags": [...],
+                "id": "88fd3f75a0d24b0380ddc40484d7a31b",
+                "permission_level": "CAN_MANAGE",
+                "route_optimized": False,
+                "task": "llm/v1/chat",
+                "endpoint_type": "EXTERNAL_MODEL",
+                "creator_display_name": "Alice",
+                "creator_kind": "User",
+            }
+
+        """
+        warnings_list = []
+
+        if config and "config" in config:
+            # Using new style: full API request payload
+            payload = config.copy()
+
+            # Validate name conflicts
+            if "name" in payload:
+                if name is not None:
+                    if payload["name"] == name:
+                        warnings_list.append(
+                            "Passing 'name' as a parameter is deprecated. "
+                            "Please specify 'name' only within the config dictionary."
+                        )
+                    else:
+                        raise MlflowException(
+                            f"Name mismatch. Found '{name}' as parameter and '{payload['name']}' "
+                            "in config. Please specify 'name' only within the config dictionary "
+                            "as this parameter is deprecated."
+                        )
+            else:
+                if name is None:
+                    raise MlflowException(
+                        "The 'name' field is required. Please specify it within the config "
+                        "dictionary."
+                    )
+                payload["name"] = name
+                warnings_list.append(
+                    "Passing 'name' as a parameter is deprecated. "
+                    "Please specify 'name' within the config dictionary."
+                )
+
+            # Validate route_optimized conflicts
+            if "route_optimized" in payload:
+                if route_optimized is not None:
+                    if payload["route_optimized"] != route_optimized:
+                        raise MlflowException(
+                            "Conflicting 'route_optimized' values found. "
+                            "Please specify 'route_optimized' only within the config dictionary "
+                            "as this parameter is deprecated."
+                        )
+                    warnings_list.append(
+                        "Passing 'route_optimized' as a parameter is deprecated. "
+                        "Please specify 'route_optimized' only within the config dictionary."
+                    )
+            else:
+                if route_optimized:
+                    payload["route_optimized"] = route_optimized
+                    warnings_list.append(
+                        "Passing 'route_optimized' as a parameter is deprecated. "
+                        "Please specify 'route_optimized' within the config dictionary."
+                    )
+        else:
+            # Handle legacy format (backwards compatibility)
+            warnings_list.append(
+                "Passing 'name', 'config', and 'route_optimized' as separate parameters is "
+                "deprecated. Please pass the full API request payload as a single dictionary "
+                "in the 'config' parameter."
+            )
+            config = config.copy() if config else {}  # avoid mutating config
+            extras = {}
+            for key in ("tags", "rate_limits"):
+                if tags := config.pop(key, None):
+                    extras[key] = tags
+            payload = {"name": name, "config": config, "route_optimized": route_optimized, **extras}
+
+        if warnings_list:
+            warnings.warn("\n".join(warnings_list), UserWarning)
+
+        return self._call_endpoint(method="POST", json_body=payload)
+
+    @deprecated(
+        alternative=(
+            "update_endpoint_config, update_endpoint_tags, update_endpoint_rate_limits, "
+            "or update_endpoint_ai_gateway"
+        )
+    )
+    def update_endpoint(self, endpoint, config=None):
+        """
+        Update a specified serving endpoint with the provided configuration.
+        See https://docs.databricks.com/api/workspace/servingendpoints/updateconfig for
+        request/response schema.
+
+        Args:
+            endpoint: The name of the serving endpoint to update.
+            config: A dictionary containing the configuration of the serving endpoint to update.
+
+        Returns:
+            A :py:class:`DatabricksEndpoint` object containing the request response.
+
+        Example:
+
+        .. code-block:: python
+
+            from mlflow.deployments import get_deploy_client
+
+            client = get_deploy_client("databricks")
+            endpoint = client.update_endpoint(
+                endpoint="chat",
+                config={
+                    "served_entities": [
+                        {
+                            "name": "test",
+                            "external_model": {
+                                "name": "gpt-4",
+                                "provider": "openai",
+                                "task": "llm/v1/chat",
+                                "openai_config": {
+                                    "openai_api_key": "{{secrets/scope/key}}",
+                                },
+                            },
+                        }
+                    ],
+                },
+            )
+            assert endpoint == {
+                "name": "chat",
+                "creator": "alice@company.com",
+                "creation_timestamp": 0,
+                "last_updated_timestamp": 0,
+                "state": {...},
+                "config": {...},
+                "tags": [...],
+                "id": "88fd3f75a0d24b0380ddc40484d7a31b",
+            }
+
+            rate_limits = client.update_endpoint(
+                endpoint="chat",
+                config={
+                    "rate_limits": [
+                        {
+                            "key": "user",
+                            "renewal_period": "minute",
+                            "calls": 10,
+                        }
+                    ],
+                },
+            )
+            assert rate_limits == {
+                "rate_limits": [
+                    {
+                        "key": "user",
+                        "renewal_period": "minute",
+                        "calls": 10,
+                    }
+                ],
+            }
+        """
+        warnings.warn(
+            "The `update_endpoint` method is deprecated. Use the specific update methods—"
+            "`update_endpoint_config`, `update_endpoint_tags`, `update_endpoint_rate_limits`, "
+            "`update_endpoint_ai_gateway`—instead.",
+            UserWarning,
+        )
+
+        if list(config) == ["rate_limits"]:
+            return self._call_endpoint(
+                method="PUT", route=posixpath.join(endpoint, "rate-limits"), json_body=config
+            )
+        else:
+            return self._call_endpoint(
+                method="PUT", route=posixpath.join(endpoint, "config"), json_body=config
+            )
+
+    @experimental
+    def update_endpoint_config(self, endpoint, config):
+        """
+        Update the configuration of a specified serving endpoint. See
+        https://docs.databricks.com/api/workspace/servingendpoints/updateconfig for request/response
+        request/response schema.
+
+        Args:
+            endpoint: The name of the serving endpoint to update.
+            config: A dictionary containing the configuration of the serving endpoint to update.
+
+        Returns:
+            A :py:class:`DatabricksEndpoint` object containing the request response.
+
+        Example:
+
+        .. code-block:: python
+
+            from mlflow.deployments import get_deploy_client
+
+            client = get_deploy_client("databricks")
+            updated_endpoint = client.update_endpoint_config(
+                endpoint="test",
+                config={
+                    "served_entities": [
+                        {
+                            "name": "gpt-4o-mini",
+                            "external_model": {
+                                "name": "gpt-4o-mini",
+                                "provider": "openai",
+                                "task": "llm/v1/chat",
+                                "openai_config": {
+                                    "openai_api_key": "{{secrets/scope/key}}",
+                                },
+                            },
+                        }
+                    ]
+                },
+            )
+            assert updated_endpoint == {
+                "name": "test",
+                "creator": "alice@company.com",
+                "creation_timestamp": 1729527763000,
+                "last_updated_timestamp": 1729530896000,
+                "state": {"ready": "READY", "config_update": "NOT_UPDATING"},
+                "config": {...},
+                "id": "44b258fb39804564b37603d8d14b853e",
+                "permission_level": "CAN_MANAGE",
+                "route_optimized": False,
+                "task": "llm/v1/chat",
+                "endpoint_type": "EXTERNAL_MODEL",
+                "creator_display_name": "Alice",
+                "creator_kind": "User",
+            }
+        """
+
+        return self._call_endpoint(
+            method="PUT", route=posixpath.join(endpoint, "config"), json_body=config
+        )
+
+    @experimental
+    def update_endpoint_tags(self, endpoint, config):
+        """
+        Update the tags of a specified serving endpoint. See
+        https://docs.databricks.com/api/workspace/servingendpoints/patch for request/response
+        schema.
+
+        Args:
+            endpoint: The name of the serving endpoint to update.
+            config: A dictionary containing tags to add and/or remove.
+
+        Returns:
+            A :py:class:`DatabricksEndpoint` object containing the request response.
+
+        Example:
+
+        .. code-block:: python
+
+            from mlflow.deployments import get_deploy_client
+
+            client = get_deploy_client("databricks")
+            updated_tags = client.update_endpoint_tags(
+                endpoint="test", config={"add_tags": [{"key": "project", "value": "test"}]}
+            )
+            assert updated_tags == {"tags": [{"key": "project", "value": "test"}]}
+        """
+        return self._call_endpoint(
+            method="PATCH", route=posixpath.join(endpoint, "tags"), json_body=config
+        )
+
+    @experimental
+    def update_endpoint_rate_limits(self, endpoint, config):
+        """
+        Update the rate limits of a specified serving endpoint.
+        See https://docs.databricks.com/api/workspace/servingendpoints/put for request/response
+        schema.
+
+        Args:
+            endpoint: The name of the serving endpoint to update.
+            config: A dictionary containing the updated rate limit configuration.
+
+        Returns:
+            A :py:class:`DatabricksEndpoint` object containing the updated rate limits.
+
+        Example:
+
+        .. code-block:: python
+
+            from mlflow.deployments import get_deploy_client
+
+            client = get_deploy_client("databricks")
+            name = "databricks-dbrx-instruct"
+            rate_limits = {
+                "rate_limits": [{"calls": 10, "key": "endpoint", "renewal_period": "minute"}]
+            }
+            updated_rate_limits = client.update_endpoint_rate_limits(
+                endpoint=name, config=rate_limits
+            )
+            assert updated_rate_limits == {
+                "rate_limits": [{"calls": 10, "key": "endpoint", "renewal_period": "minute"}]
+            }
+        """
+        return self._call_endpoint(
+            method="PUT", route=posixpath.join(endpoint, "rate-limits"), json_body=config
+        )
+
+    @experimental
+    def update_endpoint_ai_gateway(self, endpoint, config):
+        """
+        Update the AI Gateway configuration of a specified serving endpoint.
+
+        Args:
+            endpoint (str): The name of the serving endpoint to update.
+            config (dict): A dictionary containing the AI Gateway configuration to update.
+
+        Returns:
+            dict: A dictionary containing the updated AI Gateway configuration.
+
+        Example:
+
+        .. code-block:: python
+
+            from mlflow.deployments import get_deploy_client
+
+            client = get_deploy_client("databricks")
+            name = "test"
+
+            gateway_config = {
+                "usage_tracking_config": {"enabled": True},
+                "inference_table_config": {
+                    "enabled": True,
+                    "catalog_name": "my_catalog",
+                    "schema_name": "my_schema",
+                },
+            }
+
+            updated_gateway = client.update_endpoint_ai_gateway(
+                endpoint=name, config=gateway_config
+            )
+            assert updated_gateway == {
+                "usage_tracking_config": {"enabled": True},
+                "inference_table_config": {
+                    "catalog_name": "my_catalog",
+                    "schema_name": "my_schema",
+                    "table_name_prefix": "test",
+                    "enabled": True,
+                },
+            }
+        """
+        return self._call_endpoint(
+            method="PUT", route=posixpath.join(endpoint, "ai-gateway"), json_body=config
+        )
+
+    @experimental
+    def delete_endpoint(self, endpoint):
+        """
+        Delete a specified serving endpoint.
+        See https://docs.databricks.com/api/workspace/servingendpoints/delete for request/response
+        schema.
+
+        Args:
+            endpoint: The name of the serving endpoint to delete.
+
+        Returns:
+            A DatabricksEndpoint object containing the request response.
+
+        Example:
+
+        .. code-block:: python
+
+            from mlflow.deployments import get_deploy_client
+
+            client = get_deploy_client("databricks")
+            client.delete_endpoint(endpoint="chat")
+        """
+        return self._call_endpoint(method="DELETE", route=endpoint)
+
+    @experimental
+    def list_endpoints(self):
+        """
+        Retrieve all serving endpoints.
+
+        See https://docs.databricks.com/api/workspace/servingendpoints/list for request/response
+        schema.
+
+        Returns:
+            A list of :py:class:`DatabricksEndpoint` objects containing the request response.
+
+        Example:
+
+        .. code-block:: python
+
+            from mlflow.deployments import get_deploy_client
+
+            client = get_deploy_client("databricks")
+            endpoints = client.list_endpoints()
+            assert endpoints == [
+                {
+                    "name": "chat",
+                    "creator": "alice@company.com",
+                    "creation_timestamp": 0,
+                    "last_updated_timestamp": 0,
+                    "state": {...},
+                    "config": {...},
+                    "tags": [...],
+                    "id": "88fd3f75a0d24b0380ddc40484d7a31b",
+                },
+            ]
+
+        """
+        return self._call_endpoint(method="GET").endpoints
+
+    @experimental
+    def get_endpoint(self, endpoint):
+        """
+        Get a specified serving endpoint.
+        See https://docs.databricks.com/api/workspace/servingendpoints/get for request/response
+        schema.
+
+        Args:
+            endpoint: The name of the serving endpoint to get.
+
+        Returns:
+            A DatabricksEndpoint object containing the request response.
+
+        Example:
+
+        .. code-block:: python
+
+            from mlflow.deployments import get_deploy_client
+
+            client = get_deploy_client("databricks")
+            endpoint = client.get_endpoint(endpoint="chat")
+            assert endpoint == {
+                "name": "chat",
+                "creator": "alice@company.com",
+                "creation_timestamp": 0,
+                "last_updated_timestamp": 0,
+                "state": {...},
+                "config": {...},
+                "tags": [...],
+                "id": "88fd3f75a0d24b0380ddc40484d7a31b",
+            }
+        """
+        return self._call_endpoint(method="GET", route=endpoint)
+
+
+def run_local(name, model_uri, flavor=None, config=None):
+    pass
+
+
+def target_help():
+    pass

venv/lib/python3.9/site-packages/mlflow/deployments/interface.py

@@ -0,0 +1,102 @@
+import inspect
+from logging import Logger
+
+from mlflow.deployments.base import BaseDeploymentClient
+from mlflow.deployments.plugin_manager import DeploymentPlugins
+from mlflow.deployments.utils import get_deployments_target, parse_target_uri
+from mlflow.exceptions import MlflowException
+
+plugin_store = DeploymentPlugins()
+plugin_store.register("sagemaker", "mlflow.sagemaker")
+
+_logger = Logger(__name__)
+
+
+def get_deploy_client(target_uri=None):
+    """Returns a subclass of :py:class:`mlflow.deployments.BaseDeploymentClient` exposing standard
+    APIs for deploying models to the specified target. See available deployment APIs
+    by calling ``help()`` on the returned object or viewing docs for
+    :py:class:`mlflow.deployments.BaseDeploymentClient`. You can also run
+    ``mlflow deployments help -t <target-uri>`` via the CLI for more details on target-specific
+    configuration options.
+
+    Args:
+        target_uri: Optional URI of target to deploy to. If no target URI is provided, then
+            MLflow will attempt to get the deployments target set via `get_deployments_target()` or
+            `MLFLOW_DEPLOYMENTS_TARGET` environment variable.
+
+    .. code-block:: python
+        :caption: Example
+
+        from mlflow.deployments import get_deploy_client
+        import pandas as pd
+
+        client = get_deploy_client("redisai")
+        # Deploy the model stored at artifact path 'myModel' under run with ID 'someRunId'. The
+        # model artifacts are fetched from the current tracking server and then used for deployment.
+        client.create_deployment("spamDetector", "runs:/someRunId/myModel")
+        # Load a CSV of emails and score it against our deployment
+        emails_df = pd.read_csv("...")
+        prediction_df = client.predict_deployment("spamDetector", emails_df)
+        # List all deployments, get details of our particular deployment
+        print(client.list_deployments())
+        print(client.get_deployment("spamDetector"))
+        # Update our deployment to serve a different model
+        client.update_deployment("spamDetector", "runs:/anotherRunId/myModel")
+        # Delete our deployment
+        client.delete_deployment("spamDetector")
+    """
+    if not target_uri:
+        try:
+            target_uri = get_deployments_target()
+        except MlflowException:
+            _logger.info(
+                "No deployments target has been set. Please either set the MLflow deployments "
+                "target via `mlflow.deployments.set_deployments_target()` or set the environment "
+                "variable MLFLOW_DEPLOYMENTS_TARGET to the running deployment server's uri"
+            )
+            return None
+    target = parse_target_uri(target_uri)
+    plugin = plugin_store[target]
+    for _, obj in inspect.getmembers(plugin):
+        if inspect.isclass(obj):
+            if issubclass(obj, BaseDeploymentClient) and not obj == BaseDeploymentClient:
+                return obj(target_uri)
+
+
+def run_local(target, name, model_uri, flavor=None, config=None):
+    """Deploys the specified model locally, for testing. Note that models deployed locally cannot
+    be managed by other deployment APIs (e.g. ``update_deployment``, ``delete_deployment``, etc).
+
+    Args:
+        target: Target to deploy to.
+        name: Name to use for deployment
+        model_uri: URI of model to deploy
+        flavor: (optional) Model flavor to deploy. If unspecified, a default flavor
+            will be chosen.
+        config: (optional) Dict containing updated target-specific configuration for
+            the deployment
+
+    Returns:
+        None
+    """
+    return plugin_store[target].run_local(name, model_uri, flavor, config)
+
+
+def _target_help(target):
+    """
+    Return a string containing detailed documentation on the current deployment target,
+    to be displayed when users invoke the ``mlflow deployments help -t <target-name>`` CLI.
+    This method should be defined within the module specified by the plugin author.
+    The string should contain:
+    * An explanation of target-specific fields in the ``config`` passed to ``create_deployment``,
+      ``update_deployment``
+    * How to specify a ``target_uri`` (e.g. for AWS SageMaker, ``target_uri``s have a scheme of
+      "sagemaker:/<aws-cli-profile-name>", where aws-cli-profile-name is the name of an AWS
+      CLI profile https://docs.aws.amazon.com/cli/latest/userguide/cli-configure-profiles.html)
+    * Any other target-specific details.
+
+    Args:
+        target: Which target to use. This information is used to call the appropriate plugin.
+    """
+    return plugin_store[target].target_help()

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

@@ -0,0 +1,334 @@
+from typing import TYPE_CHECKING, Any, Optional
+
+import requests
+
+from mlflow import MlflowException
+from mlflow.deployments import BaseDeploymentClient
+from mlflow.deployments.constants import (
+    MLFLOW_DEPLOYMENT_CLIENT_REQUEST_RETRY_CODES,
+)
+from mlflow.deployments.server.constants import (
+    MLFLOW_DEPLOYMENTS_CRUD_ENDPOINT_BASE,
+    MLFLOW_DEPLOYMENTS_ENDPOINTS_BASE,
+    MLFLOW_DEPLOYMENTS_QUERY_SUFFIX,
+)
+from mlflow.deployments.utils import resolve_endpoint_url
+from mlflow.environment_variables import (
+    MLFLOW_DEPLOYMENT_PREDICT_TIMEOUT,
+    MLFLOW_HTTP_REQUEST_TIMEOUT,
+)
+from mlflow.protos.databricks_pb2 import BAD_REQUEST
+from mlflow.store.entities.paged_list import PagedList
+from mlflow.utils.annotations import experimental
+from mlflow.utils.credentials import get_default_host_creds
+from mlflow.utils.rest_utils import augmented_raise_for_status, http_request
+from mlflow.utils.uri import join_paths
+
+if TYPE_CHECKING:
+    from mlflow.deployments.server.config import Endpoint
+
+
+@experimental
+class MlflowDeploymentClient(BaseDeploymentClient):
+    """
+    Client for interacting with the MLflow AI Gateway.
+
+    Example:
+
+    First, start the MLflow AI Gateway:
+
+    .. code-block:: bash
+
+        mlflow gateway start --config-path path/to/config.yaml
+
+    Then, create a client and use it to interact with the server:
+
+    .. code-block:: python
+
+        from mlflow.deployments import get_deploy_client
+
+        client = get_deploy_client("http://localhost:5000")
+        endpoints = client.list_endpoints()
+        assert [e.dict() for e in endpoints] == [
+            {
+                "name": "chat",
+                "endpoint_type": "llm/v1/chat",
+                "model": {"name": "gpt-4o-mini", "provider": "openai"},
+                "endpoint_url": "http://localhost:5000/gateway/chat/invocations",
+            },
+        ]
+    """
+
+    def create_deployment(self, name, model_uri, flavor=None, config=None, endpoint=None):
+        """
+        .. warning::
+            This method is not implemented for `MlflowDeploymentClient`.
+        """
+        raise NotImplementedError
+
+    def update_deployment(self, name, model_uri=None, flavor=None, config=None, endpoint=None):
+        """
+        .. warning::
+            This method is not implemented for `MlflowDeploymentClient`.
+        """
+        raise NotImplementedError
+
+    def delete_deployment(self, name, config=None, endpoint=None):
+        """
+        .. warning::
+            This method is not implemented for `MlflowDeploymentClient`.
+        """
+        raise NotImplementedError
+
+    def list_deployments(self, endpoint=None):
+        """
+        .. warning::
+            This method is not implemented for `MlflowDeploymentClient`.
+        """
+        raise NotImplementedError
+
+    def get_deployment(self, name, endpoint=None):
+        """
+        .. warning::
+            This method is not implemented for `MLflowDeploymentClient`.
+        """
+        raise NotImplementedError
+
+    def create_endpoint(self, name, config=None):
+        """
+        .. warning::
+            This method is not implemented for `MlflowDeploymentClient`.
+        """
+        raise NotImplementedError
+
+    def update_endpoint(self, endpoint, config=None):
+        """
+        .. warning::
+            This method is not implemented for `MlflowDeploymentClient`.
+        """
+        raise NotImplementedError
+
+    def delete_endpoint(self, endpoint):
+        """
+        .. warning::
+            This method is not implemented for `MlflowDeploymentClient`.
+        """
+        raise NotImplementedError
+
+    def _call_endpoint(
+        self,
+        method: str,
+        route: str,
+        json_body: Optional[str] = None,
+        timeout: Optional[int] = None,
+    ):
+        call_kwargs = {}
+        if method.lower() == "get":
+            call_kwargs["params"] = json_body
+        else:
+            call_kwargs["json"] = json_body
+
+        response = http_request(
+            host_creds=get_default_host_creds(self.target_uri),
+            endpoint=route,
+            method=method,
+            timeout=MLFLOW_HTTP_REQUEST_TIMEOUT.get() if timeout is None else timeout,
+            retry_codes=MLFLOW_DEPLOYMENT_CLIENT_REQUEST_RETRY_CODES,
+            raise_on_status=False,
+            **call_kwargs,
+        )
+        augmented_raise_for_status(response)
+        return response.json()
+
+    @experimental
+    def get_endpoint(self, endpoint) -> "Endpoint":
+        """
+        Gets a specified endpoint configured for the MLflow AI Gateway.
+
+        Args:
+            endpoint: The name of the endpoint to retrieve.
+
+        Returns:
+            An `Endpoint` object representing the endpoint.
+
+        Example:
+
+        .. code-block:: python
+
+            from mlflow.deployments import get_deploy_client
+
+            client = get_deploy_client("http://localhost:5000")
+            endpoint = client.get_endpoint(endpoint="chat")
+            assert endpoint.dict() == {
+                "name": "chat",
+                "endpoint_type": "llm/v1/chat",
+                "model": {"name": "gpt-4o-mini", "provider": "openai"},
+                "endpoint_url": "http://localhost:5000/gateway/chat/invocations",
+            }
+        """
+        # Delayed import to avoid importing mlflow.gateway in the module scope
+        from mlflow.deployments.server.config import Endpoint
+
+        route = join_paths(MLFLOW_DEPLOYMENTS_CRUD_ENDPOINT_BASE, endpoint)
+        response = self._call_endpoint("GET", route)
+        return Endpoint(
+            **{
+                **response,
+                "endpoint_url": resolve_endpoint_url(self.target_uri, response["endpoint_url"]),
+            }
+        )
+
+    def _list_endpoints(self, page_token=None) -> "PagedList[Endpoint]":
+        # Delayed import to avoid importing mlflow.gateway in the module scope
+        from mlflow.deployments.server.config import Endpoint
+
+        params = None if page_token is None else {"page_token": page_token}
+        response_json = self._call_endpoint(
+            "GET", MLFLOW_DEPLOYMENTS_CRUD_ENDPOINT_BASE, json_body=params
+        )
+        routes = [
+            Endpoint(
+                **{
+                    **resp,
+                    "endpoint_url": resolve_endpoint_url(
+                        self.target_uri,
+                        resp["endpoint_url"],
+                    ),
+                }
+            )
+            for resp in response_json.get("endpoints", [])
+        ]
+        next_page_token = response_json.get("next_page_token")
+        return PagedList(routes, next_page_token)
+
+    @experimental
+    def list_endpoints(self) -> "list[Endpoint]":
+        """
+        List endpoints configured for the MLflow AI Gateway.
+
+        Returns:
+            A list of ``Endpoint`` objects.
+
+        Example:
+
+        .. code-block:: python
+
+            from mlflow.deployments import get_deploy_client
+
+            client = get_deploy_client("http://localhost:5000")
+
+            endpoints = client.list_endpoints()
+            assert [e.dict() for e in endpoints] == [
+                {
+                    "name": "chat",
+                    "endpoint_type": "llm/v1/chat",
+                    "model": {"name": "gpt-4o-mini", "provider": "openai"},
+                    "endpoint_url": "http://localhost:5000/gateway/chat/invocations",
+                },
+            ]
+
+        """
+        endpoints = []
+        next_page_token = None
+        while True:
+            page = self._list_endpoints(next_page_token)
+            endpoints.extend(page)
+            next_page_token = page.token
+            if next_page_token is None:
+                break
+        return endpoints
+
+    @experimental
+    def predict(self, deployment_name=None, inputs=None, endpoint=None) -> dict[str, Any]:
+        """
+        Submit a query to a configured provider endpoint.
+
+        Args:
+            deployment_name: Unused.
+            inputs: The inputs to the query, as a dictionary.
+            endpoint: The name of the endpoint to query.
+
+        Returns:
+            A dictionary containing the response from the endpoint.
+
+        Example:
+
+        .. code-block:: python
+
+            from mlflow.deployments import get_deploy_client
+
+            client = get_deploy_client("http://localhost:5000")
+
+            response = client.predict(
+                endpoint="chat",
+                inputs={"messages": [{"role": "user", "content": "Hello"}]},
+            )
+            assert response == {
+                "id": "chatcmpl-8OLoQuaeJSLybq3NBoe0w5eyqjGb9",
+                "object": "chat.completion",
+                "created": 1700814410,
+                "model": "gpt-4o-mini",
+                "choices": [
+                    {
+                        "index": 0,
+                        "message": {
+                            "role": "assistant",
+                            "content": "Hello! How can I assist you today?",
+                        },
+                        "finish_reason": "stop",
+                    }
+                ],
+                "usage": {
+                    "prompt_tokens": 9,
+                    "completion_tokens": 9,
+                    "total_tokens": 18,
+                },
+            }
+
+        Additional parameters that are valid for a given provider and endpoint configuration can be
+        included with the request as shown below, using an openai completions endpoint request as
+        an example:
+
+        .. code-block:: python
+
+            from mlflow.deployments import get_deploy_client
+
+            client = get_deploy_client("http://localhost:5000")
+            client.predict(
+                endpoint="completions",
+                inputs={
+                    "prompt": "Hello!",
+                    "temperature": 0.3,
+                    "max_tokens": 500,
+                },
+            )
+        """
+        query_route = join_paths(
+            MLFLOW_DEPLOYMENTS_ENDPOINTS_BASE, endpoint, MLFLOW_DEPLOYMENTS_QUERY_SUFFIX
+        )
+        try:
+            return self._call_endpoint(
+                "POST", query_route, inputs, MLFLOW_DEPLOYMENT_PREDICT_TIMEOUT.get()
+            )
+        except MlflowException as e:
+            if isinstance(e.__cause__, requests.exceptions.Timeout):
+                raise MlflowException(
+                    message=(
+                        "The provider has timed out while generating a response to your "
+                        "query. Please evaluate the available parameters for the query "
+                        "that you are submitting. Some parameter values and inputs can "
+                        "increase the computation time beyond the allowable route "
+                        f"timeout of {MLFLOW_DEPLOYMENT_PREDICT_TIMEOUT} "
+                        "seconds."
+                    ),
+                    error_code=BAD_REQUEST,
+                )
+            raise e
+
+
+def run_local(name, model_uri, flavor=None, config=None):
+    pass
+
+
+def target_help():
+    pass

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

@@ -0,0 +1,252 @@
+import os
+
+from mlflow.deployments import BaseDeploymentClient
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INVALID_PARAMETER_VALUE
+from mlflow.utils.openai_utils import (
+    _OAITokenHolder,
+    _OpenAIApiConfig,
+    _OpenAIEnvVar,
+)
+from mlflow.utils.rest_utils import augmented_raise_for_status
+
+
+class OpenAIDeploymentClient(BaseDeploymentClient):
+    """
+    Client for interacting with OpenAI endpoints.
+
+    Example:
+
+    First, set up credentials for authentication:
+
+    .. code-block:: bash
+
+        export OPENAI_API_KEY=...
+
+    .. seealso::
+
+        See https://mlflow.org/docs/latest/python_api/openai/index.html for other authentication
+        methods.
+
+    Then, create a deployment client and use it to interact with OpenAI endpoints:
+
+    .. code-block:: python
+
+        from mlflow.deployments import get_deploy_client
+
+        client = get_deploy_client("openai")
+        client.predict(
+            endpoint="gpt-4o-mini",
+            inputs={
+                "messages": [
+                    {"role": "user", "content": "Hello!"},
+                ],
+            },
+        )
+    """
+
+    def create_deployment(self, name, model_uri, flavor=None, config=None, endpoint=None):
+        """
+        .. warning::
+
+            This method is not implemented for `OpenAIDeploymentClient`.
+        """
+        raise NotImplementedError
+
+    def update_deployment(self, name, model_uri=None, flavor=None, config=None, endpoint=None):
+        """
+        .. warning::
+
+            This method is not implemented for `OpenAIDeploymentClient`.
+        """
+        raise NotImplementedError
+
+    def delete_deployment(self, name, config=None, endpoint=None):
+        """
+        .. warning::
+
+            This method is not implemented for `OpenAIDeploymentClient`.
+        """
+        raise NotImplementedError
+
+    def list_deployments(self, endpoint=None):
+        """
+        .. warning::
+
+            This method is not implemented for `OpenAIDeploymentClient`.
+        """
+        raise NotImplementedError
+
+    def get_deployment(self, name, endpoint=None):
+        """
+        .. warning::
+
+            This method is not implemented for `OpenAIDeploymentClient`.
+        """
+        raise NotImplementedError
+
+    def predict(self, deployment_name=None, inputs=None, endpoint=None):
+        """Query an OpenAI endpoint.
+        See https://platform.openai.com/docs/api-reference for more information.
+
+        Args:
+            deployment_name: Unused.
+            inputs: A dictionary containing the model inputs to query.
+            endpoint: The name of the endpoint to query.
+
+        Returns:
+            A dictionary containing the model outputs.
+
+        """
+        _check_openai_key()
+
+        api_config = _get_api_config_without_openai_dep()
+        api_token = _OAITokenHolder(api_config.api_type)
+        api_token.refresh()
+
+        if api_config.api_type in ("azure", "azure_ad", "azuread"):
+            from openai import AzureOpenAI
+
+            client = AzureOpenAI(
+                api_key=api_token.token,
+                azure_endpoint=api_config.api_base,
+                api_version=api_config.api_version,
+                azure_deployment=api_config.deployment_id,
+                max_retries=api_config.max_retries,
+                timeout=api_config.timeout,
+            )
+        else:
+            from openai import OpenAI
+
+            client = OpenAI(
+                api_key=api_token.token,
+                base_url=api_config.api_base,
+                max_retries=api_config.max_retries,
+                timeout=api_config.timeout,
+            )
+
+        return client.chat.completions.create(
+            messages=inputs["messages"], model=endpoint
+        ).model_dump()
+
+    def create_endpoint(self, name, config=None):
+        """
+        .. warning::
+
+            This method is not implemented for `OpenAIDeploymentClient`.
+        """
+        raise NotImplementedError
+
+    def update_endpoint(self, endpoint, config=None):
+        """
+        .. warning::
+
+            This method is not implemented for `OpenAIDeploymentClient`.
+        """
+        raise NotImplementedError
+
+    def delete_endpoint(self, endpoint):
+        """
+        .. warning::
+
+            This method is not implemented for `OpenAIDeploymentClient`.
+        """
+        raise NotImplementedError
+
+    def list_endpoints(self):
+        """
+        List the currently available models.
+        """
+
+        _check_openai_key()
+
+        api_config = _get_api_config_without_openai_dep()
+        import requests
+
+        if api_config.api_type in ("azure", "azure_ad", "azuread"):
+            raise NotImplementedError(
+                "List endpoints is not implemented for Azure OpenAI API",
+            )
+        else:
+            api_key = os.environ["OPENAI_API_KEY"]
+            request_header = {"Authorization": f"Bearer {api_key}"}
+
+            response = requests.get(
+                "https://api.openai.com/v1/models",
+                headers=request_header,
+            )
+
+            augmented_raise_for_status(response)
+
+            return response.json()
+
+    def get_endpoint(self, endpoint):
+        """
+        Get information about a specific model.
+        """
+
+        _check_openai_key()
+
+        api_config = _get_api_config_without_openai_dep()
+        import requests
+
+        if api_config.api_type in ("azure", "azure_ad", "azuread"):
+            raise NotImplementedError(
+                "Get endpoint is not implemented for Azure OpenAI API",
+            )
+        else:
+            api_key = os.environ["OPENAI_API_KEY"]
+            request_header = {"Authorization": f"Bearer {api_key}"}
+
+            response = requests.get(
+                f"https://api.openai.com/v1/models/{endpoint}",
+                headers=request_header,
+            )
+
+            augmented_raise_for_status(response)
+
+            return response.json()
+
+
+def run_local(name, model_uri, flavor=None, config=None):
+    pass
+
+
+def target_help():
+    pass
+
+
+def _get_api_config_without_openai_dep() -> _OpenAIApiConfig:
+    """
+    Gets the parameters and configuration of the OpenAI API connected to.
+    """
+    api_type = os.getenv(_OpenAIEnvVar.OPENAI_API_TYPE.value)
+    api_version = os.getenv(_OpenAIEnvVar.OPENAI_API_VERSION.value)
+    api_base = os.getenv(_OpenAIEnvVar.OPENAI_API_BASE.value, None)
+    deployment_id = os.getenv(_OpenAIEnvVar.OPENAI_DEPLOYMENT_NAME.value, None)
+    if api_type in ("azure", "azure_ad", "azuread"):
+        batch_size = 16
+        max_tokens_per_minute = 60_000
+    else:
+        # The maximum batch size is 2048:
+        # https://github.com/openai/openai-python/blob/b82a3f7e4c462a8a10fa445193301a3cefef9a4a/openai/embeddings_utils.py#L43
+        # We use a smaller batch size to be safe.
+        batch_size = 1024
+        max_tokens_per_minute = 90_000
+    return _OpenAIApiConfig(
+        api_type=api_type,
+        batch_size=batch_size,
+        max_requests_per_minute=3_500,
+        max_tokens_per_minute=max_tokens_per_minute,
+        api_base=api_base,
+        api_version=api_version,
+        deployment_id=deployment_id,
+    )
+
+
+def _check_openai_key():
+    if "OPENAI_API_KEY" not in os.environ:
+        raise MlflowException(
+            "OPENAI_API_KEY environment variable not set",
+            error_code=INVALID_PARAMETER_VALUE,
+        )

venv/lib/python3.9/site-packages/mlflow/deployments/plugin_manager.py

@@ -0,0 +1,143 @@
+import abc
+import importlib.metadata
+import inspect
+
+import importlib_metadata
+
+from mlflow.deployments.base import BaseDeploymentClient
+from mlflow.deployments.utils import parse_target_uri
+from mlflow.exceptions import MlflowException
+from mlflow.protos.databricks_pb2 import INTERNAL_ERROR, RESOURCE_DOES_NOT_EXIST
+from mlflow.utils.annotations import developer_stable
+from mlflow.utils.plugins import get_entry_points
+
+# TODO: refactor to have a common base class for all the plugin implementation in MLflow
+#   mlflow/tracking/context/registry.py
+#   mlflow/tracking/registry
+#   mlflow/store/artifact/artifact_repository_registry.py
+
+
+@developer_stable
+class PluginManager(abc.ABC):
+    """
+    Abstract class defining a entrypoint based plugin registration.
+
+    This class allows the registration of a function or class to provide an implementation
+    for a given key/name. Implementations declared though the entrypoints can be automatically
+    registered through the `register_entrypoints` method.
+    """
+
+    def __init__(self, group_name):
+        self._registry = {}
+        self.group_name = group_name
+        self._has_registered = None
+
+    @abc.abstractmethod
+    def __getitem__(self, item):
+        # Letting the child class create this function so that the child
+        # can raise custom exceptions if it needs to
+        pass
+
+    @property
+    def registry(self):
+        """
+        Registry stores the registered plugin as a key value pair where key is the
+        name of the plugin and value is the plugin object
+        """
+        return self._registry
+
+    @property
+    def has_registered(self):
+        """
+        Returns bool representing whether the "register_entrypoints" has run or not. This
+        doesn't return True if `register` method is called outside of `register_entrypoints`
+        to register plugins
+        """
+        return self._has_registered
+
+    def register(self, target_name, plugin_module):
+        """Register a deployment client given its target name and module
+        Args:
+            target_name: The name of the deployment target. This name will be used by
+                `get_deploy_client()` to retrieve a deployment client from
+                the plugin store.
+            plugin_module: The module that implements the deployment plugin interface.
+        """
+        self.registry[target_name] = importlib.metadata.EntryPoint(
+            target_name, plugin_module, self.group_name
+        )
+
+    def register_entrypoints(self):
+        """
+        Runs through all the packages that has the `group_name` defined as the entrypoint
+        and register that into the registry
+        """
+        for entrypoint in get_entry_points(self.group_name):
+            self.registry[entrypoint.name] = entrypoint
+        self._has_registered = True
+
+
+@developer_stable
+class DeploymentPlugins(PluginManager):
+    def __init__(self):
+        super().__init__("mlflow.deployments")
+        self.register_entrypoints()
+
+    def __getitem__(self, item):
+        """Override __getitem__ so that we can directly look up plugins via dict-like syntax"""
+        try:
+            target_name = parse_target_uri(item)
+            plugin_like = self.registry[target_name]
+        except KeyError:
+            msg = (
+                f'No plugin found for managing model deployments to "{item}". '
+                f'In order to deploy models to "{item}", find and install an appropriate '
+                "plugin from "
+                "https://mlflow.org/docs/latest/plugins.html#community-plugins using "
+                "your package manager (pip, conda etc)."
+            )
+            raise MlflowException(msg, error_code=RESOURCE_DOES_NOT_EXIST)
+
+        if isinstance(plugin_like, (importlib_metadata.EntryPoint, importlib.metadata.EntryPoint)):
+            try:
+                plugin_obj = plugin_like.load()
+            except (AttributeError, ImportError) as exc:
+                raise RuntimeError(f'Failed to load the plugin "{item}": {exc}')
+            self.registry[item] = plugin_obj
+        else:
+            plugin_obj = plugin_like
+
+        # Testing whether the plugin is valid or not
+        expected = {"target_help", "run_local"}
+        deployment_classes = []
+        for name, obj in inspect.getmembers(plugin_obj):
+            if name in expected:
+                expected.remove(name)
+            elif (
+                inspect.isclass(obj)
+                and issubclass(obj, BaseDeploymentClient)
+                and not obj == BaseDeploymentClient
+            ):
+                deployment_classes.append(name)
+        if len(expected) > 0:
+            raise MlflowException(
+                f"Plugin registered for the target {item} does not have all "
+                "the required interfaces. Raise an issue with the "
+                "plugin developers.\n"
+                f"Missing interfaces: {expected}",
+                error_code=INTERNAL_ERROR,
+            )
+        if len(deployment_classes) > 1:
+            raise MlflowException(
+                f"Plugin registered for the target {item} has more than one "
+                "child class of BaseDeploymentClient. Raise an issue with"
+                " the plugin developers. "
+                f"Classes found are {deployment_classes}"
+            )
+        elif len(deployment_classes) == 0:
+            raise MlflowException(
+                f"Plugin registered for the target {item} has no child class"
+                " of BaseDeploymentClient. Raise an issue with the "
+                "plugin developers"
+            )
+        return plugin_obj

venv/lib/python3.9/site-packages/mlflow/deployments/server/app.py

@@ -0,0 +1,31 @@
+"""
+TODO: Remove this module once after Deployments Server deprecation window elapses
+"""
+
+from mlflow.environment_variables import (
+    MLFLOW_DEPLOYMENTS_CONFIG,
+)
+from mlflow.exceptions import MlflowException
+from mlflow.gateway.app import GatewayAPI
+from mlflow.gateway.app import (
+    create_app_from_config as gateway_create_app_from_config,
+)
+from mlflow.gateway.app import (
+    create_app_from_path as gateway_create_app_from_path,
+)
+
+create_app_from_config = gateway_create_app_from_config
+create_app_from_path = gateway_create_app_from_path
+
+
+def create_app_from_env() -> GatewayAPI:
+    """
+    Load the path from the environment variable and generate the GatewayAPI app instance.
+    """
+    if config_path := MLFLOW_DEPLOYMENTS_CONFIG.get():
+        return create_app_from_path(config_path)
+
+    raise MlflowException(
+        f"Environment variable {MLFLOW_DEPLOYMENTS_CONFIG!r} is not set. "
+        "Please set it to the path of the gateway configuration file."
+    )

venv/lib/python3.9/site-packages/mlflow/deployments/server/config.py

@@ -0,0 +1,26 @@
+from typing import Optional
+
+from mlflow.gateway.base_models import ResponseModel
+from mlflow.gateway.config import Limit, RouteModelInfo
+
+
+class Endpoint(ResponseModel):
+    name: str
+    endpoint_type: str
+    model: RouteModelInfo
+    endpoint_url: str
+    limit: Optional[Limit]
+
+    class Config:
+        schema_extra = {
+            "example": {
+                "name": "openai-completions",
+                "endpoint_type": "llm/v1/completions",
+                "model": {
+                    "name": "gpt-4o-mini",
+                    "provider": "openai",
+                },
+                "endpoint_url": "/endpoints/completions/invocations",
+                "limit": {"calls": 1, "key": None, "renewal_period": "minute"},
+            }
+        }

venv/lib/python3.9/site-packages/mlflow/deployments/server/constants.py

@@ -0,0 +1,6 @@
+MLFLOW_DEPLOYMENTS_HEALTH_ENDPOINT = "/health"
+MLFLOW_DEPLOYMENTS_CRUD_ENDPOINT_BASE = "/api/2.0/endpoints/"
+MLFLOW_DEPLOYMENTS_LIMITS_BASE = "/api/2.0/endpoints/limits/"
+MLFLOW_DEPLOYMENTS_ENDPOINTS_BASE = "/endpoints/"
+MLFLOW_DEPLOYMENTS_QUERY_SUFFIX = "/invocations"
+MLFLOW_DEPLOYMENTS_LIST_ENDPOINTS_PAGE_SIZE = 3000

venv/lib/python3.9/site-packages/mlflow/deployments/utils.py

@@ -0,0 +1,97 @@
+import urllib
+from typing import Optional
+from urllib.parse import urlparse
+
+from mlflow.environment_variables import MLFLOW_DEPLOYMENTS_TARGET
+from mlflow.exceptions import MlflowException
+from mlflow.utils.uri import append_to_uri_path
+
+_deployments_target: Optional[str] = None
+
+
+def parse_target_uri(target_uri):
+    """Parse out the deployment target from the provided target uri"""
+    parsed = urllib.parse.urlparse(target_uri)
+    if not parsed.scheme:
+        if parsed.path:
+            # uri = 'target_name' (without :/<path>)
+            return parsed.path
+        raise MlflowException(
+            f"Not a proper deployment URI: {target_uri}. "
+            + "Deployment URIs must be of the form 'target' or 'target:/suffix'"
+        )
+    return parsed.scheme
+
+
+def _is_valid_uri(uri: str) -> bool:
+    """
+    Evaluates the basic structure of a provided uri to determine if the scheme and
+    netloc are provided
+    """
+    try:
+        parsed = urlparse(uri)
+        return bool(parsed.scheme and parsed.netloc)
+    except ValueError:
+        return False
+
+
+def resolve_endpoint_url(base_url: str, endpoint: str) -> str:
+    """Performs a validation on whether the returned value is a fully qualified url
+    or requires the assembly of a fully qualified url by appending `endpoint`.
+
+    Args:
+        base_url: The base URL. Should include the scheme and domain, e.g.,
+            ``http://127.0.0.1:6000``.
+        endpoint: The endpoint to be appended to the base URL, e.g., ``/api/2.0/endpoints/`` or,
+            in the case of Databricks, the fully qualified url.
+
+    Returns:
+        The complete URL, either directly returned or formed and returned by joining the
+        base URL and the endpoint path.
+
+    """
+    return endpoint if _is_valid_uri(endpoint) else append_to_uri_path(base_url, endpoint)
+
+
+def set_deployments_target(target: str):
+    """Sets the target deployment client for MLflow deployments
+
+    Args:
+        target: The full uri of a running MLflow AI Gateway or, if running on
+            Databricks, "databricks".
+    """
+    if not _is_valid_target(target):
+        raise MlflowException.invalid_parameter_value(
+            "The target provided is not a valid uri or 'databricks'"
+        )
+
+    global _deployments_target
+    _deployments_target = target
+
+
+def get_deployments_target() -> str:
+    """
+    Returns the currently set MLflow deployments target iff set.
+    If the deployments target has not been set by using ``set_deployments_target``, an
+    ``MlflowException`` is raised.
+    """
+    if _deployments_target is not None:
+        return _deployments_target
+    elif uri := MLFLOW_DEPLOYMENTS_TARGET.get():
+        return uri
+    else:
+        raise MlflowException(
+            "No deployments target has been set. Please either set the MLflow deployments target"
+            " via `mlflow.deployments.set_deployments_target()` or set the environment variable "
+            f"{MLFLOW_DEPLOYMENTS_TARGET} to the running deployment server's uri"
+        )
+
+
+def _is_valid_target(target: str):
+    """
+    Evaluates the basic structure of a provided target to determine if the scheme and
+    netloc are provided
+    """
+    if target == "databricks":
+        return True
+    return _is_valid_uri(target)