Reduce scheduler spikes when expanding large mapped tasks

Large mapped expansions can overload one scheduler pass, increasing scheduling latency and making unrelated Dags wait. This change smooths mapped fan-out handling so high-cardinality expansions do not monopolize scheduler heartbeats, while preserving custom task-instance mutation-hook behavior.

Author: hussein-awala

airflow-core/src/airflow/models/dagrun.py

@@ -1685,6 +1685,9 @@ def _expand_mapped_task_if_needed(ti: TI) -> Iterable[TI] | None:
         expansion_happened = False
         # Set of task ids for which was already done _revise_map_indexes_if_mapped
         revised_map_index_task_ids: set[str] = set()
+        max_tis_per_query = airflow_conf.getint("scheduler", "max_tis_per_query")
+        if max_tis_per_query <= 0:
+            max_tis_per_query = airflow_conf.getint("core", "parallelism")
         for schedulable in itertools.chain(schedulable_tis, additional_tis):
             if TYPE_CHECKING:
                 assert isinstance(schedulable.task, Operator)
@@ -1702,7 +1705,21 @@ def _expand_mapped_task_if_needed(ti: TI) -> Iterable[TI] | None:
             if schedulable.map_index < 0:
                 new_tis = _expand_mapped_task_if_needed(schedulable)
                 if new_tis is not None:
-                    additional_tis.extend(new_tis)
+                    expanded_tis = list(new_tis)
+                    # Avoid evaluating a huge number of newly expanded TIs in the same pass.
+                    # They are persisted already and picked up in subsequent loops.
+                    remaining_budget = max(max_tis_per_query - len(additional_tis), 0)
+                    if remaining_budget:
+                        additional_tis.extend(expanded_tis[:remaining_budget])
+                    dropped_tis = len(expanded_tis) - remaining_budget
+                    if dropped_tis > 0:
+                        self.log.debug(
+                            "Deferring dependency checks for expanded TIs to a later scheduler pass",
+                            task_id=schedulable.task_id,
+                            dag_id=self.dag_id,
+                            run_id=self.run_id,
+                            deferred_count=dropped_tis,
+                        )
                     expansion_happened = True
                     # Expansion changes a mapped task's instance count, which invalidates the
                     # trigger-rule upstream-count memo on this DepContext (a downstream evaluated
@@ -2145,17 +2162,50 @@ def _revise_map_indexes_if_mapped(
             )
             session.flush()
 
-        new_tis: list[TI] = []
-        for index in range(total_length):
-            if index in existing_indexes:
-                continue
+        from airflow.settings import task_instance_mutation_hook
+
+        new_indexes = [index for index in range(total_length) if index not in existing_indexes]
+        if not new_indexes:
+            return []
+
+        hook_is_noop = getattr(task_instance_mutation_hook, "is_noop", False) is True
+        if hook_is_noop:
+            ti_mappings = [
+                TI.insert_mapping(
+                    self.run_id,
+                    task,
+                    map_index=index,
+                    dag_version_id=dag_version_id,
+                    dag_run=self,
+                )
+                for index in new_indexes
+            ]
+            session.bulk_insert_mappings(TI.__mapper__, ti_mappings)
+            session.flush()
+            inserted_tis = list(
+                session.scalars(
+                    select(TI)
+                    .where(
+                        TI.dag_id == self.dag_id,
+                        TI.task_id == task.task_id,
+                        TI.run_id == self.run_id,
+                        TI.map_index.in_(new_indexes),
+                    )
+                    .order_by(TI.map_index)
+                ).all()
+            )
+            for ti in inserted_tis:
+                ti.task = task
+            return inserted_tis
+
+        created_tis: list[TI] = []
+        for index in new_indexes:
             ti = TI(task, run_id=self.run_id, map_index=index, state=None, dag_version_id=dag_version_id)
             self.log.debug("Expanding TIs upserted %s", ti)
             _add_and_prime_mapped_ti(ti, task, self, session=session)
-            new_tis.append(ti)
-        if new_tis:
-            session.flush()
-        return new_tis
+            created_tis.append(ti)
+        session.flush()
+        return created_tis
 
     @classmethod
     @provide_session

airflow-core/src/airflow/models/taskmap.py

@@ -256,22 +256,58 @@ def expand_mapped_task(
                 )
             )
 
+        from airflow.settings import task_instance_mutation_hook
+
+        hook_is_noop = getattr(task_instance_mutation_hook, "is_noop", False) is True
         new_tis: list[TaskInstance] = []
-        for index in indexes_to_map:
-            ti = TaskInstance(
-                task,
-                run_id=run_id,
-                map_index=index,
-                state=state,
-                dag_version_id=dag_version_id,
-            )
-            task.log.debug("Expanding TIs upserted %s", ti)
-            _add_and_prime_mapped_ti(
-                ti, task, dr, session=session, context_carrier=new_task_run_carrier(dr.context_carrier)
-            )
-            new_tis.append(ti)
-        if new_tis:
-            session.flush()
+        if hook_is_noop and isinstance(indexes_to_map, range):
+            ti_mappings = [
+                TaskInstance.insert_mapping(
+                    run_id,
+                    task,
+                    map_index=index,
+                    dag_version_id=dag_version_id,
+                    dag_run=dr,
+                )
+                for index in indexes_to_map
+            ]
+            if state is not None:
+                for ti_mapping in ti_mappings:
+                    ti_mapping["state"] = state
+            if ti_mappings:
+                session.bulk_insert_mappings(TaskInstance.__mapper__, ti_mappings)
+                session.flush()
+                new_tis = list(
+                    session.scalars(
+                        select(TaskInstance)
+                        .where(
+                            TaskInstance.dag_id == task.dag_id,
+                            TaskInstance.task_id == task.task_id,
+                            TaskInstance.run_id == run_id,
+                            TaskInstance.map_index >= indexes_to_map.start,
+                            TaskInstance.map_index < indexes_to_map.stop,
+                        )
+                        .order_by(TaskInstance.map_index)
+                    ).all()
+                )
+                for ti in new_tis:
+                    ti.task = task
+        else:
+            for index in indexes_to_map:
+                ti = TaskInstance(
+                    task,
+                    run_id=run_id,
+                    map_index=index,
+                    state=state,
+                    dag_version_id=dag_version_id,
+                )
+                task.log.debug("Expanding TIs upserted %s", ti)
+                _add_and_prime_mapped_ti(
+                    ti, task, dr, session=session, context_carrier=new_task_run_carrier(dr.context_carrier)
+                )
+                new_tis.append(ti)
+            if new_tis:
+                session.flush()
         all_expanded_tis.extend(new_tis)
 
         # Coerce the None case to 0 -- these two are almost treated identically,

airflow-core/tests/unit/models/test_dagrun.py

@@ -2139,6 +2139,37 @@ def task_2(arg2): ...
     ]
 
 
+def test_mapped_expansion_defers_some_tis_to_later_scheduler_pass(dag_maker, session):
+    @task
+    def task_1(): ...
+
+    with dag_maker(session=session):
+
+        @task
+        def task_2(arg2): ...
+
+        task_2.expand(arg2=task_1())
+
+    dr: DagRun = dag_maker.create_dagrun()
+    ti = dr.get_task_instance(task_id="task_1", session=session)
+    assert ti
+    ti.state = TaskInstanceState.SUCCESS
+    session.add(TaskMap.from_task_instance_xcom(ti, [1, 2, 3, 4]))
+    session.flush()
+
+    with conf_vars({("scheduler", "max_tis_per_query"): "2"}):
+        decision = dr.task_instance_scheduling_decisions(session=session)
+
+    indices = [(ti.task_id, ti.map_index) for ti in decision.schedulable_tis]
+    assert indices == [("task_2", 0), ("task_2", 1)]
+
+    with conf_vars({("scheduler", "max_tis_per_query"): "2"}):
+        decision = dr.task_instance_scheduling_decisions(session=session)
+
+    indices = [(ti.task_id, ti.map_index) for ti in decision.schedulable_tis]
+    assert indices == [("task_2", 0), ("task_2", 1), ("task_2", 2), ("task_2", 3)]
+
+
 def test_mapped_literal_length_reduction_at_runtime_adds_removed_state(dag_maker, session):
     """
     Test that when the length of mapped literal reduces at runtime, the missing task instances

airflow-core/tests/unit/models/test_mappedoperator.py

@@ -501,6 +501,47 @@ def test_expand_mapped_task_task_instance_mutation_hook(dag_maker, session, crea
             assert call.args[0].map_index == expected_map_index[index]
 
 
+def test_expand_mapped_task_uses_bulk_insert_when_mutation_hook_is_noop(dag_maker, session) -> None:
+    with dag_maker(session=session, serialized=True) as dag:
+        task1 = BaseOperator(task_id="op1")
+        mapped = MockOperator.partial(task_id="task_2").expand(arg2=task1.output)
+
+    dr = dag_maker.create_dagrun()
+
+    class NoopHook:
+        is_noop = True
+
+        def __call__(self, *_, **__):
+            return None
+
+    noop_hook = NoopHook()
+
+    with (
+        mock.patch("airflow.settings.task_instance_mutation_hook", noop_hook),
+        mock.patch.object(session, "bulk_insert_mappings", wraps=session.bulk_insert_mappings) as bulk_insert,
+    ):
+        expand_mapped_task(
+            dag.task_dict[mapped.task_id],
+            dr.run_id,
+            task1.task_id,
+            length=4,
+            session=session,
+        )
+
+    assert bulk_insert.called
+    mapped_indexes = session.scalars(
+        select(TaskInstance.map_index)
+        .where(
+            TaskInstance.dag_id == dag.dag_id,
+            TaskInstance.task_id == mapped.task_id,
+            TaskInstance.run_id == dr.run_id,
+            TaskInstance.map_index >= 0,
+        )
+        .order_by(TaskInstance.map_index)
+    ).all()
+    assert mapped_indexes == [0, 1, 2, 3]
+
+
 class TestMappedSetupTeardown:
     @staticmethod
     def get_states(dr):