From 1c381fb917a5349cf2df6d64f66247a775d349ff Mon Sep 17 00:00:00 2001 From: Collin McNulty Date: Wed, 4 Jun 2025 13:16:13 -0500 Subject: [PATCH 1/7] Filter log events to after last running --- .../src/airflow/jobs/scheduler_job_runner.py | 16 ++- .../tests/unit/jobs/test_scheduler_job.py | 97 +++++++++++++++++++ 2 files changed, 112 insertions(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 12a71275e715e..64731d5504280 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -32,7 +32,7 @@ from itertools import groupby from typing import TYPE_CHECKING, Any, Callable -from sqlalchemy import and_, delete, exists, func, select, text, tuple_, update +from sqlalchemy import and_, delete, desc, exists, func, select, text, tuple_, update from sqlalchemy.exc import OperationalError from sqlalchemy.orm import joinedload, lazyload, load_only, make_transient, selectinload from sqlalchemy.sql import expression @@ -2019,6 +2019,19 @@ def _get_num_times_stuck_in_queued(self, ti: TaskInstance, session: Session = NE We can then use this information to determine whether to reschedule a task or fail it. """ + last_running_time = session.scalar( + select(Log.dttm) + .where( + 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 == "running", + ) + .orderby(desc(Log.dttm)) + .limit(1) + ) + return ( session.query(Log) .where( @@ -2028,6 +2041,7 @@ def _get_num_times_stuck_in_queued(self, ti: TaskInstance, session: Session = NE Log.map_index == ti.map_index, Log.try_number == ti.try_number, Log.event == TASK_STUCK_IN_QUEUED_RESCHEDULE_EVENT, + Log.dttm > last_running_time, ) .count() ) diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index ac36a12f27b90..4e3bd2f30497b 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -2083,6 +2083,103 @@ def _queue_tasks(tis): states = [x.state for x in dr.get_task_instances(session=session)] assert states == ["failed", "failed"] + @conf_vars({("scheduler", "num_stuck_in_queued_retries"): "2"}) + def test_handle_stuck_queued_tasks_reschedule_sensors(self, dag_maker, session, mock_executors): + """Reschedule sensors go in and out of running repeatedly using the same try_number + Make sure that they get three attempts per reschedule, not 3 attempts per try_number""" + with dag_maker("test_fail_stuck_queued_tasks_multiple_executors"): + EmptyOperator(task_id="op1") + EmptyOperator(task_id="op2", executor="default_exec") + + def _queue_tasks(tis): + for ti in tis: + ti.state = "queued" + ti.queued_dttm = timezone.utcnow() + session.commit() + + def _add_running_event(tis): + for ti in tis: + updated_entry = Log( + dttm=timezone.utcnow(), + dag_id=ti.dag_id, + task_id=ti.task_id, + map_index=ti.map_index, + event="running", + run_id=ti.run_id, + try_number=ti.try_number, + ) + session.add(updated_entry) + + run_id = str(uuid4()) + dr = dag_maker.create_dagrun(run_id=run_id) + + tis = dr.get_task_instances(session=session) + _queue_tasks(tis=tis) + scheduler_job = Job() + scheduler = SchedulerJobRunner(job=scheduler_job, num_runs=0) + # job_runner._reschedule_stuck_task = MagicMock() + scheduler._task_queued_timeout = -300 # always in violation of timeout + + with _loader_mock(mock_executors): + scheduler._handle_tasks_stuck_in_queued() + # If the task gets stuck in queued once, we reset it to scheduled + tis = dr.get_task_instances(session=session) + assert [x.state for x in tis] == ["scheduled", "scheduled"] + assert [x.queued_dttm for x in tis] == [None, None] + + _queue_tasks(tis=tis) + log_events = [ + x.event for x in session.scalars(select(Log).where(Log.run_id == run_id).order_by(Log.id)).all() + ] + assert log_events == [ + "stuck in queued reschedule", + "stuck in queued reschedule", + ] + + with _loader_mock(mock_executors): + scheduler._handle_tasks_stuck_in_queued() + + log_events = [ + x.event for x in session.scalars(select(Log).where(Log.run_id == run_id).order_by(Log.id)).all() + ] + assert log_events == [ + "stuck in queued reschedule", + "stuck in queued reschedule", + "stuck in queued reschedule", + "stuck in queued reschedule", + ] + mock_executors[0].fail.assert_not_called() + tis = dr.get_task_instances(session=session) + assert [x.state for x in tis] == ["scheduled", "scheduled"] + + _add_running_event(tis) # This should "reset" the count of stuck queued + _queue_tasks(tis=tis) + + for _ in range(3): # should need to get stuck 3 more times before finally failing + with _loader_mock(mock_executors): + scheduler._handle_tasks_stuck_in_queued() + log_events = [ + x.event for x in session.scalars(select(Log).where(Log.run_id == run_id).order_by(Log.id)).all() + ] + assert log_events == [ + "stuck in queued reschedule", + "stuck in queued reschedule", + "stuck in queued reschedule", + "stuck in queued reschedule", + "running", + "running", + "stuck in queued reschedule", + "stuck in queued reschedule", + "stuck in queued reschedule", + "stuck in queued reschedule", + "stuck in queued tries exceeded", + "stuck in queued tries exceeded", + ] + + mock_executors[0].fail.assert_not_called() # just demoing that we don't fail with executor method + states = [x.state for x in dr.get_task_instances(session=session)] + assert states == ["failed", "failed"] + def test_revoke_task_not_imp_tolerated(self, dag_maker, session, caplog): """Test that if executor no implement revoke_task then we don't blow up.""" with dag_maker("test_fail_stuck_queued_tasks"): From 3cc12f5db24ae3e49d0ce88f69dcc3eb8bdeec36 Mon Sep 17 00:00:00 2001 From: Collin McNulty Date: Wed, 4 Jun 2025 13:26:15 -0500 Subject: [PATCH 2/7] Remove dttm check if never in running --- .../src/airflow/jobs/scheduler_job_runner.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 64731d5504280..5d3ae37478b07 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -2032,6 +2032,20 @@ def _get_num_times_stuck_in_queued(self, ti: TaskInstance, session: Session = NE .limit(1) ) + if last_running_time: + 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, + Log.dttm > last_running_time, + ) + .count() + ) return ( session.query(Log) .where( @@ -2041,7 +2055,6 @@ def _get_num_times_stuck_in_queued(self, ti: TaskInstance, session: Session = NE Log.map_index == ti.map_index, Log.try_number == ti.try_number, Log.event == TASK_STUCK_IN_QUEUED_RESCHEDULE_EVENT, - Log.dttm > last_running_time, ) .count() ) From f993947ea36e367dba1bdae5277908e4c1fbaddb Mon Sep 17 00:00:00 2001 From: Collin McNulty Date: Wed, 4 Jun 2025 13:47:26 -0500 Subject: [PATCH 3/7] since running logic made DRY --- .../src/airflow/jobs/scheduler_job_runner.py | 18 +++++------------- 1 file changed, 5 insertions(+), 13 deletions(-) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 5d3ae37478b07..c2eb489c8df17 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -2033,19 +2033,10 @@ def _get_num_times_stuck_in_queued(self, ti: TaskInstance, session: Session = NE ) if last_running_time: - 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, - Log.dttm > last_running_time, - ) - .count() - ) + since_running = Log.dttm > last_running_time + else: + since_running = True + return ( session.query(Log) .where( @@ -2055,6 +2046,7 @@ def _get_num_times_stuck_in_queued(self, ti: TaskInstance, session: Session = NE Log.map_index == ti.map_index, Log.try_number == ti.try_number, Log.event == TASK_STUCK_IN_QUEUED_RESCHEDULE_EVENT, + since_running, ) .count() ) From 5dbcb65ffe3ede74d7e6d89772e486a2389953fb Mon Sep 17 00:00:00 2001 From: Collin McNulty Date: Wed, 4 Jun 2025 14:29:55 -0500 Subject: [PATCH 4/7] typo in order_by --- airflow-core/src/airflow/jobs/scheduler_job_runner.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index c2eb489c8df17..9c346b8ebcfac 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -2028,7 +2028,7 @@ def _get_num_times_stuck_in_queued(self, ti: TaskInstance, session: Session = NE Log.try_number == ti.try_number, Log.event == "running", ) - .orderby(desc(Log.dttm)) + .order_by(desc(Log.dttm)) .limit(1) ) From b85a0fb387a0af9bee7f82a0e05df27a4cdf99ac Mon Sep 17 00:00:00 2001 From: Collin McNulty Date: Wed, 4 Jun 2025 15:32:59 -0500 Subject: [PATCH 5/7] tests for reschedule sensors corrected --- airflow-core/tests/unit/jobs/test_scheduler_job.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 4e3bd2f30497b..123c1afefd98f 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -2153,11 +2153,13 @@ def _add_running_event(tis): assert [x.state for x in tis] == ["scheduled", "scheduled"] _add_running_event(tis) # This should "reset" the count of stuck queued - _queue_tasks(tis=tis) - for _ in range(3): # should need to get stuck 3 more times before finally failing + for _ in range(3): # Should be able to be stuck 3 more times before failing + _queue_tasks(tis=tis) with _loader_mock(mock_executors): scheduler._handle_tasks_stuck_in_queued() + tis = dr.get_task_instances(session=session) + log_events = [ x.event for x in session.scalars(select(Log).where(Log.run_id == run_id).order_by(Log.id)).all() ] From 73fd35de0f00d330a908ddd3ae9206c38b7d7fc9 Mon Sep 17 00:00:00 2001 From: Collin McNulty Date: Wed, 4 Jun 2025 19:09:04 -0500 Subject: [PATCH 6/7] Include task_id in log message check Co-authored-by: Jed Cunningham <66968678+jedcunningham@users.noreply.github.com> --- airflow-core/src/airflow/jobs/scheduler_job_runner.py | 1 + 1 file changed, 1 insertion(+) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 9c346b8ebcfac..f1077c54b0982 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -2023,6 +2023,7 @@ def _get_num_times_stuck_in_queued(self, ti: TaskInstance, session: Session = NE select(Log.dttm) .where( Log.dag_id == ti.dag_id, + Log.task_id == ti.task_id, Log.run_id == ti.run_id, Log.map_index == ti.map_index, Log.try_number == ti.try_number, From 88b8b9179ed90d36bce62ccb36e4f06c2c127cd0 Mon Sep 17 00:00:00 2001 From: Collin McNulty Date: Mon, 30 Jun 2025 14:24:24 -0500 Subject: [PATCH 7/7] simplify boolean logic for log query --- .../src/airflow/jobs/scheduler_job_runner.py | 27 ++++++++----------- 1 file changed, 11 insertions(+), 16 deletions(-) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index f1077c54b0982..2a5de48b488b5 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -2033,24 +2033,19 @@ def _get_num_times_stuck_in_queued(self, ti: TaskInstance, session: Session = NE .limit(1) ) + query = 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, + ) + if last_running_time: - since_running = Log.dttm > last_running_time - else: - since_running = True + query = query.where(Log.dttm > last_running_time) - 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, - since_running, - ) - .count() - ) + return query.count() previous_ti_running_metrics: dict[tuple[str, str, str], int] = {}