代码之家  ›  专栏  ›  技术社区  ›  norbjd

气流:如果一个任务失败,则DAG标记为“成功”,因为触发规则“所有已完成”

  •  2
  • norbjd  · 技术社区  · 7 年前

    我有以下三个任务的DAG:

    start --> special_task --> end
    

    中间的任务可以成功也可以失败,但是 end 必须永远 执行(假设这是一个干净地关闭资源的任务)。为此,我用了 trigger rule ALL_DONE :

    end.trigger_rule = trigger_rule.TriggerRule.ALL_DONE
    

    利用这个, 结束 如果 special_task 失败。但是,自从 结束 是最后一个成功的任务,DAG总是标记为 SUCCESS .

    如何配置DAG,以便在其中一个任务失败时,将整个DAG标记为 FAILED ?

    复制示例

    import datetime
    
    from airflow import DAG
    from airflow.operators.bash_operator import BashOperator
    from airflow.utils import trigger_rule
    
    dag = DAG(
        dag_id='my_dag',
        start_date=datetime.datetime.today(),
        schedule_interval=None
    )
    
    start = BashOperator(
        task_id='start',
        bash_command='echo start',
        dag=dag
    )
    
    special_task = BashOperator(
        task_id='special_task',
        bash_command='exit 1', # force failure
        dag=dag
    )
    
    end = BashOperator(
        task_id='end',
        bash_command='echo end',
        dag=dag
    )
    end.trigger_rule = trigger_rule.TriggerRule.ALL_DONE
    
    start.set_downstream(special_task)
    special_task.set_downstream(end)
    

    This post 似乎是相关的,但答案不适合我的需要,因为下游的任务 结束 必须执行(因此 trigger_rule ).

    2 回复  |  直到 7 年前
        1
  •  4
  •   Bas Harenslak    7 年前

    我认为这是一个有趣的问题,花了一些时间来研究如何在没有额外的虚拟任务的情况下实现它。这成了一个有点多余的任务,但最终的结果是:

    这是完整的DAG:

    import airflow
    from airflow import AirflowException
    from airflow.models import DAG, TaskInstance, BaseOperator
    from airflow.operators.bash_operator import BashOperator
    from airflow.operators.dummy_operator import DummyOperator
    from airflow.operators.python_operator import PythonOperator
    from airflow.utils.db import provide_session
    from airflow.utils.state import State
    from airflow.utils.trigger_rule import TriggerRule
    
    default_args = {"owner": "airflow", "start_date": airflow.utils.dates.days_ago(3)}
    
    dag = DAG(
        dag_id="finally_task_set_end_state",
        default_args=default_args,
        schedule_interval="0 0 * * *",
        description="Answer for question https://stackoverflow.com/questions/51728441",
    )
    
    start = BashOperator(task_id="start", bash_command="echo start", dag=dag)
    failing_task = BashOperator(task_id="failing_task", bash_command="exit 1", dag=dag)
    
    
    @provide_session
    def _finally(task, execution_date, dag, session=None, **_):
        upstream_task_instances = (
            session.query(TaskInstance)
            .filter(
                TaskInstance.dag_id == dag.dag_id,
                TaskInstance.execution_date == execution_date,
                TaskInstance.task_id.in_(task.upstream_task_ids),
            )
            .all()
        )
        upstream_states = [ti.state for ti in upstream_task_instances]
        fail_this_task = State.FAILED in upstream_states
    
        print("Do logic here...")
    
        if fail_this_task:
            raise AirflowException("Failing task because one or more upstream tasks failed.")
    
    
    finally_ = PythonOperator(
        task_id="finally",
        python_callable=_finally,
        trigger_rule=TriggerRule.ALL_DONE,
        provide_context=True,
        dag=dag,
    )
    
    succesful_task = DummyOperator(task_id="succesful_task", dag=dag)
    
    start >> [failing_task, succesful_task] >> finally_
    

    看看 _finally 函数,由PythonOperator调用。这里有几个要点:

    1. 用注释 @provide_session 并添加参数 session=None ,以便可以使用 session .
    2. 查询当前任务的所有上游任务实例:
    upstream_task_instances = (
        session.query(TaskInstance)
        .filter(
            TaskInstance.dag_id == dag.dag_id,
            TaskInstance.execution_date == execution_date,
            TaskInstance.task_id.in_(task.upstream_task_ids),
        )
        .all()
    )
    
    1. 从返回的任务实例中,获取状态并检查 State.FAILED 在里面:
    upstream_states = [ti.state for ti in upstream_task_instances]
    fail_this_task = State.FAILED in upstream_states
    
    1. 执行自己的逻辑:
    print("Do logic here...")
    
    1. 最后,如果 fail_this_task=True :
    if fail_this_task:
        raise AirflowException("Failing task because one or more upstream tasks failed.")
    

    最终结果:

    enter image description here

        2
  •  1
  •   norbjd    7 年前

    作为 @JustinasMarozas 解释为 comment ,解决方案是创建一个虚拟任务,如:

    dummy = DummyOperator(
        task_id='test',
        dag=dag
    )
    

    把它绑在下游 special_task :

    failing_task.set_downstream(dummy)
    

    因此,DAG被标记为失败,并且 dummy 任务标记为 upstream_failed .

    希望有一个开箱即用的解决方案,但等待着,这个解决方案做的工作。

        3
  •  0
  •   GuD    6 年前

    要扩展Bas Harenslak答案,一个更简单的最终功能将检查所有任务(不仅是上游任务)的状态,可以是:

    def _finally(**kwargs):
        for task_instance in kwargs['dag_run'].get_task_instances():
            if task_instance.current_state() != State.SUCCESS and \
                    task_instance.task_id != kwargs['task_instance'].task_id:
                raise Exception("Task {} failed. Failing this DAG run".format(task_instance.task_id))
    
    推荐文章