AI模型注册中心实践笔记
最近把团队的几个机器学习项目从零散的服务器迁移到了统一的 MLOps 平台。
项目少的时候还能应付,当项目数量到了 10+,团队到了 5 人,这套流程就彻底崩了。
背景
最近把团队的几个机器学习项目从零散的服务器迁移到了统一的 MLOps 平台。迁移过程中发现了个大问题——模型管理简直是一团糟。
现状是这样的:
- 开发环境的模型文件散落在 5 台服务器的不同目录
- 模型命名随心所欲,有叫
model.pkl的,也有叫best_model_v3.h5的 - 想回滚到上个版本的模型,得翻遍每个人的聊天记录才能找到对应的文件
- 每次部署到生产环境,都要手动拷贝文件,容易出错
这就是典型的"先跑起来再说"带来的技术债。项目少的时候还能应付,当项目数量到了 10+,团队到了 5 人,这套流程就彻底崩了。
需求
到底要解决什么问题?我们把场景具体化:
- 模型去哪儿了:需要一个地方统一存放所有模型,别再到处找了
- 版本控制:每个模型需要明确的版本号,能追溯历史版本
- 元数据追踪:知道模型是用什么数据训练的,超参数是什么,性能指标如何
- 快速回滚:新模型上线有问题,能立刻切回上一个版本
- 多环境支持:开发、测试、生产环境能共享同一个注册中心,但权限隔离
- CI/CD 集成:模型训练和部署要自动化
说白了,就是需要一个 Git 版本控制系统,但是针对的是模型文件。
实现
技术选型
调研了几个方案:
| 方案 | 优点 | 缺点 | 评分 |
|---|---|---|---|
| 手动 Git 管理 | 简单直接 | 模型文件大,克隆慢,二进制文件版本控制不友好 | ⭐⭐ |
| S3 + 元数据库 | 灵活可控 | 需要自己写很多胶水代码,维护成本高 | ⭐⭐⭐ |
| MLflow Model Registry | 成熟完整,开箱即用 | 学习曲线,需要部署服务 | ⭐⭐⭐⭐⭐ |
| Kubeflow Model Registry | K8s 原生 | 复杂度高,资源消耗大 | ⭐⭐⭐⭐ |
最终选了 MLflow,理由很实际:团队已经用了 MLflow 跟踪实验,Model Registry 是无缝集成的。而且它的概念清晰,不会因为抽象层太多而让团队成员迷惑。
架构设计
整体架构长这样:
核心思想是:模型训练完成后自动注册到 Registry,通过 Stage(None -> Staging -> Production)来管理模型的生命周期。
具体实现
1. 注册中心部署
用 Docker 部署一个最小化的 MLflow 服务:
docker run -d \
--name mlflow-server \
-p 5000:5000 \
-v mlflow-data:/mlflow \
-e BACKEND_STORE_URI=postgresql://user:pass@postgres:5432/mlflow \
-e DEFAULT_ARTIFACT_ROOT=s3://mlflow-artifacts \
mlflow server \
--backend-store-uri postgresql://user:pass@postgres:5432/mlflow \
--default-artifact-root s3://mlflow-artifacts \
--host 0.0.0.0
这里踩了个坑:一开始用的 SQLite 存储,数据量上来了之后查询特别慢。换成 PostgreSQL 好了,但要注意版本兼容性——MLflow 2.x 对 PostgreSQL 的最低版本有要求。
2. 模型训练集成
在训练脚本中集成 MLflow:
import mlflow
import mlflow.sklearn
from sklearn.ensemble import RandomForestClassifier
# 设置实验
mlflow.set_experiment("fraud-detection")
with mlflow.start_run():
# 记录超参数
params = {
"n_estimators": 100,
"max_depth": 10,
"random_state": 42
}
mlflow.log_params(params)
# 训练模型
model = RandomForestClassifier(**params)
model.fit(X_train, y_train)
# 记录指标
metrics = {
"train_accuracy": model.score(X_train, y_train),
"val_accuracy": model.score(X_val, y_val)
}
mlflow.log_metrics(metrics)
# 记录模型文件
mlflow.sklearn.log_model(model, "model")
# 获取 run ID,后面注册模型要用
run_id = mlflow.active_run().info.run_id
print(f"Run ID: {run_id}")
关键点:
log_params和log_metrics会自动出现在 MLflow UIlog_model会把模型文件和依赖信息都保存下来
3. 模型注册与版本管理
训练完成后,用脚本自动注册模型:
from mlflow.tracking import MlflowClient
from mlflow.entities import ViewType
client = MlflowClient()
def register_model(model_name, run_id, description=""):
"""注册模型到 Registry"""
# 获取模型 URI
model_uri = f"runs:/{run_id}/model"
# 注册模型(如果已存在会自动创建新版本)
model_version = mlflow.register_model(
model_uri=model_uri,
name=model_name
)
# 添加版本描述
client.update_model_version(
name=model_name,
version=model_version.version,
description=description
)
print(f"Registered {model_name} version {model_version.version}")
return model_version
# 使用示例
model_version = register_model(
model_name="fraud-detection-model",
run_id="a1b2c3d4e5f6",
description="Baseline model with default hyperparameters"
)
MLflow 的版本号是自动递增的整数,不需要自己管理。
4. Stage 迁移
部署流程通过 Stage 迁移来控制:
def transition_to_staging(model_name, version):
"""迁移到 Staging 环境"""
client.transition_model_version_stage(
name=model_name,
version=version,
stage="Staging",
archive_existing_versions=True # 把旧的 Staging 版本归档
)
def transition_to_production(model_name, version):
"""迁移到 Production 环境"""
# 先在 Staging 做一些自动化测试
if not run_staging_tests(model_name, version):
raise Exception("Staging tests failed!")
client.transition_model_version_stage(
name=model_name,
version=version,
stage="Production",
archive_existing_versions=True
)
# 完整流程
def deploy_to_production(model_name, version):
"""完整的部署流程"""
print(f"Deploying {model_name} v{version}...")
# 1. 迁移到 Staging
transition_to_staging(model_name, version)
print("✓ Deployed to Staging")
# 2. 运行自动化测试(这里可以集成到 CI/CD)
if run_smoke_tests():
# 3. 迁移到 Production
transition_to_production(model_name, version)
print("✓ Deployed to Production")
else:
print("✗ Staging tests failed, aborting deployment")
# 回滚 Staging
client.transition_model_version_stage(
name=model_name,
version=version,
stage="None"
)
这里的 archive_existing_versions=True 很重要,它会自动把当前环境下的旧版本归档,避免同时有多个版本在同一环境。
5. 模型加载
加载模型很简单,只需要模型名和 stage:
def load_model(model_name, stage="Production"):
"""加载指定 stage 的模型"""
model_uri = f"models:/{model_name}/{stage}"
model = mlflow.sklearn.load_model(model_uri)
return model
# 使用示例
model = load_model("fraud-detection-model", "Production")
predictions = model.predict(new_data)
不需要知道具体的版本号,这大大简化了代码。
踩坑
1. 模型文件太大
一开始把完整的大模型(几 GB)直接存到 MLflow,结果:
- 上传和下载特别慢
- 数据库里记录的 metadata 也很多,查询变慢
解决方案:把大模型拆分:
- 小模型(< 100MB):直接存 MLflow
- 大模型:存到对象存储(S3/OSS),MLflow 只存引用和元数据
# 对于大模型,只保存配置和指针
mlflow.log_dict(
{"model_path": "s3://my-models/large-model-v1"},
"model_reference.json"
)
2. 并发冲突
团队多人同时注册模型时,偶尔会出现版本号错乱的情况。
解决方案:加锁。MLflow 本身没有提供分布式锁,我们用了 Redis:
import redis
import time
redis_client = redis.Redis(host='localhost', port=6379)
def register_model_with_lock(model_name, run_id):
lock_key = f"model_registry_lock:{model_name}"
acquired = False
try:
acquired = redis_client.set(lock_key, "1", nx=True, ex=30)
if not acquired:
raise Exception("Another registration in progress, please wait")
# 注册模型
return register_model(model_name, run_id)
finally:
if acquired:
redis_client.delete(lock_key)
3. 权限管理
MLflow 自身的权限管理比较简陋(早期版本甚至没有)。
解决方案:
- 用反向代理(Nginx)做基本的认证
- 敏感操作(删除模型、切换 Production)需要二次确认
- 审计日志:记录所有模型的注册和 Stage 变更
def audit_log(action, model_name, version, user):
"""记录审计日志"""
log_entry = {
"timestamp": time.time(),
"action": action, # "register", "transition", "delete"
"model_name": model_name,
"version": version,
"user": user
}
# 写到日志系统(ELK、Loki 等)
4. 模型依赖丢失
有一次部署到生产环境后报错,发现缺少某个依赖库。
解决方案:MLflow 会自动记录 conda.yaml 或 requirements.txt,部署时一定要用这个重建环境:
# 获取 Production 模型的依赖
mlflow models serve -m models:/fraud-detection-model/Production --env-manager conda
或者在代码中显式加载环境:
import mlflow.pyfunc
model = mlflow.pyfunc.load_model(model_uri="models:/fraud-detection-model/Production")
# 这会自动处理依赖
结果
这套系统上线 3 个月后的效果:
| 指标 | 之前 | 之后 | 改善 |
|---|---|---|---|
| 模型查找时间 | 平均 15 分钟 | 10 秒 | 90 倍 |
| 部署到生产时间 | 2-4 小时 | 30 分钟 | 4-8 倍 |
| 回滚耗时 | 1-2 小时 | 5 分钟 | 12-24 倍 |
| 模型版本混乱 | 经常发生 | 0 次 | - |
| 生产事故率 | 月均 2 次 | 月均 0.3 次 | 85% |
更重要的变化是团队的工作方式:
- 开发人员不再担心"我的模型在哪里"
- 数据科学家可以专注于模型本身,而不是文件管理
- 运维人员有了一个标准的部署流程
- 出问题时能快速定位到是哪个模型版本引入的 bug
结语
模型注册中心听起来是个高大上的 MLOps 概念,但实际上它解决的是一个非常朴素的问题:别再把东西乱放了。
从技术层面看,MLflow 并不是最先进的方案,但它胜在:
- 概念简单,学习成本低
- 和现有工作流集成容易
- 社区活跃,问题容易解决
对于大多数中小团队来说,够用就是最好的。先把模型管起来,把流程规范化,再去考虑更高级的功能,这才是务实的做法。
毕竟,我们是在解决实际工程问题,不是在玩技术栈。
版权声明: 本文首发于 指尖魔法屋-AI模型注册中心实践笔记(https://blog.thinkmoon.cn/post/288-ai-model-registry-chaos-standard-practice/) 转载或引用必须申明原指尖魔法屋来源及源地址!
评论
使用 GitHub 账号登录后即可留言,支持 Markdown。