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. (Python 函数模型格式是 MLflow Python 模型的默认模型接口。任何 MLflow Python 模型都应能被加载为 python_function 模型。)

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. (此外，mlflow.pyfunc 模块为 Python 模型定义了一种通用的 文件系统格式，并提供了保存到此格式和从中加载的实用工具。该格式是自包含的，意味着它包含了任何人加载和使用它所需的所有信息。依赖项要么直接存储在模型中，要么通过 Conda 环境引用。)

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 Models From Code for Custom Models. (此外，mlflow.pyfunc 模块还定义了使用 MLflow 未原生包含的框架和推理逻辑来创建自定义 pyfunc 模型的实用工具。请参阅 代码构建的自定义模型。)

推理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 (Python 函数模型被加载为 PyFuncModel 的实例，这是 MLflow 对模型实现和模型元数据（MLmodel 文件）的包装器。您可以通过调用具有以下签名的 predict() 方法来对模型进行评分：)

predict(
  model_input: [pandas.DataFrame, numpy.ndarray, scipy.sparse.(csc_matrix | csr_matrix),
  List[Any], Dict[str, Any], pyspark.sql.DataFrame]
) -> [numpy.ndarray | pandas.(Series | DataFrame) | List | Dict | 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). (所有 PyFunc 模型都将支持 pandas.DataFrame 作为输入，而 PyFunc 深度学习模型还将支持 Dict[str, numpy.ndarray]（命名张量）和 numpy.ndarray（未命名张量）形式的张量输入。)

Here are some examples of supported inference types, assuming we have the correct model object loaded. (以下是一些支持的推理类型的示例，假设我们已加载了正确的 model 对象。)

输入类型

示例

pandas.DataFrame

 
import pandas as pd

x_new = pd.DataFrame(dict(x1=[1, 2, 3], x2=[4, 5, 6]))
model.predict(x_new)

numpy.ndarray

 
import numpy as np

x_new = np.array([[1, 4][2, 5], [3, 6]])
model.predict(x_new)

scipy.sparse.csc_matrix or scipy.sparse.csr_matrix

 
import scipy

x_new = scipy.sparse.csc_matrix([[1, 2, 3], [4, 5, 6]])
model.predict(x_new)

x_new = scipy.sparse.csr_matrix([[1, 2, 3], [4, 5, 6]])
model.predict(x_new)

python List

 
x_new = [[1, 4], [2, 5], [3, 6]]
model.predict(x_new)

python Dict

 
x_new = dict(x1=[1, 2, 3], x2=[4, 5, 6])
model.predict(x_new)

pyspark.sql.DataFrame

 
from pyspark.sql import SparkSession

spark = SparkSession.builder.getOrCreate()

data = [(1, 4), (2, 5), (3, 6)]  # List of tuples
x_new = spark.createDataFrame(data, ["x1", "x2"])  # Specify column name
model.predict(x_new)

文件系统格式

The Pyfunc format is defined as a directory structure containing all required data, code, and configuration (Pyfunc 格式定义为包含所有必需数据、代码和配置的目录结构。)

./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 配置引用。)

MLModel 配置

A Python model contains an MLmodel file in python_function format in its root with the following parameters (Python 模型在其根目录下包含一个 python_function 格式的 MLmodel 文件，具有以下参数：)

  • 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 (可加载模型的 Python 模块。期望作为模块标识符，例如 mlflow.sklearn，它将使用 importlib.import_module 导入。导入的模块必须包含一个具有以下签名的函数：)

    _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 (path 参数由 data 参数指定，可以指向文件或目录。模型实现应为一个具有 predict 方法的对象，该方法的签名如下：)

    predict(
  model_input: [pandas.DataFrame, numpy.ndarray,
  scipy.sparse.(csc_matrix | csr_matrix), List[Any], Dict[str, Any]],
  pyspark.sql.DataFrame
) -> [numpy.ndarray | pandas.(Series | DataFrame) | List | Dict | 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. (指向包含与此模型打包在一起的代码的目录的相对路径。导入模型加载器之前，此目录内的所有文件和目录都将被添加到 Python 路径中。)

  • 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. (指向导出的 Conda 环境的相对路径。如果存在，应在运行模型之前激活此环境。)

  • Optionally, any additional parameters necessary for interpreting the serialized model in pyfunc format. (可选，解释 pyfunc 格式的序列化模型所需的任何其他参数。)

示例

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

代码构建的自定义模型

提示

MLflow 2.12.2 introduced the feature “models from code”, which greatly simplifies the process of serializing and deploying custom models through the use of script serialization. It is strongly recommended to migrate custom model implementations to this new paradigm to avoid the limitations and complexity of serializing with cloudpickle. You can learn more about models from code within the Models From Code Guide. (MLflow 2.12.2 引入了“代码构建的模型”功能，通过使用脚本序列化极大地简化了自定义模型的序列化和部署过程。强烈建议将自定义模型实现迁移到这个新范例，以避免使用 cloudpickle 进行序列化的限制和复杂性。您可以在 代码构建模型指南 中了解有关代码构建模型的更多信息。)

The section below illustrates the process of using the legacy serializer for custom Pyfunc models. Models from code will provide a far simpler experience for logging of your models. (下面的部分说明了为自定义 Pyfunc 模型使用旧版序列化器的过程。代码构建的模型将为您的模型日志记录提供更简单的体验。)

创建自定义 Pyfunc 模型

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. (MLflow 的持久化模块提供了方便的函数，用于在各种机器学习框架（scikit-learn、Keras、Pytorch 等）中创建具有 pyfunc 格式的模型；但是，它们并不适用于所有用例。例如，您可能希望使用 MLflow 不原生支持的框架，以 pyfunc 格式创建 MLflow 模型。或者，您可能希望构建一个在评估查询时执行自定义逻辑的 MLflow 模型，例如预处理和后处理例程。因此，mlflow.pyfunc 提供了从任意代码和模型数据创建 pyfunc 模型的实用工具。)

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. ( save_model() 和 log_model() 方法旨在支持多种工作流，用于创建包含自定义推理逻辑和逻辑可能需要的工件的自定义 pyfunc 模型。)

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. (一个 artifact 是一个文件或目录，例如序列化模型或 CSV。例如，序列化的 TensorFlow 图是一个工件。MLflow 模型目录也是一个工件。)

工作流

save_model() and log_model() support the following workflows ( save_model() 和 log_model() 支持以下工作流：)

  1. Programmatically defining a new MLflow model, including its attributes and artifacts. (以编程方式定义新的 MLflow 模型，包括其属性和工件。)

    Given a set of artifact URIs, save_model() and log_model() can automatically download artifacts from their URIs and create an MLflow model directory. (给定一组工件 URI，save_model() 和 log_model() 可以自动从它们的 URI 下载工件并创建一个 MLflow 模型目录。)

    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. (在这种情况下，您必须定义一个继承自 PythonModel 的 Python 类，该类定义了 predict() 方法，以及可选的 load_context() 方法。此类的实例通过 python_model 参数指定；它会被自动序列化和反序列化为一个 Python 类，包括其所有属性。)

  2. Interpreting pre-existing data as an MLflow model. (将预先存在的数据解释为 MLflow 模型。)

    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. (如果您已经有一个包含模型数据的目录，save_model() 和 log_model() 可以将数据导入为 MLflow 模型。data_path 参数指定包含模型数据的目录的本地文件系统路径。)

    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 (在这种情况下，您必须提供一个 Python 模块，称为 loader module。loader 模块定义了一个 _load_pyfunc() 方法，该方法执行以下任务：)

    • Load data from the specified data_path. For example, this process may include deserializing pickled Python objects or models or parsing CSV files. (从指定的 data_path 加载数据。例如，此过程可能包括反序列化 pickled Python 对象或模型，或解析 CSV 文件。)

    • 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. (构建并返回一个 pyfunc 兼容的模型包装器。与第一个用例一样，此包装器必须定义一个 predict() 方法，该方法用于评估查询。predict() 必须遵循 推理 API。)

    The loader_module parameter specifies the name of your loader module. ( loader_module 参数指定您的 loader 模块的名称。)

    For an example loader module implementation, refer to the loader module implementation in mlflow.sklearn. (有关 loader 模块实现的示例，请参阅 mlflow.sklearn 中的 loader 模块实现。)

哪种工作流适合我的用例？

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 (它会自动序列化和反序列化 python_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. (您可以使用在 __main__ 范围内定义的逻辑来创建模型。这允许在交互式环境（如笔记本和 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. (推理逻辑始终以代码形式持久化，而不是 Python 对象。这使得逻辑以后更容易检查和修改。)

  • 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. (如果您已经将所有模型数据收集到一个位置，第二种工作流允许直接将其保存为 MLflow 格式，而无需枚举构成工件。)

基于函数的模型 vs 基于类的模型

When creating custom PyFunc models, you can choose between two different interfaces: a function-based model and a class-based model. In short, a function-based model is simply a python function that does not take additional params. The class-based model, on the other hand, is subclass of PythonModel that supports several required and optional methods. If your use case is simple and fits within a single predict function, a function-based approach is recommended. If you need more power, such as custom serialization, custom data processing, or to override additional methods, you should use the class-based implementation. (在创建自定义 PyFunc 模型时，您可以选择两种不同的接口：基于函数的模型和基于类的模型。简而言之，基于函数的模型只是一个不带额外参数的 Python 函数。而基于类的模型是 PythonModel 的子类，它支持几个必需的和可选的方法。如果您的用例很简单，并且可以包含在单个 predict 函数中，则建议使用基于函数的方法。如果您需要更多功能，例如自定义序列化、自定义数据处理或覆盖其他方法，则应使用基于类的实现。)

Before looking at code examples, it’s important to note that both methods are serialized via cloudpickle. cloudpickle can serialize Python functions, lambda functions, and locally defined classes and functions inside other functions. This makes cloudpickle especially useful for parallel and distributed computing where code objects need to be sent over network to execute on remote workers, which is a common deployment paradigm for MLflow. (在查看代码示例之前，需要注意的是，这两种方法都通过 cloudpickle 进行序列化。cloudpickle 可以序列化 Python 函数、lambda 函数以及在其他函数中本地定义的类和函数。这使得 cloudpickle 在并行和分布式计算中特别有用，因为在这些计算中需要将代码对象通过网络发送到远程工作节点执行，这是 MLflow 的一种常见部署模式。)

That said, cloudpickle has some limitations. (尽管如此，cloudpickle 仍存在一些限制。)

  • Environment Dependency: cloudpickle does not capture the full execution environment, so in MLflow we must pass pip_requirements, extra_pip_requirements, or an input_example, the latter of which is used to infer environment dependencies. For more, refer to the model dependency docs. (环境依赖：cloudpickle 无法捕获完整的执行环境，因此在 MLflow 中，我们必须传递 pip_requirements、extra_pip_requirements 或 input_example，后者用于推断环境依赖。有关更多信息，请参阅 模型依赖文档。)

  • Object Support: cloudpickle does not serialize objects outside of the Python data model. Some relevant examples include raw files and database connections. If your program depends on these, be sure to log ways to reference these objects along with your model. (对象支持：cloudpickle 不序列化 Python 数据模型之外的对象。一些相关的例子包括原始文件和数据库连接。如果您的程序依赖于这些，请务必在模型中记录引用这些对象的方式。)

基于函数的模型

If you’re looking to serialize a simple python function without additional dependent methods, you can simply log a predict method via the keyword argument python_model. (如果您想序列化一个不带额外依赖方法的简单 Python 函数，可以通过关键字参数 python_model 直接记录一个 predict 方法。)

注意

Function-based model only supports a function with a single input argument. If you would like to pass more arguments or additional inference parameters, please use the class-based model below. (基于函数的模型仅支持带有单个输入参数的函数。如果您想传递更多参数或额外的推理参数，请使用下面的基于类的模型。)

import mlflow
import pandas as pd


# Define a simple function to log
def predict(model_input):
    return model_input.apply(lambda x: x * 2)


# Save the function as a model
with mlflow.start_run():
    mlflow.pyfunc.log_model(
        name="model", python_model=predict, pip_requirements=["pandas"]
    )
    run_id = mlflow.active_run().info.run_id

# Load the model from the tracking server and perform inference
model = mlflow.pyfunc.load_model(f"runs:/{run_id}/model")
x_new = pd.Series([1, 2, 3])

prediction = model.predict(x_new)
print(prediction)

基于类的模型

If you’re looking to serialize a more complex object, for instance a class that handles preprocessing, complex prediction logic, or custom serialization, you should subclass the PythonModel class. MLflow has tutorials on building custom PyFunc models, as shown here, so instead of duplicating that information, in this example we’ll recreate the above functionality to highlight the differences. Note that this PythonModel implementation is overly complex and we would recommend using the functional-based Model instead for this simple case. (如果您想序列化一个更复杂的对象，例如一个处理预处理、复杂预测逻辑或自定义序列化的类，您应该继承 PythonModel 类。MLflow 提供了构建自定义 PyFunc 模型的教程，如 此处 所示，因此，为了避免重复信息，在本例中，我们将重现上述功能以突出区别。请注意，这个 PythonModel 实现过于复杂，对于这个简单的用例，我们建议使用基于函数的模型。)

import mlflow
import pandas as pd


class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, context, model_input, params=None):
        return [x * 2 for x in model_input]


# Save the function as a model
with mlflow.start_run():
    mlflow.pyfunc.log_model(
        name="model", python_model=MyModel(), pip_requirements=["pandas"]
    )
    run_id = mlflow.active_run().info.run_id

# Load the model from the tracking server and perform inference
model = mlflow.pyfunc.load_model(f"runs:/{run_id}/model")
x_new = pd.Series([1, 2, 3])

print(f"Prediction:
    {model.predict(x_new)}")

The primary difference between the this implementation and the function-based implementation above is that the predict method is wrapped with a class, has the self parameter, and has the params parameter that defaults to None. Note that function-based models don’t support additional params. (此实现与上述基于函数的实现之间的主要区别在于，predict 方法被包装在一个类中，具有 self 参数，并且具有默认为 None 的 params 参数。请注意，基于函数的模型不支持额外参数。)

In summary, use the function-based Model when you have a simple function to serialize. If you need more power, use the class-based model. (总而言之，当您有一个简单的函数需要序列化时，请使用基于函数的模型。如果您需要更多功能，请使用基于类的模型。)

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', predict_stream_fn: Optional[str] = None, model_id: Optional[str] = None)[source]

Bases: object

MLflow ‘python function’ model. (MLflow ‘python 函数’模型。)

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(). (模型实现和元数据的包装器。此类不应直接构造。相反，此类实例由 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_impl 可以是实现 Pyfunc 接口的任何 Python 对象，并通过调用模型的 loader_module 返回。)

model_meta contains model metadata loaded from the MLmodel file. ( model_meta 包含从 MLmodel 文件加载的模型元数据。)

get_raw_model()[source]

Get the underlying raw model if the model wrapper implemented get_raw_model function. (如果模型包装器实现了 get_raw_model 函数，则获取底层原始模型。)

property input_example: Optional[Any]

The input example provided when the model was saved. (保存模型时提供的输入示例。)

property loader_module

Model’s flavor configuration (模型的格式配置。)

property metadata: mlflow.models.model.Model

Model metadata. (模型元数据。)

property model_config

Model’s flavor configuration (模型的格式配置。)

property model_id: str | None

The model ID of the model. (模型的 model ID。)

返回

The model ID of the model. (模型的 model ID。)

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: dict[str, typing.Any] | None = None) → pandas.core.frame.DataFrame | pandas.core.series.Series | numpy.ndarray | list | str | pyspark.sql.dataframe.DataFrame[source]
predict_stream(data: dict[str, typing.Any] | bool | bytes | float | int | str, params: Optional[dict[str, typing.Any]] = None) → Iterator[dict[str, typing.Any] | str][source]
unwrap_python_model()[source]

Unwrap the underlying Python model object. (解开底层 Python 模型对象。)

This method is useful for accessing custom model functions, while still being able to leverage the MLflow designed workflow through the predict() method. (此方法对于访问自定义模型函数很有用，同时仍然可以通过 predict() 方法利用 MLflow 设计的工作流。)

返回

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(name="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, model_code_path=None, **kwargs)[source]

Add a pyfunc spec to the model configuration. (将 pyfunc 规范添加到模型配置中。)

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. (定义 pyfunc 配置模式。调用者可以使用此模式从现有目录结构创建有效的 pyfunc 模型格式。例如，其他模型格式可以使用此来指定如何将其输出用作 pyfunc。)

注意

All paths are relative to the exported model root directory. (所有路径都相对于导出的模型根目录。)

参数

  • model – Existing model. (model – 现有模型。)

  • loader_module – The module to be used to load the model. (loader_module – 用于加载模型的模块。)

  • data – Path to the model data. (data – 模型数据的路径。)

  • code – Path to the code dependencies. (code – 代码依赖项的路径。)

  • conda_env – Conda environment. (conda_env – Conda 环境。)

  • python_env – Python environment. (python_env – Python 环境。)

  • model_config –

    The model configuration to apply to the model. This configuration is available during model loading. (要应用于模型的模型配置。此配置在模型加载期间可用。)

    注意

    Experimental: This parameter may change or be removed in a future release without warning. (实验性：此参数可能在未来的版本中更改或移除，恕不另行通知。)

  • model_code_path – Path to the model code. (model_code_path – 模型代码的路径。)

  • kwargs – Additional key-value pairs to include in the pyfunc flavor specification. Values must be YAML-serializable. (kwargs – 要包含在 pyfunc 格式规范中的其他键值对。值必须是 YAML 可序列化的。)

返回

Updated model configuration. (更新后的模型配置。)

mlflow.pyfunc.build_model_env(model_uri, save_path, env_manager='virtualenv')[source]

Prebuild model python environment and generate an archive file saved to provided save_path. (预构建模型的 Python 环境，并生成一个保存到指定 save_path 的存档文件。)

Typical usages (典型用法)

  • Pre-build a model’s environment in Databricks Runtime and then download the prebuilt python environment archive file. This pre-built environment archive can then be used in mlflow.pyfunc.spark_udf for remote inference execution when using Databricks Connect to remotely connect to a Databricks environment for code execution. (在 Databricks Runtime 中预构建模型的环境，然后下载预构建的 Python 环境存档文件。此预构建环境存档随后可在 mlflow.pyfunc.spark_udf 中用于远程推理执行，当使用 Databricks Connect 远程连接到 Databricks 环境进行代码执行时。)

注意

The build_model_env API is intended to only work when executed within Databricks runtime, serving the purpose of capturing the required execution environment that is needed for remote code execution when using DBConnect. The environment archive is designed to be used when performing remote execution using mlflow.pyfunc.spark_udf in Databricks runtime or Databricks Connect client and has no other purpose. The prebuilt env archive file cannot be used across different Databricks runtime versions or different platform machines. As such, if you connect to a different cluster that is running a different runtime version on Databricks, you will need to execute this API in a notebook and retrieve the generated archive to your local machine. Each environment snapshot is unique to the the model, the runtime version of your remote Databricks cluster, and the specification of the udf execution environment. When using the prebuilt env in mlflow.pyfunc.spark_udf, MLflow will verify whether the spark UDF sandbox environment matches the prebuilt env requirements and will raise Exceptions if there are compatibility issues. If these occur, simply re-run this API in the cluster that you are attempting to attach to. ( build_model_env API 仅旨在 Databricks Runtime 中执行时有效，其目的是捕获使用 DBConnect 进行远程代码执行所需的执行环境。环境存档旨在用于在 Databricks Runtime 或 Databricks Connect 客户端中使用 mlflow.pyfunc.spark_udf 进行远程执行时，没有其他用途。预构建的环境存档文件不能在不同的 Databricks Runtime 版本或不同的平台机器之间使用。因此，如果您连接到 Databricks 上运行不同 Runtime 版本的不同集群，您需要在笔记本中执行此 API，并将生成的存档检索到本地机器。每个环境快照都唯一对应于模型、远程 Databricks 集群的 Runtime 版本以及 UDF 执行环境的规范。在 mlflow.pyfunc.spark_udf 中使用预构建环境时，MLflow 将验证 Spark UDF 沙箱环境是否与预构建环境要求匹配，并在存在兼容性问题时引发异常。如果发生这种情况，只需在您尝试连接的集群中重新运行此 API。)

示例 
from mlflow.pyfunc import build_model_env

# Create a python environment archive file at the path `prebuilt_env_uri`
prebuilt_env_uri = build_model_env(f"runs:/{run_id}/model", "/path/to/save_directory")
参数

  • model_uri – URI to the model that is used to build the python environment. (model_uri – 用于构建 Python 环境的模型 URI。)

  • save_path – The directory path that is used to save the prebuilt model environment archive file path. The path can be either local directory path or mounted DBFS path such as ‘/dbfs/…’ or mounted UC volume path such as ‘/Volumes/…’. (save_path – 用于保存预构建模型环境存档文件路径的目录路径。该路径可以是本地目录路径，也可以是挂载的 DBFS 路径（例如 ‘/dbfs/…’）或挂载的 UC 卷路径（例如 ‘/Volumes/…’）。)

  • env_manager – The environment manager to use in order to create the python environment for model inference, the value can be either ‘virtualenv’ or ‘uv’, the default value is ‘virtualenv’. (env_manager – 用于创建模型推理的 Python 环境的环境管理器，该值可以是 ‘virtualenv’ 或 ‘uv’，默认值为 ‘virtualenv’。)

返回

Return the path of an archive file containing the python environment data. (返回包含 Python 环境数据的存档文件的路径。)

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. (下载模型依赖项并返回 requirements.txt 或 conda.yaml 文件的路径。)

警告

此 API 将所有模型构件下载到本地文件系统。对于大型模型，这可能需要很长时间。为避免此开销，请改用 mlflow.artifacts.download_artifacts("<model_uri>/requirements.txt") 或 mlflow.artifacts.download_artifacts("<model_uri>/conda.yaml")。

参数

  • model_uri – 从中获取依赖项的模型的 URI。

  • format – 返回的依赖文件格式。如果指定了 "pip" 格式，则返回 pip requirements.txt 文件的路径。如果指定了 "conda" 格式，则返回 "conda.yaml" 文件的路径。如果指定了 "pip" 格式但模型未使用 requirements.txt 文件保存，则改用提取模型 conda.yaml 文件的 pip 部分，并忽略任何其他 conda 依赖项。默认值为 "pip"。

返回

指定模型依赖项的 pip requirements.txt 文件（如果 format="pip"）或 conda.yaml 文件（如果 format="conda"）的本地文件系统路径。

mlflow.pyfunc.load_model(model_uri: str, suppress_warnings: bool = False, dst_path: str | None = None, model_config: str | pathlib.Path | dict[str, typing.Any] | None = None) → mlflow.pyfunc.PyFuncModel[source]

加载以 Python 函数格式存储的模型。

参数

  • model_uri –

    MLflow 模型在 URI 格式中的位置。例如：

    • /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 – 如果为 True，则会抑制与模型加载过程相关的非致命警告消息。如果为 False，则会发出这些警告消息。

  • 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 –

    应用于模型的模型配置。该配置将作为 PythonModel.load_context() 和 PythonModel.predict() 的 context 参数的 model_config 属性。配置可以作为文件路径或带有字符串键的字典传递。

    注意

    Experimental: This parameter may change or be removed in a future release without warning. (实验性：此参数可能在未来的版本中更改或移除，恕不另行通知。)

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

警告

mlflow.pyfunc.load_pyfunc 已弃用自 1.0 版本起。此方法将在未来版本中移除。请使用 mlflow.pyfunc.load_model。

加载以 Python 函数格式存储的模型。

参数

  • model_uri –

    MLflow 模型在 URI 格式中的位置。例如：

    • /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 – 如果为 True，则会抑制与模型加载过程相关的非致命警告消息。如果为 False，则会发出这些警告消息。

mlflow.pyfunc.log_model(artifact_path=None, loader_module=None, data_path=None, code_paths=None, infer_code_paths=False, 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, streamable=None, resources: str | list[mlflow.models.resources.Resource] | None = None, auth_policy: mlflow.models.auth_policy.AuthPolicy | None = None, prompts: list[str | Prompt] | None = None, name=None, params: dict[str, typing.Any] | None = None, tags: dict[str, typing.Any] | None = None, model_type: str | None = None, step: int = 0, model_id: str | None = None)[source]

将具有自定义推理逻辑和可选数据依赖项的 Pyfunc 模型记录为当前运行的 MLflow 构件。

有关此方法支持的工作流的信息，请参阅 创建自定义 pyfunc 模型的工作流 和 哪种工作流适合我的用例？。您不能同时指定第二个工作流的参数：loader_module、data_path 以及第一个工作流的参数：python_model、artifacts。

参数

  • artifact_path – Deprecated. Use name instead.

  • loader_module –

    用于从 data_path 加载模型的 Python 模块的名称。此模块必须定义一个原型为 _load_pyfunc(data_path) 的方法。如果不是 None，则此模块及其依赖项必须包含在以下位置之一：

    • MLflow 库。

    • conda_env 参数指定的模型 Conda 环境中列出的程序包。

    • code_paths 参数指定的​​一个或多个文件。

  • data_path – 包含模型数据的​​文件或目录的路径。

  • code_paths –

    A list of local filesystem paths to Python file dependencies (or directories containing file dependencies). These files are prepended to the system path when the model is loaded. Files declared as dependencies for a given model should have relative imports declared from a common root path if multiple files are defined with import dependencies between them to avoid import errors when loading the model.

    您可以将 code_paths 参数留空，但将 infer_code_paths 设置为 True，让 MLflow 推断模型代码路径。请参阅 infer_code_paths 参数文档了解详情。

    For a detailed explanation of code_paths functionality, recommended usage patterns and limitations, see the code_paths usage guide.

  • infer_code_paths –

    如果设置为 True，MLflow 将自动推断模型代码路径。推断出的

    代码路径文件仅包含必要的 Python 模块文件。只有当前工作目录下的 Python 代码文件可以自动推断。默认值为 False。

    警告

    请确保自定义 Python 模块代码不包含敏感数据（如凭证令牌字符串），否则它们可能会包含在自动推断的代码路径文件中并被记录到 MLflow 构件存储库。

    如果自定义 Python 模块依赖于非 Python 文件（例如 JSON 文件），且这些文件与模块代码文件路径相对，则这些非 Python 文件无法自动推断为代码路径文件。要解决此问题，您应该将所有使用的非 Python 文件放在自定义代码目录之外。

    如果 Python 代码文件被加载为 Python __main__ 模块，则此代码文件不能被推断为代码路径文件。如果您的模型依赖于在 __main__ 模块中定义的类/函数，您应该使用 cloudpickle 来转储您的模型实例，以便序列化 __main__ 中的类/函数。

    注意

    Experimental: This parameter may change or be removed in a future release without warning. (实验性：此参数可能在未来的版本中更改或移除，恕不另行通知。)

  • conda_env –

    Conda 环境的字典表示形式或 conda 环境 yaml 文件的路径。如果提供，它将描述此模型应运行的环境。至少，它应指定 get_default_conda_env() 中包含的依赖项。如果为 None，则模型会添加一个由 mlflow.models.infer_pip_requirements() 推断的 pip 依赖项的 Conda 环境。如果依赖项推断失败，则回退使用 get_default_pip_requirements。来自 conda_env 的 pip 依赖项会写入 pip requirements.txt 文件，完整的 Conda 环境会写入 conda.yaml。以下是 Conda 环境的示例字典表示形式：

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

  • python_model –

    PythonModel 的子类实例或具有单个参数的可调用对象（请参阅下面的示例）。传入的对象使用 CloudPickle 库进行序列化。python_model 也可以是 PythonModel 的文件路径，它从代码构件定义模型，而不是序列化模型对象。该类的任何依赖项都应包含在以下位置之一：

    • MLflow 库。

    • conda_env 参数指定的模型 Conda 环境中列出的程序包。

    • code_paths 参数指定的​​一个或多个文件。

    注意：如果类是从另一个模块导入的（而不是定义在 __main__ 作用域中），则定义模块也应包含在列出的位置之一。

    示例

    类模型

    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(
        name="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"]

    函数模型

    注意

    实验性：函数模型支持是实验性的，可能会在未来版本中更改或移除，恕不另行通知。

    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(
        name="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"]

    来自代码的模型

    注意

    实验性：来自代码的模型支持是实验性的，可能会在未来版本中更改或移除，恕不另行通知。

    # code.py
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]


mlflow.models.set_model(MyModel())

# log_model.py
import mlflow

with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        name="model",
        python_model="code.py",
    )

    如果 predict 方法或函数具有类型注释，MLflow 会根据类型注释自动构建模型签名（除非显式指定 signature 参数），并在将输入值传递给函数之前将其转换为指定的类型。当前支持以下类型注释：

    • List[str]

    • List[Dict[str, str]]

  • artifacts –

    包含 <name, artifact_uri> 条目的字典。远程构件 URI 会解析为绝对文件系统路径，生成一个 <name, absolute_path> 条目的字典。 python_model 可以通过 PythonModel.load_context() 和 PythonModel.predict() 的 context 参数的 artifacts 属性引用这些已解析的条目。例如，考虑以下 artifacts 字典：

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

    在这种情况下，"my_file" 构件从 S3 下载。然后 python_model 可以通过 context.artifacts["my_file"] 引用 "my_file" 作为绝对文件系统路径。

    如果为 None，则不会将任何构件添加到模型。

  • registered_model_name – 如果提供，则在 registered_model_name 下创建一个模型版本，如果给定名称的注册模型不存在，也会创建该注册模型。

  • signature –

    ModelSignature 描述了模型输入和输出 Schema。模型签名可以从具有有效模型输入的​​数据集（例如，省略目标列的训练数据集）和有效模型输出的​​数据集（例如，在训练数据集上生成的模型预测）中推断，例如：

    from mlflow.models import infer_signature

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

  • input_example – 一个或多个有效的模型输入实例。输入示例用作要馈送给模型的数据的提示。它将被转换为 Pandas DataFrame，然后使用 Pandas 的面向拆分（split-oriented）格式序列化为 json，或者转换为 numpy 数组，其中示例将通过转换为列表来序列化为 json。字节将进行 base64 编码。当 signature 参数为 None 时，输入示例用于推断模型签名。

  • await_registration_for – 等待模型版本完成创建并处于 READY 状态的秒数。默认情况下，函数等待五分钟。指定 0 或 None 可跳过等待。

  • pip_requirements – pip 依赖项字符串的可迭代对象（例如 ["scikit-learn", "-r requirements.txt", "-c constraints.txt"]）或本地文件系统上的 pip 依赖项文件的字符串路径（例如 "requirements.txt"）。如果提供，它将描述此模型应运行的环境。如果为 None，则 mlflow.models.infer_pip_requirements() 会从当前软件环境中推断出默认的依赖项列表。如果依赖项推断失败，则回退使用 get_default_pip_requirements。依赖项和约束都会被自动解析并分别写入 requirements.txt 和 constraints.txt 文件，并作为模型的一部分进行存储。依赖项也会被写入模型 Conda 环境（conda.yaml）文件的 pip 部分。

  • extra_pip_requirements –

    pip 依赖项字符串的可迭代对象（例如 ["pandas", "-r requirements.txt", "-c constraints.txt"]）或本地文件系统上的 pip 依赖项文件的字符串路径（例如 "requirements.txt"）。如果提供，它将描述附加到根据用户当前软件环境自动生成的默认 pip 依赖项集中的其他 pip 依赖项。依赖项和约束都会被自动解析并分别写入 requirements.txt 和 constraints.txt 文件，并作为模型的一部分进行存储。依赖项也会被写入模型 Conda 环境（conda.yaml）文件的 pip 部分。

    警告

    以下参数不能同时指定

    • conda_env

    • pip_requirements

    • extra_pip_requirements

    此示例演示了如何使用 pip_requirements 和 extra_pip_requirements 指定 pip requirements。

  • metadata – 传递给模型并存储在 MLmodel 文件中的自定义元数据字典。

  • model_config –

    应用于模型的模型配置。该配置将作为 PythonModel.load_context() 和 PythonModel.predict() 的 context 参数的 model_config 属性。配置可以作为文件路径或带有字符串键的字典传递。

    注意

    Experimental: This parameter may change or be removed in a future release without warning. (实验性：此参数可能在未来的版本中更改或移除，恕不另行通知。)

  • streamable – 一个布尔值，指示模型是否支持流式预测。如果为 None，MLflow 将尝试通过检查是否存在 predict_stream 方法来检查模型是否支持流式处理。默认为 None。

  • resources –

    模型资源列表或包含模型所需资源列表的 resources.yaml 文件。

    资源列表。

    注意

    Experimental: This parameter may change or be removed in a future release without warning. (实验性：此参数可能在未来的版本中更改或移除，恕不另行通知。)

  • auth_policy –

    指定模型的身份验证策略，其中包含两个关键组件。

    请注意，只能定义 auth_policy 或 resources 其中之一。

    • 系统身份验证策略：用于服务此模型所需的资源列表。

    • 用户身份验证策略：用户应访问的最小范围列表

      ，以便调用此模型。

    注意

    Experimental: This parameter may change or be removed in a future release without warning. (实验性：此参数可能在未来的版本中更改或移除，恕不另行通知。)

  • prompts –

    与模型关联的 MLflow 提示注册表中已注册的提示 URI 列表。每个提示 URI 格式应为 prompt:/<name>/<version>。应在将提示与模型关联之前在 MLflow 提示注册表中注册这些提示。

    这将创建模型和提示之间的双向链接。关联的提示可以在 MLmodel 文件中存储的模型元数据中看到。从提示注册表 UI，您也可以导航到模型。

    import mlflow

prompt_template = "Hi, {name}! How are you doing today?"

# Register a prompt in the MLflow Prompt Registry
mlflow.prompts.register_prompt("my_prompt", prompt_template, description="A simple prompt")

# Log a model with the registered prompt
with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        name=MyModel(),
        name="model",
        prompts=["prompt:/my_prompt/1"]
    )

print(model_info.prompts)
# Output: ['prompt:/my_prompt/1']

# Load the prompt
prompt = mlflow.genai.load_prompt(model_info.prompts[0])

  • name – 模型名称。

  • params – 要与模型一起记录的参数字典。

  • tags – 要与模型一起记录的标签字典。

  • model_type – 模型的类型。

  • step – 记录模型输出和指标的步骤

  • model_id – 模型的 ID。

返回

一个 ModelInfo 实例，其中包含已记录模型的元数据。

mlflow.pyfunc.save_model(path, loader_module=None, data_path=None, code_paths=None, infer_code_paths=False, conda_env=None, mlflow_model=None, python_model=None, artifacts=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, pip_requirements=None, extra_pip_requirements=None, metadata=None, model_config=None, streamable=None, resources: str | list[mlflow.models.resources.Resource] | None = None, auth_policy: mlflow.models.auth_policy.AuthPolicy | None = None, **kwargs)[source]

将具有自定义推理逻辑和可选数据依赖项的 Pyfunc 模型保存到本地文件系统上的路径。

有关此方法支持的工作流的信息，请参阅 “创建自定义 pyfunc 模型的工作流” 和 “哪种工作流适合我的用例？”。请注意，第二个工作流的参数：loader_module、data_path 以及第一个工作流的参数：python_model、artifacts，不能一起指定。

参数

  • path – 要将 Python 模型保存到的路径。

  • loader_module –

    用于从 data_path 加载模型的 Python 模块的名称。此模块必须定义一个原型为 _load_pyfunc(data_path) 的方法。如果不是 None，则此模块及其依赖项必须包含在以下位置之一：

    • MLflow 库。

    • conda_env 参数指定的模型 Conda 环境中列出的程序包。

    • code_paths 参数指定的​​一个或多个文件。

  • data_path – 包含模型数据的​​文件或目录的路径。

  • code_paths –

    A list of local filesystem paths to Python file dependencies (or directories containing file dependencies). These files are prepended to the system path when the model is loaded. Files declared as dependencies for a given model should have relative imports declared from a common root path if multiple files are defined with import dependencies between them to avoid import errors when loading the model.

    您可以将 code_paths 参数留空，但将 infer_code_paths 设置为 True，让 MLflow 推断模型代码路径。请参阅 infer_code_paths 参数文档了解详情。

    For a detailed explanation of code_paths functionality, recommended usage patterns and limitations, see the code_paths usage guide.

  • infer_code_paths –

    如果设置为 True，MLflow 将自动推断模型代码路径。推断出的

    代码路径文件仅包含必要的 Python 模块文件。只有当前工作目录下的 Python 代码文件可以自动推断。默认值为 False。

    警告

    请确保自定义 Python 模块代码不包含敏感数据（如凭证令牌字符串），否则它们可能会包含在自动推断的代码路径文件中并被记录到 MLflow 构件存储库。

    如果自定义 Python 模块依赖于非 Python 文件（例如 JSON 文件），且这些文件与模块代码文件路径相对，则这些非 Python 文件无法自动推断为代码路径文件。要解决此问题，您应该将所有使用的非 Python 文件放在自定义代码目录之外。

    如果 Python 代码文件被加载为 Python __main__ 模块，则此代码文件不能被推断为代码路径文件。如果您的模型依赖于在 __main__ 模块中定义的类/函数，您应该使用 cloudpickle 来转储您的模型实例，以便序列化 __main__ 中的类/函数。

    注意

    Experimental: This parameter may change or be removed in a future release without warning. (实验性：此参数可能在未来的版本中更改或移除，恕不另行通知。)

  • conda_env –

    Conda 环境的字典表示形式或 conda 环境 yaml 文件的路径。如果提供，它将描述此模型应运行的环境。至少，它应指定 get_default_conda_env() 中包含的依赖项。如果为 None，则模型会添加一个由 mlflow.models.infer_pip_requirements() 推断的 pip 依赖项的 Conda 环境。如果依赖项推断失败，则回退使用 get_default_pip_requirements。来自 conda_env 的 pip 依赖项会写入 pip requirements.txt 文件，完整的 Conda 环境会写入 conda.yaml。以下是 Conda 环境的示例字典表示形式：

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

  • mlflow_model – 要向其添加 **python_function** flavor 的 mlflow.models.Model 配置。

  • python_model –

    PythonModel 的子类实例或具有单个参数的可调用对象（请参阅下面的示例）。传入的对象使用 CloudPickle 库进行序列化。python_model 也可以是 PythonModel 的文件路径，它从代码构件定义模型，而不是序列化模型对象。该类的任何依赖项都应包含在以下位置之一：

    • MLflow 库。

    • conda_env 参数指定的模型 Conda 环境中列出的程序包。

    • code_paths 参数指定的​​一个或多个文件。

    注意：如果类是从另一个模块导入的（而不是定义在 __main__ 作用域中），则定义模块也应包含在列出的位置之一。

    示例

    类模型

    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"]

    函数模型

    注意

    实验性：函数模型支持是实验性的，可能会在未来版本中更改或移除，恕不另行通知。

    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"]

    来自代码的模型

    注意

    实验性：来自代码的模型支持是实验性的，可能会在未来版本中更改或移除，恕不另行通知。

    # code.py
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]


mlflow.models.set_model(MyModel())

# log_model.py
import mlflow

with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        name="model",
        python_model="code.py",
    )

    如果 predict 方法或函数具有类型注释，MLflow 会根据类型注释自动构建模型签名（除非显式指定 signature 参数），并在将输入值传递给函数之前将其转换为指定的类型。当前支持以下类型注释：

    • List[str]

    • List[Dict[str, str]]

  • artifacts –

    包含 <name, artifact_uri> 条目的字典。远程构件 URI 会解析为绝对文件系统路径，生成一个 <name, absolute_path> 条目的字典。 python_model 可以通过 PythonModel.load_context() 和 PythonModel.predict() 的 context 参数的 artifacts 属性引用这些已解析的条目。例如，考虑以下 artifacts 字典：

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

    在这种情况下，"my_file" 构件从 S3 下载。然后 python_model 可以通过 context.artifacts["my_file"] 引用 "my_file" 作为绝对文件系统路径。

    如果为 None，则不会将任何构件添加到模型。

  • signature –

    ModelSignature 描述了模型输入和输出 Schema。模型签名可以从具有有效模型输入的​​数据集（例如，省略目标列的训练数据集）和有效模型输出的​​数据集（例如，在训练数据集上生成的模型预测）中推断，例如：

    from mlflow.models import infer_signature

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

  • input_example – 一个或多个有效的模型输入实例。输入示例用作要馈送给模型的数据的提示。它将被转换为 Pandas DataFrame，然后使用 Pandas 的面向拆分（split-oriented）格式序列化为 json，或者转换为 numpy 数组，其中示例将通过转换为列表来序列化为 json。字节将进行 base64 编码。当 signature 参数为 None 时，输入示例用于推断模型签名。

  • pip_requirements – pip 依赖项字符串的可迭代对象（例如 ["scikit-learn", "-r requirements.txt", "-c constraints.txt"]）或本地文件系统上的 pip 依赖项文件的字符串路径（例如 "requirements.txt"）。如果提供，它将描述此模型应运行的环境。如果为 None，则 mlflow.models.infer_pip_requirements() 会从当前软件环境中推断出默认的依赖项列表。如果依赖项推断失败，则回退使用 get_default_pip_requirements。依赖项和约束都会被自动解析并分别写入 requirements.txt 和 constraints.txt 文件，并作为模型的一部分进行存储。依赖项也会被写入模型 Conda 环境（conda.yaml）文件的 pip 部分。

  • extra_pip_requirements –

    pip 依赖项字符串的可迭代对象（例如 ["pandas", "-r requirements.txt", "-c constraints.txt"]）或本地文件系统上的 pip 依赖项文件的字符串路径（例如 "requirements.txt"）。如果提供，它将描述附加到根据用户当前软件环境自动生成的默认 pip 依赖项集中的其他 pip 依赖项。依赖项和约束都会被自动解析并分别写入 requirements.txt 和 constraints.txt 文件，并作为模型的一部分进行存储。依赖项也会被写入模型 Conda 环境（conda.yaml）文件的 pip 部分。

    警告

    以下参数不能同时指定

    • conda_env

    • pip_requirements

    • extra_pip_requirements

    此示例演示了如何使用 pip_requirements 和 extra_pip_requirements 指定 pip requirements。

  • metadata – 传递给模型并存储在 MLmodel 文件中的自定义元数据字典。

  • model_config –

    应用于模型的模型配置。该配置将作为 PythonModel.load_context() 和 PythonModel.predict() 的 context 参数的 model_config 属性。配置可以作为文件路径或带有字符串键的字典传递。

    注意

    Experimental: This parameter may change or be removed in a future release without warning. (实验性：此参数可能在未来的版本中更改或移除，恕不另行通知。)

  • streamable – 一个布尔值，指示模型是否支持流式预测。如果为 None，MLflow 将尝试通过检查是否存在 predict_stream 方法来检查模型是否支持流式处理。默认为 None。

  • resources –

    模型资源列表或包含模型所需资源列表的 resources.yaml 文件。

    资源列表。

    注意

    Experimental: This parameter may change or be removed in a future release without warning. (实验性：此参数可能在未来的版本中更改或移除，恕不另行通知。)

  • auth_policy –

    指定模型的身份验证策略，其中包含两个关键组件。

    请注意，只能定义 auth_policy 或 resources 其中之一。

    • 系统身份验证策略：用于服务此模型所需的资源列表。

    • 用户身份验证策略：用户应访问的最小范围列表

      ，以便调用此模型。

    注意

    Experimental: This parameter may change or be removed in a future release without warning. (实验性：此参数可能在未来的版本中更改或移除，恕不另行通知。)

  • kwargs – 额外的关键字参数。

mlflow.pyfunc.spark_udf(spark, model_uri, result_type=None, env_manager=None, params: Optional[dict[str, typing.Any]] = None, extra_env: Optional[dict[str, str]] = None, prebuilt_env_uri: Optional[str] = None, model_config: Optional[Union[str, pathlib.Path, dict[str, typing.Any]]] = None)[source]

可用于调用 Python 函数格式化模型的 Spark UDF。

传递给 UDF 的参数将作为 DataFrame 转发给模型，其中列名是序数（0, 1, ...）。在某些 Spark 版本（3.0 及以上）中，也可以将输入包装在 struct 中。在这种情况下，数据将作为 DataFrame 传递，其列名由 struct 定义给出（例如，当调用为 my_udf(struct(‘x’, ‘y’)) 时，模型将以带有 2 列‘x’和‘y’的 pandas DataFrame 形式获取数据）。

如果模型包含带有张量规格输入的签名，您需要将数组类型的列作为相应的 UDF 参数传递。该列的值必须是一维数组。UDF 将使用“C”顺序（即，使用类似 C 的索引顺序读取/写入元素）来重塑列值，并将值转换为所需的张量规格类型。

如果模型包含签名，则可以在不指定列名参数的情况下调用 UDF。在这种情况下，将使用签名中的列名调用 UDF，因此评估 DataFrame 的列名必须与模型签名中的列名匹配。

预测结果将被过滤，仅包含可以表示为 result_type 的列。如果 result_type 是字符串或字符串数组，则所有预测都将转换为字符串。如果结果类型不是数组类型，则返回最左边的匹配类型列。

注意

在早期版本的 Spark（2.4 及以下版本）上不支持 pyspark.sql.types.DateType 类型的输入。

注意

使用 Databricks Connect 连接到远程 Databricks 集群时，Databricks 集群必须使用运行时版本 >= 16，并且如果设置了 'prebuilt_env_uri' 参数，则不应设置 'env_manager' 参数。Databricks 集群必须使用运行时版本 >= 15.4，并且如果设置了 'prebuilt_env_uri' 参数，则不应设置 'env_manager' 参数，如果运行时版本为 15.4 且集群是标准访问模式，则需要将“spark.databricks.safespark.archive.artifact.unpack.disabled”配置为“false”。

注意

请注意，在 Databricks Serverless 中运行时，spark 任务在 Databricks Serverless UDF 沙箱的限制范围内运行。此环境的总容量限制为 1GB，结合了可用内存和本地磁盘容量。此外，此设置中没有 GPU 设备。因此，包含大型权重或需要 GPU 的任何深度学习模型均不适合在 Databricks Serverless 上部署。

Example 
from pyspark.sql.functions import struct

predict = mlflow.pyfunc.spark_udf(spark, "/my/local/model")
df.withColumn("prediction", predict(struct("name", "age"))).show()
参数

  • spark – 一个 SparkSession 对象。

  • model_uri –

    MLflow 模型的位置（URI 格式），具有 mlflow.pyfunc 格式。例如

    • /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 –

    用户定义函数的返回类型。该值可以是 pyspark.sql.types.DataType 对象或 DDL 格式的类型字符串。只允许原始类型、原始类型的数组 pyspark.sql.types.ArrayType，或包含上述两种类型字段的 struct 类型。如果未指定，它将尝试从模型签名输出架构推断结果类型；如果模型输出架构不可用，则回退使用 double 类型。

    支持以下结果类型的类别

    • “int” 或 pyspark.sql.types.IntegerType：最左边的可以适合 int32 的整数，或者如果不存在则抛出异常。

    • “long” 或 pyspark.sql.types.LongType：最左边的可以适合 int64 的长整数，或者如果不存在则抛出异常。

    • ArrayType(IntegerType|LongType)：所有可以适合请求大小的整数列。

    • “float” 或 pyspark.sql.types.FloatType：最左边的数值结果，转换为 float32，或者如果不存在则抛出异常。

    • “double” 或 pyspark.sql.types.DoubleType：最左边的数值结果，转换为 double，或者如果不存在则抛出异常。

    • ArrayType(FloatType|DoubleType)：所有数值列，转换为请求的类型，或者如果不存在数值列则抛出异常。

    • “string” 或 pyspark.sql.types.StringType：最左边的列，转换为 string。

    • “boolean” 或 “bool” 或 pyspark.sql.types.BooleanType：最左边的列，转换为 bool，或者如果不存在则抛出异常。

    • ArrayType(StringType)：所有列，转换为 string。

    • “field1 FIELD1_TYPE, field2 FIELD2_TYPE, …”：一个包含多个由逗号分隔的字段的 struct 类型，每个字段类型必须是上面列出的类型之一。

  • env_manager –

    用于创建模型推理的 Python 环境的环境管理器。请注意，环境仅在 PySpark UDF 的上下文中恢复；UDF 外部的软件环境不受影响。如果未设置 prebuilt_env_uri 参数，则默认值为 local，并支持以下值

    • virtualenv：使用 virtualenv 恢复训练模型时使用的 Python 环境。如果未设置 env_manager，则这是默认选项。

    • uv：使用 uv 恢复训练模型时使用的 Python 环境。

    • conda: 使用 Conda 恢复训练模型时使用的软件环境。

    • local: 使用当前 Python 环境进行模型推理，这可能与训练模型时使用的环境不同，并可能导致错误或无效的预测。

    如果设置了 prebuilt_env_uri 参数，则不应设置 env_manager 参数。

  • params – 传递给模型的额外参数，用于推理。

  • extra_env – 要传递给 UDF 执行器的额外环境变量。用于需要传播到 Spark 工作者的覆盖（例如，通过 MLFLOW_SCORING_SERVER_REQUEST_TIMEOUT 覆盖评分服务器超时）。

  • prebuilt_env_uri –

    由 mlflow.pyfunc.build_model_env API 创建的预构建环境存档文件的路径。此参数只能在 Databricks Serverless notebook REPL、Databricks Shared cluster notebook REPL 和 Databricks Connect 客户端环境中使用。路径可以是本地文件路径或 DBFS 路径，如 ‘dbfs:/Volumes/…’，在这种情况下，MLflow 会自动将其下载到本地临时目录，“MLFLOW_MODEL_ENV_DOWNLOADING_TEMP_DIR”环境变量可用于指定要使用的临时目录。

    如果设置了此参数，则 env_manger 参数不得设置。

  • model_config – 加载模型时设置的模型配置。有关详细信息，请参阅 mlflow.pyfunc.load_model API 中的 'model_config' 参数。

返回

Spark UDF，它将模型 的 predict 方法应用于数据，并返回由 result_type 指定的类型，默认为 double。

mlflow.pyfunc.update_signature_for_type_hint_from_example(input_example: Any, signature: mlflow.models.signature.ModelSignature)[source]
mlflow.pyfunc.get_default_pip_requirements()[source]
返回

MLflow Models 由此格式生成的默认 pip 要求列表。调用 save_model() 和 log_model() 会生成一个 pip 环境，该环境至少包含这些要求。

mlflow.pyfunc.get_default_conda_env()[source]
返回

MLflow Models 由调用 save_model() 和 log_model() 生成的默认 Conda 环境，当提供 PythonModel 的用户定义子类时。

class mlflow.pyfunc.PythonModelContext[source]

一个工件集合，PythonModel 可以使用它来进行推理。PythonModelContext 对象是由 save_model() 和 log_model() 持久化方法隐式创建的，使用的是这些方法 artifacts 参数指定的内容。

property artifacts

一个包含 <name, artifact_path> 条目的字典，其中 artifact_path 是到工件的绝对文件系统路径。

property model_config

一个包含 <config, value> 条目的字典，其中 config 是模型配置键的名称，value 是给定配置的值。

class mlflow.pyfunc.PythonModel[source]

表示一个通用的 Python 模型，该模型可以评估输入并产生 API 兼容的输出。通过继承 PythonModel，用户可以创建具有“python_function”（“pyfunc”）格式的自定义 MLflow 模型，利用自定义推理逻辑和工件依赖项。

load_context(context)[source]

从指定的 PythonModelContext 加载工件，这些工件可供 predict() 在评估输入时使用。加载 load_model() 的 MLflow 模型时，将在构造 PythonModel 后立即调用此方法。

在调用 predict() 时，相同的 PythonModelContext 也将可用，但在模型加载时重写此方法并从上下文加载工件可能会更有效。

参数

context – 一个 PythonModelContext 实例，其中包含模型可用于执行推理的工件。

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

评估一个 pyfunc 兼容的输入并产生一个 pyfunc 兼容的输出。有关 pyfunc 输入/输出 API 的更多信息，请参阅 推理 API。

参数

  • context – 一个 PythonModelContext 实例，其中包含模型可用于执行推理的工件。

  • model_input – 模型要评估的 pyfunc 兼容输入。

  • params – 传递给模型的额外参数，用于推理。

提示

自 MLflow 2.20.0 起，如果 context 参数未使用，则可以从 predict 函数签名中移除。 def predict(self, model_input, params=None) 是有效的。

predict_stream(context, model_input, params: Optional[dict[str, typing.Any]] = None)[source]

评估一个 pyfunc 兼容的输入并产生一个输出迭代器。有关 pyfunc 输入 API 的更多信息，请参阅 推理 API。

参数

  • context – 一个 PythonModelContext 实例，其中包含模型可用于执行推理的工件。

  • model_input – 模型要评估的 pyfunc 兼容输入。

  • params – 传递给模型的额外参数，用于推理。

提示

自 MLflow 2.20.0 起，如果 context 参数未使用，则可以从 predict_stream 函数签名中移除。 def predict_stream(self, model_input, params=None) 是有效的。

property predict_type_hints: mlflow.models.signature._TypeHints

获取 predict 函数签名中类型提示的内部方法。

class mlflow.pyfunc.ChatModel(*args, **kwargs)[source]

警告

mlflow.pyfunc.model.ChatModel 自 3.0.0 起已弃用。此方法将在未来版本中移除。请使用 ResponsesAgent。

提示

自 MLflow 3.0.0 起，我们建议使用 ResponsesAgent 而不是 ChatModel，除非您需要与 OpenAI ChatCompletion API 严格兼容。

一个 PythonModel 的子类，它使得实现与流行 LLM 聊天 API 兼容的模型更加方便。通过继承 ChatModel，用户可以创建具有比通用 PythonModel API 更方便的 predict() 方法的 MLflow 模型。ChatModels 会自动定义输入/输出签名和输入示例，因此在调用 mlflow.pyfunc.save_model() 时无需手动指定这些值。

有关 ChatModel API 预期参数和输出的详细信息，请参阅下方 predict() 方法的文档。

ChatModel

PythonModel

何时使用

当您想开发和部署一个与 OpenAI 规范兼容的**标准**聊天模式的对话模型时使用。

当您希望对模型的接口进行**完全控制**或自定义模型行为的每个方面时使用。

接口

**固定**为 OpenAI 的聊天模式。

**完全控制**模型的输入和输出模式。

设置
**快速**。开箱即用，适用于对话应用程序，具有预定义的

模型签名和输入示例。

**自定义**。您需要自己定义模型签名或输入示例。

复杂性

**低**。标准化接口简化了模型部署和集成。
**高**。部署和集成自定义 PythonModel 可能并不简单。

例如，模型需要处理 Pandas DataFrame，因为 MLflow 在将输入数据传递给 PythonModel 之前会将其转换为 DataFrame。
abstract predict(context, messages: list[mlflow.types.llm.ChatMessage], params: mlflow.types.llm.ChatParams) → mlflow.types.llm.ChatCompletionResponse[source]

评估聊天输入并产生聊天输出。重写此函数以实现真实的流式预测。

参数

  • context – 一个 PythonModelContext 实例，其中包含模型可用于执行推理的工件。

  • messages (List[ChatMessage]) – 一个 ChatMessage 对象列表，表示聊天历史记录。

  • params (ChatParams) – 一个 ChatParams 对象，包含在推理期间修改模型行为的各种参数。

提示

自 MLflow 2.20.0 起，如果 context 参数未使用，则可以从 predict 函数签名中移除。 def predict(self, messages: list[ChatMessage], params: ChatParams) 是有效的。

返回

一个 ChatCompletionResponse 对象，包含模型的响应以及其他元数据。

predict_stream(context, messages: list[mlflow.types.llm.ChatMessage], params: mlflow.types.llm.ChatParams) → Generator[mlflow.types.llm.ChatCompletionChunk, None, None][source]

评估聊天输入并产生聊天输出。重写此函数以实现真实的流式预测。

参数

  • context – 一个 PythonModelContext 实例，其中包含模型可用于执行推理的工件。

  • messages (List[ChatMessage]) – 一个 ChatMessage 对象列表，表示聊天历史记录。

  • params (ChatParams) – 一个 ChatParams 对象，包含在推理期间修改模型行为的各种参数。

提示

自 MLflow 2.20.0 起，如果 context 参数未使用，则可以从 predict_stream 函数签名中移除。 def predict_stream(self, messages: list[ChatMessage], params: ChatParams) 是有效的。

返回

一个 ChatCompletionChunk 对象的生成器，包含模型的响应以及其他元数据。

class mlflow.pyfunc.ChatAgent[source]

提示

自 MLflow 3.0.0 起，我们建议使用 ResponsesAgent 而不是 ChatAgent。

ChatAgent 接口是什么？

ChatAgent 接口是一个聊天模式规范，专为编写对话代理而设计。ChatAgent 允许您的代理执行以下操作：

  • 返回多条消息

  • 为工具调用代理返回中间步骤

  • 确认工具调用

  • 支持多代理场景

编写代理时应始终使用 ChatAgent。即使对于简单的聊天模型（例如，提示工程的 LLM），我们也建议使用 ChatAgent 而不是 ChatModel，以便您将来能够灵活地支持更多代理功能。

ChatAgentRequest 模式与 OpenAI ChatCompletion 模式相似，但并非严格兼容。ChatAgent 增加了附加功能，并与 OpenAI ChatCompletionRequest 在以下方面有所不同：

  • 为每个输入/输出消息添加一个可选的 attachments 属性，用于工具和内部代理调用，以便它们可以返回可视化和进度指示器等附加输出

  • 添加一个 context 属性，其中包含 conversation_id 和 user_id 属性，以启用根据查询代理的用户来修改代理的行为

  • 添加 custom_inputs 属性，这是一个任意的 dict[str, Any]，用于传递任何附加信息以修改代理的行为

ChatAgentResponse 模式与 ChatCompletionResponse 模式在以下方面有所不同：

  • 添加 custom_outputs 键，这是一个任意的 dict[str, Any]，用于返回任何附加信息

  • 允许输出中有多条消息，以改进最终答案的内部工具调用和代理间通信的显示和评估。

以下是详细说明工具调用的 ChatAgentResponse 示例

{
    "messages": [
        {
            "role": "assistant",
            "content": "",
            "id": "run-04b46401-c569-4a4a-933e-62e38d8f9647-0",
            "tool_calls": [
                {
                    "id": "call_15ca4fcc-ffa1-419a-8748-3bea34b9c043",
                    "type": "function",
                    "function": {
                        "name": "generate_random_ints",
                        "arguments": '{"min": 1, "max": 100, "size": 5}',
                    },
                }
            ],
        },
        {
            "role": "tool",
            "content": '{"content": "Generated array of 2 random ints in [1, 100]."',
            "name": "generate_random_ints",
            "id": "call_15ca4fcc-ffa1-419a-8748-3bea34b9c043",
            "tool_call_id": "call_15ca4fcc-ffa1-419a-8748-3bea34b9c043",
        },
        {
            "role": "assistant",
            "content": "The new set of generated random numbers are: 93, 51, 12, 7, and 25",
            "name": "llm",
            "id": "run-70c7c738-739f-4ecd-ad18-0ae232df24e8-0",
        },
    ],
    "custom_outputs": {"random_nums": [93, 51, 12, 7, 25]},
}

使用 ChatAgent 流式输出代理

有关如何流式输出代理的详细信息，请阅读 ChatAgent.predict_stream 的文档字符串。

编写 ChatAgent

使用 ChatAgent 接口编写代理是一种框架无关的方式，用于创建具有标准化接口的模型，该接口可以使用 MLflow pyfunc 格式进行记录，可在客户端之间重用，并为服务工作负载做好准备。

要编写自己的代理，请继承 ChatAgent，实现 predict 方法（可选地实现 predict_stream 方法）来定义代理的非流式和流式行为。您可以使用任何代理编写框架 - 唯一硬性要求是实现 predict 接口。

def predict(
    self,
    messages: list[ChatAgentMessage],
    context: Optional[ChatContext] = None,
    custom_inputs: Optional[dict[str, Any]] = None,
) -> ChatAgentResponse: ...

除了调用与类型提示匹配的输入/输出的 predict 和 predict_stream 方法外，您还可以传递一个与 ChatAgentRequest 模式匹配的单个输入字典，以便于测试。

chat_agent = MyChatAgent()
chat_agent.predict(
    {
        "messages": [{"role": "user", "content": "What is 10 + 10?"}],
        "context": {"conversation_id": "123", "user_id": "456"},
    }
)

有关 ChatAgentState 文档字符串中 LangGraph 代理的 predict 和 predict_stream 的示例实现，请参阅。

记录 ChatAgent

由于 LLM 框架的格局不断发展，并且并非所有格式都能被 MLflow 原生支持，因此我们建议使用 代码生成模型（Models-from-Code）记录方法。

with mlflow.start_run():
    logged_agent_info = mlflow.pyfunc.log_model(
        name="agent",
        python_model=os.path.join(os.getcwd(), "agent"),
        # Add serving endpoints, tools, and vector search indexes here
        resources=[],
    )

记录模型后，您可以使用与 ChatAgentRequest 模式匹配的单个字典来查询模型。在底层，它将被转换为您的 predict 和 predict_stream 方法所期望的 Python 对象。

loaded_model = mlflow.pyfunc.load_model(tmp_path)
loaded_model.predict(
    {
        "messages": [{"role": "user", "content": "What is 10 + 10?"}],
        "context": {"conversation_id": "123", "user_id": "456"},
    }
)

为了使记录 ChatAgent 模型尽可能简单，MLflow 内置了以下功能：

  • 自动模型签名推断
  • Metadata

    • 将自动将 {"task": "agent/v2/chat"} 追加到您在记录模型时可能传入的任何元数据中

  • 输入示例

    • 提供输入示例是可选的，默认会提供 mlflow.types.agent.CHAT_AGENT_INPUT_EXAMPLE

    • 如果您提供输入示例，请确保它是符合 ChatAgentRequest 模式的字典

    • input_example = {
    "messages": [{"role": "user", "content": "What is MLflow?"}],
    "context": {"conversation_id": "123", "user_id": "456"},
}

从 ChatModel 迁移到 ChatAgent

要将现有的 ChatModel 转换为接受 List[ChatMessage] 和 ChatParams 并输出 ChatCompletionResponse 的 ChatModel，请执行以下操作：

  • 继承 ChatAgent 而不是 ChatModel

  • 将 ChatModel 的 load_context 实现中的任何功能移至新 ChatAgent 的 __init__ 方法中。

  • 在将模型的输入转换为字典时，请使用 .model_dump() 而不是 .to_dict()。例如，使用 [msg.model_dump() for msg in messages] 而不是 [msg.to_dict() for msg in messages]。

  • 返回 ChatAgentResponse 而不是 ChatCompletionResponse。

例如，我们可以将“Chat Model Intro”中的 ChatModel 转换为 ChatAgent。

class SimpleOllamaModel(ChatModel):
    def __init__(self):
        self.model_name = "llama3.2:1b"
        self.client = None

    def load_context(self, context):
        self.client = ollama.Client()

    def predict(
        self, context, messages: list[ChatMessage], params: ChatParams = None
    ) -> ChatCompletionResponse:
        ollama_messages = [msg.to_dict() for msg in messages]
        response = self.client.chat(model=self.model_name, messages=ollama_messages)
        return ChatCompletionResponse(
            choices=[{"index": 0, "message": response["message"]}],
            model=self.model_name,
        )

class SimpleOllamaModel(ChatAgent):
    def __init__(self):
        self.model_name = "llama3.2:1b"
        self.client = None
        self.client = ollama.Client()

    def predict(
        self,
        messages: list[ChatAgentMessage],
        context: Optional[ChatContext] = None,
        custom_inputs: Optional[dict[str, Any]] = None,
    ) -> ChatAgentResponse:
        ollama_messages = self._convert_messages_to_dict(messages)
        response = self.client.chat(model=self.model_name, messages=ollama_messages)
        return ChatAgentResponse(**{"messages": [response["message"]]})

ChatAgent 连接器

MLflow 提供了便捷的 API，用于将流行的创作框架中编写的代理包装到 ChatAgent 中。有关 LangGraph 的示例，请参阅 ChatAgentState 的文档字符串。

abstract predict(messages: list[mlflow.types.agent.ChatAgentMessage], context: Optional[mlflow.types.agent.ChatContext] = None, custom_inputs: Optional[dict[str, typing.Any]] = None) → mlflow.types.agent.ChatAgentResponse[source]

给定一个 ChatAgent 输入，返回一个 ChatAgent 输出。除了使用与类型提示匹配的输入调用 predict 外，您还可以传递一个与 ChatAgentRequest schema 匹配的单个输入字典，以便于测试。

chat_agent = ChatAgent()
chat_agent.predict(
    {
        "messages": [{"role": "user", "content": "What is 10 + 10?"}],
        "context": {"conversation_id": "123", "user_id": "456"},
    }
)
参数

  • messages (List[ChatAgentMessage]) – 一个 ChatAgentMessage 对象列表，表示聊天记录。

  • context (ChatContext) – 一个包含 conversation_id 和 user_id 的 ChatContext 对象。**可选**，默认为 None。

  • custom_inputs (Dict[str, Any]) – 一个可选参数，用于向模型提供任意的附加输入。字典值必须是 JSON 可序列化的。**可选**，默认为 None。

返回

一个 ChatAgentResponse 对象，包含模型的响应以及其他元数据。

predict_stream(messages: list[mlflow.types.agent.ChatAgentMessage], context: Optional[mlflow.types.agent.ChatContext] = None, custom_inputs: Optional[dict[str, typing.Any]] = None) → Generator[mlflow.types.agent.ChatAgentChunk, None, None][source]

给定一个 ChatAgent 输入，返回一个包含流式 ChatAgent 输出块的生成器。除了使用与类型提示匹配的输入调用 predict_stream 外，您还可以传递一个与 ChatAgentRequest schema 匹配的单个输入字典，以便于测试。

chat_agent = ChatAgent()
for event in chat_agent.predict_stream(
    {
        "messages": [{"role": "user", "content": "What is 10 + 10?"}],
        "context": {"conversation_id": "123", "user_id": "456"},
    }
):
    print(event)

要支持代理的流式传输，请在 ChatAgent 的子类中重写此方法。在实现 predict_stream 时，请牢记以下要求：

  • 确保您的实现符合 predict_stream 的类型签名。例如，流式消息必须是 ChatAgentChunk 类型，其中每个块包含来自单个响应消息的部分输出。

  • 特定响应中的块最多只能包含一个 custom_outputs 键。

  • 包含单个响应消息部分内容的块必须具有相同的 id。消息的内容字段和 ChatAgentChunk 的使用情况统计信息应由消耗方客户端聚合。请参阅下面的示例。

{"delta": {"role": "assistant", "content": "Born", "id": "123"}}
{"delta": {"role": "assistant", "content": " in", "id": "123"}}
{"delta": {"role": "assistant", "content": " data", "id": "123"}}
参数

  • messages (List[ChatAgentMessage]) – 一个 ChatAgentMessage 对象列表，表示聊天记录。

  • context (ChatContext) – 一个包含 conversation_id 和 user_id 的 ChatContext 对象。**可选**，默认为 None。

  • custom_inputs (Dict[str, Any]) – 一个可选参数，用于向模型提供任意的附加输入。字典值必须是 JSON 可序列化的。**可选**，默认为 None。

返回

一个包含模型响应以及其他元数据的 ChatAgentChunk 对象的生成器。

class mlflow.pyfunc.ResponsesAgent[source]

创建 ResponsesAgent 模型的基础类。它可以作为任何代理框架的包装器，以创建可以部署到 MLflow 的代理模型。它有一些辅助方法，可以帮助创建作为 ResponsesAgentResponse 或 ResponsesAgentStreamEvent 的一部分的输出项。

有关更多详细信息，请参阅 https://mlflow.org.cn/docs/latest/genai/flavors/responses-agent-intro。

static create_annotation_added(item_id: str, annotation: dict[str, typing.Any], annotation_index: int | None = 0) → dict[str, typing.Any][source]
static create_function_call_item(id: str, call_id: str, name: str, arguments: str) → dict[str, typing.Any][source]

用于创建符合函数调用项 schema 的字典的辅助方法。

有关更多信息，请参阅 https://mlflow.org.cn/docs/latest/genai/flavors/responses-agent-intro#creating-agent-output。

参数

  • id (str) – 输出项的 id。

  • call_id (str) – 函数调用的 id。

  • name (str) – 要调用的函数的名称。

  • arguments (str) – 要传递给函数的参数。

static create_function_call_output_item(call_id: str, output: str) → dict[str, typing.Any][source]

用于创建符合函数调用输出项 schema 的字典的辅助方法。

有关更多信息，请参阅 https://mlflow.org.cn/docs/latest/genai/flavors/responses-agent-intro#creating-agent-output。

参数

  • call_id (str) – 函数调用的 id。

  • output (str) – 函数调用的输出。

static create_reasoning_item(id: str, reasoning_text: str) → dict[str, typing.Any][source]

用于创建符合推理项 schema 的字典的辅助方法。

有关更多信息，请参阅 https://www.mlflow.org/docs/latest/llms/responses-agent-intro/#creating-agent-output。

static create_text_delta(delta: str, item_id: str) → dict[str, typing.Any][source]

用于创建符合用于流式传输的文本 delta schema 的字典的辅助方法。

有关更多信息，请参阅 https://mlflow.org.cn/docs/latest/genai/flavors/responses-agent-intro#streaming-agent-output。

static create_text_output_item(text: str, id: str, annotations: Optional[list[dict[str, typing.Any]]] = None) → dict[str, typing.Any][source]

用于创建符合文本输出项 schema 的字典的辅助方法。

有关更多信息，请参阅 https://mlflow.org.cn/docs/latest/genai/flavors/responses-agent-intro#creating-agent-output。

参数

  • text (str) – 要输出的文本。

  • id (str) – 输出项的 id。

  • annotations (Optional[list[dict]]) – 输出项的注释。

static output_to_responses_items_stream(chunks: Iterator[dict[str, typing.Any]], aggregator: Optional[list[dict[str, typing.Any]]] = None) → Generator[mlflow.types.responses.ResponsesAgentStreamEvent, None, None][source]

对于流式传输，从各种消息格式的字典转换为 Responses 输出项，返回一个 ResponsesAgentStreamEvent 对象的生成器。

如果提供了 aggregator，它将被扩展为包含聚合后的输出项字典。

目前，仅处理 Chat Completion 块的流。

abstract predict(request: mlflow.types.responses.ResponsesAgentRequest) → mlflow.types.responses.ResponsesAgentResponse[source]

给定一个 ResponsesAgentRequest，返回一个 ResponsesAgentResponse。

您可以在 https://mlflow.org.cn/docs/latest/genai/flavors/responses-agent-intro#simple-chat-example 和 https://mlflow.org.cn/docs/latest/genai/flavors/responses-agent-intro#tool-calling-example 中看到示例实现。

predict_stream(request: mlflow.types.responses.ResponsesAgentRequest) → Generator[mlflow.types.responses.ResponsesAgentStreamEvent, None, None][source]

给定一个 ResponsesAgentRequest，返回一个 ResponsesAgentStreamEvent 对象的生成器。

有关更多详细信息，请参阅 https://mlflow.org.cn/docs/latest/genai/flavors/responses-agent-intro#streaming-agent-output。

您可以在 https://mlflow.org.cn/docs/latest/genai/flavors/responses-agent-intro#simple-chat-example 和 https://mlflow.org.cn/docs/latest/genai/flavors/responses-agent-intro#tool-calling-example 中看到示例实现。

static prep_msgs_for_cc_llm(responses_input: list[dict[str, typing.Any] | mlflow.types.responses_helpers.Message | mlflow.types.responses_helpers.OutputItem]) → list[dict[str, typing.Any]][source]

从 Responses 输入项转换为 ChatCompletion 字典。

static responses_agent_output_reducer(chunks: list[mlflow.types.responses.ResponsesAgentStreamEvent | dict[str, typing.Any]])[source]