mlflow.pyfunc

The python_function model flavor serves as a default model interface for MLflow Python models. Any MLflow Python model is expected to be loadable as a python_function model.

In addition, the mlflow.pyfunc module defines a generic filesystem format for Python models and provides utilities for saving to and loading from this format. The format is self contained in the sense that it includes all necessary information for anyone to load it and use it. Dependencies are either stored directly with the model or referenced via a Conda environment.

The mlflow.pyfunc module also defines utilities for creating custom pyfunc models using frameworks and inference logic that may not be natively included in MLflow. See Creating custom Pyfunc models.

Inference API

Python function models are loaded as an instance of PyFuncModel, which is an MLflow wrapper around the model implementation and model metadata (MLmodel file). You can score the model by calling the predict() method, which has the following signature:

predict(
  model_input: [pandas.DataFrame, numpy.ndarray, scipy.sparse.(csc.csc_matrix | csr.csr_matrix),
  List[Any], Dict[str, Any], pyspark.sql.DataFrame]
) -> [numpy.ndarray | pandas.(Series | DataFrame) | List | pyspark.sql.DataFrame]

All PyFunc models will support pandas.DataFrame as input and PyFunc deep learning models will also support tensor inputs in the form of Dict[str, numpy.ndarray] (named tensors) and numpy.ndarrays (unnamed tensors).

Filesystem format

The Pyfunc format is defined as a directory structure containing all required data, code, and configuration:

./dst-path/
    ./MLmodel: configuration
    <code>: code packaged with the model (specified in the MLmodel file)
    <data>: data packaged with the model (specified in the MLmodel file)
    <env>: Conda environment definition (specified in the MLmodel file)

The directory structure may contain additional contents that can be referenced by the MLmodel configuration.

MLModel configuration

A Python model contains an MLmodel file in python_function format in its root with the following parameters:

  • loader_module [required]:

    Python module that can load the model. Expected as module identifier e.g. mlflow.sklearn, it will be imported using importlib.import_module. The imported module must contain a function with the following signature:

    _load_pyfunc(path: string) -> <pyfunc model implementation>
    

    The path argument is specified by the data parameter and may refer to a file or directory. The model implementation is expected to be an object with a predict method with the following signature:

    predict(
      model_input: [pandas.DataFrame, numpy.ndarray,
      scipy.sparse.(csc.csc_matrix | csr.csr_matrix), List[Any], Dict[str, Any]],
      pyspark.sql.DataFrame
    ) -> [numpy.ndarray | pandas.(Series | DataFrame) | List | pyspark.sql.DataFrame]
    
  • code [optional]:

    Relative path to a directory containing the code packaged with this model. All files and directories inside this directory are added to the Python path prior to importing the model loader.

  • data [optional]:

    Relative path to a file or directory containing model data. The path is passed to the model loader.

  • env [optional]:

    Relative path to an exported Conda environment. If present this environment should be activated prior to running the model.

  • Optionally, any additional parameters necessary for interpreting the serialized model in pyfunc format.

Example

tree example/sklearn_iris/mlruns/run1/outputs/linear-lr
├── MLmodel
├── code
│   ├── sklearn_iris.py
│
├── data
│   └── model.pkl
└── mlflow_env.yml
cat example/sklearn_iris/mlruns/run1/outputs/linear-lr/MLmodel
python_function:
  code: code
  data: data/model.pkl
  loader_module: mlflow.sklearn
  env: mlflow_env.yml
  main: sklearn_iris

Creating custom Pyfunc models

MLflow’s persistence modules provide convenience functions for creating models with the pyfunc flavor in a variety of machine learning frameworks (scikit-learn, Keras, Pytorch, and more); however, they do not cover every use case. For example, you may want to create an MLflow model with the pyfunc flavor using a framework that MLflow does not natively support. Alternatively, you may want to build an MLflow model that executes custom logic when evaluating queries, such as preprocessing and postprocessing routines. Therefore, mlflow.pyfunc provides utilities for creating pyfunc models from arbitrary code and model data.

The save_model() and log_model() methods are designed to support multiple workflows for creating custom pyfunc models that incorporate custom inference logic and artifacts that the logic may require.

An artifact is a file or directory, such as a serialized model or a CSV. For example, a serialized TensorFlow graph is an artifact. An MLflow model directory is also an artifact.

Workflows

save_model() and log_model() support the following workflows:

  1. Programmatically defining a new MLflow model, including its attributes and artifacts.

    Given a set of artifact URIs, save_model() and log_model() can automatically download artifacts from their URIs and create an MLflow model directory.

    In this case, you must define a Python class which inherits from PythonModel, defining predict() and, optionally, load_context(). An instance of this class is specified via the python_model parameter; it is automatically serialized and deserialized as a Python class, including all of its attributes.

  2. Interpreting pre-existing data as an MLflow model.

    If you already have a directory containing model data, save_model() and log_model() can import the data as an MLflow model. The data_path parameter specifies the local filesystem path to the directory containing model data.

    In this case, you must provide a Python module, called a loader module. The loader module defines a _load_pyfunc() method that performs the following tasks:

    • Load data from the specified data_path. For example, this process may include deserializing pickled Python objects or models or parsing CSV files.

    • Construct and return a pyfunc-compatible model wrapper. As in the first use case, this wrapper must define a predict() method that is used to evaluate queries. predict() must adhere to the Inference API.

    The loader_module parameter specifies the name of your loader module.

    For an example loader module implementation, refer to the loader module implementation in mlflow.sklearn.

Which workflow is right for my use case?

We consider the first workflow to be more user-friendly and generally recommend it for the following reasons:

  • It automatically resolves and collects specified model artifacts.

  • It automatically serializes and deserializes the python_model instance and all of its attributes, reducing the amount of user logic that is required to load the model

  • You can create Models using logic that is defined in the __main__ scope. This allows custom models to be constructed in interactive environments, such as notebooks and the Python REPL.

You may prefer the second, lower-level workflow for the following reasons:

  • Inference logic is always persisted as code, rather than a Python object. This makes logic easier to inspect and modify later.

  • If you have already collected all of your model data in a single location, the second workflow allows it to be saved in MLflow format directly, without enumerating constituent artifacts.

class mlflow.pyfunc.EnvType[source]

Bases: object

CONDA = 'conda'
VIRTUALENV = 'virtualenv'
class mlflow.pyfunc.PyFuncModel(model_meta: mlflow.models.model.Model, model_impl: Any, predict_fn: str = 'predict')[source]

Bases: object

MLflow ‘python function’ model.

Wrapper around model implementation and metadata. This class is not meant to be constructed directly. Instead, instances of this class are constructed and returned from load_model().

model_impl can be any Python object that implements the Pyfunc interface, and is returned by invoking the model’s loader_module.

model_meta contains model metadata loaded from the MLmodel file.

property loader_module

Model’s flavor configuration

Note

Experimental: This property may change or be removed in a future release without warning.

property metadata

Model metadata.

property model_config

Model’s flavor configuration

Note

Experimental: This property may change or be removed in a future release without warning.

predict(data: Union[pandas.core.frame.DataFrame, pandas.core.series.Series, numpy.ndarray, csc_matrix, csr_matrix, List[Any], Dict[str, Any], datetime.datetime, bool, bytes, float, int, str, pyspark.sql.dataframe.DataFrame], params: Optional[Dict[str, Any]] = None)Union[pandas.core.frame.DataFrame, pandas.core.series.Series, numpy.ndarray, list, str, pyspark.sql.dataframe.DataFrame][source]

Generates model predictions.

If the model contains signature, enforce the input schema first before calling the model implementation with the sanitized input. If the pyfunc model does not include model schema, the input is passed to the model implementation as is. See Model Signature Enforcement for more details.

Parameters
  • data – Model input as one of pandas.DataFrame, numpy.ndarray, scipy.sparse.(csc.csc_matrix | csr.csr_matrix), List[Any], or Dict[str, numpy.ndarray]. For model signatures with tensor spec inputs (e.g. the Tensorflow core / Keras model), the input data type must be one of numpy.ndarray, List[numpy.ndarray], Dict[str, numpy.ndarray] or pandas.DataFrame. If data is of pandas.DataFrame type and the model contains a signature with tensor spec inputs, the corresponding column values in the pandas DataFrame will be reshaped to the required shape with ‘C’ order (i.e. read / write the elements using C-like index order), and DataFrame column values will be cast as the required tensor spec type. For Pyspark DataFrame inputs, MLflow will only enforce the schema on a subset of the data rows.

  • params

    Additional parameters to pass to the model for inference.

    Note

    Experimental: This parameter may change or be removed in a future release without warning.

Returns

Model predictions as one of pandas.DataFrame, pandas.Series, numpy.ndarray or list.

unwrap_python_model()[source]

Note

Experimental: This function may change or be removed in a future release without warning.

Unwrap the underlying Python model object.

This method is useful for accessing custom model functions, while still being able to leverage the MLflow designed workflow through the predict() method.

Returns

The underlying wrapped model object

Example
import mlflow


# define a custom model
class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input, params=None):
        return self.my_custom_function(model_input, params)

    def my_custom_function(self, model_input, params=None):
        # do something with the model input
        return 0


some_input = 1
# save the model
with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(artifact_path="model", python_model=MyModel())

# load the model
loaded_model = mlflow.pyfunc.load_model(model_uri=model_info.model_uri)
print(type(loaded_model))  # <class 'mlflow.pyfunc.model.PyFuncModel'>
unwrapped_model = loaded_model.unwrap_python_model()
print(type(unwrapped_model))  # <class '__main__.MyModel'>

# does not work, only predict() is exposed
# print(loaded_model.my_custom_function(some_input))
print(unwrapped_model.my_custom_function(some_input))  # works
print(loaded_model.predict(some_input))  # works

# works, but None is needed for context arg
print(unwrapped_model.predict(None, some_input))
mlflow.pyfunc.add_to_model(model, loader_module, data=None, code=None, conda_env=None, python_env=None, model_config=None, **kwargs)[source]

Add a pyfunc spec to the model configuration.

Defines pyfunc configuration schema. Caller can use this to create a valid pyfunc model flavor out of an existing directory structure. For example, other model flavors can use this to specify how to use their output as a pyfunc.

Note

All paths are relative to the exported model root directory.

Parameters
  • model – Existing model.

  • loader_module – The module to be used to load the model.

  • data – Path to the model data.

  • code – Path to the code dependencies.

  • conda_env – Conda environment.

  • python_env – Python environment.

  • req – pip requirements file.

  • kwargs – Additional key-value pairs to include in the pyfunc flavor specification. Values must be YAML-serializable.

  • model_config

    The model configuration to apply to the model. This configuration is available during model loading.

    Note

    Experimental: This parameter may change or be removed in a future release without warning.

Returns

Updated model configuration.

mlflow.pyfunc.get_model_dependencies(model_uri, format='pip')[source]

Downloads the model dependencies and returns the path to requirements.txt or conda.yaml file.

Warning

This API downloads all the model artifacts to the local filesystem. This may take a long time for large models. To avoid this overhead, use mlflow.artifacts.download_artifacts("<model_uri>/requirements.txt") or mlflow.artifacts.download_artifacts("<model_uri>/conda.yaml") instead.

Parameters
  • model_uri – The uri of the model to get dependencies from.

  • format – The format of the returned dependency file. If the "pip" format is specified, the path to a pip requirements.txt file is returned. If the "conda" format is specified, the path to a "conda.yaml" file is returned . If the "pip" format is specified but the model was not saved with a requirements.txt file, the pip section of the model’s conda.yaml file is extracted instead, and any additional conda dependencies are ignored. Default value is "pip".

Returns

The local filesystem path to either a pip requirements.txt file (if format="pip") or a conda.yaml file (if format="conda") specifying the model’s dependencies.

mlflow.pyfunc.load_model(model_uri: str, suppress_warnings: bool = False, dst_path: Optional[str] = None, model_config: Optional[Dict[str, Any]] = None)mlflow.pyfunc.PyFuncModel[source]

Load a model stored in Python function format.

Parameters
  • model_uri

    The location, in URI format, of the MLflow model. For example:

    • /Users/me/path/to/local/model

    • relative/path/to/local/model

    • s3://my_bucket/path/to/model

    • runs:/<mlflow_run_id>/run-relative/path/to/model

    • models:/<model_name>/<model_version>

    • models:/<model_name>/<stage>

    • mlflow-artifacts:/path/to/model

    For more information about supported URI schemes, see Referencing Artifacts.

  • suppress_warnings – If True, non-fatal warning messages associated with the model loading process will be suppressed. If False, these warning messages will be emitted.

  • dst_path – The local filesystem path to which to download the model artifact. This directory must already exist. If unspecified, a local output path will be created.

  • model_config

    The model configuration to apply to the model. This configuration is available during model loading.

    Note

    Experimental: This parameter may change or be removed in a future release without warning.

mlflow.pyfunc.load_pyfunc(model_uri, suppress_warnings=False)[source]

Warning

mlflow.pyfunc.load_pyfunc is deprecated since 1.0. This method will be removed in a future release. Use mlflow.pyfunc.load_model instead.

Load a model stored in Python function format.

Parameters
  • model_uri

    The location, in URI format, of the MLflow model. For example:

    • /Users/me/path/to/local/model

    • relative/path/to/local/model

    • s3://my_bucket/path/to/model

    • runs:/<mlflow_run_id>/run-relative/path/to/model

    • models:/<model_name>/<model_version>

    • models:/<model_name>/<stage>

    • mlflow-artifacts:/path/to/model

    For more information about supported URI schemes, see Referencing Artifacts.

  • suppress_warnings – If True, non-fatal warning messages associated with the model loading process will be suppressed. If False, these warning messages will be emitted.

mlflow.pyfunc.log_model(artifact_path, loader_module=None, data_path=None, code_path=None, conda_env=None, python_model=None, artifacts=None, registered_model_name=None, signature: mlflow.models.signature.ModelSignature = None, input_example: Union[pandas.core.frame.DataFrame, numpy.ndarray, dict, list, csr_matrix, csc_matrix, str, bytes, tuple] = None, await_registration_for=300, pip_requirements=None, extra_pip_requirements=None, metadata=None, model_config=None, example_no_conversion=False)[source]

Log a Pyfunc model with custom inference logic and optional data dependencies as an MLflow artifact for the current run.

For information about the workflows that this method supports, see Workflows for creating custom pyfunc models and Which workflow is right for my use case?. You cannot specify the parameters for the second workflow: loader_module, data_path and the parameters for the first workflow: python_model, artifacts together.

Parameters
  • artifact_path – The run-relative artifact path to which to log the Python model.

  • loader_module

    The name of the Python module that is used to load the model from data_path. This module must define a method with the prototype _load_pyfunc(data_path). If not None, this module and its dependencies must be included in one of the following locations:

    • The MLflow library.

    • Package(s) listed in the model’s Conda environment, specified by the conda_env parameter.

    • One or more of the files specified by the code_path parameter.

  • data_path – Path to a file or directory containing model data.

  • code_path – A list of local filesystem paths to Python file dependencies (or directories containing file dependencies). These files are prepended to the system path before the model is loaded.

  • conda_env

    Either a dictionary representation of a Conda environment or the path to a conda environment yaml file. If provided, this describes the environment this model should be run in. At a minimum, it should specify the dependencies contained in get_default_conda_env(). If None, a conda environment with pip requirements inferred by mlflow.models.infer_pip_requirements() is added to the model. If the requirement inference fails, it falls back to using get_default_pip_requirements(). pip requirements from conda_env are written to a pip requirements.txt file and the full conda environment is written to conda.yaml. The following is an example dictionary representation of a conda environment:

    {
        "name": "mlflow-env",
        "channels": ["conda-forge"],
        "dependencies": [
            "python=3.8.15",
            {
                "pip": [
                    "scikit-learn==x.y.z"
                ],
            },
        ],
    }
    

  • python_model

    An instance of a subclass of PythonModel or a callable object with a single argument (see the examples below). The passed-in object is serialized using the CloudPickle library. Any dependencies of the class should be included in one of the following locations:

    • The MLflow library.

    • Package(s) listed in the model’s Conda environment, specified by the conda_env parameter.

    • One or more of the files specified by the code_path parameter.

    Note: If the class is imported from another module, as opposed to being defined in the __main__ scope, the defining module should also be included in one of the listed locations.

    Examples

    Class model

    from typing import List
    import mlflow
    
    
    class MyModel(mlflow.pyfunc.PythonModel):
        def predict(self, context, model_input: List[str], params=None) -> List[str]:
            return [i.upper() for i in model_input]
    
    
    with mlflow.start_run():
        model_info = mlflow.pyfunc.log_model(
            artifact_path="model",
            python_model=MyModel(),
        )
    
    
    loaded_model = mlflow.pyfunc.load_model(model_uri=model_info.model_uri)
    print(loaded_model.predict(["a", "b", "c"]))  # -> ["A", "B", "C"]
    

    Functional model

    Note

    Experimental: Functional model support is experimental and may change or be removed in a future release without warning.

    from typing import List
    import mlflow
    
    
    def predict(model_input: List[str]) -> List[str]:
        return [i.upper() for i in model_input]
    
    
    with mlflow.start_run():
        model_info = mlflow.pyfunc.log_model(
            artifact_path="model", python_model=predict, input_example=["a"]
        )
    
    
    loaded_model = mlflow.pyfunc.load_model(model_uri=model_info.model_uri)
    print(loaded_model.predict(["a", "b", "c"]))  # -> ["A", "B", "C"]
    

    If the predict method or function has type annotations, MLflow automatically constructs a model signature based on the type annotations (unless the signature argument is explicitly specified), and converts the input value to the specified type before passing it to the function. Currently, the following type annotations are supported:

    • List[str]

    • List[Dict[str, str]]

  • artifacts

    A dictionary containing <name, artifact_uri> entries. Remote artifact URIs are resolved to absolute filesystem paths, producing a dictionary of <name, absolute_path> entries. python_model can reference these resolved entries as the artifacts property of the context parameter in PythonModel.load_context() and PythonModel.predict(). For example, consider the following artifacts dictionary:

    {"my_file": "s3://my-bucket/path/to/my/file"}
    

    In this case, the "my_file" artifact is downloaded from S3. The python_model can then refer to "my_file" as an absolute filesystem path via context.artifacts["my_file"].

    If None, no artifacts are added to the model.

  • registered_model_name – This argument may change or be removed in a future release without warning. If given, create a model version under registered_model_name, also creating a registered model if one with the given name does not exist.

  • signature

    ModelSignature describes model input and output Schema. The model signature can be inferred from datasets with valid model input (e.g. the training dataset with target column omitted) and valid model output (e.g. model predictions generated on the training dataset), for example:

    from mlflow.models import infer_signature
    
    train = df.drop_column("target_label")
    predictions = ...  # compute model predictions
    signature = infer_signature(train, predictions)
    

  • input_example – one or several instances of valid model input. The input example is used as a hint of what data to feed the model. It will be converted to a Pandas DataFrame and then serialized to json using the Pandas split-oriented format, or a numpy array where the example will be serialized to json by converting it to a list. Bytes are base64-encoded. When the signature parameter is None, the input example is used to infer a model signature.

  • await_registration_for – Number of seconds to wait for the model version to finish being created and is in READY status. By default, the function waits for five minutes. Specify 0 or None to skip waiting.

  • pip_requirements – Either an iterable of pip requirement strings (e.g. ["scikit-learn", "-r requirements.txt", "-c constraints.txt"]) or the string path to a pip requirements file on the local filesystem (e.g. "requirements.txt"). If provided, this describes the environment this model should be run in. If None, a default list of requirements is inferred by mlflow.models.infer_pip_requirements() from the current software environment. If the requirement inference fails, it falls back to using get_default_pip_requirements(). Both requirements and constraints are automatically parsed and written to requirements.txt and constraints.txt files, respectively, and stored as part of the model. Requirements are also written to the pip section of the model’s conda environment (conda.yaml) file.

  • extra_pip_requirements

    Either an iterable of pip requirement strings (e.g. ["pandas", "-r requirements.txt", "-c constraints.txt"]) or the string path to a pip requirements file on the local filesystem (e.g. "requirements.txt"). If provided, this describes additional pip requirements that are appended to a default set of pip requirements generated automatically based on the user’s current software environment. Both requirements and constraints are automatically parsed and written to requirements.txt and constraints.txt files, respectively, and stored as part of the model. Requirements are also written to the pip section of the model’s conda environment (conda.yaml) file.

    Warning

    The following arguments can’t be specified at the same time:

    • conda_env

    • pip_requirements

    • extra_pip_requirements

    This example demonstrates how to specify pip requirements using pip_requirements and extra_pip_requirements.

  • metadata

    Custom metadata dictionary passed to the model and stored in the MLmodel file.

    Note

    Experimental: This parameter may change or be removed in a future release without warning.

  • model_config

    The model configuration to apply to the model. This configuration is available during model loading.

    Note

    Experimental: This parameter may change or be removed in a future release without warning.

  • example_no_conversion – If True, the input example will not be converted to a Pandas DataFrame format when saving. This is useful when the model expects a non-DataFrame input and the input example could be passed directly to the model. Defaults to False for backwards compatibility.

Returns

A ModelInfo instance that contains the metadata of the logged model.

mlflow.pyfunc.save_model(path, loader_module=None, data_path=None, code_path=None, conda_env=None, mlflow_model=Model(), python_model=None, artifacts=None)[source]

Save a Pyfunc model with custom inference logic and optional data dependencies to a path on the local filesystem.

For information about the workflows that this method supports, please see “workflows for creating custom pyfunc models” and “which workflow is right for my use case?”. Note that the parameters for the second workflow: loader_module, data_path and the parameters for the first workflow: python_model, artifacts, cannot be specified together.

Parameters
  • path – The path to which to save the Python model.

  • loader_module

    The name of the Python module that is used to load the model from data_path. This module must define a method with the prototype _load_pyfunc(data_path). If not None, this module and its dependencies must be included in one of the following locations:

    • The MLflow library.

    • Package(s) listed in the model’s Conda environment, specified by the conda_env parameter.

    • One or more of the files specified by the code_path parameter.

  • data_path – Path to a file or directory containing model data.

  • code_path – A list of local filesystem paths to Python file dependencies (or directories containing file dependencies). These files are prepended to the system path before the model is loaded.

  • conda_env

    Either a dictionary representation of a Conda environment or the path to a conda environment yaml file. If provided, this describes the environment this model should be run in. At a minimum, it should specify the dependencies contained in get_default_conda_env(). If None, a conda environment with pip requirements inferred by mlflow.models.infer_pip_requirements() is added to the model. If the requirement inference fails, it falls back to using get_default_pip_requirements(). pip requirements from conda_env are written to a pip requirements.txt file and the full conda environment is written to conda.yaml. The following is an example dictionary representation of a conda environment:

    {
        "name": "mlflow-env",
        "channels": ["conda-forge"],
        "dependencies": [
            "python=3.8.15",
            {
                "pip": [
                    "scikit-learn==x.y.z"
                ],
            },
        ],
    }
    

  • mlflow_modelmlflow.models.Model configuration to which to add the python_function flavor.

  • python_model

    An instance of a subclass of PythonModel or a callable object with a single argument (see the examples below). The passed-in object is serialized using the CloudPickle library. Any dependencies of the class should be included in one of the following locations:

    • The MLflow library.

    • Package(s) listed in the model’s Conda environment, specified by the conda_env parameter.

    • One or more of the files specified by the code_path parameter.

    Note: If the class is imported from another module, as opposed to being defined in the __main__ scope, the defining module should also be included in one of the listed locations.

    Examples

    Class model

    from typing import List, Dict
    import mlflow
    
    
    class MyModel(mlflow.pyfunc.PythonModel):
        def predict(self, context, model_input: List[str], params=None) -> List[str]:
            return [i.upper() for i in model_input]
    
    
    mlflow.pyfunc.save_model("model", python_model=MyModel(), input_example=["a"])
    model = mlflow.pyfunc.load_model("model")
    print(model.predict(["a", "b", "c"]))  # -> ["A", "B", "C"]
    

    Functional model

    Note

    Experimental: Functional model support is experimental and may change or be removed in a future release without warning.

    from typing import List
    import mlflow
    
    
    def predict(model_input: List[str]) -> List[str]:
        return [i.upper() for i in model_input]
    
    
    mlflow.pyfunc.save_model("model", python_model=predict, input_example=["a"])
    model = mlflow.pyfunc.load_model("model")
    print(model.predict(["a", "b", "c"]))  # -> ["A", "B", "C"]
    

    If the predict method or function has type annotations, MLflow automatically constructs a model signature based on the type annotations (unless the signature argument is explicitly specified), and converts the input value to the specified type before passing it to the function. Currently, the following type annotations are supported:

    • List[str]

    • List[Dict[str, str]]

  • artifacts

    A dictionary containing <name, artifact_uri> entries. Remote artifact URIs are resolved to absolute filesystem paths, producing a dictionary of <name, absolute_path> entries. python_model can reference these resolved entries as the artifacts property of the context parameter in PythonModel.load_context() and PythonModel.predict(). For example, consider the following artifacts dictionary:

    {
        "my_file": "s3://my-bucket/path/to/my/file"
    }
    

    In this case, the "my_file" artifact is downloaded from S3. The python_model can then refer to "my_file" as an absolute filesystem path via context.artifacts["my_file"].

    If None, no artifacts are added to the model.

  • signature

    ModelSignature describes model input and output Schema. The model signature can be inferred from datasets with valid model input (e.g. the training dataset with target column omitted) and valid model output (e.g. model predictions generated on the training dataset), for example:

    from mlflow.models import infer_signature
    
    train = df.drop_column("target_label")
    predictions = ...  # compute model predictions
    signature = infer_signature(train, predictions)
    

  • input_example – one or several instances of valid model input. The input example is used as a hint of what data to feed the model. It will be converted to a Pandas DataFrame and then serialized to json using the Pandas split-oriented format, or a numpy array where the example will be serialized to json by converting it to a list. Bytes are base64-encoded. When the signature parameter is None, the input example is used to infer a model signature.

  • pip_requirements – Either an iterable of pip requirement strings (e.g. ["scikit-learn", "-r requirements.txt", "-c constraints.txt"]) or the string path to a pip requirements file on the local filesystem (e.g. "requirements.txt"). If provided, this describes the environment this model should be run in. If None, a default list of requirements is inferred by mlflow.models.infer_pip_requirements() from the current software environment. If the requirement inference fails, it falls back to using get_default_pip_requirements(). Both requirements and constraints are automatically parsed and written to requirements.txt and constraints.txt files, respectively, and stored as part of the model. Requirements are also written to the pip section of the model’s conda environment (conda.yaml) file.

  • extra_pip_requirements

    Either an iterable of pip requirement strings (e.g. ["pandas", "-r requirements.txt", "-c constraints.txt"]) or the string path to a pip requirements file on the local filesystem (e.g. "requirements.txt"). If provided, this describes additional pip requirements that are appended to a default set of pip requirements generated automatically based on the user’s current software environment. Both requirements and constraints are automatically parsed and written to requirements.txt and constraints.txt files, respectively, and stored as part of the model. Requirements are also written to the pip section of the model’s conda environment (conda.yaml) file.

    Warning

    The following arguments can’t be specified at the same time:

    • conda_env

    • pip_requirements

    • extra_pip_requirements

    This example demonstrates how to specify pip requirements using pip_requirements and extra_pip_requirements.

  • metadata

    Custom metadata dictionary passed to the model and stored in the MLmodel file.

    Note

    Experimental: This parameter may change or be removed in a future release without warning.

  • model_config

    The model configuration to apply to the model. This configuration is available during model loading.

    Note

    Experimental: This parameter may change or be removed in a future release without warning.

  • example_no_conversion – If True, the input example will not be converted to a Pandas DataFrame format when saving. This is useful when the model expects a non-DataFrame input and the input example could be passed directly to the model. Defaults to False for backwards compatibility.

mlflow.pyfunc.spark_udf(spark, model_uri, result_type=None, env_manager='local', params: Optional[Dict[str, Any]] = None, extra_env: Optional[Dict[str, str]] = None)[source]

A Spark UDF that can be used to invoke the Python function formatted model.

Parameters passed to the UDF are forwarded to the model as a DataFrame where the column names are ordinals (0, 1, …). On some versions of Spark (3.0 and above), it is also possible to wrap the input in a struct. In that case, the data will be passed as a DataFrame with column names given by the struct definition (e.g. when invoked as my_udf(struct(‘x’, ‘y’)), the model will get the data as a pandas DataFrame with 2 columns ‘x’ and ‘y’).

If a model contains a signature with tensor spec inputs, you will need to pass a column of array type as a corresponding UDF argument. The column values of which must be one dimensional arrays. The UDF will reshape the column values to the required shape with ‘C’ order (i.e. read / write the elements using C-like index order) and cast the values as the required tensor spec type.

If a model contains a signature, the UDF can be called without specifying column name arguments. In this case, the UDF will be called with column names from signature, so the evaluation dataframe’s column names must match the model signature’s column names.

The predictions are filtered to contain only the columns that can be represented as the result_type. If the result_type is string or array of strings, all predictions are converted to string. If the result type is not an array type, the left most column with matching type is returned.

NOTE: Inputs of type pyspark.sql.types.DateType are not supported on earlier versions of Spark (2.4 and below).

Example
from pyspark.sql.functions import struct

predict = mlflow.pyfunc.spark_udf(spark, "/my/local/model")
df.withColumn("prediction", predict(struct("name", "age"))).show()
Parameters
  • spark – A SparkSession object.

  • model_uri

    The location, in URI format, of the MLflow model with the mlflow.pyfunc flavor. For example:

    • /Users/me/path/to/local/model

    • relative/path/to/local/model

    • s3://my_bucket/path/to/model

    • runs:/<mlflow_run_id>/run-relative/path/to/model

    • models:/<model_name>/<model_version>

    • models:/<model_name>/<stage>

    • mlflow-artifacts:/path/to/model

    For more information about supported URI schemes, see Referencing Artifacts.

  • result_type

    the return type of the user-defined function. The value can be either a pyspark.sql.types.DataType object or a DDL-formatted type string. Only a primitive type, an array pyspark.sql.types.ArrayType of primitive type, or a struct type containing fields of above 2 kinds of types are allowed. If unspecified, it tries to infer result type from model signature output schema, if model output schema is not available, it fallbacks to use double type.

    The following classes of result type are supported:

    • ”int” or pyspark.sql.types.IntegerType: The leftmost integer that can fit in an int32 or an exception if there is none.

    • ”long” or pyspark.sql.types.LongType: The leftmost long integer that can fit in an int64 or an exception if there is none.

    • ArrayType(IntegerType|LongType): All integer columns that can fit into the requested size.

    • ”float” or pyspark.sql.types.FloatType: The leftmost numeric result cast to float32 or an exception if there is none.

    • ”double” or pyspark.sql.types.DoubleType: The leftmost numeric result cast to double or an exception if there is none.

    • ArrayType(FloatType|DoubleType): All numeric columns cast to the requested type or an exception if there are no numeric columns.

    • ”string” or pyspark.sql.types.StringType: The leftmost column converted to string.

    • ”boolean” or “bool” or pyspark.sql.types.BooleanType: The leftmost column converted to bool or an exception if there is none.

    • ArrayType(StringType): All columns converted to string.

    • ”field1 FIELD1_TYPE, field2 FIELD2_TYPE, …”: A struct type containing multiple fields separated by comma, each field type must be one of types listed above.

  • env_manager

    The environment manager to use in order to create the python environment for model inference. Note that environment is only restored in the context of the PySpark UDF; the software environment outside of the UDF is unaffected. Default value is local, and the following values are supported:

    • virtualenv: Use virtualenv to restore the python environment that was used to train the model.

    • conda: (Recommended) Use Conda to restore the software environment that was used to train the model.

    • local: Use the current Python environment for model inference, which may differ from the environment used to train the model and may lead to errors or invalid predictions.

  • params

    Additional parameters to pass to the model for inference.

    Note

    Experimental: This parameter may change or be removed in a future release without warning.

  • extra_env – Extra environment variables to pass to the UDF executors.

Returns

Spark UDF that applies the model’s predict method to the data and returns a type specified by result_type, which by default is a double.

mlflow.pyfunc.get_default_pip_requirements()[source]
Returns

A list of default pip requirements for MLflow Models produced by this flavor. Calls to save_model() and log_model() produce a pip environment that, at minimum, contains these requirements.

mlflow.pyfunc.get_default_conda_env()[source]
Returns

The default Conda environment for MLflow Models produced by calls to save_model() and log_model() when a user-defined subclass of PythonModel is provided.

class mlflow.pyfunc.PythonModelContext[source]

A collection of artifacts that a PythonModel can use when performing inference. PythonModelContext objects are created implicitly by the save_model() and log_model() persistence methods, using the contents specified by the artifacts parameter of these methods.

property artifacts

A dictionary containing <name, artifact_path> entries, where artifact_path is an absolute filesystem path to the artifact.

property model_config

A dictionary containing <config, value> entries, where config is the name of the model configuration keys and value is the value of the given configuration.

Note

Experimental: This property may change or be removed in a future release without warning.

class mlflow.pyfunc.PythonModel[source]

Represents a generic Python model that evaluates inputs and produces API-compatible outputs. By subclassing PythonModel, users can create customized MLflow models with the “python_function” (“pyfunc”) flavor, leveraging custom inference logic and artifact dependencies.

load_context(context)[source]

Loads artifacts from the specified PythonModelContext that can be used by predict() when evaluating inputs. When loading an MLflow model with load_model(), this method is called as soon as the PythonModel is constructed.

The same PythonModelContext will also be available during calls to predict(), but it may be more efficient to override this method and load artifacts from the context at model load time.

Parameters

context – A PythonModelContext instance containing artifacts that the model can use to perform inference.

abstract predict(context, model_input, params: Optional[Dict[str, Any]] = None)[source]

Evaluates a pyfunc-compatible input and produces a pyfunc-compatible output. For more information about the pyfunc input/output API, see the Inference API.

Parameters
  • context – A PythonModelContext instance containing artifacts that the model can use to perform inference.

  • model_input – A pyfunc-compatible input for the model to evaluate.

  • params

    Additional parameters to pass to the model for inference.

    Note

    Experimental: This parameter may change or be removed in a future release without warning.

class mlflow.pyfunc.ChatModel[source]

Note

Experimental: This class may change or be removed in a future release without warning.

A subclass of PythonModel that makes it more convenient to implement models that are compatible with popular LLM chat APIs. By subclassing ChatModel, users can create MLflow models with a predict() method that is more convenient for chat tasks than the generic PythonModel API. ChatModels automatically define input/output signatures and an input example, so manually specifying these values when calling mlflow.pyfunc.save_model() is not necessary.

See the documentation of the predict() method below for details on that parameters and outputs that are expected by the ChatModel API.

abstract predict(context, messages: List[mlflow.types.llm.ChatMessage], params: mlflow.types.llm.ChatParams)mlflow.types.llm.ChatResponse[source]

Evaluates a chat input and produces a chat output.

Parameters
  • messages (List[ChatMessage]) – A list of ChatMessage objects representing chat history.

  • params (ChatParams) – A ChatParams object containing various parameters used to modify model behavior during inference.

Returns

A ChatResponse object containing the model’s response(s), as well as other metadata.