kaxil commented on code in PR #67672:
URL: https://github.com/apache/airflow/pull/67672#discussion_r3673565046
##########
airflow-core/src/airflow/models/dagrun.py:
##########
@@ -1520,16 +1520,26 @@ def _expand_mapped_task_if_needed(ti: TI) ->
Iterable[TI] | None:
if new_tis is not None:
additional_tis.extend(new_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
+ # later in this same pass must see the post-expansion
count).
+ dep_context.upstream_task_id_counts.clear()
if new_tis is None and schedulable.state in SCHEDULEABLE_STATES:
# It's enough to revise map index once per task id,
# checking the map index for each mapped task significantly
slows down scheduling
if schedulable.task.task_id not in revised_map_index_task_ids:
- ready_tis.extend(
+ revised_tis = list(
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)
+ if revised_tis:
Review Comment:
`_revise_map_indexes_if_mapped` returns a `list[TI]` now, it stopped being a
generator in #69565, so the `list()` here was a redundant copy and there's no
need to compare `len(ready_tis)` before and after either. Dropped the wrapper:
```python
revised_tis = self._revise_map_indexes_if_mapped(
schedulable.task, dag_version_id=schedulable.dag_version_id,
session=session
)
ready_tis.extend(revised_tis)
```
Chasing this turned up a real one, though: the memo field was `init=False`,
and `attrs.evolve` only carries over fields that `__init__` accepts, so the
`UP_FOR_RESCHEDULE` evolve in `are_dependencies_met` handed every
reschedule-mode ti a fresh empty dict. Those tis never read the memo and never
warmed it for anything else, which is the exact fan-out this was meant to
collapse. Fixed and covered by a test.
--
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]
