How to control concurrency of a subdag or taskgroup in Airflow?

Viewed 1178

I have a simple ETL workflow composed of three tasks for each table I want to extract :

1_extract_to_tmp >> 2_push_to_s3 >> 3_delete_tmp

As I want to reproduce the same steps for multiple tables, I was thinking of grouping these task in a TaskGroup or a subDAG and instanciate it dynamically for each table to extract. My final DAG will look like this : My DAG

From what I read, the TaskGroup is the prefered solution now to do that.

The problem is that I also need to control how much extractions are done in parallel because I don't have enough local disk space for all data and apparently the subDag operator does not comply to pool setting and the TaskGroup does not have one.

Do you know a way to achieve this ?

Am I doing it wrong ?

Maybe my DAG should not be designed like that.

1 Answers

TaskGroup is just a UI feature. It doesn't really contains logic so any parallelism limitation you want to enforce can be done regardless to the usage of TaskGroup.

You didn't explain what exactly you want to limit. If you want to limit the overall tasks that can run in parallel with on your dag (overwrite the airflow.cfg default) then set concurrency in your DAG contractor:

dag = DAG(dag_id='my_dag', concurrency=5, ...)

If you are looking to limit concurrency usage to protect a specific resources (database, api, etc...) then use pools:

my_op = PythonOperator(python_callable=func,
                          task_id='my_task',
                          pool='my_pool',
                          dag=dag)

The task will be executed only if there are free slots in the pool. The number of available slots are defined when you create the pool.

As for SubDags I wouldn't recommand using it at all. While it's not officially deprecated it probably will be. See https://github.com/apache/airflow/issues/12292

Related