Apache Airflow 动态任务映射并发控制问题解析

在Apache Airflow 3.0.0版本中，开发人员发现了一个关于动态任务映射并发控制的重要问题。该问题涉及max_active_tis_per_dag参数在动态映射任务中未能正确生效的情况，导致任务并发执行数量超出预期限制。

问题背景

动态任务映射是Airflow中一项强大的功能，它允许根据上游任务的输出动态生成多个任务实例。在实际应用中，我们经常需要控制这些动态生成任务的并发执行数量，以避免资源争用或API调用频率限制等问题。

在Airflow 2.10版本中，通过max_active_tis_per_dag参数可以有效地限制任务实例的并发数量。然而，升级到3.0.0版本后，这一机制出现了异常，导致动态映射任务无法遵守设定的并发限制。

问题复现

通过一个简单的DAG示例可以清晰地复现这个问题：

from airflow.sdk import dag, task

@dag
def example_simplest_dag():
    @task
    def my_task():
        return [1, 2, 3, 4, 5, 6, 7]

    @task(max_active_tis_per_dag=1)
    def map_me_but_slowly(a):
        import time
        time.sleep(10)
        print(a + 1)

    map_me_but_slowly.expand(a=my_task())

example_simplest_dag()

按照预期，设置了max_active_tis_per_dag=1后，map_me_but_slowly任务的多个映射实例应该串行执行，每次只运行一个实例。然而在实际运行中，多个实例却同时并行执行，完全忽略了并发限制。

问题根源

经过深入分析，发现问题出在Airflow 3.0.0版本中对动态任务映射的处理逻辑上。具体来说：

1. 对于动态映射任务，并发控制参数需要同时在partial_kwargs中进行检查
2. 在3.0.0版本中，这一检查逻辑存在遗漏，导致参数无法正确生效
3. 此外，还发现max_active_tis_per_dagrun参数也存在类似问题，该参数专门用于控制单个DAG运行中动态任务的并发数量

解决方案

针对这一问题，社区开发人员提出了修复方案：

1. 完善partial_kwargs中的参数检查逻辑
2. 同时修复max_active_tis_per_dagrun参数的实现
3. 确保两种并发控制参数都能在动态映射任务中正确工作

修复后的行为如下：

• max_active_tis_per_dag：限制任务在所有DAG运行中的并发实例总数
• max_active_tis_per_dagrun：限制单个DAG运行中动态任务的并发数量

验证结果

通过修改后的测试DAG验证修复效果：

from airflow.sdk import dag, task

@dag
def test_dag():
    @task
    def my_task():
        return [1, 2, 3, 4, 5, 6, 7]

    @task(max_active_tis_per_dag=1)
    def strictly_serial_task(a):
        import time
        time.sleep(20)
        print(a + 1)

    @task(max_active_tis_per_dagrun=1)
    def dagrun_serial_task(a):
        import time
        time.sleep(20)
        print(a + 1)

    dagrun_serial_task.expand(a=my_task())
    strictly_serial_task.expand(a=my_task())

test_dag()

测试结果表明：

1. strictly_serial_task在所有DAG运行中始终保持最多一个实例运行
2. dagrun_serial_task在每个DAG运行中保持最多一个实例运行，但不同DAG运行间的实例可以并行

总结

这个问题的修复对于需要精确控制动态任务并发执行的Airflow用户至关重要。通过正确实现这两个参数，用户可以：

1. 防止API调用频率限制
2. 避免数据库连接耗尽
3. 控制资源使用量
4. 实现更精细的任务调度策略

对于从Airflow 2.x升级到3.0的用户，如果依赖动态任务映射的并发控制功能，建议关注此问题的修复进展，或暂时回退到2.10版本。