Airflow中如何实现条件不满足时跳过指定DAG任务跳转至后续任务

需要实现Airflow DAG分支逻辑：校验任务生成1-10的随机数，若随机数大于5，则跳过紧邻的run_beagle_script任务，直接执行wait_for_result任务；否则按init -> 校验 -> run_beagle_script -> wait_for_result -> update_status的顺序逐任务执行。
逻辑流程图如下：

原代码运行时，随机数大于5的场景下无法按预期跳过任务跳转。

原代码存在4处错误导致分支逻辑失效：

• 分支汇合节点未配置正确触发规则：Airflow任务默认触发规则为all_success，要求所有直接上游全部执行成功才会启动当前任务。跳转分支下run_beagle_script被标记为跳过，wait_for_result不满足默认触发规则，会被直接标记为上游失败，无法执行。
• XCOM拉取逻辑不兼容分支场景：wait_for_result、update_status两个任务硬编码拉取run_beagle_script推送的XCOM值，一旦该任务被跳过，没有对应XCOM数据，会直接抛出KeyError导致任务崩溃。
• 算子类型误用：update_status任务不需要分支判断逻辑，原代码错用BranchPythonOperator，该算子要求调用函数必须返回下游任务ID，原函数无返回值，本身就会触发执行错误。
• 依赖链配置缺失：顺序执行场景下，run_beagle_script没有配置指向wait_for_result的下游依赖，就算不触发分支，链路也是断裂的。

直接替换为以下修复后的代码即可：

from airflow import DAG
from airflow.operators.python_operator import PythonOperator, BranchPythonOperator
from random import randint
from datetime import datetime

def _training_model():
    return randint(1, 10)

def check_condition_func(ti, **kwargs):
    rand_num = _training_model()
    print('check_condition_func result: ', rand_num)
    if rand_num > 5:
        return 'wait_for_result'
    else:
        return 'run_beagle_script'

def init_params(**kwargs):
    parsed_dag_input = {
        'file_name': 'Imputation',
        'input_format': 'pdf'
    }
    kwargs['ti'].xcom_push(key='value', value=parsed_dag_input)

def run_beagle_script(**kwargs):
    imputed_file_name = kwargs['ti'].xcom_pull(key='value')['file_name']
    print('imputed_file_name function received file: ', imputed_file_name)
    run_beagle_json = {
        'result': imputed_file_name + '_Beagle'
    }
    kwargs['ti'].xcom_push(key='run_beagle_res', value=run_beagle_json)

def wait_for_result_file(**kwargs):
    # 兼容run_beagle_script被跳过的场景，拉取不到值时用初始参数兜底
    beagle_res = kwargs['ti'].xcom_pull(key='run_beagle_res', task_ids='run_beagle_script')
    if beagle_res:
        wait_for_res_val = beagle_res['result']
    else:
        init_val = kwargs['ti'].xcom_pull(key='value', task_ids='init')['file_name']
        wait_for_res_val = f"{init_val}_Beagle"
    print('wait_for_result_file function received value: ', wait_for_res_val)
    wait_for_json = {
        'result': f"{wait_for_res_val}.pdf"
    }
    kwargs['ti'].xcom_push(key='wait_for_res', value=wait_for_json)

def update_status(**kwargs):
    # 同样兼容跳过场景
    beagle_res = kwargs['ti'].xcom_pull(key='run_beagle_res', task_ids='run_beagle_script')
    if beagle_res:
        update_status_res_val = beagle_res['result']
    else:
        init_val = kwargs['ti'].xcom_pull(key='value', task_ids='init')['file_name']
        update_status_res_val = f"{init_val}_Beagle"
    print('update_status function received value: ', update_status_res_val)


with DAG("imputation_beagle", start_date=datetime(2021, 1, 1),
    schedule_interval="*/5 * * * *", catchup=False) as dag:

        input_parameters = PythonOperator(
            task_id="init",
            python_callable=init_params,
            provide_context=True
        )

        checker = BranchPythonOperator(
            task_id="check_condition",
            python_callable=check_condition_func,
            provide_context=True
        )

        run_beagle_script_op = PythonOperator(
            task_id="run_beagle_script",
            python_callable=run_beagle_script,
            provide_context=True
        )

        wait_for_result_file_op = PythonOperator(
            task_id="wait_for_result",
            python_callable=wait_for_result_file,
            provide_context=True,
            # 配置触发规则：上游无失败、至少一个上游成功即可执行，跳过状态不算失败
            trigger_rule="none_failed_min_one_success"
        )

        update_job_status = PythonOperator(
            task_id="update_status",
            python_callable=update_status,
            provide_context=True,
            trigger_rule="none_failed_min_one_success"
        )

        # 配置完整依赖链路
        input_parameters >> checker
        checker >> [run_beagle_script_op, wait_for_result_file_op]
        run_beagle_script_op >> wait_for_result_file_op
        wait_for_result_file_op >> update_job_status

• 给wait_for_result、update_status两个汇合节点配置trigger_rule="none_failed_min_one_success"，只要上游没有失败任务、至少一个上游执行成功就可以启动，兼容分支跳过时部分上游被跳过的场景。
• 调整两个汇合节点的XCOM拉取逻辑，拉取不到run_beagle_script的输出时，用init节点推送的初始参数兜底计算结果，避免KeyError。
• 将update_status的算子从BranchPythonOperator改回普通PythonOperator，匹配该任务无分支逻辑的实际用途。
• 补全依赖链路：checker同时指向两个分支节点，run_beagle_script执行完成后流向wait_for_result，保证顺序执行场景链路通畅。

• 随机数<=5时：checker路由到run_beagle_script，脚本执行完成后进入wait_for_result，最后执行状态更新，全链路按顺序执行。
• 随机数>5时：checker直接路由到wait_for_result，run_beagle_script被标记为跳过，后续任务用兜底逻辑正常取值执行，完全符合预期。

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