使用Pulumi在AWS上部署MLflow跟踪服务器
1. 项目概述
在机器学习项目开发过程中,实验跟踪一直是个令人头疼的问题。我见过太多团队还在用Excel表格手动记录实验参数和结果,这不仅效率低下,还容易出错。当项目规模扩大、团队成员增加时,这种手工管理方式很快就会变得难以维护。
MLflow作为开源的机器学习生命周期管理平台,提供了完整的实验跟踪、模型版本管理和协作工具。但要让整个团队都能方便地使用,我们需要一个集中部署的MLflow跟踪服务器。本文将详细介绍如何使用Pulumi在AWS上部署一个完整的MLflow跟踪服务器,包括:
- 使用EC2作为服务器主机
- 用RDS PostgreSQL存储实验元数据
- 通过S3存储模型和实验结果
- 配置Nginx反向代理和基础认证
- 实战演示如何从本地提交实验
这个方案最大的优势在于完全自动化部署,避免了手动配置AWS资源可能出现的错误,同时也方便后续扩展和维护。我在实际项目中采用这个架构后,团队协作效率提升了至少3倍。
2. 环境准备与工具安装
2.1 基础环境配置
在开始部署前,我们需要准备好本地开发环境。以下是必须安装的组件及其作用:
-
Python 3.8+ :Pulumi支持多种语言,但我们选择Python作为基础设施代码(Infrastructure as Code)的编写语言,因为MLflow生态主要基于Python。
-
Pulumi CLI :这是核心工具,允许我们用Python代码定义和部署云资源。安装后需要运行
pulumi login配置访问凭证。 -
AWS CLI :用于本地与AWS服务的交互。安装后需要通过
aws configure设置访问密钥和默认区域。
提示:建议使用IAM用户而非根账户凭证,并为该用户分配适当的权限(如AdministratorAccess)。生产环境中应进一步细化权限。
2.2 Pulumi项目初始化
创建一个新目录并初始化Pulumi项目:
mkdir mlflow-pulumi && cd mlflow-pulumi
pulumi new aws-python -y
这会生成基本的项目结构:
mlflow-pulumi/
├── Pulumi.yaml # 项目配置文件
├── __main__.py # 主部署脚本
├── requirements.txt # Python依赖
└── venv/ # 虚拟环境(如使用)
安装必要的Python包:
pip install pulumi pulumi_aws boto3
2.3 AWS资源规划
我们需要部署以下AWS服务,建议提前规划好命名和配置:
| 服务类型 | 用途 | 推荐配置 |
|---|---|---|
| VPC | 网络隔离 | CIDR: 10.0.0.0/16 |
| EC2 | MLflow服务器 | t2.micro (免费层) |
| RDS | 元数据存储 | PostgreSQL 13.8, db.t3.micro |
| S3 | 实验数据存储 | 标准存储类 |
| IAM | 访问控制 | 最小权限原则 |
3. 基础设施代码实现
3.1 项目结构设计
我们采用模块化设计,将不同AWS服务的代码分离:
services/
├── ec2.py # EC2相关资源
├── rds.py # RDS数据库
└── s3.py # S3存储桶
3.2 S3存储桶实现
在
services/s3.py
中定义创建S3 bucket的函数:
import logging
from pulumi_aws import s3
logger = logging.getLogger(__name__)
def create_bucket(name: str) -> s3.Bucket:
"""
创建用于存储MLflow实验数据的S3桶
参数:
name: 存储桶名称,需全局唯一
返回:
s3.Bucket对象
"""
try:
logger.info(f"正在创建S3存储桶: {name}")
bucket = s3.Bucket(
name,
acl="private",
force_destroy=True, # 允许非空删除
versioning={"enabled": True} # 启用版本控制
)
return bucket
except Exception as e:
logger.error(f"创建S3存储桶失败: {str(e)}")
raise
关键点说明:
-
force_destroy允许在删除堆栈时自动清理桶内数据 - 版本控制可防止实验数据意外覆盖
- ACL设为private确保数据安全
3.3 RDS数据库实现
在
services/rds.py
中定义PostgreSQL数据库的创建逻辑:
from pulumi_aws import rds
def create_postgresql(
name: str,
db_name: str,
username: str,
password: str,
subnet_ids: list[str],
security_group_id: str
) -> rds.Instance:
"""
创建PostgreSQL数据库实例
参数:
name: 实例标识名
db_name: 数据库名称
username: 管理员账号
password: 管理员密码
subnet_ids: 子网ID列表
security_group_id: 安全组ID
返回:
rds.Instance对象
"""
return rds.Instance(
name,
allocated_storage=20, # GB
engine="postgresql",
engine_version="13.8",
instance_class="db.t3.micro",
db_name=db_name,
username=username,
password=password,
skip_final_snapshot=True, # 简化演示,生产环境应为False
db_subnet_group_name=rds.SubnetGroup(
f"{name}-subnet-group",
subnet_ids=subnet_ids
).name,
vpc_security_group_ids=[security_group_id],
publicly_accessible=False, # 禁止公网直接访问
parameter_group_name="default.postgres13"
)
安全建议:
- 生产环境应启用final snapshot
- 密码应使用Pulumi的Secret配置
- 通过安全组严格控制访问源
3.4 EC2服务器实现
在
services/ec2.py
中定义MLflow服务器的创建逻辑。首先是网络基础架构:
from pulumi_aws import ec2
def create_vpc(name: str, cidr_block: str) -> ec2.Vpc:
"""创建VPC虚拟网络"""
return ec2.Vpc(
name,
cidr_block=cidr_block,
enable_dns_support=True,
enable_dns_hostnames=True,
tags={"Name": name}
)
def create_subnet(name: str, vpc_id: str, cidr_block: str, az: str) -> ec2.Subnet:
"""在指定可用区创建子网"""
return ec2.Subnet(
name,
vpc_id=vpc_id,
cidr_block=cidr_block,
availability_zone=az,
map_public_ip_on_launch=True, # 自动分配公网IP
tags={"Name": name}
)
然后是安全组配置,需要开放MLflow服务端口(默认5000)和SSH管理端口:
def create_security_group(name: str, vpc_id: str) -> ec2.SecurityGroup:
"""创建安全组,控制入站/出站流量"""
return ec2.SecurityGroup(
name,
vpc_id=vpc_id,
description="MLflow服务器安全组",
ingress=[
{"protocol": "tcp", "from_port": 22, "to_port": 22, "cidr_blocks": ["0.0.0.0/0"]},
{"protocol": "tcp", "from_port": 5000, "to_port": 5000, "cidr_blocks": ["0.0.0.0/0"]}
],
egress=[{"protocol": "-1", "from_port": 0, "to_port": 0, "cidr_blocks": ["0.0.0.0/0"]}],
tags={"Name": name}
)
最后是EC2实例创建,我们选择Ubuntu Server 20.04 LTS作为操作系统:
def create_instance(
name: str,
instance_type: str,
subnet_id: str,
security_group_ids: list[str],
key_name: str
) -> ec2.Instance:
"""创建EC2实例"""
return ec2.Instance(
name,
instance_type=instance_type,
ami="ami-0c55b159cbfafe1f0", # Ubuntu 20.04 LTS us-east-1
subnet_id=subnet_id,
vpc_security_group_ids=security_group_ids,
key_name=key_name, # 需提前在AWS创建密钥对
tags={"Name": name},
user_data="""#!/bin/bash
apt-get update -y
apt-get install -y python3-pip
pip3 install mlflow boto3 psycopg2-binary
"""
)
4. 资源编排与部署
4.1 主部署脚本
在
__main__.py
中整合所有资源创建逻辑:
import pulumi
from services.ec2 import create_vpc, create_subnet, create_security_group, create_instance
from services.rds import create_postgresql
from services.s3 import create_bucket
# 1. 创建网络基础架构
vpc = create_vpc("mlflow-vpc", "10.0.0.0/16")
subnet = create_subnet("mlflow-subnet", vpc.id, "10.0.1.0/24", "us-east-1a")
sg = create_security_group("mlflow-sg", vpc.id)
# 2. 创建计算资源
ec2_instance = create_instance(
"mlflow-server",
"t2.micro",
subnet.id,
[sg.id],
"mlflow-key" # 替换为你的密钥对名称
)
# 3. 创建数据存储
s3_bucket = create_bucket("mlflow-artifacts-unique-name") # 名称需全局唯一
db = create_postgresql(
"mlflow-db",
"mlflow",
"admin",
"SecurePassword123!", # 生产环境应使用pulumi.Config().require_secret()
[subnet.id],
sg.id
)
# 4. 导出关键信息
pulumi.export("instance_public_ip", ec2_instance.public_ip)
pulumi.export("s3_bucket_name", s3_bucket.id)
pulumi.export("db_endpoint", db.endpoint)
4.2 部署执行
运行以下命令部署基础设施:
pulumi up
部署过程会显示预览,确认无误后选择yes继续。部署完成后会输出我们导出的资源信息。
4.3 自动化部署进阶
对于团队协作场景,建议将部署流程集成到CI/CD系统中。以下是GitHub Actions的示例配置:
name: Deploy MLflow Infrastructure
on:
push:
branches: [ main ]
jobs:
deploy:
runs-on: ubuntu-latest
steps:
- uses: actions/checkout@v2
- uses: actions/setup-python@v2
with:
python-version: '3.8'
- run: pip install pulumi pulumi_aws
- run: pulumi login
env:
PULUMI_ACCESS_TOKEN: ${{ secrets.PULUMI_ACCESS_TOKEN }}
- run: pulumi up --yes --stack production
env:
AWS_ACCESS_KEY_ID: ${{ secrets.AWS_ACCESS_KEY_ID }}
AWS_SECRET_ACCESS_KEY: ${{ secrets.AWS_SECRET_ACCESS_KEY }}
5. MLflow服务器配置
5.1 服务器初始化
通过SSH连接到EC2实例:
ssh -i "mlflow-key.pem" ubuntu@<public-ip>
安装必要组件:
sudo apt update
sudo apt install -y nginx apache2-utils
sudo pip install mlflow boto3 psycopg2-binary
5.2 Nginx配置
创建Nginx配置文件
/etc/nginx/sites-available/mlflow
:
server {
listen 80;
server_name _;
auth_basic "MLflow Access";
auth_basic_user_file /etc/nginx/.htpasswd;
location / {
proxy_pass http://localhost:5000;
proxy_set_header Host $host;
proxy_set_header X-Real-IP $remote_addr;
}
}
启用配置并重启Nginx:
sudo ln -s /etc/nginx/sites-available/mlflow /etc/nginx/sites-enabled/
sudo rm /etc/nginx/sites-enabled/default
sudo systemctl restart nginx
创建认证用户:
sudo htpasswd -c /etc/nginx/.htpasswd mlflow-user
5.3 启动MLflow服务
使用systemd创建持久化服务。创建
/etc/systemd/system/mlflow.service
:
[Unit]
Description=MLflow Tracking Server
After=network.target
[Service]
User=ubuntu
Group=ubuntu
ExecStart=/usr/local/bin/mlflow server \
--host 0.0.0.0 \
--port 5000 \
--backend-store-uri postgresql://admin:SecurePassword123!@<db-endpoint>:5432/mlflow \
--default-artifact-root s3://mlflow-artifacts-unique-name/
Restart=always
[Install]
WantedBy=multi-user.target
启用并启动服务:
sudo systemctl daemon-reload
sudo systemctl enable mlflow
sudo systemctl start mlflow
6. 使用MLflow跟踪实验
6.1 基础配置
在Python脚本中设置跟踪URI:
import mlflow
mlflow.set_tracking_uri("http://<ec2-public-ip>")
mlflow.set_experiment("my-experiment")
6.2 完整实验示例
以下是一个完整的分类实验示例:
import mlflow
from sklearn.datasets import load_iris
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, classification_report
# 设置MLflow
mlflow.set_tracking_uri("http://<ec2-public-ip>")
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)
# 开始实验
with mlflow.start_run():
# 训练模型
clf = RandomForestClassifier(n_estimators=100, max_depth=3)
clf.fit(X_train, y_train)
# 评估模型
y_pred = clf.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)
# 记录参数和指标
mlflow.log_params({"n_estimators": 100, "max_depth": 3})
mlflow.log_metric("accuracy", accuracy)
# 记录分类报告
report = classification_report(y_test, y_pred, output_dict=True)
mlflow.log_dict(report, "classification_report.json")
# 记录模型
mlflow.sklearn.log_model(clf, "model")
6.3 高级功能
MLflow还支持:
- 模型注册表
- 项目打包
- 模型服务化
例如注册模型:
mlflow.register_model(
"runs:/<run-id>/model",
"iris-classifier"
)
7. 维护与优化建议
7.1 监控设置
建议配置以下监控:
- EC2实例的CPU/内存使用率
- RDS的连接数和存储使用
- S3存储桶的容量增长
可以使用AWS CloudWatch设置警报,当资源使用超过80%时触发通知。
7.2 成本优化
- RDS :考虑使用Serverless版本,根据负载自动扩展
- EC2 :使用Spot实例或ARM架构实例(Graviton)降低成本
- S3 :设置生命周期策略,将旧实验数据转移到低频访问层
7.3 安全加固
- 启用VPC流日志监控网络流量
- 定期轮换数据库密码
- 限制S3桶的访问权限
- 考虑使用AWS Certificate Manager为Nginx添加HTTPS支持
8. 常见问题排查
8.1 连接问题
症状 :无法访问MLflow UI
- 检查安全组是否开放了80端口
-
确认Nginx服务正在运行:
sudo systemctl status nginx -
查看Nginx错误日志:
sudo tail -f /var/log/nginx/error.log
8.2 数据库问题
症状 :MLflow无法连接PostgreSQL
- 确认RDS实例状态为"available"
- 检查安全组是否允许EC2访问5432端口
- 验证连接字符串中的用户名密码是否正确
8.3 存储问题
症状 :无法记录实验数据到S3
- 确认EC2实例角色有S3写入权限
- 检查S3桶名称拼写是否正确
- 验证桶的ACL和策略设置
9. 架构演进方向
随着团队规模扩大,可以考虑以下优化:
-
高可用架构 :
- 在多可用区部署RDS
- 使用EC2自动伸缩组
- 设置S3跨区域复制
-
性能优化 :
- 为RDS添加只读副本
- 使用ElastiCache缓存频繁访问的数据
- 考虑使用EFS共享存储
-
访问控制 :
- 集成AWS Cognito进行用户管理
- 实现基于角色的访问控制(RBAC)
- 设置详细的S3桶策略
这套架构在我参与的多个机器学习项目中表现稳定,特别是在团队协作场景下,显著提升了实验的可复现性和协作效率。实际部署时,建议根据具体需求调整资源配置和安全设置。
更多推荐
所有评论(0)