Apache Airflow и XCom

XCom или Cross-Communication, это механизм Apache Airflow для передачи параметров из одного оператора в другой. Если просто, то это таблица в базе данных, хранящая значения, записанные операторами Airflow. У этой таблицы есть несколько столбцов:

  • key — ключ записи
  • value — значение, данные хранятся в двоичном формате. Для MySQL это тип BLOB, а для Postgres BYTEA.
  • timestamp — время создания записи в базе
  • execution_date — дата выполнения DAG
  • task_id — ID оператора, записавшего данные в XCom
  • dag_id — ID DAG

Несмотря на то, что BYTEA в Postgres может хранить до 1 Гб двоичных данных, сообщество не рекомендует передавать через XCom большие данные. Это связано с тем, что перед укладкой в таблицу они подвергаются сериализации (JSON/pickle), а при чтении происходит этап десериализации. Также не стоит забывать, что данные передаются по сети, а это значительно замедляет работу пайплайна. Цитата из официальной доки по PostgreSQL:

The BYTEA data type is not well suited for storing very large amounts of binary data. While a column of type BYTEA can hold up to 1 GB of binary data, it would require a huge amount of memory to process such a large value.

Зачем всё так сложно? Почему нельзя передавать аргументы как это происходит при написании обычного кода на Python?

Главная причина в том, что операторы могут исполняться в разных адресных пространствах (Local Executor) и даже на отдельных физических машинах (Celery/Dask Executors, Kubernetes Executor). В такой ситуации нужен механизм для передачи сообщений от одной машины к другой. Центральным хранилищем в этом случае является база данных Airflow (рекомендуется использовать PostgreSQL).

Как работать с XCom

По умолчанию если оператор в методе .execute возвращает значение, это значение автоматически попадает в XCom. Это поведение можно изменить, передав в аргумент do_xcom_push значение False:

from airflow.operators.bash import BashOperator

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

Если callable, переданный в аргументе python_callable у PythonOperator возвращает какое-либо значение, то это значение также автоматически будет записано в XCom.

Пример кода:

from airflow.operators.python import PythonOperator

def return_something():
    return 1 + 1  # автоматически запишется в XCom

PythonOperator(
    task_id='return_something',
    python_callable=return_something,
)

Помимо этого с XCom можно работать вручную, записывая или получая данные путём вызова методов xcom_push и xcom_pull у производных от TaskInstance (кстати, тоже таблица в БД).

TaskInstance доступен в контексте выполнения DAG, а также как переменная в шаблоне под именем ti или task_instance. Например, если вам необходимо прочитать значение предыдущего таска внутри функции PythonOperator, то необходимо из контекста получить TaskInstance и вызвать метод xcom_pull, передав ему название оператора, который уложил результат в XCom:

from datetime import datetime

from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.operators.bash import BashOperator

def print_xcom(**kwargs):
    ti = kwargs['ti']
    print('The value is: {}'.format(
        ti.xcom_pull(task_ids='hello_world')
    ))

with DAG(
    dag_id='xcom_dag',
    start_date=datetime(2021, 3, 1),
    schedule_interval='@once',
) as dag:

    cmd = BashOperator(
        task_id='hello_world',
        bash_command='echo "Hello world"',
    )

    printer = PythonOperator(
        task_id='printer',
        python_callable=print_xcom,
    )

    cmd >> printer

Все переменные из контекста DagRun передаются в python_callable как именованные аргументы. Также контекст внутри функции можно получить путём вызова функции get_current_context, это удобно при использовании TaskFlow API:

from airflow.operators.python import get_current_context

@task
def download_file():
    context = get_current_context()
    return download_dataset(context['execution_date'].strftime('%Y-%m'))

Может возникнуть вопрос, а что если одновременно выполняются несколько DagRun, не получится ли так, что мы прочитаем "чужое" значение в XCom? Нет. Записи привязываются к dag_id, task_id и execution_date. У, скажем, 3-х одновременно выполняющихся DagRun будут разные значения execution_date.

Просмотр XCom доступен через веб-интерфейс Apache Airflow в разделе Admin → XComs:



На скрине пример запуска DAG чей код был выше.

Обратите внимание на колонку Key, в неё по умолчанию записывается значение return_value, что, видимо, является индикатором того, что оператор в методе execute вернул значение. При самостоятельной записи в XCom можно указать свой key:

from datetime import datetime

from airflow import DAG
from airflow.operators.python import PythonOperator
from airflow.operators.bash import BashOperator

def push_cmd(**kwargs):
    ti = kwargs['ti']
    ti.xcom_push(value='cat /etc/passwd | wc -l', key='passwd_len')

def pull_result(**kwargs):
    ti = kwargs['ti']
    print(ti.xcom_pull(task_ids='bash_executor'))

with DAG(
    dag_id='xcom_pull_dag',
    start_date=datetime(2021, 3, 1),
    schedule_interval='@once',
) as dag:

    push_cmd_task = PythonOperator(
        task_id='push_cmd',
        python_callable=push_cmd,
    )

    cmd = BashOperator(
        task_id='bash_executor',
        bash_command='{{ ti.xcom_pull(task_ids="push_cmd", key="passwd_len") }}',
    )

    pull_result_task = PythonOperator(
        task_id='pull_result',
        python_callable=pull_result,
    )

    push_cmd_task >> cmd >> pull_result_task

Также обратите внимание, что метод xcom_pull может принимать список в аргументе task_ids. Туда можно передать несколько task_id по которым хочется получить значения. Взгляните на следующий граф:



Каждый push таск укладывает в XCom значение, а calculator таск собирает их все и суммирует. Результат записывается обратно в XCom через возвращаемое значение в функции. Вот код этого DAGа:

from datetime import datetime

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

def calculator_func(**kwargs):
    ti = kwargs['ti']
    tasks = [f'push_{i}' for i in range(1, 10)]
    values = ti.xcom_pull(task_ids=tasks)
    return sum(values)

with DAG(
    dag_id='xcom_multiple_tasks',
    start_date=datetime(2021, 3, 1),
    schedule_interval='@once',
) as dag:

    tasks = []

    for i in range(1, 10):
        task = PythonOperator(
            task_id=f'push_{i}',
            python_callable=lambda i=i: i,
        )

        tasks.append(task)

    calculator = PythonOperator(
        task_id='calculator',
        python_callable=calculator_func,
    )

    calculator.set_upstream(tasks)

А вот так это выглядит в веб-интерфейсе:

На этом, пожалуй, всё. Главное помните, что не стоит передавать через XCom большие объекты, например, файлы или DataFrame объекты какого-нибудь pandas.

Примеры всех DAG можно найти в репозитории.

💌 Присоединяйтесь к рассылке

Понравился контент? Пожалуйста, подпишись на рассылку.