构建高效MLOps工具链：基于Airflow与MLFlow的CI/CD实践指南

1次阅读

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

在传统机器学习项目中，团队常面临以下挑战：

手工操作易错 ：从数据预处理到模型部署需要大量手工操作，容易因人为失误导致流程中断
环境不一致 ：开发、测试、生产环境差异导致 ” 在我机器上能跑 ” 的问题频发
实验不可复现 ：超参数、数据版本和代码变更缺乏系统追踪，难以回溯最佳模型
部署效率低下 ：从实验到生产通常需要数周时间，无法快速响应业务需求变化

工作流编排 ：通过 DAG（有向无环图）定义任务依赖关系
调度控制 ：支持定时触发、手动触发和外部事件触发
错误处理 ：内置重试机制和报警通知功能
扩展性 ：丰富的 Operator 库支持各种数据源和计算平台

实验跟踪 ：记录参数、指标、代码和模型的全生命周期
模型管理 ：版本控制、阶段转换（Staging/Production）
部署抽象 ：统一模型打包格式（MLmodel）支持多种部署方式
协作支持 ：中央化的模型注册表实现团队共享

graph TD
    A[代码仓库] -->| 提交触发 | B[CI/CD 系统]
    B --> C[Airflow 调度器]
    C --> D[训练集群]
    D --> E[MLFlow Tracking Server]
    E --> F[模型注册表]
    F --> G[生产推理服务]
    G --> H[监控告警]
    H --> C

关键组件说明：

CI/CD 系统 ：监听代码变更，触发 Airflow DAG 运行
Airflow：编排数据获取、特征工程、模型训练全流程
MLFlow：记录实验指标，管理模型版本和元数据
模型服务 ：从注册表加载指定版本模型提供服务

from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime

def train_model(**kwargs):
    import mlflow
    # 从 Airflow 参数获取超参数
    params = kwargs['params']

    with mlflow.start_run():
        mlflow.log_params(params)
        # 训练代码...
        model = train(params)
        # 记录模型
        mlflow.pyfunc.log_model("model", python_model=model)

# 定义 DAG
with DAG(
    'ml_pipeline',
    schedule_interval='@weekly',
    default_args={
        'retries': 2,
        'retry_delay': timedelta(minutes=5)
    }
) as dag:

    preprocess = PythonOperator(
        task_id='data_preprocess',
        python_callable=preprocess_data
    )

    training = PythonOperator(
        task_id='model_training',
        python_callable=train_model,
        op_kwargs={'params': {'max_depth': 5, 'learning_rate': 0.1}}
    )

    validate = PythonOperator(
        task_id='model_validation',
        python_callable=validate_model
    )

    preprocess >> training >> validate

import mlflow.pyfunc

# 定义可部署的 Python 模型
class MyModel(mlflow.pyfunc.PythonModel):
    def __init__(self, trained_model):
        self.model = trained_model

    def predict(self, context, inputs):
        return self.model.predict(inputs)

# 加载生产环境模型
def load_production_model():
    client = mlflow.tracking.MlflowClient()
    model_uri = f"models:/my_model/production"
    return mlflow.pyfunc.load_model(model_uri)

import pytest
from mlflow import MlflowClient

def test_model_deployment():
    # 加载测试版本模型
    model = mlflow.pyfunc.load_model("models:/my_model/Staging")

    # 构造测试输入
    test_input = pd.DataFrame([[1,2,3]])

    # 验证预测输出
    output = model.predict(test_input)
    assert output.shape[0] == 1

    # 验证指标达标
    client = MlflowClient()
    run = client.get_run(model.metadata.run_id)
    assert run.data.metrics['accuracy'] > 0.85