循环运行 Airflow 任务直到完成

Question

我正在构建一个基本的 ETL 管道，该管道会访问一个主端点，该端点包含要处理的 ID（每次调用的可变数量）列表。我目前的想法是使用RabbitMQ作为队列系统，并从RabbitMQ消耗三个任务（提取，转换，加载）。我在网上看到的大多数教程都展示了任务在退出之前的简单顺序执行。我尝试构建一个 DAG，对我们收到的每个 ID 执行此顺序操作。但是当我不知道存在多少 ID 时，我在试图弄清楚如何通过气流安排所有这些任务时遇到了问题。

这是有问题的 DAG 的一般树视图：

我已经开始使用 RabbitMQ 将这些 ID 推送到队列中，并让 celery 后台处理可变数量的工作人员来处理负载。我遇到的问题是我不知道如何跳出“消费”循环。例如（我使用伪代码对 RabbitMQ 进行一些抽象）：

def extract():
    # callback function when messages are sent to this worker
    def _extract(channel, url, rest):
        resp = request.get(url)
        channel.publish('transform_queue', resp)

    # attach the callback to the queue
    channel.basic_consume('extract_queue', callback=_extract)
    channel.start_consuming() # runs a pseudo loop waiting for messages here

请注意，一些变量（例如 _extract 下的通道是隐式的，但很可能会包含在自定义运算符中。

Load 和 Transform 函数的工作方式类似。我遇到的问题是，当函数开始使用它时，它直到关闭才停止。我已经能够发送哨兵消息以允许函数“退出”，但是这会导致任务被标记为失败，并发送重试。例如这里是哨兵关闭的代码。

def extract():
    # callback function when messages are sent to this worker
    def _extract(channel, message, rest):
        if message == SHUTDOWN:
            exit()
        
        resp = requests.get(message.url)
        channel.publish('transform_queue', resp)

    # attach the callback to the queue
    channel.basic_consume('extract_queue', callback=_extract)
    channel.start_consuming() # runs a pseudo loop waiting for messages here

还可以选择有选择地取消消费者，但这只会增加更多的复杂性，因为仍然存在轮询取消的问题，然后该任务最终会遇到上述相同的问题。

我的主要问题是：

有没有办法在这个设置中成功退出？

这是解决这个问题的最佳方法吗？我想这是气流的常见用例，因此必须有一些最佳实践或常见设置。然而，我一直没能找到它。

Answer 1

我可以从您的问题中理解以下内容，因此建议探索以下内容。

<块引用>

您不确定输入的数量，因此不确定要运行流程的次数

您可以创建一个自定义操作符（比如 FindIDs），首先找出您需要执行多少个 ID 并将值推送到 XComs。然后可以将这些消息用于您的其他功能（例如提取器、转换器和加载器），它们可以按如下顺序设置

start >> findIds >> extractor >> transformer >> loader >> end

检查https://airflow.apache.org/docs/apache-airflow/stable/concepts.html?#xcoms

<块引用>

如果没有输入（或您的情况下的 ID），您需要跳过某些执行

在这种情况下，我将使用 ShortCircuitOperator 并有条件地跳过 DAG 的执行。 检查：https://github.com/apache/airflow/blob/master/airflow/example_dags/example_short_circuit_operator.py