Fix pre-mature evaluation of tasks in mapped task group

[ephraimbuddy]

Getting the relevant upstream indexes of a task instance in a mapped task group should only be done when the task has expanded. If the task has not expanded yet, we should return None so that the task can wait for the upstreams before trying to run.
This issue is more noticeable when the trigger rule is ONE_FAILED because then, the task instance is marked as SKIPPED.
This commit fixes this issue.
closes: #34023

[uranusjr]

I wonder why this case does not raise NotFullyPopulated

[ephraimbuddy]

From this

airflow/airflow/ti_deps/deps/trigger_rule_dep.py

Line 128 in 1b26cc0

"""Get how many tis the current task is supposed to be expanded into.

, it seems like it's the supposed number it could expand into but I think this case is because the group has expanded but the task hasn't?

[uranusjr]

I digged around a bit, this counts as expanded because there is indeed enough information to expand the operator, but the other depended task (t2) has not finished. This makes me think the fix is not right; we should investigate why t2 is not considered relevant instead. We need to somehow distinguish upstreams that are depended on as a DAG structure (e.g. with >>), or depended on for expansion (what _get_relevant_upstream_map_indexes is for); for >> dependencies we need to depend on the entire task regardless of map indexes instead.

[uranusjr]

I think the fix to this would be we should check whether an upstream is a dependency used for task expansion. If it is, we go through the normal “relevant map index” check; if not, we need to depend on its entirety i.e. _get_relevant_upstream_map_indexes should return None.

This change seems to work for me:

def _iter_expansion_dependencies() -> Iterator[str]:
    from airflow.models.mappedoperator import MappedOperator

    if isinstance(ti.task, MappedOperator):
        for op in ti.task.iter_mapped_dependencies():
            yield op.task_id
    for tg in ti.task.task_group.iter_mapped_task_groups():
        for op in tg.iter_mapped_dependencies():
            yield op.task_id

def _get_relevant_upstream_map_indexes(upstream_id: str) -> int | range | None:
    if upstream_id not in set(_iter_expansion_dependencies()):
        return None
    # Same as before...

This is obviously not optimised and I want to add a better abstraction for this. But @ephraimbuddy maybe you can help check if this can fix the actual DAG in the reported issue.

[ephraimbuddy]

I think the fix to this would be we should check whether an upstream is a dependency used for task expansion. If it is, we go through the normal “relevant map index” check; if not, we need to depend on its entirety i.e. _get_relevant_upstream_map_indexes should return None.

This change seems to work for me:

def _iter_expansion_dependencies() -> Iterator[str]:
    from airflow.models.mappedoperator import MappedOperator

    if isinstance(ti.task, MappedOperator):
        for op in ti.task.iter_mapped_dependencies():
            yield op.task_id
    for tg in ti.task.task_group.iter_mapped_task_groups():
        for op in tg.iter_mapped_dependencies():
            yield op.task_id

def _get_relevant_upstream_map_indexes(upstream_id: str) -> int | range | None:
    if upstream_id not in set(_iter_expansion_dependencies()):
        return None
    # Same as before...

This is obviously not optimised and I want to add a better abstraction for this. But @ephraimbuddy maybe you can help check if this can fix the actual DAG in the reported issue.

It works. Feel free to optimize it on this PR. Also I modified the code to make mypy happy

Edit:
@uranusjr, this works however, all the upstream tasks have to be completed before the task can run. This might not go well with the fast fail feature where if we have multiple upstream then the first fail will trigger a failure on its downstream even if some upstream tasks are still running. If this behaviour is ok for the mapped task group, I can go ahead and address the failing test.

[uranusjr]

all the upstream tasks have to be completed before the task can run. This might not go well with the fast fail feature where if we have multiple upstream then the first fail will trigger a failure on its downstream even if some upstream tasks are still running.

Does this mean with this PR the downstream will run when the first failure happens, or when all tasks (success or not) finish?

[ephraimbuddy]

when all tasks (success or not) finish?

Yes. When all tasks are completed before the downstream

[ephraimbuddy]

@uranusjr, Do you think I should fix the tests and get this ready?