LabHub

博客

MLflow 完全指南:从实验追踪到 Model Registry、生产部署

한국어English日本語中文

MLflow 是什么?

MLflow 是一个用于管理 ML 生命周期的开源平台,由四个核心组件构成:

安装与服务器配置

基础安装

# pip 安装
pip install mlflow

# 支持额外框架
pip install mlflow[extras]  # sklearn、tensorflow、pytorch 等

# 启动服务器(本地)
mlflow server --host 0.0.0.0 --port 5000

# 使用 PostgreSQL + S3 后端的生产服务器
mlflow server \
  --backend-store-uri postgresql://mlflow:password@localhost:5432/mlflow \
  --default-artifact-root s3://mlflow-artifacts/ \
  --host 0.0.0.0 --port 5000

用 Docker Compose 部署

# docker-compose.yml
services:
  mlflow:
    image: ghcr.io/mlflow/mlflow:v3.15.1
    ports:
      - '5000:5000'
    environment:
      - MLFLOW_BACKEND_STORE_URI=postgresql://mlflow:password@postgres:5432/mlflow
      - MLFLOW_DEFAULT_ARTIFACT_ROOT=s3://mlflow-artifacts/
      - AWS_ACCESS_KEY_ID=${AWS_ACCESS_KEY_ID}
      - AWS_SECRET_ACCESS_KEY=${AWS_SECRET_ACCESS_KEY}
    command: >
      mlflow server
      --backend-store-uri postgresql://mlflow:password@postgres:5432/mlflow
      --default-artifact-root s3://mlflow-artifacts/
      --host 0.0.0.0 --port 5000
    depends_on:
      - postgres

  postgres:
    image: postgres:16
    environment:
      POSTGRES_USER: mlflow
      POSTGRES_PASSWORD: password
      POSTGRES_DB: mlflow
    volumes:
      - pgdata:/var/lib/postgresql/data

volumes:
  pgdata:

最初的五分钟:启动服务器,记录第一个 Run

本文的代码基于 MLflow 3.15.1 核对。从 2.x 到 3.x 有几个参数名发生了变化,所以先把版本对齐会更快。

在终端里只敲 mlflow server 就能把服务器拉起来。从 MLflow 3.7.0 开始 SQLite 成为默认的后端存储,因此一个参数都不给时,它会在你运行命令的目录下自动创建 sqlite:///mlflow.db。也就是说不再必须加 --backend-store-uri。不过官方文档提示,面向高并发的生产部署应当考虑 PostgreSQL 或 MySQL。后端存储支持的方言有四种:sqlite、postgresql、mysql、mssql。

当前的 CLI 文档里只有 mlflow server。老文章中频繁出现的 mlflow ui 在命令索引里找不到。文档没有写明它在哪个版本被移除,所以我不下断言,但如果这个命令已经形成肌肉记忆,迁到 mlflow server 更稳妥。

# 不带参数 — MLflow 3.7.0 及以上会自动生成 ./mlflow.db
mlflow server

# 想显式指定保存位置时
mlflow server \
  --backend-store-uri sqlite:///mlflow.db \
  --artifacts-destination ./mlartifacts \
  --host 127.0.0.1 --port 5000

# 客户端一侧(运行脚本的 shell)
export MLFLOW_TRACKING_URI=http://127.0.0.1:5000
export MLFLOW_EXPERIMENT_NAME=iris-classification

用浏览器打开 5000 端口就能看到 UI。一开始只有一个 Default 实验,Run 列表是空的。如果脚本跑完列表仍然是空的,多半是客户端没有指向服务器。客户端通过 mlflow.set_tracking_uri() 调用或 MLFLOW_TRACKING_URI 环境变量得知服务器地址,而这个环境变量的默认值是 None。什么都不配置,记录就不会送到服务器。实验名用 MLFLOW_EXPERIMENT_NAME,注册表地址用 MLFLOW_REGISTRY_URI 分别指定。

产物路径也是第一天容易混淆的地方。--serve-artifacts 默认是开启的。这意味着客户端不直接连存储,而是经由追踪服务器上传下载产物,所以 UI 上显示的产物 URI 以 mlflow-artifacts:/ 开头。实际保存位置由 --artifacts-destination 决定。想让客户端直接访问存储就用 --no-serve-artifacts,反过来想单独立一台产物代理服务器就用 --artifacts-only。把这三者分清楚,遇到训练日志留下了但模型文件没上传的情况时,能立刻缩小排查范围。

实验追踪(Tracking)

基本用法

import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, f1_score, precision_score

# 配置追踪服务器
mlflow.set_tracking_uri("http://localhost:5000")

# 创建/设置实验
mlflow.set_experiment("iris-classification")

# 准备数据
X, y = load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 运行实验
with mlflow.start_run(run_name="rf-baseline"):
    # 记录参数
    params = {
        "n_estimators": 100,
        "max_depth": 5,
        "min_samples_split": 2,
        "random_state": 42
    }
    mlflow.log_params(params)

    # 训练模型
    model = RandomForestClassifier(**params)
    model.fit(X_train, y_train)

    # 预测与指标
    y_pred = model.predict(X_test)
    metrics = {
        "accuracy": accuracy_score(y_test, y_pred),
        "f1_macro": f1_score(y_test, y_pred, average="macro"),
        "precision_macro": precision_score(y_test, y_pred, average="macro")
    }
    mlflow.log_metrics(metrics)

    # 标签
    mlflow.set_tag("model_type", "random_forest")
    mlflow.set_tag("dataset", "iris")

    # 保存模型(MLflow 3 起用 name= 取代 artifact_path=)
    mlflow.sklearn.log_model(
        model,
        name="model",
        registered_model_name="iris-classifier"
    )

    # 自定义产物(图表、报告等)
    import matplotlib.pyplot as plt
    from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

    cm = confusion_matrix(y_test, y_pred)
    fig, ax = plt.subplots()
    ConfusionMatrixDisplay(cm).plot(ax=ax)
    fig.savefig("confusion_matrix.png")
    mlflow.log_artifact("confusion_matrix.png")

    print(f"Run ID: {mlflow.active_run().info.run_id}")
    print(f"Metrics: {metrics}")

超参数调优追踪

import optuna
import mlflow

def objective(trial):
    params = {
        "n_estimators": trial.suggest_int("n_estimators", 50, 500),
        "max_depth": trial.suggest_int("max_depth", 2, 20),
        "min_samples_split": trial.suggest_int("min_samples_split", 2, 10),
        "min_samples_leaf": trial.suggest_int("min_samples_leaf", 1, 5),
    }

    with mlflow.start_run(nested=True, run_name=f"trial-{trial.number}"):
        mlflow.log_params(params)

        model = RandomForestClassifier(**params, random_state=42)
        model.fit(X_train, y_train)
        y_pred = model.predict(X_test)

        accuracy = accuracy_score(y_test, y_pred)
        mlflow.log_metric("accuracy", accuracy)

        return accuracy

# 运行 Optuna Study
with mlflow.start_run(run_name="hyperparameter-tuning"):
    study = optuna.create_study(direction="maximize")
    study.optimize(objective, n_trials=50)

    # 记录最优结果
    mlflow.log_params(study.best_params)
    mlflow.log_metric("best_accuracy", study.best_value)
    mlflow.set_tag("best_trial", study.best_trial.number)

PyTorch 模型追踪

import torch
import torch.nn as nn
import mlflow.pytorch

class SimpleNet(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super().__init__()
        self.fc1 = nn.Linear(input_dim, hidden_dim)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_dim, output_dim)

    def forward(self, x):
        return self.fc2(self.relu(self.fc1(x)))

with mlflow.start_run(run_name="pytorch-model"):
    model = SimpleNet(4, 32, 3)
    optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
    criterion = nn.CrossEntropyLoss()

    mlflow.log_params({
        "hidden_dim": 32,
        "learning_rate": 0.001,
        "optimizer": "Adam",
        "epochs": 100
    })

    for epoch in range(100):
        # 训练逻辑...
        loss = criterion(model(X_tensor), y_tensor)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 按 epoch 记录指标
        mlflow.log_metric("train_loss", loss.item(), step=epoch)

    # 保存 PyTorch 模型
    mlflow.pytorch.log_model(model, "model")

一个 Run 实际留下了什么

mlflow.start_run() 打开的 Run 保存四类记录:参数、指标、标签和产物。参数以字符串保存,指标是数值,同时传入 step 就会变成带时间轴的曲线。

start_run() 接受的参数是 run_idexperiment_idrun_namenestedparent_run_idtagsdescriptionlog_system_metrics。调优代码里用到的 nested=True 是在父 Run 内创建子 Run 的开关,子 Run 会自动带上 mlflow.parentRunId 系统标签。UI 里能折叠成树状,也是靠这个标签。

MLflow 3 变化最明显的是模型日志。第一,mlflow.sklearn.log_model()mlflow.pyfunc.log_model() 的参数改成了 name=。文档明确写着 artifact_path= 已弃用、应改用 name,这个变更从 MLflow 3.0 开始。第二,不带 mlflow.start_run() 上下文也可以调用 log_model()。返回的 ModelInfo 里带有 model_id,可以用它拼出 models:/ 形式的 URI 重新加载。第三,mlflow.sklearn.log_model()serialization_format 默认值现在是 skops。如果你的加载环境是按 cloudpickle 准备的,需要确认一下。

import mlflow
from mlflow.models import infer_signature

signature = infer_signature(model_input=X_train, model_output=model.predict(X_train))

# MLflow 3: 无需 start_run() 也能记录,并返回 ModelInfo
info = mlflow.sklearn.log_model(
    model,
    name="model",                 # artifact_path= 已弃用
    signature=signature,
    input_example=X_train[:2],
    registered_model_name="iris-classifier",
)

print(info.model_id)
loaded = mlflow.pyfunc.load_model(f"models:/{info.model_id}")

infer_signature(model_input=None, model_output=None, params=None) 会推断输入输出的 schema 并与模型一起保存。带上签名后,服务阶段列数或类型对不上时会在请求阶段被拦下,而不是让预测悄悄变得不对。

记录攒起来之后用 mlflow.search_runs() 取出。这个函数返回 pandas DataFrame,也就是说实验结果可以直接当成表来排序和 groupby。列名沿用存储层的前缀。

run_id  status  start_time  params.n_estimators  params.max_depth  metrics.accuracy  tags.model_type

先不加过滤条件调用一次并打印 .columns,然后再挑需要的列,这个顺序比较省事。具体有哪些列取决于你在 Run 里记录了什么。

Model Registry

模型注册与版本管理

from mlflow import MlflowClient

client = MlflowClient()

# 注册模型(在 log_model 中使用 registered_model_name 时会自动注册)
# 或手动注册:
result = client.create_registered_model(
    name="iris-classifier",
    description="鸢尾花分类模型"
)

# 将特定 run 的模型注册为一个版本
model_version = client.create_model_version(
    name="iris-classifier",
    source=f"runs:/{run_id}/model",
    run_id=run_id,
    description="RandomForest baseline v1"
)

print(f"Model Version: {model_version.version}")

用 Alias 进行部署管理

# MLflow 2.x 使用 Alias(Stage 已弃用)
client = MlflowClient()

# 设置生产环境 alias
client.set_registered_model_alias(
    name="iris-classifier",
    alias="champion",
    version=3
)

# 设置挑战者模型
client.set_registered_model_alias(
    name="iris-classifier",
    alias="challenger",
    version=5
)

# 通过 Alias 加载模型
champion_model = mlflow.pyfunc.load_model("models:/iris-classifier@champion")
challenger_model = mlflow.pyfunc.load_model("models:/iris-classifier@challenger")

# A/B 测试
champion_pred = champion_model.predict(X_test)
challenger_pred = challenger_model.predict(X_test)

print(f"Champion accuracy: {accuracy_score(y_test, champion_pred)}")
print(f"Challenger accuracy: {accuracy_score(y_test, challenger_pred)}")

模型标签的使用

# 为模型版本添加标签
client.set_model_version_tag(
    name="iris-classifier",
    version=3,
    key="validation_status",
    value="approved"
)

client.set_model_version_tag(
    name="iris-classifier",
    version=3,
    key="approved_by",
    value="data-science-lead"
)

# 按标签搜索模型
from mlflow import search_model_versions

approved_versions = search_model_versions(
    "name='iris-classifier' AND tag.validation_status='approved'"
)

模型服务(Serving)

MLflow 内置服务

# 本地 REST API 服务
mlflow models serve \
  -m "models:/iris-classifier@champion" \
  --port 8080 \
  --env-manager local

# 测试请求
curl -X POST http://localhost:8080/invocations \
  -H "Content-Type: application/json" \
  -d '{"inputs": [[5.1, 3.5, 1.4, 0.2]]}'

服务命令的默认值

mlflow models serve 的默认值影响比想象中大。-p/--port 是 5000,-h/--host 是 127.0.0.1,-w/--workers 是 1,-t/--timeout 是 60 秒。主机默认值是回环地址,所以在容器里原样启动时外部连不上;worker 只有一个,压测数字也会低于预期。

环境管理器用 --env-manager 选择。有效值是 localvirtualenvuvconda,默认是 virtualenv。默认为 virtualenv 意味着每次启动服务都会重新创建隔离环境,所以首次启动比预想的慢。如果镜像里依赖已经装好,--env-manager local 最快;既要可复现又想要速度,uv 值得一试。

/invocations 端点接受的负载键有五个:dataframe_splitdataframe_recordsinstancesinputsparams。五个当前都有效,没有被标记为弃用的。

mlflow models serve \
  -m "models:/iris-classifier@champion" \
  --host 0.0.0.0 --port 8080 \
  --workers 4 \
  --env-manager local

用 FastAPI 自定义服务

from fastapi import FastAPI
import mlflow.pyfunc
import numpy as np

app = FastAPI()

# 加载模型(服务器启动时执行一次)
model = mlflow.pyfunc.load_model("models:/iris-classifier@champion")

@app.post("/predict")
async def predict(features: list[list[float]]):
    predictions = model.predict(np.array(features))
    return {
        "predictions": predictions.tolist(),
        "model_version": "champion"
    }

@app.get("/health")
async def health():
    return {"status": "healthy", "model": "iris-classifier@champion"}

实验比较与分析

在 MLflow UI 中比较

# 搜索实验(CLI)
mlflow runs list --experiment-id 1

# 按指标搜索
mlflow runs list \
  --experiment-id 1 \
  --filter "metrics.accuracy > 0.95" \
  --order-by "metrics.accuracy DESC"

用 Python API 分析

import mlflow
import pandas as pd

# 查询实验中的全部 run
runs = mlflow.search_runs(
    experiment_ids=["1"],
    filter_string="metrics.accuracy > 0.9",
    order_by=["metrics.accuracy DESC"],
    max_results=10
)

# 以 DataFrame 形式分析
print(runs[["run_id", "params.n_estimators", "params.max_depth", "metrics.accuracy"]])

# 查找最优 run
best_run = runs.iloc[0]
print(f"Best run: {best_run.run_id}, Accuracy: {best_run['metrics.accuracy']}")

搜索语法:几乎所有人都会卡一次的地方

mlflow.search_runs() 和 UI 的搜索框用的是同一套过滤语法。语法本身很短,但第一次用几乎一定会撞上几条规则。

前缀对象示例
metrics.数值指标metrics.accuracy > 0.72
params.超参数(以字符串保存)params.n_estimators = "100"
tags.用户标签与系统标签tags.environment IS NOT NULL
datasets.数据集信息datasets.name = "iris"
attributes.Run 自身的属性attributes.status = "FINISHED"

通过 attributes. 可以访问 statususer_idrun_namerun_idstart_timeend_time

卡住的地方有三个。第一,支持 AND,但不支持 OR。想把两个条件用 or 连起来,只能查询两次然后在 DataFrame 层面合并。第二,参数一律以字符串保存,所以即使看起来是数字也要用双引号括起来。第三,LIKE 区分大小写,ILIKE 不区分。IS NULLIS NOT NULL 只能用于参数和标签。

runs = mlflow.search_runs(
    experiment_ids=["1"],
    filter_string='metrics.accuracy > 0.72 AND metrics.loss <= 0.15',
    order_by=["metrics.accuracy DESC"],
)

# 参数是字符串 — 不加引号就匹配不上
exact = mlflow.search_runs(filter_string='params.n_estimators = "100"')

# 没有 OR,只能查两次再拼起来
import pandas as pd

merged = pd.concat([
    mlflow.search_runs(filter_string='tags.model_type = "random_forest"'),
    mlflow.search_runs(filter_string='tags.model_type = "xgboost"'),
]).drop_duplicates(subset="run_id")

生产环境检查清单

□ 将后端存储设置为 PostgreSQL/MySQL
□ 将产物存储设置为 S3/GCS/MinIO
□ 配置认证/授权(OIDC、Basic Auth)
□ 设置自动实验记录(autolog)
□ 制定 Model Registry 的 alias 规范
□ 在 CI/CD 中自动化模型验证
□ 配置模型服务的健康检查
□ 制定实验清理策略(归档旧的 run)

失败案例与陷阱

先写症状,因为实际遇到的顺序就是这样。

症状:训练已经结束,Run 却一直停在 RUNNING。 诊断:autolog 和手动 Run 撞上了。没有活跃 Run 时,mlflow.autolog() 会自己创建一个,并在训练结束后自己关闭。但如果已经有打开的 Run,按文档的说法,它会记录到那个 Run,却不会在训练结束后自动关闭它。如果你没有用 with 块就调用了 start_run(),就必须自己调用 mlflow.end_run()

症状:参数搜索跑了 50 次,子 Run 却只有 5 个。 诊断:sklearn 的 autolog 会为参数搜索 estimator 创建一个父 Run 和多个嵌套子 Run,而子 Run 的数量受 max_tuning_runs 限制,默认值是 5。其余的试验压根不会被记录成独立的 Run。想全部看到就把这个值调大。autolog 支持的 flavor 有 Keras/TensorFlow、LightGBM、Paddle、PySpark、PyTorch、scikit-learn、Spark、statsmodels、XGBoost。

症状:log_model() 抛出弃用警告。 诊断:你在用 artifact_path=。从 MLflow 3.0 起改成了 name=,文档里也直接写明 artifact_path= 已弃用、应改用 name。现在只是警告,不会立刻出错,但新写的代码应统一用 name=

症状:mlflow.register_model() 调用好几分钟不返回。 诊断:这个函数有 await_registration_for 参数,默认值是 300 秒。它会等待模型版本就绪,最多等五分钟。CI 流水线莫名其妙多出五分钟,通常就在这里。

症状:旧教程里的 Stage 代码报警告或者不工作。 诊断:文档写明 Model Stage 已弃用,并将在未来的主版本中移除。对应的 API 是 transition_model_version_stage()。移除版本尚未公布,但现在新写的代码应该迁到 alias 和标签。文档举出的、大致对应旧 Production 阶段的名字是 champion。设置 alias 用 client.set_registered_model_alias(),读取用 client.get_model_version_by_alias(),删除用 client.delete_registered_model_alias()。加载模型时用 models:/iris-classifier@champion 形式的 URI,就不用把版本号写死在代码里。

症状:多人同时跑训练后,记录开始间歇性失败。 诊断:先确认后端存储是不是 SQLite。SQLite 基于文件锁,并发写入集中时会排队或失败。官方指引很明确:面向高并发的生产部署应当考虑 PostgreSQL 或 MySQL。一个人用的笔记本上 SQLite 足够,团队接入之后就是该迁移的时候了。

# autolog 会替你关闭 Run 的情况,以及不会的情况
mlflow.autolog()

model.fit(X_train, y_train)          # 没有活跃 Run -> 替你创建并关闭

with mlflow.start_run(run_name="manual"):
    model.fit(X_train, y_train)      # 记录到这里,离开代码块时关闭

run = mlflow.start_run(run_name="leaky")
model.fit(X_train, y_train)          # 会记录,但不会关闭
mlflow.end_run()                     # 必须自己关闭

什么时候不该用 MLflow

入门文章很少写这一段,所以单独立一节。

如果是一个 notebook 就能收尾的一次性分析,MLflow 属于过度配置。给一个判断标准:如果你没打算把同一份代码换着参数跑三次以上,那还太早。反过来,只要有人问过一次上周那次跑的是什么配置,那就是引入的时机。

MLflow 不做的事情也值得说清楚。

总结起来,MLflow 是一本账:记录你用什么配置跑了什么、结果如何。

参考资料

以下链接为 2026-08-16 核对。正文中的参数与默认值依据 MLflow 3.15.1 文档,版本不同可能会有差异。准确的 API 请在你所使用版本的文档中确认。


📝 确认测验(6 题)

Q1. MLflow 的四个核心组件是什么?

Tracking、Projects、Models、Model Registry

Q2. mlflow.log_params 与 mlflow.log_metrics 的区别是什么?

log_params 记录训练超参数(字符串),log_metrics 记录性能指标(数值)。指标可以通过 step 参数按 epoch 进行追踪。

Q3. MLflow 2.x 中用于模型部署管理的概念是什么?

Alias(例如 @champion、@challenger)。Stage 已被弃用。

Q4. nested=True 参数在什么时候使用?

用于像超参数调优这样,需要在父 run 内部记录多个子 run 的场景。

Q5. 为什么用 S3 作为产物存储?

把模型文件、图表等大体积产物保存在可扩展的对象存储中,便于团队间共享与版本管理。

Q6. mlflow.autolog() 的优点和缺点是什么?

优点:无需修改代码即可自动记录参数/指标/模型。缺点:可能会记录过多不必要的信息,自定义指标仍需单独记录。

测验

Q1:《MLflow 完全指南:从实验追踪到 Model Registry、生产部署》一文主要讨论的内容是什么?

通过动手实践,使用 MLflow 完成 ML 实验管理的完整工作流程。用 Tracking 记录实验,用 Model Registry 进行版本管理,直至完成生产部署。

Q2:MLflow 是什么? MLflow 是一个用于管理 ML 生命周期的开源平台,由四个核心组件构成:MLflow Tracking:记录实验参数、 指标和产物 MLflow Projects:将可复现的 ML 代码打包 MLflow Models:将各种框架的模型打包为统一格式 MLflow Model Registry:模型版本管理与部署工作流

Q3:安装与服务器配置的关键步骤是什么? 基础安装 用 Docker Compose 部署

Q4:实验追踪(Tracking)有哪些要点? 基本用法 超参数调优追踪 PyTorch 模型追踪

Q5:Model Registry 是如何工作的? 模型注册与版本管理 用 Alias 进行部署管理 模型标签的使用

评论

还没有评论。

登录后即可发表评论