Add max_active_tis_per_dagrun for Dynamic Task Mapping

[hussein-awala]

closes: #29084

---

This PR adds a new operator parameter max_active_tis_per_dagrun to control the number of task instances of the task can be running in parallel in the same Dag Run.
This conf is useless for the normal tasks when the number of task instances of the task is always <= 1 in the same DAG run, but very useful and important for the Dynamic Task Mapping to control its parallelism in each different DAG run.

[hussein-awala]

The PR is ready for review, I tried to cover all the changes in my tests.
And here is a simple dag to test it visually:

import pendulum
import time

from airflow.decorators import dag, task


@dag(
    dag_id='max_active_tis_per_dagrun',
    default_args={},
    start_date=pendulum.datetime(2021, 1, 1, tz="UTC"),
    schedule=None
)
def processing_dag():

    @task
    def get_numbers():
        return list(range(20))

    @task(max_active_tis_per_dagrun=2)
    def process(number):
        print(number)
        time.sleep(5)

    numbers = get_numbers()

    process.expand(number=numbers)


my_dag = processing_dag()

You can trigger 4 DAG runs via the UI, and check how many mapped task will be running in parallel in each DAG run.

[potiuk]

I think this one is great - though I do not understand all the details of mapped tasks and concurrency there. @uranusjr @ashb - maybe you can take a look, sounds like a great feature of Dynamic Task Mapping.

[uranusjr]

Looks good to me in general, with some nitpicks.

[VladimirYushkevich]

@hussein-awala, in your example for me always one task from []Mapped Tasks is running. I'm using airflow 2.8.1 and trying to find setting that allows to run all task in []Mapped Tasks in parallel.

[hussein-awala]

Could you provide the dag and the task configurations?

[VladimirYushkevich]

airflow dags details example_dag:

property_name                   | property_value                                                                                                                                 
================================+================================================================================================================================================
dag_id                          | example_dag                                                                                                                                    
root_dag_id                     | None                                                                                                                                           
is_paused                       | False                                                                                                                                          
is_active                       | True                                                                                                                                           
is_subdag                       | False                                                                                                                                          
last_parsed_time                | 2024-03-15T23:34:18.998039+00:00                                                                                                               
last_pickled                    | None                                                                                                                                           
last_expired                    | None                                                                                                                                           
scheduler_lock                  | None                                                                                                                                           
pickle_id                       | None                                                                                                                                           
default_view                    | grid                                                                                                                                           
owners                          | airflow                                                                                                                                        
description                     | None                                                                                                                                           
schedule_interval               | {'__type': 'CronExpression', 'value': '@hourly'}                                                                                               
timetable_description           | Every hour                                                                                                                                                                                                
max_active_tasks                | 16                                                                                                                                             
max_active_runs                 | 4                                                                                                                                              
has_task_concurrency_limits     | True                                                                                                                                           
has_import_errors               | False                                                                                                                                          
next_dagrun                     | 2024-03-15T02:00:00+00:00                                                                                                                      
next_dagrun_data_interval_start | 2024-03-15T02:00:00+00:00                                                                                                                      
next_dagrun_data_interval_end   | 2024-03-15T03:00:00+00:00                                                                                                                      
next_dagrun_create_after        | None

task:

Attribute | Value
-- | --
custom_operator_name | None
dag_model | <DAG: example_dag>
dag_run | <DagRun example_dag @ 2024-03-15 02:00:00+00:00: scheduled__2024-03-15T02:00:00+00:00, state:running, queued_at: 2024-03-15 23:35:19.813859+00:00. externally triggered: False>
duration | None
end_date | None
execution_date | 2024-03-15 02:00:00+00:00
executor_config | {}
external_executor_id | None
hostname |  
is_premature | False
is_trigger_log_context | False
job_id | None
map_index | 0
max_tries | 0
metadata | MetaData()
next_kwargs | None
next_method | None
next_try_number | 1
note | None
operator | _PythonDecoratedOperator
operator_name | _PythonDecoratedOperator
pid | None
pool | default_pool
pool_slots | 1
prev_attempted_tries | 0
previous_start_date_success | None
previous_ti | None
previous_ti_success | None
priority_weight | 1
queue | default
queued_by_job | None
queued_by_job_id | None
queued_dttm | None
registry | <sqlalchemy.orm.decl_api.registry object at 0x1219edb50>
rendered_task_instance_fields | None
run_as_user | None
run_id | scheduled__2024-03-15T02:00:00+00:00
start_date | None
state | scheduled
task_id | process
task_instance_note | None
test_mode | False
trigger_id | None
trigger_timeout | None
try_number | 0
unixname | vladimir
updated_at | 2024-03-15, 23:35:21

[VladimirYushkevich]

it is pretty much the same as example here:

@dag(
    dag_id="example_dag",
    schedule="@hourly",
    start_date=pendulum.today("UTC").add(hours=-1),
    is_paused_upon_creation=False,
    max_active_runs=4,
    default_args={},
)
def example_dag():
    @task
    def get_numbers():
        return list(range(20))

    @task(max_active_tis_per_dagrun=4, max_active_tis_per_dag=4)
    def process(number):
        print(number)
        time.sleep(5)

    numbers = get_numbers()

    process.expand(number=numbers)

[hussein-awala]

I tested the same dag with main branch, and I had 4 TIs running at once:

Could you check if you have more than 1 free worker slot, and that the pool default_pool size is bigger than 1 and has more than 1 free slot?

[VladimirYushkevich]

[VladimirYushkevich]

I guess my issue because of SequentialExecutor (I'm using SQLite as a backend).

[hussein-awala]

You have 3 queued TIs so everything is fine, and yes, with SequentialExecutor, you cannot have more than 1 running TI.
From the documentation:

This executor will only run one task instance at a time.