How to use Dynamic Task Mapping with TaskGroups

Question

In my actual DAG, I need to first get a list of IDs and then for each ID run a set of tasks.

I have used Dynamic Task Mapping to pass a list to a single task or operator to have it process the list, but can we do this using a TaskGroup as well?

If I can figure out how to pass a variable value at the TaskGroup level, so it uses that value in all sub tasks, then I should be able to meet my requirement.

The below should give you an idea of what I am looking for, just need help getting it working.

from airflow import DAG, XComArg
from datetime import datetime
from airflow.decorators import task
from airflow.utils.task_group import TaskGroup
from airflow.operators.python import PythonOperator


with DAG(
    'dtm_tg_test',
    schedule_interval = None,
    start_date = datetime(2022, 1, 1)
) as dag:

    def getList():
        return [ "Hello", "World" ]

    def printText(text):
        print(text)

    get_list = PythonOperator(
                        task_id = "get_list",
                        python_callable = getList,
                        dag = dag
                    )

    with TaskGroup.partial(
                            group_id = "task_group"
    ).expand(
        list = XComArg(get_list)
    ) as task_group:
        print_text = PythonOperator(
                            task_id = "print_output",
                            python_callable = printText,
                            op_kwargs = { "text": list }
                            dag = dag
                        )
        
        print_again = PythonOperator(
                            task_id = "print_output",
                            python_callable = printText,
                            op_kwargs = { "text": list }
                            dag = dag
                        )

        print_text >> print_again

    get_list >> task_group

Answer 1

You can achieve it with the following example :

list_ids = ['45', '48']

@task_group()
def parent_group(list_ids: List[str]) -> List[TaskGroup]:
    return list(map(build_group_for_id, list_ids))

def build_group_for_id(current_id: str) -> TaskGroup:
    with TaskGroup(group_id=f'group_for_id_{current_id}') as group:
        print_text = PythonOperator(
                            task_id = f"print_output_{current_id}",
                            python_callable = printText,
                            op_kwargs = { "text": current_id }
                            dag = dag
                        )
        
        print_again = PythonOperator(
                            task_id = f"print_output_other_{current_id}",
                            python_callable = printText,
                            op_kwargs = { "text":  current_id}
                            dag = dag

        print_text >> print_again

    return group

with airflow.DAG(
        "my_dag", default_args=args, schedule_interval=None,
) as dag:

   DummyOperator(task_id='start_dag') >> parent_group(list_ids())

Some explanations :

• I create a parent taskGroup called parent_group
• This parent group takes the list of IDs
• I add a loop and for each parent ID, I create a TaskGroup containing your 2 Aiflow tasks (print operators)
• For the TaskGroup related to a parent ID, the TaskGroup ID is built from it in order to be unique in the DAG
• For the print operators inside the TaskGroup, I generated again the task IDs by the current parent ID