Question

我正在使用SubDAG在Airflow中创建动态DAG。我需要的是SubDAG内的任务数量由前一个任务的结果决定（subtask_ids函数的middle_section变量应该是initial_task的同一个变量功能）。

问题是我无法在xcom的子标记函数中访问SubDagOperator因为我没有任何上下文。此外，由于调度程序的自动发现DAG功能，我无法访问任何数据库以读取某些值：middle_section每隔几秒执行一次。

你们怎么解决这个问题？根据上一个任务的结果在SubDAG中创建动态数量的任务？

以下是我正在开发的代码：

import airflow
from airflow.models import DAG
from airflow.operators.python_operator import PythonOperator
from airflow.operators.subdag_operator import SubDagOperator

args = {
    'owner': 'airflow',
    'start_date': airflow.utils.dates.days_ago(2),
}


def initial_task(**context):
    subtask_ids = [0, 1, 2]
    task_instance = context['ti']
    task_instance.xcom_push(key='depot_ids', value=subtask_ids)


def middle_section_task(subtask_id):
    print(subtask_id)


def middle_section(parent_dag, arg):
    subdag = DAG(dag_id=f'{dag.dag_id}.middle',
                 default_args=args, schedule_interval='@once')

    subtask_ids = ''  # Read from xcom

    for subtask_id in subtask_ids:
        PythonOperator(task_id=f'{dag.dag_id}.middle_section_task_{subtask_id}',
                       python_callable=middle_section_task,
                       op_kwargs={'subtask_id': subtask_id}, dag=subdag)

    return subdag


def end_task(**context):
    print('Finished')


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

initial = PythonOperator(task_id='start_task', python_callable=initial_task,
                         provide_context=True, dag=dag)

middle = SubDagOperator(task_id='middle', subdag=middle_section(dag, args),
                        default_args=args, dag=dag)

end = PythonOperator(task_id='end_task', python_callable=end_task,
                     provide_context=True, dag=dag)

initial >> middle >> end

Answer 1

我遇到了同样的问题，因为我认为气流任务和子任务的数量是在DAG验证时刻定义的，所以我无法以“Airflow方式”正确解决100％的问题。并且在验证时没有运行任务，因此气流无法预先知道将安排多少个子标记。

我绕过这个问题的方式可能不是最好的（我愿意接受建议），但它有效：

main_dag.py

# imports omitted for brevity
def get_info_from_db():
    # get info from db or somewhere else, this info will define the number of subdag tasks to run
    return urls, names

dag = DAG(...)

urls, names = get_info_from_db()

# You may ignore the dummy operators
start = DummyOperator(task_id='start', default_args=args, dag=dag)
sub_section = SubDagOperator(
    task_id='import-file',
    subdag=imported_subdag(DAG_NAME, 'subdag-name', args, urls=urls, file_names=names),
    default_args=args,
    dag=dag,
)
end = DummyOperator(task_id='end', default_args=args, dag=dag)

start.set_downstream(sub_section)
section_1.set_downstream(end)

然后最后我有了我的subdag.py（确保它可以从气流中发现），以防它在一个单独的文件中

# imports omitted for brevity
def fetch_files(file_url, file_name):
    # get file and save it to disk
    return file_location

# this is how I get info returned from the previous task: fetch_files
def validate_file(task_id, **kwargs):
    ti = kwargs['ti']
    task = 'fetch_file-{}'.format(task_id)
    file_location = ti.xcom_pull(task_ids=task)

def imported_subdag(parent_dag_name, child_dag_name, args, urls, file_names):
    dag_subdag = DAG(
        dag_id='%s.%s' % (parent_dag_name, child_dag_name),
        default_args=args,
        schedule_interval="@daily",
    )
    for i in range(len(urls)):
        # the task name should also be dynamic in order not to have duplicates
        validate_file_operator = PythonOperator(task_id='validate_file-{}'.format(i+1),
                                                python_callable=validate_file,
                                                provide_context=True, dag=dag_subdag, op_kwargs={'task_id': i + 1})
        fetch_operator = PythonOperator(task_id='fetch_file-{}'.format(i+1),
                                        python_callable=fetch_zip, dag=dag_subdag,
                                        op_kwargs={'file_url': urls[i], 'file_name': file_names[i]})
        fetch_operator.set_downstream(validate_file_operator)
    return dag_subdag

基本上我的逻辑是，在Airflow验证的那一刻get_info_from_db()被执行，所有dags和子标签都被动态地正确调度。如果我在db中添加或删除内容，则在下一次dag验证中将更新要运行的任务数。

这种方法适合我的用例，但我希望将来Airflow本身支持这个功能（动态数量的任务/ subdag.tasks）。

如何根据上一个任务的结果在SubDAG中真正创建n个任务

1 个答案: