This is an automated email from the ASF dual-hosted git repository.
hussein-awala pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new 0fb5bfc4df1 Fill scheduler task batches after concurrency rejection
(#70307)
0fb5bfc4df1 is described below
commit 0fb5bfc4df1da8e4f2ae1e436321a8d228630ed3
Author: Arthur <[email protected]>
AuthorDate: Thu Oct 22 11:12:12 2026 +0200
Fill scheduler task batches after concurrency rejection (#70307)
* Fill scheduler candidate batches after concurrency rejection
A candidate query could stop after selecting any executable task, leaving
remaining capacity unused when higher-priority candidates were rejected by
in-memory concurrency checks. This could repeatedly starve lower-priority tasks.
* Clarify scheduler candidate retry logging
* Avoid one-row scheduler refill queries
When a candidate batch is nearly full, reducing the query limit to the
remaining capacity can require one database round trip for every distinct
blocked task.
* Bound scheduler refill queries inside the critical section
* Exclude only the saturated Dag run from scheduler refill queries
* Test that refill cap ignores queries before the first selection
---
.../src/airflow/jobs/scheduler_job_runner.py | 90 +++++++----
airflow-core/tests/unit/jobs/test_scheduler_job.py | 170 +++++++++++++++++++++
.../observability/metrics/metrics_template.yaml | 8 +
3 files changed, 240 insertions(+), 28 deletions(-)
diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
index 7449c09e348..5a56f7bfc33 100644
--- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py
+++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py
@@ -174,6 +174,17 @@ TASK_STUCK_IN_QUEUED_RESCHEDULE_EVENT = "stuck in queued
reschedule"
# safety bound, not a behavioural knob operators need to tune.
MAX_PARTITION_DAG_RUNS_PER_LOOP = 500
+MAX_TI_REFILL_QUERIES_PER_LOOP = 3
+"""
+Max extra task instance queries per critical section to fill a partially
selected batch.
+
+Each query runs while holding the pool row locks (and the advisory lock on
PostgreSQL), which other
+HA schedulers wait on. Queries issued before any task instance is selected are
not capped, so a loop
+still selects at least one runnable task instance queued behind blocked ones.
+
+:meta private:
+"""
+
def _eager_load_dag_run_for_validation() -> tuple[LoaderOption, LoaderOption]:
"""
@@ -734,28 +745,38 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
concurrency_map = ConcurrencyMap()
concurrency_map.load(session=session)
+ executor_slots_available: dict[ExecutorName, int] = {}
+ for executor in self.executors:
+ if TYPE_CHECKING:
+ # All executors should have a name if they are initted from
the executor_loader.
+ # But we need to check for None to make mypy happy.
+ assert executor.name
+ executor_slots_available[executor.name] = executor.slots_available
+
# Number of tasks that cannot be scheduled because of no open slot in
pool
num_starving_tasks_total = 0
- # dag and task ids that can't be queued because of concurrency limits
- starved_dags: set[str] = set()
+ # Dag runs and tasks that can't be queued because of concurrency limits
+ starved_dag_runs: set[tuple[str, str]] = set()
starved_tasks: set[tuple[str, str]] = set()
starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]] =
set()
pool_num_starving_tasks: dict[str, int] = Counter()
+ num_refill_queries = 0
for loop_count in itertools.count(start=1):
num_starved_pools = len(starved_pools)
- num_starved_dags = len(starved_dags)
+ num_starved_dag_runs = len(starved_dag_runs)
num_starved_tasks = len(starved_tasks)
num_starved_tasks_task_dagrun_concurrency =
len(starved_tasks_task_dagrun_concurrency)
query = self._build_schedulable_tis_query(
starved_pools,
- starved_dags,
+ starved_dag_runs,
starved_tasks,
starved_tasks_task_dagrun_concurrency,
max_tis,
+ excluded_ti_ids=[ti.id for ti in executable_tis],
)
timer = stats.timer("scheduler.critical_section_query_duration")
@@ -805,15 +826,6 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
if team := dag_id_to_team_name.get(ti.dag_id):
ti.dag_run._team_name = team
- executor_slots_available: dict[ExecutorName, int] = {}
- # First get a mapping of executor names to slots they have
available
- for executor in self.executors:
- if TYPE_CHECKING:
- # All executors should have a name if they are initted
from the executor_loader.
- # But we need to check for None to make mypy happy.
- assert executor.name
- executor_slots_available[executor.name] =
executor.slots_available
-
for task_instance in task_instances_to_examine:
pool_name = task_instance.pool
@@ -890,20 +902,20 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
current_active_tasks_per_dag_run =
concurrency_map.dag_run_active_tasks_map[dag_run_key]
dag_max_active_tasks = task_instance.dag_model.max_active_tasks
self.log.info(
- "DAG %s has %s/%s running and queued tasks",
+ "Dag run %s of Dag %s has %s/%s running and queued tasks",
+ task_instance.run_id,
dag_id,
current_active_tasks_per_dag_run,
dag_max_active_tasks,
)
if current_active_tasks_per_dag_run >= dag_max_active_tasks:
self.log.info(
- "Not executing %s since the number of tasks running or
queued "
- "from DAG %s is >= to the DAG's max_active_tasks limit
of %s",
+ "Not executing %s since its Dag run has reached "
+ "the Dag's max_active_tasks limit of %s",
task_instance,
- dag_id,
dag_max_active_tasks,
)
- starved_dags.add(dag_id)
+ starved_dag_runs.add(dag_run_key)
continue
# Many DAGs do not define task concurrency limits, so avoid
@@ -956,21 +968,37 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
pool_stats["open"] = open_slots
- is_done = executable_tis or len(task_instances_to_examine) <
max_tis
+ if len(executable_tis) >= max_tis:
+ break
+
# Check this to avoid accidental infinite loops
found_new_filters = (
len(starved_pools) > num_starved_pools
- or len(starved_dags) > num_starved_dags
+ or len(starved_dag_runs) > num_starved_dag_runs
or len(starved_tasks) > num_starved_tasks
or len(starved_tasks_task_dagrun_concurrency) >
num_starved_tasks_task_dagrun_concurrency
)
- if is_done or not found_new_filters:
+ if len(executable_tis) >= max_tis or not found_new_filters:
break
+ if executable_tis:
+ if num_refill_queries >= MAX_TI_REFILL_QUERIES_PER_LOOP:
+ self.log.debug(
+ "Selected %s of %s task instances; reached the limit
of %s refill queries.",
+ len(executable_tis),
+ max_tis,
+ MAX_TI_REFILL_QUERIES_PER_LOOP,
+ )
+
stats.incr("scheduler.critical_section_refill_limit_reached")
+ break
+ num_refill_queries += 1
+
self.log.info(
- "Found no task instances to queue on query iteration %s "
- "but there could be more candidate task instances to check.",
+ "Selected %s of %s task instances after iteration %s; "
+ "retrying after excluding known blocked candidates.",
+ len(executable_tis),
+ max_tis,
loop_count,
)
@@ -993,15 +1021,17 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
def _build_schedulable_tis_query(
self,
starved_pools: set[str],
- starved_dags: set[str],
+ starved_dag_runs: set[tuple[str, str]],
starved_tasks: set[tuple[str, str]],
starved_tasks_task_dagrun_concurrency: set[tuple[str, str, str]],
max_tis: int,
+ *,
+ excluded_ti_ids: Collection[UUID] = (),
) -> Select[tuple[TI]]:
"""
Build a query that fetches SCHEDULED TIs eligible for execution this
cycle.
- Applies current starvation exclusions so that saturated pools, DAGs,
or tasks
+ Applies current starvation exclusions so that saturated pools, Dag
runs, or tasks
don't re-appear in the candidate set. Row-number windowing enforces
``max_active_tasks`` per DagRun. The returned query is ready to be
wrapped
with ``with_row_locks`` and executed by the caller; no session is
required here.
@@ -1035,13 +1065,13 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
)
# Starvation filters should be applied before computing the row_num
based on the
- # max_active_tasks limit. That way, starved dags and tasks that
shouldn't run,
+ # max_active_tasks limit. That way, starved Dag runs and tasks that
shouldn't run,
# won't occupy a slot.
if starved_pools:
query = query.where(TI.pool.not_in(starved_pools))
- if starved_dags:
- query = query.where(TI.dag_id.not_in(starved_dags))
+ if starved_dag_runs:
+ query = query.where(tuple_(TI.dag_id,
TI.run_id).not_in(starved_dag_runs))
if starved_tasks:
query = query.where(tuple_(TI.dag_id,
TI.task_id).not_in(starved_tasks))
@@ -1051,6 +1081,10 @@ class SchedulerJobRunner(BaseJobRunner, LoggingMixin):
tuple_(TI.dag_id, TI.run_id,
TI.task_id).not_in(starved_tasks_task_dagrun_concurrency)
)
+ if excluded_ti_ids:
+ # Selected TIs stay SCHEDULED in the database until the selection
loop marks them QUEUED.
+ query = query.where(TI.id.not_in(excluded_ti_ids))
+
# Create a subquery with row numbers partitioned by dag_id and run_id.
# Different dags can have the same run_id but
# the dag_id combined with the run_id uniquely identify a run.
diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py
b/airflow-core/tests/unit/jobs/test_scheduler_job.py
index a90b8bd00fa..84bb06cc768 100644
--- a/airflow-core/tests/unit/jobs/test_scheduler_job.py
+++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py
@@ -69,6 +69,7 @@ from airflow.executors.workloads import WorkloadType
from airflow.executors.workloads.types import TaskInstanceUuid
from airflow.jobs.job import Job, run_job
from airflow.jobs.scheduler_job_runner import (
+ MAX_TI_REFILL_QUERIES_PER_LOOP,
TASK_STUCK_IN_QUEUED_RESCHEDULE_EVENT,
SchedulerJobRunner,
)
@@ -3537,6 +3538,175 @@ class TestSchedulerJob:
session.rollback()
+ @pytest.fixture
+ def interleaved_blocked_and_runnable_tis(self, dag_maker, session):
+ """
+ Create SCHEDULED TIs ordered so each 3-row candidate page holds one
runnable TI until the last.
+
+ Candidate order: blocked_0, runnable_0, blocked_1 | blocked_2,
runnable_1, blocked_3 |
+ blocked_4, runnable_2, runnable_3. Blocked TIs are at their
``max_active_tis_per_dag`` limit.
+ """
+ with dag_maker(dag_id="interleaved_blocked_and_runnable",
session=session):
+ blocked_tasks = [
+ EmptyOperator(task_id=f"blocked_{index}",
priority_weight=weight, max_active_tis_per_dag=1)
+ for index, weight in enumerate((9, 7, 6, 4, 3))
+ ]
+ runnable_tasks = [
+ EmptyOperator(task_id=f"runnable_{index}",
priority_weight=weight)
+ for index, weight in enumerate((8, 5, 2, 1))
+ ]
+ first_run = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED)
+ second_run = dag_maker.create_dagrun_after(first_run,
run_type=DagRunType.SCHEDULED)
+ for blocked_task in blocked_tasks:
+ first_run.get_task_instance(blocked_task.task_id,
session=session).state = State.RUNNING
+ second_run.get_task_instance(blocked_task.task_id,
session=session).state = State.SCHEDULED
+ for runnable_task in runnable_tasks:
+ first_run.get_task_instance(runnable_task.task_id,
session=session).state = State.SCHEDULED
+ session.flush()
+
+ @pytest.mark.usefixtures("interleaved_blocked_and_runnable_tis")
+ @pytest.mark.parametrize(
+ ("executor_slots", "expected_task_ids"),
+ [
+ pytest.param(8, ["runnable_0", "runnable_1", "runnable_2"],
id="fills_batch_up_to_max_tis"),
+ pytest.param(2, ["runnable_0", "runnable_1"],
id="executor_slots_shared_across_queries"),
+ ],
+ )
+ def
test_select_task_instances_to_queue_refills_batch_after_concurrency_rejection(
+ self, mock_executors, session, executor_slots, expected_task_ids
+ ):
+ self.job_runner = SchedulerJobRunner(job=Job())
+ self.job_runner.executor.slots_available = executor_slots
+
+ queued_tis = self.job_runner._select_task_instances_to_queue(
+ 3, make_pool_stats(), set(), session=session
+ )
+
+ assert [ti.task_id for ti in queued_tis] == expected_task_ids
+
+ @pytest.mark.usefixtures("interleaved_blocked_and_runnable_tis")
+
@mock.patch("airflow.jobs.scheduler_job_runner.MAX_TI_REFILL_QUERIES_PER_LOOP",
1)
+ @mock.patch("airflow._shared.observability.metrics.stats._get_backend")
+ def test_select_task_instances_to_queue_stops_at_refill_query_limit(self,
mock_get_backend, session):
+ mock_stats = mock.MagicMock(spec=StatsLogger)
+ mock_get_backend.return_value = mock_stats
+ self.job_runner = SchedulerJobRunner(job=Job())
+
+ queued_tis = self.job_runner._select_task_instances_to_queue(
+ 3, make_pool_stats(), set(), session=session
+ )
+
+ assert [ti.task_id for ti in queued_tis] == ["runnable_0",
"runnable_1"]
+
mock_stats.incr.assert_any_call("scheduler.critical_section_refill_limit_reached")
+
+ def
test_select_task_instances_to_queue_does_not_cap_queries_before_first_selection(
+ self, dag_maker, session
+ ):
+ """More blocked candidates than the refill cap must not hide a
lower-priority runnable TI."""
+ self.job_runner = SchedulerJobRunner(job=Job())
+ with dag_maker(dag_id="blocked_tis_exceed_refill_cap",
session=session):
+ blocked_tasks = [
+ EmptyOperator(task_id=f"blocked_{index}", priority_weight=10,
max_active_tis_per_dag=1)
+ for index in range(MAX_TI_REFILL_QUERIES_PER_LOOP + 2)
+ ]
+ runnable_task = EmptyOperator(task_id="runnable",
priority_weight=1)
+ first_run = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED)
+ second_run = dag_maker.create_dagrun_after(first_run,
run_type=DagRunType.SCHEDULED)
+ for blocked_task in blocked_tasks:
+ first_run.get_task_instance(blocked_task.task_id,
session=session).state = State.RUNNING
+ second_run.get_task_instance(blocked_task.task_id,
session=session).state = State.SCHEDULED
+ runnable_ti = second_run.get_task_instance(runnable_task.task_id,
session=session)
+ runnable_ti.state = State.SCHEDULED
+ session.flush()
+
+ queued_tis = self.job_runner._select_task_instances_to_queue(
+ 1, make_pool_stats(), set(), session=session
+ )
+
+ assert [ti.key for ti in queued_tis] == [runnable_ti.key]
+
+ def
test_select_task_instances_to_queue_finds_runnable_hidden_by_max_active_tasks_window(
+ self, dag_maker, session
+ ):
+ """Blocked TIs filling a run's ``max_active_tasks`` window must not
hide its runnable TIs."""
+ self.job_runner = SchedulerJobRunner(job=Job())
+ with dag_maker(
+ dag_id="blocked_tis_fill_max_active_tasks_window",
max_active_tasks=2, session=session
+ ):
+ blocked_tasks = [
+ EmptyOperator(task_id=f"blocked_{index}",
priority_weight=weight, max_active_tis_per_dag=1)
+ for index, weight in enumerate((3, 2))
+ ]
+ runnable_task = EmptyOperator(task_id="runnable",
priority_weight=1)
+ first_run = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED)
+ second_run = dag_maker.create_dagrun_after(first_run,
run_type=DagRunType.SCHEDULED)
+ for blocked_task in blocked_tasks:
+ first_run.get_task_instance(blocked_task.task_id,
session=session).state = State.RUNNING
+ second_run.get_task_instance(blocked_task.task_id,
session=session).state = State.SCHEDULED
+ runnable_ti = second_run.get_task_instance(runnable_task.task_id,
session=session)
+ runnable_ti.state = State.SCHEDULED
+ session.flush()
+
+ queued_tis = self.job_runner._select_task_instances_to_queue(
+ 32, make_pool_stats(), set(), session=session
+ )
+
+ assert [ti.key for ti in queued_tis] == [runnable_ti.key]
+
+ def test_select_task_instances_to_queue_preserves_priority_across_dag_runs(
+ self, dag_maker, mock_executors, session
+ ):
+ with dag_maker(dag_id="multiple_runs", max_active_tasks=2,
session=session):
+ EmptyOperator(task_id="running")
+ EmptyOperator(task_id="fills_run", priority_weight=100)
+ EmptyOperator(task_id="blocked_in_run", priority_weight=99)
+ EmptyOperator(task_id="other_run", priority_weight=98)
+ first_run = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED)
+ second_run = dag_maker.create_dagrun_after(first_run,
run_type=DagRunType.SCHEDULED)
+
+ with dag_maker(dag_id="lower_priority", session=session):
+ EmptyOperator(task_id="competitor", priority_weight=1)
+ competing_run = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED)
+
+ first_run.get_task_instance("running", session=session).state =
State.RUNNING
+ fills_run_ti = first_run.get_task_instance("fills_run",
session=session)
+ fills_run_ti.state = State.SCHEDULED
+ first_run.get_task_instance("blocked_in_run", session=session).state =
State.SCHEDULED
+ other_run_ti = second_run.get_task_instance("other_run",
session=session)
+ other_run_ti.state = State.SCHEDULED
+ competing_run.get_task_instance("competitor", session=session).state =
State.SCHEDULED
+ session.flush()
+
+ expected_keys = [fills_run_ti.key, other_run_ti.key]
+ pools = make_pool_stats(total=3, running=1)
+ self.job_runner = SchedulerJobRunner(job=Job())
+ queued_tis = []
+ # Allow another pass so this checks priority, not whether one pass
fills the batch.
+ for _ in range(2):
+ max_tis = int(pools["default_pool"]["open"])
+ if not max_tis:
+ break
+ queued_tis.extend(
+ self.job_runner._select_task_instances_to_queue(max_tis,
pools, set(), session=session)
+ )
+
+ assert [ti.key for ti in queued_tis] == expected_keys
+
+ def test_build_schedulable_tis_query_excludes_only_starved_dag_run(self,
dag_maker, session):
+ with dag_maker(dag_id="starved_run", session=session):
+ EmptyOperator(task_id="task")
+ first_run = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED)
+ second_run = dag_maker.create_dagrun_after(first_run,
run_type=DagRunType.SCHEDULED)
+ for dag_run in (first_run, second_run):
+ dag_run.get_task_instance("task", session=session).state =
State.SCHEDULED
+ session.flush()
+
+ query = SchedulerJobRunner(job=Job())._build_schedulable_tis_query(
+ set(), {("starved_run", first_run.run_id)}, set(), set(), 32
+ )
+
+ assert [ti.run_id for ti in session.scalars(query)] ==
[second_run.run_id]
+
def
test_find_executable_task_instances_task_concurrency_per_dagrun_for_first(self,
dag_maker):
scheduler_job = Job()
self.job_runner = SchedulerJobRunner(job=scheduler_job)
diff --git
a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
index 45ce30f648c..0f9a796f480 100644
---
a/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
+++
b/shared/observability/src/airflow_shared/observability/metrics/metrics_template.yaml
@@ -207,6 +207,14 @@ metrics:
legacy_name: "-"
name_variables: []
+ - name: "scheduler.critical_section_refill_limit_reached"
+ description: "Count of times the scheduler stopped querying for more task
instances to fill
+ its batch because it reached the per-loop limit of refill queries. The
remaining runnable
+ task instances are considered in the next scheduler loop."
+ type: "counter"
+ legacy_name: "-"
+ name_variables: []
+
- name: "ti.start"
description: "Number of started task in a given Dag. Similar to
{job_name}_start but for task.
Metric with dag_id and task_id tagging."
