mlflow.pyfunc

python_function 模型风格作为 MLflow Python 模型的默认模型接口。任何 MLflow Python 模型都应能以 python_function 模型的形式加载。

此外，mlflow.pyfunc 模块为 Python 模型定义了一种通用的 filesystem format，并提供了用于将模型保存到该格式以及从该格式加载模型的实用工具。该格式是自包含的，意味着它包含任何人加载并使用它所需的所有必要信息。依赖项要么与模型一起直接存储，要么通过 Conda 环境引用。

该 mlflow.pyfunc 模块还定义了用于创建自定义 pyfunc 模型的实用工具，适用于可能不被 MLflow 原生包含的框架和推理逻辑。请参见 Models From Code for Custom Models。

推理 API

Python 函数模型以 PyFuncModel 的实例形式加载，它是围绕模型实现和模型元数据（MLmodel 文件）的 MLflow 包装器。您可以通过调用 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]

所有 PyFunc 模型都将支持 pandas.DataFrame 作为输入，且 PyFunc 深度学习模型还将以 Dict[str, numpy.ndarray]（命名张量）和 numpy.ndarrays（未命名张量）的形式支持张量输入。

下面是一些受支持的推理类型示例，假设我们已经加载了正确的 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)

文件系统格式

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)

目录结构可能包含可被 MLmodel 配置引用的其他内容。

MLModel 配置

一个 Python 模型在其根目录包含一个 MLmodel 文件，格式为 python_function，具有以下参数：

  • loader_module [required]:

    可以加载模型的 Python 模块。预期以模块标识符的形式提供，例如 mlflow.sklearn，它将使用 importlib.import_module 导入。导入的模块必须包含具有以下签名的函数：

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

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

    包含随此模型打包代码的目录的相对路径。此目录内的所有文件和子目录在导入模型加载器之前都会被加入到 Python 路径中。

  • data [optional]:

    相对于包含模型数据的文件或目录的路径。该路径将传递给模型加载器。

  • env [optional]:

    导出的 Conda 环境的相对路径。如果存在，应在运行模型之前激活此环境。

  • 可选的：任何在解释以 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 引入了“models from code”功能，该功能通过脚本序列化大大简化了自定义模型的序列化和部署过程。强烈建议将自定义模型实现迁移到这一新范式，以避免使用 cloudpickle 进行序列化时的局限性和复杂性。您可以在 Models From Code Guide 中了解有关“models from code”的更多信息。

下面的部分说明了为自定义 Pyfunc 模型使用旧版序列化器的过程。通过代码创建的模型在记录模型时会提供更简单的体验。

创建自定义 Pyfunc 模型

MLflow 的持久化模块提供便捷函数，用于在多种机器学习框架中创建带有 pyfunc 风格的模型（scikit-learn、Keras、Pytorch 等）；但它们并不能涵盖所有使用场景。例如，您可能希望使用 MLflow 原生不直接支持的框架来创建带有 pyfunc 风格的 MLflow 模型。或者，您可能希望构建一个在评估查询时执行自定义逻辑（例如预处理和后处理例程）的 MLflow 模型。因此，mlflow.pyfunc 提供了用于从任意代码和模型数据创建 pyfunc 模型的实用工具。

save_model() 和 log_model() 方法旨在支持多种工作流，以创建将自定义推理逻辑及该逻辑可能需要的工件纳入的自定义 pyfunc 模型。

一个工件是一个文件或目录，例如序列化的模型或 CSV 文件。例如，序列化的 TensorFlow 图就是一个工件。MLflow 模型目录也是一个工件。

工作流

save_model() 和 log_model() 支持以下工作流程：

  1. 以编程方式定义一个新的 MLflow 模型，包括其属性和工件。

    给定一组工件 URI，save_model() 和 log_model() 可以自动从其 URI 下载工件并创建一个 MLflow 模型目录。

    在这种情况下，您必须定义一个继承自 PythonModel 的 Python 类，该类需要定义 predict()，可选地定义 load_context()。该类的一个实例通过 python_model 参数指定；作为一个 Python 类，它会被自动序列化和反序列化，包括其所有属性。

  2. 将预先存在的数据解释为 MLflow 模型。

    如果您已经有一个包含模型数据的目录，save_model() 和 log_model() 可以将这些数据导入为 MLflow 模型。data_path 参数 指定包含模型数据的目录在本地文件系统中的路径。

    在这种情况下，您必须提供一个 Python 模块，称为 加载器模块。该加载器模块定义了一个 _load_pyfunc() 方法，执行以下任务：

    • 从指定的 data_path 加载数据。例如，这一过程可能包括反序列化用 pickle 序列化的 Python 对象或模型，或解析 CSV 文件。

    • 构造并返回一个与 pyfunc 兼容的模型包装器。与第一个用例相同，该包装器必须定义一个 predict() 方法，用于评估查询。 predict() 必须遵守 Inference API。

    参数 loader_module 指定了您加载器模块的名称。

    要查看示例加载器模块的实现，请参阅 loader module implementation in mlflow.sklearn.

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

我们认为第一种工作流更符合用户使用习惯，并且通常推荐，原因如下：

  • 它会自动解析并收集指定的模型产物。

  • 它会自动序列化和反序列化 python_model 实例及其所有属性，从而减少加载模型所需的用户逻辑量

  • 您可以使用在 __main__ 作用域中定义的逻辑来创建模型。这允许在交互式环境中（例如笔记本和 Python REPL）构建自定义模型。

出于以下原因，您可能更倾向于第二种较低级别的工作流：

  • 推理逻辑始终以代码形式保存，而不是作为 Python 对象。这使得以后更容易检查和修改逻辑。

  • 如果您已将所有模型数据收集到单个位置，第二种工作流允许直接以 MLflow 格式保存，而无需枚举组成的工件。

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

在创建自定义 PyFunc 模型时，你可以在两种不同的接口之间进行选择：基于函数的模型和基于类的模型。简而言之，基于函数的模型仅仅是一个不接受额外参数的 Python 函数。另一方面，基于类的模型是 PythonModel 的子类，支持若干必需和可选的方法。如果你的用例简单并且可以在单个 predict 函数内完成，推荐使用基于函数的方法。如果你需要更强的能力，例如自定义序列化、自定义数据处理或重写其他方法，则应使用基于类的实现。

在查看示例代码之前，重要的是要注意这两种方法都是通过cloudpickle进行序列化的。cloudpickle 可以序列化 Python 函数、lambda 函数，以及在其他函数内部本地定义的类和函数。这使得 cloudpickle 在并行和分布式计算中尤其有用，因为代码对象需要通过网络发送到远程工作节点上执行，这也是 MLflow 的一种常见部署范式。

话虽如此，cloudpickle 有一些限制。

  • 环境依赖：cloudpickle 无法捕获完整的执行环境，因此在 MLflow 中我们必须传入 pip_requirements、extra_pip_requirements 或一个 input_example，后者用于推断环境依赖。更多信息请参阅 the model dependency docs。

  • 对象支持: cloudpickle 不会序列化 Python 数据模型之外的对象。一些相关示例包括原始文件和数据库连接。如果您的程序依赖这些，请务必将引用这些对象的方法与您的模型一起记录。

函数式模型

如果你想序列化一个简单的 python 函数而不包含额外的依赖方法，你可以通过关键字参数 python_model 简单地记录一个 predict 方法。

注意

基于函数的模型仅支持具有单一输入参数的函数。如果您想传递更多参数或额外的推理参数，请使用下面的基于类的模型。

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)

基于类的模型

如果您想序列化更复杂的对象，例如负责预处理、复杂预测逻辑或自定义序列化的类，则应继承 PythonModel 类。MLflow 有关于构建自定义 PyFunc 模型的教程，如 here 所示，因此我们不会重复这些内容，在本示例中我们将重新创建上述功能以突出差异。请注意，此 PythonModel 实现过于复杂，对于这种简单情况，我们建议改用基于函数的 Model。

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)}")

此实现与上面的基于函数的实现之间的主要区别在于 predict 方法被封装在一个类中，具有 self 参数，并且具有默认值为 None 的 params 参数。注意，基于函数的模型不支持额外的 params。

总之，当你有一个需要序列化的简单函数时，使用基于函数的模型。 如果你需要更强大的功能，请使用基于类的模型。

class mlflow.pyfunc.EnvType[source]

基类： 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]

基类： object

MLflow ‘python function’ 模型。

作为模型实现和元数据的包装器。此类并非用于直接构造。相反，此类的实例由load_model()构造并返回。

model_impl 可以是实现 Pyfunc interface 的任意 Python 对象，并且由调用模型的 loader_module 返回。

model_meta 包含从 MLmodel 文件加载的模型元数据。

get_raw_model()[source]

如果模型包装器实现了get_raw_model函数，则获取底层原始模型。

property input_example: Optional[Any]

模型保存时提供的输入示例。

property loader_module

模型的 flavor 配置

property metadata: mlflow.models.model.Model

模型元数据。

property model_config

模型的 flavor 配置

property model_id: str | None

模型的 model ID。

Returns

该模型的模型 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]

解包底层的 Python 模型对象。

此方法对于访问自定义模型函数很有用，同时仍然能够通过 predict() 方法利用 MLflow 设计的工作流程。

Returns

底层封装的模型对象

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]

将一个pyfunc规范添加到模型配置中。

定义 pyfunc 的配置模式。调用者可以使用它从现有的目录结构创建有效的 pyfunc 模型 flavor。例如，其他模型 flavor 可以使用它来指定如何将其输出作为 pyfunc 使用。

注意

所有路径均相对于导出的模型根目录。

Parameters

  • model – 现有模型。

  • loader_module – 用于加载模型的模块。

  • data – 模型数据的路径。

  • code – 代码依赖的路径。

  • conda_env – Conda 环境.

  • python_env – Python 环境。

  • model_config –

    要应用于模型的配置。此配置在模型加载时可用。

    注意

    实验性：此参数可能在未来版本中更改或被移除，恕不另行通知。

  • model_code_path – 模型代码的路径。

  • kwargs – 额外的键值对，包含在 pyfunc flavor 规范中。值必须可被 YAML 序列化。

Returns

已更新模型配置。

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

预构建模型 Python 环境并生成一个归档文件，保存到提供的 save_path。

Typical usages:

  • 在 Databricks Runtime 中预先构建模型的环境，然后下载预构建的 Python 环境归档文件。该预构建环境归档随后可以在 mlflow.pyfunc.spark_udf 中使用，以便在使用 Databricks Connect 远程连接到 Databricks 环境执行代码时进行远程推理执行。

注意

build_model_env API 仅在 Databricks 运行时中执行时有效，目的是捕捉在使用 DBConnect 进行远程代码执行时所需的执行环境。该环境归档旨在在 Databricks 运行时或 Databricks Connect 客户端中使用 mlflow.pyfunc.spark_udf 执行远程执行时使用，没有其他用途。预构建的环境归档文件不能跨不同的 Databricks 运行时版本或不同平台的机器使用。因此，如果您连接到在 Databricks 上运行不同运行时版本的不同集群，则需要在该集群的 notebook 中执行此 API，并将生成的归档检索到本地机器。每个环境快照都与模型、远程 Databricks 集群的运行时版本以及 UDF 执行环境的规范唯一对应。在 mlflow.pyfunc.spark_udf 中使用预构建环境时，MLflow 会验证 spark UDF 沙箱环境是否满足预构建环境的要求，如果存在兼容性问题将引发异常。如果发生此类情况，只需在您尝试附加的集群中重新运行此 API。

Example 
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")
Parameters

  • model_uri – 用于构建 python 环境的模型的 URI。

  • save_path – 用于保存预构建模型环境归档文件的目录路径。该路径可以是本地目录路径，或挂载的 DBFS 路径（例如 ‘/dbfs/…’）或挂载的 UC 卷路径（例如 ‘/Volumes/…’）。

  • env_manager – 用于为模型推理创建 python 环境的环境管理器，取值可以是 ‘virtualenv’ 或 ‘uv’，默认值为 ‘virtualenv’。

Returns

返回包含 Python 环境数据的归档文件的路径。

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

下载模型的依赖项并返回 requirements.txt 或 conda.yaml 文件的路径。

警告

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

Parameters

  • model_uri – 用于获取依赖项的模型的 uri。

  • format – 返回的依赖文件的格式。如果指定了 "pip" 格式，则返回 pip requirements.txt 文件的路径。如果指定了 "conda" 格式，则返回 "conda.yaml" 文件的路径。如果指定了 "pip" 格式，但模型未随带 requirements.txt 文件，则改为从模型的 conda.yaml 文件中提取 pip 部分，并忽略任何额外的 conda 依赖。默认值为 "pip"。

Returns

本地文件系统路径，指向一个 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 函数格式存储的模型。

Parameters

  • 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

    有关受支持的 URI 方案的更多信息，请参阅 Referencing Artifacts。

  • suppress_warnings – 如果 True，与模型加载过程相关的非致命警告消息将被抑制。如果 False，这些警告消息将被发出。

  • dst_path – 要将模型工件下载到的本地文件系统路径。该目录必须已存在。如果未指定，将创建一个本地输出路径。

  • model_config –

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

    注意

    实验性：该参数可能在将来的版本中更改或被移除，且不会另行通知。

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

警告

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

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

Parameters

  • 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

    有关受支持的 URI 方案的更多信息，请参阅 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 工件。

有关此方法所支持的工作流的信息，请参见 Workflows for creating custom pyfunc models 和 Which workflow is right for my use case?。您不能同时指定第二种工作流的参数：loader_module、data_path，以及第一种工作流的参数：python_model、artifacts。

Parameters

  • artifact_path – 已弃用。请改用 name。

  • loader_module –

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

    • The MLflow library.

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

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

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

  • code_paths –

    本地文件系统中指向 Python 文件依赖项的路径列表（或包含文件依赖项的目录）。这些文件在模型加载时会被预先添加到系统路径。若多个文件之间存在导入依赖，则应为给定模型声明为依赖项的文件从一个公共根路径声明相对导入，以避免在加载模型时出现导入错误。

    可以将 code_paths 参数留空，但将 infer_code_paths 设置为 True，让 MLflow 推断模型代码路径。有关详细信息，请参阅 infer_code_paths 参数文档。

    有关 code_paths 功能、推荐使用模式和限制的详细说明，请参阅 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__ 中的类/函数进行 pickle。

    注意

    实验性：该参数可能在将来的版本中更改或被移除，且不会另行通知。

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

    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. The python_model can also be a file path to the PythonModel which defines the model from code artifact rather than serializing the model object. 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_paths 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(
        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"]

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

    Model from code

    Note

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

    # 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",
    )

    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 –

    一个包含 artifact_uri> 条目的字典。远程 artifact URIs 会解析为绝对文件系统路径，生成一个包含 absolute_path> 条目的字典。python_model 可以在 context 参数的 artifacts 属性中引用这些解析后的条目，位于 PythonModel.load_context() 和 PythonModel.predict() 中。

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

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

    如果 None，则不会向模型添加任何 artifacts。

  • registered_model_name – 如果提供，则在 registered_model_name 下创建一个模型版本，如果不存在具有该名称的注册模型，则同时创建一个注册模型。

  • signature –

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

    from mlflow.models import infer_signature

train = df.drop_column("target_label")
predictions = ...  # 计算模型预测
signature = infer_signature(train, predictions)

  • input_example – 一个或多个有效模型输入实例。输入示例用于提示应向模型提供何种数据。它将被转换为一个 Pandas DataFrame，然后使用 Pandas 的 split-oriented 格式序列化为 json，或者转换为一个 numpy array，其中示例将通过将其转换为列表的方式序列化为 json。字节使用 base64 编码。当 signature 参数为 None 时，输入示例用于推断模型签名。

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

  • pip_requirements – 要么是一个可迭代的 pip 需求字符串（例如 ["scikit-learn", "-r requirements.txt", "-c constraints.txt"]）要么是本地文件系统上 pip requirements 文件的字符串路径（例如 "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 文件的字符串路径（例如 "requirements.txt"）。如果提供，该参数描述了附加的 pip 依赖，这些依赖会被追加到基于用户当前软件环境自动生成的默认 pip 依赖集合中。requirements 和 constraints 会被自动解析并分别写入 requirements.txt 和 constraints.txt 文件，并作为模型的一部分存储。依赖项也会被写入模型的 conda 环境（conda.yaml）文件的 pip 部分。

    警告

    以下参数不能同时指定：

    • conda_env

    • pip_requirements

    • extra_pip_requirements

    This example 演示了如何使用 pip_requirements 和 extra_pip_requirements 指定 pip 依赖。

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

  • model_config –

    应用于模型的模型配置。该配置将作为 model_config 属性，位于 context 参数中在 PythonModel.load_context() 和 PythonModel.predict() 中可用。该配置可以作为文件路径传入，或作为键为字符串的 dict。

    注意

    实验性：此参数可能在未来版本中更改或被移除，恕不另行通知。

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

  • resources –

    模型资源的列表，或包含该列表的 resources.yaml 文件

    用于为模型提供服务所需的 resources。

    注意

    实验性：此参数可能在将来的版本中更改或被移除，恕不另行通知。

  • auth_policy –

    指定模型的认证策略，包含两个关键组成部分。

    请注意，只有 auth_policy 或 resources 之一定义。

    • 系统认证策略：用于提供此模型服务所需的资源列表。

    • 用户认证策略：用户应该拥有访问权限的最小范围列表

      以便调用此模型。

    注意

    实验性：此参数可能在未来版本中更改或被移除，恕不另行通知。

  • prompts –

    一个在 MLflow Prompt Registry 注册的 prompt URI 列表，用于与模型关联。 Each prompt URI should be in the form prompt:/<name>/<version>. 在与模型关联之前，这些 prompt 应该先在 MLflow Prompt Registry 中注册。

    这将在模型和 prompt 之间创建相互链接。关联的 prompt 可以在存储于 MLmodel 文件中的模型元数据中看到。通过 Prompt Registry UI，你也可以导航到该模型。

    import mlflow

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

# 在 MLflow Prompt Registry 中注册一个 prompt
mlflow.prompts.register_prompt("my_prompt", prompt_template, description="A simple prompt")

# 使用已注册的 prompt 记录模型
with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        name=MyModel(),
        name="model",
        prompts=["prompt:/my_prompt/1"]
    )

print(model_info.prompts)
# 输出：['prompt:/my_prompt/1']

# 加载 prompt
prompt = mlflow.genai.load_prompt(model_info.prompts[0])

  • name – 模型名称。

  • params – 一个用于与模型一同记录的参数字典。

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

  • model_type – 模型的类型。

  • step – 在该步记录模型输出和指标

  • model_id – 模型的 ID。

Returns

一个 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 模型与自定义推理逻辑和可选的数据依赖项保存到本地文件系统上的路径。

关于该方法支持的工作流的信息，请参见 “workflows for creating custom pyfunc models” 和 “which workflow is right for my use case?”。 注意，针对第二个工作流的参数：loader_module、data_path，以及针对第一个工作流的参数：python_model、artifacts，不能同时指定。

Parameters

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

  • loader_module –

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

    • The MLflow library.

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

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

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

  • code_paths –

    一系列指向本地文件系统中 Python 文件依赖项（或包含文件依赖项的目录）的路径。当模型加载时，这些文件会被预先添加到系统路径。对于为给定模型声明的依赖文件，如果多个文件之间存在导入依赖关系，应从共同的根路径声明相对导入，以避免在加载模型时出现导入错误。

    您可以不设置code_paths参数，但将infer_code_paths设置为True，让 MLflow 推断模型代码路径。有关详细信息，请参阅infer_code_paths参数文档。

    有关code_paths功能、推荐的使用模式和限制的详细说明，请参阅 code_paths usage guide。

  • infer_code_paths –

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

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

    警告

    请确保自定义的 python 模块代码不包含诸如凭证令牌字符串等敏感数据，否则这些内容可能会被包含在自动推断的代码路径文件中并记录到 MLflow artifact 存储库中。

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

    如果某个 python 代码文件作为 python __main__ 模块被加载，则该代码文件无法被推断为代码路径文件。如果你的模型依赖于在 __main__ 模块中定义的类/函数，你应使用 cloudpickle 来 dump 你的模型实例，以便 pickle 在 __main__ 中的类/函数。

    注意

    实验性：此参数可能在未来的版本中在未警告的情况下更改或被删除。

  • 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 – mlflow.models.Model 配置，向其添加 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. The python_model can also be a file path to the PythonModel which defines the model from code artifact rather than serializing the model object. 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_paths 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"]

    Model from code

    Note

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

    # 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",
    )

    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 –

    一个包含 artifact_uri> 条目的字典。远程 artifact URIs 会被解析为绝对文件系统路径，从而生成一个包含 absolute_path> 条目的字典。python_model 可以将这些已解析的条目作为 artifacts 属性引用，该属性位于 context 参数中，出现在 PythonModel.load_context() 和 PythonModel.predict() 中。例如，考虑以下 artifacts 字典：

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

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

    如果 None，不会将任何 artifacts 添加到模型。

  • signature –

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

    from mlflow.models import infer_signature

train = df.drop_column("target_label")
predictions = ...  # 计算模型预测
signature = infer_signature(train, predictions)

  • input_example – 一个或多个有效模型输入实例。输入示例用于提示应向模型提供何种数据。它将被转换为一个 Pandas DataFrame，然后使用 Pandas 的 split-oriented 格式序列化为 json，或者转换为一个 numpy array，其中示例将通过将其转换为列表的方式序列化为 json。字节使用 base64 编码。当 signature 参数为 None 时，输入示例用于推断模型签名。

  • pip_requirements – 要么是一个可迭代的 pip 需求字符串（例如 ["scikit-learn", "-r requirements.txt", "-c constraints.txt"]）要么是本地文件系统上 pip requirements 文件的字符串路径（例如 "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 文件的字符串路径（例如 "requirements.txt"）。如果提供，该参数描述了附加的 pip 依赖，这些依赖会被追加到基于用户当前软件环境自动生成的默认 pip 依赖集合中。requirements 和 constraints 会被自动解析并分别写入 requirements.txt 和 constraints.txt 文件，并作为模型的一部分存储。依赖项也会被写入模型的 conda 环境（conda.yaml）文件的 pip 部分。

    警告

    以下参数不能同时指定：

    • conda_env

    • pip_requirements

    • extra_pip_requirements

    This example 演示了如何使用 pip_requirements 和 extra_pip_requirements 指定 pip 依赖。

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

  • model_config –

    要应用到模型的配置。该配置将在 model_config 属性中通过 context 参数在 PythonModel.load_context() 和 PythonModel.predict() 中可用。该配置可以作为文件路径传入，或作为键为字符串的 dict 传入。

    注意

    实验性：该参数可能在未来的版本中更改或被移除，恕不另行通知。

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

  • resources –

    一组模型资源，或包含以下列表的 resources.yaml 文件

    用于提供模型服务所需的资源。

    注意

    实验性：该参数可能在将来的版本中更改或被移除，

  • auth_policy –

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

    注意，应该只定义 auth_policy 或 resources 中的一个。

    • 系统授权策略：提供此模型服务所需的资源列表。

    • 用户授权策略：用户为了调用该模型应具有访问权限的最小范围列表

      。

    注意

    实验性：此参数可能在将来的版本中被更改或移除，且不另行通知。

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

一个 Spark UDF，可用于调用 Python 函数格式的模型。

传递给 UDF 的参数会作为一个 DataFrame 转发到模型，其中列名为序号（0, 1, …）。在某些版本的 Spark（3.0 及以上）中，也可以将输入包装在 struct 中。在那种情况下，数据将作为一个列名由 struct 定义的 DataFrame 传递（例如，当以 my_udf(struct(‘x’, ‘y’)) 调用时，模型会得到一个具有 2 列 ‘x’ 和 ‘y’ 的 pandas DataFrame）。

如果模型的签名包含带 tensor spec 的输入，您需要传递一个数组类型的列作为对应的 UDF 参数。该列的值必须是一维数组。UDF 会以 ‘C’ 顺序（即使用类 C 的索引顺序读取/写入元素）将列值重塑为所需形状，并将这些值转换为所需的 tensor spec 类型。

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

预测结果被筛选，只保留能够表示为 result_type 的列。如果 result_type 是字符串或字符串数组，则所有预测都会被转换为字符串。如果结果类型不是数组类型，则返回最左侧的第一个匹配类型的列。

注意

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

注意

当使用 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()
Parameters

  • spark – 一个 SparkSession 对象.

  • model_uri –

    使用 mlflow.pyfunc flavor 的 MLflow 模型的位置，采用 URI 格式。例如：

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

    • relative/path/to/local/model

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

    • runs://run-relative/path/to/model

    • models://

    • models://

    • mlflow-artifacts:/path/to/model

    有关受支持的 URI 方案的更多信息，请参见 Referencing Artifacts。

  • result_type –

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

    支持以下几类返回类型：

    • ”int” or pyspark.sql.types.IntegerType：左起第一个能适配到 int32 的整数，否则抛出异常。

    • ”long” or pyspark.sql.types.LongType：左起第一个能适配到 int64 的长整数，否则抛出异常。

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

    • ”float” or pyspark.sql.types.FloatType：左起第一个数值结果被转换为 float32，否则抛出异常。

    • ”double” or pyspark.sql.types.DoubleType：左起第一个数值结果被转换为 double，否则抛出异常。

    • ArrayType(FloatType|DoubleType)：所有数值列被转换为所请求的类型；如果没有数值列则抛出异常。

    • ”string” or pyspark.sql.types.StringType：左起第一个列被转换为 string。

    • ”boolean” or “bool” or 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 创建的预构建 env 归档文件的路径。此参数仅可在 Databricks Serverless notebook REPL、Databricks Shared cluster notebook REPL 和 Databricks Connect 客户端环境中使用。该路径可以是本地文件路径或诸如 ‘dbfs:/Volumes/…’ 的 DBFS 路径，在这种情况下，MLflow 会自动将其下载到本地临时目录，可以通过设置 “MLFLOW_MODEL_ENV_DOWNLOADING_TEMP_DIR” 环境变量来指定要使用的临时目录。

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

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

Returns

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

此 flavor 生成的 MLflow Models 的默认 pip 依赖项列表。调用 save_model() 和 log_model() 会生成一个 pip 环境，该环境至少包含这些依赖项。

mlflow.pyfunc.get_default_conda_env()[source]
Returns

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

class mlflow.pyfunc.PythonModelContext[source]

一组工件，供 PythonModel 在执行推理时使用。 PythonModelContext 对象由持久化方法 save_model() 和 log_model() 隐式地 创建，使用这些方法的 artifacts 参数指定的内容。

property artifacts

一个字典，包含 artifact_path> 条目，其中 artifact_path 是指向该工件的绝对文件系统路径。

property model_config

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

class mlflow.pyfunc.PythonModel[source]

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

load_context(context)[source]

从指定的 PythonModelContext 加载可供 predict() 在评估输入时使用的工件。当使用 load_model() 加载 MLflow 模型时，一旦 PythonModel 被构建，就会调用此方法。

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

Parameters

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

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

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

Parameters

  • 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 的更多信息，请参见 Inference API。

Parameters

  • 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，用户可以创建具有 predict() 方法的 MLflow 模型，该方法比通用的 PythonModel API 更适合聊天任务。ChatModels 会自动定义输入/输出签名和一个输入示例，因此在调用 mlflow.pyfunc.save_model() 时不需要手动指定这些值。

请参阅下面 predict() 方法的文档，了解 ChatModel API 所期望的参数和输出的详细信息。

ChatModel

PythonModel

何时使用

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

当您想对模型的界面拥有完全控制或自定义模型行为的每一个方面时使用。

接口

固定 为 OpenAI 的聊天架构。

完全控制 模型的输入和输出架构。

设置
快速. 开箱即可用于对话应用，具有预定义

模型签名和输入示例。

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

复杂度

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

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

评估聊天输入并生成聊天输出。

Parameters

提示

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

Returns

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

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

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

Parameters

提示

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

Returns

一个针对 ChatCompletionChunk 对象的生成器，包含模型的响应（可能有多个）以及其他元数据。

class mlflow.pyfunc.ChatAgent[source]

提示

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

ChatAgent 智能体接口是什么？

ChatAgent interface 是一种为编写对话智能体而设计的聊天模式规范。ChatAgent 允许你的智能体执行以下操作：

  • 返回多条消息

  • 为调用工具的智能体返回中间步骤

  • 确认工具调用

  • 支持多智能体场景

ChatAgent 在编写智能体时应始终使用。我们还建议使用 ChatAgent 而不是 ChatModel，即使在像简单聊天模型（例如 prompt-engineered LLMs）这样的用例中，也是为了在将来为您提供支持更多智能体功能的灵活性。

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 接口创建智能体是一种与框架无关的方式，用于创建具有标准化接口的模型，该模型可以用 MLflow pyfunc flavor 记录，能够在不同客户端之间重用，并已准备好用于服务工作负载。

要编写你自己的智能体，请继承 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 模式匹配的单个输入 dict，以便于测试。

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

查看 predict 和 predict_stream 在 ChatAgentState 的文档字符串中，用于 LangGraph 智能体的示例实现。

记录 Chat智能体

由于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 内置了以下功能：

  • Automatic Model Signature Inference
  • Metadata

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

  • Input Example

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

    • 如果确实提供输入示例，请确保它是一个符合 ChatAgentRequest 模式的 dict

    • input_example = {
    "messages": [{"role": "user", "content": "什么是 MLflow？"}],
    "context": {"conversation_id": "123", "user_id": "456"},
}

从 ChatModel 迁移到 ChatAgent

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

  • 对 ChatAgent 进行子类化，而不是 ChatModel

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

  • 在将您的模型的输入转换为字典时，使用 .model_dump_compat() 而不是 .to_dict()。例如： [msg.model_dump_compat() 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 提供了便捷的 APIs，用于将用流行的创作框架编写的智能体封装为 ChatAgent。示例包括：

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 模式的单个输入字典，以便于测试。

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

  • messages (List[ChatAgentMessage]) – 表示对话历史的一组ChatAgentMessage对象。

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

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

Returns

一个 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 模式匹配的单个输入 dict，以便于测试。

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。消息的 content 字段和 ChatAgentChunk 的使用统计应由消费端客户端聚合。参见下面的示例。

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

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

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

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

Returns

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

class mlflow.pyfunc.ResponsesAgent[source]

注意

实验性：此类可能在将来的版本中更改或在未发出警告的情况下被移除。

用于创建 ResponsesAgent 模型的基类。它可以作为任何智能体框架的包装器，用于创建可以部署到 MLflow 的智能体模型。包含一些辅助方法，用于创建可作为 ResponsesAgentResponse 或 ResponsesAgentStreamEvent 一部分的输出项。

详情请见 https://www.mlflow.org/docs/latest/llms/responses-agent-intro/。

create_function_call_item(id: str, call_id: str, name: str, arguments: str) → dict[str, typing.Any][source]

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

欲了解更多，请访问 https://www.mlflow.org/docs/latest/llms/responses-agent-intro/#creating-agent-output.

Parameters

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

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

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

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

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

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

欲了解更多，请访问 https://www.mlflow.org/docs/latest/llms/responses-agent-intro/#creating-agent-output.

Parameters

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

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

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

辅助方法，用于创建符合用于流式传输的文本增量（text delta）模式的字典。

阅读更多请访问 https://www.mlflow.org/docs/latest/llms/responses-agent-intro/#streaming-agent-output。

create_text_output_item(text: str, id: str) → dict[str, typing.Any][source]

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

欲了解更多，请访问 https://www.mlflow.org/docs/latest/llms/responses-agent-intro/#creating-agent-output.

Parameters

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

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

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

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

你可以在 https://www.mlflow.org/docs/latest/llms/responses-agent-intro#simple-chat-example 和 https://www.mlflow.org/docs/latest/llms/responses-agent-intro#tool-calling-example 看到示例实现。

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

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

有关更多详情，请参见 https://www.mlflow.org/docs/latest/llms/responses-agent-intro#streaming-agent-output.

您可以在 https://www.mlflow.org/docs/latest/llms/responses-agent-intro#simple-chat-example 和 https://www.mlflow.org/docs/latest/llms/responses-agent-intro#tool-calling-example 查看示例实现。

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