Configure Apache Airflow DAG performance optimization best practices

What this solves

Apache Airflow DAGs can become performance bottlenecks without proper optimization. This tutorial covers advanced techniques to maximize DAG throughput, reduce task latency, and scale workflows efficiently in production environments.

Prerequisites and environment setup

airflow version
airflow config get-value core executor

airflow config get-value core parallelism
airflow config get-value core dag_concurrency
airflow config get-value core max_active_runs_per_dag

DAG structure optimization and parallelism tuning

# Core parallelism settings
[core]
Maximum number of task instances running simultaneously
parallelism = 64

Maximum number of task instances per DAG
dag_concurrency = 16

Maximum number of active DAG runs per DAG
max_active_runs_per_dag = 4

Maximum number of active tasks per DAG run
max_active_tasks_per_dag = 16

from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime, timedelta

DAG-level optimization settings
default_args = {
    'owner': 'data-team',
    'depends_on_past': False,
    'start_date': datetime(2024, 1, 1),
    'email_on_failure': True,
    'email_on_retry': False,
    'retries': 2,
    'retry_delay': timedelta(minutes=5),
    'max_active_tis_per_dag': 8,
    'pool': 'high_priority_pool'
}

dag = DAG(
    'optimized_data_pipeline',
    default_args=default_args,
    description='Performance optimized DAG',
    schedule_interval='@hourly',
    catchup=False,
    max_active_runs=2,
    max_active_tasks=12,
    concurrency=8,
    tags=['production', 'optimized']
)

from airflow.decorators import task, dag
from airflow.utils.task_group import TaskGroup
from datetime import datetime
import concurrent.futures

@dag(
    schedule_interval='@daily',
    start_date=datetime(2024, 1, 1),
    catchup=False,
    max_active_runs=1,
    max_active_tasks=20
)
def parallel_processing_dag():
    
    @task
    def extract_data():
        # Data extraction logic
        return ['dataset1', 'dataset2', 'dataset3', 'dataset4']
    
    @task
    def process_dataset(dataset_name: str):
        # Process individual dataset
        print(f"Processing {dataset_name}")
        return f"processed_{dataset_name}"
    
    # Use dynamic task mapping for parallel processing
    with TaskGroup('parallel_processing') as processing_group:
        datasets = extract_data()
        processed = process_dataset.expand(dataset_name=datasets)
    
    @task
    def aggregate_results(processed_datasets):
        # Aggregate all processed datasets
        return f"Aggregated: {len(processed_datasets)} datasets"
    
    aggregate_results(processed)

parallel_dag = parallel_processing_dag()

# Create resource pools via Airflow CLI
airflow pools set high_priority_pool 8 "High priority tasks pool"
airflow pools set database_pool 4 "Database connection pool"
airflow pools set api_pool 6 "External API calls pool"
airflow pools set compute_intensive_pool 2 "CPU intensive tasks pool"

Task dependencies and resource allocation strategies

from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.utils.trigger_rule import TriggerRule
from datetime import datetime, timedelta

Optimized dependency structure
def create_optimized_dag():
    dag = DAG(
        'dependency_optimized_dag',
        start_date=datetime(2024, 1, 1),
        schedule_interval='@hourly',
        max_active_runs=2,
        catchup=False
    )
    
    # Parallel extraction tasks
    extract_users = PythonOperator(
        task_id='extract_users',
        python_callable=lambda: print("Extracting users"),
        pool='database_pool',
        dag=dag
    )
    
    extract_orders = PythonOperator(
        task_id='extract_orders',
        python_callable=lambda: print("Extracting orders"),
        pool='database_pool',
        dag=dag
    )
    
    extract_products = PythonOperator(
        task_id='extract_products',
        python_callable=lambda: print("Extracting products"),
        pool='database_pool',
        dag=dag
    )
    
    # Parallel transformation tasks
    transform_users = PythonOperator(
        task_id='transform_users',
        python_callable=lambda: print("Transforming users"),
        pool='compute_intensive_pool',
        dag=dag
    )
    
    transform_orders = PythonOperator(
        task_id='transform_orders',
        python_callable=lambda: print("Transforming orders"),
        pool='compute_intensive_pool',
        dag=dag
    )
    
    # Final aggregation with appropriate trigger rule
    aggregate_data = PythonOperator(
        task_id='aggregate_data',
        python_callable=lambda: print("Aggregating all data"),
        trigger_rule=TriggerRule.ALL_SUCCESS,
        pool='high_priority_pool',
        dag=dag
    )
    
    # Set up dependencies for maximum parallelism
    extract_users >> transform_users >> aggregate_data
    extract_orders >> transform_orders >> aggregate_data
    extract_products >> aggregate_data  # Direct dependency if no transformation needed
    
    return dag

from airflow.operators.python import PythonOperator
from airflow.providers.docker.operators.docker import DockerOperator
from airflow.configuration import conf

Task with resource constraints
heavy_computation_task = PythonOperator(
    task_id='heavy_computation',
    python_callable=compute_intensive_function,
    pool='compute_intensive_pool',
    pool_slots=2,  # Use 2 pool slots for this task
    priority_weight=10,  # Higher priority
    weight_rule='upstream',
    max_active_tis_per_dag=2,
    execution_timeout=timedelta(hours=2),
    dag=dag
)

Docker task with resource limits
docker_task = DockerOperator(
    task_id='docker_processing',
    image='python:3.11-slim',
    command='python /app/process_data.py',
    docker_url='unix://var/run/docker.sock',
    network_mode='bridge',
    pool='docker_pool',
    # Resource limits for Docker container
    mem_limit='2g',
    cpus=1.0,
    auto_remove=True,
    mount_tmp_dir=False,
    dag=dag
)

from airflow.operators.python import PythonOperator
from airflow.utils.dates import days_ago
from datetime import timedelta
import random

def task_with_smart_retries():
    # Simulate task that might fail
    if random.random() < 0.3:  # 30% chance of failure
        raise Exception("Temporary failure")
    return "Task completed successfully"

Task with optimized retry configuration
resilient_task = PythonOperator(
    task_id='resilient_processing',
    python_callable=task_with_smart_retries,
    retries=3,
    retry_delay=timedelta(minutes=2),
    retry_exponential_backoff=True,
    max_retry_delay=timedelta(minutes=10),
    execution_timeout=timedelta(minutes=30),
    # Don't retry on certain error types
    retry_exceptions=[ConnectionError, TimeoutError],
    dag=dag
)

Executor configuration and worker scaling

[core]
executor = CeleryExecutor

[celery]
Redis broker configuration
broker_url = redis://redis-cluster:6379/0
result_backend = redis://redis-cluster:6379/1

Worker configuration
worker_concurrency = 8
worker_prefetch_multiplier = 1
worker_max_tasks_per_child = 1000
worker_disable_rate_limits = True

Task routing
task_routes = {
    'airflow.executors.celery_executor.execute_command': {
        'queue': 'default',
        'routing_key': 'default'
    }
}

Queue configuration
default_queue = default
flower_host = 0.0.0.0
flower_port = 5555

#!/bin/bash

Start Celery worker with autoscaling
celery -A airflow.executors.celery_executor worker \
    --autoscale=16,4 \
    --queues=default,high_priority,compute_intensive \
    --hostname=worker-$(hostname)-%i \
    --loglevel=info \
    --max-tasks-per-child=1000 \
    --prefetch-multiplier=1 \
    --optimization=fair

from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime

Tasks routed to different queues based on resource requirements
dag = DAG(
    'queue_routing_dag',
    start_date=datetime(2024, 1, 1),
    schedule_interval='@hourly',
    catchup=False
)

Light processing task - default queue
light_task = PythonOperator(
    task_id='light_processing',
    python_callable=lambda: print("Light processing"),
    queue='default',
    pool='default_pool',
    dag=dag
)

CPU intensive task - dedicated queue
heavy_task = PythonOperator(
    task_id='heavy_processing',
    python_callable=lambda: print("Heavy processing"),
    queue='compute_intensive',
    pool='compute_intensive_pool',
    dag=dag
)

High priority task - priority queue
priority_task = PythonOperator(
    task_id='priority_processing',
    python_callable=lambda: print("Priority processing"),
    queue='high_priority',
    pool='high_priority_pool',
    priority_weight=100,
    dag=dag
)

[core]
executor = KubernetesExecutor

[kubernetes]
namespace = airflow
worker_container_repository = apache/airflow
worker_container_tag = 2.7.3
delete_worker_pods = True
delete_worker_pods_on_failure = False

Resource configuration for worker pods
worker_pods_creation_batch_size = 5
multi_namespace_mode = False
in_cluster = True
cluster_context = None
config_file = None

Worker pod resource limits
worker_container_image_pull_policy = IfNotPresent
worker_service_account_name = airflow-worker

Monitoring and performance troubleshooting

[metrics]
statsd_on = True
statsd_host = localhost
statsd_port = 8125
statsd_prefix = airflow

Enable detailed metrics
statsd_datadog_enabled = True
statsd_datadog_tags = env:production,service:airflow

Metrics to collect
statsd_allow_list = airflow.dag_processing.processes,
                   airflow.executor.open_slots,
                   airflow.executor.queued_tasks,
                   airflow.executor.running_tasks,
                   airflow.pool.open_slots,
                   airflow.pool.queued_slots,
                   airflow.pool.running_slots,
                   airflow.scheduler.tasks,
                   airflow.task_instance_duration

-- Query to identify slow-running tasks
SELECT 
    task_id,
    dag_id,
    AVG(EXTRACT(EPOCH FROM (end_date - start_date))) as avg_duration_seconds,
    COUNT(*) as execution_count
FROM task_instance 
WHERE start_date >= NOW() - INTERVAL '7 days'
    AND state = 'success'
GROUP BY task_id, dag_id
HAVING AVG(EXTRACT(EPOCH FROM (end_date - start_date))) > 300
ORDER BY avg_duration_seconds DESC;

-- Query to find frequently failing tasks
SELECT 
    task_id,
    dag_id,
    COUNT(*) as failure_count,
    MAX(end_date) as last_failure
FROM task_instance 
WHERE start_date >= NOW() - INTERVAL '7 days'
    AND state = 'failed'
GROUP BY task_id, dag_id
HAVING COUNT(*) > 5
ORDER BY failure_count DESC;

from airflow.plugins_manager import AirflowPlugin
from airflow.models import BaseOperator
from airflow.utils.decorators import apply_defaults
from airflow.hooks.base_hook import BaseHook
import requests
import json

class PerformanceMonitorOperator(BaseOperator):
    @apply_defaults
    def __init__(self, threshold_seconds=600, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.threshold_seconds = threshold_seconds
    
    def execute(self, context):
        # Check task duration
        task_instance = context['task_instance']
        duration = (task_instance.end_date - task_instance.start_date).total_seconds()
        
        if duration > self.threshold_seconds:
            self.send_alert({
                'task_id': task_instance.task_id,
                'dag_id': task_instance.dag_id,
                'duration': duration,
                'threshold': self.threshold_seconds,
                'execution_date': str(task_instance.execution_date)
            })
    
    def send_alert(self, alert_data):
        # Send alert to monitoring system
        webhook_url = "https://hooks.slack.com/services/YOUR/WEBHOOK/URL"
        
        message = {
            "text": f"⚠️ Performance Alert: Task {alert_data['task_id']} exceeded threshold",
            "blocks": [{
                "type": "section",
                "text": {
                    "type": "mrkdwn",
                    "text": f"Task: {alert_data['task_id']}\nDAG: {alert_data['dag_id']}\nDuration: {alert_data['duration']:.1f}s\nThreshold: {alert_data['threshold']}s"
                }
            }]
        }
        
        requests.post(webhook_url, json=message)

class PerformanceMonitoringPlugin(AirflowPlugin):
    name = "performance_monitoring"
    operators = [PerformanceMonitorOperator]

from airflow.decorators import task, dag
from datetime import datetime
import time
import psutil
import logging

logger = logging.getLogger(__name__)

@dag(
    schedule_interval='@daily',
    start_date=datetime(2024, 1, 1),
    catchup=False
)
def performance_profiled_dag():
    
    @task
    def profile_task_performance():
        start_time = time.time()
        start_memory = psutil.virtual_memory().used
        start_cpu = psutil.cpu_percent()
        
        # Your task logic here
        time.sleep(10)  # Simulate work
        
        end_time = time.time()
        end_memory = psutil.virtual_memory().used
        end_cpu = psutil.cpu_percent()
        
        performance_metrics = {
            'execution_time': end_time - start_time,
            'memory_used': end_memory - start_memory,
            'cpu_usage': (start_cpu + end_cpu) / 2,
            'timestamp': datetime.now().isoformat()
        }
        
        logger.info(f"Performance metrics: {performance_metrics}")
        return performance_metrics
    
    profile_task_performance()

performance_dag = performance_profiled_dag()

Advanced optimization techniques

from airflow.decorators import task, dag
from datetime import datetime, timedelta
import hashlib
import json
import redis

Redis client for caching
redis_client = redis.Redis(host='localhost', port=6379, db=2)

@dag(
    schedule_interval='@hourly',
    start_date=datetime(2024, 1, 1),
    catchup=False
)
def cached_processing_dag():
    
    @task
    def expensive_computation(input_data: str):
        # Create cache key from input
        cache_key = f"computation_{hashlib.md5(input_data.encode()).hexdigest()}"
        
        # Check cache first
        cached_result = redis_client.get(cache_key)
        if cached_result:
            print(f"Cache hit for {cache_key}")
            return json.loads(cached_result)
        
        # Perform expensive computation
        print(f"Computing result for {input_data}")
        result = {'processed': input_data, 'timestamp': datetime.now().isoformat()}
        
        # Cache result for 1 hour
        redis_client.setex(cache_key, 3600, json.dumps(result))
        return result
    
    expensive_computation("sample_input_data")

cached_dag = cached_processing_dag()

[database]
Connection pool settings
sql_alchemy_pool_enabled = True
sql_alchemy_pool_size = 10
sql_alchemy_max_overflow = 20
sql_alchemy_pool_recycle = 3600
sql_alchemy_pool_pre_ping = True
sql_alchemy_schema = airflow

Query optimization
sql_alchemy_connect_args = {
    'connect_timeout': 10,
    'application_name': 'airflow',
    'options': '-c statement_timeout=300000'
}

Enable connection pooling for better performance
load_default_connections = True

Verify your setup

Run these commands to validate your Airflow performance optimizations are working correctly.

# Check current configuration
airflow config get-value core parallelism
airflow config get-value core dag_concurrency
airflow config get-value core max_active_runs_per_dag

Verify executor configuration
airflow config get-value core executor

Check pool configuration
airflow pools list

Test DAG performance
airflow dags test optimized_data_pipeline 2024-01-01

Monitor task execution
airflow tasks list optimized_data_pipeline
airflow tasks states-for-dag-run optimized_data_pipeline 2024-01-01

Common issues

Next steps

• Set up Airflow high availability with CeleryExecutor clustering
• Implement comprehensive DAG monitoring with DataDog integration
• Configure Kubernetes custom metrics autoscaling for Airflow workers
• Implement DAG security policies and data governance
• Configure data lineage tracking with OpenLineage

Running this in production?