跳至主内容

MLflow 模型服务

使用 MLflow 的全面 serving 功能，将训练好的模型转换为可用于生产的推理服务器。可在本地、云端或通过托管端点部署，并使用标准化的 REST APIs。

REST API 端点

为模型推理自动生成标准化的 REST 端点，具有一致的请求/响应格式。

多框架支持

通过 MLflow 的 flavor 系统，以统一的部署模式为来自任何机器学习框架的模型提供服务。

自定义应用

构建具有自定义逻辑、预处理和业务规则的复杂服务应用程序。

可扩展的部署

将应用部署到各种目标，从本地开发服务器到云平台和 Kubernetes 集群。

快速开始​

通过以下简单步骤，您可以在几分钟内让模型开始提供服务：

选择你的服务方式：

# Serve a logged model
mlflow models serve -m "models:/<model-id>" -p 5000

# Serve a registered model
mlflow models serve -m "models:/<model-name>/<model-version>" -p 5000

# Serve a model from local path
mlflow models serve -m ./path/to/model -p 5000

你的模型可以通过 http://localhost:5000 访问

模型服务的工作原理​

MLflow 将经过训练的模型转换为可投入生产的 HTTP 服务器，通过一个精心编排的过程来处理从模型加载到请求处理的所有环节。

服务器启动和模型加载​

当您运行 mlflow models serve 时，MLflow 会先分析模型的元数据以确定如何加载它。每个模型都包含一个 MLmodel 文件，该文件指定使用哪种“flavor”——是 scikit-learn、PyTorch、TensorFlow 还是自定义的 PyFunc 模型。

MLflow 将模型工件下载到本地目录，并创建具有标准化端点的 FastAPI 服务器。该服务器使用适当的、特定于 flavor 的加载逻辑来加载您的模型。例如，scikit-learn 模型使用 pickle 加载，而 PyTorch 模型则加载其状态字典和模型类。

服务器暴露了四个关键端点：

  • POST /invocations - 主要的预测端点
  • GET /ping 和 GET /health - 用于监控的健康检查
  • GET /version - 返回服务器和模型信息

请求处理流水线​

当预测请求到达 /invocations 时，MLflow 会通过多个验证和转换步骤对其进行处理：

输入格式检测: MLflow 自动检测您正在使用的输入格式。它支持多种格式以适应不同的用例：

  • dataframe_split: Pandas DataFrame 具有独立的列和数据数组
  • dataframe_records: 表示行的字典列表
  • instances: TensorFlow Serving 用于单个预测的格式
  • inputs: 用于更复杂输入的命名张量格式

模式验证: 如果你的模型包含签名（输入/输出模式），MLflow 会将传入的数据与其进行验证。这样可以在数据到达你的模型之前发现类型不匹配和缺失的列。

参数提取: MLflow 将预测数据与可选参数分离。像 temperature 用于语言模型，或像 threshold 用于分类器的参数会被提取并单独传递给支持它们的模型。

模型预测与响应​

在输入经过验证和格式化后，MLflow 调用模型的 predict() 方法。框架会自动检测模型是否接受参数，并据此正确调用：

# For models that accept parameters
raw_predictions = model.predict(data, params=params)

# For traditional models
raw_predictions = model.predict(data)

MLflow 然后将预测序列化回 JSON，处理各种数据类型，包括 NumPy 数组、pandas DataFrame 和 Python 列表。响应格式取决于您的输入格式——传统请求会被包装在一个 predictions 对象中，而 LLM 风格的请求则返回未包装的结果。

Flavor 系统​

MLflow 的 flavor 系统使服务在不同的 ML 框架上能够一致地运行。每个 flavor 实现了特定于框架的加载与预测逻辑，同时暴露出统一的接口。

当您使用 mlflow.sklearn.log_model() 或 mlflow.pytorch.log_model() 记录模型时，MLflow 会同时创建一个特定于 flavor 的表示和一个 PyFunc 包装器。PyFunc 包装器提供了服务层所期望的标准化 predict() 接口，而 flavor 则处理框架特有的细节，例如张量操作或数据预处理。

这种架构意味着你可以使用相同的服务命令和 API 来为 scikit-learn、PyTorch、TensorFlow 以及自定义模型提供服务。

错误处理与调试​

MLflow 的 serving 基础设施包括全面的错误处理，以帮助您调试问题：

  • 模式错误: 关于数据类型不匹配或缺失列的详细消息
  • 格式错误: 在输入格式不正确或存在歧义时提供清晰的指导
  • 模型错误: 来自您模型的预测代码的完整堆栈跟踪
  • 服务器错误: 超时和资源相关的错误处理

服务器记录所有请求和错误，便于诊断生产环境问题。

输入格式示例​

以下是 MLflow 接受的主要输入格式：

// dataframe_split format
{
  "dataframe_split": {
    "columns": ["feature1", "feature2", "feature3"],
    "data": [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
  }
}

// dataframe_records format
{
  "dataframe_records": [
    {"feature1": 1.0, "feature2": 2.0, "feature3": 3.0},
    {"feature1": 4.0, "feature2": 5.0, "feature3": 6.0}
  ]
}

// instances format (for simple models)
{
  "instances": [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]
}

所有格式都会返回一个一致的响应结构，其中包含您的预测以及模型提供的任何其他元数据。

关键实现概念​

为成功的模型服务做好准备：

  • 模型签名: 定义用于自动请求验证的输入/输出模式
  • 环境管理: 捕获依赖项以确保可重现的部署
  • 模型注册表: 使用别名实现无缝的生产更新
  • 元数据: 包括用于调试和监控的相关上下文
import mlflow
from mlflow.models.signature import infer_signature
from mlflow.tracking import MlflowClient

# Log model with comprehensive serving metadata
signature = infer_signature(X_train, model.predict(X_train))
mlflow.sklearn.log_model(
    sk_model=model,
    name="my_model",
    signature=signature,
    registered_model_name="production_model",
    input_example=X_train[:5],  # Visible example for the MLflow UI
)

# Use aliases for production deployment
client = MlflowClient()
client.set_registered_model_alias(
    name="production_model", alias="production", version="1"
)

完整示例：从训练到生产​

按照此逐步指南，从模型训练到部署为可用的 REST API：

使用自动记录训练一个简单的模型：

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
import pandas as pd

# Load sample data
iris = load_iris()
X = iris.data
y = iris.target
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

# Enable sklearn autologging with model registration
mlflow.sklearn.autolog(registered_model_name="iris_classifier")

# Train model - MLflow automatically logs everything
with mlflow.start_run() as run:
    model = RandomForestClassifier(n_estimators=10, random_state=42)
    model.fit(X_train, y_train)

    # Autologging automatically captures:
    # - Model artifacts
    # - Training parameters (n_estimators, random_state, etc.)
    # - Training metrics (score on training data)
    # - Model signature (inferred from training data)
    # - Input example

    # Optional: Log additional custom metrics
    accuracy = model.score(X_test, y_test)
    mlflow.log_metric("test_accuracy", accuracy)

    print(f"Run ID: {run.info.run_id}")
    print("Model automatically logged and registered!")

下一步​

准备构建更高级的服务应用？探索这些专业主题：

Get Started

每个部分的示例都旨在实用并可直接使用。先从上方的快速开始入手，然后探索与您的部署需求相匹配的用例。

自定义应用

通过自定义预处理、路由和业务规则构建复杂的服务逻辑

构建自定义应用 →

响应智能体

处理复杂的响应模式和多步骤推理工作流

了解智能体 →