1. 项目概述

在机器学习项目开发过程中,实验跟踪一直是个令人头疼的问题。我见过太多团队还在用Excel表格手动记录实验参数和结果,这不仅效率低下,还容易出错。当项目规模扩大、团队成员增加时,这种手工管理方式很快就会变得难以维护。

MLflow作为开源的机器学习生命周期管理平台,提供了完整的实验跟踪、模型版本管理和协作工具。但要让整个团队都能方便地使用,我们需要一个集中部署的MLflow跟踪服务器。本文将详细介绍如何使用Pulumi在AWS上部署一个完整的MLflow跟踪服务器,包括:

  • 使用EC2作为服务器主机
  • 用RDS PostgreSQL存储实验元数据
  • 通过S3存储模型和实验结果
  • 配置Nginx反向代理和基础认证
  • 实战演示如何从本地提交实验

这个方案最大的优势在于完全自动化部署,避免了手动配置AWS资源可能出现的错误,同时也方便后续扩展和维护。我在实际项目中采用这个架构后,团队协作效率提升了至少3倍。

2. 环境准备与工具安装

2.1 基础环境配置

在开始部署前,我们需要准备好本地开发环境。以下是必须安装的组件及其作用:

  1. Python 3.8+ :Pulumi支持多种语言,但我们选择Python作为基础设施代码(Infrastructure as Code)的编写语言,因为MLflow生态主要基于Python。

  2. Pulumi CLI :这是核心工具,允许我们用Python代码定义和部署云资源。安装后需要运行 pulumi login 配置访问凭证。

  3. 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"
    )

安全建议:

  1. 生产环境应启用final snapshot
  2. 密码应使用Pulumi的Secret配置
  3. 通过安全组严格控制访问源

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 监控设置

建议配置以下监控:

  1. EC2实例的CPU/内存使用率
  2. RDS的连接数和存储使用
  3. S3存储桶的容量增长

可以使用AWS CloudWatch设置警报,当资源使用超过80%时触发通知。

7.2 成本优化

  1. RDS :考虑使用Serverless版本,根据负载自动扩展
  2. EC2 :使用Spot实例或ARM架构实例(Graviton)降低成本
  3. S3 :设置生命周期策略,将旧实验数据转移到低频访问层

7.3 安全加固

  1. 启用VPC流日志监控网络流量
  2. 定期轮换数据库密码
  3. 限制S3桶的访问权限
  4. 考虑使用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. 架构演进方向

随着团队规模扩大,可以考虑以下优化:

  1. 高可用架构

    • 在多可用区部署RDS
    • 使用EC2自动伸缩组
    • 设置S3跨区域复制
  2. 性能优化

    • 为RDS添加只读副本
    • 使用ElastiCache缓存频繁访问的数据
    • 考虑使用EFS共享存储
  3. 访问控制

    • 集成AWS Cognito进行用户管理
    • 实现基于角色的访问控制(RBAC)
    • 设置详细的S3桶策略

这套架构在我参与的多个机器学习项目中表现稳定,特别是在团队协作场景下,显著提升了实验的可复现性和协作效率。实际部署时,建议根据具体需求调整资源配置和安全设置。

更多推荐