如何在AWS MWAA v2.7中限制Airflow并行任务并发数？

问题背景 #

需要限制Airflow任务的并行运行数量，防止压垮AWS S3等资源。但现有代码中，athena_section任务组的20个任务会同时启动，而非预期的10个。尝试过concurrency、task_concurrency和max_active_tasks_per_dag参数均未生效。

错误原因分析 #

• max_active_tasks_per_dag是DAG级配置，不能在TaskGroup内部修改DAG属性生效，且Airflow 2.7中该参数已被concurrency替代（功能一致，concurrency为通用名称）。
• task_concurrency是Operator级参数，作用是限制单个任务ID在不同DAG run中的并发实例数，而非同个DAG run内多个任务的并行数。
• 事后修改athena_section.dag.max_active_tasks_per_dag的方式不符合Airflow配置逻辑，DAG参数需在初始化时定义。

正确解决方案 #

方案1：控制整个DAG的并发任务数 #

在DAG初始化时设置concurrency参数，限制整个DAG同时运行的任务总数：

with DAG(
        dag_id="parallel_lanes_with_limited_first_task",
        default_args=DEFAULT_ARGS,
        start_date=datetime(2024, 1, 1, 1, 0, 0),
        schedule_interval="0 * * * *",
        max_active_runs=1,
        concurrency=10,  # 整个DAG最多同时运行10个任务
        tags=[],
        catchup=False,
) as dag:
    # 后续TaskGroup代码不变

方案2：控制单个TaskGroup的并发任务数（推荐，更精细） #

Airflow 2.7及以上版本支持TaskGroup的max_active_tasks参数，可单独限制某个任务组内的并行数。修改后的完整代码：

import time
from datetime import datetime

from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.utils.task_group import TaskGroup

DEFAULT_ARGS = {
    "owner": "airflow",
    "depends_on_past": False,
}

ACTIVE_PARTNERS = [{"id": n} for n in range(1, 21)]

def simulate_task_execution(sleep_time):
    print(f"模拟任务执行，休眠{sleep_time}秒。")
    time.sleep(sleep_time)

with DAG(
        dag_id="parallel_lanes_with_limited_first_task",
        default_args=DEFAULT_ARGS,
        start_date=datetime(2024, 1, 1, 1, 0, 0),
        schedule_interval="0 * * * *",
        max_active_runs=1,
        tags=[],
        catchup=False,
) as dag:
    # 在TaskGroup初始化时设置max_active_tasks，限制组内并行数
    with TaskGroup("athena_section", max_active_tasks=10) as athena_section:
        for partner in ACTIVE_PARTNERS:
            athena_insert = PythonOperator(
                task_id=f"partner_{partner['id']}_athena_insert",
                python_callable=simulate_task_execution,
                op_args=[30],  # 休眠30秒
            )

    with TaskGroup("ecs_section", max_active_tasks=100) as ecs_section:
        for partner in ACTIVE_PARTNERS:
            ecs_operators = PythonOperator(
                task_id=f"data_to_dynamodb_ecs_task_{partner['id']}",
                python_callable=simulate_task_execution,
                op_args=[5],  # 休眠5秒
            )

    athena_section >> ecs_section

参数说明 #

• max_active_tasks（TaskGroup级）：限制当前TaskGroup内，同一DAG run中最多同时运行的任务数量。
• concurrency（DAG级）：限制整个DAG中，同一DAG run中最多同时运行的任务总数。
• max_active_runs：限制该DAG同时运行的DAG实例数（即不同时间触发的run）。

验证注意事项 #

• 确认MWAA环境为Airflow 2.7及以上版本，max_active_tasks是2.7版本新增的TaskGroup参数。
• 测试时可缩短任务休眠时间，快速验证并行数是否符合预期。

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