apache / apache/airflow

Pool Support for Task Groups in Apache Airflow

Open
#35,689 6 comments 5 reactions 0 assignees View on GitHub
area:scheduler area:TaskGroup kind:feature
Dominant language
Python
Stars
46.9k
Forks
17.8k
Avg merge
2d 7h
Merged PRs (30d)
484

Description

### Description

I would like to propose a new feature for Apache Airflow: the ability to assign a pool to an entire task group, enabling better control over the execution of tasks within the group. This feature would be particularly useful for ensuring that tasks within a group are executed sequentially before another group starts, especially when multiple groups of tasks are designed to run in parallel.

I propose adding the functionality to assign a pool directly to a task group, which would allow all tasks within the group to be managed by the pool's slot limitations. This feature would provide a much simpler and more efficient way to ensure that tasks within a group are executed in sequence before moving on to another group, especially in cases where parallel execution of multiple groups is required.

This enhancement would greatly improve the usability and flexibility of task groups in Airflow, especially for complex workflows where task sequencing and parallel execution control are crucial.

Thank you for considering this feature request.

### Use case/motivation

In my current workflow, I have multiple task groups structured to run in parallel, each containing a sequence of tasks that should be executed in a specific order. The task groups are similar in structure and purpose, resembling the following format:

Task Group 1: a >> b >> c
Task Group 2: d >> e >> f
Task Group 3: g >> h >> i
The desired behavior is to have two task groups running in parallel at any given time, with each group executing its tasks in the defined sequence. However, I'm encountering a challenge with the current Airflow features:

When using Concurrency: Airflow schedules the first task of each group (a, d, g) to run simultaneously, followed by the second tasks of each group (b, e, h), and so on. This approach does not maintain the sequential integrity of each group.

When using a Pool: Assigning all tasks in a group to a specific pool still results in the same behavior as above, where Airflow schedules one task from each group in parallel, rather than executing all tasks in a single group sequentially before starting the next group.

### Related issues

Current Workarounds and Limitations:

SubDAGs: One solution is to use SubDagOperator to encapsulate each task group into a SubDAG. While this approach ensures sequential execution within each group, it introduces performance and complexity issues, as SubDAGs are known to be less efficient and more complex to manage.

### Are you willing to submit a PR?

- [ ] Yes I am willing to submit a PR!

### Code of Conduct

- [X] I agree to follow this project's [Code of Conduct](https://github.com/apache/airflow/blob/main/CODE_OF_CONDUCT.md)

Contributor guide

Open the contributing guide

Research direction

Review the existing TaskGroup, Pool, and scheduler behavior described in the issue, comparing it with the SubDagOperator workaround. Define how assigning a pool to a whole task group should enforce per-group sequencing while allowing multiple groups to run in parallel; done means this behavior is implemented and covered by relevant scheduler tests.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
data-engineering
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.