将参数从 Airflow 中的 BranchPythonOperator 传递给函数

Question

我在下面运行代码来创建 DAG。 创建了 Dag，但choose_best_model Dag 失败了。 错误是： ERROR - _choose_best_model() missing 1 required positional argument: 'ti' 。 我的气流版本是： 1.10.3 。 如何解决此错误？

my_dag.py

from airflow import DAG
from airflow.operators.python_operator import PythonOperator, BranchPythonOperator
from airflow.operators.bash_operator import BashOperator

from random import randint
from datetime import datetime

def _choose_best_model(ti):
    accuracies = ti.xcom_pull(task_ids=[
        'training_model_A',
        'training_model_B',
        'training_model_C'
    ])
    best_accuracy = max(accuracies)
    if (best_accuracy > 8):
        return 'accurate'
    return 'inaccurate'


def _training_model():
    return randint(1, 10)

with DAG("my_dag", start_date=datetime(2021, 1, 1),
    schedule_interval="@daily", catchup=False) as dag:

        training_model_A = PythonOperator(
            task_id="training_model_A",
            python_callable=_training_model
        )

        training_model_B = PythonOperator(
            task_id="training_model_B",
            python_callable=_training_model
        )

        training_model_C = PythonOperator(
            task_id="training_model_C",
            python_callable=_training_model
        )

        choose_best_model = BranchPythonOperator(
            task_id="choose_best_model",
            python_callable=_choose_best_model
        )

        accurate = BashOperator(
            task_id="accurate",
            bash_command="echo 'accurate'"
        )

        inaccurate = BashOperator(
            task_id="inaccurate",
            bash_command="echo 'inaccurate'"
        )

        [training_model_A, training_model_B, training_model_C] >> choose_best_model >> [accurate, inaccurate]

Answer 1

您需要将provide_context参数传递给您的操作员（它扩展了定义它的PythonOperator ）。 您还需要将kwargs添加到函数的签名中。

可以在此处 (v.1.10.15)找到可以传递给您的python_callable的上下文中参数的完整列表。

完成此操作后，您还可以使用op_kwargs参数将其他自定义参数传递给您的函数。

PythonOperator气流文档

[...]    
def _choose_best_model(ti, **kwargs):   # <-- here 
    accuracies = ti.xcom_pull(task_ids=[
        'training_model_A',
        'training_model_B',
        'training_model_C'
    ])
    [...]    
    
with DAG("my_dag", start_date=datetime(2021, 1, 1),
    schedule_interval="@daily", catchup=False) as dag:

        [...]    

        choose_best_model = BranchPythonOperator(
            task_id="choose_best_model",
            python_callable=_choose_best_model, 
            provide_context=True,   # <-- here 
        )

        [...]