control state of individual taskflow in mapped task-group

[ntnhaatj]

Description

Hi, as my issue was raised here which also included depth-first execution model context, I would like to open an issue to have more discussion on this matter.

Use case/motivation

motivated from the depth-first execution model implemented in mapped TaskGroup:

                     ┌──────────────┐ ┌────────────────┐ ┌───────────┐
                 ┌──►│extract_file_1│►│transform_file_1│►│load_file_1│
                 │   └──────────────┘ └────────────────┘ └───────────┘
                 │
                 │   ┌──────────────┐ ┌────────────────┐ ┌───────────┐
 ┌─────────────┐ ├──►│extract_file_2│►│transform_file_2│►│load_file_2│
 │get_file_list├─┤   └──────────────┘ └────────────────┘ └───────────┘
 └─────────────┘ │
                 │         ...               ...              ...
                 │
                 │   ┌──────────────┐ ┌────────────────┐ ┌───────────┐
                 └──►│extract_file_N│►│transform_file_n│►│load_file_N│
                     └──────────────┘ └────────────────┘ └───────────┘

At present, the upstream / downstream list dependencies now only applied on DAGNode (which is the whole TaskGroup or Operator)

There might be better control over desired TaskFlow if we could enhance support for deeper upstream/downstream dependencies at the granularity of mapped task instances, instead of applying them to the entire TaskGroup.

In the model above, extract_file_1 task downstream list should be only transform_file_1 and load_file_1 in order, rather than encompassing the whole mapped group transform_file[] and load_file[] as in the current implementation.

For instance, I scheduled my test DAG on Airflow 2.7.0:

from airflow.decorators import dag, task_group, task
from airflow.operators.empty import EmptyOperator
from pendulum import datetime

files = ["a", "b", "c"]

@dag(start_date=datetime(2022, 12, 1), schedule=None, catchup=False)
def task_group_mapping_example():
    @task_group(group_id="etl")
    def etl_pipeline(file):
        e = EmptyOperator(task_id="e")
        t = EmptyOperator(task_id="t")
        l = EmptyOperator(task_id="l")

        e >> t >> l

    etl = etl_pipeline.expand(file=files)
    etl

task_group_mapping_example()

Clear etl.e[1] state with downstream will trigger all mapped task in etl.t[] and etl.l[] group.

Thanks,

Related issues

No response

Are you willing to submit a PR?

• Yes I am willing to submit a PR!

Code of Conduct

• I agree to follow this project's Code of Conduct

[boring-cyborg[bot]]

Thanks for opening your first issue here! Be sure to follow the issue template! If you are willing to raise PR to address this issue please do so, no need to wait for approval.

[potiuk]

Converted it into discussion, since it is, well, discussion.

[potiuk]

I believe this is a bug @uranusjr ? From what I understand how mapping task groups work, clearing only a single mapped task instance in such case should not clear all the other mapped-indexed downstream tasks, just the tasks within the same task group index ?

[ntnhaatj]

hi @potiuk and @uranusjr, is there any update on this?

[uranusjr]

Yeah I think this is not right, probably something in the logic that calculates what tasks should be affected.

[ntnhaatj]

Hi, apologize for bumping this again.

Retrying an individual task within a mapped task group would require retrying the entire task group. This becomes cumbersome when my task group contains hundreds of pipelines.

Any improvement would be greatly appreciated. Thank you a lot!