airflow-ha

High Availability (HA) DAG Utility

Overview

This library provides an operator called HighAvailabilityOperator, which inherits from PythonSensor and runs a user-provided python_callable. The return value can trigger the following actions:

Callbacks passed to HighAvailabilityOperator also apply to its generated decide, retrigger, and stop tasks, including the task that forces a failed DAG run after retriggering. Callbacks receive each executing task's context, so on_success_callback can run more than once per DAG run. Omitted callbacks inherit from DAG default_args. Callbacks in pass_trigger_kwargs or fail_trigger_kwargs override the corresponding retrigger task's callback.

Limiters

Arguments to HighAvailabilityOperator can be used to configure finishing behavior outside of the callable:

• runtime: A timedelta or int (seconds). The operator will turn off cleanly after dag.start_date + runtime ((PASS, STOP))
• endtime: A time or str (isoformat time). The operator will turn off cleanly after today + endtime ((PASS, STOP))
• maxretrigger: An integer. The operator will turn off after maxretrigger retriggers ((<previous status, STOP))

Example - Always On

Consider the following DAG:

with DAG(
    dag_id="test-high-availability",
    description="Test HA Operator",
    schedule=timedelta(days=1),
    start_date=datetime(2024, 1, 1),
    catchup=False,
):
    ha = HighAvailabilityOperator(
        task_id="ha",
        timeout=30,
        poke_interval=5,
        python_callable=lambda **kwargs: choice(
            (
                (Result.PASS, Action.CONTINUE),
                (Result.PASS, Action.RETRIGGER),
                (Result.PASS, Action.STOP),
                (Result.FAIL, Action.CONTINUE),
                (Result.FAIL, Action.RETRIGGER),
                (Result.FAIL, Action.STOP),
            )
        ),
    )
    
    pre = PythonOperator(task_id="pre", python_callable=lambda **kwargs: "test")
    pre >> ha
    
    retrigger_fail = PythonOperator(task_id="retrigger_fail", python_callable=lambda **kwargs: "test")
    ha.retrigger_fail >> retrigger_fail

    stop_fail = PythonOperator(task_id="stop_fail", python_callable=lambda **kwargs: fail_, trigger_rule="all_failed")
    ha.stop_fail >> stop_fail
    
    retrigger_pass = PythonOperator(task_id="retrigger_pass", python_callable=lambda **kwargs: "test")
    ha.retrigger_pass >> retrigger_pass

    stop_pass = PythonOperator(task_id="stop_pass", python_callable=lambda **kwargs: "test")
    ha.stop_pass >> stop_pass

This produces a DAG with the following topology:

This DAG exhibits cool behavior. If the check returns CONTINUE, the DAG will continue to run the sensor. If the check returns RETRIGGER or the interval elapses, the DAG will re-trigger itself and finish. If the check returns STOP, the DAG will finish and not retrigger itself. If the check returns PASS, the current DAG run will end in a successful state. If the check returns FAIL, the current DAG run will end in a failed state.

This allows the one to build "always-on" DAGs without having individual long blocking tasks.

This library is used to build airflow-supervisor, which uses supervisor as a process-monitor while checking and restarting jobs via airflow-ha.

Example - Recursive

You can also use this library to build recursive DAGs - or "Cyclic DAGs", despite the oxymoronic name.

The following code makes a DAG that triggers itself with some decrementing counter, starting with value 3:

with DAG(
    dag_id="test-ha-counter",
    description="Test HA Countdown",
    schedule=timedelta(days=1),
    start_date=datetime(2024, 1, 1),
    catchup=False,
):
    
    def _get_count(**kwargs):
        # The default is 3
        return kwargs['dag_run'].conf.get('counter', 3) - 1

    get_count = PythonOperator(task_id="get-count", python_callable=_get_count)

    def _keep_counting(**kwargs):
        count = kwargs["task_instance"].xcom_pull(key="return_value", task_ids="get-count")
        return (Result.PASS, Action.RETRIGGER) if count > 0 else (Result.PASS, Action.STOP) if count == 0 else (Result.FAIL, Action.STOP)

    keep_counting = HighAvailabilityOperator(
        task_id="ha",
        timeout=30,
        poke_interval=5,
        python_callable=_keep_counting,
        pass_trigger_kwargs={"conf": '''{"counter": {{ ti.xcom_pull(key="return_value", task_ids="get-count") }}}'''},
    )

    get_count >> keep_counting

License

This software is licensed under the Apache 2.0 license. See the LICENSE file for details.