Copilot commented on code in PR #71967:
URL: https://github.com/apache/airflow/pull/71967#discussion_r3857042924
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -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])
Review Comment:
The budget only applies while the expansion is created. On the next call,
`task_instance_scheduling_decisions()` reloads every persisted `None`-state TI
and `_get_ready_tis` dependency-checks the entire set without this cap; the
added test even expects all four TIs on its second call. For a 10,000-way map
this therefore postpones the same 10,000-TI spike by one heartbeat instead of
smoothing it. Apply a deterministic per-pass budget to the full schedulable set
(including already-persisted and revised mapped TIs), then test that successive
passes remain bounded.
---
Drafted-by: GitHub Copilot (no human review before posting)
##########
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,
Review Comment:
This newly added mock has neither `spec` nor `autospec`, so it can accept
calls or attributes that the real session method would reject. Please give the
spy a spec while retaining `wraps`.
---
Drafted-by: GitHub Copilot (no human review before posting)
##########
airflow-core/src/airflow/models/taskmap.py:
##########
@@ -256,22 +257,63 @@ def expand_mapped_task(
)
)
- 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)
+ from airflow.settings import get_policy_plugin_manager,
task_instance_mutation_hook
+
+ policy_hook =
get_policy_plugin_manager().hook.task_instance_mutation_hook
+ hook_is_noop = False
+ if not isinstance(policy_hook, Mock):
+ hook_is_noop = getattr(task_instance_mutation_hook, "is_noop",
False) is True and all(
+ getattr(hook.function, "is_noop", False) is True for hook in
policy_hook.get_hookimpls()
Review Comment:
This production path changes behavior based on whether the Pluggy hook
object is a `unittest.mock.Mock`, leaking test implementation details into
runtime logic and requiring a testing module import in `taskmap.py`. Make the
centrally maintained no-op marker/hook contract authoritative, and have tests
explicitly configure that marker when replacing a hook rather than adding a
Mock-specific branch here.
---
Drafted-by: GitHub Copilot (no human review before posting)
##########
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")
Review Comment:
The newly introduced `max_tis_per_query <= 0` fallback is not exercised by
the added test, which only sets the value to `2`. Add a case with a
non-positive query limit and a small `[core] parallelism` value so a regression
that removes or misreads the fallback fails the test.
---
Drafted-by: GitHub Copilot (no human review before posting)
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -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,
Review Comment:
The new bulk branch in `_revise_map_indexes_if_mapped` has no regression
test that distinguishes it from the previous per-TI path. The existing
mapped-length test only checks resulting indexes and would still pass if this
branch were reverted. Add tests that verify the noop hook uses
`bulk_insert_mappings` and that a non-noop mutation hook is still invoked and
persisted for every newly revised index.
---
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]
