star

venv/lib/python3.9/site-packages/mlflow/pyfunc/_mlflow_pyfunc_backend_predict.py

@@ -0,0 +1,61 @@
+"""
+This script should be executed in a fresh python interpreter process using `subprocess`.
+"""
+
+import argparse
+
+from mlflow.pyfunc.scoring_server import _predict
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model-uri", required=True)
+    parser.add_argument("--input-path", required=False)
+    parser.add_argument("--output-path", required=False)
+    parser.add_argument("--content-type", required=True)
+    return parser.parse_args()
+
+
+# Guidance for fixing missing module error
+_MISSING_MODULE_HELP_MSG = (
+    "Exception occurred while running inference: {e}"
+    "\n\n"
+    "\033[93m[Hint] It appears that your MLflow Model doesn't contain the required "
+    "dependency '{missing_module}' to run model inference. When logging a model, MLflow "
+    "detects dependencies based on the model flavor, but it is possible that some "
+    "dependencies are not captured. In this case, you can manually add dependencies "
+    "using the `extra_pip_requirements` parameter of `mlflow.pyfunc.log_model`.\033[0m"
+    """
+
+\033[1mSample code:\033[0m
+    ----
+    mlflow.pyfunc.log_model(
+        artifact_path="model",
+        python_model=your_model,
+        extra_pip_requirements=["{missing_module}==x.y.z"]
+    )
+    ----
+
+    For mode guidance on fixing missing dependencies, please refer to the MLflow docs:
+    https://www.mlflow.org/docs/latest/deployment/index.html#how-to-fix-dependency-errors-when-serving-my-model
+"""
+)
+
+
+def main():
+    args = parse_args()
+
+    try:
+        _predict(
+            model_uri=args.model_uri,
+            input_path=args.input_path if args.input_path else None,
+            output_path=args.output_path if args.output_path else None,
+            content_type=args.content_type,
+        )
+    except ModuleNotFoundError as e:
+        message = _MISSING_MODULE_HELP_MSG.format(e=str(e), missing_module=e.name)
+        raise RuntimeError(message) from e
+
+
+if __name__ == "__main__":
+    main()

venv/lib/python3.9/site-packages/mlflow/pyfunc/backend.py

@@ -0,0 +1,518 @@
+import ctypes
+import json
+import logging
+import os
+import pathlib
+import posixpath
+import shlex
+import signal
+import subprocess
+import sys
+import warnings
+from pathlib import Path
+
+from mlflow import pyfunc
+from mlflow.exceptions import MlflowException
+from mlflow.models import FlavorBackend, Model, docker_utils
+from mlflow.models.docker_utils import PYTHON_SLIM_BASE_IMAGE, UBUNTU_BASE_IMAGE
+from mlflow.pyfunc import (
+    ENV,
+    _extract_conda_env,
+    _mlflow_pyfunc_backend_predict,
+    mlserver,
+    scoring_server,
+)
+from mlflow.tracking.artifact_utils import _download_artifact_from_uri
+from mlflow.utils import env_manager as em
+from mlflow.utils.conda import get_conda_bin_executable, get_or_create_conda_env
+from mlflow.utils.environment import Environment, _get_pip_install_mlflow, _PythonEnv
+from mlflow.utils.file_utils import (
+    TempDir,
+    get_or_create_nfs_tmp_dir,
+    get_or_create_tmp_dir,
+    path_to_local_file_uri,
+)
+from mlflow.utils.model_utils import _get_all_flavor_configurations
+from mlflow.utils.nfs_on_spark import get_nfs_cache_root_dir
+from mlflow.utils.os import is_windows
+from mlflow.utils.process import ShellCommandException, cache_return_value_per_process
+from mlflow.utils.virtualenv import _get_or_create_virtualenv
+from mlflow.version import VERSION
+
+_logger = logging.getLogger(__name__)
+
+_STDIN_SERVER_SCRIPT = Path(__file__).parent.joinpath("stdin_server.py")
+
+# Flavors that require Java to be installed in the environment
+JAVA_FLAVORS = {"johnsnowlabs", "h2o", "mleap", "spark"}
+
+# Some flavor requires additional packages to be installed in the environment
+FLAVOR_SPECIFIC_APT_PACKAGES = {
+    "lightgbm": ["libgomp1"],
+    "paddle": ["libgomp1"],
+}
+
+# Directory to store loaded model inside the Docker context directory
+_MODEL_DIR_NAME = "model_dir"
+LOCAL_ENV_MANAGER_ERROR_MESSAGE = "We cannot use 'LOCAL' environment manager "
+"for your model configuration. Please specify a virtualenv or conda environment "
+"manager instead with `--env-manager` argument."
+
+
+def _set_mlflow_config_env(command_env, model_config):
+    if model_config:
+        command_env[scoring_server.SERVING_MODEL_CONFIG] = json.dumps(model_config)
+    return command_env
+
+
+class PyFuncBackend(FlavorBackend):
+    """
+    Flavor backend implementation for the generic python models.
+    """
+
+    def __init__(  # noqa: D417
+        self,
+        config,
+        env_manager,
+        workers=1,
+        install_mlflow=False,
+        create_env_root_dir=False,
+        env_root_dir=None,
+        **kwargs,
+    ):
+        """
+        Args:
+            env_manager: Environment manager to use for preparing the environment. If None,
+                MLflow will automatically pick the env manager based on the model's flavor
+                configuration for generate_dockerfile. It can't be None for other methods.
+            env_root_dir: Root path for conda env. If None, use Conda's default environments
+                directory. Note if this is set, conda package cache path becomes
+                "{env_root_dir}/conda_cache_pkgs" instead of the global package cache
+                path, and pip package cache path becomes
+                "{env_root_dir}/pip_cache_pkgs" instead of the global package cache
+                path.
+        """
+        super().__init__(config=config, **kwargs)
+        self._nworkers = workers or 1
+        if env_manager == em.CONDA and ENV not in config:
+            warnings.warn(
+                "Conda environment is not specified in config `env`. Using local environment."
+            )
+            env_manager = em.LOCAL
+        self._env_manager = env_manager
+        self._install_mlflow = install_mlflow
+        self._env_id = os.environ.get("MLFLOW_HOME", VERSION) if install_mlflow else None
+        self._create_env_root_dir = create_env_root_dir
+        self._env_root_dir = env_root_dir
+        self._environment = None
+
+    def prepare_env(
+        self, model_uri, capture_output=False, pip_requirements_override=None, extra_envs=None
+    ):
+        if self._environment is not None:
+            return self._environment
+
+        @cache_return_value_per_process
+        def _get_or_create_env_root_dir(should_use_nfs):
+            if should_use_nfs:
+                root_tmp_dir = get_or_create_nfs_tmp_dir()
+            else:
+                root_tmp_dir = get_or_create_tmp_dir()
+
+            envs_root_dir = os.path.join(root_tmp_dir, "envs")
+            os.makedirs(envs_root_dir, exist_ok=True)
+            return envs_root_dir
+
+        local_path = _download_artifact_from_uri(model_uri)
+        if self._create_env_root_dir:
+            if self._env_root_dir is not None:
+                raise Exception("env_root_dir can not be set when create_env_root_dir=True")
+            nfs_root_dir = get_nfs_cache_root_dir()
+            env_root_dir = _get_or_create_env_root_dir(nfs_root_dir is not None)
+        else:
+            env_root_dir = self._env_root_dir
+
+        if self._env_manager in {em.VIRTUALENV, em.UV}:
+            activate_cmd = _get_or_create_virtualenv(
+                local_path,
+                self._env_id,
+                env_root_dir=env_root_dir,
+                capture_output=capture_output,
+                pip_requirements_override=pip_requirements_override,
+                env_manager=self._env_manager,
+            )
+            self._environment = Environment(activate_cmd, extra_env=extra_envs)
+        elif self._env_manager == em.CONDA:
+            conda_env_path = os.path.join(local_path, _extract_conda_env(self._config[ENV]))
+            self._environment = get_or_create_conda_env(
+                conda_env_path,
+                env_id=self._env_id,
+                capture_output=capture_output,
+                env_root_dir=env_root_dir,
+                pip_requirements_override=pip_requirements_override,
+                extra_envs=extra_envs,
+            )
+
+        elif self._env_manager == em.LOCAL:
+            raise Exception("Prepare env should not be called with local env manager!")
+        else:
+            raise Exception(f"Unexpected env manager value '{self._env_manager}'")
+
+        if self._install_mlflow:
+            self._environment.execute(_get_pip_install_mlflow())
+        else:
+            self._environment.execute('python -c ""')
+
+        return self._environment
+
+    def predict(
+        self,
+        model_uri,
+        input_path,
+        output_path,
+        content_type,
+        pip_requirements_override=None,
+        extra_envs=None,
+    ):
+        """
+        Generate predictions using generic python model saved with MLflow. The expected format of
+        the input JSON is the MLflow scoring format.
+        Return the prediction results as a JSON.
+        """
+        local_path = _download_artifact_from_uri(model_uri)
+        # NB: Absolute windows paths do not work with mlflow apis, use file uri to ensure
+        # platform compatibility.
+        local_uri = path_to_local_file_uri(local_path)
+
+        if self._env_manager != em.LOCAL:
+            predict_cmd = [
+                "python",
+                _mlflow_pyfunc_backend_predict.__file__,
+                "--model-uri",
+                str(local_uri),
+                "--content-type",
+                shlex.quote(str(content_type)),
+            ]
+            if input_path:
+                predict_cmd += ["--input-path", shlex.quote(str(input_path))]
+            if output_path:
+                predict_cmd += ["--output-path", shlex.quote(str(output_path))]
+
+            if pip_requirements_override and self._env_manager == em.CONDA:
+                # Conda use = instead of == for version pinning
+                pip_requirements_override = [
+                    pip_req.replace("==", "=") for pip_req in pip_requirements_override
+                ]
+
+            environment = self.prepare_env(
+                local_path,
+                pip_requirements_override=pip_requirements_override,
+                extra_envs=extra_envs,
+            )
+
+            try:
+                environment.execute(" ".join(predict_cmd))
+            except ShellCommandException as e:
+                raise MlflowException(
+                    f"{e}\n\nAn exception occurred while running model prediction within a "
+                    f"{self._env_manager} environment. You can find the error message "
+                    f"from the prediction subprocess by scrolling above."
+                ) from None
+        else:
+            if pip_requirements_override:
+                raise MlflowException(
+                    "`pip_requirements_override` is not supported for local env manager."
+                    "Please use conda or virtualenv instead."
+                )
+            scoring_server._predict(local_uri, input_path, output_path, content_type)
+
+    def serve(
+        self,
+        model_uri,
+        port,
+        host,
+        timeout,
+        enable_mlserver,
+        synchronous=True,
+        stdout=None,
+        stderr=None,
+        model_config=None,
+    ):
+        """
+        Serve pyfunc model locally.
+        """
+        local_path = _download_artifact_from_uri(model_uri)
+
+        server_implementation = mlserver if enable_mlserver else scoring_server
+        command, command_env = server_implementation.get_cmd(
+            local_path, port, host, timeout, self._nworkers
+        )
+        _set_mlflow_config_env(command_env, model_config)
+
+        if sys.platform.startswith("linux"):
+
+            def setup_sigterm_on_parent_death():
+                """
+                Uses prctl to automatically send SIGTERM to the command process when its parent is
+                dead.
+
+                This handles the case when the parent is a PySpark worker process.
+                If a user cancels the PySpark job, the worker process gets killed, regardless of
+                PySpark daemon and worker reuse settings.
+                We use prctl to ensure the command process receives SIGTERM after spark job
+                cancellation.
+                The command process itself should handle SIGTERM properly.
+                This is a no-op on macOS because prctl is not supported.
+
+                Note:
+                When a pyspark job canceled, the UDF python process are killed by signal "SIGKILL",
+                This case neither "atexit" nor signal handler can capture SIGKILL signal.
+                prctl is the only way to capture SIGKILL signal.
+                """
+                try:
+                    libc = ctypes.CDLL("libc.so.6")
+                    # Set the parent process death signal of the command process to SIGTERM.
+                    libc.prctl(1, signal.SIGTERM)  # PR_SET_PDEATHSIG, see prctl.h
+                except OSError as e:
+                    # TODO: find approach for supporting MacOS/Windows system which does
+                    #  not support prctl.
+                    warnings.warn(f"Setup libc.prctl PR_SET_PDEATHSIG failed, error {e!r}.")
+
+        else:
+            setup_sigterm_on_parent_death = None
+
+        if not is_windows():
+            # Add "exec" before the starting scoring server command, so that the scoring server
+            # process replaces the bash process, otherwise the scoring server process is created
+            # as a child process of the bash process.
+            # Note we in `mlflow.pyfunc.spark_udf`, use prctl PR_SET_PDEATHSIG to ensure scoring
+            # server process being killed when UDF process exit. The PR_SET_PDEATHSIG can only
+            # send signal to the bash process, if the scoring server process is created as a
+            # child process of the bash process, then it cannot receive the signal sent by prctl.
+            # TODO: For Windows, there's no equivalent things of Unix shell's exec. Windows also
+            #  does not support prctl. We need to find an approach to address it.
+            command = "exec " + command
+
+        if self._env_manager != em.LOCAL:
+            return self.prepare_env(local_path).execute(
+                command,
+                command_env,
+                stdout=stdout,
+                stderr=stderr,
+                preexec_fn=setup_sigterm_on_parent_death,
+                synchronous=synchronous,
+            )
+        else:
+            _logger.info("=== Running command '%s'", command)
+
+            if not is_windows():
+                command = ["bash", "-c", command]
+
+            child_proc = subprocess.Popen(
+                command,
+                env=command_env,
+                preexec_fn=setup_sigterm_on_parent_death,
+                stdout=stdout,
+                stderr=stderr,
+            )
+
+            if synchronous:
+                rc = child_proc.wait()
+                if rc != 0:
+                    raise Exception(
+                        f"Command '{command}' returned non zero return code. Return code = {rc}"
+                    )
+                return 0
+            else:
+                return child_proc
+
+    def serve_stdin(
+        self,
+        model_uri,
+        stdout=None,
+        stderr=None,
+        model_config=None,
+    ):
+        local_path = _download_artifact_from_uri(model_uri)
+
+        command_env = os.environ.copy()
+        _set_mlflow_config_env(command_env, model_config)
+
+        return self.prepare_env(local_path).execute(
+            command=f"python {_STDIN_SERVER_SCRIPT} --model-uri {local_path}",
+            command_env=command_env,
+            stdin=subprocess.PIPE,
+            stdout=stdout,
+            stderr=stderr,
+            synchronous=False,
+        )
+
+    def can_score_model(self):
+        if self._env_manager == em.LOCAL:
+            # noconda => already in python and dependencies are assumed to be installed.
+            return True
+        conda_path = get_conda_bin_executable("conda")
+        try:
+            p = subprocess.Popen(
+                [conda_path, "--version"], stdout=subprocess.PIPE, stderr=subprocess.PIPE
+            )
+            _, _ = p.communicate()
+            return p.wait() == 0
+        except FileNotFoundError:
+            # Can not find conda
+            return False
+
+    def build_image(
+        self,
+        model_uri,
+        image_name,
+        install_java=False,
+        install_mlflow=False,
+        mlflow_home=None,
+        enable_mlserver=False,
+        base_image=None,
+    ):
+        with TempDir() as tmp:
+            cwd = tmp.path()
+            self.generate_dockerfile(
+                model_uri=model_uri,
+                output_dir=cwd,
+                install_java=install_java,
+                install_mlflow=install_mlflow,
+                mlflow_home=mlflow_home,
+                enable_mlserver=enable_mlserver,
+                base_image=base_image,
+            )
+
+            _logger.info("Building docker image with name %s", image_name)
+            docker_utils.build_image_from_context(context_dir=cwd, image_name=image_name)
+
+    def generate_dockerfile(
+        self,
+        model_uri,
+        output_dir,
+        install_java=False,
+        install_mlflow=False,
+        mlflow_home=None,
+        enable_mlserver=False,
+        base_image=None,
+    ):
+        os.makedirs(output_dir, exist_ok=True)
+        _logger.debug("Created all folders in path", extra={"output_directory": output_dir})
+
+        if model_uri:
+            model_cwd = os.path.join(output_dir, _MODEL_DIR_NAME)
+            pathlib.Path(model_cwd).mkdir(parents=True, exist_ok=True)
+            model_path = _download_artifact_from_uri(model_uri, output_path=model_cwd)
+            base_image = base_image or self._get_base_image(model_path, install_java)
+            env_manager = self._env_manager or em.LOCAL
+
+            if base_image.startswith("python"):
+                # we can directly use local env for python image
+                if env_manager in [em.CONDA, em.VIRTUALENV]:
+                    # we can directly use ubuntu image for conda and virtualenv
+                    base_image = UBUNTU_BASE_IMAGE
+            elif base_image == UBUNTU_BASE_IMAGE:
+                env_manager = self._env_manager or em.VIRTUALENV
+                # installing python on ubuntu image is problematic and not recommended officially
+                # , so we recommend using conda or virtualenv instead on ubuntu image
+                if env_manager == em.LOCAL:
+                    raise MlflowException.invalid_parameter_value(LOCAL_ENV_MANAGER_ERROR_MESSAGE)
+
+            model_install_steps = self._model_installation_steps(
+                model_path, env_manager, install_mlflow, enable_mlserver
+            )
+            entrypoint = f"from mlflow.models import container as C; C._serve('{env_manager}')"
+
+        # if no model_uri specified, user must use virtualenv or conda env based on ubuntu image
+        else:
+            base_image = base_image or UBUNTU_BASE_IMAGE
+            env_manager = self._env_manager or em.VIRTUALENV
+            if env_manager == em.LOCAL:
+                raise MlflowException.invalid_parameter_value(LOCAL_ENV_MANAGER_ERROR_MESSAGE)
+
+            model_install_steps = ""
+            # If model_uri is not specified, dependencies are installed at runtime
+            entrypoint = (
+                self._get_install_pyfunc_deps_cmd(env_manager, install_mlflow, enable_mlserver)
+                + f" C._serve('{env_manager}')"
+            )
+
+        dockerfile_text = docker_utils.generate_dockerfile(
+            output_dir=output_dir,
+            base_image=base_image,
+            model_install_steps=model_install_steps,
+            entrypoint=entrypoint,
+            env_manager=env_manager,
+            mlflow_home=mlflow_home,
+            enable_mlserver=enable_mlserver,
+            # always disable env creation at runtime for pyfunc
+            disable_env_creation_at_runtime=True,
+        )
+        _logger.debug("generated dockerfile at {output_dir}", extra={"dockerfile": dockerfile_text})
+
+    def _get_base_image(self, model_path: str, install_java: bool) -> str:
+        """
+        Determine the base image to use for the Dockerfile.
+
+        We use Python slim base image when all the following conditions are met:
+          1. Model URI is specified by the user
+          2. Model flavor does not require Java
+          3. Python version is specified in the model
+
+        Returns:
+            Either the Ubuntu base image or the Python slim base image.
+        """
+        # Check if the model requires Java
+        if not install_java:
+            flavors = _get_all_flavor_configurations(model_path).keys()
+            if java_flavors := JAVA_FLAVORS & flavors:
+                _logger.info(f"Detected java flavors {java_flavors}, installing Java in the image")
+                install_java = True
+
+        # Use ubuntu base image if Java is required
+        if install_java:
+            return UBUNTU_BASE_IMAGE
+
+        # Get Python version from MLmodel
+        try:
+            env_conf = Model.load(model_path).flavors[pyfunc.FLAVOR_NAME][pyfunc.ENV][em.VIRTUALENV]
+            python_env_config_path = os.path.join(model_path, env_conf)
+
+            python_env = _PythonEnv.from_yaml(python_env_config_path)
+            return PYTHON_SLIM_BASE_IMAGE.format(version=python_env.python)
+        except Exception as e:
+            _logger.warning(
+                f"Failed to determine Python version from {model_path}. "
+                f"Defaulting to {UBUNTU_BASE_IMAGE}. Error: {e}"
+            )
+            return UBUNTU_BASE_IMAGE
+
+    def _model_installation_steps(self, model_path, env_manager, install_mlflow, enable_mlserver):
+        model_dir = str(posixpath.join(_MODEL_DIR_NAME, os.path.basename(model_path)))
+        # Copy model to image if model_uri is specified
+        steps = (
+            "# Copy model to image and install dependencies\n"
+            f"COPY {model_dir} /opt/ml/model\nRUN python -c "
+        )
+        steps += (
+            f'"{self._get_install_pyfunc_deps_cmd(env_manager, install_mlflow, enable_mlserver)}"'
+        )
+
+        # Install flavor-specific dependencies if needed
+        flavors = _get_all_flavor_configurations(model_path).keys()
+        for flavor in flavors:
+            if flavor in FLAVOR_SPECIFIC_APT_PACKAGES:
+                packages = " ".join(FLAVOR_SPECIFIC_APT_PACKAGES[flavor])
+                steps += f"\nRUN apt-get install -y --no-install-recommends {packages}"
+
+        return steps
+
+    def _get_install_pyfunc_deps_cmd(
+        self, env_manager: str, install_mlflow: bool, enable_mlserver: bool
+    ):
+        return (
+            "from mlflow.models import container as C; "
+            f"C._install_pyfunc_deps('/opt/ml/model', install_mlflow={install_mlflow}, "
+            f"enable_mlserver={enable_mlserver}, env_manager='{env_manager}');"
+        )

venv/lib/python3.9/site-packages/mlflow/pyfunc/context.py

@@ -0,0 +1,70 @@
+import contextlib
+from contextvars import ContextVar
+from dataclasses import dataclass
+from typing import Optional
+
+# A thread local variable to store the context of the current prediction request.
+# This is particularly used to associate logs/traces with a specific prediction request in the
+# caller side. The context variable is intended to be set by the called before invoking the
+# predict method, using the set_prediction_context context manager.
+_PREDICTION_REQUEST_CTX = ContextVar("mlflow_prediction_request_context", default=None)
+
+
+@dataclass
+class Context:
+    # A unique identifier for the current prediction request.
+    request_id: Optional[str] = None
+    # Whether the current prediction request is as a part of MLflow model evaluation.
+    is_evaluate: bool = False
+    # The schema of the dependencies to be added into the tag of trace info.
+    dependencies_schemas: Optional[dict] = None
+
+    def update(self, **kwargs):
+        for key, value in kwargs.items():
+            if hasattr(self, key):
+                setattr(self, key, value)
+            else:
+                raise AttributeError(f"Context has no attribute named '{key}'")
+
+
+@contextlib.contextmanager
+def set_prediction_context(context: Optional[Context]):
+    """
+    Set the context for the current prediction request. The context will be set as a thread-local
+    variable and will be accessible globally within the same thread.
+
+    Args:
+        context: The context for the current prediction request.
+    """
+    if context and not isinstance(context, Context):
+        raise TypeError(f"Expected context to be an instance of Context, but got: {context}")
+
+    token = _PREDICTION_REQUEST_CTX.set(context)
+    try:
+        yield
+    finally:
+        _PREDICTION_REQUEST_CTX.reset(token)
+
+
+def get_prediction_context() -> Optional[Context]:
+    """
+    Get the context for the current prediction request. The context is thread-local and is set
+    using the set_prediction_context context manager.
+
+    Returns:
+        The context for the current prediction request, or None if no context is set.
+    """
+    return _PREDICTION_REQUEST_CTX.get()
+
+
+@contextlib.contextmanager
+def maybe_set_prediction_context(context: Optional[Context]):
+    """
+    Set the prediction context if the given context
+    is not None. Otherwise no-op.
+    """
+    if context:
+        with set_prediction_context(context):
+            yield
+    else:
+        yield

venv/lib/python3.9/site-packages/mlflow/pyfunc/dbconnect_artifact_cache.py

@@ -0,0 +1,166 @@
+import json
+import os
+import subprocess
+import tarfile
+
+from mlflow.utils.databricks_utils import is_in_databricks_runtime
+from mlflow.utils.file_utils import get_or_create_tmp_dir
+
+_CACHE_MAP_FILE_NAME = "db_connect_artifact_cache.json"
+
+
+class DBConnectArtifactCache:
+    """
+    Manages Databricks Connect artifacts cache.
+    Note it doesn't support OSS Spark Connect.
+
+    This class can be used in the following environment:
+     - Databricks shared cluster python notebook REPL
+     - Databricks Serverless python notebook REPL
+     - Databricks connect client python REPL that connects to remote Databricks Serverless
+     - Databricks connect client python REPL that connects to remote Databricks shared cluster
+
+    .. code-block:: python
+        :caption: Example
+
+        # client side code
+        db_artifact_cache = DBConnectArtifactCache.get_or_create()
+        db_artifact_cache.add_artifact_archive("archive1", "/tmp/archive1.tar.gz")
+
+
+        @pandas_udf(...)
+        def my_udf(x):
+            # we can get the unpacked archive files in `archive1_unpacked_dir`
+            archive1_unpacked_dir = db_artifact_cache.get("archive1")
+    """
+
+    _global_cache = None
+
+    @staticmethod
+    def get_or_create(spark):
+        if (
+            DBConnectArtifactCache._global_cache is None
+            or spark is not DBConnectArtifactCache._global_cache._spark
+        ):
+            DBConnectArtifactCache._global_cache = DBConnectArtifactCache(spark)
+            cache_file = os.path.join(get_or_create_tmp_dir(), _CACHE_MAP_FILE_NAME)
+            if is_in_databricks_runtime() and os.path.exists(cache_file):
+                # In databricks runtime (shared cluster or Serverless), when you restart the
+                # notebook REPL by %restart_python or dbutils.library.restartPython(), the
+                # DBConnect session is still preserved. So in this case, we can reuse the cached
+                # artifact files.
+                # So that when adding artifact, the cache map is serialized to local disk file
+                # `db_connect_artifact_cache.json` and after REPL restarts,
+                # `DBConnectArtifactCache` restores the cache map by loading data from the file.
+                with open(cache_file) as f:
+                    DBConnectArtifactCache._global_cache._cache = json.load(f)
+        return DBConnectArtifactCache._global_cache
+
+    def __init__(self, spark):
+        self._spark = spark
+        self._cache = {}
+
+    def __getstate__(self):
+        """
+        The `DBConnectArtifactCache` instance is created in Databricks Connect client side,
+        and it will be pickled to Databricks Connect UDF sandbox
+        (see `get_unpacked_artifact_dir` method), but Spark Connect client object is
+        not pickle-able, we need to skip this field.
+        """
+        state = self.__dict__.copy()
+        # Don't pickle `_spark`
+        del state["_spark"]
+        return state
+
+    def __setstate__(self, state):
+        self.__dict__.update(state)
+        self._spark = None
+
+    def has_cache_key(self, cache_key):
+        return cache_key in self._cache
+
+    def add_artifact_archive(self, cache_key, artifact_archive_path):
+        """
+        Add an artifact archive file to Databricks connect cache.
+        The archive file must be 'tar.gz' format.
+        You can only call this method in Databricks Connect client side.
+        """
+        if not artifact_archive_path.endswith(".tar.gz"):
+            raise RuntimeError(
+                "'add_artifact_archive' only supports archive file in 'tar.gz' format."
+            )
+
+        archive_file_name = os.path.basename(artifact_archive_path)
+        if cache_key not in self._cache:
+            self._spark.addArtifact(artifact_archive_path, archive=True)
+            self._cache[cache_key] = archive_file_name
+
+        if is_in_databricks_runtime():
+            with open(os.path.join(get_or_create_tmp_dir(), _CACHE_MAP_FILE_NAME), "w") as f:
+                json.dump(self._cache, f)
+
+    def get_unpacked_artifact_dir(self, cache_key):
+        """
+        Get unpacked artifact directory path, you can only call this method
+        inside Databricks Connect spark UDF sandbox.
+        """
+        if cache_key not in self._cache:
+            raise RuntimeError(f"The artifact '{cache_key}' does not exist.")
+        archive_file_name = self._cache[cache_key]
+
+        session_id = os.environ.get("DB_SESSION_UUID")
+        if not session_id:
+            # If 'DB_SESSION_UUID' environment variable does not exist, it means it is running
+            # in a dedicated mode Spark cluster.
+            return os.path.join(os.getcwd(), archive_file_name)
+
+        relative_path = os.path.join("artifacts", session_id, "archives", archive_file_name)
+        single_driver_root = "/local_disk0/.ephemeral_nfs"
+        single_candidate = os.path.join(single_driver_root, relative_path)
+        if os.path.exists(single_candidate):
+            return single_candidate
+
+        multi_driver_root = "/local_disk0/.ephemeral_nfs_multi_driver"
+        if (
+            os.environ.get("AETHER_MULTI_DRIVER_ENABLED", "false") == "true"
+            and os.environ.get("AETHER_MULTI_DRIVER_NOTEBOOK_LIBRARY_ENABLED", "false") == "true"
+            and os.path.isdir(multi_driver_root)
+        ):
+            try:
+                children = sorted(os.listdir(multi_driver_root))
+            except OSError:
+                children = []
+
+            for child in children:
+                child_candidate = os.path.join(multi_driver_root, child, relative_path)
+                if os.path.exists(child_candidate):
+                    return child_candidate
+
+        # Fall back to the original single-driver location to preserve previous behaviour.
+        return single_candidate
+
+
+def archive_directory(input_dir, archive_file_path):
+    """
+    Archive the `input_dir` directory, save the archive file to `archive_file_path`,
+    the generated archive file is 'tar.gz' format.
+    Note: all symlink files in the input directory are kept as it is in the archive file.
+    """
+
+    archive_file_path = os.path.abspath(archive_file_path)
+    # Note: `shutil.make_archive` doesn't work because it replaces symlink files with
+    #  the file symlink pointing to, which is not the expected behavior in our usage.
+    #  We need to pack the python and virtualenv environment, which contains a bunch of
+    #  symlink files.
+    subprocess.check_call(
+        ["tar", "-czf", archive_file_path, *os.listdir(input_dir)],
+        cwd=input_dir,
+    )
+    return archive_file_path
+
+
+def extract_archive_to_dir(archive_path, dest_dir):
+    os.makedirs(dest_dir, exist_ok=True)
+    with tarfile.open(archive_path, "r") as tar:
+        tar.extractall(path=dest_dir)
+    return dest_dir

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

@@ -0,0 +1,7 @@
+import mlflow.pyfunc.loaders.chat_agent
+import mlflow.pyfunc.loaders.chat_model
+import mlflow.pyfunc.loaders.code_model
+from mlflow.utils.pydantic_utils import IS_PYDANTIC_V2_OR_NEWER
+
+if IS_PYDANTIC_V2_OR_NEWER:
+    import mlflow.pyfunc.loaders.responses_agent  # noqa: F401

venv/lib/python3.9/site-packages/mlflow/pyfunc/loaders/chat_agent.py

@@ -0,0 +1,117 @@
+from typing import Any, Generator, Optional
+
+import pydantic
+
+from mlflow.exceptions import MlflowException
+from mlflow.models.utils import _convert_llm_ndarray_to_list
+from mlflow.protos.databricks_pb2 import INTERNAL_ERROR
+from mlflow.pyfunc.model import (
+    _load_context_model_and_signature,
+)
+from mlflow.types.agent import (
+    ChatAgentChunk,
+    ChatAgentMessage,
+    ChatAgentResponse,
+    ChatContext,
+)
+from mlflow.types.type_hints import model_validate
+from mlflow.utils.annotations import experimental
+
+
+def _load_pyfunc(model_path: str, model_config: Optional[dict[str, Any]] = None):
+    _, chat_agent, _ = _load_context_model_and_signature(model_path, model_config)
+    return _ChatAgentPyfuncWrapper(chat_agent)
+
+
+@experimental
+class _ChatAgentPyfuncWrapper:
+    """
+    Wrapper class that converts dict inputs to pydantic objects accepted by :class:`~ChatAgent`.
+    """
+
+    def __init__(self, chat_agent):
+        """
+        Args:
+            chat_agent: An instance of a subclass of :class:`~ChatAgent`.
+        """
+        self.chat_agent = chat_agent
+
+    def get_raw_model(self):
+        """
+        Returns the underlying model.
+        """
+        return self.chat_agent
+
+    def _convert_input(
+        self, model_input
+    ) -> tuple[list[ChatAgentMessage], Optional[ChatContext], Optional[dict[str, Any]]]:
+        import pandas
+
+        if isinstance(model_input, dict):
+            dict_input = model_input
+        elif isinstance(model_input, pandas.DataFrame):
+            dict_input = {
+                k: _convert_llm_ndarray_to_list(v[0])
+                for k, v in model_input.to_dict(orient="list").items()
+            }
+        else:
+            raise MlflowException(
+                "Unsupported model input type. Expected a dict or pandas.DataFrame, but got "
+                f"{type(model_input)} instead.",
+                error_code=INTERNAL_ERROR,
+            )
+
+        messages = [ChatAgentMessage(**message) for message in dict_input.get("messages", [])]
+        context = ChatContext(**dict_input["context"]) if "context" in dict_input else None
+        custom_inputs = dict_input.get("custom_inputs", None)
+
+        return messages, context, custom_inputs
+
+    def _response_to_dict(self, response, pydantic_class) -> dict[str, Any]:
+        if isinstance(response, pydantic_class):
+            return response.model_dump_compat(exclude_none=True)
+        try:
+            model_validate(pydantic_class, response)
+        except pydantic.ValidationError as e:
+            raise MlflowException(
+                message=(
+                    f"Model returned an invalid response. Expected a {pydantic_class.__name__} "
+                    f"object or dictionary with the same schema. Pydantic validation error: {e}"
+                ),
+                error_code=INTERNAL_ERROR,
+            ) from e
+        return response
+
+    def predict(self, model_input: dict[str, Any], params=None) -> dict[str, Any]:
+        """
+        Args:
+            model_input: A dict with the
+                :py:class:`ChatAgentRequest <mlflow.types.agent.ChatAgentRequest>` schema.
+            params: Unused in this function, but required in the signature because
+                `load_model_and_predict` in `utils/_capture_modules.py` expects a params field
+
+        Returns:
+            A dict with the (:py:class:`ChatAgentResponse <mlflow.types.agent.ChatAgentResponse>`)
+                schema.
+        """
+        messages, context, custom_inputs = self._convert_input(model_input)
+        response = self.chat_agent.predict(messages, context, custom_inputs)
+        return self._response_to_dict(response, ChatAgentResponse)
+
+    def predict_stream(
+        self, model_input: dict[str, Any], params=None
+    ) -> Generator[dict[str, Any], None, None]:
+        """
+        Args:
+            model_input: A dict with the
+                :py:class:`ChatAgentRequest <mlflow.types.agent.ChatAgentRequest>` schema.
+            params: Unused in this function, but required in the signature because
+                `load_model_and_predict` in `utils/_capture_modules.py` expects a params field
+
+        Returns:
+            A generator over dicts with the
+                (:py:class:`ChatAgentChunk <mlflow.types.agent.ChatAgentChunk>`) schema.
+        """
+        messages, context, custom_inputs = self._convert_input(model_input)
+        for response in self.chat_agent.predict_stream(messages, context, custom_inputs):
+            yield self._response_to_dict(response, ChatAgentChunk)

venv/lib/python3.9/site-packages/mlflow/pyfunc/loaders/chat_model.py

@@ -0,0 +1,127 @@
+import inspect
+import logging
+from typing import Any, Generator, Optional
+
+from mlflow.exceptions import MlflowException
+from mlflow.models.utils import _convert_llm_ndarray_to_list
+from mlflow.protos.databricks_pb2 import INTERNAL_ERROR
+from mlflow.pyfunc.model import (
+    _load_context_model_and_signature,
+)
+from mlflow.types.llm import ChatCompletionChunk, ChatCompletionResponse, ChatMessage, ChatParams
+from mlflow.utils.annotations import experimental
+
+_logger = logging.getLogger(__name__)
+
+
+def _load_pyfunc(model_path: str, model_config: Optional[dict[str, Any]] = None):
+    context, chat_model, signature = _load_context_model_and_signature(model_path, model_config)
+    return _ChatModelPyfuncWrapper(chat_model=chat_model, context=context, signature=signature)
+
+
+@experimental
+class _ChatModelPyfuncWrapper:
+    """
+    Wrapper class that converts dict inputs to pydantic objects accepted by :class:`~ChatModel`.
+    """
+
+    def __init__(self, chat_model, context, signature):
+        """
+        Args:
+            chat_model: An instance of a subclass of :class:`~ChatModel`.
+            context: A :class:`~PythonModelContext` instance containing artifacts that
+                        ``chat_model`` may use when performing inference.
+            signature: :class:`~ModelSignature` instance describing model input and output.
+        """
+        self.chat_model = chat_model
+        self.context = context
+        self.signature = signature
+
+    def get_raw_model(self):
+        """
+        Returns the underlying model.
+        """
+        return self.chat_model
+
+    def _convert_input(self, model_input):
+        import pandas
+
+        if isinstance(model_input, dict):
+            dict_input = model_input
+        elif isinstance(model_input, pandas.DataFrame):
+            dict_input = {
+                k: _convert_llm_ndarray_to_list(v[0])
+                for k, v in model_input.to_dict(orient="list").items()
+            }
+        else:
+            raise MlflowException(
+                "Unsupported model input type. Expected a dict or pandas.DataFrame, "
+                f"but got {type(model_input)} instead.",
+                error_code=INTERNAL_ERROR,
+            )
+
+        messages = [ChatMessage.from_dict(message) for message in dict_input.pop("messages", [])]
+        params = ChatParams.from_dict(dict_input)
+        return messages, params
+
+    def predict(
+        self, model_input: dict[str, Any], params: Optional[dict[str, Any]] = None
+    ) -> dict[str, Any]:
+        """
+        Args:
+            model_input: Model input data in the form of a chat request.
+            params: Additional parameters to pass to the model for inference.
+                       Unused in this implementation, as the params are handled
+                       via ``self._convert_input()``.
+
+        Returns:
+            Model predictions in :py:class:`~ChatCompletionResponse` format.
+        """
+        messages, params = self._convert_input(model_input)
+        parameters = inspect.signature(self.chat_model.predict).parameters
+        if "context" in parameters or len(parameters) == 3:
+            response = self.chat_model.predict(self.context, messages, params)
+        else:
+            response = self.chat_model.predict(messages, params)
+        return self._response_to_dict(response)
+
+    def _response_to_dict(self, response: ChatCompletionResponse) -> dict[str, Any]:
+        if not isinstance(response, ChatCompletionResponse):
+            raise MlflowException(
+                "Model returned an invalid response. Expected a ChatCompletionResponse, but "
+                f"got {type(response)} instead.",
+                error_code=INTERNAL_ERROR,
+            )
+        return response.to_dict()
+
+    def _streaming_response_to_dict(self, response: ChatCompletionChunk) -> dict[str, Any]:
+        if not isinstance(response, ChatCompletionChunk):
+            raise MlflowException(
+                "Model returned an invalid response. Expected a ChatCompletionChunk, but "
+                f"got {type(response)} instead.",
+                error_code=INTERNAL_ERROR,
+            )
+        return response.to_dict()
+
+    def predict_stream(
+        self, model_input: dict[str, Any], params: Optional[dict[str, Any]] = None
+    ) -> Generator[dict[str, Any], None, None]:
+        """
+        Args:
+            model_input: Model input data in the form of a chat request.
+            params: Additional parameters to pass to the model for inference.
+                       Unused in this implementation, as the params are handled
+                       via ``self._convert_input()``.
+
+        Returns:
+            Generator over model predictions in :py:class:`~ChatCompletionChunk` format.
+        """
+        messages, params = self._convert_input(model_input)
+        parameters = inspect.signature(self.chat_model.predict_stream).parameters
+        if "context" in parameters or len(parameters) == 3:
+            stream = self.chat_model.predict_stream(self.context, messages, params)
+        else:
+            stream = self.chat_model.predict_stream(messages, params)
+
+        for response in stream:
+            yield self._streaming_response_to_dict(response)

venv/lib/python3.9/site-packages/mlflow/pyfunc/loaders/code_model.py

@@ -0,0 +1,31 @@
+from typing import Any, Optional
+
+from mlflow.pyfunc.loaders.chat_agent import _ChatAgentPyfuncWrapper
+from mlflow.pyfunc.loaders.chat_model import _ChatModelPyfuncWrapper
+from mlflow.pyfunc.model import (
+    ChatAgent,
+    ChatModel,
+    _load_context_model_and_signature,
+    _PythonModelPyfuncWrapper,
+)
+
+try:
+    from mlflow.pyfunc.model import ResponsesAgent
+
+    IS_RESPONSES_AGENT_AVAILABLE = True
+except ImportError:
+    IS_RESPONSES_AGENT_AVAILABLE = False
+
+
+def _load_pyfunc(local_path: str, model_config: Optional[dict[str, Any]] = None):
+    context, model, signature = _load_context_model_and_signature(local_path, model_config)
+    if isinstance(model, ChatModel):
+        return _ChatModelPyfuncWrapper(model, context, signature)
+    elif isinstance(model, ChatAgent):
+        return _ChatAgentPyfuncWrapper(model)
+    elif IS_RESPONSES_AGENT_AVAILABLE and isinstance(model, ResponsesAgent):
+        from mlflow.pyfunc.loaders.responses_agent import _ResponsesAgentPyfuncWrapper
+
+        return _ResponsesAgentPyfuncWrapper(model)
+    else:
+        return _PythonModelPyfuncWrapper(model, context, signature)

venv/lib/python3.9/site-packages/mlflow/pyfunc/loaders/responses_agent.py

@@ -0,0 +1,108 @@
+from typing import Any, Generator, Optional
+
+import pydantic
+
+from mlflow.exceptions import MlflowException
+from mlflow.models.utils import _convert_llm_ndarray_to_list
+from mlflow.protos.databricks_pb2 import INTERNAL_ERROR
+from mlflow.pyfunc.model import _load_context_model_and_signature
+from mlflow.types.type_hints import model_validate
+from mlflow.utils.annotations import experimental
+from mlflow.utils.pydantic_utils import IS_PYDANTIC_V2_OR_NEWER
+
+if not IS_PYDANTIC_V2_OR_NEWER:
+    raise ImportError(
+        "ResponsesAgent and its pydantic classes are not supported in pydantic v1. "
+        "Please upgrade to pydantic v2 or newer to use ResponsesAgent.",
+    )
+
+from mlflow.types.responses import ResponsesRequest, ResponsesResponse, ResponsesStreamEvent
+
+
+def _load_pyfunc(model_path: str, model_config: Optional[dict[str, Any]] = None):
+    _, responses_agent, _ = _load_context_model_and_signature(model_path, model_config)
+    return _ResponsesAgentPyfuncWrapper(responses_agent)
+
+
+@experimental
+class _ResponsesAgentPyfuncWrapper:
+    """
+    Wrapper class that converts dict inputs to pydantic objects accepted by
+    :class:`~ResponsesAgent`.
+    """
+
+    def __init__(self, responses_agent):
+        self.responses_agent = responses_agent
+
+    def get_raw_model(self):
+        """
+        Returns the underlying model.
+        """
+        return self.responses_agent
+
+    def _convert_input(self, model_input) -> ResponsesRequest:
+        import pandas
+
+        if isinstance(model_input, pandas.DataFrame):
+            model_input = {
+                k: _convert_llm_ndarray_to_list(v[0])
+                for k, v in model_input.to_dict(orient="list").items()
+            }
+        elif not isinstance(model_input, dict):
+            raise MlflowException(
+                "Unsupported model input type. Expected a dict or pandas.DataFrame, but got "
+                f"{type(model_input)} instead.",
+                error_code=INTERNAL_ERROR,
+            )
+        return ResponsesRequest(**model_input)
+
+    def _response_to_dict(self, response, pydantic_class) -> dict[str, Any]:
+        if isinstance(response, pydantic_class):
+            return response.model_dump_compat(exclude_none=True)
+        try:
+            model_validate(pydantic_class, response)
+        except pydantic.ValidationError as e:
+            raise MlflowException(
+                message=(
+                    f"Model returned an invalid response. Expected a {pydantic_class.__name__} "
+                    f"object or dictionary with the same schema. Pydantic validation error: {e}"
+                ),
+                error_code=INTERNAL_ERROR,
+            ) from e
+        return response
+
+    def predict(self, model_input: dict[str, Any], params=None) -> dict[str, Any]:
+        """
+        Args:
+            model_input: A dict with the
+                :py:class:`ResponsesRequest <mlflow.types.responses.ResponsesRequest>` schema.
+            params: Unused in this function, but required in the signature because
+                `load_model_and_predict` in `utils/_capture_modules.py` expects a params field
+
+        Returns:
+            A dict with the
+            (:py:class:`ResponsesResponse <mlflow.types.responses.ResponsesResponse>`)
+            schema.
+        """
+        request = self._convert_input(model_input)
+        response = self.responses_agent.predict(request)
+        return self._response_to_dict(response, ResponsesResponse)
+
+    def predict_stream(
+        self, model_input: dict[str, Any], params=None
+    ) -> Generator[dict[str, Any], None, None]:
+        """
+        Args:
+            model_input: A dict with the
+                :py:class:`ResponsesRequest <mlflow.types.responses.ResponsesRequest>` schema.
+            params: Unused in this function, but required in the signature because
+                `load_model_and_predict` in `utils/_capture_modules.py` expects a params field
+
+        Returns:
+            A generator over dicts with the
+                (:py:class:`ResponsesStreamEvent <mlflow.types.responses.ResponsesStreamEvent>`)
+                schema.
+        """
+        request = self._convert_input(model_input)
+        for response in self.responses_agent.predict_stream(request):
+            yield self._response_to_dict(response, ResponsesStreamEvent)

venv/lib/python3.9/site-packages/mlflow/pyfunc/mlserver.py

@@ -0,0 +1,46 @@
+import logging
+import os
+from typing import Optional
+
+_logger = logging.getLogger(__name__)
+
+MLServerMLflowRuntime = "mlserver_mlflow.MLflowRuntime"
+MLServerDefaultModelName = "mlflow-model"
+
+
+def get_cmd(
+    model_uri: str,
+    port: Optional[int] = None,
+    host: Optional[str] = None,
+    timeout: Optional[int] = None,
+    nworkers: Optional[int] = None,
+    model_name: Optional[str] = None,
+    model_version: Optional[str] = None,
+) -> tuple[str, dict[str, str]]:
+    cmd = f"mlserver start {model_uri}"
+
+    cmd_env = os.environ.copy()
+
+    if port:
+        cmd_env["MLSERVER_HTTP_PORT"] = str(port)
+
+    if host:
+        cmd_env["MLSERVER_HOST"] = host
+
+    if timeout:
+        _logger.warning("Timeout is not yet supported in MLServer.")
+
+    if nworkers:
+        cmd_env["MLSERVER_PARALLEL_WORKERS"] = str(nworkers)
+
+    # give precedence to user env var input
+    cmd_env["MLSERVER_MODEL_NAME"] = (
+        cmd_env.get("MLSERVER_MODEL_NAME") or model_name or MLServerDefaultModelName
+    )
+    if model_version and not cmd_env.get("MLSERVER_MODEL_VERSION"):
+        cmd_env["MLSERVER_MODEL_VERSION"] = model_version
+
+    cmd_env["MLSERVER_MODEL_IMPLEMENTATION"] = MLServerMLflowRuntime
+    cmd_env["MLSERVER_MODEL_URI"] = model_uri
+
+    return cmd, cmd_env

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

@@ -0,0 +1,606 @@
+"""
+Scoring server for python model format.
+The passed int model is expected to have function:
+   predict(pandas.Dataframe) -> pandas.DataFrame
+
+Input, expected in text/csv or application/json format,
+is parsed into pandas.DataFrame and passed to the model.
+
+Defines four endpoints:
+    /ping used for health check
+    /health (same as /ping)
+    /version used for getting the mlflow version
+    /invocations used for scoring
+"""
+
+import asyncio
+import inspect
+import json
+import logging
+import os
+import shlex
+import sys
+import traceback
+from functools import wraps
+from typing import Any, NamedTuple, Optional
+
+from mlflow.environment_variables import (
+    _MLFLOW_IS_IN_SERVING_ENVIRONMENT,
+    MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT,
+)
+
+# NB: We need to be careful what we import form mlflow here. Scoring server is used from within
+# model's conda environment. The version of mlflow doing the serving (outside) and the version of
+# mlflow in the model's conda environment (inside) can differ. We should therefore keep mlflow
+# dependencies to the minimum here.
+# ALl of the mlflow dependencies below need to be backwards compatible.
+from mlflow.exceptions import MlflowException
+from mlflow.pyfunc.model import _log_warning_if_params_not_in_predict_signature
+from mlflow.types import ParamSchema, Schema
+from mlflow.utils import reraise
+from mlflow.utils.annotations import deprecated
+from mlflow.utils.file_utils import path_to_local_file_uri
+from mlflow.utils.proto_json_utils import (
+    MlflowInvalidInputException,
+    NumpyEncoder,
+    _get_jsonable_obj,
+    dataframe_from_parsed_json,
+    parse_tf_serving_input,
+)
+from mlflow.version import VERSION
+
+try:
+    from mlflow.pyfunc import PyFuncModel, load_model
+except ImportError:
+    from mlflow.pyfunc import load_pyfunc as load_model
+from io import StringIO
+
+from mlflow.protos.databricks_pb2 import BAD_REQUEST, INVALID_PARAMETER_VALUE
+from mlflow.pyfunc.utils.serving_data_parser import is_unified_llm_input
+
+_SERVER_MODEL_PATH = "__pyfunc_model_path__"
+SERVING_MODEL_CONFIG = "SERVING_MODEL_CONFIG"
+
+CONTENT_TYPE_CSV = "text/csv"
+CONTENT_TYPE_JSON = "application/json"
+
+CONTENT_TYPES = [
+    CONTENT_TYPE_CSV,
+    CONTENT_TYPE_JSON,
+]
+
+_logger = logging.getLogger(__name__)
+
+DF_RECORDS = "dataframe_records"
+DF_SPLIT = "dataframe_split"
+INSTANCES = "instances"
+INPUTS = "inputs"
+
+SUPPORTED_FORMATS = {DF_RECORDS, DF_SPLIT, INSTANCES, INPUTS}
+SERVING_PARAMS_KEY = "params"
+
+REQUIRED_INPUT_FORMAT = (
+    f"The input must be a JSON dictionary with exactly one of the input fields {SUPPORTED_FORMATS}"
+)
+SCORING_PROTOCOL_CHANGE_INFO = (
+    "IMPORTANT: The MLflow Model scoring protocol has changed in MLflow version 2.0. If you are"
+    " seeing this error, you are likely using an outdated scoring request format. To resolve the"
+    " error, either update your request format or adjust your MLflow Model's requirements file to"
+    " specify an older version of MLflow (for example, change the 'mlflow' requirement specifier"
+    " to 'mlflow==1.30.0'). If you are making a request using the MLflow client"
+    " (e.g. via `mlflow.pyfunc.spark_udf()`), upgrade your MLflow client to a version >= 2.0 in"
+    " order to use the new request format. For more information about the updated MLflow"
+    " Model scoring protocol in MLflow 2.0, see"
+    " https://mlflow.org/docs/latest/models.html#deploy-mlflow-models."
+)
+
+
+def load_model_with_mlflow_config(model_uri):
+    extra_kwargs = {}
+    if model_config_json := os.environ.get(SERVING_MODEL_CONFIG):
+        extra_kwargs["model_config"] = json.loads(model_config_json)
+
+    return load_model(model_uri, **extra_kwargs)
+
+
+# Keep this method to maintain compatibility with MLServer
+# https://github.com/SeldonIO/MLServer/blob/caa173ab099a4ec002a7c252cbcc511646c261a6/runtimes/mlflow/mlserver_mlflow/runtime.py#L13C5-L13C31
+@deprecated("infer_and_parse_data", "2.6.0")
+def infer_and_parse_json_input(json_input, schema: Schema = None):
+    """
+    Args:
+        json_input: A JSON-formatted string representation of TF serving input or a Pandas
+                    DataFrame, or a stream containing such a string representation.
+        schema: Optional schema specification to be used during parsing.
+    """
+    if isinstance(json_input, dict):
+        decoded_input = json_input
+    else:
+        try:
+            decoded_input = json.loads(json_input)
+        except json.decoder.JSONDecodeError as ex:
+            raise MlflowException(
+                message=(
+                    "Failed to parse input from JSON. Ensure that input is a valid JSON"
+                    f" formatted string. Error: '{ex}'. Input: \n{json_input}\n"
+                ),
+                error_code=BAD_REQUEST,
+            )
+    if isinstance(decoded_input, dict):
+        format_keys = set(decoded_input.keys()).intersection(SUPPORTED_FORMATS)
+        if len(format_keys) != 1:
+            message = f"Received dictionary with input fields: {list(decoded_input.keys())}"
+            raise MlflowException(
+                message=f"{REQUIRED_INPUT_FORMAT}. {message}. {SCORING_PROTOCOL_CHANGE_INFO}",
+                error_code=BAD_REQUEST,
+            )
+        input_format = format_keys.pop()
+        if input_format in (INSTANCES, INPUTS):
+            return parse_tf_serving_input(decoded_input, schema=schema)
+
+        elif input_format in (DF_SPLIT, DF_RECORDS):
+            # NB: skip the dataframe_ prefix
+            pandas_orient = input_format[10:]
+            return dataframe_from_parsed_json(
+                decoded_input[input_format], pandas_orient=pandas_orient, schema=schema
+            )
+    elif isinstance(decoded_input, list):
+        message = "Received a list"
+        raise MlflowException(
+            message=f"{REQUIRED_INPUT_FORMAT}. {message}. {SCORING_PROTOCOL_CHANGE_INFO}",
+            error_code=BAD_REQUEST,
+        )
+    else:
+        message = f"Received unexpected input type '{type(decoded_input)}'"
+        raise MlflowException(
+            message=f"{REQUIRED_INPUT_FORMAT}. {message}.", error_code=BAD_REQUEST
+        )
+
+
+def _decode_json_input(json_input):
+    """
+    Args:
+        json_input: A JSON-formatted string representation of TF serving input or a Pandas
+                    DataFrame, or a stream containing such a string representation.
+
+    Returns:
+        A dictionary representation of the JSON input.
+    """
+    if isinstance(json_input, dict):
+        return json_input
+
+    try:
+        decoded_input = json.loads(json_input)
+    except json.decoder.JSONDecodeError as ex:
+        raise MlflowInvalidInputException(
+            "Ensure that input is a valid JSON formatted string. "
+            f"Error: '{ex!r}'\nInput: \n{json_input}\n"
+        ) from ex
+
+    if isinstance(decoded_input, dict):
+        return decoded_input
+    if isinstance(decoded_input, list):
+        raise MlflowInvalidInputException(f"{REQUIRED_INPUT_FORMAT}. Received a list.")
+
+    raise MlflowInvalidInputException(
+        f"{REQUIRED_INPUT_FORMAT}. Received unexpected input type '{type(decoded_input)}."
+    )
+
+
+def _split_data_and_params_for_llm_input(json_input, param_schema: Optional[ParamSchema]):
+    data = {}
+    params = {}
+    schema_params = {param.name for param in param_schema.params} if param_schema else {}
+
+    for key, value in json_input.items():
+        # if the model defines a param schema, then we can add
+        # it to the params dict. otherwise, add it to the data
+        # dict to prevent it from being ignored at inference time
+        if key in schema_params:
+            params[key] = value
+        else:
+            data[key] = value
+
+    return data, params
+
+
+def _split_data_and_params(json_input):
+    input_dict = _decode_json_input(json_input)
+    data = {k: v for k, v in input_dict.items() if k in SUPPORTED_FORMATS}
+    params = input_dict.pop(SERVING_PARAMS_KEY, None)
+    return data, params
+
+
+def infer_and_parse_data(data, schema: Schema = None):
+    """
+    Args:
+        data: A dictionary representation of TF serving input or a Pandas
+            DataFrame, or a stream containing such a string representation.
+        schema: Optional schema specification to be used during parsing.
+    """
+
+    format_keys = set(data.keys()).intersection(SUPPORTED_FORMATS)
+    if len(format_keys) != 1:
+        message = f"Received dictionary with input fields: {list(data.keys())}"
+        raise MlflowException(
+            message=f"{REQUIRED_INPUT_FORMAT}. {message}. {SCORING_PROTOCOL_CHANGE_INFO}",
+            error_code=BAD_REQUEST,
+        )
+    input_format = format_keys.pop()
+    if input_format in (INSTANCES, INPUTS):
+        return parse_tf_serving_input(data, schema=schema)
+
+    if input_format in (DF_SPLIT, DF_RECORDS):
+        pandas_orient = input_format[10:]  # skip the dataframe_ prefix
+        return dataframe_from_parsed_json(
+            data[input_format], pandas_orient=pandas_orient, schema=schema
+        )
+
+
+def parse_csv_input(csv_input, schema: Schema = None):
+    """
+    Args:
+        csv_input: A CSV-formatted string representation of a Pandas DataFrame, or a stream
+                   containing such a string representation.
+        schema: Optional schema specification to be used during parsing.
+    """
+    import pandas as pd
+
+    try:
+        if schema is None:
+            return pd.read_csv(csv_input)
+        else:
+            dtypes = dict(zip(schema.input_names(), schema.pandas_types()))
+            return pd.read_csv(csv_input, dtype=dtypes)
+    except Exception as e:
+        _handle_serving_error(
+            error_message=(
+                "Failed to parse input as a Pandas DataFrame. Ensure that the input is"
+                " a valid CSV-formatted Pandas DataFrame produced using the"
+                f" `pandas.DataFrame.to_csv()` method. Error: '{e}'"
+            ),
+            error_code=BAD_REQUEST,
+        )
+
+
+def unwrapped_predictions_to_json(raw_predictions, output):
+    predictions = _get_jsonable_obj(raw_predictions, pandas_orient="records")
+    return json.dump(predictions, output, cls=NumpyEncoder)
+
+
+def predictions_to_json(raw_predictions, output, metadata=None):
+    if metadata and "predictions" in metadata:
+        raise MlflowException(
+            "metadata cannot contain 'predictions' key", error_code=INVALID_PARAMETER_VALUE
+        )
+    predictions = _get_jsonable_obj(raw_predictions, pandas_orient="records")
+    return json.dump({"predictions": predictions, **(metadata or {})}, output, cls=NumpyEncoder)
+
+
+def _handle_serving_error(error_message, error_code, include_traceback=True):
+    """
+    Logs information about an exception thrown by model inference code that is currently being
+    handled and reraises it with the specified error message. The exception stack trace
+    is also included in the reraised error message.
+
+    Args:
+        error_message: A message for the reraised exception.
+        error_code: An appropriate error code for the reraised exception. This should be one of
+            the codes listed in the `mlflow.protos.databricks_pb2` proto.
+        include_traceback: Whether to include the current traceback in the returned error.
+    """
+    if include_traceback:
+        traceback_buf = StringIO()
+        traceback.print_exc(file=traceback_buf)
+        traceback_str = traceback_buf.getvalue()
+        e = MlflowException(message=error_message, error_code=error_code, stack_trace=traceback_str)
+    else:
+        e = MlflowException(message=error_message, error_code=error_code)
+    reraise(MlflowException, e)
+
+
+class InvocationsResponse(NamedTuple):
+    response: str
+    status: int
+    mimetype: str
+
+
+def invocations(data, content_type, model, input_schema):
+    type_parts = list(map(str.strip, content_type.split(";")))
+    mime_type = type_parts[0]
+    parameter_value_pairs = type_parts[1:]
+    parameter_values = {
+        key: value for pair in parameter_value_pairs for key, _, value in [pair.partition("=")]
+    }
+
+    charset = parameter_values.get("charset", "utf-8").lower()
+    if charset != "utf-8":
+        return InvocationsResponse(
+            response="The scoring server only supports UTF-8",
+            status=415,
+            mimetype="text/plain",
+        )
+
+    unexpected_content_parameters = set(parameter_values.keys()).difference({"charset"})
+    if unexpected_content_parameters:
+        return InvocationsResponse(
+            response=(
+                f"Unrecognized content type parameters: "
+                f"{', '.join(unexpected_content_parameters)}. "
+                f"{SCORING_PROTOCOL_CHANGE_INFO}"
+            ),
+            status=415,
+            mimetype="text/plain",
+        )
+
+    # The traditional JSON request/response format, wraps the data with one of the supported keys
+    # like "dataframe_split" and "predictions". For LLM use cases, we also support unwrapped JSON
+    # payload, to provide unified prediction interface.
+    should_parse_as_unified_llm_input = False
+
+    if mime_type == CONTENT_TYPE_CSV:
+        # Convert from CSV to pandas
+        if isinstance(data, bytes):
+            data = data.decode("utf-8")
+        csv_input = StringIO(data)
+        data = parse_csv_input(csv_input=csv_input, schema=input_schema)
+        params = None
+    elif mime_type == CONTENT_TYPE_JSON:
+        parsed_json_input = _parse_json_data(data, model.metadata, input_schema)
+        data = parsed_json_input.data
+        params = parsed_json_input.params
+        should_parse_as_unified_llm_input = parsed_json_input.is_unified_llm_input
+    else:
+        return InvocationsResponse(
+            response=(
+                "This predictor only supports the following content types:"
+                f" Types: {CONTENT_TYPES}."
+                f" Got '{content_type}'."
+            ),
+            status=415,
+            mimetype="text/plain",
+        )
+
+    # Do the prediction
+    # NB: utils._validate_serving_input mimic the scoring process here to validate input_example
+    # work for serving, so any changes here should be reflected there as well
+    try:
+        if "params" in inspect.signature(model.predict).parameters:
+            raw_predictions = model.predict(data, params=params)
+        else:
+            _log_warning_if_params_not_in_predict_signature(_logger, params)
+            raw_predictions = model.predict(data)
+    except MlflowException as e:
+        if "Failed to enforce schema" in e.message:
+            _logger.warning(
+                "If using `instances` as input key, we internally convert "
+                "the data type from `records` (List[Dict]) type to "
+                "`list` (Dict[str, List]) type if the data is a pandas "
+                "dataframe representation. This might cause schema changes. "
+                "Please use `inputs` to avoid this conversion.\n"
+            )
+        e.message = f"Failed to predict data '{data}'. \nError: {e.message}"
+        raise e
+    except Exception:
+        raise MlflowException(
+            message=(
+                "Encountered an unexpected error while evaluating the model. Verify"
+                " that the serialized input Dataframe is compatible with the model for"
+                " inference."
+            ),
+            error_code=BAD_REQUEST,
+            stack_trace=traceback.format_exc(),
+        )
+    result = StringIO()
+
+    # if the data was formatted using the unified LLM format,
+    # then return the data without the "predictions" key
+    if should_parse_as_unified_llm_input:
+        unwrapped_predictions_to_json(raw_predictions, result)
+    else:
+        predictions_to_json(raw_predictions, result)
+
+    return InvocationsResponse(response=result.getvalue(), status=200, mimetype="application/json")
+
+
+class ParsedJsonInput(NamedTuple):
+    data: Any
+    params: Optional[dict]
+    is_unified_llm_input: bool
+
+
+def _parse_json_data(data, metadata, input_schema):
+    json_input = _decode_json_input(data)
+    _is_unified_llm_input = is_unified_llm_input(json_input)
+    # no data parsing for unified LLM input format
+    if _is_unified_llm_input:
+        # Unified LLM input format
+        if hasattr(metadata, "get_params_schema"):
+            params_schema = metadata.get_params_schema()
+        else:
+            params_schema = None
+        data, params = _split_data_and_params_for_llm_input(json_input, params_schema)
+    else:
+        # Traditional json input format
+        data, params = _split_data_and_params(data)
+        # data only needs to be parsed if the model signature is not from type hint
+        # default to True for backwards compatibility
+        should_parse_data = (
+            not metadata._is_signature_from_type_hint()
+            if hasattr(metadata, "_is_signature_from_type_hint")
+            else True
+        )
+        if should_parse_data:
+            data = infer_and_parse_data(data, input_schema)
+        else:
+            if INPUTS not in data:
+                raise MlflowException.invalid_parameter_value(
+                    "Request payload must be a dictionary with 'inputs' key when "
+                    f"the model contains a valid type hint. Found keys in payload: {data.keys()}."
+                )
+            data = data[INPUTS]
+    return ParsedJsonInput(data, params, _is_unified_llm_input)
+
+
+def _async_catch_mlflow_exception(func):
+    from fastapi.responses import Response
+
+    @wraps(func)
+    async def wrapper(*args, **kwargs):
+        try:
+            return await func(*args, **kwargs)
+        except MlflowException as e:
+            return Response(
+                content=e.serialize_as_json(),
+                status_code=e.get_http_status_code(),
+                media_type="application/json",
+            )
+
+    return wrapper
+
+
+def init(model: PyFuncModel):
+    """
+    Initialize the server. Loads pyfunc model from the path.
+    """
+    from fastapi import FastAPI, Request
+    from fastapi.responses import Response
+
+    app = FastAPI()
+    input_schema = model.metadata.get_input_schema()
+    # set the environment variable to indicate that we are in a serving environment
+    os.environ[_MLFLOW_IS_IN_SERVING_ENVIRONMENT.name] = "true"
+    timeout = MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT.get()
+
+    @app.middleware("http")
+    async def timeout_middleware(request: Request, call_next):
+        try:
+            return await asyncio.wait_for(call_next(request), timeout=timeout)
+        except (asyncio.TimeoutError, TimeoutError):
+            return Response(
+                content="Request processing time exceeded limit",
+                status_code=504,
+                media_type="application/json",
+            )
+
+    @app.route("/ping", methods=["GET"])
+    @app.route("/health", methods=["GET"])
+    async def ping(request: Request):
+        """
+        Determine if the container is working and healthy.
+        We declare it healthy if we can load the model successfully.
+        """
+        health = model is not None
+        status = 200 if health else 404
+        return Response(content="\n", status_code=status, media_type="application/json")
+
+    @app.route("/version", methods=["GET"])
+    async def version(request: Request):
+        """
+        Returns the current mlflow version.
+        """
+        return Response(content=VERSION, status_code=200, media_type="application/json")
+
+    @app.route("/invocations", methods=["POST"])
+    @_async_catch_mlflow_exception
+    async def transformation(request: Request):
+        """
+        Do an inference on a single batch of data. In this sample server,
+        we take data as CSV or json, convert it to a Pandas DataFrame or Numpy,
+        generate predictions and convert them back to json.
+        """
+
+        data = await request.body()
+        content_type = request.headers.get("content-type")
+        # TODO: convert "invocations" to an async method to make internal logic fully non-blocking.
+        result = await asyncio.to_thread(invocations, data, content_type, model, input_schema)
+
+        return Response(
+            content=result.response, status_code=result.status, media_type=result.mimetype
+        )
+
+    return app
+
+
+def _predict(model_uri, input_path, output_path, content_type):
+    from mlflow.pyfunc.utils.environment import _simulate_serving_environment
+
+    with _simulate_serving_environment():
+        pyfunc_model = load_model(model_uri)
+
+        should_parse_as_unified_llm_input = False
+        if content_type == "json":
+            if input_path is None:
+                input_str = sys.stdin.read()
+            else:
+                with open(input_path) as f:
+                    input_str = f.read()
+            parsed_json_input = _parse_json_data(
+                data=input_str,
+                metadata=pyfunc_model.metadata,
+                input_schema=pyfunc_model.metadata.get_input_schema(),
+            )
+            df = parsed_json_input.data
+            params = parsed_json_input.params
+            should_parse_as_unified_llm_input = parsed_json_input.is_unified_llm_input
+        elif content_type == "csv":
+            df = (
+                parse_csv_input(input_path)
+                if input_path is not None
+                else parse_csv_input(sys.stdin)
+            )
+            params = None
+        else:
+            raise Exception(f"Unknown content type '{content_type}'")
+
+        if "params" in inspect.signature(pyfunc_model.predict).parameters:
+            raw_predictions = pyfunc_model.predict(df, params=params)
+        else:
+            _log_warning_if_params_not_in_predict_signature(_logger, params)
+            raw_predictions = pyfunc_model.predict(df)
+
+        parse_output_func = (
+            unwrapped_predictions_to_json
+            if should_parse_as_unified_llm_input
+            else predictions_to_json
+        )
+
+        if output_path is None:
+            parse_output_func(raw_predictions, sys.stdout)
+        else:
+            with open(output_path, "w") as fout:
+                parse_output_func(raw_predictions, fout)
+
+
+def _serve(model_uri, port, host):
+    pyfunc_model = load_model(model_uri)
+    init(pyfunc_model).run(port=port, host=host)
+
+
+def get_cmd(
+    model_uri: str,
+    port: Optional[int] = None,
+    host: Optional[int] = None,
+    timeout: Optional[int] = None,
+    nworkers: Optional[int] = None,
+) -> tuple[str, dict[str, str]]:
+    local_uri = path_to_local_file_uri(model_uri)
+    timeout = timeout or MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT.get()
+
+    args = []
+    if host:
+        args.append(f"--host {shlex.quote(host)}")
+
+    if port:
+        args.append(f"--port {port}")
+
+    if nworkers:
+        args.append(f"--workers {nworkers}")
+
+    command = f"uvicorn {' '.join(args)} mlflow.pyfunc.scoring_server.app:app"
+
+    command_env = os.environ.copy()
+    command_env[_SERVER_MODEL_PATH] = local_uri
+    command_env[MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT.name] = str(timeout)
+
+    return command, command_env

venv/lib/python3.9/site-packages/mlflow/pyfunc/scoring_server/app.py

@@ -0,0 +1,7 @@
+import os
+
+from mlflow.pyfunc import scoring_server
+
+app = scoring_server.init(
+    scoring_server.load_model_with_mlflow_config(os.environ[scoring_server._SERVER_MODEL_PATH])
+)

venv/lib/python3.9/site-packages/mlflow/pyfunc/scoring_server/client.py

@@ -0,0 +1,146 @@
+import json
+import logging
+import tempfile
+import time
+import uuid
+from abc import ABC, abstractmethod
+from pathlib import Path
+from typing import Any, Optional
+
+import requests
+
+from mlflow.deployments import PredictionsResponse
+from mlflow.environment_variables import MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT
+from mlflow.exceptions import MlflowException
+from mlflow.pyfunc import scoring_server
+from mlflow.utils.proto_json_utils import dump_input_data
+
+_logger = logging.getLogger(__name__)
+
+
+class BaseScoringServerClient(ABC):
+    @abstractmethod
+    def wait_server_ready(self, timeout=30, scoring_server_proc=None):
+        """
+        Wait until the scoring server is ready to accept requests.
+        """
+
+    @abstractmethod
+    def invoke(self, data, params: Optional[dict[str, Any]] = None):
+        """
+        Invoke inference on input data. The input data must be pandas dataframe or numpy array or
+        a dict of numpy arrays.
+
+        Args:
+            data: Model input data.
+            params: Additional parameters to pass to the model for inference.
+
+        Returns:
+            Prediction result.
+        """
+
+
+class ScoringServerClient(BaseScoringServerClient):
+    def __init__(self, host, port):
+        self.url_prefix = f"http://{host}:{port}"
+
+    def ping(self):
+        ping_status = requests.get(url=self.url_prefix + "/ping")
+        if ping_status.status_code != 200:
+            raise Exception(f"ping failed (error code {ping_status.status_code})")
+
+    def get_version(self):
+        resp_status = requests.get(url=self.url_prefix + "/version")
+        if resp_status.status_code != 200:
+            raise Exception(f"version failed (error code {resp_status.status_code})")
+        return resp_status.text
+
+    def wait_server_ready(self, timeout=30, scoring_server_proc=None):
+        begin_time = time.time()
+
+        while True:
+            time.sleep(0.3)
+            try:
+                self.ping()
+                return
+            except Exception:
+                pass
+            if time.time() - begin_time > timeout:
+                break
+            if scoring_server_proc is not None:
+                return_code = scoring_server_proc.poll()
+                if return_code is not None:
+                    raise RuntimeError(f"Server process already exit with returncode {return_code}")
+        raise RuntimeError("Wait scoring server ready timeout.")
+
+    def invoke(self, data, params: Optional[dict[str, Any]] = None):
+        """
+        Args:
+            data: Model input data.
+            params: Additional parameters to pass to the model for inference.
+
+        Returns:
+            :py:class:`PredictionsResponse <mlflow.deployments.PredictionsResponse>` result.
+        """
+        response = requests.post(
+            url=self.url_prefix + "/invocations",
+            data=dump_input_data(data, params=params),
+            headers={"Content-Type": scoring_server.CONTENT_TYPE_JSON},
+        )
+        if response.status_code != 200:
+            raise Exception(
+                f"Invocation failed (error code {response.status_code}, response: {response.text})"
+            )
+        return PredictionsResponse.from_json(response.text)
+
+
+class StdinScoringServerClient(BaseScoringServerClient):
+    def __init__(self, process):
+        super().__init__()
+        self.process = process
+        self.tmpdir = Path(tempfile.mkdtemp())
+        self.output_json = self.tmpdir.joinpath("output.json")
+
+    def wait_server_ready(self, timeout=30, scoring_server_proc=None):
+        return_code = self.process.poll()
+        if return_code is not None:
+            raise RuntimeError(f"Server process already exit with returncode {return_code}")
+
+    def invoke(self, data, params: Optional[dict[str, Any]] = None):
+        """
+        Invoke inference on input data. The input data must be pandas dataframe or numpy array or
+        a dict of numpy arrays.
+
+        Args:
+            data: Model input data.
+            params: Additional parameters to pass to the model for inference.
+
+        Returns:
+            :py:class:`PredictionsResponse <mlflow.deployments.PredictionsResponse>` result.
+        """
+        if not self.output_json.exists():
+            self.output_json.touch()
+
+        request_id = str(uuid.uuid4())
+        request = {
+            "id": request_id,
+            "data": dump_input_data(data, params=params),
+            "output_file": str(self.output_json),
+        }
+        self.process.stdin.write(json.dumps(request) + "\n")
+        self.process.stdin.flush()
+
+        begin_time = time.time()
+        while True:
+            _logger.info("Waiting for scoring to complete...")
+            try:
+                with self.output_json.open(mode="r+") as f:
+                    resp = PredictionsResponse.from_json(f.read())
+                    if resp.get("id") == request_id:
+                        f.truncate(0)
+                        return resp
+            except Exception as e:
+                _logger.debug("Exception while waiting for scoring to complete: %s", e)
+            if time.time() - begin_time > MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT.get():
+                raise MlflowException("Scoring timeout")
+            time.sleep(1)

venv/lib/python3.9/site-packages/mlflow/pyfunc/spark_model_cache.py

@@ -0,0 +1,48 @@
+from mlflow.utils._spark_utils import _SparkDirectoryDistributor
+
+
+class SparkModelCache:
+    """Caches models in memory on Spark Executors, to avoid continually reloading from disk.
+
+    This class has to be part of a different module than the one that _uses_ it. This is
+    because Spark will pickle classes that are defined in the local scope, but relies on
+    Python's module loading behavior for classes in different modules. In this case, we
+    are relying on the fact that Python will load a module at-most-once, and can therefore
+    store per-process state in a static map.
+    """
+
+    # Map from unique name --> (loaded model, local_model_path).
+    _models = {}
+
+    # Number of cache hits we've had, for testing purposes.
+    _cache_hits = 0
+
+    def __init__(self):
+        pass
+
+    @staticmethod
+    def add_local_model(spark, model_path):
+        """Given a SparkSession and a model_path which refers to a pyfunc directory locally,
+        we will zip the directory up, enable it to be distributed to executors, and return
+        the "archive_path", which should be used as the path in get_or_load().
+        """
+        return _SparkDirectoryDistributor.add_dir(spark, model_path)
+
+    @staticmethod
+    def get_or_load(archive_path):
+        """Given a path returned by add_local_model(), this method will return a tuple of
+        (loaded_model, local_model_path).
+        If this Python process ever loaded the model before, we will reuse that copy.
+        """
+        if archive_path in SparkModelCache._models:
+            SparkModelCache._cache_hits += 1
+            return SparkModelCache._models[archive_path]
+
+        local_model_dir = _SparkDirectoryDistributor.get_or_extract(archive_path)
+
+        # We must rely on a supposed cyclic import here because we want this behavior
+        # on the Spark Executors (i.e., don't try to pickle the load_model function).
+        from mlflow.pyfunc import load_model
+
+        SparkModelCache._models[archive_path] = (load_model(local_model_dir), local_model_dir)
+        return SparkModelCache._models[archive_path]

venv/lib/python3.9/site-packages/mlflow/pyfunc/stdin_server.py

@@ -0,0 +1,44 @@
+import argparse
+import inspect
+import json
+import logging
+import sys
+
+from mlflow.pyfunc import scoring_server
+from mlflow.pyfunc.model import _log_warning_if_params_not_in_predict_signature
+
+_logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--model-uri")
+args = parser.parse_args()
+
+_logger.info("Loading model from %s", args.model_uri)
+
+model = scoring_server.load_model_with_mlflow_config(args.model_uri)
+input_schema = model.metadata.get_input_schema()
+_logger.info("Loaded model")
+
+_logger.info("Waiting for request")
+for line in sys.stdin:
+    _logger.info("Received request")
+    request = json.loads(line)
+
+    _logger.info("Parsing input data")
+    data = request["data"]
+    data, params = scoring_server._split_data_and_params(data)
+    data = scoring_server.infer_and_parse_data(data, input_schema)
+
+    _logger.info("Making predictions")
+    if "params" in inspect.signature(model.predict).parameters:
+        preds = model.predict(data, params=params)
+    else:
+        _log_warning_if_params_not_in_predict_signature(_logger, params)
+        preds = model.predict(data)
+
+    _logger.info("Writing predictions")
+    with open(request["output_file"], "a") as f:
+        scoring_server.predictions_to_json(preds, f, {"id": request["id"]})
+
+    _logger.info("Done")

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

@@ -0,0 +1,3 @@
+from mlflow.pyfunc.utils.data_validation import pyfunc
+
+__all__ = ["pyfunc"]

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

@@ -0,0 +1,224 @@
+import inspect
+import warnings
+from functools import lru_cache, wraps
+from typing import Any, NamedTuple, Optional
+
+import pydantic
+
+from mlflow.exceptions import MlflowException
+from mlflow.models.signature import (
+    _extract_type_hints,
+    _is_context_in_predict_function_signature,
+)
+from mlflow.types.type_hints import (
+    InvalidTypeHintException,
+    _convert_data_to_type_hint,
+    _infer_schema_from_list_type_hint,
+    _is_type_hint_from_example,
+    _signature_cannot_be_inferred_from_type_hint,
+    _validate_data_against_type_hint,
+    model_validate,
+)
+from mlflow.utils.annotations import filter_user_warnings_once
+from mlflow.utils.warnings_utils import color_warning
+
+_INVALID_SIGNATURE_ERROR_MSG = (
+    "Model's `{func_name}` method contains invalid parameters: {invalid_params}. "
+    "Only the following parameter names are allowed: context, model_input, and params. "
+    "Note that invalid parameters will no longer be permitted in future versions."
+)
+
+
+class FuncInfo(NamedTuple):
+    input_type_hint: Optional[type[Any]]
+    input_param_name: str
+
+
+def pyfunc(func):
+    """
+    A decorator that forces data validation against type hint of the input data
+    in the wrapped method. It is no-op if the type hint is not supported by MLflow.
+
+    .. note::
+        The function that applies this decorator must be a valid `predict` function
+        of `mlflow.pyfunc.PythonModel`, or a callable that takes a single input.
+    """
+
+    func_info = _get_func_info_if_type_hint_supported(func)
+    return _wrap_predict_with_pyfunc(func, func_info)
+
+
+def _wrap_predict_with_pyfunc(func, func_info: Optional[FuncInfo]):
+    if func_info is not None:
+        model_input_index = _model_input_index_in_function_signature(func)
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            try:
+                args, kwargs = _validate_model_input(
+                    args,
+                    kwargs,
+                    model_input_index,
+                    func_info.input_type_hint,
+                    func_info.input_param_name,
+                )
+            except Exception as e:
+                if isinstance(e, MlflowException):
+                    raise e
+                raise MlflowException(
+                    "Failed to validate the input data against the type hint "
+                    f"`{func_info.input_type_hint}`. Error: {e}"
+                )
+            return func(*args, **kwargs)
+    else:
+
+        @wraps(func)
+        def wrapper(*args, **kwargs):
+            return func(*args, **kwargs)
+
+    wrapper._is_pyfunc = True
+    return wrapper
+
+
+def wrap_non_list_predict_pydantic(func, input_pydantic_model, validation_error_msg, unpack=False):
+    """
+    Used by MLflow defined subclasses of PythonModel that have non-list a pydantic model as input.
+    Takes in a dict input, validates it against `input_pydantic_model`, and then creates
+    the pydantic model.
+
+    If `unpack` is True, the validated dict is parsed into the function arguments.
+    Otherwise, the whole pydantic object is passed to the function.
+
+    Args:
+        func: The predict/predict_stream method of the PythonModel subclass.
+        input_pydantic_model: The pydantic model that the input should be validated against.
+        validation_error_msg: The error message to raise if the dict input fails to validate.
+        unpack: Whether to unpack the validated dict into the function arguments. Defaults to False.
+
+    Raises:
+        MlflowException: If the input fails to validate against the pydantic model.
+
+    Returns:
+        A function that can take either a dict input or a pydantic object as input.
+    """
+
+    @wraps(func)
+    def wrapper(self, *args, **kwargs):
+        if len(args) == 1 and isinstance(args[0], dict):
+            try:
+                model_validate(input_pydantic_model, args[0])
+                pydantic_obj = input_pydantic_model(**args[0])
+            except pydantic.ValidationError as e:
+                raise MlflowException(
+                    f"{validation_error_msg} Pydantic validation error: {e}"
+                ) from e
+            else:
+                if unpack:
+                    param_names = inspect.signature(func).parameters.keys() - {"self"}
+                    kwargs = {k: getattr(pydantic_obj, k) for k in param_names}
+                    return func(self, **kwargs)
+                else:
+                    return func(self, pydantic_obj)
+        else:
+            # Before logging, this is equivalent to the behavior from the raw predict method
+            # After logging, signature enforcement happens in the _convert_input method
+            # of the wrapper class
+            return func(self, *args, **kwargs)
+
+    wrapper._is_pyfunc = True
+    return wrapper
+
+
+def _check_func_signature(func, func_name) -> list[str]:
+    parameters = inspect.signature(func).parameters
+    param_names = [name for name in parameters.keys() if name != "self"]
+    if invalid_params := set(param_names) - {"self", "context", "model_input", "params"}:
+        warnings.warn(
+            _INVALID_SIGNATURE_ERROR_MSG.format(func_name=func_name, invalid_params=invalid_params),
+            FutureWarning,
+            stacklevel=2,
+        )
+    return param_names
+
+
+@lru_cache
+@filter_user_warnings_once
+def _get_func_info_if_type_hint_supported(func) -> Optional[FuncInfo]:
+    """
+    Internal method to check if the predict function has type hints and if they are supported
+    by MLflow.
+    For PythonModel, the signature must be one of below:
+        - predict(self, context, model_input, params=None)
+        - predict(self, model_input, params=None)
+    For callables, the function must contain only one input argument.
+    """
+    param_names = _check_func_signature(func, "predict")
+    input_arg_index = 1 if _is_context_in_predict_function_signature(func=func) else 0
+    type_hint = _extract_type_hints(func, input_arg_index=input_arg_index).input
+    input_param_name = param_names[input_arg_index]
+    if type_hint is not None:
+        if _signature_cannot_be_inferred_from_type_hint(type_hint) or _is_type_hint_from_example(
+            type_hint
+        ):
+            return
+        try:
+            _infer_schema_from_list_type_hint(type_hint)
+        except InvalidTypeHintException as e:
+            raise MlflowException(
+                f"{e.message} To disable data validation, remove the type hint from the "
+                "predict function. Otherwise, fix the type hint."
+            )
+        # catch other exceptions to avoid breaking model usage
+        except Exception as e:
+            color_warning(
+                message="Type hint used in the model's predict function is not supported "
+                f"for MLflow's schema validation. {e} "
+                "Remove the type hint to disable this warning. "
+                "To enable validation for the input data, specify input example "
+                "or model signature when logging the model. ",
+                category=UserWarning,
+                stacklevel=3,
+                color="red",
+            )
+        else:
+            return FuncInfo(input_type_hint=type_hint, input_param_name=input_param_name)
+    else:
+        color_warning(
+            "Add type hints to the `predict` method to enable data validation "
+            "and automatic signature inference during model logging. "
+            "Check https://mlflow.org/docs/latest/model/python_model.html#type-hint-usage-in-pythonmodel"
+            " for more details.",
+            stacklevel=1,
+            color="yellow",
+            category=UserWarning,
+        )
+
+
+def _model_input_index_in_function_signature(func):
+    parameters = inspect.signature(func).parameters
+    # we need to exclude the first argument if "self" is in the parameters
+    index = 1 if "self" in parameters else 0
+    if _is_context_in_predict_function_signature(parameters=parameters):
+        index += 1
+    return index
+
+
+def _validate_model_input(
+    args, kwargs, model_input_index_in_sig, type_hint, model_input_param_name
+):
+    model_input = None
+    input_pos = None
+    if model_input_param_name in kwargs:
+        model_input = kwargs[model_input_param_name]
+        input_pos = "kwargs"
+    elif len(args) >= model_input_index_in_sig + 1:
+        model_input = args[model_input_index_in_sig]
+        input_pos = model_input_index_in_sig
+    if input_pos is not None:
+        data = _convert_data_to_type_hint(model_input, type_hint)
+        data = _validate_data_against_type_hint(data, type_hint)
+        if input_pos == "kwargs":
+            kwargs[model_input_param_name] = data
+        else:
+            args = args[:input_pos] + (data,) + args[input_pos + 1 :]
+    return args, kwargs

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

@@ -0,0 +1,22 @@
+import os
+from contextlib import contextmanager
+
+from mlflow.environment_variables import _MLFLOW_IS_IN_SERVING_ENVIRONMENT
+
+
+@contextmanager
+def _simulate_serving_environment():
+    """
+    Some functions (e.g. validate_serving_input) replicate the data transformation logic
+    that happens in the model serving environment to validate data before model deployment.
+    This context manager can be used to simulate the serving environment for such functions.
+    """
+    original_value = _MLFLOW_IS_IN_SERVING_ENVIRONMENT.get_raw()
+    try:
+        _MLFLOW_IS_IN_SERVING_ENVIRONMENT.set("true")
+        yield
+    finally:
+        if original_value is not None:
+            os.environ[_MLFLOW_IS_IN_SERVING_ENVIRONMENT.name] = original_value
+        else:
+            del os.environ[_MLFLOW_IS_IN_SERVING_ENVIRONMENT.name]

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

@@ -0,0 +1,47 @@
+from dataclasses import fields, is_dataclass
+from typing import Union, get_args, get_origin
+
+from mlflow.utils.annotations import experimental
+
+
+def _is_optional_dataclass(field_type) -> bool:
+    """
+    Check if the field type is an Optional containing a dataclass.
+    Currently, ... | None (in Python 3.10) is not supported.
+    """
+    if get_origin(field_type) is Union:
+        inner_types = get_args(field_type)
+        # Check if it's a Union[Dataclass, NoneType] (i.e., Optional[Dataclass])
+        if len(inner_types) == 2 and any(t is type(None) for t in inner_types):
+            effective_type = next(t for t in get_args(field_type) if t is not type(None))
+            return is_dataclass(effective_type)
+    return False
+
+
+@experimental
+def _hydrate_dataclass(dataclass_type, data):
+    """Recursively create an instance of the dataclass_type from data."""
+    if not (is_dataclass(dataclass_type) or _is_optional_dataclass(dataclass_type)):
+        raise ValueError(f"{dataclass_type.__name__} is not a dataclass")
+    if data is None:
+        return None
+
+    field_names = {f.name: f.type for f in fields(dataclass_type)}
+    kwargs = {}
+    for key, field_type in field_names.items():
+        if key in data:
+            value = data[key]
+            if is_dataclass(field_type):
+                kwargs[key] = _hydrate_dataclass(field_type, value)
+            elif _is_optional_dataclass(field_type):
+                effective_type = next(t for t in get_args(field_type) if t is not type(None))
+                kwargs[key] = _hydrate_dataclass(effective_type, value)
+            elif get_origin(field_type) == list:
+                item_type = get_args(field_type)[0]
+                if is_dataclass(item_type):
+                    kwargs[key] = [_hydrate_dataclass(item_type, item) for item in value]
+                else:
+                    kwargs[key] = value
+            else:
+                kwargs[key] = value
+    return dataclass_type(**kwargs)

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

@@ -0,0 +1,9 @@
+# Support unwrapped JSON with these keys for LLM use cases of Chat, Completions, Embeddings tasks
+LLM_CHAT_KEY = "messages"
+LLM_COMPLETIONS_KEY = "prompt"
+LLM_EMBEDDINGS_KEY = "input"
+SUPPORTED_LLM_FORMATS = {LLM_CHAT_KEY, LLM_COMPLETIONS_KEY, LLM_EMBEDDINGS_KEY}
+
+
+def is_unified_llm_input(json_input: dict):
+    return any(x in json_input for x in SUPPORTED_LLM_FORMATS)