跳到主要内容

MLflow 模型服务

使用 MLflow 全面的服务功能,将您训练好的模型转化为生产级的推理服务器。通过标准化的 REST API,支持本地、云端或托管端点的部署。

REST API 端点

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

多框架支持

通过 MLflow 的“flavor”系统,支持任何 ML 框架的模型服务,并采用统一的部署模式。

自定义应用程序

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

可扩展部署

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

快速入门

通过以下简单步骤,即可在几分钟内完成模型服务

选择您的服务方式

bash
# 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

您的模型将可以通过 https://:5000 访问

模型服务工作原理

MLflow 通过一个精心协调的过程,将您训练好的模型转化为生产级的 HTTP 服务器,该过程负责从模型加载到请求处理的所有事务。

服务器启动和模型加载

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

MLflow 将模型构件下载到本地目录,并创建一个带有标准化端点的 FastAPI 服务器。该服务器使用适合的模型 flavor 加载逻辑来加载您的模型。例如,scikit-learn 模型使用 pickle 加载,而 PyTorch 模型则加载其 state_dict 和模型类。

服务器公开了四个关键端点

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

请求处理流水线

当预测请求到达 /invocations 端点时,MLflow 会通过多个验证和转换步骤来处理它

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

  • dataframe_split:Pandas DataFrame,包含分离的列和数据数组
  • dataframe_records:表示行的字典列表
  • instances:TensorFlow Serving 格式,用于单个预测
  • inputs:命名张量格式,适用于更复杂的输入

模式验证:如果您的模型包含签名(输入/输出模式),MLflow 会根据该签名验证传入的数据。这可以在数据到达您的模型之前捕获类型不匹配和缺失的列。

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

模型预测和响应

一旦输入被验证并格式化,MLflow 就会调用您模型的 predict() 方法。框架会自动检测您的模型是否接受参数并进行相应调用。

python
# 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 的服务基础设施包括全面的错误处理,以帮助您调试问题

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

服务器会记录所有请求和错误,从而更容易诊断生产问题。

输入格式示例

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

json
// 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]]
}

所有格式都返回一致的响应结构,包含您的预测以及模型提供的任何额外元数据。

关键实现概念

准备您的模型以成功服务

  • 模型签名:定义输入/输出模式以实现自动请求验证
  • 环境管理:捕获依赖项以确保可复现的部署
  • 模型注册表:使用别名进行无缝的生产更新
  • 元数据:包含用于调试和监控的相关上下文
python
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

使用自动日志记录训练一个简单模型

python
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!")

后续步骤

准备构建更高级的服务应用程序?探索这些专门的主题

开始使用

每个部分中的示例都旨在实用且可直接使用。从上面的快速入门开始,然后探索符合您部署需求的用例。