怎样利用MLFlow/Vertex AI Model Registry管理模型版本与元数据？

在AI模型进入生产环境的过程中，版本混乱、元数据丢失和部署流程不规范是常见的挑战。MLflow Model Registry提供了一个集中的存储库，用于管理MLflow Tracking中记录的模型，使其从实验阶段平滑过渡到部署阶段。本文将深入探讨如何使用MLflow来实现这一目标。

1. 环境准备与MLflow服务启动

首先，确保安装了必要的库，并启动MLflow Tracking Server，该服务器包含了Model Registry的功能。

1
2
3
4
5
pip install mlflow scikit-learn pandas



# 启动MLflow Tracking Server

# 使用SQLite作为后端存储，/mlruns作为artifact存储

mlflow server --backend-store-uri sqlite:///mlruns.db --default-artifact-root ./mlruns -h 0.0.0.0 -p 5000

我们将假定MLflow服务运行在本地 http://127.0.0.1:5000。

2. 模型训练、追踪与自动注册

我们使用一个简单的Scikit-learn模型作为示例。通过设置Tracking URI并启用autolog，我们可以简化模型的记录过程。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
import mlflow

import mlflow.sklearn

from sklearn.linear_model import LinearRegression

from sklearn.model_selection import train_test_split

import pandas as pd

import numpy as np



# 设置Tracking Server URI

mlflow.set_tracking_uri("http://127.0.0.1:5000")



MODEL_NAME = "MyLinearModel"



# 模拟数据

X = pd.DataFrame(np.random.rand(100, 3), columns=['F1', 'F2', 'F3'])

y = 2 * X['F1'] - 0.5 * X['F2'] + np.random.normal(0, 0.1, 100)



X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2)



# 启用autologging，自动记录参数、指标和模型

mlflow.sklearn.autolog()



with mlflow.start_run(run_name="training_run_v1") as run:

    # 训练模型

    model = LinearRegression()

    model.fit(X_train, y_train)



    # 记录额外的元数据

    mlflow.log_metric("R2_score", model.score(X_test, y_test))



    # 显式将模型注册到Model Registry，命名为MODEL_NAME

    run_id = run.info.run_id

    model_path = "model"



    # 关键步骤：将模型注册到指定的名称下

    model_uri = f"runs:/{run_id}/{model_path}"

    mlflow.register_model(

        model_uri=model_uri,

        name=MODEL_NAME

    )



print(f"Model {MODEL_NAME} registered successfully.")

执行上述代码后，一个新的注册模型（MyLinearModel）将被创建，包含第一个版本（V1），阶段为None。

3. 通过API管理模型版本阶段

Model Registry的核心功能是阶段管理（Staging, Production, Archived）。我们使用MlflowClient来程序化地控制模型的生命周期。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
from mlflow.tracking import MlflowClient



client = MlflowClient()



# 假设我们注册的模型名称是 "MyLinearModel"

registered_model_name = "MyLinearModel"



# 获取最新的版本号

latest_versions = client.get_latest_versions(registered_model_name, stages=["None"])

if latest_versions:

    version_number = latest_versions[0].version

    print(f"最新版本号: {version_number}")



    # 步骤 A: 将V1从 None 迁移到 Staging (暂存/测试阶段)

    client.transition_model_version_stage(

        name=registered_model_name,

        version=version_number,

        stage="Staging",

        archive_existing_versions=False # 不归档旧版本

    )

    print(f"版本 {version_number} 已迁移至 Staging 阶段")



    # 步骤 B: 假设测试通过，将其迁移到 Production (生产阶段)

    client.transition_model_version_stage(

        name=registered_model_name,

        version=version_number,

        stage="Production",

        # 通常我们会选择归档所有旧的生产版本，确保只有一个 Production 模型

        archive_existing_versions=True 

    )

    print(f"版本 {version_number} 已迁移至 Production 阶段")

else:

    print("未找到任何版本。")

通过上述操作，我们成功地将特定的模型版本提升到了生产环境，同时确保了模型注册表中始终只有一个处于生产状态的模型（通过archive_existing_versions=True）。

4. 生产环境部署：加载最新的生产模型

部署服务只需要知道模型的名称和所需的阶段，而不需要关心具体的版本号（例如V3, V5）。这是Model Registry带来的最大便利。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
import mlflow.sklearn

import numpy as np



# 定义模型URI，使用阶段别名

production_model_uri = f"models:/{MODEL_NAME}/Production"



print(f"尝试加载URI: {production_model_uri}")



# 加载当前处于 'Production' 阶段的模型

loaded_model = mlflow.sklearn.load_model(production_model_uri)



# 使用加载的模型进行预测

new_data = pd.DataFrame([[0.9, 0.1, 0.5]], columns=['F1', 'F2', 'F3'])

prediction = loaded_model.predict(new_data)



print(f"成功加载生产模型，并进行预测: {prediction[0]:.4f}")

总结

MLflow Model Registry通过提供清晰的版本命名、集中的元数据存储和灵活的阶段转换API，极大地简化了AI模型的管理和部署流程。无论是本地部署还是集成到如Vertex AI Model Registry这样的云服务中，理解和掌握MLflow的这一功能都是构建稳健MLOps管线的关键一步。