Task groups for visual organization in Apache Airflow - Commands & Configuration

Introduction

When you have many tasks in an Airflow workflow, it can get messy and hard to understand. Task groups help by grouping related tasks together visually, making the workflow easier to read and manage.

When your workflow has many tasks that belong to the same step or process.

When you want to make your Airflow UI cleaner and easier to navigate.

When you want to reuse or organize tasks logically without changing their execution.

When you want to show a clear structure of your workflow to your team.

When you want to simplify complex dependencies by grouping tasks.

Config File - my_dag.py

my_dag.py

from airflow import DAG
from airflow.operators.bash import BashOperator
from airflow.utils.task_group import TaskGroup
from datetime import datetime

default_args = {
    'start_date': datetime(2024, 1, 1),
}

with DAG('example_task_group_dag', default_args=default_args, schedule_interval='@daily', catchup=False) as dag:
    start = BashOperator(
        task_id='start',
        bash_command='echo Start'
    )

    with TaskGroup('group1', tooltip='Tasks for step 1') as group1:
        task1 = BashOperator(
            task_id='task1',
            bash_command='echo Task 1'
        )
        task2 = BashOperator(
            task_id='task2',
            bash_command='echo Task 2'
        )

    with TaskGroup('group2', tooltip='Tasks for step 2') as group2:
        task3 = BashOperator(
            task_id='task3',
            bash_command='echo Task 3'
        )
        task4 = BashOperator(
            task_id='task4',
            bash_command='echo Task 4'
        )

    end = BashOperator(
        task_id='end',
        bash_command='echo End'
    )

    start >> group1 >> group2 >> end

This DAG defines a simple workflow with a start and end task.

Two TaskGroup blocks group related tasks: group1 contains task1 and task2, and group2 contains task3 and task4.

The groups help visually organize tasks in the Airflow UI, showing them as collapsible sections with tooltips.

The dependencies connect the start task to group1, then group2, then the end task, preserving execution order.

Commands

List all available DAGs to verify that the new DAG is recognized by Airflow.

Terminal

airflow dags list

Expected Output

example_task_group_dag

Trigger the DAG run manually to start the workflow and see the task groups in action.

Terminal

airflow dags trigger example_task_group_dag

Expected Output

Created <DagRun example_task_group_dag @ 2024-06-01 12:00:00: manual__2024-06-01T12:00:00+00:00, externally triggered: True>

List all tasks in the DAG, including those inside task groups, to confirm their presence.

Terminal

airflow tasks list example_task_group_dag

Expected Output

start

group1.task1

group1.task2

group2.task3

group2.task4

end

Test run a single task inside a task group to verify it runs correctly without executing the whole DAG.

Terminal

airflow tasks test example_task_group_dag group1.task1 2024-06-01

Expected Output

[2024-06-01 00:00:00,000] {bash.py:123} INFO - Running command: echo Task 1

[2024-06-01 00:00:00,001] {bash.py:130} INFO - Output: Task 1

[2024-06-01 00:00:00,002] {bash.py:135} INFO - Command exited with return code 0

Key Concept

If you remember nothing else from this pattern, remember: Task groups help you organize and visually simplify complex Airflow workflows without changing task execution.

Common Mistakes

Using task groups but not prefixing task IDs inside groups.

Task IDs inside groups must be unique and are prefixed automatically; forgetting this can cause confusion or conflicts.

Always refer to tasks inside groups with the full ID including the group prefix, like 'group1.task1'.

Trying to set dependencies directly between tasks in different groups without using full task IDs.

Dependencies must use full task IDs including group prefixes to work correctly.

Use full task IDs like 'group1.task2 >> group2.task3' to set dependencies between tasks in different groups.

Summary

Define TaskGroup blocks in your DAG to group related tasks visually.

Use full task IDs with group prefixes to manage dependencies and task references.

Trigger and test your DAG to see task groups as collapsible sections in the Airflow UI.

Practice

1. What is the main purpose of using TaskGroup in Airflow DAGs?

easy

A. To visually group related tasks for better organization

B. To execute tasks in parallel automatically

C. To replace the need for defining individual tasks

D. To schedule DAG runs at specific times

Solution

1. Step 1: Understand TaskGroup purpose

   TaskGroup is designed to group related tasks visually in the Airflow UI.

2. Step 2: Differentiate from other features

   It does not change execution behavior or scheduling but improves readability.

3. Final Answer:

   To visually group related tasks for better organization -> Option A

4. Quick Check:

   TaskGroup = Visual grouping [OK]

Hint: TaskGroup groups tasks visually, not for execution control [OK]

Common Mistakes:

• Thinking TaskGroup runs tasks in parallel
• Confusing TaskGroup with scheduling
• Assuming TaskGroup replaces tasks

2. Which of the following is the correct syntax to create a task group named 'processing' in Airflow?

easy

A. with TaskGroup('processing') as tg:

B. TaskGroup('processing'):

C. group = TaskGroup('processing')

D. with task_group('processing') as tg:

Solution

1. Step 1: Recall TaskGroup syntax

   The correct way to create a task group is using a context manager: with TaskGroup('name') as variable:

2. Step 2: Check option correctness

   with TaskGroup('processing') as tg: uses correct capitalization and syntax. TaskGroup('processing'): misses 'with'. group = TaskGroup('processing') misses 'with' and context manager. with task_group('processing') as tg: uses wrong function name casing.

3. Final Answer:

   with TaskGroup('processing') as tg: -> Option A

4. Quick Check:

   Use 'with TaskGroup(...) as ...:' [OK]

Hint: Use 'with TaskGroup("name") as var:' to create groups [OK]

Common Mistakes:

• Omitting 'with' keyword
• Using wrong capitalization like 'task_group'
• Not using 'as' to assign the group

3. Given the code snippet:

with TaskGroup('group1') as tg1:
    task_a = DummyOperator(task_id='a')
    task_b = DummyOperator(task_id='b')
task_c = DummyOperator(task_id='c')
task_a >> task_c

What will be the visual relationship between tasks in the Airflow UI?

medium

A. Tasks 'a', 'b', and 'c' are all grouped together

B. Tasks 'a' and 'b' are grouped; 'a' runs before 'c'; 'b' is independent

C. Task 'c' runs before 'a' and 'b' inside the group

D. Tasks 'a' and 'b' run after 'c'

Solution

1. Step 1: Analyze TaskGroup and tasks

   Tasks 'a' and 'b' are inside TaskGroup 'group1', so they appear grouped visually. Task 'c' is outside the group.

2. Step 2: Check dependencies

   Only 'task_a' has a downstream dependency to 'task_c', so 'a' runs before 'c'. 'b' has no dependencies and is independent.

3. Final Answer:

   Tasks 'a' and 'b' are grouped; 'a' runs before 'c'; 'b' is independent -> Option B

4. Quick Check:

   Group contains a,b; a >> c; b independent [OK]

Hint: Only tasks inside TaskGroup are grouped visually [OK]

Common Mistakes:

• Assuming all tasks are grouped
• Thinking 'b' depends on 'c'
• Confusing execution order

4. Identify the error in this Airflow DAG snippet using TaskGroup:

with TaskGroup('group1'):
    task1 = DummyOperator(task_id='task1')
task2 = DummyOperator(task_id='task2')
tg >> task2

medium

A. DummyOperator cannot be used inside TaskGroup

B. TaskGroup name must be uppercase

C. Task dependencies cannot cross TaskGroup boundaries

D. Missing 'as' keyword to assign TaskGroup to a variable

Solution

1. Step 1: Check TaskGroup usage

   The 'with TaskGroup('group1'):' statement should assign the group to a variable using 'as', e.g. 'with TaskGroup('group1') as tg:'

2. Step 2: Validate other options

   TaskGroup names can be lowercase. DummyOperator is allowed inside TaskGroup. Dependencies can cross groups.

3. Final Answer:

   Missing 'as' keyword to assign TaskGroup to a variable -> Option D

4. Quick Check:

   Use 'with TaskGroup(...) as var:' [OK]

Hint: Always use 'as' to assign TaskGroup in 'with' statement [OK]

Common Mistakes:

• Omitting 'as' in 'with TaskGroup(...) as ...:'
• Thinking TaskGroup names must be uppercase
• Believing DummyOperator is disallowed inside groups

5. You want to create a DAG with two task groups: 'extract' and 'transform'. Tasks in 'extract' must run before any in 'transform'. Which code snippet correctly sets this up?

hard

A. with TaskGroup('extract') as extract:
t1 = DummyOperator(task_id='t1')
t1b = DummyOperator(task_id='t1b')
with TaskGroup('transform') as transform:
t2 = DummyOperator(task_id='t2')
t2b = DummyOperator(task_id='t2b')
t1 >> t2

B. extract = TaskGroup('extract')
transform = TaskGroup('transform')
extract >> transform

C. with TaskGroup('extract') as extract:
t1 = DummyOperator(task_id='t1')
t1b = DummyOperator(task_id='t1b')
with TaskGroup('transform') as transform:
t2 = DummyOperator(task_id='t2')
t2b = DummyOperator(task_id='t2b')
extract >> transform

D. with TaskGroup('extract'):
t1 = DummyOperator(task_id='t1')
t1b = DummyOperator(task_id='t1b')
with TaskGroup('transform'):
t2 = DummyOperator(task_id='t2')
t2b = DummyOperator(task_id='t2b')
t1 >> t2

Solution

1. Step 1: Assign TaskGroups to variables

   Use 'with TaskGroup(...) as var:' to create groups and assign them to variables for referencing in dependencies.

2. Step 2: Set group-level dependencies

   Use extract >> transform to ensure every task in 'extract' (t1, t1b) completes before any task in 'transform' (t2, t2b) starts.

3. Step 3: Check options

   with TaskGroup('extract') as extract: t1 = ... t1b... with TaskGroup('transform') as transform: t2 = ... t2b... t1 >> t2 only links t1 to t2; t1b independent of transform tasks. extract = TaskGroup('extract') transform = TaskGroup('transform') extract >> transform misses 'with' context; no tasks added to groups. with TaskGroup('extract'): t1 = ... t1b... with TaskGroup('transform'): t2 = ... t2b... t1 >> t2 misses group variables and only links specific tasks. with TaskGroup('extract') as extract: t1 = ... t1b... with TaskGroup('transform') as transform: t2 = ... t2b... extract >> transform correctly uses group dependency.

4. Final Answer:

   with TaskGroup('extract') as extract: t1 = DummyOperator(task_id='t1') t1b = DummyOperator(task_id='t1b') with TaskGroup('transform') as transform: t2 = DummyOperator(task_id='t2') t2b = DummyOperator(task_id='t2b') extract >> transform -> Option C

5. Quick Check:

   tg1 >> tg2 for complete group ordering [OK]

Hint: Assign 'with TaskGroup(name) as tg:' then tg1 >> tg2 [OK]

Common Mistakes:

• Using individual task deps (t1 >> t2) with multiple tasks per group
• Missing 'with' context manager
• Omitting 'as var' when using group dependencies
• No tasks added without context manager