Copilot commented on code in PR #71967:
URL: https://github.com/apache/airflow/pull/71967#discussion_r3857294633
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1715,8 +1735,17 @@ def _expand_mapped_task_if_needed(ti: TI) ->
Iterable[TI] | None:
revised_tis = self._revise_map_indexes_if_mapped(
schedulable.task,
dag_version_id=schedulable.dag_version_id, session=session
)
- ready_tis.extend(revised_tis)
- revised_map_index_task_ids.add(schedulable.task.task_id)
+ remaining_budget = max(max_tis_per_query - len(ready_tis),
0)
+ ready_tis.extend(revised_tis[:remaining_budget])
+ if remaining_budget < len(revised_tis):
Review Comment:
The new deferral branch for map-index revision is not exercised. The
expansion tests cover newly expanded and already-persisted TIs, while the
existing length-increase test adds only one revised TI under the default
budget, so it never proves that excess `revised_tis` are persisted and returned
on a later pass. Add a low-budget map-length-growth regression test covering
both passes.
---
Drafted-by: GitHub Copilot (no human review before posting)
##########
airflow-core/tests/unit/models/test_dagrun.py:
##########
@@ -2139,6 +2139,174 @@ 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)]
+ for ti in decision.schedulable_tis:
+ ti.state = TaskInstanceState.SCHEDULED
+ 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", 2), ("task_2", 3)]
+
+
+def
test_mapped_expansion_defers_some_tis_with_non_positive_tis_query_limit(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"): "0", ("core",
"parallelism"): "1"}):
+ 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)]
+ (decision.schedulable_tis[0]).state = TaskInstanceState.SCHEDULED
+ session.flush()
+
+ with conf_vars({("scheduler", "max_tis_per_query"): "0", ("core",
"parallelism"): "1"}):
+ 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", 1)]
+
+
+def
test_revise_map_indexes_if_mapped_uses_bulk_insert_when_mutation_hook_is_noop(dag_maker,
session) -> None:
+ dag, mapped, dr = _make_mapped_dag_for_expansion(
+ dag_maker, session, dag_id="test_revise_map_indexes_bulk"
+ )
+
+ expand_mapped_task(
+ dag.task_dict[mapped.task_id],
+ dr.run_id,
+ "op1",
+ length=2,
+ session=session,
+ )
+
+ upstream_ti = dr.get_task_instance(task_id="op1", session=session)
+ assert upstream_ti
+ session.merge(TaskMap.from_task_instance_xcom(upstream_ti, [1, 2, 3, 4]))
+ session.flush()
+
+ dag_version_id_row = DagVersion.get_latest_version(dag_id=dr.dag_id,
session=session)
+ assert dag_version_id_row is not None
+ dag_version_id = dag_version_id_row.id
+
+ class NoopHook:
+ is_noop = True
+
+ def __call__(self, *_, **__):
+ return None
+
+ with (
+ mock.patch("airflow.settings.task_instance_mutation_hook", NoopHook()),
+ mock.patch.object(
+ session,
+ "bulk_insert_mappings",
+ wraps=session.bulk_insert_mappings,
+ spec=session.bulk_insert_mappings,
+ ) as bulk_insert,
+ ):
+ created_tis = dr._revise_map_indexes_if_mapped(mapped,
dag_version_id=dag_version_id, session=session)
+
+ assert bulk_insert.called
+ assert [ti.map_index for ti in created_tis] == [2, 3]
+
+
+def
test_revise_map_indexes_if_mapped_calls_mutation_hook_for_new_tis(dag_maker,
session) -> None:
+ dag, mapped, dr = _make_mapped_dag_for_expansion(
+ dag_maker, session, dag_id="test_revise_map_indexes_hooked"
+ )
+
+ expand_mapped_task(
+ dag.task_dict[mapped.task_id],
+ dr.run_id,
+ "op1",
+ length=2,
+ session=session,
+ )
+
+ upstream_ti = dr.get_task_instance(task_id="op1", session=session)
+ assert upstream_ti
+ session.merge(TaskMap.from_task_instance_xcom(upstream_ti, [1, 2, 3, 4]))
+ session.flush()
+
+ dag_version_id_row = DagVersion.get_latest_version(dag_id=dr.dag_id,
session=session)
+ assert dag_version_id_row is not None
+ dag_version_id = dag_version_id_row.id
+
+ existing_indexes = session.scalars(
+ select(TI.map_index)
+ .where(TI.dag_id == dr.dag_id, TI.task_id == mapped.task_id, TI.run_id
== dr.run_id)
+ .order_by(TI.map_index)
+ ).all()
+ assert existing_indexes == [0, 1]
+ called_indexes: list[int] = []
+
+ class MutationHook:
+ is_noop = False
+
+ def __call__(self, task_instance, dag_run=None):
+ called_indexes.append(task_instance.map_index)
+ task_instance.queue = f"q_{task_instance.map_index}"
+
+ mutation_hook = MutationHook()
+
+ with (
+ mock.patch("airflow.settings.task_instance_mutation_hook",
mutation_hook),
+ mock.patch("airflow.models.taskinstance.task_instance_mutation_hook",
autospec=True) as hook,
+ mock.patch.object(
+ session,
+ "bulk_insert_mappings",
+ wraps=session.bulk_insert_mappings,
+ spec=session.bulk_insert_mappings,
+ ) as bulk_insert,
+ ):
+ hook.side_effect = mutation_hook
+ created_tis = dr._revise_map_indexes_if_mapped(mapped,
dag_version_id=dag_version_id, session=session)
+
+ assert set(called_indexes) == {2, 3}
+ assert len(created_tis) == 2
+ assert not bulk_insert.called
+ assert [ti.queue for ti in sorted(created_tis, key=lambda ti:
ti.map_index)] == ["q_2", "q_3"]
Review Comment:
These assertions all pass against the pre-PR implementation: it already
invokes `_add_and_prime_mapped_ti` for every new index and never calls
`bulk_insert_mappings`. As written, this added test cannot detect a regression
in the changed behavior. Repurpose it to exercise the newly introduced
guard—for example, a registered policy hook while the wrapper's `is_noop`
marker remains true—or remove it as coverage padding.
---
Drafted-by: GitHub Copilot (no human review before posting)
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
