Airflow 使用DB动态生成气流任务

我想像这样运行气流dag->

• 我有两个气流工人W1和W2
• 在W1中，我计划了一个任务（W1-1），但在W2中，我想创建X个任务（W2-1，W2-2…W2-X）
• 每个任务的数字X和bash命令将从DB调用中派生
• worker W2的所有任务应在W1完成后并行运行

这是我的密码

dag = DAG('deploy_single', catchup=False, default_args=default_args, schedule_interval='16 15 * * *')

t1 = BashOperator(
        task_id='dummy_task',
        bash_command='echo hi > /tmp/hi',
        queue='W1_queue',
        dag=dag)

get_all_engines = "select full_command, queue_name from internal_airflow_hosts where logical_group = 'live_engines';"

db_creds = json.loads(open('/opt/airflow/db_creds.json').read())
conn_dict = db_creds["airflowdb_local"]
connection = psycopg2.connect(**conn_dict)

cursor = connection.cursor()

cursor.execute(get_all_engines)
records = cursor.fetchall()
i = 1
for record in records:
    t = BashOperator(
        task_id='script_test_'+str(i),
        bash_command="{full_command} ".format(full_command=str(record[0])),
        queue=str(record[1]),
        dag=dag)
    t.set_upstream(t1)
    i += 1

cursor.close()
connection.close()

但是，当我运行此命令时，W1上的任务成功完成，但W2上的所有任务都失败。在airflow UI中，我可以看到它可以解决正确数量的任务（本例中为10个），但这10个任务中的每一个都失败了

查看日志，我发现在W2（另一台机器上）上，airflow找不到

db_creds.json

文件

我不想将DB creds文件提供给W2

我的问题是，在这种情况下，如何动态创建气流任务？

---

基本上，我希望在airflow服务器上运行一个DB查询，并根据该查询的结果将任务分配给一个或多个Worker。数据库将包含关于哪些引擎处于活动状态等的更新信息，我希望DAG反映这一点。从日志来看，似乎每个工人都在运行DB查询。向每个工作人员提供对数据库的访问不是一个选项。

一种方法是将信息存储在数据库中

您可以获取在变量中动态生成DAG（以及必要的配置）所需的信息，并让W2从那里访问它

---

变量是可用于存储所有任务都可以访问的静态信息（没有关联时间戳的信息）的

谢谢@viraj parekh和@cwertz

经过多次尝试和错误，找到了在这种情况下使用气流变量的正确方法

步骤1）我们创建另一个名为

gen_var.py

的脚本，并将其放在dag文件夹中。这样，调度器将拾取并生成变量。如果生成变量的代码在

deploy\u single

dag中，那么我们会遇到与工作人员尝试处理dag相同的依赖性问题

"""
Code that goes along with the Airflow tutorial located at:
https://github.com/airbnb/airflow/blob/master/airflow/example_dags/tutorial.py
"""
import json
import psycopg2
from airflow.models import Variable
from psycopg2.extensions import AsIs

get_all_engines = "select full_command, queue_name from internal_airflow_hosts where logical_group = 'live_engines';"

db_creds = json.loads(open('/opt/airflow/db_creds.json').read())
conn_dict = db_creds["airflowdb_local"]
connection = psycopg2.connect(**conn_dict)

cursor = connection.cursor()

cursor.execute(get_all_engines)
records = cursor.fetchall()

hosts = {}
i = 1
for record in records:
    comm_dict = {}
    comm_dict['full_command'] = str(record[0])
    comm_dict['queue_name'] = str(record[1])
    hosts[i] = comm_dict
    i += 1

cursor.close()
connection.close()

Variable.set("hosts",hosts,serialize_json=True)

请注意对

serialize_json

的调用。气流将尝试将变量存储为字符串。如果希望将其存储为dict，则使用

serialize\u json=True

。气流仍将通过

json.dumps将其存储为字符串

步骤2）简化dag并调用此
“hosts”
变量（现在反序列化以获取dict），如下所示-

hoztz = Variable.get("hosts",deserialize_json=True)
for key in hoztz:
    host = hoztz.get(key)
    t = BashOperator(
        task_id='script_test_'+str(key),
        bash_command="{full_command} ".format(full_command=str(host.get('full_command'))),
        queue=str(host.get('queue_name')),
        dag=dag)
    t.set_upstream(t1)

希望它能帮助其他人。
W2任务失败的原因是正在运行的任务必须存在于DAG中。当它被硬编码时，这不是一个问题，但是由于您正在动态创建任务，workers/scheduler/webserver都需要访问构建DAG所需的任何依赖项。在您的情况下，这是一个db连接@Viraj Parekh的建议允许您将该依赖关系转换为一个气流变量，工作人员将有权访问该变量。@cwurtz-作为回答添加的更新谢谢。我将尝试使用气流变量和更新。