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(