Task-level failure callbacks are not triggered when tasks are SIGKILLed

[manipatnam]

Apache Airflow version

3.1.6

If "Other Airflow 3 version" selected, which one?

3.1.2

What happened?

When a task is OOM-killed (container/process terminated with SIGKILL), the task appears to fail silently:

Task-level on_failure_callback is not executed

In logs/behavior, the executor reports the task as SUCCESS

Because the scheduler never observes a FAILED state (and no mismatch is detected), it does not trigger task-level failure handling/callbacks

What you think should happen instead?

When a task process is terminated with SIGKILL:

The executor should report the task as FAILED (or otherwise indicate abnormal termination)
Airflow should then trigger task-level on_failure_callback (and any related failure handling)

How to reproduce

Create a task which can OOM and has task level failure call back and it wont be triggered

from airflow.operators.python import PythonOperator
from datetime import datetime

def func():
    a = "asd"
    while True:
        a += a*100000

def task_failure_callback(context):
    """Callback function called when a task fails."""
    task_instance = context.get('task_instance')
    dag_run = context.get('dag_run')
    print(f"Task {task_instance.task_id} failed in DAG {dag_run.dag_id} with run_id {dag_run.run_id}")


with DAG(
    dag_id="oom_example",
    start_date=datetime(2024, 1, 1),
    schedule="@daily",
    catchup=False,
    doc_md=__doc__,
    tags=["oom"],
) as dag:
    oom_task = PythonOperator(
        task_id="oom_task",
        python_callable=func,
        retries=0,
        on_failure_callback=task_failure_callback,
    )

Operating System

debian

Versions of Apache Airflow Providers

No response

Deployment

Astronomer

Deployment details

No response

Anything else?

No response

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[Prajwal7842]

I think on_failure_callback is meant to be used for a graceful exit, but given SIGKILL is meant to be non-graceful and we cannot catch it in python, the behavior seems consistent with force exit.

Update: When trying out this dag, I am able to see that event logs show a state mismatch.

[jroachgolf84]

@RNHTTR might be able to provide a bit more detail around this.

[Prajwal7842]

I would like to investigate and work on this. I found this similar discussion related to task retries: #55753

[potiuk]

There is not much to investigate - the on_failure_callback is not going to be called when task is killed - this is a given because the interpreter is killed without any chance for reaction.

The only way to handle that is to make completely new flow where the supervisor should see if the task has been killed forecefully and **start a "run on_failure_callback" new task that will only do that. Supervisor cannot do it on it's own because supervisor does not parse the Dag that has on_failure_callback defined, so effectively it will need to run a new child task that will again parse the dag and only runs on_failure_callback. This has a number of edge cases to look at - for example handling the case where OOM happens during parsing the dag,

This is quite clear what's going on here, it just needs quite complex implementation of handling this special flow.

[jroachgolf84]

@potiuk, I've always wondered about this. What do you think about an on_kill_callback? I assume that would follow the framework of what you mentioned above. But curious your thoughts.

[potiuk]

@potiuk, I've always wondered about this. What do you think about an on_kill_callback? I assume that would follow the framework of what you mentioned above. But curious your thoughts.

Yes it could be posible, though I would rather extend the on_failure_callback and add type of failure there if possible.

Actually .. I think it will not work in all cases anyway. If the task process is just killed with -9 or SEGV, this might work - but it is not likely to work with OOM - because generally OOM kills the whole Pod - which means that both task and supervisor wil get -9 (there might be some races) and neither of them will have a chance to react. What you really want to do in this case is to have an executor and deployment dependent behaviour - for example K8S executor should monitor reason for POD failue - and then not even start another task, but even start another POD to run on_task_failure (or on_task_kill) - and in case of Celery Executor for example you would have to track which celery workers run which tasks, and do similar thing for alltasks that were runnning in the same worker .. Etc. etc.

And it will look differently if airflow is run using different cluster technoloty - for example Nomad.

There are lots and lots of flows that would need to be considered to make it robust and it's I think impossible to have a working solution without close coupling the solutions with deployment or limiting to single deployment configuration (say K8S executor).

I would say that this one is "won't fix" kind of issues - and we need to simply accept there is no on_failure_callback generated in case of SIGKILL