Airflow中EMR Steps传感器失败后如何触发前置Operator重试？

完全可以实现这个需求，核心是通过Airflow的任务状态控制和回调机制，让传感器失败时主动触发前置任务重试，以下是具体实现方案：

核心思路 #

当EMR Steps传感器（a_sensor）失败时，通过自定义回调函数修改前置Operator（a）的状态为UP_FOR_RETRY，触发调度器重新执行a；同时关闭传感器自身的重试逻辑，让其等待a重试完成后再次执行。

具体实现代码 #

from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.sensors.emr import EmrStepSensor
from airflow.utils.state import State
from airflow.utils.trigger_rule import TriggerRule
from datetime import datetime, timedelta

# 传感器失败时触发前置任务重试的回调函数
def trigger_previous_task_retry(context):
    ti = context["ti"]
    # 获取前置任务a的TaskInstance对象
    prev_task_ti = ti.get_previous_ti(task_id="task_a")
    if prev_task_ti:
        # 将前置任务状态设置为UP_FOR_RETRY，触发调度器重试
        prev_task_ti.set_state(State.UP_FOR_RETRY)

with DAG(
    dag_id="emr_step_failover_flow",
    start_date=datetime(2024, 1, 1),
    schedule_interval=None,
    catchup=False,
    default_args={"retries": 0}
) as dag:
    # 前置EMR任务Operator
    task_a = PythonOperator(
        task_id="task_a",
        python_callable=lambda: print("执行EMR任务步骤"),
        retries=3,  # 设置前置任务的重试次数
        retry_delay=timedelta(minutes=5)
    )

    # EMR步骤传感器
    task_a_sensor = EmrStepSensor(
        task_id="task_a_sensor",
        job_flow_id="your_emr_cluster_id",
        step_id="your_emr_step_id",
        trigger_rule=TriggerRule.ALL_SUCCESS,
        on_failure_callback=trigger_previous_task_retry,
        retries=0  # 关闭传感器自身重试，等待前置任务重试后再执行
    )

    task_a >> task_a_sensor

关键配置说明 #

1. 传感器回调函数：通过on_failure_callback捕获传感器失败事件，调用Airflow的TaskInstance API修改前置任务状态，触发重试。
2. 前置任务重试设置：给task_a配置足够的retries和retry_delay，确保故障时能多次重试。
3. 传感器重试关闭：将task_a_sensor的retries设为0，避免传感器自身重复重试，而是依赖前置任务重试完成后重新执行。

替代方案：分支任务控制流程 #

如果需要更灵活的流程判断，可以使用BranchPythonOperator实现分支逻辑，根据传感器状态决定是否重试前置任务：

from airflow.operators.python import BranchPythonOperator

def check_sensor_status(**context):
    ti = context["ti"]
    sensor_ti = ti.get_task_instance(task_id="task_a_sensor")
    # 传感器失败则返回前置任务ID，否则返回结束任务ID
    return "task_a" if sensor_ti.state == State.FAILED else "end_flow"

branch_check = BranchPythonOperator(
    task_id="check_sensor_status",
    python_callable=check_sensor_status,
    trigger_rule=TriggerRule.ONE_FAILED
)

end_flow = PythonOperator(
    task_id="end_flow",
    python_callable=lambda: print("流程执行完成"),
    trigger_rule=TriggerRule.ALL_SUCCESS
)

task_a >> task_a_sensor >> branch_check
branch_check >> task_a
branch_check >> end_flow

注意事项 #

• 确保Airflow元数据库权限允许回调函数修改任务状态。
• 根据业务需求合理设置前置任务的重试次数，避免无限循环。
• Airflow 2.x版本对TaskInstance API的兼容性更好，建议使用该版本以上的环境。

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