From 999ca766a0618c54e08f9fbff092e713a6af5ef9 Mon Sep 17 00:00:00 2001 From: rjgoyln Date: Thu, 6 Aug 2026 01:45:03 +0800 Subject: [PATCH] Stop failing reschedule-mode sensors on a stale executor success A sensor running with mode="reschedule" was intermittently marked failed when the executor success for one poke reached the scheduler after the task instance had already been put back for the next poke. The existing guards for this race only cover the defer path: a reschedule exit leaves next_method unset and does not bump try_number, so none of them applied and the event was treated as an external kill. --- .../src/airflow/jobs/scheduler_job_runner.py | 38 ++++++++++++-- .../tests/unit/jobs/test_scheduler_job.py | 50 +++++++++++++++++++ 2 files changed, 83 insertions(+), 5 deletions(-) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index c6b555cf238eb..c7e127e9f8bd6 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -26,7 +26,7 @@ import sys import time from collections import Counter, defaultdict, deque -from collections.abc import Callable, Collection, Iterable, Iterator +from collections.abc import Callable, Collection, Iterable, Iterator, Sequence from contextlib import ExitStack from datetime import datetime, timedelta from functools import lru_cache, partial @@ -107,6 +107,7 @@ from airflow.models.serialized_dag import SerializedDagModel from airflow.models.taskinstance import TaskInstance from airflow.models.taskinstancekey import TaskInstanceKey +from airflow.models.taskreschedule import TaskReschedule from airflow.models.team import Team from airflow.models.trigger import TRIGGER_FAIL_REPR, Trigger, TriggerFailureReason, handle_event_submit from airflow.observability.metrics import stats_utils @@ -1319,6 +1320,23 @@ def _emit_executor_events_batch_metrics(num_events: int) -> None: stats.gauge("scheduler.executor_events.batch_size", num_events) stats.incr("scheduler.executor_events.processed", num_events) + @staticmethod + def _find_ti_ids_rescheduled_this_try(tis: Sequence[TI], *, session: Session) -> set[UUID]: + """Find which of the given task instances have already been rescheduled during their current try.""" + candidate_ids = [ + ti.id + for ti in tis + if ti.state in (TaskInstanceState.SCHEDULED, TaskInstanceState.QUEUED) and ti.next_method is None + ] + if not candidate_ids: + return set() + # Resolved once for the whole batch to keep the reschedule lookup off the per-event hot path. + return set( + session.scalars( + select(TaskReschedule.ti_id).where(TaskReschedule.ti_id.in_(candidate_ids)).distinct() + ) + ) + @classmethod def process_executor_events( cls, @@ -1452,7 +1470,8 @@ def process_executor_events( # row lock this entire set of taskinstances to make sure the scheduler doesn't fail when we have # multi-schedulers locked_query = with_row_locks(query, of=TI, session=session, skip_locked=True) - tis: Iterator[TI] = session.scalars(locked_query) + tis: Sequence[TI] = session.scalars(locked_query).all() + ti_ids_rescheduled_this_try = cls._find_ti_ids_rescheduled_this_try(tis, session=session) for ti in tis: try_number = ti_primary_key_to_try_number_map[ti.key.primary] buffer_key = ti.key.with_try_number(try_number) @@ -1516,6 +1535,8 @@ def process_executor_events( # from the worker exit after defer() has not been processed yet - should not fail it. # 4) the trigger already put the TI back to queued (resume after defer) but the executor success # from the worker exit after defer() has not been processed yet - should not fail it. + # 5) a sensor in reschedule mode was put back to scheduled or queued for its next poke but the + # executor success from the previous poke has not been processed yet - should not fail it. # All of this could also happen if the state is "running", # but that is handled by the scheduler detecting task instances without heartbeats. @@ -1530,11 +1551,18 @@ def process_executor_events( ti.queued_by_job_id != job_id # Another scheduler has queued this task again or executor.has_task(ti) # This scheduler has this task already or ( - # Resume-after-defer: trigger moved TI to scheduled or queued (next_method set) - # before we saw the executor success from the defer exit for the same try_number. ti.state in (TaskInstanceState.SCHEDULED, TaskInstanceState.QUEUED) and state == TaskInstanceState.SUCCESS - and ti.next_method is not None + and ( + # Resume-after-defer: trigger moved TI to scheduled or queued (next_method set) + # before we saw the executor success from the defer exit for the same try_number. + ti.next_method is not None + # Sensor in reschedule mode: the TI was put back for its next poke before we saw + # the executor success from the previous poke. A reschedule exit leaves + # next_method unset and keeps the same try_number, so an earlier reschedule of + # this try is the only thing that tells this apart from an externally killed task. + or ti.id in ti_ids_rescheduled_this_try + ) ) ) diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index d344df97c200c..4443ebf8705c9 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -95,6 +95,7 @@ from airflow.models.pool import Pool from airflow.models.serialized_dag import SerializedDagModel from airflow.models.taskinstance import TaskInstance +from airflow.models.taskreschedule import TaskReschedule from airflow.models.team import Team from airflow.models.trigger import Trigger from airflow.partition_mappers.base import ( @@ -1006,6 +1007,55 @@ def test_process_executor_events_stale_success_when_queued_after_defer( tags={"dag_id": dag_id, "task_id": ti1.task_id}, ) + @pytest.mark.parametrize("ti_state", [State.SCHEDULED, State.QUEUED]) + @mock.patch("airflow._shared.observability.metrics.stats._get_backend") + def test_process_executor_events_stale_success_between_reschedule_pokes( + self, mock_get_backend, ti_state, dag_maker + ): + """Stale success from an earlier poke of a reschedule-mode sensor must not read as a state mismatch.""" + mock_stats = mock.MagicMock(spec=StatsLogger) + mock_get_backend.return_value = mock_stats + dag_id = "test_process_executor_events_stale_success_between_reschedule_pokes" + + session = settings.Session() + with dag_maker(dag_id=dag_id, fileloc="/test_path1/"): + task1 = EmptyOperator(task_id="dummy_task") + ti1 = dag_maker.create_dagrun().get_task_instance(task1.task_id) + + executor = MockExecutor(do_update=False) + scheduler_job = Job() + session.add(scheduler_job) + session.flush() + self.job_runner = SchedulerJobRunner(scheduler_job, executors=[executor]) + + # A reschedule exit leaves next_method unset and keeps the same try_number, so the reschedule + # row is the only thing standing between this and the externally-killed branch. + ti1.state = ti_state + ti1.next_method = None + ti1.queued_by_job_id = scheduler_job.id + ti1.try_number = 1 + session.merge(ti1) + poke_end = DEFAULT_DATE + datetime.timedelta(seconds=5) + session.add( + TaskReschedule( + ti_id=ti1.id, + start_date=DEFAULT_DATE, + end_date=poke_end, + reschedule_date=poke_end + datetime.timedelta(seconds=60), + ) + ) + session.commit() + + executor.event_buffer[ti1.key] = State.SUCCESS, None + executor.has_task = mock.MagicMock(spec=executor.has_task, return_value=False) + mock_stats.incr.reset_mock() + + self.job_runner._process_executor_events(executor=executor, session=session) + ti1.refresh_from_db(session=session) + assert ti1.state == ti_state + self.job_runner.executor.callback_sink.send.assert_not_called() + mock_stats.incr.assert_called_once_with("scheduler.executor_events.processed", count=1) + @mock.patch("airflow.jobs.scheduler_job_runner.TaskCallbackRequest") @mock.patch("airflow._shared.observability.metrics.stats._get_backend") def test_process_executor_events_multiple_try_numbers_warns(