Conditional retry on callback

[cmpeburak]

Hi everyone,

I have a DAG that is used for multiple customers that subscribed to an app. The dag has multiple tasks of various operators. When a customer deletes the app, corresponding resources are removed from the database, and it's expected that the tasks will fail for that customer. What I want to do is to check the activity status of a customer on the failure of any task and skip the retry if the customer is not subscribed anymore.

I know that AirflowFailException can be used to skip retries on failure; however, I want this behavior to be generic so that it applies on failures/retries of all tasks and operators (not only the python operator)

I've created a function as below and passed it as on_retry_callback argument of the DAG.

def dag_retry(context):
    print("Running dag retry callback.")

    # Do not retry if a customer is not active.
    customer = context["params"]["customer"]
    if not check_if_customer_active(customer):
        raise AirflowFailException(
            f"Customer {customer} is not active anymore! "
            f"No more retry is needed."
        )

When a task of an unsubscribed customer fails for the first time, this function is called, and AirflowFailException is raised as expected. However, the task is still retried. I think the retry is already triggered when it comes to that function.

Do you have any idea how to solve this?

Thanks in advance.

[Taragolis]

You could try to set status of Task Instance to failed

from airflow.utils.state import TaskInstanceState


def retry_callback(context) -> None:
    if True:  # some condition
        ti = context["ti"]
        print(f"Set task instance {ti} state to {TaskInstanceState.FAILED}")
        ti.set_state(TaskInstanceState.FAILED)
        return

[cmpeburak]

That worked, thanks a lot!

[edualvess]

This is a true life safer to this day, a very clever and elegant solution. Thanks for sharing!

[rwhitt2049]

I made a slight tweak on the above that will ensure that any on_failure_callbacks still get called. Just setting the instance state prevented retries, but any failure callbacks didn't get executed. (Tested in Airflow 2.6.3)

def retry_callback(context):
    if True:  # some condition
        ti = context["ti"]
        ti.handle_failure(error=None, force_fail=True).  # this will set the state to failed and call the on_failure_callback

[ndellosa95]

Is there any way to do this in Airflow 3? The handle_failure isn't available on the RuntimeTaskInstanceProtocol type in Airflow 3.