Task does not retry when worker is killed due to OOM

[rawwar]

Apache Airflow version

main

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

No response

What happened?

When using CeleryExecutor, when a task gets killed due to OOM, it directly goes to failed state instead of retrying.

What you think should happen instead?

If retries are configured, it should go into "up_for_retry" state.

How to reproduce

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

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

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

Operating System

macos

Versions of Apache Airflow Providers

No response

Deployment

Astronomer

Deployment details

I tested this on Astronomer deployment as well as locally with breeze(on main branch) and was able to replicate this.

breeze command i used - breeze start-airflow -b postgres -P 17 --executor CeleryExecutor

Anything else?

On worker, I noticed this log:

2025-09-17 05:48:03.348301 [info     ] Task finished                  [supervisor] duration=3.201495791989146 exit_code=-9 final_state=failed

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

[amoghrajesh]

Let me assign this to myself, I will investigate a bit and try to get this fixed.

[amoghrajesh]

Investigated a bit and it seems like the TI gets failed from here: https://github.com/apache/airflow/blob/main/task-sdk/src/airflow/sdk/execution_time/supervisor.py#L1150.

My proposal would be to handle this slightly better, if a task has retries defined, we should retry if the signals are retry-able like the ones the hardware / containers can throw due to OOM / worker restart etc: SIGKILL (OOM), SIGTERM (worker restart) (all this obviously if the task is qualified to retry).

[ashb]

Under normal operation which process/component (task, supervisor, apiserver) is responsible for deciding in retries?

I know a task can say "fail and don't retry" but isn't that a different API endpoint we hit? If so shouldn't this "failed with signal" already be handled by the API server? 🤷🏻

Part of the issue is right now the Supervisor know nothing about the task. Now we could send that information from task to supervisor (but not on to API server) and then go "oh it failed with a signal/without an explicit final state". We could also use that "TaskInformation" message from the task to tell the supervisor about timeouts too for #53337

[amoghrajesh]

Under a normal operation (not killed my system: OOM or whatever), the API server sends in the decision if a task can retry or not with the run context response: https://github.com/apache/airflow/blob/main/airflow-core/src/airflow/api_fastapi/execution_api/routes/task_instances.py#L270, this is later tried on by the task runner itself.

So -- in short, api server decides it and the task sdk enforces it.

I know a task can say "fail and don't retry" but isn't that a different API endpoint we hit?

Nope, it is part of the same TI state update endpoint. We either send failed / up_for_retry to the API server through TI update API and it marks it as such.

The Problem with Signal based failures is that:

• API Server already decided should_retry in the run context
• Task runner never gets to enforce this decision (process killed)
• Leading to supervisor bypasses the decision and always returns failed for TI

[amoghrajesh]

I know a task can say "fail and don't retry" but isn't that a different API endpoint we hit? If so shouldn't this "failed with signal" already be handled by the API server? 🤷🏻

I do not think so, because it is the supervisor detecting such a failure directly (task doesn't tell it out loud), so the supervisor ends up sending a "failed" instead of an indication that it failed with a signal.