From 1330e1832ded0252698a1ad29b286cdfe3fb3019 Mon Sep 17 00:00:00 2001 From: avolant Date: Thu, 23 Jul 2026 18:04:26 +0200 Subject: [PATCH 1/6] 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. --- .../src/airflow/jobs/scheduler_job_runner.py | 35 ++++++++++------- .../tests/unit/jobs/test_scheduler_job.py | 38 +++++++++++++++++++ 2 files changed, 59 insertions(+), 14 deletions(-) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 7729ca789e5f8..b530a5b68d685 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -754,6 +754,14 @@ def _select_task_instances_to_queue( 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 @@ -775,7 +783,8 @@ def _select_task_instances_to_queue( starved_dags, starved_tasks, starved_tasks_task_dagrun_concurrency, - max_tis, + max_tis - len(executable_tis), + excluded_ti_ids=[ti.id for ti in executable_tis], ) timer = stats.timer("scheduler.critical_section_query_duration") @@ -825,15 +834,6 @@ def _select_task_instances_to_queue( 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 @@ -976,7 +976,6 @@ def _select_task_instances_to_queue( pool_stats["open"] = open_slots - is_done = executable_tis or len(task_instances_to_examine) < max_tis # Check this to avoid accidental infinite loops found_new_filters = ( len(starved_pools) > num_starved_pools @@ -985,12 +984,14 @@ def _select_task_instances_to_queue( 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 self.log.info( - "Found no task instances to queue on query iteration %s " - "but there could be more candidate task instances to check.", + "Found %s of %s task instances for the batch after query iteration %s " + "and checking for more candidates.", + len(executable_tis), + max_tis, loop_count, ) @@ -1017,6 +1018,8 @@ def _build_schedulable_tis_query( 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. @@ -1071,6 +1074,10 @@ def _build_schedulable_tis_query( tuple_(TI.dag_id, TI.run_id, TI.task_id).not_in(starved_tasks_task_dagrun_concurrency) ) + if excluded_ti_ids: + # TIs already selected in this critical section stay SCHEDULED until it ends. + 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 8b575a8369abd..d792a4ce5d160 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -3537,6 +3537,44 @@ def test_find_executable_task_instances_not_enough_task_concurrency_for_first(se session.rollback() + def test_find_executable_task_instances_fills_batch_after_task_concurrency_rejection( + self, dag_maker, session + ): + job_runner = SchedulerJobRunner(job=Job()) + dag_id = ( + "SchedulerJobTest" + ".test_find_executable_task_instances_fills_batch_after_task_concurrency_rejection" + ) + with dag_maker(dag_id=dag_id, session=session): + blocked = EmptyOperator( + task_id="blocked", + priority_weight=3, + max_active_tis_per_dag=1, + ) + first_runnable = EmptyOperator(task_id="first_runnable", priority_weight=2) + second_runnable = EmptyOperator(task_id="second_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) + + running_ti = first_run.get_task_instance(blocked.task_id, session=session) + first_runnable_ti = first_run.get_task_instance(first_runnable.task_id, session=session) + second_runnable_ti = first_run.get_task_instance(second_runnable.task_id, session=session) + blocked_ti = second_run.get_task_instance(blocked.task_id, session=session) + + running_ti.state = State.RUNNING + first_runnable_ti.state = State.SCHEDULED + second_runnable_ti.state = State.SCHEDULED + blocked_ti.state = State.SCHEDULED + session.flush() + + queued_tis = job_runner._select_task_instances_to_queue(2, make_pool_stats(), set(), session=session) + + assert {ti.key for ti in queued_tis} == { + first_runnable_ti.key, + second_runnable_ti.key, + } + 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) From 4c7adf0996d82367d64cb80b685da5a4f7e57a55 Mon Sep 17 00:00:00 2001 From: avolant Date: Thu, 23 Jul 2026 18:16:45 +0200 Subject: [PATCH 2/6] Clarify scheduler candidate retry logging --- airflow-core/src/airflow/jobs/scheduler_job_runner.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index b530a5b68d685..5abb9e720714b 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -988,8 +988,8 @@ def _select_task_instances_to_queue( break self.log.info( - "Found %s of %s task instances for the batch after query iteration %s " - "and checking for more candidates.", + "Selected %s of %s task instances after iteration %s; " + "retrying after excluding known blocked candidates.", len(executable_tis), max_tis, loop_count, From 414be5da74dc1b08251e9699f6274bd57bfb6099 Mon Sep 17 00:00:00 2001 From: avolant Date: Mon, 27 Jul 2026 16:39:59 +0200 Subject: [PATCH 3/6] 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. --- .../src/airflow/jobs/scheduler_job_runner.py | 5 +- .../tests/unit/jobs/test_scheduler_job.py | 57 ++++++++++--------- 2 files changed, 33 insertions(+), 29 deletions(-) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 5abb9e720714b..6d9f87ecff864 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -783,7 +783,7 @@ def _select_task_instances_to_queue( starved_dags, starved_tasks, starved_tasks_task_dagrun_concurrency, - max_tis - len(executable_tis), + max_tis, excluded_ti_ids=[ti.id for ti in executable_tis], ) @@ -976,6 +976,9 @@ def _select_task_instances_to_queue( pool_stats["open"] = open_slots + if len(executable_tis) >= max_tis: + break + # Check this to avoid accidental infinite loops found_new_filters = ( len(starved_pools) > num_starved_pools diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index d792a4ce5d160..3e26ca23f2c7e 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -3537,43 +3537,44 @@ def test_find_executable_task_instances_not_enough_task_concurrency_for_first(se session.rollback() - def test_find_executable_task_instances_fills_batch_after_task_concurrency_rejection( - self, dag_maker, session - ): + def test_find_executable_task_instances_uses_full_sized_refill_queries(self, dag_maker, session): job_runner = SchedulerJobRunner(job=Job()) - dag_id = ( - "SchedulerJobTest" - ".test_find_executable_task_instances_fills_batch_after_task_concurrency_rejection" - ) + dag_id = "SchedulerJobTest.test_find_executable_task_instances_uses_full_sized_refill_queries" with dag_maker(dag_id=dag_id, session=session): - blocked = EmptyOperator( - task_id="blocked", - priority_weight=3, - max_active_tis_per_dag=1, - ) - first_runnable = EmptyOperator(task_id="first_runnable", priority_weight=2) - second_runnable = EmptyOperator(task_id="second_runnable", priority_weight=1) + blocked_tasks = [ + EmptyOperator( + task_id=f"blocked_{index}", + priority_weight=priority_weight, + max_active_tis_per_dag=1, + ) + for index, priority_weight in enumerate((6, 3, 2)) + ] + runnable_tasks = [ + EmptyOperator(task_id=f"runnable_{index}", priority_weight=priority_weight) + for index, priority_weight in enumerate((5, 4, 1)) + ] first_run = dag_maker.create_dagrun(run_type=DagRunType.SCHEDULED) second_run = dag_maker.create_dagrun_after(first_run, run_type=DagRunType.SCHEDULED) - running_ti = first_run.get_task_instance(blocked.task_id, session=session) - first_runnable_ti = first_run.get_task_instance(first_runnable.task_id, session=session) - second_runnable_ti = first_run.get_task_instance(second_runnable.task_id, session=session) - blocked_ti = second_run.get_task_instance(blocked.task_id, session=session) - - running_ti.state = State.RUNNING - first_runnable_ti.state = State.SCHEDULED - second_runnable_ti.state = State.SCHEDULED - blocked_ti.state = State.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_tis = [ + first_run.get_task_instance(runnable_task.task_id, session=session) + for runnable_task in runnable_tasks + ] + for runnable_ti in runnable_tis: + runnable_ti.state = State.SCHEDULED session.flush() - queued_tis = job_runner._select_task_instances_to_queue(2, make_pool_stats(), set(), session=session) + with mock.patch.object(session, "scalars", wraps=session.scalars) as scalars: + queued_tis = job_runner._select_task_instances_to_queue( + 3, make_pool_stats(), set(), session=session + ) - assert {ti.key for ti in queued_tis} == { - first_runnable_ti.key, - second_runnable_ti.key, - } + assert {ti.key for ti in queued_tis} == {ti.key for ti in runnable_tis} + assert scalars.call_count == 2 def test_find_executable_task_instances_task_concurrency_per_dagrun_for_first(self, dag_maker): scheduler_job = Job() From 6b72e2ab7cdd8a87fc92ad34327945d147608532 Mon Sep 17 00:00:00 2001 From: Arthur Date: Mon, 5 Oct 2026 17:37:11 +0200 Subject: [PATCH 4/6] Bound scheduler refill queries inside the critical section --- .../src/airflow/jobs/scheduler_job_runner.py | 26 ++++- .../tests/unit/jobs/test_scheduler_job.py | 100 +++++++++++++----- .../metrics/metrics_template.yaml | 8 ++ 3 files changed, 108 insertions(+), 26 deletions(-) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 6d9f87ecff864..7eb4e4a9e83de 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 @@ # 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]: """ @@ -771,6 +782,7 @@ def _select_task_instances_to_queue( 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) @@ -990,6 +1002,18 @@ def _select_task_instances_to_queue( 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( "Selected %s of %s task instances after iteration %s; " "retrying after excluding known blocked candidates.", @@ -1078,7 +1102,7 @@ def _build_schedulable_tis_query( ) if excluded_ti_ids: - # TIs already selected in this critical section stay SCHEDULED until it ends. + # 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. diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 3e26ca23f2c7e..0b40b64109978 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -3537,44 +3537,94 @@ def test_find_executable_task_instances_not_enough_task_concurrency_for_first(se session.rollback() - def test_find_executable_task_instances_uses_full_sized_refill_queries(self, dag_maker, session): - job_runner = SchedulerJobRunner(job=Job()) - dag_id = "SchedulerJobTest.test_find_executable_task_instances_uses_full_sized_refill_queries" - with dag_maker(dag_id=dag_id, session=session): + @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=priority_weight, - max_active_tis_per_dag=1, - ) - for index, priority_weight in enumerate((6, 3, 2)) + 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=priority_weight) - for index, priority_weight in enumerate((5, 4, 1)) + 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_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_tis = [ - first_run.get_task_instance(runnable_task.task_id, session=session) - for runnable_task in runnable_tasks - ] - for runnable_ti in runnable_tis: - runnable_ti.state = State.SCHEDULED + runnable_ti = second_run.get_task_instance(runnable_task.task_id, session=session) + runnable_ti.state = State.SCHEDULED session.flush() - with mock.patch.object(session, "scalars", wraps=session.scalars) as scalars: - queued_tis = job_runner._select_task_instances_to_queue( - 3, make_pool_stats(), set(), session=session - ) + 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} == {ti.key for ti in runnable_tis} - assert scalars.call_count == 2 + assert [ti.key for ti in queued_tis] == [runnable_ti.key] def test_find_executable_task_instances_task_concurrency_per_dagrun_for_first(self, dag_maker): scheduler_job = 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 45ce30f648cae..0f9a796f4807f 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." From e0726a48ecd5ff043b8ec5fda93aa8456f831a71 Mon Sep 17 00:00:00 2001 From: Arthur Date: Thu, 8 Oct 2026 14:24:50 +0200 Subject: [PATCH 5/6] Exclude only the saturated Dag run from scheduler refill queries --- .../src/airflow/jobs/scheduler_job_runner.py | 30 +++++------ .../tests/unit/jobs/test_scheduler_job.py | 54 +++++++++++++++++++ 2 files changed, 69 insertions(+), 15 deletions(-) diff --git a/airflow-core/src/airflow/jobs/scheduler_job_runner.py b/airflow-core/src/airflow/jobs/scheduler_job_runner.py index 7eb4e4a9e83de..eda8409c38d1e 100644 --- a/airflow-core/src/airflow/jobs/scheduler_job_runner.py +++ b/airflow-core/src/airflow/jobs/scheduler_job_runner.py @@ -776,8 +776,8 @@ def _select_task_instances_to_queue( # 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() @@ -786,13 +786,13 @@ def _select_task_instances_to_queue( 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, @@ -922,20 +922,20 @@ def _select_task_instances_to_queue( 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 @@ -994,7 +994,7 @@ def _select_task_instances_to_queue( # 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 ) @@ -1041,7 +1041,7 @@ def _select_task_instances_to_queue( 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, @@ -1051,7 +1051,7 @@ def _build_schedulable_tis_query( """ 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. @@ -1085,13 +1085,13 @@ def _build_schedulable_tis_query( ) # 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)) diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 0b40b64109978..15f01d9eb5c47 100644 --- a/airflow-core/tests/unit/jobs/test_scheduler_job.py +++ b/airflow-core/tests/unit/jobs/test_scheduler_job.py @@ -3626,6 +3626,60 @@ def test_select_task_instances_to_queue_finds_runnable_hidden_by_max_active_task 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) From 5c083773ed236cd4209caec16612d0a31f4f99c9 Mon Sep 17 00:00:00 2001 From: Arthur Date: Mon, 19 Oct 2026 15:13:37 +0200 Subject: [PATCH 6/6] Test that refill cap ignores queries before the first selection --- .../tests/unit/jobs/test_scheduler_job.py | 27 +++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/airflow-core/tests/unit/jobs/test_scheduler_job.py b/airflow-core/tests/unit/jobs/test_scheduler_job.py index 15f01d9eb5c47..89afbfebc08a1 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.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, ) @@ -3598,6 +3599,32 @@ def test_select_task_instances_to_queue_stops_at_refill_query_limit(self, mock_g 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 ):