如何限制Airflow并行任务失败时重复触发on_failure_callback

当Airflow中并行执行的task1和task2均失败时，每个失败任务都会触发on_failure_callback，导致向同一Kafka主题发送两条错误消息，但接收端仅支持处理单条错误场景，需实现首次触发回调后阻止后续失败任务发送消息的逻辑。

现有DAG代码如下：

from datetime import datetime, timedelta
from airflow.models import Variable
from airflow import DAG
from airflow.operators.dummy import DummyOperator
from airflow.operators.python_operator import PythonOperator


def task_failure_callback(context):
    ti = context['task_instance']
    print(f"task {ti.task_id } failed in dag { ti.dag_id }, error: {ti.xcom_pull(key='error')} ")
    #call function to publish kafka message


def task_success_callback(context):
    ti = context['task_instance']
    print(f"Task {ti.task_id } has succeeded in dag { ti.dag_id }.")
    #call function to publish kafka message    

def dag_success_callback(context):
    dag_status = f"DAG has succeeded, run_id: {context['run_id']}"
    print(dag_status)
    Variable.set("TEST_CALLBACK_DAG_STATUS", dag_status)
    #call function to publish kafka message

def dag_failure_callback(context):
    ti = context['task_instance']
    dag_status = f"DAG has failed, run_id: {context['run_id']}, task id: {ti.task_id}"
    print(dag_status)
    Variable.set("TEST_CALLBACK_DAG_STATUS", dag_status)
    #call function to publish kafka message

def user_func1(ti):
    try:
        input_val = int(Variable.get("TEST_CALLBACK_INPUT", 0))
        if input_val % 10 == 0:
            raise ValueError("Invalid Input")
    except Exception as e:
        ti.xcom_push(key="error", value=str(e))
        raise e

def user_func2(ti):
    try:
        input_val = int(Variable.get("TEST_CALLBACK_INPUT", 0))
        if input_val % 2 == 0:
            raise ValueError("Invalid Input")
    except Exception as e:
        ti.xcom_push(key="error", value=str(e))
        raise e
    # pass

default_args = {
    "on_success_callback": None,
    "on_failure_callback": dag_failure_callback,
}

with DAG(
    dag_id="test_callbacks_dag",
    default_args=default_args,
    schedule_interval=None,
    start_date=datetime(2021, 1, 1),
    dagrun_timeout=timedelta(minutes=60),
    catchup=False,
) as dag:

    task1 = PythonOperator(task_id="task1", python_callable=user_func1)
    task2 = PythonOperator(task_id="task2", python_callable=user_func2)
    task3 = DummyOperator(task_id="task3", on_success_callback=task_success_callback)
    [task1, task2] >> task3

并行任务失败时的日志：

[2022-10-08, 00:10:51 IST] {logging_mixin.py:115} INFO - DAG has failed, run_id: manual__2022-10-07T18:40:50.355282+00:00, task id: task1
[2022-10-08, 00:10:51 IST] {logging_mixin.py:115} INFO - DAG has failed, run_id: manual__2022-10-07T18:40:50.355282+00:00, task id: task2

核心思路是利用XCom为每个DAG Run存储一个「是否已发送失败消息」的标记，在回调函数中先检查该标记：若未发送则执行Kafka消息推送并设置标记，若已发送则直接跳过。XCom天然与DAG Run绑定，能避免不同运行实例之间的干扰。

修改后的完整代码如下：

from datetime import datetime, timedelta
from airflow.models import Variable, XCom
from airflow import DAG
from airflow.operators.dummy import DummyOperator
from airflow.operators.python_operator import PythonOperator
from airflow.utils.session import provide_session


def task_failure_callback(context):
    ti = context['task_instance']
    print(f"task {ti.task_id } failed in dag { ti.dag_id }, error: {ti.xcom_pull(key='error')} ")
    #call function to publish kafka message


def task_success_callback(context):
    ti = context['task_instance']
    print(f"Task {ti.task_id } has succeeded in dag { ti.dag_id }.")
    #call function to publish kafka message    

def dag_success_callback(context):
    dag_status = f"DAG has succeeded, run_id: {context['run_id']}"
    print(dag_status)
    Variable.set("TEST_CALLBACK_DAG_STATUS", dag_status)
    #call function to publish kafka message

@provide_session
def dag_failure_callback(context, session=None):
    ti = context['task_instance']
    run_id = context['run_id']
    dag_id = context['dag'].dag_id
    # 定义唯一标记键，绑定DAG ID和Run ID
    marker_key = f"kafka_failure_msg_sent_{run_id}"
    
    # 检查XCom中是否已存在标记
    existing_marker = session.query(XCom).filter(
        XCom.dag_id == dag_id,
        XCom.run_id == run_id,
        XCom.key == marker_key
    ).first()
    
    if not existing_marker:
        # 未发送过消息，执行Kafka推送逻辑
        dag_status = f"DAG has failed, run_id: {run_id}, task id: {ti.task_id}"
        print(dag_status)
        Variable.set("TEST_CALLBACK_DAG_STATUS", dag_status)
        # call function to publish kafka message
        
        # 写入XCom标记，标记已发送
        ti.xcom_push(key=marker_key, value=True)
    else:
        # 已发送过消息，跳过推送
        print(f"Kafka failure message already sent for run_id {run_id}, skipping...")

def user_func1(ti):
    try:
        input_val = int(Variable.get("TEST_CALLBACK_INPUT", 0))
        if input_val % 10 == 0:
            raise ValueError("Invalid Input")
    except Exception as e:
        ti.xcom_push(key="error", value=str(e))
        raise e

def user_func2(ti):
    try:
        input_val = int(Variable.get("TEST_CALLBACK_INPUT", 0))
        if input_val % 2 == 0:
            raise ValueError("Invalid Input")
    except Exception as e:
        ti.xcom_push(key="error", value=str(e))
        raise e
    # pass

default_args = {
    "on_success_callback": None,
    "on_failure_callback": dag_failure_callback,
}

with DAG(
    dag_id="test_callbacks_dag",
    default_args=default_args,
    schedule_interval=None,
    start_date=datetime(2021, 1, 1),
    dagrun_timeout=timedelta(minutes=60),
    catchup=False,
) as dag:

    task1 = PythonOperator(task_id="task1", python_callable=user_func1)
    task2 = PythonOperator(task_id="task2", python_callable=user_func2)
    task3 = DummyOperator(task_id="task3", on_success_callback=task_success_callback)
    [task1, task2] >> task3

1. 标记唯一性：通过kafka_failure_msg_sent_{run_id}作为XCom的键，确保每个DAG Run的标记独立，不会互相干扰。
2. 原子性检查：使用Airflow的@provide_session装饰器直接操作数据库会话，避免并发场景下的竞态条件（多个任务同时检查标记时，确保只有第一个任务能写入标记并发送消息）。
3. XCom清理：Airflow会自动清理旧的XCom数据（默认保留30天），无需手动处理过期标记。

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