Allow for retry when tasks are stuck in queued

Author: dimberman

airflow/executors/base_executor.py

@@ -478,6 +478,15 @@ def queued(self, key: TaskInstanceKey, info=None) -> None:
         """
         self.change_state(key, TaskInstanceState.QUEUED, info)
 
+    def retry_stuck_task(self, key: TaskInstanceKey, info=None) -> None:
+        """
+        Set queued state for the event.
+
+        :param info: Executor information for the task instance
+        :param key: Unique key for the task instance
+        """
+        self.change_state(key, TaskInstanceState.STUCK_RESCHEDULE, info)
+
     def running_state(self, key: TaskInstanceKey, info=None) -> None:
         """
         Set running state for the event.

airflow/jobs/scheduler_job_runner.py

@@ -1075,7 +1075,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(
@@ -1782,7 +1782,7 @@ 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.
 
@@ -1825,6 +1825,8 @@ def _fail_tasks_stuck_in_queued(self, session: Session = NEW_SESSION) -> None:
                         num_times_stuck = self._get_num_times_stuck(ti, session)
                         if num_times_stuck < conf.getint("core", "num_stuck_retries", fallback=3):
                             self._reset_task_instance(ti, session)
+                        else:
+                            executor.fail(ti.key)
 
 
             except NotImplementedError:
@@ -1840,6 +1842,7 @@ def _get_num_times_stuck(self, ti: TaskInstance, session: Session = NEW_SESSION)
                         .where(Log.event == "stuck in queued")
                         ))
 
+    @provide_session
     def _reset_task_instance(self, ti: TaskInstance, session: Session = NEW_SESSION):
         ti.external_executor_id = None
         ti.state = State.SCHEDULED

airflow/utils/state.py

@@ -53,6 +53,7 @@ class TaskInstanceState(str, Enum):
     SUCCESS = "success"  # Task completed
     RESTARTING = "restarting"  # External request to restart (e.g. cleared when running)
     FAILED = "failed"  # Task errored out
+    STUCK_RESCHEDULE = "stuck_reschedule"  # Task got stuck in queued but is up for resch
     UP_FOR_RETRY = "up_for_retry"  # Task failed but has retries left
     UP_FOR_RESCHEDULE = "up_for_reschedule"  # A waiting `reschedule` sensor
     UPSTREAM_FAILED = "upstream_failed"  # One or more upstream deps failed

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

@@ -450,7 +450,7 @@ def cleanup_stuck_queued_tasks(self, tis: list[TaskInstance]) -> list[str]:
         for ti in tis:
             readable_tis.append(repr(ti))
             task_instance_key = ti.key
-            self.fail(task_instance_key, None)
+            self.retry_stuck_task(task_instance_key, None)
             celery_async_result = self.tasks.pop(task_instance_key, None)
             if celery_async_result:
                 try:

tests/jobs/test_scheduler_job.py

@@ -2202,7 +2202,7 @@ def test_adopt_or_reset_orphaned_tasks_multiple_executors(self, dag_maker, mock_
         # Second executor called for ti3
         mock_executors[1].try_adopt_task_instances.assert_called_once_with([ti3])
 
-    def test_fail_stuck_queued_tasks(self, dag_maker, session, mock_executors):
+    def test_handle_stuck_queued_tasks(self, dag_maker, session, mock_executors):
         with dag_maker("test_fail_stuck_queued_tasks_multiple_executors"):
             op1 = EmptyOperator(task_id="op1")
             op2 = EmptyOperator(task_id="op2", executor="default_exec")
@@ -2228,7 +2228,7 @@ def test_fail_stuck_queued_tasks(self, dag_maker, session, mock_executors):
                 (None,): mock_executors[0],
                 ("secondary_exec",): mock_executors[1],
             }[x]
-            job_runner._fail_tasks_stuck_in_queued()
+            job_runner._handle_tasks_stuck_in_queued()
 
         # Default executor is called for ti1 (no explicit executor override uses default) and ti2 (where we
         # explicitly marked that for execution by the default executor)
@@ -2238,7 +2238,7 @@ def test_fail_stuck_queued_tasks(self, dag_maker, session, mock_executors):
             mock_executors[0].cleanup_stuck_queued_tasks.assert_called_once_with(tis=[ti2, ti1])
         mock_executors[1].cleanup_stuck_queued_tasks.assert_called_once_with(tis=[ti3])
 
-    def test_fail_stuck_queued_tasks_raises_not_implemented(self, dag_maker, session, caplog):
+    def test_handle_stuck_queued_tasks_raises_not_implemented(self, dag_maker, session, caplog):
         with dag_maker("test_fail_stuck_queued_tasks"):
             op1 = EmptyOperator(task_id="op1")
 
@@ -2253,7 +2253,7 @@ def test_fail_stuck_queued_tasks_raises_not_implemented(self, dag_maker, session
         job_runner = SchedulerJobRunner(job=scheduler_job, num_runs=0)
         job_runner._task_queued_timeout = 300
         with caplog.at_level(logging.DEBUG):
-            job_runner._fail_tasks_stuck_in_queued()
+            job_runner._handle_tasks_stuck_in_queued()
         assert "Executor doesn't support cleanup of stuck queued tasks. Skipping." in caplog.text
 
     @mock.patch("airflow.dag_processing.manager.DagFileProcessorAgent")