跳到主要内容

MLflow PythonModel 指南

MLflow PythonModel 简介​

mlflow.pyfunc 模块提供了 save_model() 和

log_model() 工具，用于创建带有 python_function 格式的 MLflow Models，其中包含用户指定的代码和*artifact*（文件）依赖项。

MLflow PythonModel 使您能够实现自定义模型逻辑，同时利用 MLflow 的打包和部署功能。

定义 PythonModel 有两种方式：继承 mlflow.pyfunc.PythonModel() 或定义一个可调用对象。本指南将提供一个关于如何定义和使用自定义 PythonModel 的完整 walkthrough。

定义自定义 PythonModel​

选项 1：继承 PythonModel​

mlflow.pyfunc 模块提供了一个 通用的 PythonModel 类，可用于定义您自己的自定义模型。通过继承它，模型可以与其他 MLflow 组件无缝集成。

PythonModel 的方法

  • predict 有效的 PythonModel 必须实现 predict 方法，该方法定义了模型的预测逻辑。当模型使用 mlflow.pyfunc.load_model 加载为 PyFunc 模型并调用 predict 函数时，MLflow 会调用此方法。
  • predict_stream 如果模型 intended for use in streaming environments，则应实现 predict_stream 方法。当模型使用 mlflow.pyfunc.load_model 加载为 PyFunc 模型并调用 predict_stream 时，MLflow 会调用此方法。
  • load_context 如果模型需要额外的上下文来加载，请实现 load_context 方法。有关更多详细信息，请参阅 load_context()。
提示

从 MLflow 2.20.0 开始，如果 context 参数未使用，则可以从 predict 和 predict_stream 函数中移除。例如，def predict(self, model_input, params) 是一个有效的 predict 函数签名。

下面是一个简单的 PythonModel 示例，它接受一个字符串列表并返回它。

import mlflow


class MyModel(mlflow.pyfunc.PythonModel):
    def predict(self, model_input: list[str], params=None) -> list[str]:
        return model_input

选项 2：定义一个可调用对象​

记录 PythonModel 的另一种方法是定义一个*接受单个参数*并返回预测结果的可调用对象。此可调用对象可以通过传递给 mlflow.pyfunc.log_model 来记录为 PythonModel。

提示

从 MLflow 2.20.0 开始，您可以在可调用对象上使用 @pyfunc 装饰器，以根据类型提示启用对输入的验证。有关更多详细信息，请查看 PythonModel 中的类型提示用法。

from mlflow.pyfunc.utils import pyfunc


@pyfunc
def predict(model_input: list[str]) -> list[str]:
    return model_input

记录模型​

使用 pyfunc 模块通过 mlflow.pyfunc.log_model() 记录自定义模型。

import mlflow

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

在部署前验证模型​

在部署模型之前，使用 mlflow.models.predict() API 来验证模型依赖项和输入数据。有关更多详细信息，请参阅 MLflow 模型验证。

import mlflow

mlflow.models.predict(
    model_uri=model_info.model_uri,
    input_data=["a", "b", "c"],
    env_manager="uv",
    # ONLY SET THIS if your model dependencies include pre-release packages (e.g. mlflow==3.2.0rc0)
    extra_envs={"UV_PRERELEASE": "allow"},
)

此外，您可以本地加载模型并通过运行预测来验证它。

import mlflow

pyfunc_model = mlflow.pyfunc.load_model(model_info.model_uri)
pyfunc_model.predict(["hello", "world"])

部署模型​

在生产环境中使用模型的最后一步是将其部署。请遵循 MLflow 模型部署指南来部署模型。

PythonModel 中的类型提示用法​

从 MLflow 2.20.0 开始，类型提示现在是定义模型接口的有效方式。您可以使用类型提示来定义模型的输入和输出类型。使用类型提示会带来以下好处：

  • 数据验证：MLflow 根据模型中定义的类型提示验证输入数据。无论模型是 PythonModel 实例还是已加载的 PyFunc 模型，输入数据都会得到一致的验证。
  • 类型提示推断：MLflow 根据模型中定义的类型提示推断模型的输入和输出模式，并将该推断的结构设置为已记录的模型签名。
重要

使用类型提示还可以通过跳过将数据转换为 pandas DataFrame，而是直接根据类型提示验证数据来提高性能。

支持的类型提示​

PythonModel 输入签名中使用的类型提示必须是 list[...] 类型，因为 PythonModel 的 predict 函数期望批处理输入数据。以下类型提示支持作为 list[...] 的元素类型：

  • 基本类型：int、float、str、bool、bytes、datetime.datetime
  • 集合类型：list、dict
  • 联合类型：Union[type1, type2, ...] 或 type1 | type2 | ...
  • 可选类型：Optional[type]
  • Pydantic 模型：pydantic.BaseModel 的子类（字段必须是本节中提到的受支持类型）
  • typing.Any：Any

类型提示使用限制

  • Pydantic 模型：可选字段必须包含默认值。
  • 联合类型：超过一种有效类型的联合会被推断为 MLflow 中的 AnyType，MLflow 不会基于此进行数据验证。
  • 可选类型：可选类型不能直接用于 list[...]，因为 predict 函数的输入不应为 None。

以下是一些支持的类型提示示例

  • list[str]、list[int]、list[float]、list[bool]、list[bytes]、list[datetime.datetime]
  • list[list[str]]...
  • list[dict[str, str]]、list[dict[str, int]]、list[dict[str, list[str]]]...
  • list[Union[int, str]]、list[str | dict[str, int]]...

以下是嵌套 Pydantic 模型作为类型提示的示例

from mlflow.pyfunc.utils import pyfunc
import pydantic
from typing import Optional


class Message(pydantic.BaseModel):
    role: str
    content: str


class FunctionParams(pydantic.BaseModel):
    properties: dict[str, str]
    type: str = "object"
    required: Optional[list[str]] = None
    additionalProperties: Optional[bool] = None


class ToolDefinition(pydantic.BaseModel):
    name: str
    description: Optional[str] = None
    parameters: Optional[FunctionParams] = None
    strict: Optional[bool] = None


class ChatRequest(pydantic.BaseModel):
    messages: list[Message]
    tool: Optional[ToolDefinition] = None


@pyfunc
def predict(model_input: list[ChatRequest]) -> list[list[str]]:
    return [[msg.content for msg in request.messages] for request in model_input]


input_example = [ChatRequest(messages=[Message(role="user", content="Hello")])]
print(predict(input_example))  # Output: [['Hello']]

在 PythonModel 中使用类型提示​

要在 PythonModel 中使用类型提示，您可以在 predict 函数签名中定义输入和输出类型。下面是一个 PythonModel 的示例，它接受 Message 对象列表并返回字符串列表。

import pydantic
import mlflow


class Message(pydantic.BaseModel):
    role: str
    content: str


class CustomModel(mlflow.pyfunc.PythonModel):
    def predict(self, model_input: list[Message], params=None) -> list[str]:
        return [msg.content for msg in model_input]

PythonModel 中的类型提示数据验证​

通过继承 mlflow.pyfunc.PythonModel()，您可以免费获得基于类型提示的数据验证。数据验证对 PythonModel 实例和已加载的 PyFunc 模型都有效。

下面的示例演示了数据验证如何基于上面定义的 CustomModel 工作。

model = CustomModel()

# The input_example can be a list of Message objects as defined in the type hint
input_example = [
    Message(role="system", content="Hello"),
    Message(role="user", content="Hi"),
]
print(model.predict(input_example))  # Output: ['Hello', 'Hi']

# The input_example can also be a list of dict with the same schema as Message
input_example = [
    {"role": "system", "content": "Hello"},
    {"role": "user", "content": "Hi"},
]
print(model.predict(input_example))  # Output: ['Hello', 'Hi']

# If your input doesn't match the schema, it will raise an exception
# e.g. content field is missing here, but it's required in the Message definition
model.predict([{"role": "system"}])
# Output: 1 validation error for Message\ncontent\n  Field required [type=missing, input_value={'role': 'system'}, input_type=dict]

# The same data validation works if you log and load the model as pyfunc
model_info = mlflow.pyfunc.log_model(
    name="model",
    python_model=model,
    input_example=input_example,
)
pyfunc_model = mlflow.pyfunc.load_model(model_info.model_uri)
print(pyfunc_model.predict(input_example))

对于可调用对象，您可以使用 @pyfunc 装饰器根据类型提示启用数据验证。

from mlflow.pyfunc.utils import pyfunc


@pyfunc
def predict(model_input: list[Message]) -> list[str]:
    return [msg.content for msg in model_input]


# The input_example can be a list of Message objects as defined in the type hint
input_example = [
    Message(role="system", content="Hello"),
    Message(role="user", content="Hi"),
]
print(predict(input_example))  # Output: ['Hello', 'Hi']

# The input_example can also be a list of dict with the same schema as Message
input_example = [
    {"role": "system", "content": "Hello"},
    {"role": "user", "content": "Hi"},
]
print(predict(input_example))  # Output: ['Hello', 'Hi']

# If your input doesn't match the schema, it will raise an exception
# e.g. passing a list of string here will raise an exception
predict(["hello"])
# Output: Failed to validate data against type hint `list[Message]`, invalid elements:
# [('hello', "Expecting example to be a dictionary or pydantic model instance for Pydantic type hint, got <class 'str'>")]
注意

MLflow 不会根据类型提示验证模型输出，但输出类型提示用于模型签名推断。

Pydantic 模型类型提示数据转换​

对于 Pydantic 模型类型提示，输入数据可以是 Pydantic 对象或匹配 Pydantic 模型架构的字典。MLflow 会在将数据传递给 predict 函数之前自动将提供的数据转换为类型提示对象。与上一节中的示例相比，[{"role": "system", "content": "Hello"}] 在 predict 函数中被转换为 [Message(role="system", content="Hello")]。

下面的示例演示了如何将基类用作类型提示，同时保留子类中的字段。

from pydantic import BaseModel, ConfigDict
from mlflow.pyfunc.utils import pyfunc


class BaseMessage(BaseModel):
    # set extra='allow' to allow extra fields in the subclass
    model_config = ConfigDict(extra="allow")

    role: str
    content: str


class SystemMessage(BaseMessage):
    system_prompt: str


class UserMessage(BaseMessage):
    user_prompt: str


@pyfunc
def predict(model_input: list[BaseMessage]) -> list[str]:
    result = []
    for msg in model_input:
        if hasattr(msg, "system_prompt"):
            result.append(msg.system_prompt)
        elif hasattr(msg, "user_prompt"):
            result.append(msg.user_prompt)
    return result


input_example = [
    {"role": "system", "content": "Hello", "system_prompt": "Hi"},
    {"role": "user", "content": "Hi", "user_prompt": "Hello"},
]
print(predict(input_example))  # Output: ['Hi', 'Hello']

基于类型提示的模型签名推断​

在记录带有类型提示的 PythonModel 时，MLflow 会根据模型中定义的类型提示自动推断模型的输入和输出模式。

注意

记录带有类型提示的 PythonModel 时，请勿显式传递 signature 参数。如果您传递了 signature 参数，MLflow 仍会使用基于类型提示的推断签名，并在它们不匹配时发出警告。

下表说明了类型提示如何映射到模型签名中的给定模式

类型提示推断的模式
list[str]Schema([ColSpec(type=DataType.string)])
list[list[str]]Schema([ColSpec(type=Array(DataType.string))])
list[dict[str, str]]Schema([ColSpec(type=Map(DataType.string))])
list[Union[int, str]]Schema([ColSpec(type=AnyType())])
list[Any]Schema([ColSpec(type=AnyType())])
list[pydantic.BaseModel]Schema([ColSpec(type=Object([...]))]) # properties based on the pydantic model fields
警告

Pydantic 对象不能用于 infer_signature 函数。要将 Pydantic 对象用作模型输入，您必须在 PythonModel 的 predict 函数签名中将类型提示定义为 Pydantic 模型。

在模型记录期间使用类型提示的输入示例​

在记录 PythonModel 时，建议提供一个与模型中定义的类型提示匹配的输入示例。输入示例用于验证类型提示并检查 predict 函数是否按预期工作。

import mlflow

mlflow.pyfunc.log_model(
    name="model",
    python_model=CustomModel(),
    input_example=["a", "b", "c"],
)

查询托管带类型提示的 PythonModel 的服务终端节点​

在查询托管带类型提示的 PythonModel 的服务终端节点时，您*必须在请求正文中通过* inputs *键传递输入数据*。下面的示例演示了如何本地服务模型并查询它

mlflow models serve -m runs:/<run_id>/model --env-manager local
curl http://127.0.0.1:5000/invocations -H 'Content-Type: application/json' -d '{"inputs": [{"role": "system", "content": "Hello"}]}'

额外的允许类型提示，不支持数据验证或模式推断​

MLflow 还支持在 PythonModel 中使用以下类型提示，但它们不用于数据验证或模式推断，并且在模型记录期间需要提供有效的模型签名或 input_example。

  • pandas.DataFrame
  • pandas.Series
  • numpy.ndarray
  • scipy.sparse.csc_matrix
  • scipy.sparse.csr_matrix

TypeFromExample 类型提示用法​

MLflow 提供了一个特殊的类型提示 TypeFromExample，它有助于在 PyFunc 预测期间将输入数据转换为与输入示例的类型匹配。如果您不想为模型输入显式定义类型提示，但仍然希望在预测期间数据符合输入示例类型，这非常有用。*要使用此功能，必须在模型记录期间提供有效的输入示例*。输入示例必须是以下类型之一，因为 predict 函数期望批处理输入数据：

  • list
  • pandas.DataFrame
  • pandas.Series

下面的示例演示了如何使用 TypeFromExample 类型提示

import mlflow
from mlflow.types.type_hints import TypeFromExample


class Model(mlflow.pyfunc.PythonModel):
    def predict(self, model_input: TypeFromExample):
        return model_input


with mlflow.start_run():
    model_info = mlflow.pyfunc.log_model(
        name="model",
        python_model=Model(),
        input_example=["a", "b", "c"],
    )
pyfunc_model = mlflow.pyfunc.load_model(model_info.model_uri)
assert pyfunc_model.predict(["d", "e", "f"]) == ["d", "e", "f"]
警告

如果未使用类型提示或 TypeFromExample，MLflow 的模式强制执行将默认将输入数据转换为 pandas DataFrame。如果模型期望的类型与输入示例相同，这可能不是理想的。强烈建议使用支持的类型提示来避免此转换并启用基于指定类型提示的数据验证。