concurrency control of mapped task group

[mis98zb]

And, to control the concurrency of task group, the option max_active_tis_per_dag and max_active_tis_per_dagrun on task_instance do not work well.

Below is my simplified case:

import random
from datetime import timedelta
import logging

import pendulum
from airflow import DAG
from airflow.decorators import task, task_group
from airflow.sensors.base import PokeReturnValue

log = logging.getLogger(__name__)


@task(task_id='prepareJobParameter')
def prepare_job_parameter(ti=None) -> []:
    log.info('prepare script~~~~~~~~~~~~~~~~')
    return [
        {'parameter': 'aaaaaaaaaaaaaaa'},
        {'parameter': 'bbbbbbbbbbbbbbb'},
        {'parameter': 'ccccccccccccccc'}
    ]


@task_group(group_id=f'runScopeJob')
def run_job(task_name: str, submit_param: dict, check_interval: int, timeout: timedelta) -> None:
    @task(task_id=f'submitJob')
    def submit_job(job_param: dict, ti=None) -> str:
        log.info('submit job~~~~~~~~~~~~~~~~')
        return 'job url'

    @task.sensor(task_id=f'waitJob', poke_interval=check_interval, timeout=timeout.total_seconds(), mode="reschedule")
    def wait_job(job_url: str, ti=None) -> PokeReturnValue:
        log.info('check job status~~~~~~~~~~~~~~~~')
        status = random.randint(0, 3)
        if status > 0:
            return PokeReturnValue(is_done=False, xcom_value=status)
        else:
            return PokeReturnValue(is_done=True, xcom_value=status)

    wait_job(submit_job(submit_param))


@task(task_id="publishData")
def publish_data(ti=None) -> bool:
    log.info('publish data~~~~~~~~~~~~~~~~')
    return True


with DAG(
        dag_id='test_task_group_parallel',
        catchup=False,
        start_date=pendulum.datetime(2023, 4, 15, tz="UTC"),
        schedule=timedelta(hours=4),
        tags=['igs pipeline']
) as pipeline:
    task_pre_job = prepare_job_parameter()
    task_run_job = run_job.partial(task_name='Bingen', check_interval=180, timeout=timedelta(hours=24)).expand(
        submit_param=task_pre_job)
    task_run_job >> publish_data()


if __name__ == "__main__":
    pipeline.test()

Inside the mapped task group, there are two tasks: submit_job and wait_job.
The submit_job will finish very quickly.
The wait_job will be rescheduled till the job complete.
The max_active_tis_per_dagrun on submit_job cannot control the concurrency of the job, because it finished quickly.
All jobs will be submitted in a short time.

[mis98zb]

I'm considering to add a property, concurrent_limit to calss MappedTaskGroup.
In class TaskConcurrencyDep, if the ti.task.task_group is MappedTaskGroup, we can check if the running task group count is less than the concurrent limitation.

The problem is how to detect the tis is belong to the task group.

1. if prefix_group_id is on, we can check prefix of task_id
2. another way is to check if ti.task_id is in ti.task.task_group.children.keys(). This needs full serialized DAG, which is expensive.
3. third way maybe add a column task_group in task_instance table, and add acolumn concurrency limitation in task_map table.

[mis98zb]

Considering the case of nested task groups, the 3rd solution may be changed as below:
Instead of add a column to task_instance table, just chage the map_index to string, and store the value as: {group name}:{map index}(|{group name}:{map index})*. The top level group is on the left side.