Fill scheduler task batches after concurrency rejection

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]:
     """
@@ -754,28 +765,38 @@ 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
 
-        # 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")
@@ -825,15 +846,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
 
@@ -910,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
@@ -976,21 +988,37 @@ def _select_task_instances_to_queue(
 
                 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,
             )
 
@@ -1013,15 +1041,17 @@ 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,
+        *,
+        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.
@@ -1055,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))
@@ -1071,6 +1101,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:
+            # 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.

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,
 )
@@ -3537,6 +3538,175 @@ def test_find_executable_task_instances_not_enough_task_concurrency_for_first(se
 
         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)

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."