Allow for retry when tasks are stuck in queued

airflow/config_templates/config.yml

@@ -2331,6 +2331,14 @@ scheduler:
       type: integer
       example: ~
       default: "8974"
+    num_stuck_in_queued_retries:
+      description: |
+        Number of times Airflow will reschedule a task that gets stuck in queued before marking the
+        task as failed
+      version_added: 2.10.4
+      type: integer
+      example: ~
+      default: "3"
     orphaned_tasks_check_interval:
       description: |
         How often (in seconds) should the scheduler check for orphaned tasks and SchedulerJobs

airflow/executors/base_executor.py

@@ -22,12 +22,14 @@
 import sys
 from collections import defaultdict, deque
 from dataclasses import dataclass, field
-from typing import TYPE_CHECKING, Any, List, Optional, Sequence, Tuple
+from typing import TYPE_CHECKING, Any, Iterable, List, Optional, Sequence, Tuple
 
 import pendulum
+from deprecated import deprecated
 
 from airflow.cli.cli_config import DefaultHelpParser
 from airflow.configuration import conf
+from airflow.exceptions import RemovedInAirflow3Warning
 from airflow.executors.executor_loader import ExecutorLoader
 from airflow.models import Log
 from airflow.stats import Stats
@@ -540,7 +542,12 @@ def terminate(self):
         """Get called when the daemon receives a SIGTERM."""
         raise NotImplementedError
 
-    def cleanup_stuck_queued_tasks(self, tis: list[TaskInstance]) -> list[str]:  # pragma: no cover
+    @deprecated(
+        reason="Replaced by function `cleanup_tasks_stuck_in_queued`.",
+        category=RemovedInAirflow3Warning,
+        action="ignore",
+    )
+    def cleanup_stuck_queued_tasks(self, tis: list[TaskInstance]) -> list[str]:
         """
         Handle remnants of tasks that were failed because they were stuck in queued.
 
@@ -551,7 +558,21 @@ def cleanup_stuck_queued_tasks(self, tis: list[TaskInstance]) -> list[str]:  # p
         :param tis: List of Task Instances to clean up
         :return: List of readable task instances for a warning message
         """
-        raise NotImplementedError()
+        raise NotImplementedError
+
+    def cleanup_tasks_stuck_in_queued(self, *, tis: Iterable[TaskInstance]) -> Iterable[TaskInstance]:
+        """
+        Try to remove the tis from the executor so they can safely be requeued.
+
+        This function will be called for task instances that the scheduler detects
+        as having been in the queued state for too long.
+
+        The timeout duration is controlled by scheduler setting ``task_queued_timeout``.
+
+        :param tis: List of Task Instances to clean up
+        :return: List of task instances that were removed from the executor.
+        """
+        raise NotImplementedError
 
     def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[TaskInstance]:
         """

airflow/jobs/scheduler_job_runner.py

@@ -25,11 +25,13 @@
 import sys
 import time
 from collections import Counter, defaultdict, deque
+from contextlib import suppress
 from datetime import timedelta
 from functools import lru_cache, partial
 from pathlib import Path
 from typing import TYPE_CHECKING, Any, Callable, Collection, Iterable, Iterator
 
+from deprecated import deprecated
 from sqlalchemy import and_, delete, exists, func, not_, select, text, update
 from sqlalchemy.exc import OperationalError
 from sqlalchemy.orm import lazyload, load_only, make_transient, selectinload
@@ -39,7 +41,7 @@
 from airflow.callbacks.callback_requests import DagCallbackRequest, TaskCallbackRequest
 from airflow.callbacks.pipe_callback_sink import PipeCallbackSink
 from airflow.configuration import conf
-from airflow.exceptions import UnknownExecutorException
+from airflow.exceptions import RemovedInAirflow3Warning, UnknownExecutorException
 from airflow.executors.executor_loader import ExecutorLoader
 from airflow.jobs.base_job_runner import BaseJobRunner
 from airflow.jobs.job import Job, perform_heartbeat
@@ -98,6 +100,9 @@
 DR = DagRun
 DM = DagModel
 
+TASK_STUCK_IN_QUEUED_RESCHEDULE_EVENT = "stuck in queued reschedule"
+""":meta private:"""
+
 
 class ConcurrencyMap:
     """
@@ -1067,7 +1072,7 @@ def _run_scheduler_loop(self) -> None:
 
         timers.call_regular_interval(
             conf.getfloat("scheduler", "task_queued_timeout_check_interval"),
-            self._fail_tasks_stuck_in_queued,
+            self._handle_tasks_stuck_in_queued,
         )
 
         timers.call_regular_interval(
@@ -1115,6 +1120,7 @@ def _run_scheduler_loop(self) -> None:
                 for executor in self.job.executors:
                     try:
                         # this is backcompat check if executor does not inherit from BaseExecutor
+                        # todo: remove in airflow 3.0
                         if not hasattr(executor, "_task_event_logs"):
                             continue
                         with create_session() as session:
@@ -1776,15 +1782,23 @@ def _send_dag_callbacks_to_processor(self, dag: DAG, callback: DagCallbackReques
             self.log.debug("callback is empty")
 
     @provide_session
-    def _fail_tasks_stuck_in_queued(self, session: Session = NEW_SESSION) -> None:
+    def _handle_tasks_stuck_in_queued(self, session: Session = NEW_SESSION) -> None:
         """
-        Mark tasks stuck in queued for longer than `task_queued_timeout` as failed.
+
+        Handle the scenario where a task is queued for longer than `task_queued_timeout`.
 
         Tasks can get stuck in queued for a wide variety of reasons (e.g. celery loses
         track of a task, a cluster can't further scale up its workers, etc.), but tasks
-        should not be stuck in queued for a long time. This will mark tasks stuck in
-        queued for longer than `self._task_queued_timeout` as failed. If the task has
-        available retries, it will be retried.
+        should not be stuck in queued for a long time.
+
+        Originally, we simply marked a task as failed when it was stuck in queued for
+        too long. We found that this led to suboptimal outcomes as ideally we would like "failed"
+        to mean that a task was unable to run, instead of it meaning that we were unable to run the task.
+
+        As a compromise between always failing a stuck task and always rescheduling a stuck task (which could
+        lead to tasks being stuck in queued forever without informing the user), we have creating the config
+        ``[scheduler] num_stuck_in_queued_retries``. With this new configuration, an airflow admin can decide how
+        sensitive they would like their airflow to be WRT failing stuck tasks.
         """
         self.log.debug("Calling SchedulerJob._fail_tasks_stuck_in_queued method")
 
@@ -1796,28 +1810,108 @@ def _fail_tasks_stuck_in_queued(self, session: Session = NEW_SESSION) -> None:
             )
         ).all()
 
+        num_allowed_retries = conf.getint("scheduler", "num_stuck_in_queued_retries")
         for executor, stuck_tis in self._executor_to_tis(tasks_stuck_in_queued).items():
             try:
-                cleaned_up_task_instances = set(executor.cleanup_stuck_queued_tasks(tis=stuck_tis))
-                for ti in stuck_tis:
-                    if repr(ti) in cleaned_up_task_instances:
-                        self.log.warning(
-                            "Marking task instance %s stuck in queued as failed. "
-                            "If the task instance has available retries, it will be retried.",
-                            ti,
+                # BaseExecutor has "abstract" method `cleanup_tasks_stuck_in_queued`
+                # We are tolerant of implementers not implementing it.
+                tis = executor.cleanup_tasks_stuck_in_queued(tis=stuck_tis)
+            except NotImplementedError:
+                # todo: remove the call to _stuck_in_queued_backcompat_logic in airflow 3.0.
+                #   after 3.0, `cleanup_stuck_queued_tasks` will be removed, so we should
+                #   just continue immediately.
+                self._stuck_in_queued_backcompat_logic(executor, stuck_tis)
+                continue
+
+            # ok we know that we now have a "modern" version of the executor, which
+            # expects us to try to requeue tasks "cleaned up" by `cleanup_stuck_queued_tasks`
+            for ti in tis:
+                num_times_stuck = self._get_num_times_stuck_in_queued(ti, session)
+                if num_times_stuck < num_allowed_retries:
+                    self.log.warning("Task stuck in queued; will try to requeue. task_id=%s", ti.task_id)
+                    session.add(
+                        Log(
+                            event=TASK_STUCK_IN_QUEUED_RESCHEDULE_EVENT,
+                            task_instance=ti.key,
+                            extra=(
+                                f"Task was in queued state for longer than {self._task_queued_timeout} "
+                                "seconds; task state will be set back to scheduled."
+                            ),
                         )
-                        session.add(
-                            Log(
-                                event="stuck in queued",
-                                task_instance=ti.key,
-                                extra=(
-                                    "Task will be marked as failed. If the task instance has "
-                                    "available retries, it will be retried."
-                                ),
-                            )
+                    )
+                    self._reschedule_stuck_task(ti)
+                else:
+                    self.log.warning(
+                        "Task requeue attempts exceeded max; marking failed. task_instance=%s",
+                        ti,
+                    )
+                    session.add(
+                        Log(
+                            event="stuck in queued tries exceeded",
+                            task_instance=ti.key,
+                            extra=(
+                                f"Task was requeued more than {num_allowed_retries} times "
+                                "and will be failed."
+                            ),
                         )
-            except NotImplementedError:
-                self.log.debug("Executor doesn't support cleanup of stuck queued tasks. Skipping.")
+                    )
+                    ti.set_state(TaskInstanceState.FAILED, session=session)
+                with suppress(KeyError):
+                    executor.running.remove(ti.key)
+
+    @deprecated(
+        reason="This is backcompat layer for older executor interface. Should be removed in 3.0",
+        category=RemovedInAirflow3Warning,
+        action="ignore",
+    )
+    def _stuck_in_queued_backcompat_logic(self, executor, stuck_tis):
+        """
+        Try to invoke stuck in queued cleanup for older executor interface.
+
+        TODO: remove in airflow 3.0
+
+        Here we handle case where the executor pre-dates the interface change that
+        introduced `cleanup_tasks_stuck_in_queued` and deprecated `cleanup_stuck_queued_tasks`.
+
+        """
+        with suppress(NotImplementedError):
+            for ti_repr in executor.cleanup_stuck_queued_tasks(tis=stuck_tis):
+                self.log.warning(
+                    "Task instance %s stuck in queued.  May be set to failed.",
+                    ti_repr,
+                )
+
+    @provide_session
+    def _reschedule_stuck_task(self, ti, session=NEW_SESSION):
+        session.execute(
+            update(TI)
+            .where(TI.filter_for_tis([ti]))
+            .values(
+                state=TaskInstanceState.SCHEDULED,
+                queued_dttm=None,
+            )
+            .execution_options(synchronize_session=False)
+        )
+
+    @provide_session
+    def _get_num_times_stuck_in_queued(self, ti: TaskInstance, session: Session = NEW_SESSION) -> int:
+        """
+        Check the Log table to see how many times a taskinstance has been stuck in queued.
+
+        We can then use this information to determine whether to reschedule a task or fail it.
+        """
+        return (
+            session.query(Log)
+            .where(
+                Log.task_id == ti.task_id,
+                Log.dag_id == ti.dag_id,
+                Log.run_id == ti.run_id,
+                Log.map_index == ti.map_index,
+                Log.try_number == ti.try_number,
+                Log.event == TASK_STUCK_IN_QUEUED_RESCHEDULE_EVENT,
+            )
+            .count()
+        )
 
     @provide_session
     def _emit_pool_metrics(self, session: Session = NEW_SESSION) -> None:

docs/spelling_wordlist.txt

@@ -1375,6 +1375,7 @@ repos
 repr
 req
 reqs
+requeued
 Reserialize
 reserialize
 reserialized

providers/src/airflow/providers/celery/executors/celery_executor.py

@@ -32,9 +32,10 @@
 from collections import Counter
 from concurrent.futures import ProcessPoolExecutor
 from multiprocessing import cpu_count
-from typing import TYPE_CHECKING, Any, Optional, Sequence, Tuple
+from typing import TYPE_CHECKING, Any, Generator, Iterable, Optional, Sequence, Tuple
 
 from celery import states as celery_states
+from deprecated import deprecated
 from packaging.version import Version
 
 from airflow import __version__ as airflow_version
@@ -52,7 +53,7 @@
     lazy_load_command,
 )
 from airflow.configuration import conf
-from airflow.exceptions import AirflowTaskTimeout
+from airflow.exceptions import AirflowProviderDeprecationWarning, AirflowTaskTimeout
 from airflow.executors.base_executor import BaseExecutor
 from airflow.stats import Stats
 from airflow.utils.state import TaskInstanceState
@@ -433,31 +434,42 @@ def try_adopt_task_instances(self, tis: Sequence[TaskInstance]) -> Sequence[Task
 
         return not_adopted_tis
 
+    @deprecated(
+        reason="Replaced by function `cleanup_tasks_stuck_in_queued`. "
+        "Upgrade airflow core to make this go away.",
+        category=AirflowProviderDeprecationWarning,
+    )
     def cleanup_stuck_queued_tasks(self, tis: list[TaskInstance]) -> list[str]:
         """
-        Handle remnants of tasks that were failed because they were stuck in queued.
+        Remove tasks stuck in queued from executor and fail them.
 
-        Tasks can get stuck in queued. If such a task is detected, it will be marked
-        as `UP_FOR_RETRY` if the task instance has remaining retries or marked as `FAILED`
-        if it doesn't.
+        This method is deprecated. Use `cleanup_tasks_stuck_in_queued` instead.
+        """
+        cleaned_up_tis = self.cleanup_tasks_stuck_in_queued(tis=tis)
+        reprs = []
+        for ti in cleaned_up_tis:
+            reprs.append(repr(ti))
+            self.fail(ti.key)
+        return reprs
+
+    def cleanup_tasks_stuck_in_queued(
+        self, *, tis: Iterable[TaskInstance]
+    ) -> Generator[TaskInstance, None, None]:
+        """
+        Given list of tis that are stuck in queued, remove them from executor.
 
-        :param tis: List of Task Instances to clean up
-        :return: List of readable task instances for a warning message
+        :param tis: List of task instances to remove from this executor.
         """
-        readable_tis = []
         from airflow.providers.celery.executors.celery_executor_utils import app
 
         for ti in tis:
-            readable_tis.append(repr(ti))
-            task_instance_key = ti.key
-            self.fail(task_instance_key, None)
-            celery_async_result = self.tasks.pop(task_instance_key, None)
+            celery_async_result = self.tasks.pop(ti.key, None)
             if celery_async_result:
                 try:
                     app.control.revoke(celery_async_result.task_id)
                 except Exception as ex:
-                    self.log.error("Error revoking task instance %s from celery: %s", task_instance_key, ex)
-        return readable_tis
+                    self.log.error("Error revoking task instance %s from celery: %s", ti.key, ex)
+            yield ti
 
     @staticmethod
     def get_cli_commands() -> list[GroupCommand]: