做数据和机器学习的同学，大多会碰到一个头疼的问题：手里的大数据调度平台（比如Airflow）管着每天定时跑的数据清洗、同步，可到了要把这些数据喂给机器学习模型做预测的时候，就变成了“两张皮”——要么数据时间不对，要么模型用了旧版本，结果乱成一团。
一、为什么集成会“拧巴”
就拿之前待的电商团队来说，运营同学每周一要给10万以上的活跃用户做留存预测，用来做会员推送的依据。之前的做法是：数据岗同学每周日晚上手动把上周的用户行为数据导出到Excel，发给算法岗同学，算法岗同学再手动把数据导入模型跑预测，然后把结果再导回去。这种方式的坑数不胜数：比如有一次数据岗同学忘了导出周日的最后几个小时的数据，算法岗用了不完整的数据跑，结果预测的留存率比实际低了30%，导致运营的推送计划差点出错；还有一次模型更新了，数据岗同学还在用旧模型，结果推送出去的信息完全不准，被用户投诉了好几次。这就是Airflow和机器学习平台集成的核心痛点：数据、任务、模型三者的联动没有标准化，全靠人工衔接，容易出错，也没法追踪。
二、集成时常见的“坑”与解决办法
2.1 数据时间戳的“时差”问题
Airflow的调度是按周期来的，比如每周一跑上周的任务，这时候很容易出现时间边界错误——比如要取上周一到上周日的数据，结果取到了本周一的，或者少了周日的最后一秒。解决办法是用Airflow内置的模板变量，精准获取时间范围，不用手动写死日期。下面是带注释的示例，技术栈统一为Airflow 2.7、Python 3.10、MLflow 2.x：
# 技术栈：Airflow 2.7, Python 3.10, MLflow 2.x
from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.sensors.external_task import ExternalTaskSensor
from datetime import datetime, timedelta
import mlflow
import pandas as pd
# DAG全局配置：定义任务的基本属性
default_args = {
'owner': 'data_ml_team', # 任务负责人
'depends_on_past': True, # 上一个调度周期任务成功才跑本次
'start_date': datetime(2024, 1, 1), # DAG启动时间
'email_on_failure': True, # 任务失败发邮件通知
'retries': 2, # 任务失败重试次数
'retry_delay': timedelta(minutes=5), # 重试间隔
}
# 定义DAG：用户留存预测任务
with DAG(
'user_retention_prediction',
default_args=default_args,
schedule_interval='0 9 * * 1', # 每周一早上9点执行，对应上周数据
catchup=False, # 不补跑历史任务，避免重复计算
tags=['ml', 'user', 'retention'], # 任务标签，方便Airflow查找
) as dag:
# 1. 等待上游数据清洗任务完成
# 上游是专门做用户行为数据清洗的DAG，task_id是clean_user_behavior
wait_for_clean_data = ExternalTaskSensor(
task_id='wait_for_clean_behavior',
external_dag_id='user_behavior_clean', # 上游DAG的ID
external_task_id='clean_user_behavior', # 上游任务的ID
mode='reschedule', # 没完成就释放Worker，减少资源浪费
poke_interval=300, # 每隔5分钟检查一次上游任务状态
timeout=3600, # 最多等待1小时，避免无限等待
)
# 2. 加载上周的用户行为数据，用模板变量精准取时间
def load_target_data(**context):
# Airflow内置变量：prev_end_date是上一个调度周期的结束时间
# 比如本次调度时间是2024-05-20，prev_end_date就是2024-05-19，对应上周日
last_week_start = context['prev_ds'] # 上周一，格式是YYYY-MM-DD
last_week_end = context['prev_end_date'] # 上周日，格式同上
print(f"准备加载{last_week_start}至{last_week_end}的用户行为数据")
# 实际项目中这里会从Hive或MySQL拉数据，示例用模拟数据
# 注释：这里的return值会传给下一个任务，用XCom机制传递
return f"user_behavior_{last_week_start}.csv"
load_data_task = PythonOperator(
task_id='load_last_week_data',
python_callable=load_target_data,
provide_context=True, # 必须开启才能拿到Airflow的上下文变量
)
# 3. 调用MLflow运行留存预测模型
def run_retention_model(data_file):
# 初始化MLflow跟踪服务器，本地测试用，线上改远程地址
mlflow.set_tracking_uri("http://localhost:5000")
# 加载注册为生产环境的模型，这里用latest，自动取最新生产版本
# 注释：如果需要固定版本，改成models:/user_retention_prod/1即可
model = mlflow.pyfunc.load_model(model_uri="models:/user_retention_prod/latest")
# 模拟加载数据，实际用pd.read_csv读取data_file
# 这里的特征是用户的活跃度、消费金额等，对应模型训练时的输入
sample_features = pd.DataFrame([[0.8, 120, 30], [0.5, 50, 10], [0.9, 200, 50]])
prediction = model.predict(sample_features)
# 把预测结果存到MLflow的制品里，方便后续回溯
output_path = f"pred_result_{data_file.split('_')[-1]}"
with open(output_path, "w") as f:
f.write(str(prediction))
print(f"预测完成，结果已保存至{output_path}")
return prediction
run_pred_task = PythonOperator(
task_id='run_retention_prediction',
python_callable=run_retention_model,
op_kwargs={'data_file': load_data_task.output}, # 传递上一个任务的输出
)
# 定义任务执行顺序：等待数据清洗→加载数据→运行预测
wait_for_clean_data >> load_data_task >> run_pred_task
这个示例解决了时间戳和任务依赖的问题，每个部分都有详细注释，方便理解。
2.2 模型版本的“混乱”问题
机器学习模型不是一成不变的，旧模型可能效果差，新模型上线后如果Airflow还在用旧模型，结果就会出错。之前的电商团队就踩过这个坑，模型升级后，大家忘了改Airflow里的代码，导致连续三周的预测结果都不对，直到运营同学发现推送的效果不好才查出来。解决办法是用MLflow的模型阶段管理，把模型分成Staging（测试）和Production（生产）阶段，Airflow直接加载Production阶段的模型，不需要硬编码版本号——就像上面示例里的models:/user_retention_prod/latest，只要把新模型改成Production阶段，Airflow不用改任何代码，自动用最新的生产模型，灵活又不容易出错。
2.3 任务执行的“黑盒”问题
有时候Airflow的任务跑失败了，不知道是哪里出了问题：是数据加载错了，还是模型调用错了？Airflow本身会记录每个任务的日志，MLflow会记录每个模型的参数、输入、输出，两者结合就能把整个流程的每个环节都串起来，出问题的时候能快速找到根源。比如之前有一次预测结果为0，查Airflow日志发现是数据加载的时候把用户ID当成了特征，MLflow日志里也记录了这次的输入特征，快速定位到了问题。
三、集成的优势和要注意的地方
3.1 集成的好处
第一个是自动化，把原来手动做的数据导出、模型运行、结果推送都交给Airflow和MLflow，不用人工盯着，减少了很多重复劳动；第二个是可追溯，每个任务的执行时间、状态、日志，每个模型的输入、参数、结果，都能在Airflow和MLflow里查到，出问题能快速回溯；第三个是复用性，模型注册后可以重复用，不同的任务可以调用同一个模型，不用重复开发。
3.2 要注意的关键点
第一是资源隔离，Airflow的任务要给足够的资源，比如模型训练的时候可能占CPU或内存，要设置任务的资源限制，避免影响其他任务；第二是环境一致性，Airflow Worker的Python环境要和MLflow模型的运行环境一致，比如都用Python3.10，避免依赖冲突，比如之前有团队因为Airflow用Python3.8，模型用Python3.10，导致MLflow加载模型失败；第三是错误报警，一定要给Airflow任务设置错误报警，比如邮件、Slack，一旦任务失败，立刻通知到对应的人，避免影响业务；第四是网络连通性，线上部署的时候，Airflow的Worker要能访问MLflow的跟踪服务器和模型存储地址，不然任务会一直失败。
四、总结
Airflow和机器学习平台的集成，本质是把数据调度的“指挥权”交给Airflow，把模型的“算力仓库”交给MLflow，两者结合解决了机器学习落地时的人工干预多、容易出错、追溯难的痛点。从数据准备到模型运行，全流程自动化，适合电商、金融这类需要定期跑机器学习任务的场景。只要注意资源、环境、报警、网络这几个关键点，就能稳定运行，帮团队节省大量时间，提升效率。
评论
围绕“大数据调度平台Airflow与机器学习平台集成问题解决与优化方法”参与讨论