airflow-manual/09 - Пулы, XCom, TaskGroup и алертинг в Airflow.md

Пулы, XCom, TaskGroup и алертинг в Airflow

В этом материале мы познакомимся с управлением ресурсами через пулы задач, обменом данными между задачами через XCom, группировкой задач с помощью TaskGroup и настройкой системы оповещений (алертинга) в Airflow.

Начнем с улучшения нашего базового пайплайна — проведем рефакторинг кода для лучшей читаемости и поддержки.

Управление ресурсами с помощью пулов задач

В системах с высокой нагрузкой, где одновременно запускается множество задач и DAG-ов, может возникнуть чрезмерная нагрузка на исполнителей и серверную часть. Это может привести к ошибкам выполнения и даже к отказу системы, если не установить соответствующие ограничения.

В Airflow для решения этой проблемы существует механизм управления ресурсами — пулы задач (pools). По умолчанию в Airflow настроен один пул задач — default_pool с 128 слотами. При стандартном pool_slots=1 это верхняя граница в 128 одновременно работающих задач; ограничения executor и DAG могут уменьшить ее. Пул default_pool нельзя удалить, но можно изменить его размер — увеличить или уменьшить количество слотов.

Когда планировщик обнаруживает, что наступило время выполнения DAG, он запускает задачу согласно заданной последовательности. По умолчанию задача занимает один слот в пуле и освобождает его после завершения.

Создать новый пул задач и установить его размер можно через веб-интерфейс Airflow. Рассмотрим пример пула data_processing_pool для тяжёлых задач:

1. В верхнем меню откройте Admin → Pools.
2. Нажмите кнопку + или Create.
3. В поле Name укажите, например, data_processing_pool.
4. В поле Slots задайте количество слотов, например 5
   (это значит, что одновременно смогут выполняться не более 5 задач из этого пула).
5. При желании заполните Description — например, Пул для тяжёлых задач бэкапа.
6. Нажмите Save.

После этого в задачах можно указать параметр pool="data_processing_pool", и они будут занимать слоты именно этого пула.

Зачем нужны пулы задач? Они помогают:

• Организовать запуск процессов в системе
• Предотвратить перегрузку системы при выполнении большого количества ресурсоемких задач

Вы можете задать "вес" задачи через параметр pool_slots, чтобы оптимизировать распределение нагрузки. Планировщик учитывает сумму занятых слотов. Например, в пуле из двух слотов задачи с pool_slots=2 и pool_slots=1 не смогут работать одновременно: одна ждет завершения другой.

Пример настройки веса задач:

В приведенном примере показано, как можно настроить использование пула задач с различным весом. Задача 'data_backup_task' использует 3 слота в пуле 'data_processing_pool', что делает ее более ресурсоемкой по сравнению с задачами 'file_check_task' и 'cleanup_task', которые используют по 1 слоту. Это позволяет контролировать распределение ресурсов между различными задачами.

BashOperator(
    task_id="data_backup_task",
    bash_command="bash backup_script.sh",
    pool_slots=3,
    pool="data_processing_pool",
)

BashOperator(
    task_id="file_check_task",
    bash_command="bash validate_files.sh",
    pool_slots=1,
    pool="data_processing_pool",
)

BashOperator(
    task_id="cleanup_task",
    bash_command="bash cleanup_files.sh",
    pool_slots=1,
    pool="data_processing_pool",
)

Более подробную информацию о механизме пулов можно найти в официальной документации.

Обмен данными между задачами: XCom и контекст выполнения

XCom (от англ. cross-communications — «межзадачная коммуникация») — это механизм обмена сообщениями между задачами внутри одного DAG. Архитектурно каждая задача изолирована от других и работает в собственном контексте.

💡 Контекст выполнения — это набор параметров, передаваемых при запуске задачи, а также метаданные, генерируемые во время выполнения: время старта, время завершения, имя DAG и другие. Контекст можно найти следующим образом:

1. Откройте веб-интерфейс Airflow и перейдите на вкладку DAGs.
2. Найдите нужный DAG и кликните по его имени.
3. На странице DAG убедитесь, что открыта вкладка Grid.
4. В сетке выберите нужный запуск DAG (колонка) и задачу (строка) и кликните по цветному квадратику задачи.
5. Справа откроется панель Task Instance (детали экземпляра задачи). В ней можно увидеть:
   • dag_id и task_id;
   • логическую дату запуска (logical_date / execution_date);
   • текущий статус, время старта и завершения;
   • ссылки на лог, XCom и другую служебную информацию.

Именно эти поля и составляют большую часть «контекста выполнения», который доступен в Jinja-шаблонах через объекты вроде {{ dag_run }} и {{ task_instance }}.

В XCom можно передавать сериализованные объекты. Значения XCom хранятся в базе данных Airflow и доступны через интерфейс.

Посмотреть XCom можно двумя способами.

1. Через конкретную задачу в DAG

1. Откройте нужный DAG и вкладку Grid.
2. Найдите нужный запуск и кликните по квадратику задачи.
3. В правой панели Task Instance перейдите на вкладку XCom.
4. В таблице вы увидите все XCom-записи для этого экземпляра задачи: ключ (key), значение (value), время создания и т. д.

2. Через общий список XCom

1. В верхнем меню выберите Browse → XComs.
2. Отфильтруйте записи по dag_id, task_id или другим полям, если нужно.
3. Откройте интересующую запись, чтобы увидеть её содержимое.

Важное правило: XCom предназначен для обмена небольшими сообщениями. Данные проходят сериализацию/десериализацию при чтении и записи в таблицу.

💡 Сериализация — процесс преобразования структуры данных в последовательность байтов. Десериализация — восстановление структуры данных из байтовой последовательности.

Для передачи больших объемов данных используйте внешние средства: файловую систему, базы данных (чаще всего PostgreSQL) или очереди сообщений (например, Kafka).

Механизм XCom похож на работу функций в Python. Многие операторы (например, PythonOperator) по умолчанию возвращают результат выполнения задачи. За это отвечает параметр do_xcom_push, который во многих случаях равен True по умолчанию.

Чтобы прочитать сообщения из XCom, используйте метод xcom_pull в контексте задачи:

В этом примере мы получаем результат выполнения задачи с идентификатором 'data_processing_task' с помощью метода xcom_pull. Это позволяет передавать небольшие объемы данных между задачами в рамках одного DAG.

result = task_instance.xcom_pull(task_ids='data_processing_task')

Также можно обращаться к сообщениям через Jinja-шаблоны:

SELECT * FROM {{ task_instance.xcom_pull(task_ids='foo', key='table_name') }}

Если оператор возвращает значение и параметр do_xcom_push установлен в True (по умолчанию), это значение автоматически записывается в XCom.

Пример явного запрета записи в XCom:

В этом примере задача 'show_directory_contents' создается с параметром do_xcom_push=False, что означает, что результат выполнения этой задачи не будет автоматически сохранен в XCom. Это полезно, когда вы не хотите, чтобы задача передавала какие-либо данные другим задачам через XCom.

list_files = BashOperator(
    task_id='show_directory_contents',
    bash_command='ls -la',
    do_xcom_push=False
)

def calculate_sum():
    return 2 + 3

Пример автоматической записи в XCom (параметр do_xcom_push по умолчанию True):

В этом примере задача 'calculate_sum' автоматически записывает результат выполнения функции calculate_sum в XCom, так как параметр do_xcom_push по умолчанию установлен в True. Это позволяет использовать результат этой задачи в других задачах DAG через XCom.

sum_result = PythonOperator(
    task_id='calculate_sum',
    python_callable=calculate_sum,
)

XCom похожи на переменные (variables) в Airflow, но предназначены именно для взаимодействия между задачами в рамках одного DAG, а не для глобальных настроек.

XCom упрощает взаимодействие между задачами и применяется в различных сценариях.

Группировка задач с помощью TaskGroup

Для удобства визуализации задачи можно группировать в веб-интерфейсе Airflow (начиная с версии 2.0). Повторяющиеся или логически связанные задачи можно объединить в группы:

В этом примере задачи сгруппированы в три логические группы: 'data_extraction', 'data_transformation' и 'data_loading'. Каждая группа содержит несколько задач, которые выполняются последовательно внутри группы. Затем группы связаны между собой, чтобы показать общий порядок выполнения этапов обработки данных.

with TaskGroup("data_extraction") as extraction_group:
    extract_1 = DummyOperator(task_id="extract_source_1")
    extract_2 = DummyOperator(task_id="extract_source_2")
    extract_3 = DummyOperator(task_id="extract_source_3")
    extract_1 >> extract_2 >> extract_3

with TaskGroup("data_transformation") as transformation_group:
    transform_1 = DummyOperator(task_id="transform_step_1")
    transform_2 = DummyOperator(task_id="transform_step_2")
    transform_1 >> transform_2

with TaskGroup("data_loading") as loading_group:
    load_1 = DummyOperator(task_id="load_to_target_1")
    load_2 = DummyOperator(task_id="load_to_target_2")
    load_1 >> load_2

[extraction_group, transformation_group] >> loading_group

Обратите внимание: при использовании TaskGroup последовательность задач указывается внутри группы после объявления всех задач, а в конце DAG описывается последовательность выполнения самих групп.

Ниже приведена упрощённая схема зависимостей между тремя группами задач:

flowchart LR

  %% group1
  subgraph G1["group1"]
    g1_t1["task1"]
    g1_t2["task2"]
    g1_t3["task3"]

    g1_t1 --> g1_t2
    g1_t1 --> g1_t3
  end

  %% group2
  subgraph G2["group2"]
    g2_t1["task1"]
    g2_t2["task2"]

    g2_t1 --> g2_t2
  end

  %% group3
  subgraph G3["group3"]
    g3_t1["task1"]
    g3_t2["task2"]

    g3_t1 --> g3_t2
  end

  %% зависимости между группами
  g1_t2 --> g3_t1
  g1_t3 --> g3_t1
  g2_t2 --> g3_t1

Визуально в интерфейсе Airflow группы задач отображаются как один узел с небольшим индикатором. Клик по нему разворачивает или сворачивает вложенные задачи, что значительно улучшает восприятие DAG с большим количеством задач и связей, особенно когда в них десятки и сотни задач.

TaskGroup — это удобный способ логической группировки задач, который помогает упростить код и представить сложные пайплайны более компактно.

Система оповещений (алертинг)

Для первого опыта достаточно callback с записью в лог. Он получает контекст задачи и помогает связать сообщение с конкретным запуском:

import logging


def notify_failure(context):
    ti = context["ti"]
    logging.error(
        "Ошибка: dag=%s task=%s run=%s причина=%s",
        ti.dag_id, ti.task_id, context["run_id"], context["exception"],
    )

Укажите on_failure_callback=notify_failure в рабочем операторе. Функция вызывается после окончательного отказа, когда повторы исчерпаны. Для сообщения о предстоящем повторе существует on_retry_callback. Например, при retries=2 задача может выполниться три раза: исходная попытка и два повтора.

Callback запускается при реальном выполнении задачи. Ручная смена статуса в UI его не проверяет. Смотрите лог самой задачи и логи scheduler; ошибки callback ищите в логах scheduler. Документация Airflow 2.9.2 о callbacks.

Отправка почты по желанию

В учебном стенде SMTP не настроен. Для отправки нужен SMTP-сервер и адрес получателя, которому можно отправлять тестовые сообщения. Укажите настройки в общем окружении сервисов Airflow: AIRFLOW__SMTP__SMTP_HOST, AIRFLOW__SMTP__SMTP_PORT, AIRFLOW__SMTP__SMTP_MAIL_FROM, параметры TLS/SSL и учетные данные вашего сервера. Секреты храните локально, вне Git. После изменения окружения пересоздайте сервисы командой docker compose up -d. Полный список параметров находится в конфигурации Airflow 2.9.2.

После настройки замените запись в лог отправкой сообщения:

from html import escape

from airflow.utils.email import send_email_smtp

MAIL_TO = ["student@example.com"]  # Замените своим тестовым адресом.


def notify_failure(context):
    ti = context["ti"]
    send_email_smtp(
        to=MAIL_TO,
        subject=f"Ошибка {ti.dag_id}: {ti.task_id}",
        html_content=(
            f"<p>Запуск: {escape(context['run_id'])}</p>"
            f"<p>Причина: {escape(str(context['exception']))}</p>"
            f'<p><a href="{escape(ti.log_url, quote=True)}">Лог задачи</a></p>'
        ),
    )


def notify_success(context):
    send_email_smtp(
        to=MAIL_TO,
        subject=f"DAG {context['dag'].dag_id} завершен успешно",
        html_content=f"<p>Запуск: {escape(context['run_id'])}</p>",
    )

Для примера с ошибками задайте on_failure_callback=notify_failure у двух рабочих задач, а on_success_callback=notify_success у самого DAG(...). Если упадут обе задачи, придут два сообщения об отказе. Сообщение об успехе отправится только после успешного DAG Run. Не включайте одновременно email_on_failure=True, если не хотите получать еще и встроенное письмо о той же ошибке.

Не добавляйте ради уведомления единственную завершающую задачу с trigger_rule="all_done": если она успешна, Airflow может признать весь запуск успешным, несмотря на отказ внутри графа. Callback не меняет граф и итог обработки. Ошибка отправки означает проблему уведомления; она не превращает успешно обработанные данные в неуспешные и не скрывает исходный отказ. Подробности определения результата - в описании DAG Run.

Проверка: сначала воспроизведите окончательный отказ, затем успешный запуск. Сверьте текст каждого письма с состоянием соответствующего DAG Run. Упражнения находятся в блоке ошибок и задании на почту.

Практическое применение: улучшенный DAG

Теперь применим изученные концепции для усовершенствования нашего DAG. Добавим переменные и разобьем задачи на логические группы.

Сначала создадим переменную DATABASE_URL со строкой подключения к базе данных через веб-интерфейс Airflow и импортируем ее в коде DAG:

В этом примере используется переменная 'database_connection_string', предварительно созданная в интерфейсе Airflow, для хранения строки подключения к базе данных. Это позволяет избежать жесткого кодирования конфиденциальной информации в коде DAG и упрощает настройку подключения для разных окружений.

import os
import datetime as dt
import pandas as pd
from airflow.models import DAG
from airflow.operators.bash import BashOperator
from airflow.operators.python import PythonOperator
from airflow.operators.dummy import DummyOperator
from airflow.utils.task_group import TaskGroup
from airflow.models import Variable
from sqlalchemy import create_engine

DATABASE_URL = Variable.get('database_connection_string')

args = {
    'owner': 'analytics_team',
    'start_date': dt.datetime(2021, 6, 15),
    'retries': 2,
    'retry_delay': dt.timedelta(minutes=2),
}

# функции для обработки данных
def get_file_path(file_name):
    return os.path.join(os.path.expanduser('~/data'), file_name)

def load_customer_data():
    file_path = get_file_path('customer_data.csv')
    df = pd.read_csv(file_path)
    engine = create_engine(DATABASE_URL)
    df.to_sql('customers', engine, index=False, if_exists='replace', schema='staging')

def aggregate_customer_data():
    engine = create_engine(DATABASE_URL)
    customer_df = pd.read_sql('select * from staging.customers', con=engine)
    
    df = customer_df.groupby(['region', 'category']).agg(
            total_orders=('orders', 'sum'),
            avg_amount=('amount', 'mean')
        ).reset_index()
    
    df.to_sql('customer_summary', engine, index=False, if_exists='replace', schema='analytics')

dag = DAG(
    dag_id='customer_pipeline_enhanced',
    schedule_interval=None,
    default_args=args,
)

Теперь разобьем задачи на логические группы и добавим Jinja-шаблоны для доступа к контексту выполнения:

В этом примере создается начальная задача 'pipeline_start', которая использует Jinja-шаблоны для вывода информации о запуске DAG, включая идентификатор запуска (run_id) и информацию о DAG Run. Затем задачи группируются в логическую группу 'data_processing_stage', что улучшает структуру и читаемость DAG.

# Начальная задача с информацией о запуске
start_task = BashOperator(
    task_id='pipeline_start',
    bash_command='echo "Pipeline started! Run ID: {{ run_id }} | DAG Run: {{ dag_run }}"',
    dag=dag,
)

# Группа задач по предварительной обработке данных
with TaskGroup(group_id="data_processing_stage") as data_processing:
    # Загрузка данных
    load_customer_dataset = PythonOperator(
        task_id='load_customer_data',
        python_callable=load_customer_data,
        dag=dag,
    )
    # Агрегация и запись данных
    aggregate_customer_dataset = PythonOperator(
        task_id='aggregate_customer_data',
        python_callable=aggregate_customer_data,
        dag=dag,
    )
    load_customer_dataset >> aggregate_customer_dataset

# Установка последовательности выполнения
start_task >> data_processing

Поскольку последовательность задач внутри групп указывается при их создании, в конце необходимо определить порядок выполнения самих групп, чтобы планировщик понимал общую логику выполнения.

Итоговый DAG будет выглядеть следующим образом:

import os
import datetime as dt
import pandas as pd
from airflow.models import DAG
from airflow.operators.bash import BashOperator
from airflow.operators.python import PythonOperator
from airflow.operators.dummy import DummyOperator
from airflow.utils.task_group import TaskGroup
from airflow.models import Variable
from sqlalchemy import create_engine

DATABASE_URL = Variable.get('database_connection_string')

args = {
    'owner': 'analytics_team',
    'start_date': dt.datetime(2021, 6, 15),
    'retries': 2,
    'retry_delay': dt.timedelta(minutes=2),
}

def get_file_path(file_name):
    return os.path.join(os.path.expanduser('~/data'), file_name)

def load_customer_data():
    file_path = get_file_path('customer_data.csv')
    df = pd.read_csv(file_path)
    engine = create_engine(DATABASE_URL)
    df.to_sql('customers', engine, index=False, if_exists='replace', schema='staging')

def aggregate_customer_data():
    engine = create_engine(DATABASE_URL)
    customer_df = pd.read_sql('select * from staging.customers', con=engine)
    
    df = customer_df.groupby(['region', 'category']).agg(
            total_orders=('orders', 'sum'),
            avg_amount=('amount', 'mean')
        ).reset_index()
    
    df.to_sql('customer_summary', engine, index=False, if_exists='replace', schema='analytics')

dag = DAG(
    dag_id='customer_pipeline_enhanced',
    schedule_interval=None,
    default_args=args,
)

# Начальная задача
start_task = BashOperator(
    task_id='pipeline_start',
    bash_command='echo "Pipeline started! Run ID: {{ run_id }} | DAG Run: {{ dag_run }}"',
    dag=dag,
)

# Группа предварительной обработки
with TaskGroup(group_id="data_processing_stage") as data_processing:
    load_customer_dataset = PythonOperator(
        task_id='load_customer_data',
        python_callable=load_customer_data,
        dag=dag,
    )
    aggregate_customer_dataset = PythonOperator(
        task_id='aggregate_customer_data',
        python_callable=aggregate_customer_data,
        dag=dag,
    )
    load_customer_dataset >> aggregate_customer_dataset

start_task >> data_processing

В этом материале мы рассмотрели расширенные возможности Airflow, которые помогут улучшить работу ваших пайплайнов:

• Управление ресурсами с помощью пулов задач
• Обмен данными между задачами через XCom
• Логическая группировка задач с TaskGroup
• Настройка системы оповещений

Помните: не стоит использовать все доступные функции сразу. Выбирайте инструменты последовательно и находите оптимальный набор возможностей под конкретную задачу.