MLflow PythonModel 指南

MLflow PythonModel 简介

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

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

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

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

定义自定义 PythonModel

选项 1：继承 PythonModel

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

PythonModel 的方法：

predict 一个有效的 PythonModel 必须实现 predict 方法，该方法定义了模型的预测逻辑。当模型作为 PyFunc 模型使用 mlflow.pyfunc.load_model 加载并调用 predict 函数时，MLflow 会调用此方法。
predict_stream 该方法应在模型用于流式环境时实现。MLflow 在模型作为 PyFunc 模型通过 mlflow.pyfunc.load_model 加载并调用 predict_stream 时会调用此方法。
load_context 实现 load_context 方法，如果模型需要加载额外的上下文。有关更多详细信息，请参阅 load_context().

tip

从 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。

tip

从 MLflow 2.20.0 起，您可以在可调用对象上使用 @pyfunc 装饰器，以根据类型提示对输入进行数据验证。查看 type hint usage in 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 Model Validation 获取更多详细信息。

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 Model Deployment 指南部署该模型。

在 PythonModel 中使用类型提示

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

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

支持的类型提示

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

原始类型: int, float, str, bool, bytes, datetime.datetime
集合类型: list, dict
联合类型: Union[type1, type2, ...] 或 type1 | type2 | ...
可选类型: Optional[type]
Pydantic 模型: Subclass of pydantic.BaseModel（字段必须为本节中提到的受支持类型）
typing.Any: 任意类型

类型提示使用的限制：

Pydantic 模型: 可选字段必须包含默认值.
联合类型：多个有效类型的联合在 MLflow 中会被推断为 AnyType，且 MLflow 不会基于此进行数据验证。
Optional types: 可选类型不能直接用于 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'>")]

note

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

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

对于 Pydantic 模型的类型提示，输入数据可以是一个 Pydantic 对象，或者是一个与该 Pydantic 模型的 schema 相匹配的字典。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 会根据模型中定义的类型提示自动推断模型的输入和输出架构。

note

在使用类型提示记录 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([...]))]) # 基于 pydantic 模型字段的属性

warning

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 函数期望批量输入数据：

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

warning

如果既没有使用类型提示，也没有使用 TypeFromExample，MLflow 的 schema 强制将默认把输入数据转换为 pandas DataFrame。这在模型期望与输入示例相同类型时可能并不理想。强烈建议使用受支持的类型提示以避免此转换，并基于指定的类型提示启用数据验证。

MLflow PythonModel 简介​

定义自定义 PythonModel​

选项 1：继承 PythonModel​

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

记录模型​

在部署前验证模型​

部署模型​

在 PythonModel 中使用类型提示​

支持的类型提示​

在 PythonModel 中使用类型提示​

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

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

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

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

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

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

TypeFromExample 类型提示的用法​

MLflow PythonModel 简介

定义自定义 PythonModel

选项 1：继承 PythonModel

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

记录模型

在部署前验证模型

部署模型

在 PythonModel 中使用类型提示

支持的类型提示

在 PythonModel 中使用类型提示

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

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

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

在模型记录期间的输入示例及类型提示

查询托管带有类型提示的 PythonModel 的服务端点

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

TypeFromExample 类型提示的用法