如何基于用户输入或DAG参数动态创建Airflow Task Group？

核心思路是利用Airflow DAG的params配置项定义可动态调整的参数，在DAG解析阶段读取该参数值，循环生成对应数量的Task Group，替代硬编码的固定范围。

实现步骤 #

1. 在DAG定义中添加可配置参数：设置num_task_groups参数并指定默认值，允许触发DAG时自定义数量
2. 读取并校验参数：获取参数值，确保为合法的正整数，避免无效输入导致DAG加载失败
3. 循环生成Task Group：用参数值替代硬编码的range(1,3)，动态创建对应数量的子Task Group

完整代码示例 #

from airflow import DAG
from airflow.operators.empty import EmptyOperator
from airflow.utils.task_group import TaskGroup
from datetime import datetime

dag = DAG(
    dag_id="dynamic_task_groups_dag",
    start_date=datetime(2024, 1, 1),
    schedule_interval=None,
    params={
        "num_task_groups": 2,  # 默认生成2个Task Group
    },
    catchup=False,
)

with dag:
    t1 = EmptyOperator(task_id="start")
    t2 = EmptyOperator(task_id="end")
    groups = []

    # 读取DAG参数中的Task Group数量，添加合法性校验
    num_groups = dag.params.get("num_task_groups", 2)
    if not isinstance(num_groups, int) or num_groups < 1:
        raise ValueError("num_task_groups必须是大于等于1的整数")

    sub_groups = []
    # 用参数值替代硬编码的range范围
    for s_id in range(1, num_groups + 1):
        @task_group(group_id=f"sub_group{s_id}")
        def tg2():
            st1 = EmptyOperator(task_id="task1")
            st2 = EmptyOperator(task_id="task2")
            st1 >> st2
        sub_groups.append(tg2())
    
    t1 >> sub_groups >> t2

    # 保留原代码中的tg1逻辑
    @task_group(group_id="tg1")
    def tg1():
        # 此处可添加原tg1的任务逻辑
        pass
    groups.append(tg1())
    groups[0] >> groups[1]

使用方式 #

在Airflow UI触发DAG时，展开Config面板，修改num_task_groups的值即可动态调整Task Group的数量。例如输入3，触发后会生成sub_group1、sub_group2、sub_group3三个Task Group。

注意事项 #

• 参数校验不可省略：必须确保num_task_groups是正整数，否则会导致DAG加载失败
• 若需从外部系统获取数量（如数据库、配置中心），可在DAG解析阶段执行查询并赋值给num_groups，但需注意DAG解析的性能影响
• DAG解析阶段会读取参数值，因此参数需在DAG加载时即可获取，避免使用仅在任务运行时才可用的变量

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