Failing to execute successive Dynamic Task Group Mapping

[pmercatoris]

Apache Airflow version

2.5.1

What happened

I am trying to map a task group within another task group which is already mapped. However, when launching the dag, all tasks of the containing task group successfully finish.
However, the following occurs:

• The contained task group is not mapped and has a state of upstream failed.
  
  

• The graph UI doesn't show after launching the dag as in Empty Graph view for Consecutive Dynamic Task Group Mapping #29287
  

What you think should happen instead

I would expect the print_group and print task to start as soon as 1 of the load task finishes

How to reproduce

I am currently using the docker-compose of the version 2.5.1
FROM apache/airflow:2.5.1-python3.10

import pendulum
from airflow.decorators import dag, task, task_group

@task
def get_symbols():
    res = [('A', 1, 111), ('B', 2, 222)]
    return res


@task
def print(symbol_info, data_interval_end=None):
    # Do some work...
    print(symbol_info)
    return symbol_info

@task_group()
def print_group(symbol):
    return print(symbol_info=symbol)


@task
def extract(symbol_info, data_interval_end=None):
    # Do some work...
    return symbol_info
@task
def transform(symbol_info, data_interval_end=None):
    # Do some work...
    return symbol_info

@task
def load(symbol_info, data_interval_end=None):
    # Do some work...
    return 2*[symbol_info]

@task_group
def etl_tg(symbol):
    raw_symbols_data = extract(symbol_info=symbol)
    clean_symbols_data = transform(symbol_info=raw_symbols_data)
    loaded_symbols = load(symbol_info=clean_symbols_data)
    return print_group.expand(symbol=loaded_symbols)

@dag(
    dag_id=f"task_group_mapping_2",
    tags=["sandbox"],
    schedule=None,
    start_date=pendulum.datetime(2023, 1, 1, tz="UTC"),
    catchup=False,
    max_active_runs=1,
)
def etl_dag():

    # DAG
    symbols = get_symbols()
    etl_tg.expand(symbol=symbols)

etl_dag() 

Operating System

Ubuntu 20.04.5 LTS

Versions of Apache Airflow Providers

No response

Deployment

Docker-Compose

Deployment details

No response

Anything else

No response

Are you willing to submit PR?

• Yes I am willing to submit a PR!

Code of Conduct

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

[hussein-awala]

Operator and task group expansion in an expanded task group is not supported yet. We raise NotImplementedError exception for the operators but I am not sure if it is raised in your case.

To implement the mapped tasks, we have added a new column map_index which represent the id of each mapped task, and we set it to -1 for the normal tasks. Supporting the nested expansion is not possible with the same logic, because we need to add a new column for each level, and the number of possible levels is infinite in this case (nested expansion in nested expansion ...), so I don't think we will support it soon.

[badrisrinivas9]

Can we have map_index in the task id. Something like group1[0].group2.task3 and then we can reuse the map_index column to specify if its normal task under last group it has expanded and go on for the nested expansion.
Task_id map_index
task1 -1
group1[0].task2 0
group1[1].task2 1
group1[0].group2.task3 -1
group1[0].group2.task4 0
group1[0].group2.task4 1
group1[1].group2.task3 -1
group1[1].group2.task4 0
group1[1].group2.task4 1

in group1 is expanded then in group2 task 4 is expanded. Just checking if it's feasible?

[potiuk]

as @hussein-awala explained, this is not supported yet.

[potiuk]

Looking at the code it's an automatically added dependency (not ment to be used by users) and it is added automatically when some of the max_active kwargs are used. Here is copy of the test that tests it:

class TestTaskConcurrencyDep:
    def _get_task(self, **kwargs):
        return BaseOperator(task_id="test_task", dag=DAG("test_dag"), **kwargs)

    @pytest.mark.parametrize(
        "kwargs, num_running_tis, is_task_concurrency_dep_met",
        [
            ({}, None, True),
            ({"max_active_tis_per_dag": 1}, 0, True),
            ({"max_active_tis_per_dag": 2}, 1, True),
            ({"max_active_tis_per_dag": 2}, 2, False),
            ({"max_active_tis_per_dagrun": 2}, 1, True),
            ({"max_active_tis_per_dagrun": 2}, 2, False),
            ({"max_active_tis_per_dag": 2, "max_active_tis_per_dagrun": 2}, 1, True),
            ({"max_active_tis_per_dag": 1, "max_active_tis_per_dagrun": 2}, 1, False),
            ({"max_active_tis_per_dag": 2, "max_active_tis_per_dagrun": 1}, 1, False),
            ({"max_active_tis_per_dag": 1, "max_active_tis_per_dagrun": 1}, 1, False),
        ],
    )
    def test_concurrency(self, kwargs, num_running_tis, is_task_concurrency_dep_met):
        task = self._get_task(start_date=datetime(2016, 1, 1), **kwargs)
        dep_context = DepContext()
        ti = Mock(task=task, execution_date=datetime(2016, 1, 1))
        if num_running_tis is not None:
            ti.get_num_running_task_instances.return_value = num_running_tis
        assert TaskConcurrencyDep().is_met(ti=ti, dep_context=dep_context) == is_task_concurrency_dep_met

[mis98zb]

Thank you for the help!

I still have question:
Since it is checked during scheduling the TIs (ScheduleJobRunner._schedule_dag_run()), why it is checked again when queue these selected TIs (ScheduleJobRunner._executable_task_instances_to_queued())?
Is this duplicated?

sequenceDiagram

ScheduleJobRunner ->>+ ScheduleJobRunner: _execute()/_run_scheduler_loop()/_do_scheduling()
ScheduleJobRunner ->>+ ScheduleJobRunner: _start_queued_dagruns()
ScheduleJobRunner ->> ScheduleJobRunner: _get_next_dagruns_to_examine(QUEUED)
ScheduleJobRunner ->>- DagRun: notify_dagrun_state_changed()
ScheduleJobRunner ->> ScheduleJobRunner: _get_next_dagruns_to_examine(RUNNING)
ScheduleJobRunner ->>+ ScheduleJobRunner: _schedule_all_dag_runs()
ScheduleJobRunner ->>+ ScheduleJobRunner: _schedule_dag_run()
ScheduleJobRunner ->>+ DagRun: update_state() //schedule_tasks()
DagRun ->>+ DagRun: task_instance_scheduling_decisions()
DagRun ->>+ DagRun: _get_ready_tis()
DagRun ->>+ TaskInstance: are_dependencies_met()
TaskInstance ->>+ TaskInstance: get_failed_dep_statuses()
TaskInstance ->>+ xxxxxDep: BaseTIDep::get_dep_statuses()
xxxxxDep ->>- xxxxxDep: _get_dep_statuses()
deactivate TaskInstance
deactivate TaskInstance
DagRun ->>+ DagRun: _expand_mapped_task_if_needed()
DagRun ->>- xxxxOperator: expand_mapped_task()
deactivate DagRun
deactivate DagRun
deactivate DagRun
ScheduleJobRunner ->>- DagRun: schedule_tis()
deactivate ScheduleJobRunner
ScheduleJobRunner ->> ScheduleJobRunner: _send_dag_callbacks_to_processor() 
ScheduleJobRunner ->>+ ScheduleJobRunner: _critical_section_enqueue_task_instances()
ScheduleJobRunner ->>+ ScheduleJobRunner: _executable_task_instances_to_queued()
ScheduleJobRunner ->> ScheduleJobRunner: query scheduled task instances to examine
ScheduleJobRunner ->>+ ScheduleJobRunner: __get_concurrency_maps()
ScheduleJobRunner ->>- ConcurrencyMap: from_concurrency_map()
deactivate ScheduleJobRunner
ScheduleJobRunner ->>- ScheduleJobRunner: _enqueue_task_instances_with_queued_state()
deactivate ScheduleJobRunner

Loading

[potiuk]

Scheduling happens first (creating dag run entries in the table) and then scheduler schedules tasks. You can find a talk https://www.youtube.com/watch?v=DYC4-xElccE from Airflow Summit if you want to get a deep dive how scheduler works

[mis98zb]

Thank you for the information.
But with the link https://www.youtube.com/watch?v=DYC4-xElccE I only see a 30s advertising video....
Is this link correct?

[potiuk]

Not sure. Try to search "deep dive airflow scheduler" - maybe it was a mobile copy&paste, but you will fnd it

https://www.youtube.com/watch?v=DYC4-xElccE

I try to copy & paste it again, maybe it will work - if not, just search for it.