無法在 Airflow 中的自定義運算符中通過 xcom

Question

我有一個簡單的線性 DAG（使用 Airflow 2.0 創建）有兩個任務。 我為擴展BaseOperator的每個任務都有自定義運算符。 以下是 dag 和操作符的代碼：-

class Operator1(BaseOperator):
    @apply_defaults
    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
    def execute(self, context):
        ...
        logging.info('First task')
        context['task_instance'].xcom_push(key="payload", value=data)
        return data

class Operator2(BaseOperator):
    @apply_defaults
    def __init__(self, **kwargs) -> None:
        super().__init__(**kwargs)
    def execute(self, context):
        ...
        logging.info("context is ", context)
        parameters = context['task_instance'].xcom_pull(key="payload", value=data)
 

with DAG('dag_1', default_args=DEFAULT_ARGS, schedule_interval=None) as dag:
    TASK_1 = Operator1(
        task_id='task_1',
        do_xcom_push=True)
    TASK_2 = Operator2(
        task_id='task_2',
        do_xcom_push=True)
    TASK_1 >> TASK_2 

當我運行 DAG 時，我發現用於獲取xcom值的context是空的。 我在stackoverflow上搜索了很多答案，並嘗試了其中提到的方式，但沒有奏效。

真的很感激這個問題的一些提示 - 如何在自定義運算符中推送和拉取 xcom 值？

Answer 1

我拿了你的代碼並運行它，第一個問題是start_date沒有定義，所以它最終出現了一個異常：

Exception has occurred: AirflowException (note: full exception trace is shown but execution is paused at: _run_module_as_main)
Task is missing the start_date parameter

此外，在Operator1 class 中，未定義data變量。 我想也許你在制作代碼示例時錯過了它們。

除了代碼有效之外，我認為您應該在執行 xcom_pull 操作時考慮定義task_id參數。

來自TaskInstance xcom_pull方法描述：

:param task_ids：只有具有匹配 id 的任務中的 XCom 才會被拉取。 可以通過 None 來刪除過濾器。

這是一個工作示例的代碼，請注意，我使用兩種等效方法來執行XComs操作：

from airflow import DAG
from airflow.utils.dates import days_ago
from airflow.utils.decorators import apply_defaults
from airflow.models import BaseOperator


class Operator1(BaseOperator):
    @apply_defaults
    def __init__(self,  *args, **kwargs) -> None:
        super(Operator1, self).__init__(*args, **kwargs)

    def execute(self, context):
        print('First task')
        data = "valuable_data"
        more_data = "more_valueable_data"
        context['task_instance'].xcom_push(key="payload", value=data)
        self.xcom_push(context, "more_data", more_data)
        return data


class Operator2(BaseOperator):
    @apply_defaults
    def __init__(self,  *args, **kwargs) -> None:
        super(Operator2, self).__init__(*args, **kwargs)

    def execute(self, context):
        # print(f"context is  {context}")
        data = context['task_instance'].xcom_pull(
            "task_1",
            key="payload")
        more_data = self.xcom_pull(context, "task_1", key="more_data")

        print(f"Obtained data: {data}")
        print(f"Obtained more_data: {more_data}")


with DAG('dag_1',
         default_args={'owner': 'airflow'},
         start_date=days_ago(1),
         catchup=False,
         schedule_interval=None) as dag:

    TASK_1 = Operator1(
        task_id='task_1'
    )
    TASK_2 = Operator2(
        task_id='task_2'
    )
    TASK_1 >> TASK_2

來自 Task_2 的日志：

[2021-06-15 12:55:01,206] {taskinstance.py:1255} INFO - Exporting the following env vars:
AIRFLOW_CTX_DAG_OWNER=airflow
AIRFLOW_CTX_DAG_ID=dag_1
AIRFLOW_CTX_TASK_ID=task_2
AIRFLOW_CTX_EXECUTION_DATE=2021-06-14T00:00:00+00:00
AIRFLOW_CTX_DAG_RUN_ID=backfill__2021-06-14T00:00:00+00:00
Obtained data: valuable_data
Obtained more_data: more_valueable_data
[2021-06-15 12:55:01,227] {taskinstance.py:1159} INFO - Marking task as SUCCESS. dag_id=dag_1, task_id=task_2, execution_date=20210614T000000, start_date=20210615T120402, end_date=20210615T125501

旁注：我更改了__init__方法以便也接受 *args。 我正在使用print ，但可以使用 Airflow 記錄器作為self.log.info('msg')來完成。

讓我知道這是否對您有用！