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

airflow-core/docs/administration-and-deployment/listeners.rst

@@ -103,6 +103,28 @@ of a :class:`~airflow.sdk.execution_time.task_runner.RuntimeTaskInstance` instan
     :start-after: [START howto_listen_ti_skipped_task]
     :end-before: [END howto_listen_ti_skipped_task]
 
+Failure cause
+"""""""""""""
+
+``on_task_instance_failed`` accepts an optional ``failure_kind`` argument, a
+:class:`~airflow_shared.state.TaskFailureKind` — ``INFRA`` (an infrastructure disruption),
+``APPLICATION`` (the task's own code), ``TIMEOUT``, or ``MANUAL`` (an operator set it failed). It is
+``None`` when the cause was not classified. It is a ``str`` enum, so it compares equal to
+``"infra"`` / ``"application"`` / ``"timeout"`` / ``"manual"``. For an ``INFRA`` failure the
+executor's reason token (``Evicted`` / ...) arrives as the transient ``reason`` argument. A listener that
+wants the cause declares the arguments:
+
+.. code-block:: python
+
+    @hookimpl
+    def on_task_instance_failed(self, previous_state, task_instance, error, failure_kind, reason):
+        if failure_kind == "infra":
+            # an infrastructure disruption, not the task's own code
+            alert(reason)
+
+A listener that does not declare the argument keeps working — the parameter is dispatched by name, so
+existing listeners need no migration.
+
 Asset Events
 --------------
 

airflow-core/newsfragments/66405.feature.rst

@@ -0,0 +1 @@
+Add an opt-in infrastructure retry budget (``[core] infra_failure_refund_retries``, bounded by ``max_infra_refunds``) so a failure an executor classifies as infrastructure — a Kubernetes pod ``Evicted`` or preempted by the node — does not spend a task's ``retries``. Disabled by default; a failure the executor cannot positively classify as infra spends the retry as before.

airflow-core/newsfragments/66405.significant.rst

@@ -0,0 +1,31 @@
+Add ``failure_kind`` and ``reason`` arguments to the ``on_task_instance_failed`` listener hook, so a listener can tell an infrastructure disruption from the task's own failure (AIP-97 foundation).
+
+Today the listener only sees the worker-side ``error`` exception; failure
+causes that originate outside the worker process — pod ``OOMKilled`` /
+``PodEvicted`` on Kubernetes, a lost heartbeat, an oom-killer ``SIGKILL`` on
+the local executor — look the same as a real bug.
+
+``failure_kind`` is a small ``TaskFailureKind`` enum (a ``str`` enum, so it
+compares equal to ``"infra"`` / ``"application"`` / ``"timeout"`` / ``"manual"``). The executor's
+short reason token (``OOMKilled`` / ``Evicted`` / ...) arrives as the transient
+``reason`` argument, handed to the listener at failure time rather than
+persisted, so the change needs no schema migration.
+
+.. code-block:: python
+
+    from airflow.listeners import hookimpl
+
+
+    class InfraTrackingListener:
+        @hookimpl
+        def on_task_instance_failed(self, previous_state, task_instance, error, failure_kind, reason):
+            if failure_kind == "infra":
+                # route the reason token to a capacity-planning alert
+                ...
+
+The Kubernetes executor classifies a failed pod: a node-level disruption
+(``Evicted`` / ``Preempting`` / ``NodeShutdown`` / ...) is ``"infra"``; a
+container that ended on its own — an app crash, or an ``OOMKilled`` against
+its *own* limit — is ``"application"``. pluggy dispatches by parameter name, so
+existing ``hookimpl`` implementations that don't declare ``failure_kind`` or
+``reason`` keep working unchanged.

airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py

@@ -28,7 +28,7 @@
 from sqlalchemy.orm import joinedload
 from sqlalchemy.orm.session import Session
 
-from airflow._shared.state import TaskScope
+from airflow._shared.state import TaskFailureKind, TaskScope
 from airflow.api_fastapi.app import get_auth_manager
 from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails
 from airflow.api_fastapi.common.dagbag import DagBagDep, get_latest_version_of_dag
@@ -117,6 +117,8 @@ def _emit_state_listener_hooks(updated_tis: list[TI], new_state: str | TaskInsta
                     previous_state=None,
                     task_instance=ti,
                     error=f"TaskInstance's state was manually set to `{TaskInstanceState.FAILED}`.",
+                    failure_kind=TaskFailureKind.MANUAL,
+                    reason=None,
                 )
             elif new_state == TaskInstanceState.SKIPPED:
                 get_listener_manager().hook.on_task_instance_skipped(previous_state=None, task_instance=ti)

airflow-core/src/airflow/executors/base_executor.py

@@ -72,6 +72,7 @@ def get_execution_api_server_url(conf_source: AirflowConfigParser | ExecutorConf
     from sqlalchemy.orm import Session
 
     from airflow._shared.logging.remote import StreamingLogResponse
+    from airflow._shared.state import TaskFailureKind
     from airflow.api_fastapi.auth.tokens import JWTGenerator
     from airflow.callbacks.base_callback_sink import BaseCallbackSink
     from airflow.callbacks.callback_requests import CallbackRequest
@@ -228,6 +229,7 @@ def __init__(self, parallelism: int = PARALLELISM, team_name: str | None = None)
         self.queued_connection_tests: dict[ConnectionTestKey, workloads.TestConnection] = {}
         self.running: set[WorkloadKey] = set()
         self.event_buffer: dict[WorkloadKey, EventBufferValueType] = {}
+        self.task_failure_info: dict[WorkloadKey, tuple[TaskFailureKind, str | None]] = {}
         self._task_event_logs: deque[Log] = deque()
         self.conf = ExecutorConf(team_name)
 
@@ -550,6 +552,14 @@ def get_event_buffer(self, dag_ids=None) -> dict[WorkloadKey, EventBufferValueTy
 
         return cleared_events
 
+    def get_task_failure_info(self, key: WorkloadKey) -> tuple[TaskFailureKind, str | None] | None:
+        """
+        Return and clear the ``(failure_kind, reason)`` this executor classified for ``key``.
+
+        None unless the executor can read the real cause, as the Kubernetes executor does.
+        """
+        return self.task_failure_info.pop(key, None)
+
     def get_task_log(self, ti: TaskInstance, try_number: int) -> tuple[list[str], list[str]]:
         """
         Return the task logs.

airflow-core/src/airflow/jobs/scheduler_job_runner.py

@@ -56,6 +56,7 @@
 
 from airflow import settings
 from airflow._shared.observability.metrics import stats
+from airflow._shared.state import TaskFailureKind
 from airflow._shared.timezones import timezone
 from airflow.api_fastapi.execution_api.datamodels.taskinstance import DagRun as DRDataModel, TIRunContext
 from airflow.assets.evaluation import AssetEvaluator
@@ -1468,6 +1469,11 @@ def process_executor_events(
                     job_id,
                 )
             state, info = event_buffer.pop(buffer_key)
+            # Pop unconditionally, not only on the killed-externally branch below, so a
+            # self-reporting task can't leak its entry.
+            executor_failure_kind: tuple[TaskFailureKind, str | None] | None = executor.get_task_failure_info(
+                ti.key
+            )
 
             if state in (TaskInstanceState.QUEUED, TaskInstanceState.RUNNING):
                 ti.external_executor_id = info
@@ -1663,7 +1669,19 @@ def process_executor_events(
                     executor.send_callback(email_request)
 
                 # Update task state - emails are handled by DAG processor now
-                ti.handle_failure(error=msg, session=session)
+                # Refund only on a positive infra classification. An executor that reports
+                # nothing but a state can't tell an eviction from a silent crash, so it
+                # stays unclassified and spends the retry as today.
+                if executor_failure_kind is not None:
+                    failure_kind, reason = executor_failure_kind
+                else:
+                    failure_kind, reason = None, None
+                ti.handle_failure(
+                    error=msg,
+                    session=session,
+                    failure_kind=failure_kind,
+                    reason=reason if failure_kind == TaskFailureKind.INFRA else None,
+                )
 
         cls._emit_executor_events_batch_metrics(num_events)
         return len(event_buffer)

airflow-core/src/airflow/models/taskinstance.py

@@ -75,6 +75,7 @@
     new_dagrun_trace_carrier,
     new_task_run_carrier,
 )
+from airflow._shared.state import TaskFailureKind
 from airflow._shared.timezones import timezone
 from airflow.assets.manager import asset_manager
 from airflow.configuration import conf
@@ -557,6 +558,48 @@ def _date_or_empty(*, task_instance: TaskInstance, attr: str) -> str:
     return result.strftime("%Y%m%dT%H%M%S") if result else ""
 
 
+def _maybe_refund_infra_attempt(
+    *,
+    task_instance: TaskInstance,
+    task: Operator | None,
+    failure_kind: TaskFailureKind | None,
+    reason: str | None = None,
+) -> bool:
+    """
+    Refund one retry attempt (bump ``max_tries``) for an infra-classified failure.
+
+    Opt-in and capped by config. Every other kind, and an unclassified failure,
+    spends the user's retry as before. Returns True if an attempt was refunded.
+
+    :meta private:
+    """
+    # retries may be None (unset); treat it as falsy like is_eligible_to_retry does.
+    # Refunds still apply at retries=0: an infra disruption is not the user's failure,
+    # so it earns an attempt back regardless of their retry budget.
+    retries = getattr(task, "retries", None) or 0
+    if (
+        failure_kind != TaskFailureKind.INFRA
+        or task is None
+        or not conf.getboolean("core", "infra_failure_refund_retries", fallback=False)
+    ):
+        return False
+    max_infra_refunds = conf.getint("core", "max_infra_refunds", fallback=3)
+    refunds_used = max((task_instance.max_tries or 0) - retries, 0)
+    if refunds_used >= max_infra_refunds:
+        return False
+    task_instance.max_tries += 1
+    log.info(
+        "infra failure (%s) refunded attempt for %s; max_tries now %s "
+        "(refund %s/%s), user retry budget preserved",
+        reason,
+        task_instance,
+        task_instance.max_tries,
+        refunds_used + 1,
+        max_infra_refunds,
+    )
+    return True
+
+
 def uuid7() -> UUID:
     """Generate a new UUID7."""
     return uuid6.uuid7()
@@ -650,7 +693,6 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload):
     # Cleared on task start (ti_run).  Read by next_retry_datetime().
     retry_delay_override: Mapped[float | None] = mapped_column(Float, nullable=True)
     retry_reason: Mapped[str | None] = mapped_column(String(500), nullable=True)
-
     __table_args__ = (
         Index("ti_dag_state", dag_id, state),
         Index("ti_dag_run", dag_id, run_id),
@@ -1849,6 +1891,8 @@ def fetch_handle_failure_context(
         *,
         session: Session,
         fail_fast: bool = False,
+        failure_kind: TaskFailureKind | None = None,
+        reason: str | None = None,
     ):
         """
         Fetch the context needed to handle a failure.
@@ -1858,6 +1902,12 @@ def fetch_handle_failure_context(
         :param test_mode: doesn't record success or failure in the DB if True
         :param session: SQLAlchemy ORM Session
         :param fail_fast: if True, fail all downstream tasks
+        :param failure_kind: what caused the failure (:class:`TaskFailureKind` or
+            ``None``). ``INFRA`` refunds the attempt instead of charging the user's
+            retry budget, gated by config.
+        :param reason: the executor's token for an infra failure (e.g. ``Evicted``).
+            The worker reported no error in that case, so this is passed to the
+            listener rather than persisted.
         """
         if error:
             cls.logger().error("%s", error)
@@ -1889,6 +1939,8 @@ def fetch_handle_failure_context(
         # Actual callbacks are handled by the DAG processor, not the scheduler
         task = getattr(ti, "task", None)
 
+        _maybe_refund_infra_attempt(task_instance=ti, task=task, failure_kind=failure_kind, reason=reason)
+
         if not ti.is_eligible_to_retry():
             ti.state = TaskInstanceState.FAILED
 
@@ -1905,7 +1957,11 @@ def fetch_handle_failure_context(
 
         try:
             get_listener_manager().hook.on_task_instance_failed(
-                previous_state=TaskInstanceState.RUNNING, task_instance=ti, error=error
+                previous_state=TaskInstanceState.RUNNING,
+                task_instance=ti,
+                error=error,
+                failure_kind=failure_kind,
+                reason=reason,
             )
         except Exception:
             log.exception("error calling listener")
@@ -1927,13 +1983,19 @@ def handle_failure(
         test_mode: bool | None = None,
         *,
         session: Session = NEW_SESSION,
+        failure_kind: TaskFailureKind | None = None,
+        reason: str | None = None,
     ) -> None:
         """
         Handle Failure for a task instance.
 
         :param error: if specified, log the specific exception if thrown
         :param test_mode: doesn't record success or failure in the DB if True
         :param session: SQLAlchemy ORM Session
+        :param failure_kind: what caused the failure (:class:`TaskFailureKind` or
+            ``None``), forwarded to the listener and the retry decision.
+        :param reason: the executor's token for an infra failure, passed to the
+            listener rather than persisted.
         """
         if TYPE_CHECKING:
             assert self.task
@@ -1950,6 +2012,8 @@ def handle_failure(
             test_mode=test_mode,
             session=session,
             fail_fast=fail_fast,
+            failure_kind=failure_kind,
+            reason=reason,
         )
 
         _log_state(task_instance=self)
@@ -1963,15 +2027,9 @@ def is_eligible_to_retry(self) -> bool:
             # If a task is cleared when running, it goes into RESTARTING state and is always
             # eligible for retry
             return True
-        if not getattr(self, "task", None):
-            # Couldn't load the task, don't know number of retries, guess:
-            return self.try_number <= self.max_tries
-
-        if TYPE_CHECKING:
-            assert self.task
-            assert self.task.retries
-
-        return bool(self.task.retries and self.try_number <= self.max_tries)
+        # Mirror of the execution API's _is_eligible_to_retry; the scheduler and worker
+        # retry-decision paths must not diverge. See its comment for why retries is not consulted.
+        return self.max_tries != 0 and self.try_number <= self.max_tries
 
     def set_duration(self) -> None:
         """Set task instance duration."""

airflow-core/tests/unit/jobs/test_scheduler_job.py

@@ -45,6 +45,7 @@
 from airflow import settings
 from airflow._shared.module_loading import qualname
 from airflow._shared.observability.metrics.base_stats_logger import StatsLogger
+from airflow._shared.state import TaskFailureKind
 from airflow._shared.timezones import timezone
 from airflow.api_fastapi.auth.tokens import JWTGenerator
 from airflow.assets.manager import AssetManager
@@ -536,6 +537,42 @@ def test_process_executor_events(self, mock_get_backend, mock_task_callback, dag
             any_order=True,
         )
 
+    @conf_vars({("core", "infra_failure_refund_retries"): "True", ("core", "max_infra_refunds"): "3"})
+    @mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest", spec=TaskCallbackRequest)
+    def test_process_executor_events_infra_classification(self, mock_task_callback, dag_maker):
+        """Only a positively-classified infra failure refunds; an unclassified
+        state-mismatch death (no executor bridge) spends the retry, and a bridge 'user'
+        classification (e.g. an app OOM) does not refund."""
+        session = settings.Session()
+
+        def _run(dag_id: str, stashed: tuple[TaskFailureKind, str | None] | None) -> TaskInstance:
+            with dag_maker(dag_id=dag_id, fileloc=f"/{dag_id}/"):
+                task = EmptyOperator(task_id="t", retries=1)
+            ti = dag_maker.create_dagrun().get_task_instance(task.task_id)
+            ti.state = State.QUEUED
+            session.merge(ti)
+            session.commit()
+            executor = MockExecutor(do_update=False)
+            job_runner = SchedulerJobRunner(job=Job(), executors=[executor])
+            executor.event_buffer[ti.key] = State.FAILED, None
+            if stashed is not None:
+                executor.task_failure_info[ti.key] = stashed
+            job_runner._process_executor_events(executor=executor, session=session)
+            ti.refresh_from_db(session=session)
+            return ti
+
+        # Unclassified: no refund, exactly as today.
+        ti = _run("aip97_unclassified", stashed=None)
+        assert ti.max_tries == 1
+
+        # Classified infra (a pod Evicted): refunded.
+        ti = _run("aip97_infra", stashed=(TaskFailureKind.INFRA, "Evicted"))
+        assert ti.max_tries == 2
+
+        # Classified application (an OOM against its own limit): no refund.
+        ti = _run("aip97_app_oom", stashed=(TaskFailureKind.APPLICATION, "OOMKilled"))
+        assert ti.max_tries == 1
+
     @mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest", spec=TaskCallbackRequest)
     def test_process_executor_events_restarting_cleared_task(self, mock_task_callback, dag_maker):
         """
@@ -9623,6 +9660,7 @@ def test_scheduler_passes_context_from_server_on_task_failure(self, dag_maker, s
         # Mock the executor to simulate a task failure
         mock_executor = MagicMock(spec=BaseExecutor)
         mock_executor.has_task = mock.MagicMock(return_value=False)
+        mock_executor.get_task_failure_info.return_value = None
         scheduler_job = Job()
         self.job_runner = SchedulerJobRunner(scheduler_job, executors=[mock_executor])