AIP-97: Disruption Readiness (Transparent Infra Failure Retries through Context Capture and Propagation)

[1fanwang]

The implementation behind AIP-97: Disruption Readiness.

Why

When a task is killed from outside (eviction, OOM, force-delete, lost heartbeat), the worker dies without raising. An eviction and a real ValueError therefore look the same to on_task_instance_failed, to lineage, and to the retry counter. Airflow already requeues a pre-start pod without spending a retry; this extends that behavior to running tasks.

What

Pillar 1 adds an optional failure_kind enum (infra, application, timeout, manual) and a short reason token to the hook. Pluggy matches by name, so listeners that do not declare the new arguments are unchanged. The Kubernetes executor's classify_pod_failure() reads pod details it already collects, with its reason set pinned to the kubelet's Go definitions. Node drains and preemptions leave that reason empty, so it also reads pod.status.conditions[type=DisruptionTarget].

Pillar 2 is opt-in and off by default: on an infra failure, refund the attempt instead of charging the user's retries, reusing max_tries and capping at max_infra_refunds.

Retry eligibility

The refund lifts max_tries, so is_eligible_to_retry now reads max_tries rather than the user's retries. That matches the Execution API's _is_eligible_to_retry, which already dropped retries, so the two retry paths stop disagreeing and a retries=0 task gets the refund. Eligibility differs only once a refund lifts max_tries past retries.

Alternatives set aside: an infra_retry_count column would bring back the migration this avoids; a handle_failure bypass leaves the eligibility paths divergent; checkpoint-resume needs AIP-103 state this design does not keep.

No DB change

The reason reaches the listener as a transient argument, not a column, so there is no migration and it backports cleanly. Whether it rides reason or folds into error is open (#23 vs #22).

Related POCs

failure_kind also reaches callbacks, metrics (live OTel e2e), and scheduler logs. Other executor paths are POC'd for Celery and heartbeat, with config and listener docs.

Testing

Unit: AIP-97 tests pass across core, Kubernetes, the scheduler, and the task SDK; existing scheduler retry-eligibility tests are unchanged; mypy is clean.

Live on real Postgres 16, through the real (not mocked) handle_failure(failure_kind=...) path on 3ce5dbc184; core and task SDK runtime are unchanged since. A registered listener confirms it receives kind and reason. airflow db migrate stamps to the existing head:

retries=0 scenario	state	max_tries	eligible
infra, refund on	up_for_retry	0 to 1	yes
infra, refund off	failed	0	no
application, refund on	failed	0	no

Commands and raw logs

$ docker run -d --name pg -e POSTGRES_USER=airflow -e POSTGRES_PASSWORD=airflow     -e POSTGRES_DB=airflow -p 55432:5432 postgres:16
$ export AIRFLOW__DATABASE__SQL_ALCHEMY_CONN=    "******localhost:55432/airflow"

$ airflow db migrate
  Running stamp_revision  -> 02ca36b0235b
  Database migration done!                       # existing head, no new column

$ pytest airflow-core/tests/unit/models/test_e2e_retries0_live.py -q -s
[e2e] backing store: localhost:55432/airflow
  start (retries=0)      state=None          max_tries=0 eligible=False
  after infra kill       state=up_for_retry  max_tries=1 eligible=True
  flag off, infra kill   state=failed        max_tries=0 eligible=False
  app failure            state=failed        max_tries=0 eligible=False

# the same run, real state-machine transitions (pytest -s):
  Marking task as UP_FOR_RETRY. dag_id=e2e_infra_on   task_id=op
  Marking task as FAILED.       dag_id=e2e_infra_off  task_id=op
  Marking task as FAILED.       dag_id=e2e_app        task_id=op
1 passed

Kubernetes, live on a v1.35.0 cluster. A node drain and a preemption are indistinguishable from a crash except for the condition:

disruption	pod.status.reason	container	DisruptionTarget	kind
node drain (taint)	empty	Error, exit 143	True/DeletionByTaintManager	infra
preemption	empty	Error, exit 143	True/PreemptionByScheduler	infra

Commands and raw logs

$ kind create cluster --name aip97 --image kindest/node:v1.35.0   # 1 control-plane, 2 workers
$ kubectl taint node aip97-worker node.kubernetes.io/unschedulable=:NoExecute   # what a drain does

# every watch event on the victim pod, node drain:
  12.4s  MODIFIED  phase=Running  deletionTS=-    status.reason=None  DisruptionTarget=True/DeletionByTaintManager  container=-
  17.8s  MODIFIED  phase=Failed   deletionTS=SET  status.reason=None  DisruptionTarget=True/DeletionByTaintManager  container=term:Error:exit=143
  18.1s  DELETED   phase=Failed   deletionTS=SET  status.reason=None  DisruptionTarget=True/DeletionByTaintManager  container=term:Error:exit=143

# scheduler preemption, same shape:
  12.5s  MODIFIED  phase=Running  deletionTS=-    status.reason=None  DisruptionTarget=True/PreemptionByScheduler   container=-
  17.7s  MODIFIED  phase=Failed   deletionTS=SET  status.reason=None  DisruptionTarget=True/PreemptionByScheduler   container=term:Error:exit=143
  18.2s  DELETED   phase=Failed   deletionTS=SET  status.reason=None  DisruptionTarget=True/PreemptionByScheduler   container=term:Error:exit=143

# the pod object recorded above, replayed through the classifier:
$ pytest providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_classify_pod_failure.py -q
24 passed
# on the pre-fix source the 8 disruption cases fail; both above classify as application.

Both providers still import on Airflow 3.3.0; classify_pod_failure is gated to 3.4+.