From 188bb01c8a91f88085b3dfa7fb3aceeaf89da35a Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Mon, 27 Jul 2026 16:48:52 -0700 Subject: [PATCH 1/6] AIP-97: Disruption Readiness: classify why a task failed and refund infra retries (no DB) Pillar 1: add a `failure_kind` argument (infra / application / timeout / manual) and a transient `reason` token to the `on_task_instance_failed` listener hook. The worker sets application or timeout for the failures it catches; the scheduler sets infra for a task killed from outside. The Kubernetes executor classifies a failed pod (Evicted / preempted is infra; an app crash or an OOMKill against the container's own limit is application). pluggy dispatches by name, so a listener that does not declare the arguments is unchanged. Pillar 2: a failure classified as infra gives back the retry it used (reuses `max_tries`, bounded by `[core] max_infra_refunds`), so infra churn does not spend a task's `retries`. Opt-in via `[core] infra_failure_refund_retries`, off by default. This holds even at `retries=0`: `is_eligible_to_retry` now gates on `max_tries` (the retry ceiling the refund bumps), matching the execution API's `_is_eligible_to_retry`, rather than short-circuiting on the user's `retries`. The two retry-decision paths no longer disagree, and no per-task counter or column is needed. No schema change: the reason is handed to the listener at failure time rather than persisted, so this backports cleanly to older release branches with no upgrade or rollback step. How the reason token best reaches the listener (this transient `reason` argument, or folding it into the existing `error` argument) is AIP-97 open question 3. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Signed-off-by: 1fanwang <1fannnw@gmail.com> --- .../listeners.rst | 22 +++ airflow-core/newsfragments/66405.feature.rst | 1 + .../newsfragments/66405.significant.rst | 31 ++++ .../services/public/task_instances.py | 5 +- .../src/airflow/executors/base_executor.py | 12 ++ .../src/airflow/jobs/scheduler_job_runner.py | 21 ++- .../src/airflow/models/taskinstance.py | 79 ++++++++-- .../tests/unit/jobs/test_scheduler_job.py | 38 +++++ .../unit/listeners/test_listener_types.py | 75 +++++++++ .../test_taskinstance_failure_kind_sources.py | 70 +++++++++ .../models/test_taskinstance_infra_refund.py | 132 ++++++++++++++++ .../tests_common/test_utils/version_compat.py | 2 + .../executors/kubernetes_executor.py | 11 ++ .../executors/kubernetes_executor_utils.py | 56 ++++++- .../cncf/kubernetes/version_compat.py | 2 + .../executors/test_classify_pod_failure.py | 147 ++++++++++++++++++ .../listeners/spec/taskinstance.py | 27 +++- .../src/airflow_shared/state/__init__.py | 23 +++ .../airflow/sdk/execution_time/task_runner.py | 15 +- 19 files changed, 751 insertions(+), 18 deletions(-) create mode 100644 airflow-core/newsfragments/66405.feature.rst create mode 100644 airflow-core/newsfragments/66405.significant.rst create mode 100644 airflow-core/tests/unit/listeners/test_listener_types.py create mode 100644 airflow-core/tests/unit/models/test_taskinstance_failure_kind_sources.py create mode 100644 airflow-core/tests/unit/models/test_taskinstance_infra_refund.py create mode 100644 providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_classify_pod_failure.py diff --git a/airflow-core/docs/administration-and-deployment/listeners.rst b/airflow-core/docs/administration-and-deployment/listeners.rst index 70dde0b7fd2af..fa39a4c5d2164 100644 --- a/airflow-core/docs/administration-and-deployment/listeners.rst +++ b/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 -------------- diff --git a/airflow-core/newsfragments/66405.feature.rst b/airflow-core/newsfragments/66405.feature.rst new file mode 100644 index 0000000000000..5450d2b33ee5b --- /dev/null +++ b/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. diff --git a/airflow-core/newsfragments/66405.significant.rst b/airflow-core/newsfragments/66405.significant.rst new file mode 100644 index 0000000000000..7aa5d38568229 --- /dev/null +++ b/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. diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py b/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py index b00ba4effdbf6..fc2764cceaad0 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py +++ b/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 @@ -113,10 +113,13 @@ def _emit_state_listener_hooks(updated_tis: list[TI], new_state: str | TaskInsta if new_state == TaskInstanceState.SUCCESS: get_listener_manager().hook.on_task_instance_success(previous_state=None, task_instance=ti) elif new_state == TaskInstanceState.FAILED: + # An operator set this failed — a manual action, never refunded. get_listener_manager().hook.on_task_instance_failed( 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) diff --git a/airflow-core/src/airflow/executors/base_executor.py b/airflow-core/src/airflow/executors/base_executor.py index 8030e9c14455e..c430cfc90d52a 100644 --- a/airflow-core/src/airflow/executors/base_executor.py +++ b/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,8 @@ 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] = {} + # An executor's (failure_kind, reason) for a key, read once by the scheduler. + self.task_failure_info: dict[WorkloadKey, tuple[TaskFailureKind, str | None]] = {} self._task_event_logs: deque[Log] = deque() self.conf = ExecutorConf(team_name) @@ -550,6 +553,15 @@ def get_event_buffer(self, dag_ids=None) -> dict[WorkloadKey, EventBufferValueTy return cleared_events + def get_task_failure_info(self, key: WorkloadKey): + """ + Return and clear the ``(failure_kind, reason)`` this executor classified for ``key``. + + Populated by executors that can read the real failure cause (e.g. the Kubernetes + executor telling an ``Evicted`` pod from an ``OOMKilled`` one). Returns None otherwise. + """ + 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. diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index daef05b2b7c07..a702493b530af 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/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 @@ -1467,6 +1468,10 @@ def process_executor_events( job_id, ) state, info = event_buffer.pop(buffer_key) + # Pop the executor's (failure_kind, reason) for this event now, so the entry + # is always cleared — not only on the killed-externally branch below — which + # keeps self-reporting tasks from leaking entries. + executor_failure_kind = executor.get_task_failure_info(ti.key) if state in (TaskInstanceState.QUEUED, TaskInstanceState.RUNNING): ti.external_executor_id = info @@ -1662,7 +1667,21 @@ 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) + # The worker died from outside without reporting an error. Refund only + # when an executor positively classified this as infra (its bridge read + # the real cause); an executor that only reports 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) diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index 0dd10507fc863..0dfe60466f8da 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/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,44 @@ 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, task, failure_kind, reason=None) -> bool: + """ + Refund one retry attempt (bump ``max_tries``) for an infra-classified failure. + + Fires only when ``failure_kind == INFRA``, ``[core] infra_failure_refund_retries`` + is on, and the ``[core] max_infra_refunds`` cap is not yet reached. Every other + kind — and an unclassified (``None``) failure — spends the user's retry. Returns + True if an attempt was refunded. + + :meta private: + """ + # retries may be None (unset); treat it as falsy like is_eligible_to_retry does. + # The refund still applies at retries=0: an infra disruption is not the user's + # failure, so it earns an attempt back regardless of the user's 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 +689,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), @@ -1842,6 +1880,8 @@ def fetch_handle_failure_context( *, session: Session, fail_fast: bool = False, + failure_kind: str | None = None, + reason: str | None = None, ): """ Fetch the context needed to handle a failure. @@ -1851,6 +1891,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 never reported an error in that case, so it is surfaced as the + failure's ``reason`` argument to the listener rather than persisted. """ if error: cls.logger().error("%s", error) @@ -1882,6 +1928,11 @@ def fetch_handle_failure_context( # Actual callbacks are handled by the DAG processor, not the scheduler task = getattr(ti, "task", None) + # An infra-classified failure refunds the attempt instead of spending the + # user's retry budget (opt-in, capped). Extends the Kubernetes executor's + # pre-execution requeue-without-consuming-a-retry to the mid-execution case. + _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 @@ -1898,7 +1949,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") @@ -1920,6 +1975,8 @@ def handle_failure( test_mode: bool | None = None, *, session: Session = NEW_SESSION, + failure_kind: str | None = None, + reason: str | None = None, ) -> None: """ Handle Failure for a task instance. @@ -1927,6 +1984,10 @@ def handle_failure( :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, surfaced to the + listener as the failure's ``error`` (not persisted). """ if TYPE_CHECKING: assert self.task @@ -1943,6 +2004,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) @@ -1956,15 +2019,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.""" diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 4add965a6abda..62ada334f23b4 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -43,6 +43,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 @@ -534,6 +535,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, stashed): + 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 (executor only reports a state) → no refund, exactly as today. + ti = _run("aip97_unclassified", stashed=None) + assert ti.max_tries == 1 + + # Bridge classified it infra (e.g. a pod Evicted) → refunded. + ti = _run("aip97_infra", stashed=(TaskFailureKind.INFRA, "Evicted")) + assert ti.max_tries == 2 + + # Bridge classified it user (e.g. an app 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): """ @@ -9281,6 +9318,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]) diff --git a/airflow-core/tests/unit/listeners/test_listener_types.py b/airflow-core/tests/unit/listeners/test_listener_types.py new file mode 100644 index 0000000000000..09cf01d229985 --- /dev/null +++ b/airflow-core/tests/unit/listeners/test_listener_types.py @@ -0,0 +1,75 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from airflow.listeners import hookimpl +from airflow.listeners.listener import get_listener_manager + + +class TestOnTaskInstanceFailedAcceptsFailureKind: + """The on_task_instance_failed hookspec accepts the optional ``failure_kind`` + argument (AIP-97) so listener authors can opt in to the failure source + (``infra`` / ``application`` / ``timeout`` / ``manual``) without parsing the error.""" + + def test_listener_with_failure_kind_receives_it(self, listener_manager): + received: dict[str, str | None] = {"failure_kind": None} + + # Per the hookspec docstring, listener implementations must declare + # failure_kind WITHOUT a default value — pluggy treats the impl + # default as authoritative and silently overrides the caller's value. + class InfraListener: + @hookimpl + def on_task_instance_failed( + self, + previous_state, + task_instance, + error, + failure_kind, + ): + received["failure_kind"] = failure_kind + + listener_manager(InfraListener()) + + get_listener_manager().hook.on_task_instance_failed( + previous_state=None, + task_instance=None, + error=RuntimeError("boom"), + failure_kind="infra", + ) + + assert received["failure_kind"] == "infra" + + def test_listener_without_failure_kind_param_keeps_working(self, listener_manager): + """Pluggy dispatches by parameter name, so existing hookimpls that + don't declare ``failure_kind`` continue to work unchanged.""" + called = {"count": 0} + + class LegacyListener: + @hookimpl + def on_task_instance_failed(self, previous_state, task_instance, error): + called["count"] += 1 + + listener_manager(LegacyListener()) + + get_listener_manager().hook.on_task_instance_failed( + previous_state=None, + task_instance=None, + error=RuntimeError("boom"), + failure_kind="infra", + ) + + assert called["count"] == 1 diff --git a/airflow-core/tests/unit/models/test_taskinstance_failure_kind_sources.py b/airflow-core/tests/unit/models/test_taskinstance_failure_kind_sources.py new file mode 100644 index 0000000000000..564fccec53c9d --- /dev/null +++ b/airflow-core/tests/unit/models/test_taskinstance_failure_kind_sources.py @@ -0,0 +1,70 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +""" +AIP-97 SIGTERM-source classification safety. + +A SIGTERM to a running task is ambiguous — a user mark-failed and an infra +eviction both deliver one. The refund must never fire for a user-initiated +stop. The single gate is ``failure_kind``: only ``INFRA`` refunds. These tests +pin that the kind each termination source carries maps to the right refund +decision, independent of the signal. +""" + +from __future__ import annotations + +import pytest + +from airflow._shared.state import TaskFailureKind +from airflow.models.taskinstance import _maybe_refund_infra_attempt + +from tests_common.test_utils.config import conf_vars + +ENABLED = {("core", "infra_failure_refund_retries"): "True", ("core", "max_infra_refunds"): "3"} + + +class _TI: + def __init__(self, max_tries=1): + self.max_tries = max_tries + + def __str__(self): + return "" + + +class _Task: + retries = 1 + + +@pytest.mark.parametrize( + ("source", "failure_kind", "should_refund"), + [ + # non-infra causes — never refunded, whatever signal did the killing + ("mark_task_failed", TaskFailureKind.MANUAL, False), + ("mark_dagrun_failed", TaskFailureKind.MANUAL, False), + ("app_exception", TaskFailureKind.APPLICATION, False), + ("own_limit_oom", TaskFailureKind.APPLICATION, False), + ("execution_timeout", TaskFailureKind.TIMEOUT, False), + ("unclassified", None, False), + # infrastructure disruption — the only refunded cause + ("pod_evicted", TaskFailureKind.INFRA, True), + ], +) +@conf_vars(ENABLED) +def test_only_infra_source_refunds(source, failure_kind, should_refund): + ti = _TI(max_tries=1) + refunded = _maybe_refund_infra_attempt(task_instance=ti, task=_Task(), failure_kind=failure_kind) + assert refunded is should_refund, source + assert ti.max_tries == (2 if should_refund else 1) diff --git a/airflow-core/tests/unit/models/test_taskinstance_infra_refund.py b/airflow-core/tests/unit/models/test_taskinstance_infra_refund.py new file mode 100644 index 0000000000000..87b7d14b99ef9 --- /dev/null +++ b/airflow-core/tests/unit/models/test_taskinstance_infra_refund.py @@ -0,0 +1,132 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from airflow.models.taskinstance import TaskInstance, _maybe_refund_infra_attempt +from airflow.utils.state import TaskInstanceState + +from tests_common.test_utils.config import conf_vars + + +class _FakeTI: + def __init__(self, max_tries: int): + self.max_tries = max_tries + + def __str__(self) -> str: + return "" + + +class _FakeTask: + def __init__(self, retries: int): + self.retries = retries + + +ENABLED = {("core", "infra_failure_refund_retries"): "True", ("core", "max_infra_refunds"): "3"} + + +class TestMaybeRefundInfraAttempt: + """The single safety gate: only ``failure_kind == "infra"``, with the flag on and under the cap, refunds.""" + + @conf_vars(ENABLED) + def test_infra_failure_refunds_one_attempt(self): + ti, task = _FakeTI(max_tries=1), _FakeTask(retries=1) + assert _maybe_refund_infra_attempt(task_instance=ti, task=task, failure_kind="infra") is True + assert ti.max_tries == 2 + + @conf_vars(ENABLED) + def test_app_failure_does_not_refund(self): + # failure_kind=None is the ordinary worker-exception path — a real bug must spend the budget. + ti, task = _FakeTI(max_tries=1), _FakeTask(retries=1) + assert _maybe_refund_infra_attempt(task_instance=ti, task=task, failure_kind=None) is False + assert ti.max_tries == 1 + + @conf_vars(ENABLED) + @pytest.mark.parametrize("failure_kind", ["application", "manual", "timeout"]) + def test_non_infra_kind_does_not_refund(self, failure_kind): + ti, task = _FakeTI(max_tries=1), _FakeTask(retries=1) + assert _maybe_refund_infra_attempt(task_instance=ti, task=task, failure_kind=failure_kind) is False + assert ti.max_tries == 1 + + @conf_vars(ENABLED) + def test_none_retries_refunds_without_crashing(self): + ti, task = _FakeTI(max_tries=0), _FakeTask(retries=None) + assert _maybe_refund_infra_attempt(task_instance=ti, task=task, failure_kind="infra") is True + assert ti.max_tries == 1 + + @conf_vars(ENABLED) + def test_zero_retries_still_refunds(self): + ti, task = _FakeTI(max_tries=0), _FakeTask(retries=0) + assert _maybe_refund_infra_attempt(task_instance=ti, task=task, failure_kind="infra") is True + assert ti.max_tries == 1 + + @conf_vars(ENABLED) + def test_zero_retries_refunds_are_capped(self): + ti, task = _FakeTI(max_tries=0), _FakeTask(retries=0) + assert [ + _maybe_refund_infra_attempt(task_instance=ti, task=task, failure_kind="infra") for _ in range(5) + ] == [True, True, True, False, False] + assert ti.max_tries == 3 + + @conf_vars({("core", "infra_failure_refund_retries"): "False"}) + def test_disabled_by_default_does_not_refund(self): + ti, task = _FakeTI(max_tries=1), _FakeTask(retries=1) + assert _maybe_refund_infra_attempt(task_instance=ti, task=task, failure_kind="infra") is False + assert ti.max_tries == 1 + + @conf_vars(ENABLED) + def test_cap_bounds_the_refunds(self): + # retries=1, cap=3 → refunds allowed only while (max_tries - retries) < 3, i.e. max_tries in {1,2,3}. + ti, task = _FakeTI(max_tries=1), _FakeTask(retries=1) + assert [ + _maybe_refund_infra_attempt(task_instance=ti, task=task, failure_kind="infra") for _ in range(5) + ] == [ + True, + True, + True, + False, + False, + ] + assert ti.max_tries == 4 # 1 + three refunds, then capped + + +class TestIsEligibleToRetryUsesMaxTries: + """is_eligible_to_retry gates on max_tries, not task.retries, so an infra refund that bumps + max_tries grants a retry even at retries=0 (and the two retry-decision paths stay in sync).""" + + @staticmethod + def _eligible(*, max_tries: int, try_number: int, state=None) -> bool: + stub = SimpleNamespace(state=state, max_tries=max_tries, try_number=try_number) + return TaskInstance.is_eligible_to_retry(stub) # type: ignore[arg-type] + + def test_zero_budget_not_bumped_is_not_eligible(self): + assert self._eligible(max_tries=0, try_number=1) is False + + def test_zero_budget_bumped_by_refund_is_eligible(self): + assert self._eligible(max_tries=1, try_number=1) is True + + def test_normal_budget_exhausted_is_not_eligible(self): + assert self._eligible(max_tries=2, try_number=3) is False + + def test_normal_budget_remaining_is_eligible(self): + assert self._eligible(max_tries=2, try_number=2) is True + + def test_restarting_always_eligible(self): + assert self._eligible(max_tries=0, try_number=9, state=TaskInstanceState.RESTARTING) is True diff --git a/devel-common/src/tests_common/test_utils/version_compat.py b/devel-common/src/tests_common/test_utils/version_compat.py index 7eb25dec2b3cb..e1b51ebe0347c 100644 --- a/devel-common/src/tests_common/test_utils/version_compat.py +++ b/devel-common/src/tests_common/test_utils/version_compat.py @@ -42,6 +42,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_2_PLUS = get_base_airflow_version_tuple() >= (3, 2, 0) AIRFLOW_V_3_2_2_PLUS = get_base_airflow_version_tuple() >= (3, 2, 2) AIRFLOW_V_3_3_PLUS = get_base_airflow_version_tuple() >= (3, 3, 0) +AIRFLOW_V_3_4_PLUS = get_base_airflow_version_tuple() >= (3, 4, 0) if AIRFLOW_V_3_1_PLUS: from airflow.sdk import PokeReturnValue, timezone @@ -64,6 +65,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: "AIRFLOW_V_3_1_PLUS", "AIRFLOW_V_3_2_PLUS", "AIRFLOW_V_3_3_PLUS", + "AIRFLOW_V_3_4_PLUS", "NOTSET", "XCOM_RETURN_KEY", "ArgNotSet", diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py index 8aaa430b92108..bf8644d2e16e5 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py @@ -596,6 +596,17 @@ def _change_state( container_message, exit_code, ) + + # Classify the pod failure so the scheduler can tell an infra + # disruption from the task's own failure; it reads this back via + # get_task_failure_info() when it processes the failure event. + from airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils import ( + classify_pod_failure, + ) + + task_failure_kind = classify_pod_failure(failure_details) + if task_failure_kind is not None: + self.task_failure_info[key] = task_failure_kind else: task_key_str = f"{key.dag_id}.{key.task_id}.{key.try_number}" self.log.warning( diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py index 2c4bd306e52d7..b9a1d3473f2c5 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py @@ -45,6 +45,7 @@ create_unique_id, ) from airflow.providers.cncf.kubernetes.pod_generator import PodGenerator, workload_to_command_args +from airflow.providers.cncf.kubernetes.version_compat import AIRFLOW_V_3_4_PLUS from airflow.providers.common.compat.sdk import AirflowException, Stats from airflow.utils.helpers import prune_dict from airflow.utils.log.logging_mixin import LoggingMixin @@ -329,10 +330,13 @@ def process_status( # since kube server have received request to delete pod set TI state failed if event["type"] == "DELETED" and pod.metadata.deletion_timestamp: self.log.info( - "Event: Pod %s deleted before it could complete, annotations: %s", + "Event: Pod %s deleted while running, annotations: %s", pod_name, annotations_string, ) + # A running task's pod removed by the platform (drain, preemption, spot + # reclaim, force-delete) is an infra disruption. An Airflow-initiated stop + # sets the TI terminal first, so the scheduler skips that case by state. self.watcher_queue.put( KubernetesWatch( pod_name, @@ -340,7 +344,7 @@ def process_status( TaskInstanceState.FAILED, annotations, resource_version, - None, + {"pod_status": "Running", "pod_reason": POD_DELETED_REASON}, ) ) else: @@ -409,6 +413,54 @@ def collect_pod_failure_details(pod: k8s.V1Pod, logger) -> FailureDetails | None } +# Reason token for a pod removed by the platform while its task was still running +# (node drain, preemption, spot reclaim, or a force-delete) — an infra disruption, +# not the task's own code. Distinct from a container that ran and exited on its own. +POD_DELETED_REASON = "PodDeleted" + +# Pod/node-level reasons that mean the platform ended the pod (an infrastructure +# disruption), as opposed to the container terminating on its own. +_INFRA_FAILURE_REASONS = frozenset( + { + "Evicted", + "Preempting", + "NodeShutdown", + "Shutdown", + "NodeLost", + "TerminationByKubelet", + "DeletionByTaintManager", + "DisruptionTarget", + POD_DELETED_REASON, + } +) + + +def classify_pod_failure(failure_details: FailureDetails | None): + """ + Classify a pod failure into a ``(TaskFailureKind, reason)`` pair. + + A node-level disruption (eviction, preemption, node shutdown) is ``INFRA``; + a container that ended on its own, an app crash or an ``OOMKilled`` against + its own limit, is ``APPLICATION``, so it earns no infra refund. Returns None + when there is nothing to classify, or when running against an Airflow that + predates the ``TaskFailureKind`` failure-context API (3.4). + """ + if not failure_details or not AIRFLOW_V_3_4_PLUS: + return None + + from airflow._shared.state import TaskFailureKind + + pod_reason = failure_details.get("pod_reason") + container_reason = failure_details.get("container_reason") + + if pod_reason in _INFRA_FAILURE_REASONS or container_reason in _INFRA_FAILURE_REASONS: + failure_kind = TaskFailureKind.INFRA + else: + failure_kind = TaskFailureKind.APPLICATION + + return failure_kind, (pod_reason or container_reason) + + def _analyze_containers( container_statuses: list[k8s.V1ContainerStatus] | None, container_type: Literal["init", "main"] ) -> FailureDetails | None: diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/version_compat.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/version_compat.py index dab50860d19b1..6a2f0ddebf2be 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/version_compat.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/version_compat.py @@ -36,6 +36,7 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: AIRFLOW_V_3_1_PLUS = get_base_airflow_version_tuple() >= (3, 1, 0) AIRFLOW_V_3_2_PLUS = get_base_airflow_version_tuple() >= (3, 2, 0) AIRFLOW_V_3_3_PLUS = get_base_airflow_version_tuple() >= (3, 3, 0) +AIRFLOW_V_3_4_PLUS = get_base_airflow_version_tuple() >= (3, 4, 0) __all__ = [ @@ -43,4 +44,5 @@ def get_base_airflow_version_tuple() -> tuple[int, int, int]: "AIRFLOW_V_3_1_PLUS", "AIRFLOW_V_3_2_PLUS", "AIRFLOW_V_3_3_PLUS", + "AIRFLOW_V_3_4_PLUS", ] diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_classify_pod_failure.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_classify_pod_failure.py new file mode 100644 index 0000000000000..dc14859a055c2 --- /dev/null +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_classify_pod_failure.py @@ -0,0 +1,147 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. +from __future__ import annotations + +import logging + +import pytest +from kubernetes import client as k8s + +from airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils import ( + classify_pod_failure, + collect_pod_failure_details, +) + +from tests_common.test_utils.version_compat import AIRFLOW_V_3_4_PLUS + +pytestmark = pytest.mark.skipif( + not AIRFLOW_V_3_4_PLUS, + reason="TaskFailureKind failure-context classification is Airflow 3.4+ only", +) + + +class TestClassifyPodFailure: + """The K8s bridge maps a pod/container failure to a (failure_kind, infra_reason) the retry decision can trust.""" + + @pytest.mark.parametrize( + ("pod_reason", "container_reason", "expected_kind"), + [ + # Node/platform ended the pod -> infra (earns a refund). + ("Evicted", None, "infra"), + ("Preempting", None, "infra"), + ("NodeShutdown", None, "infra"), + ("NodeLost", None, "infra"), + ("DisruptionTarget", None, "infra"), + ("TerminationByKubelet", None, "infra"), + # A running task's pod removed by the platform (drain, preempt, spot reclaim, + # force-delete) -> infra. The task didn't fail on its own; its pod was taken. + ("PodDeleted", None, "infra"), + # Container ended on its own -> user (no refund). The key one: an OOM against + # the container's OWN limit is the app's memory problem, not an infra disruption. + (None, "OOMKilled", "application"), + (None, "Error", "application"), + (None, "ContainerCannotRun", "application"), + # A node eviction that also shows a container OOM is still infra (the node acted). + ("Evicted", "OOMKilled", "infra"), + ], + ) + def test_kind_classification(self, pod_reason, container_reason, expected_kind): + details = {"pod_status": "Failed", "pod_reason": pod_reason, "container_reason": container_reason} + result = classify_pod_failure(details) + assert result is not None + failure_kind, infra_reason = result + assert failure_kind == expected_kind + assert infra_reason == (pod_reason or container_reason) + + def test_none_when_nothing_to_classify(self): + assert classify_pod_failure(None) is None + assert classify_pod_failure({}) is None + + def test_end_to_end_oomkilled_pod_is_application(self): + # Mirrors the live kind result: a real OOMKilled container (exit 137) flows through + # collect_pod_failure_details -> classify_pod_failure and classifies as application, not infra, + # so an app OOM does not earn an infra refund. + pod = k8s.V1Pod( + metadata=k8s.V1ObjectMeta(name="aip97-oom"), + status=k8s.V1PodStatus( + phase="Failed", + container_statuses=[ + k8s.V1ContainerStatus( + name="base", + image="python:3.11-slim", + image_id="", + ready=False, + restart_count=0, + state=k8s.V1ContainerState( + terminated=k8s.V1ContainerStateTerminated(reason="OOMKilled", exit_code=137) + ), + ) + ], + ), + ) + details = collect_pod_failure_details(pod, logging.getLogger("test")) + assert details is not None + assert details["container_reason"] == "OOMKilled" + assert details["exit_code"] == 137 + + failure_kind, infra_reason = classify_pod_failure(details) + assert failure_kind == "application" + assert infra_reason == "OOMKilled" + + def test_end_to_end_evicted_pod_is_infra(self): + # A real node-pressure eviction: the kubelet sets the pod phase=Failed with + # status.reason=Evicted (the pod object persists). This flows through + # collect_pod_failure_details -> classify_pod_failure and classifies as infra, + # so the disruption earns a refund. This is the positive-infra path the + # conservative default requires (an unclassified death does not refund). + pod = k8s.V1Pod( + metadata=k8s.V1ObjectMeta(name="aip97-evicted"), + status=k8s.V1PodStatus( + phase="Failed", reason="Evicted", message="The node was low on resource: memory." + ), + ) + details = collect_pod_failure_details(pod, logging.getLogger("test")) + assert details is not None + assert details["pod_reason"] == "Evicted" + + failure_kind, infra_reason = classify_pod_failure(details) + assert failure_kind == "infra" + assert infra_reason == "Evicted" + + def test_pod_deleted_while_running_is_infra(self): + # A pod force-deleted / drained / preempted while its task was running has no + # container-exit status; the watcher emits {pod_reason: PodDeleted}. It must + # classify infra so the disruption earns a refund (the scheduler still gates on + # the TI being non-terminal, so an Airflow-initiated stop is excluded upstream). + failure_kind, infra_reason = classify_pod_failure( + {"pod_status": "Running", "pod_reason": "PodDeleted"} + ) + assert failure_kind == "infra" + assert infra_reason == "PodDeleted" + + """The executor stashes a (failure_kind, infra_reason); the scheduler reads it once.""" + + def test_base_executor_failure_info_round_trips_once(self): + from airflow.executors.local_executor import LocalExecutor + + ex = LocalExecutor() + key = ("dag", "task", "run", 1, -1) + classified = classify_pod_failure({"pod_status": "Failed", "pod_reason": "Evicted"}) + ex.task_failure_info[key] = classified + # first read returns it, and clears it so a later event can't reuse stale context + assert ex.get_task_failure_info(key) is classified + assert ex.get_task_failure_info(key) is None diff --git a/shared/listeners/src/airflow_shared/listeners/spec/taskinstance.py b/shared/listeners/src/airflow_shared/listeners/spec/taskinstance.py index d3450d6b05aa7..34ccc7512cdc7 100644 --- a/shared/listeners/src/airflow_shared/listeners/spec/taskinstance.py +++ b/shared/listeners/src/airflow_shared/listeners/spec/taskinstance.py @@ -27,6 +27,7 @@ from airflow.models.taskinstance import TaskInstance from airflow.sdk.execution_time.task_runner import RuntimeTaskInstance from airflow.utils.state import TaskInstanceState + from airflow_shared.state import TaskFailureKind hookspec = HookspecMarker("airflow") @@ -52,8 +53,32 @@ def on_task_instance_failed( previous_state: TaskInstanceState | None, task_instance: RuntimeTaskInstance | TaskInstance, error: None | str | BaseException, + failure_kind: TaskFailureKind | None, + reason: str | None, ): - """Execute when task state changes to FAIL. previous_state can be None.""" + """ + Execute when task state changes to FAIL. previous_state can be None. + + :param previous_state: Previous state of the task instance (can be None) + :param task_instance: The task instance object + :param error: The exception that caused the failure (or human-readable + message string for API-driven manual transitions) + :param failure_kind: What caused the failure: a :class:`TaskFailureKind` + (``INFRA`` / ``APPLICATION`` / ``TIMEOUT`` / ``MANUAL``), or ``None`` when + the cause was not classified. + :param reason: For an ``INFRA`` failure, the executor's short token for what + happened (e.g. ``Evicted``, ``PodDeleted``), since the worker left no + ``error``. ``None`` otherwise. Transient — persist it yourself if you need + it to outlive the failure. + + Pluggy dispatches by parameter name, so a ``hookimpl`` that doesn't declare + ``failure_kind`` / ``reason`` keeps working. One that does must declare it + *without* a default — pluggy treats an impl-side default as authoritative and + overrides the caller's value:: + + @hookimpl + def on_task_instance_failed(self, previous_state, task_instance, error, failure_kind, reason): ... + """ @hookspec diff --git a/shared/state/src/airflow_shared/state/__init__.py b/shared/state/src/airflow_shared/state/__init__.py index e2c094d84818d..74966080ee488 100644 --- a/shared/state/src/airflow_shared/state/__init__.py +++ b/shared/state/src/airflow_shared/state/__init__.py @@ -72,6 +72,29 @@ def __post_init__(self) -> None: StoreScope = TaskScope | AssetScope +class TaskFailureKind(str, Enum): + """ + What caused a task instance to fail. + + Set on ``on_task_instance_failed`` so a listener can act on the cause without + parsing the error. + + ``INFRA`` — an infrastructure disruption (eviction, preemption, node loss). + ``APPLICATION`` — the task's own code (a raised exception, an OOM against its + own limit, a non-zero exit). + ``TIMEOUT`` — the task exceeded its ``execution_timeout``. + ``MANUAL`` — an operator set it failed (mark task or dag run failed). + """ + + INFRA = "infra" + APPLICATION = "application" + TIMEOUT = "timeout" + MANUAL = "manual" + + def __str__(self) -> str: + return self.value + + class AssetStateStoreWriterKind(str, Enum): """ Identifies what kind of writer last updated an asset state store entry. diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index 31d0831bbadec..b76eefc97ea12 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -47,6 +47,7 @@ from airflow.sdk._shared.observability.metrics import stats from airflow.sdk._shared.observability.metrics.stats import build_dag_metric_tags from airflow.sdk._shared.observability.traces import get_task_span_detail_level +from airflow.sdk._shared.state import TaskFailureKind from airflow.sdk._shared.template_rendering import truncate_rendered_value from airflow.sdk.api.client import get_hostname, getuser from airflow.sdk.api.datamodels._generated import ( @@ -2298,7 +2299,12 @@ def finalize( _run_task_state_change_callbacks(task, "on_retry_callback", context, log) 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=TaskFailureKind.TIMEOUT + if isinstance(error, AirflowTaskTimeout) + else TaskFailureKind.APPLICATION, ) except Exception: log.exception("error calling listener") @@ -2308,7 +2314,12 @@ def finalize( _run_task_state_change_callbacks(task, "on_failure_callback", context, log) 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=TaskFailureKind.TIMEOUT + if isinstance(error, AirflowTaskTimeout) + else TaskFailureKind.APPLICATION, ) except Exception: log.exception("error calling listener") From 79a6ebd9f342e8726d514615cf02943f896916f6 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Wed, 29 Jul 2026 14:55:21 -0700 Subject: [PATCH 2/6] Pass reason to the failure listener from the worker path, trim comments finalize() omitted reason= on both on_task_instance_failed calls. pluggy raises HookCallError when a hookspec argument is missing, and the call is wrapped in a bare except, so a listener using the signature the hookspec documents silently never fired on the worker path. Also trims narrating comments and docstrings down to the reasons the code can't carry, and corrects handle_failure's docstring, which said reason reaches the listener as error. Signed-off-by: 1fanwang <1fannnw@gmail.com> --- .../services/public/task_instances.py | 1 - .../src/airflow/executors/base_executor.py | 4 +- .../src/airflow/jobs/scheduler_job_runner.py | 13 +++--- .../src/airflow/models/taskinstance.py | 21 ++++------ .../tests/unit/jobs/test_scheduler_job.py | 6 +-- .../unit/listeners/test_listener_types.py | 13 +++--- .../test_taskinstance_failure_kind_sources.py | 6 +-- .../models/test_taskinstance_infra_refund.py | 4 +- .../executors/kubernetes_executor.py | 5 +-- .../executors/kubernetes_executor_utils.py | 22 +++++----- .../executors/test_classify_pod_failure.py | 31 ++++++-------- .../listeners/spec/taskinstance.py | 7 ++-- .../src/airflow_shared/state/__init__.py | 9 ++--- .../airflow/sdk/execution_time/task_runner.py | 2 + .../execution_time/test_task_runner.py | 40 ++++++++++++++++++- 15 files changed, 102 insertions(+), 82 deletions(-) diff --git a/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py b/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py index fc2764cceaad0..011cb228d739b 100644 --- a/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py +++ b/airflow-core/src/airflow/api_fastapi/core_api/services/public/task_instances.py @@ -113,7 +113,6 @@ def _emit_state_listener_hooks(updated_tis: list[TI], new_state: str | TaskInsta if new_state == TaskInstanceState.SUCCESS: get_listener_manager().hook.on_task_instance_success(previous_state=None, task_instance=ti) elif new_state == TaskInstanceState.FAILED: - # An operator set this failed — a manual action, never refunded. get_listener_manager().hook.on_task_instance_failed( previous_state=None, task_instance=ti, diff --git a/airflow-core/src/airflow/executors/base_executor.py b/airflow-core/src/airflow/executors/base_executor.py index c430cfc90d52a..3629fe15caeeb 100644 --- a/airflow-core/src/airflow/executors/base_executor.py +++ b/airflow-core/src/airflow/executors/base_executor.py @@ -229,7 +229,6 @@ 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] = {} - # An executor's (failure_kind, reason) for a key, read once by the scheduler. self.task_failure_info: dict[WorkloadKey, tuple[TaskFailureKind, str | None]] = {} self._task_event_logs: deque[Log] = deque() self.conf = ExecutorConf(team_name) @@ -557,8 +556,7 @@ def get_task_failure_info(self, key: WorkloadKey): """ Return and clear the ``(failure_kind, reason)`` this executor classified for ``key``. - Populated by executors that can read the real failure cause (e.g. the Kubernetes - executor telling an ``Evicted`` pod from an ``OOMKilled`` one). Returns None otherwise. + None unless the executor can read the real cause, as the Kubernetes executor does. """ return self.task_failure_info.pop(key, None) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index a702493b530af..e825319017c6a 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -1468,9 +1468,8 @@ def process_executor_events( job_id, ) state, info = event_buffer.pop(buffer_key) - # Pop the executor's (failure_kind, reason) for this event now, so the entry - # is always cleared — not only on the killed-externally branch below — which - # keeps self-reporting tasks from leaking entries. + # Pop unconditionally, not only on the killed-externally branch below, so a + # self-reporting task can't leak its entry. executor_failure_kind = executor.get_task_failure_info(ti.key) if state in (TaskInstanceState.QUEUED, TaskInstanceState.RUNNING): @@ -1667,11 +1666,9 @@ def process_executor_events( executor.send_callback(email_request) # Update task state - emails are handled by DAG processor now - # The worker died from outside without reporting an error. Refund only - # when an executor positively classified this as infra (its bridge read - # the real cause); an executor that only reports a state can't tell an - # eviction from a silent crash, so it stays unclassified and spends the - # retry as today. + # 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: diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index 0dfe60466f8da..e219563870a11 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/airflow-core/src/airflow/models/taskinstance.py @@ -562,16 +562,14 @@ def _maybe_refund_infra_attempt(*, task_instance, task, failure_kind, reason=Non """ Refund one retry attempt (bump ``max_tries``) for an infra-classified failure. - Fires only when ``failure_kind == INFRA``, ``[core] infra_failure_refund_retries`` - is on, and the ``[core] max_infra_refunds`` cap is not yet reached. Every other - kind — and an unclassified (``None``) failure — spends the user's retry. Returns - True if an attempt was refunded. + 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. - # The refund still applies at retries=0: an infra disruption is not the user's - # failure, so it earns an attempt back regardless of the user's retry budget. + # 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 @@ -1895,8 +1893,8 @@ def fetch_handle_failure_context( ``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 never reported an error in that case, so it is surfaced as the - failure's ``reason`` argument to the listener rather than persisted. + 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) @@ -1928,9 +1926,6 @@ def fetch_handle_failure_context( # Actual callbacks are handled by the DAG processor, not the scheduler task = getattr(ti, "task", None) - # An infra-classified failure refunds the attempt instead of spending the - # user's retry budget (opt-in, capped). Extends the Kubernetes executor's - # pre-execution requeue-without-consuming-a-retry to the mid-execution case. _maybe_refund_infra_attempt(task_instance=ti, task=task, failure_kind=failure_kind, reason=reason) if not ti.is_eligible_to_retry(): @@ -1986,8 +1981,8 @@ def handle_failure( :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, surfaced to the - listener as the failure's ``error`` (not persisted). + :param reason: the executor's token for an infra failure, passed to the + listener rather than persisted. """ if TYPE_CHECKING: assert self.task diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 62ada334f23b4..94fc561b8cd77 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -559,15 +559,15 @@ def _run(dag_id, stashed): ti.refresh_from_db(session=session) return ti - # Unclassified (executor only reports a state) → no refund, exactly as today. + # Unclassified: no refund, exactly as today. ti = _run("aip97_unclassified", stashed=None) assert ti.max_tries == 1 - # Bridge classified it infra (e.g. a pod Evicted) → refunded. + # Classified infra (a pod Evicted): refunded. ti = _run("aip97_infra", stashed=(TaskFailureKind.INFRA, "Evicted")) assert ti.max_tries == 2 - # Bridge classified it user (e.g. an app OOM against its own limit) → no refund. + # Classified application (an OOM against its own limit): no refund. ti = _run("aip97_app_oom", stashed=(TaskFailureKind.APPLICATION, "OOMKilled")) assert ti.max_tries == 1 diff --git a/airflow-core/tests/unit/listeners/test_listener_types.py b/airflow-core/tests/unit/listeners/test_listener_types.py index 09cf01d229985..b405d63a8708f 100644 --- a/airflow-core/tests/unit/listeners/test_listener_types.py +++ b/airflow-core/tests/unit/listeners/test_listener_types.py @@ -26,11 +26,10 @@ class TestOnTaskInstanceFailedAcceptsFailureKind: (``infra`` / ``application`` / ``timeout`` / ``manual``) without parsing the error.""" def test_listener_with_failure_kind_receives_it(self, listener_manager): - received: dict[str, str | None] = {"failure_kind": None} + received: dict[str, str | None] = {"failure_kind": None, "reason": None} - # Per the hookspec docstring, listener implementations must declare - # failure_kind WITHOUT a default value — pluggy treats the impl - # default as authoritative and silently overrides the caller's value. + # Must be declared WITHOUT a default: pluggy treats an impl-side default as + # authoritative and silently overrides the caller's value. class InfraListener: @hookimpl def on_task_instance_failed( @@ -39,8 +38,10 @@ def on_task_instance_failed( task_instance, error, failure_kind, + reason, ): received["failure_kind"] = failure_kind + received["reason"] = reason listener_manager(InfraListener()) @@ -49,9 +50,10 @@ def on_task_instance_failed( task_instance=None, error=RuntimeError("boom"), failure_kind="infra", + reason="Evicted", ) - assert received["failure_kind"] == "infra" + assert received == {"failure_kind": "infra", "reason": "Evicted"} def test_listener_without_failure_kind_param_keeps_working(self, listener_manager): """Pluggy dispatches by parameter name, so existing hookimpls that @@ -70,6 +72,7 @@ def on_task_instance_failed(self, previous_state, task_instance, error): task_instance=None, error=RuntimeError("boom"), failure_kind="infra", + reason=None, ) assert called["count"] == 1 diff --git a/airflow-core/tests/unit/models/test_taskinstance_failure_kind_sources.py b/airflow-core/tests/unit/models/test_taskinstance_failure_kind_sources.py index 564fccec53c9d..14b0aaa325f55 100644 --- a/airflow-core/tests/unit/models/test_taskinstance_failure_kind_sources.py +++ b/airflow-core/tests/unit/models/test_taskinstance_failure_kind_sources.py @@ -17,7 +17,7 @@ """ AIP-97 SIGTERM-source classification safety. -A SIGTERM to a running task is ambiguous — a user mark-failed and an infra +A SIGTERM to a running task is ambiguous: a user mark-failed and an infra eviction both deliver one. The refund must never fire for a user-initiated stop. The single gate is ``failure_kind``: only ``INFRA`` refunds. These tests pin that the kind each termination source carries maps to the right refund @@ -51,14 +51,14 @@ class _Task: @pytest.mark.parametrize( ("source", "failure_kind", "should_refund"), [ - # non-infra causes — never refunded, whatever signal did the killing + # non-infra causes, never refunded whatever signal did the killing ("mark_task_failed", TaskFailureKind.MANUAL, False), ("mark_dagrun_failed", TaskFailureKind.MANUAL, False), ("app_exception", TaskFailureKind.APPLICATION, False), ("own_limit_oom", TaskFailureKind.APPLICATION, False), ("execution_timeout", TaskFailureKind.TIMEOUT, False), ("unclassified", None, False), - # infrastructure disruption — the only refunded cause + # infrastructure disruption, the only refunded cause ("pod_evicted", TaskFailureKind.INFRA, True), ], ) diff --git a/airflow-core/tests/unit/models/test_taskinstance_infra_refund.py b/airflow-core/tests/unit/models/test_taskinstance_infra_refund.py index 87b7d14b99ef9..b60afff9976c3 100644 --- a/airflow-core/tests/unit/models/test_taskinstance_infra_refund.py +++ b/airflow-core/tests/unit/models/test_taskinstance_infra_refund.py @@ -53,7 +53,7 @@ def test_infra_failure_refunds_one_attempt(self): @conf_vars(ENABLED) def test_app_failure_does_not_refund(self): - # failure_kind=None is the ordinary worker-exception path — a real bug must spend the budget. + # failure_kind=None is the ordinary worker-exception path: a real bug spends the budget. ti, task = _FakeTI(max_tries=1), _FakeTask(retries=1) assert _maybe_refund_infra_attempt(task_instance=ti, task=task, failure_kind=None) is False assert ti.max_tries == 1 @@ -93,7 +93,7 @@ def test_disabled_by_default_does_not_refund(self): @conf_vars(ENABLED) def test_cap_bounds_the_refunds(self): - # retries=1, cap=3 → refunds allowed only while (max_tries - retries) < 3, i.e. max_tries in {1,2,3}. + # retries=1, cap=3: refunds allowed only while (max_tries - retries) < 3. ti, task = _FakeTI(max_tries=1), _FakeTask(retries=1) assert [ _maybe_refund_infra_attempt(task_instance=ti, task=task, failure_kind="infra") for _ in range(5) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py index bf8644d2e16e5..95180ba8f6958 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor.py @@ -597,9 +597,8 @@ def _change_state( exit_code, ) - # Classify the pod failure so the scheduler can tell an infra - # disruption from the task's own failure; it reads this back via - # get_task_failure_info() when it processes the failure event. + # The scheduler reads this back via get_task_failure_info() when it + # processes the failure event. from airflow.providers.cncf.kubernetes.executors.kubernetes_executor_utils import ( classify_pod_failure, ) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py index b9a1d3473f2c5..e38e1d2dff8f2 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py @@ -334,9 +334,8 @@ def process_status( pod_name, annotations_string, ) - # A running task's pod removed by the platform (drain, preemption, spot - # reclaim, force-delete) is an infra disruption. An Airflow-initiated stop - # sets the TI terminal first, so the scheduler skips that case by state. + # An Airflow-initiated stop sets the TI terminal first, so the scheduler + # excludes that case by state and only a real disruption reaches here. self.watcher_queue.put( KubernetesWatch( pod_name, @@ -413,13 +412,12 @@ def collect_pod_failure_details(pod: k8s.V1Pod, logger) -> FailureDetails | None } -# Reason token for a pod removed by the platform while its task was still running -# (node drain, preemption, spot reclaim, or a force-delete) — an infra disruption, -# not the task's own code. Distinct from a container that ran and exited on its own. +# A pod taken by the platform while its task ran: node drain, preemption, spot reclaim, +# or a force-delete. Distinct from a container that ran and exited on its own. POD_DELETED_REASON = "PodDeleted" -# Pod/node-level reasons that mean the platform ended the pod (an infrastructure -# disruption), as opposed to the container terminating on its own. +# Pod/node-level reasons meaning the platform ended the pod, as opposed to the container +# terminating on its own. _INFRA_FAILURE_REASONS = frozenset( { "Evicted", @@ -439,11 +437,9 @@ def classify_pod_failure(failure_details: FailureDetails | None): """ Classify a pod failure into a ``(TaskFailureKind, reason)`` pair. - A node-level disruption (eviction, preemption, node shutdown) is ``INFRA``; - a container that ended on its own, an app crash or an ``OOMKilled`` against - its own limit, is ``APPLICATION``, so it earns no infra refund. Returns None - when there is nothing to classify, or when running against an Airflow that - predates the ``TaskFailureKind`` failure-context API (3.4). + A node-level disruption is ``INFRA``. A container that ended on its own, including + an ``OOMKilled`` against its own limit, is ``APPLICATION`` and earns no refund. + Returns None with nothing to classify, or on an Airflow predating the API (3.4). """ if not failure_details or not AIRFLOW_V_3_4_PLUS: return None diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_classify_pod_failure.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_classify_pod_failure.py index dc14859a055c2..d9a6cd9045374 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_classify_pod_failure.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_classify_pod_failure.py @@ -40,22 +40,21 @@ class TestClassifyPodFailure: @pytest.mark.parametrize( ("pod_reason", "container_reason", "expected_kind"), [ - # Node/platform ended the pod -> infra (earns a refund). + # Node/platform ended the pod -> infra. ("Evicted", None, "infra"), ("Preempting", None, "infra"), ("NodeShutdown", None, "infra"), ("NodeLost", None, "infra"), ("DisruptionTarget", None, "infra"), ("TerminationByKubelet", None, "infra"), - # A running task's pod removed by the platform (drain, preempt, spot reclaim, - # force-delete) -> infra. The task didn't fail on its own; its pod was taken. + # Pod taken while the task ran (drain, preempt, spot reclaim, force-delete). ("PodDeleted", None, "infra"), - # Container ended on its own -> user (no refund). The key one: an OOM against - # the container's OWN limit is the app's memory problem, not an infra disruption. + # Container ended on its own -> application. An OOM against the container's OWN + # limit is the app's memory problem, not a disruption. (None, "OOMKilled", "application"), (None, "Error", "application"), (None, "ContainerCannotRun", "application"), - # A node eviction that also shows a container OOM is still infra (the node acted). + # A node eviction that also shows a container OOM is still infra: the node acted. ("Evicted", "OOMKilled", "infra"), ], ) @@ -72,9 +71,8 @@ def test_none_when_nothing_to_classify(self): assert classify_pod_failure({}) is None def test_end_to_end_oomkilled_pod_is_application(self): - # Mirrors the live kind result: a real OOMKilled container (exit 137) flows through - # collect_pod_failure_details -> classify_pod_failure and classifies as application, not infra, - # so an app OOM does not earn an infra refund. + # Mirrors the live kind result: a real OOMKilled container (exit 137) classifies as + # application, so an app OOM earns no refund. pod = k8s.V1Pod( metadata=k8s.V1ObjectMeta(name="aip97-oom"), status=k8s.V1PodStatus( @@ -103,11 +101,9 @@ def test_end_to_end_oomkilled_pod_is_application(self): assert infra_reason == "OOMKilled" def test_end_to_end_evicted_pod_is_infra(self): - # A real node-pressure eviction: the kubelet sets the pod phase=Failed with - # status.reason=Evicted (the pod object persists). This flows through - # collect_pod_failure_details -> classify_pod_failure and classifies as infra, - # so the disruption earns a refund. This is the positive-infra path the - # conservative default requires (an unclassified death does not refund). + # A real node-pressure eviction: the kubelet sets phase=Failed with reason=Evicted and + # the pod object persists. This is the positive-infra path the conservative default + # requires, since an unclassified death does not refund. pod = k8s.V1Pod( metadata=k8s.V1ObjectMeta(name="aip97-evicted"), status=k8s.V1PodStatus( @@ -123,10 +119,9 @@ def test_end_to_end_evicted_pod_is_infra(self): assert infra_reason == "Evicted" def test_pod_deleted_while_running_is_infra(self): - # A pod force-deleted / drained / preempted while its task was running has no - # container-exit status; the watcher emits {pod_reason: PodDeleted}. It must - # classify infra so the disruption earns a refund (the scheduler still gates on - # the TI being non-terminal, so an Airflow-initiated stop is excluded upstream). + # Force-delete / drain / preempt leaves no container-exit status, so the watcher emits + # {pod_reason: PodDeleted}. The scheduler still gates on the TI being non-terminal, so an + # Airflow-initiated stop is excluded upstream. failure_kind, infra_reason = classify_pod_failure( {"pod_status": "Running", "pod_reason": "PodDeleted"} ) diff --git a/shared/listeners/src/airflow_shared/listeners/spec/taskinstance.py b/shared/listeners/src/airflow_shared/listeners/spec/taskinstance.py index 34ccc7512cdc7..aec552d646ec6 100644 --- a/shared/listeners/src/airflow_shared/listeners/spec/taskinstance.py +++ b/shared/listeners/src/airflow_shared/listeners/spec/taskinstance.py @@ -68,13 +68,12 @@ def on_task_instance_failed( the cause was not classified. :param reason: For an ``INFRA`` failure, the executor's short token for what happened (e.g. ``Evicted``, ``PodDeleted``), since the worker left no - ``error``. ``None`` otherwise. Transient — persist it yourself if you need - it to outlive the failure. + ``error``. ``None`` otherwise, and never persisted. Pluggy dispatches by parameter name, so a ``hookimpl`` that doesn't declare ``failure_kind`` / ``reason`` keeps working. One that does must declare it - *without* a default — pluggy treats an impl-side default as authoritative and - overrides the caller's value:: + *without* a default, since pluggy treats an impl-side default as authoritative + and overrides the caller's value:: @hookimpl def on_task_instance_failed(self, previous_state, task_instance, error, failure_kind, reason): ... diff --git a/shared/state/src/airflow_shared/state/__init__.py b/shared/state/src/airflow_shared/state/__init__.py index 74966080ee488..583cab5a0c5be 100644 --- a/shared/state/src/airflow_shared/state/__init__.py +++ b/shared/state/src/airflow_shared/state/__init__.py @@ -79,11 +79,10 @@ class TaskFailureKind(str, Enum): Set on ``on_task_instance_failed`` so a listener can act on the cause without parsing the error. - ``INFRA`` — an infrastructure disruption (eviction, preemption, node loss). - ``APPLICATION`` — the task's own code (a raised exception, an OOM against its - own limit, a non-zero exit). - ``TIMEOUT`` — the task exceeded its ``execution_timeout``. - ``MANUAL`` — an operator set it failed (mark task or dag run failed). + - ``INFRA``: an infrastructure disruption (eviction, preemption, node loss). + - ``APPLICATION``: the task's own code, including an OOM against its own limit. + - ``TIMEOUT``: the task exceeded its ``execution_timeout``. + - ``MANUAL``: an operator set it failed (mark task or dag run failed). """ INFRA = "infra" diff --git a/task-sdk/src/airflow/sdk/execution_time/task_runner.py b/task-sdk/src/airflow/sdk/execution_time/task_runner.py index b76eefc97ea12..eb37879e892c6 100644 --- a/task-sdk/src/airflow/sdk/execution_time/task_runner.py +++ b/task-sdk/src/airflow/sdk/execution_time/task_runner.py @@ -2305,6 +2305,7 @@ def finalize( failure_kind=TaskFailureKind.TIMEOUT if isinstance(error, AirflowTaskTimeout) else TaskFailureKind.APPLICATION, + reason=None, ) except Exception: log.exception("error calling listener") @@ -2320,6 +2321,7 @@ def finalize( failure_kind=TaskFailureKind.TIMEOUT if isinstance(error, AirflowTaskTimeout) else TaskFailureKind.APPLICATION, + reason=None, ) except Exception: log.exception("error calling listener") diff --git a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py index 0031974b86646..736c157f43d99 100644 --- a/task-sdk/tests/task_sdk/execution_time/test_task_runner.py +++ b/task-sdk/tests/task_sdk/execution_time/test_task_runner.py @@ -63,7 +63,7 @@ timezone, ) from airflow.sdk._shared.observability.metrics.base_stats_logger import StatsLogger -from airflow.sdk._shared.state import AssetScope, TaskScope +from airflow.sdk._shared.state import AssetScope, TaskFailureKind, TaskScope from airflow.sdk.api.datamodels._generated import ( AssetProfile, AssetResponse, @@ -4506,6 +4506,44 @@ def execute(self, context): assert listener.state == [TaskInstanceState.RUNNING, TaskInstanceState.FAILED] assert listener.error == error + @pytest.mark.parametrize( + "state", [TaskInstanceState.FAILED, TaskInstanceState.UP_FOR_RETRY], ids=["failed", "up_for_retry"] + ) + def test_listener_declaring_reason_still_fires( + self, state, mocked_parse, mock_supervisor_comms, listener_manager + ): + """A hookimpl using the full signature the hookspec documents must still be called. + + pluggy raises HookCallError when a hookspec argument is missing from the call, and + finalize() swallows it, so omitting one here silently stops the listener firing. + """ + received = {} + + class FullSignatureListener: + @hookimpl + def on_task_instance_failed(self, previous_state, task_instance, error, failure_kind, reason): + received["failure_kind"] = failure_kind + received["reason"] = reason + + listener_manager(FullSignatureListener()) + + task = BaseOperator(task_id="test_listener_declaring_reason_still_fires") + dag = get_inline_dag(dag_id="test_dag", task=task) + ti = TaskInstance( + id=uuid7(), + task_id=task.task_id, + dag_id=dag.dag_id, + run_id="test_run", + try_number=1, + dag_version_id=uuid7(), + ) + runtime_ti = RuntimeTaskInstance.model_construct( + **ti.model_dump(exclude_unset=True), task=task, start_date=timezone.utcnow() + ) + finalize(runtime_ti, state, runtime_ti.get_template_context(), mock.MagicMock(), RuntimeError("boom")) + + assert received == {"failure_kind": TaskFailureKind.APPLICATION, "reason": None} + def test_task_runner_calls_listeners_failed_when_terminal_send_fails( self, mocked_parse, mock_supervisor_comms, listener_manager ): From 3ce5dbc184f557bd4000f4ca1446b8d21694786e Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Wed, 29 Jul 2026 17:35:58 -0700 Subject: [PATCH 3/6] AIP-97: ground the K8s infra-reason set in verified kubelet behavior Four of the reasons could never match. TerminationByKubelet, DeletionByTaintManager and DisruptionTarget are pod *condition* reasons, and Shutdown is a node event reason; the classifier reads only pod.status.reason and a container's terminated/waiting reason, so none of them was reachable. The tests asserted the same wrong assumption, so they passed. Replace them with Terminated, which graceful node shutdown writes to pod.status.reason when it kills an already-running pod. That is the node-drain case this feature exists for, and it was missing. The kubelet is its sole writer, so it cannot collide with an application exit, and it never appears as a container reason. Pin every remaining reason to its Go definition, and record that a taint-manager eviction and a scheduler preemption stay unclassified until the classifier reads pod.status.conditions. Also type the signatures this change added and keyword the watcher call it modified. Signed-off-by: 1fanwang <1fannnw@gmail.com> --- .../src/airflow/executors/base_executor.py | 2 +- .../src/airflow/jobs/scheduler_job_runner.py | 4 +- .../src/airflow/models/taskinstance.py | 12 +++-- .../executors/kubernetes_executor_utils.py | 48 +++++++++++++------ .../executors/test_classify_pod_failure.py | 4 +- 5 files changed, 48 insertions(+), 22 deletions(-) diff --git a/airflow-core/src/airflow/executors/base_executor.py b/airflow-core/src/airflow/executors/base_executor.py index 3629fe15caeeb..93452ec1b794e 100644 --- a/airflow-core/src/airflow/executors/base_executor.py +++ b/airflow-core/src/airflow/executors/base_executor.py @@ -552,7 +552,7 @@ def get_event_buffer(self, dag_ids=None) -> dict[WorkloadKey, EventBufferValueTy return cleared_events - def get_task_failure_info(self, key: WorkloadKey): + 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``. diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index e825319017c6a..1827717cd4a0f 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -1470,7 +1470,9 @@ def process_executor_events( 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 = executor.get_task_failure_info(ti.key) + 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 diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index e219563870a11..e5eccd323e26b 100644 --- a/airflow-core/src/airflow/models/taskinstance.py +++ b/airflow-core/src/airflow/models/taskinstance.py @@ -558,7 +558,13 @@ 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, task, failure_kind, reason=None) -> bool: +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. @@ -1878,7 +1884,7 @@ def fetch_handle_failure_context( *, session: Session, fail_fast: bool = False, - failure_kind: str | None = None, + failure_kind: TaskFailureKind | None = None, reason: str | None = None, ): """ @@ -1970,7 +1976,7 @@ def handle_failure( test_mode: bool | None = None, *, session: Session = NEW_SESSION, - failure_kind: str | None = None, + failure_kind: TaskFailureKind | None = None, reason: str | None = None, ) -> None: """ diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py index e38e1d2dff8f2..7ad585ec0e4e3 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py @@ -54,6 +54,8 @@ if TYPE_CHECKING: from kubernetes.client import Configuration, models as k8s + from airflow._shared.state import TaskFailureKind + class ResourceVersion: """Singleton for tracking resourceVersion from Kubernetes.""" @@ -338,12 +340,12 @@ def process_status( # excludes that case by state and only a real disruption reaches here. self.watcher_queue.put( KubernetesWatch( - pod_name, - namespace, - TaskInstanceState.FAILED, - annotations, - resource_version, - {"pod_status": "Running", "pod_reason": POD_DELETED_REASON}, + pod_name=pod_name, + namespace=namespace, + state=TaskInstanceState.FAILED, + annotations=annotations, + resource_version=resource_version, + failure_details={"pod_status": "Running", "pod_reason": POD_DELETED_REASON}, ) ) else: @@ -414,26 +416,42 @@ def collect_pod_failure_details(pod: k8s.V1Pod, logger) -> FailureDetails | None # A pod taken by the platform while its task ran: node drain, preemption, spot reclaim, # or a force-delete. Distinct from a container that ran and exited on its own. -POD_DELETED_REASON = "PodDeleted" +POD_DELETED_REASON: str = "PodDeleted" -# Pod/node-level reasons meaning the platform ended the pod, as opposed to the container -# terminating on its own. -_INFRA_FAILURE_REASONS = frozenset( +# Pod-level reasons meaning the platform ended the pod, as opposed to the container +# terminating on its own. Matched against pod.status.reason and a container's +# terminated/waiting reason, which are the only two fields collect_pod_failure_details +# populates. +# +# The Python kubernetes client publishes no constants for these (it models only the +# API schema, not the kubelet's reason vocabulary), so they are pinned to their Go +# definitions at kubernetes/kubernetes 6fd1049740274a3adfa5d5be904715a780465bf2: +# Evicted pkg/kubelet/eviction/eviction_manager.go#L648 +# Preempting pkg/kubelet/preemption/preemption.go +# NodeShutdown pkg/kubelet/nodeshutdown/nodeshutdown_manager_linux.go (admission rejection) +# Terminated pkg/kubelet/nodeshutdown/nodeshutdown_manager.go#L88 (graceful node +# shutdown killing an already-running pod; sole writer of this reason) +# NodeLost pkg/controller/util/node/controller_utils.go#L108 +# +# Known gap: a taint-manager eviction and a scheduler preemption leave pod.status.reason +# empty and are only observable via pod.status.conditions[type=DisruptionTarget].reason +# ("DeletionByTaintManager" / "PreemptionByScheduler"). Reading conditions is a separate +# change; those two disruptions currently classify as APPLICATION. +_INFRA_FAILURE_REASONS: frozenset[str] = frozenset( { "Evicted", "Preempting", "NodeShutdown", - "Shutdown", + "Terminated", "NodeLost", - "TerminationByKubelet", - "DeletionByTaintManager", - "DisruptionTarget", POD_DELETED_REASON, } ) -def classify_pod_failure(failure_details: FailureDetails | None): +def classify_pod_failure( + failure_details: FailureDetails | None, +) -> tuple[TaskFailureKind, str | None] | None: """ Classify a pod failure into a ``(TaskFailureKind, reason)`` pair. diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_classify_pod_failure.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_classify_pod_failure.py index d9a6cd9045374..feed64013009c 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_classify_pod_failure.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_classify_pod_failure.py @@ -45,8 +45,8 @@ class TestClassifyPodFailure: ("Preempting", None, "infra"), ("NodeShutdown", None, "infra"), ("NodeLost", None, "infra"), - ("DisruptionTarget", None, "infra"), - ("TerminationByKubelet", None, "infra"), + # Graceful node shutdown killing an already-running pod: the node-drain case. + ("Terminated", None, "infra"), # Pod taken while the task ran (drain, preempt, spot reclaim, force-delete). ("PodDeleted", None, "infra"), # Container ended on its own -> application. An OOM against the container's OWN From 8d4d7156b98872b8421877a5dd898c908f2c8a58 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Wed, 29 Jul 2026 17:44:33 -0700 Subject: [PATCH 4/6] Type the AIP-97 test helpers consistently The stub classes and helpers this PR adds were half-typed: one _TI stub was annotated while its sibling was not, and _eligible/_run left one parameter bare. Types the remaining ones to their real domain types (_run's stashed is the executor's failure-info tuple, not a flag). Signed-off-by: 1fanwang <1fannnw@gmail.com> --- airflow-core/tests/unit/jobs/test_scheduler_job.py | 2 +- .../unit/models/test_taskinstance_failure_kind_sources.py | 4 ++-- .../tests/unit/models/test_taskinstance_infra_refund.py | 2 +- 3 files changed, 4 insertions(+), 4 deletions(-) diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 94fc561b8cd77..940cb73705345 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -543,7 +543,7 @@ def test_process_executor_events_infra_classification(self, mock_task_callback, classification (e.g. an app OOM) does not refund.""" session = settings.Session() - def _run(dag_id, stashed): + 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) diff --git a/airflow-core/tests/unit/models/test_taskinstance_failure_kind_sources.py b/airflow-core/tests/unit/models/test_taskinstance_failure_kind_sources.py index 14b0aaa325f55..042b5ef7cb472 100644 --- a/airflow-core/tests/unit/models/test_taskinstance_failure_kind_sources.py +++ b/airflow-core/tests/unit/models/test_taskinstance_failure_kind_sources.py @@ -37,10 +37,10 @@ class _TI: - def __init__(self, max_tries=1): + def __init__(self, max_tries: int = 1): self.max_tries = max_tries - def __str__(self): + def __str__(self) -> str: return "" diff --git a/airflow-core/tests/unit/models/test_taskinstance_infra_refund.py b/airflow-core/tests/unit/models/test_taskinstance_infra_refund.py index b60afff9976c3..2bb984865ed20 100644 --- a/airflow-core/tests/unit/models/test_taskinstance_infra_refund.py +++ b/airflow-core/tests/unit/models/test_taskinstance_infra_refund.py @@ -112,7 +112,7 @@ class TestIsEligibleToRetryUsesMaxTries: max_tries grants a retry even at retries=0 (and the two retry-decision paths stay in sync).""" @staticmethod - def _eligible(*, max_tries: int, try_number: int, state=None) -> bool: + def _eligible(*, max_tries: int, try_number: int, state: TaskInstanceState | None = None) -> bool: stub = SimpleNamespace(state=state, max_tries=max_tries, try_number=try_number) return TaskInstance.is_eligible_to_retry(stub) # type: ignore[arg-type] From 5e7d1ece1179d2989312e31a7551fd065edd2b76 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Wed, 29 Jul 2026 21:22:18 -0700 Subject: [PATCH 5/6] Drop the future-work note from the reason-set comment The comment editorialised about a change that now has its own PR. Keep only what stops the footgun it warned about: condition reasons are not matched against pod.status.reason, so they do not belong in this set. Signed-off-by: 1fanwang <1fannnw@gmail.com> --- .../cncf/kubernetes/executors/kubernetes_executor_utils.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py index 7ad585ec0e4e3..414a88dc524ed 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py @@ -433,10 +433,8 @@ def collect_pod_failure_details(pod: k8s.V1Pod, logger) -> FailureDetails | None # shutdown killing an already-running pod; sole writer of this reason) # NodeLost pkg/controller/util/node/controller_utils.go#L108 # -# Known gap: a taint-manager eviction and a scheduler preemption leave pod.status.reason -# empty and are only observable via pod.status.conditions[type=DisruptionTarget].reason -# ("DeletionByTaintManager" / "PreemptionByScheduler"). Reading conditions is a separate -# change; those two disruptions currently classify as APPLICATION. +# Matched against pod.status.reason and container reasons only, so condition reasons such as +# DeletionByTaintManager do not belong here; they need pod.status.conditions. _INFRA_FAILURE_REASONS: frozenset[str] = frozenset( { "Evicted", From 8ae4f32a3a9152cd32bacae03d67a7ba8ada3121 Mon Sep 17 00:00:00 2001 From: 1fanwang <1fannnw@gmail.com> Date: Wed, 29 Jul 2026 19:44:46 -0700 Subject: [PATCH 6/6] Classify control-plane pod disruptions as infra A node drain and a scheduler preemption are infrastructure taking the pod away, but Airflow reported both to the listener as an application failure and refunded nothing. Verified live on a k8s v1.35.0 cluster: both reach phase=Failed with pod.status.reason empty and the container reading only Error/exit 143, so neither the pod reason nor the container reason can see them. The signal lives in pod.status.conditions[type=DisruptionTarget], which the control plane sets before the delete and which survives onto the terminal object. Read that condition in collect_pod_failure_details and check it first in classify_pod_failure. Gated on status "True", matching Kubernetes' own podFailurePolicy matcher, since the writers update the condition in place. Signed-off-by: 1fanwang <1fannnw@gmail.com> --- .../executors/kubernetes_executor_types.py | 1 + .../executors/kubernetes_executor_utils.py | 41 ++++++++- .../executors/test_classify_pod_failure.py | 91 +++++++++++++++++++ 3 files changed, 132 insertions(+), 1 deletion(-) diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_types.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_types.py index 45d2ea301f75b..0abfe6b8ab586 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_types.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_types.py @@ -40,6 +40,7 @@ class FailureDetails(TypedDict, total=False): exit_code: int | None container_type: Literal["init", "main"] | None container_name: str | None + disruption_reason: str | None class KubernetesResults(NamedTuple): diff --git a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py index 414a88dc524ed..19981f691e21b 100644 --- a/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py +++ b/providers/cncf/kubernetes/src/airflow/providers/cncf/kubernetes/executors/kubernetes_executor_utils.py @@ -362,6 +362,22 @@ def process_status( ) +def _disruption_target_reason(pod_status: k8s.V1PodStatus) -> str | None: + """ + Return the ``DisruptionTarget`` condition's reason, which outlives the pod's phase. + + Gated on status "True" like Kubernetes' own podFailurePolicy matcher, since the writers + update the condition in place and a stale reason can survive a flip to "False". + """ + for condition in getattr(pod_status, "conditions", None) or []: + if ( + getattr(condition, "type", None) == "DisruptionTarget" + and getattr(condition, "status", None) == "True" + ): + return getattr(condition, "reason", None) + return None + + def collect_pod_failure_details(pod: k8s.V1Pod, logger) -> FailureDetails | None: """ Collect detailed failure information from a failed pod. @@ -385,6 +401,7 @@ def collect_pod_failure_details(pod: k8s.V1Pod, logger) -> FailureDetails | None "pod_status": getattr(pod.status, "phase", None), "pod_reason": getattr(pod.status, "reason", None), "pod_message": getattr(pod.status, "message", None), + "disruption_reason": _disruption_target_reason(pod.status), } # Check init containers first (they run before main containers) @@ -434,7 +451,7 @@ def collect_pod_failure_details(pod: k8s.V1Pod, logger) -> FailureDetails | None # NodeLost pkg/controller/util/node/controller_utils.go#L108 # # Matched against pod.status.reason and container reasons only, so condition reasons such as -# DeletionByTaintManager do not belong here; they need pod.status.conditions. +# DeletionByTaintManager belong in _DISRUPTION_TARGET_REASONS below, not here. _INFRA_FAILURE_REASONS: frozenset[str] = frozenset( { "Evicted", @@ -446,6 +463,24 @@ def collect_pod_failure_details(pod: k8s.V1Pod, logger) -> FailureDetails | None } ) +# Reasons for the DisruptionTarget condition. The control plane sets it, then deletes the pod +# without ever setting pod.status.reason, so the set above cannot see these. Verified live on +# k8s v1.35.0: a node-drain taint eviction and a scheduler preemption both reach phase=Failed +# with pod.status.reason empty and only "Error"/exit 143 on the container, and the condition is +# still on that terminal object. Pinned to the same commit as above: +# PreemptionByScheduler staging/src/k8s.io/api/core/v1/types.go#L3792 +# TerminationByKubelet staging/src/k8s.io/api/core/v1/types.go#L3788 +# DeletionByTaintManager pkg/controller/tainteviction/taint_eviction.go#L139 +# DeletionByPodGC pkg/controller/podgc/gc_controller.go#L257 +_DISRUPTION_TARGET_REASONS: frozenset[str] = frozenset( + { + "PreemptionByScheduler", + "TerminationByKubelet", + "DeletionByTaintManager", + "DeletionByPodGC", + } +) + def classify_pod_failure( failure_details: FailureDetails | None, @@ -464,6 +499,10 @@ def classify_pod_failure( pod_reason = failure_details.get("pod_reason") container_reason = failure_details.get("container_reason") + disruption_reason = failure_details.get("disruption_reason") + + if disruption_reason in _DISRUPTION_TARGET_REASONS: + return TaskFailureKind.INFRA, disruption_reason if pod_reason in _INFRA_FAILURE_REASONS or container_reason in _INFRA_FAILURE_REASONS: failure_kind = TaskFailureKind.INFRA diff --git a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_classify_pod_failure.py b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_classify_pod_failure.py index feed64013009c..efb8e5be5a78a 100644 --- a/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_classify_pod_failure.py +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_classify_pod_failure.py @@ -140,3 +140,94 @@ def test_base_executor_failure_info_round_trips_once(self): # first read returns it, and clears it so a later event can't reuse stale context assert ex.get_task_failure_info(key) is classified assert ex.get_task_failure_info(key) is None + + +class TestDisruptionTargetCondition: + """A control-plane disruption is only visible in the DisruptionTarget condition.""" + + @pytest.mark.parametrize( + "disruption_reason", + ["PreemptionByScheduler", "DeletionByTaintManager", "TerminationByKubelet", "DeletionByPodGC"], + ) + def test_disruption_condition_is_infra(self, disruption_reason): + details = { + "pod_status": "Failed", + "pod_reason": None, + "container_reason": "Error", + "disruption_reason": disruption_reason, + } + failure_kind, infra_reason = classify_pod_failure(details) + assert failure_kind == "infra" + assert infra_reason == disruption_reason + + def test_unrelated_condition_reason_stays_application(self): + details = {"pod_status": "Failed", "container_reason": "Error", "disruption_reason": "SomethingElse"} + failure_kind, _ = classify_pod_failure(details) + assert failure_kind == "application" + + @pytest.mark.parametrize( + ("disruption_reason", "pod_name"), + [("DeletionByTaintManager", "aip97-victim"), ("PreemptionByScheduler", "aip97-victim2")], + ) + def test_end_to_end_live_shape_is_infra(self, disruption_reason, pod_name): + # The exact object observed on a real k8s v1.35.0 cluster: a node-drain taint eviction + # and a scheduler preemption both land on phase=Failed with reason empty and the + # container reading only Error/143, so without the condition both look like an app crash. + pod = k8s.V1Pod( + metadata=k8s.V1ObjectMeta(name=pod_name, deletion_timestamp="2026-07-29T19:21:36Z"), + status=k8s.V1PodStatus( + phase="Failed", + reason=None, + conditions=[ + k8s.V1PodCondition(type="DisruptionTarget", status="True", reason=disruption_reason) + ], + container_statuses=[ + k8s.V1ContainerStatus( + name="base", + image="busybox:1.36", + image_id="", + ready=False, + restart_count=0, + state=k8s.V1ContainerState( + terminated=k8s.V1ContainerStateTerminated(reason="Error", exit_code=143) + ), + ) + ], + ), + ) + details = collect_pod_failure_details(pod, logging.getLogger("test")) + assert details is not None + assert details["pod_reason"] is None + assert details["container_reason"] == "Error" + assert details["disruption_reason"] == disruption_reason + + failure_kind, infra_reason = classify_pod_failure(details) + assert failure_kind == "infra" + assert infra_reason == disruption_reason + + @pytest.mark.parametrize( + "conditions", + [ + pytest.param( + [k8s.V1PodCondition(type="Ready", status="False", reason="DeletionByTaintManager")], + id="reason-on-a-different-condition-type", + ), + pytest.param( + [ + k8s.V1PodCondition( + type="DisruptionTarget", status="False", reason="DeletionByTaintManager" + ) + ], + id="disruption-condition-no-longer-true", + ), + ], + ) + def test_only_a_true_disruption_target_counts(self, conditions): + pod = k8s.V1Pod( + metadata=k8s.V1ObjectMeta(name="aip97-negative"), + status=k8s.V1PodStatus(phase="Failed", reason=None, conditions=conditions), + ) + details = collect_pod_failure_details(pod, logging.getLogger("test")) + assert details is not None + assert details["disruption_reason"] is None + assert classify_pod_failure(details)[0] != "infra"