如何在MWAA Airflow工作流中设置训练任务按月执行，其余步骤每周执行？

针对你的需求，这里提供两种实用的Airflow配置方案，让训练任务仅每月执行一次，同时保证数据处理和批量推理每周正常运行：

方法1：使用ShortCircuitOperator控制训练任务执行 #

ShortCircuitOperator会根据判断结果决定是否触发后续任务。如果判断为False，训练任务会被标记为"skipped"，且不影响后续批量推理任务的执行。

代码示例 #

from airflow import DAG
from airflow.operators.python import ShortCircuitOperator, PythonOperator
from datetime import datetime, timedelta

def data_processing_func():
    # 数据处理业务逻辑
    print("执行数据处理任务")

def train_func():
    # 模型训练业务逻辑
    print("执行训练任务")

def batch_inference_func():
    # 批量推理业务逻辑
    print("执行批量推理任务")

def should_run_training(**context):
    # 判断当前执行日期是否为每月1号（可根据需求调整日期规则）
    execution_date = context['execution_date']
    return execution_date.day == 1

default_args = {
    'owner': 'ml-team',
    'depends_on_past': False,
    'start_date': datetime(2024, 1, 1),
    'retries': 1,
    'retry_delay': timedelta(minutes=5),
}

with DAG(
    'ml_workflow',
    default_args=default_args,
    description='机器学习训练工作流',
    schedule_interval='@weekly',  # 每周调度一次
    catchup=False,
) as dag:

    data_processing = PythonOperator(
        task_id='data_processing',
        python_callable=data_processing_func
    )

    check_training = ShortCircuitOperator(
        task_id='check_training',
        python_callable=should_run_training,
        provide_context=True
    )

    train = PythonOperator(
        task_id='train',
        python_callable=train_func
    )

    batch_inference = PythonOperator(
        task_id='batch_inference',
        python_callable=batch_inference_func
    )

    # 任务依赖设置：数据处理后检查是否需要训练；训练完成后执行推理，同时数据处理直接触发推理（确保每周必跑）
    data_processing >> check_training >> train >> batch_inference
    data_processing >> batch_inference

逻辑说明 #

• 工作流按每周频率调度，满足数据处理和批量推理的周期需求
• should_run_training函数通过执行日期判断是否触发训练，灵活适配每月一次的要求
• 批量推理任务同时关联数据处理和训练任务，确保无论训练是否执行，每周都会正常运行

方法2：使用BranchPythonOperator实现分支逻辑 #

通过分支操作符，让工作流在每月指定日期走"数据处理→训练→批量推理"流程，非指定日期直接走"数据处理→批量推理"流程，完全跳过训练任务。

代码示例 #

from airflow import DAG
from airflow.operators.python import PythonOperator, BranchPythonOperator
from datetime import datetime, timedelta

# 数据处理、训练、批量推理函数同方法1

def choose_branch(**context):
    execution_date = context['execution_date']
    # 如果是每月1号，进入训练分支；否则直接跳转到批量推理
    if execution_date.day == 1:
        return 'train'
    else:
        return 'batch_inference'

with DAG(
    'ml_workflow_branch',
    default_args=default_args,
    description='机器学习训练工作流（分支版本）',
    schedule_interval='@weekly',
    catchup=False,
) as dag:

    data_processing = PythonOperator(
        task_id='data_processing',
        python_callable=data_processing_func
    )

    branch = BranchPythonOperator(
        task_id='branch',
        python_callable=choose_branch,
        provide_context=True
    )

    train = PythonOperator(
        task_id='train',
        python_callable=train_func
    )

    batch_inference = PythonOperator(
        task_id='batch_inference',
        python_callable=batch_inference_func,
        trigger_rule='none_failed_min_one_success'  # 确保任一上游分支成功都能触发推理
    )

    # 任务依赖设置
    data_processing >> branch >> [train, batch_inference]
    train >> batch_inference

逻辑说明 #

• choose_branch函数根据执行日期动态选择任务路径
• batch_inference任务设置trigger_rule='none_failed_min_one_success'，避免因分支路径未执行导致推理任务被阻塞
• 保持每周调度频率，每月仅一次会触发训练分支

注意事项 #

• 可根据实际需求调整日期判断逻辑，比如改为每月第一个周一、指定工作日等
• 上述代码完全兼容MWAA环境，无需额外配置
• 开启catchup=False可避免历史任务重复执行，符合生产环境需求

内容的提问来源于stack exchange，提问作者dewdrops