从零构建MLOps工具链：基于MLflow和Airflow的CI/CD实战指南

1次阅读

共计 2563 个字符，预计需要花费 7 分钟才能阅读完成。

在传统的机器学习项目中，我们经常会遇到以下问题：

实验不可复现：训练参数和代码版本没有关联记录，几周后就无法复现之前的实验结果
模型版本混乱：生产环境同时运行着多个模型版本，出现问题难以追踪
环境不一致：开发环境和生产环境的依赖版本不同，导致模型表现差异
协作困难：团队成员各自维护实验记录，难以共享和比较结果

实验管理：MLflow 支持完整的实验参数、指标、代码版本记录，TensorBoard 主要关注训练过程可视化
模型管理：MLflow 提供模型注册表功能，TensorBoard 没有
部署支持：MLflow 可以直接打包模型为不同格式，TensorBoard 需要额外工具

调度能力：Airflow 支持更复杂的时间调度和依赖关系
任务监控：Airflow 提供完整的 Web UI，Luigi 相对简单
扩展性：Airflow 的 Operator 机制更容易扩展

# 示例：MLflow artifact 存储路径
artifact_uri = f"s3://my-ml-bucket/{experiment_id}/{run_id}/artifacts"
mlflow.set_tracking_uri("http://mlflow-server:5000")
mlflow.set_experiment(experiment_name)

from airflow import DAG
from airflow.operators.python_operator import PythonOperator
from datetime import datetime
import mlflow

def train_model(**kwargs):
    try:
        with mlflow.start_run():
            # 训练代码...
            mlflow.log_param("param1", value1)
            mlflow.log_metric("accuracy", 0.95)
            mlflow.sklearn.log_model(model, "model")
    except Exception as e:
        # 错误处理
        kwargs['ti'].xcom_push(key='error', value=str(e))
        raise

# 定义 DAG
default_args = {
    'owner': 'airflow',
    'start_date': datetime(2023, 1, 1),
}

dag = DAG(
    'model_training',
    default_args=default_args,
    schedule_interval='@weekly'
)

train_task = PythonOperator(
    task_id='train_model',
    python_callable=train_model,
    dag=dag,
)

# 当模型注册表中新模型达到指定指标时触发部署
from mlflow.tracking import MlflowClient

def deploy_if_best(model_name, metric_name, threshold):
    client = MlflowClient()
    latest_versions = client.get_latest_versions(model_name)

    for version in latest_versions:
        run = client.get_run(version.run_id)
        if run.data.metrics[metric_name] > threshold:
            # 触发部署流程
            trigger_deployment(version.source)

version: '3'
services:
  mlflow:
    image: mlflow/mlflow:1.20.2
    environment:
      - MLFLOW_BACKEND_STORE_URI=postgresql://user:password@postgres/mlflow
      - MLFLOW_ARTIFACT_ROOT=s3://my-ml-bucket
    ports:
      - "5000:5000"
    depends_on:
      - postgres

  postgres:
    image: postgres:13
    environment:
      - POSTGRES_USER=user
      - POSTGRES_PASSWORD=password
      - POSTGRES_DB=mlflow
    volumes:
      - pg_data:/var/lib/postgresql/data

volumes:
  pg_data:

# airflow.cfg
[celery]
worker_concurrency = 8  # 根据机器 CPU 核心数调整
worker_autoscale = 4,12  # 动态伸缩范围

from prometheus_client import start_http_server, Summary

# 定义指标
MODEL_LATENCY = Summary('model_latency_seconds', 'Time spent processing prediction')

@MODEL_LATENCY.time()
def predict(input_data):
    # 模型预测代码
    return model.predict(input_data)

# 启动监控服务器
start_http_server(8000)