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..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 @@ -28,7 +28,7 @@ from sqlalchemy.orm import joinedload from sqlalchemy.orm.session import Session -from airflow._shared.state import TaskScope +from airflow._shared.state import TaskFailureKind, TaskScope from airflow.api_fastapi.app import get_auth_manager from airflow.api_fastapi.auth.managers.models.resource_details import DagAccessEntity, DagDetails from airflow.api_fastapi.common.dagbag import DagBagDep, get_latest_version_of_dag @@ -117,6 +117,8 @@ def _emit_state_listener_hooks(updated_tis: list[TI], new_state: str | TaskInsta previous_state=None, task_instance=ti, error=f"TaskInstance's state was manually set to `{TaskInstanceState.FAILED}`.", + failure_kind=TaskFailureKind.MANUAL, + reason=None, ) elif new_state == TaskInstanceState.SKIPPED: get_listener_manager().hook.on_task_instance_skipped(previous_state=None, task_instance=ti) diff --git a/airflow-core/src/airflow/executors/base_executor.py b/airflow-core/src/airflow/executors/base_executor.py index 8030e9c14455e..93452ec1b794e 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,7 @@ def __init__(self, parallelism: int = PARALLELISM, team_name: str | None = None) self.queued_connection_tests: dict[ConnectionTestKey, workloads.TestConnection] = {} self.running: set[WorkloadKey] = set() self.event_buffer: dict[WorkloadKey, EventBufferValueType] = {} + self.task_failure_info: dict[WorkloadKey, tuple[TaskFailureKind, str | None]] = {} self._task_event_logs: deque[Log] = deque() self.conf = ExecutorConf(team_name) @@ -550,6 +552,14 @@ def get_event_buffer(self, dag_ids=None) -> dict[WorkloadKey, EventBufferValueTy return cleared_events + def get_task_failure_info(self, key: WorkloadKey) -> tuple[TaskFailureKind, str | None] | None: + """ + Return and clear the ``(failure_kind, reason)`` this executor classified for ``key``. + + None unless the executor can read the real cause, as the Kubernetes executor does. + """ + return self.task_failure_info.pop(key, None) + def get_task_log(self, ti: TaskInstance, try_number: int) -> tuple[list[str], list[str]]: """ Return the task logs. diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 40d1651505794..d66ae5f0a3785 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 @@ -1468,6 +1469,11 @@ def process_executor_events( job_id, ) state, info = event_buffer.pop(buffer_key) + # Pop unconditionally, not only on the killed-externally branch below, so a + # self-reporting task can't leak its entry. + executor_failure_kind: tuple[TaskFailureKind, str | None] | None = executor.get_task_failure_info( + ti.key + ) if state in (TaskInstanceState.QUEUED, TaskInstanceState.RUNNING): ti.external_executor_id = info @@ -1663,7 +1669,19 @@ def process_executor_events( executor.send_callback(email_request) # Update task state - emails are handled by DAG processor now - ti.handle_failure(error=msg, session=session) + # Refund only on a positive infra classification. An executor that reports + # nothing but a state can't tell an eviction from a silent crash, so it + # stays unclassified and spends the retry as today. + if executor_failure_kind is not None: + failure_kind, reason = executor_failure_kind + else: + failure_kind, reason = None, None + ti.handle_failure( + error=msg, + session=session, + failure_kind=failure_kind, + reason=reason if failure_kind == TaskFailureKind.INFRA else None, + ) cls._emit_executor_events_batch_metrics(num_events) return len(event_buffer) diff --git a/airflow-core/src/airflow/models/taskinstance.py b/airflow-core/src/airflow/models/taskinstance.py index dbcb311afe447..753658cba2b20 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,48 @@ def _date_or_empty(*, task_instance: TaskInstance, attr: str) -> str: return result.strftime("%Y%m%dT%H%M%S") if result else "" +def _maybe_refund_infra_attempt( + *, + task_instance: TaskInstance, + task: Operator | None, + failure_kind: TaskFailureKind | None, + reason: str | None = None, +) -> bool: + """ + Refund one retry attempt (bump ``max_tries``) for an infra-classified failure. + + Opt-in and capped by config. Every other kind, and an unclassified failure, + spends the user's retry as before. Returns True if an attempt was refunded. + + :meta private: + """ + # retries may be None (unset); treat it as falsy like is_eligible_to_retry does. + # Refunds still apply at retries=0: an infra disruption is not the user's failure, + # so it earns an attempt back regardless of their retry budget. + retries = getattr(task, "retries", None) or 0 + if ( + failure_kind != TaskFailureKind.INFRA + or task is None + or not conf.getboolean("core", "infra_failure_refund_retries", fallback=False) + ): + return False + max_infra_refunds = conf.getint("core", "max_infra_refunds", fallback=3) + refunds_used = max((task_instance.max_tries or 0) - retries, 0) + if refunds_used >= max_infra_refunds: + return False + task_instance.max_tries += 1 + log.info( + "infra failure (%s) refunded attempt for %s; max_tries now %s " + "(refund %s/%s), user retry budget preserved", + reason, + task_instance, + task_instance.max_tries, + refunds_used + 1, + max_infra_refunds, + ) + return True + + def uuid7() -> UUID: """Generate a new UUID7.""" return uuid6.uuid7() @@ -650,7 +693,6 @@ class TaskInstance(Base, LoggingMixin, BaseWorkload): # Cleared on task start (ti_run). Read by next_retry_datetime(). retry_delay_override: Mapped[float | None] = mapped_column(Float, nullable=True) retry_reason: Mapped[str | None] = mapped_column(String(500), nullable=True) - __table_args__ = ( Index("ti_dag_state", dag_id, state), Index("ti_dag_run", dag_id, run_id), @@ -1849,6 +1891,8 @@ def fetch_handle_failure_context( *, session: Session, fail_fast: bool = False, + failure_kind: TaskFailureKind | None = None, + reason: str | None = None, ): """ Fetch the context needed to handle a failure. @@ -1858,6 +1902,12 @@ def fetch_handle_failure_context( :param test_mode: doesn't record success or failure in the DB if True :param session: SQLAlchemy ORM Session :param fail_fast: if True, fail all downstream tasks + :param failure_kind: what caused the failure (:class:`TaskFailureKind` or + ``None``). ``INFRA`` refunds the attempt instead of charging the user's + retry budget, gated by config. + :param reason: the executor's token for an infra failure (e.g. ``Evicted``). + The worker reported no error in that case, so this is passed to the + listener rather than persisted. """ if error: cls.logger().error("%s", error) @@ -1889,6 +1939,8 @@ def fetch_handle_failure_context( # Actual callbacks are handled by the DAG processor, not the scheduler task = getattr(ti, "task", None) + _maybe_refund_infra_attempt(task_instance=ti, task=task, failure_kind=failure_kind, reason=reason) + if not ti.is_eligible_to_retry(): ti.state = TaskInstanceState.FAILED @@ -1905,7 +1957,11 @@ def fetch_handle_failure_context( try: get_listener_manager().hook.on_task_instance_failed( - previous_state=TaskInstanceState.RUNNING, task_instance=ti, error=error + previous_state=TaskInstanceState.RUNNING, + task_instance=ti, + error=error, + failure_kind=failure_kind, + reason=reason, ) except Exception: log.exception("error calling listener") @@ -1927,6 +1983,8 @@ def handle_failure( test_mode: bool | None = None, *, session: Session = NEW_SESSION, + failure_kind: TaskFailureKind | None = None, + reason: str | None = None, ) -> None: """ Handle Failure for a task instance. @@ -1934,6 +1992,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, passed to the + listener rather than persisted. """ if TYPE_CHECKING: assert self.task @@ -1950,6 +2012,8 @@ def handle_failure( test_mode=test_mode, session=session, fail_fast=fail_fast, + failure_kind=failure_kind, + reason=reason, ) _log_state(task_instance=self) @@ -1963,15 +2027,9 @@ def is_eligible_to_retry(self) -> bool: # If a task is cleared when running, it goes into RESTARTING state and is always # eligible for retry return True - if not getattr(self, "task", None): - # Couldn't load the task, don't know number of retries, guess: - return self.try_number <= self.max_tries - - if TYPE_CHECKING: - assert self.task - assert self.task.retries - - return bool(self.task.retries and self.try_number <= self.max_tries) + # Mirror of the execution API's _is_eligible_to_retry; the scheduler and worker + # retry-decision paths must not diverge. See its comment for why retries is not consulted. + return self.max_tries != 0 and self.try_number <= self.max_tries def set_duration(self) -> None: """Set task instance duration.""" diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index c83d9039d8f6e..5d23703b157df 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -45,6 +45,7 @@ from airflow import settings from airflow._shared.module_loading import qualname from airflow._shared.observability.metrics.base_stats_logger import StatsLogger +from airflow._shared.state import TaskFailureKind from airflow._shared.timezones import timezone from airflow.api_fastapi.auth.tokens import JWTGenerator from airflow.assets.manager import AssetManager @@ -536,6 +537,42 @@ def test_process_executor_events(self, mock_get_backend, mock_task_callback, dag any_order=True, ) + @conf_vars({("core", "infra_failure_refund_retries"): "True", ("core", "max_infra_refunds"): "3"}) + @mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest", spec=TaskCallbackRequest) + def test_process_executor_events_infra_classification(self, mock_task_callback, dag_maker): + """Only a positively-classified infra failure refunds; an unclassified + state-mismatch death (no executor bridge) spends the retry, and a bridge 'user' + classification (e.g. an app OOM) does not refund.""" + session = settings.Session() + + def _run(dag_id: str, stashed: tuple[TaskFailureKind, str | None] | None) -> TaskInstance: + with dag_maker(dag_id=dag_id, fileloc=f"/{dag_id}/"): + task = EmptyOperator(task_id="t", retries=1) + ti = dag_maker.create_dagrun().get_task_instance(task.task_id) + ti.state = State.QUEUED + session.merge(ti) + session.commit() + executor = MockExecutor(do_update=False) + job_runner = SchedulerJobRunner(job=Job(), executors=[executor]) + executor.event_buffer[ti.key] = State.FAILED, None + if stashed is not None: + executor.task_failure_info[ti.key] = stashed + job_runner._process_executor_events(executor=executor, session=session) + ti.refresh_from_db(session=session) + return ti + + # Unclassified: no refund, exactly as today. + ti = _run("aip97_unclassified", stashed=None) + assert ti.max_tries == 1 + + # Classified infra (a pod Evicted): refunded. + ti = _run("aip97_infra", stashed=(TaskFailureKind.INFRA, "Evicted")) + assert ti.max_tries == 2 + + # Classified application (an OOM against its own limit): no refund. + ti = _run("aip97_app_oom", stashed=(TaskFailureKind.APPLICATION, "OOMKilled")) + assert ti.max_tries == 1 + @mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest", spec=TaskCallbackRequest) def test_process_executor_events_restarting_cleared_task(self, mock_task_callback, dag_maker): """ @@ -9623,6 +9660,7 @@ def test_scheduler_passes_context_from_server_on_task_failure(self, dag_maker, s # Mock the executor to simulate a task failure mock_executor = MagicMock(spec=BaseExecutor) mock_executor.has_task = mock.MagicMock(return_value=False) + mock_executor.get_task_failure_info.return_value = None scheduler_job = Job() self.job_runner = SchedulerJobRunner(scheduler_job, executors=[mock_executor]) 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..b405d63a8708f --- /dev/null +++ b/airflow-core/tests/unit/listeners/test_listener_types.py @@ -0,0 +1,78 @@ +# 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, "reason": None} + + # 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( + self, + previous_state, + task_instance, + error, + failure_kind, + reason, + ): + received["failure_kind"] = failure_kind + received["reason"] = reason + + listener_manager(InfraListener()) + + get_listener_manager().hook.on_task_instance_failed( + previous_state=None, + task_instance=None, + error=RuntimeError("boom"), + failure_kind="infra", + reason="Evicted", + ) + + 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 + 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", + 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 new file mode 100644 index 0000000000000..042b5ef7cb472 --- /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: int = 1): + self.max_tries = max_tries + + def __str__(self) -> str: + 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..2bb984865ed20 --- /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 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 + + @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. + 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: 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] + + 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 b75815df2fd91..7920c5ca5a0ab 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 @@ -718,6 +718,16 @@ def _change_state( container_message, exit_code, ) + + # 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, + ) + + 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_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 0bf1ed8b23e6d..9d78918aa66df 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 @@ -47,6 +47,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 @@ -57,6 +58,8 @@ from kubernetes.client import Configuration, models as k8s + from airflow._shared.state import TaskFailureKind + class ResourceVersion: """Singleton for tracking resourceVersion from Kubernetes.""" @@ -333,18 +336,20 @@ 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, ) + # 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, - namespace, - TaskInstanceState.FAILED, - annotations, - resource_version, - None, + 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: @@ -361,6 +366,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. @@ -384,6 +405,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) @@ -413,6 +435,87 @@ 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: str = "PodDeleted" + +# 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 +# +# Matched against pod.status.reason and container reasons only, so condition reasons such as +# DeletionByTaintManager belong in _DISRUPTION_TARGET_REASONS below, not here. +_INFRA_FAILURE_REASONS: frozenset[str] = frozenset( + { + "Evicted", + "Preempting", + "NodeShutdown", + "Terminated", + "NodeLost", + POD_DELETED_REASON, + } +) + +# 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, +) -> tuple[TaskFailureKind, str | None] | None: + """ + Classify a pod failure into a ``(TaskFailureKind, reason)`` pair. + + 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 + + from airflow._shared.state import TaskFailureKind + + 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 + 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..efb8e5be5a78a --- /dev/null +++ b/providers/cncf/kubernetes/tests/unit/cncf/kubernetes/executors/test_classify_pod_failure.py @@ -0,0 +1,233 @@ +# 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. + ("Evicted", None, "infra"), + ("Preempting", None, "infra"), + ("NodeShutdown", None, "infra"), + ("NodeLost", 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 + # 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. + ("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) classifies as + # application, so an app OOM earns no 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 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( + 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): + # 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"} + ) + 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 + + +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" diff --git a/shared/listeners/src/airflow_shared/listeners/spec/taskinstance.py b/shared/listeners/src/airflow_shared/listeners/spec/taskinstance.py index d3450d6b05aa7..aec552d646ec6 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,31 @@ 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, 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, 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): ... + """ @hookspec diff --git a/shared/state/src/airflow_shared/state/__init__.py b/shared/state/src/airflow_shared/state/__init__.py index e2c094d84818d..583cab5a0c5be 100644 --- a/shared/state/src/airflow_shared/state/__init__.py +++ b/shared/state/src/airflow_shared/state/__init__.py @@ -72,6 +72,28 @@ 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, 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" + 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 1ce848128d663..df3d905fbdfd1 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 ( @@ -2366,7 +2367,13 @@ 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, + reason=None, ) except Exception: log.exception("error calling listener") @@ -2376,7 +2383,13 @@ 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, + 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 cc9fb77e08921..3dc4191c8c75b 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 @@ -64,7 +64,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, @@ -4731,6 +4731,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 ):