rjgoyln commented on code in PR #72668:
URL: https://github.com/apache/airflow/pull/72668#discussion_r3963871089
##########
airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py:
##########
@@ -89,6 +89,28 @@ def calculate(cls, finished_upstreams:
Iterator[TaskInstance]) -> _UpstreamTISta
)
+def _covers_expansion_one_to_one(task, counted_task_ids: Collection[str]) ->
bool:
+ """
+ Whether the counted upstream instances map 1:1 onto this task's expansion.
+
+ Inferring instance removal from upstream state counts is only sound when
this task expands
+ over exactly one XCom reference, that reference is the entire counted
upstream set, and no
+ literal or mapped task group multiplies the length. Otherwise the task's
expansion length
+ has no relation to the counts: comparing ``map_index`` against the success
count would
+ remove valid instances (static lists, zips, multiple upstreams).
+ """
+ if task.get_closest_mapped_task_group() is not None:
+ return False
+ iter_deps = getattr(task, "iter_mapped_dependencies", None)
+ if iter_deps is None:
+ return False
Review Comment:
Minor typing nit: could we use the existing `is_mapped(task)` pattern
instead of `getattr()` here? It’s already used in
`_iter_expansion_dependencies` and provides proper `TypeGuard` narrowing.
Also, `task` is the only untyped parameter in this module, so annotating it
as `Operator` would let mypy check the function body as well.
##########
airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py:
##########
@@ -432,9 +453,12 @@ def _evaluate_direct_relatives() -> Iterator[TIDepStatus]:
new_state = TaskInstanceState.UPSTREAM_FAILED
elif skipped:
new_state = TaskInstanceState.SKIPPED
- elif removed and success and ti.map_index > -1:
- if ti.map_index >= success:
- new_state = TaskInstanceState.REMOVED
+ elif removed and ti.map_index > -1:
+ # All counted (live) instances are successes here, so
with a strict 1:1
+ # mapping the success count is the ti's own
post-shrink expansion length.
+ if upstream_done and
_covers_expansion_one_to_one(task, task.upstream_task_ids):
+ if ti.map_index >= success:
+ new_state = TaskInstanceState.REMOVED
Review Comment:
One concern: dropping `and success` changes behavior, but this case is not
mentioned or covered by a test.
With an unmapped upstream whose single TI is `REMOVED`,
`_covers_expansion_one_to_one()` can still return `True`, and `success == 0`
makes `map_index >= success` true for every downstream instance. I reproduced
this through the same scheduling-decision path:
```python
@task
def producer():
return [1, 2, 3]
@task
def work(arg):
return arg
work.expand(arg=producer())
# run producer, set its single ti to REMOVED, run the scheduling decision
again
```
`main` leaves all three `work` instances at `None`; this PR marks all three
`REMOVED`.
Removal may be the intended outcome since the expansion input is no longer
resolvable, but this should be an explicit, tested decision. Could we either
keep `and success`, or add a regression test documenting the intended behavior?
##########
airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py:
##########
@@ -89,6 +89,28 @@ def calculate(cls, finished_upstreams:
Iterator[TaskInstance]) -> _UpstreamTISta
)
+def _covers_expansion_one_to_one(task, counted_task_ids: Collection[str]) ->
bool:
+ """
+ Whether the counted upstream instances map 1:1 onto this task's expansion.
+
+ Inferring instance removal from upstream state counts is only sound when
this task expands
+ over exactly one XCom reference, that reference is the entire counted
upstream set, and no
+ literal or mapped task group multiplies the length. Otherwise the task's
expansion length
+ has no relation to the counts: comparing ``map_index`` against the success
count would
+ remove valid instances (static lists, zips, multiple upstreams).
+ """
+ if task.get_closest_mapped_task_group() is not None:
+ return False
+ iter_deps = getattr(task, "iter_mapped_dependencies", None)
+ if iter_deps is None:
+ return False
+ referenced = [op.task_id for op in iter_deps()]
+ if len(referenced) != 1 or set(referenced) != set(counted_task_ids):
+ return False
+ expand_value = task._get_specified_expand_input().value
+ return not isinstance(expand_value, Mapping) or len(expand_value) == 1
Review Comment:
One important correctness point: the helper also needs to verify that the
referenced upstream itself needs expansion. Without that, an unmapped upstream
returning a list can look like a 1:1 mapping even though it has only one TI and
produces multiple downstream instances.
I'd suggest using `get_needs_expansion()` rather than `is_mapped()`, since a
plain task inside a mapped task group can still legitimately produce one XCom
per instance.
Combined with the typing point above:
```python
def _covers_expansion_one_to_one(task: Operator, counted_task_ids:
Collection[str]) -> bool:
"""Whether the counted upstream instances map 1:1 onto this task's
expansion."""
from airflow.serialization.definitions.mappedoperator import is_mapped
if task.get_closest_mapped_task_group() is not None or not
is_mapped(task):
return False
referenced = list(task.iter_mapped_dependencies())
if len(referenced) != 1 or {op.task_id for op in referenced} !=
set(counted_task_ids):
return False
if not referenced[0].get_needs_expansion():
return False
expand_value = task._get_specified_expand_input().value
return not isinstance(expand_value, Mapping) or len(expand_value) == 1
```
I tested this version locally: ruff is clean, `ti_deps` passes, and the
unmapped-upstream case below keeps main's behavior.
##########
airflow-core/src/airflow/ti_deps/deps/trigger_rule_dep.py:
##########
@@ -318,9 +340,8 @@ def _evaluate_setup_constraint(
new_state = TaskInstanceState.UPSTREAM_FAILED
elif skipped:
new_state = TaskInstanceState.SKIPPED
- elif removed and success and ti.map_index > -1:
- if ti.map_index >= success:
- new_state = TaskInstanceState.REMOVED
+ # No REMOVED inference here: a task can only expand over a
direct upstream,
+ # and this branch counts indirect setups only.
Review Comment:
One test coverage concern: this deletion is not currently covered by a test
that fails without it.
I restored the removed branch and ran `airflow-core/tests/unit/ti_deps/`:
323 passed, with no test detecting the difference.
`test_setup_constraint_mapped_task_upstream_removed_and_success` still passes,
but through `_evaluate_direct_relatives` — `setup_3` expands 1:1 over
`setup_2`, so it no longer exercises the deleted branch.
The deletion itself looks sound to me, but since this changes behavior,
could we add a regression test for a non-1:1 downstream (e.g. mapped over a
literal list) with indirect setups? That would ensure the test fails without
this change.
##########
airflow-core/tests/unit/models/test_taskinstance.py:
##########
@@ -1246,8 +1246,8 @@ def test_check_task_dependencies(
["all_success", _UpstreamTIStates(2, 0, 0, 0, 0, 0, 0, 0), True,
None, False],
["all_success", _UpstreamTIStates(2, 0, 1, 0, 0, 0, 0, 0), True,
State.UPSTREAM_FAILED, False],
["all_success", _UpstreamTIStates(2, 1, 0, 0, 0, 0, 0, 0), True,
State.SKIPPED, False],
- # ti.map_index >= success
- ["all_success", _UpstreamTIStates(3, 0, 0, 0, 2, 0, 0, 0), True,
State.REMOVED, True],
+ # 1:1 mapped upstream shrank: map_index >= success once the
upstream is done
+ ["all_success", _UpstreamTIStates(3, 0, 0, 0, 2, 5, 0, 0), True,
State.REMOVED, True],
Review Comment:
One small test coverage point: this replaces the existing row rather than
adding a new one, so we lose coverage of the other side of the new
`upstream_done` gate.
Could we keep both cases? I checked that keeping the old `done=0` row passes
as written, and it now specifically verifies that no removal inference happens
before the upstream is complete.
--
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]
