You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@airflow.apache.org by ra...@apache.org on 2025/07/01 10:26:13 UTC
(airflow) 01/01: Correctly treat requeues on reschedule sensors as resetting after each reschedule (#51410)
This is an automated email from the ASF dual-hosted git repository.
rahulvats pushed a commit to branch backport-51410
in repository https://gitbox.apache.org/repos/asf/airflow.git
commit a36042b03354f19ada97556a3aad75709644427a
Author: Collin McNulty <co...@gmail.com>
AuthorDate: Tue Jul 1 00:55:15 2025 -0500
Correctly treat requeues on reschedule sensors as resetting after each reschedule (#51410)
(cherry picked from commit a3621016840d8dba97777954333f6b214d4f3363)
---
airflow/jobs/scheduler_job_runner.py | 27 +++++++---
tests/jobs/test_scheduler_job.py | 99 ++++++++++++++++++++++++++++++++++++
2 files changed, 120 insertions(+), 6 deletions(-)
diff --git a/airflow/jobs/scheduler_job_runner.py b/airflow/jobs/scheduler_job_runner.py
index c9afd40f719..6e6a54a1758 100644
--- a/airflow/jobs/scheduler_job_runner.py
+++ b/airflow/jobs/scheduler_job_runner.py
@@ -33,7 +33,7 @@ from pathlib import Path
from typing import TYPE_CHECKING, Any, Callable, Collection, Iterable, Iterator
from deprecated import deprecated
-from sqlalchemy import and_, delete, func, not_, or_, select, text, update
+from sqlalchemy import and_, delete, desc, func, not_, or_, select, text, update
from sqlalchemy.exc import OperationalError
from sqlalchemy.orm import lazyload, load_only, make_transient, selectinload
from sqlalchemy.sql import expression
@@ -1899,19 +1899,34 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
We can then use this information to determine whether to reschedule a task or fail it.
"""
- return (
- session.query(Log)
+ last_running_time = session.scalar(
+ select(Log.dttm)
.where(
- Log.task_id == ti.task_id,
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,
- Log.event == TASK_STUCK_IN_QUEUED_RESCHEDULE_EVENT,
+ Log.event == "running",
)
- .count()
+ .order_by(desc(Log.dttm))
+ .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:
+ query = query.where(Log.dttm > last_running_time)
+
+ return query.count()
+
@provide_session
def _emit_pool_metrics(self, session: Session = NEW_SESSION) -> None:
from airflow.models.pool import Pool
diff --git a/tests/jobs/test_scheduler_job.py b/tests/jobs/test_scheduler_job.py
index 97d3e9fe87d..f8a879145df 100644
--- a/tests/jobs/test_scheduler_job.py
+++ b/tests/jobs/test_scheduler_job.py
@@ -2314,6 +2314,105 @@ class TestSchedulerJob:
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
+
+ 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()
+ ]
+ 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"):