构建需要上下文限制的 Airflow DAG

Question

• 我有一组要作为 DAG 运行的作业单元（工人）
• Group1 有 10 个工作人员，每个工作人员从数据库中提取多个表。请注意，每个工作人员都映射到一个数据库实例，并且每个工作人员需要成功处理总共 100 个表才能成功将自己标记为完成
• Group1 有一个限制，即一次不应超过所有这 10 个工作人员的 5 个表。例如：
  • Worker1 正在提取 2 个表
  • Worker2 正在提取 2 个表
  • Worker3 正在提取 1 个表
  • Worker4...Worker10 需要等到 Worker1...Worker3 放弃线程
  • Worker4...Worker10 可以在 step1 中的线程释放后立即拿起表
  • 当每个工作人员完成所有 100 个表时，它会继续执行步骤 2，而无需等待。 Step2 没有并发限制

我应该能够创建一个满足节流的单个节点 Group1 并且还具有

• 10 个独立的工作节点，因此我可以在其中任何一个发生故障时重新启动它们

我尝试在下图中对此进行解释：

• 如果任何一个worker出现故障，我可以重新启动它而不影响其他worker。它仍然使用来自 Group1 的相同线程池，因此强制执行并发限制
• 一旦 step1 和 step2 的所有元素都完成，Group1 就会完成
• Step2 没有任何并发​​措施

如何在 Airflow 中为 Spring Boot Java 应用程序实现这样的层次结构？ 是否可以使用 Airflow 构造设计这种 DAG，并动态地告诉 Java 应用程序一次可以提取多少个表。例如，如果除 Worker1 之外的所有工作线程都已完成，则 Worker1 现在可以使用所有 5 个可用线程，而其他所有线程都将继续执行步骤 2。

Answer 1

这些约束不能被建模为有向无环图，因此不能完全按照描述在气流中实现。但是，它们可以建模为队列，因此可以使用作业队列框架来实现。这是您的两个选择：

作为气流 DAG 实现次优：

from airflow.models import DAG
from airflow.operators.subdag_operator import SubDagOperator
# Executors that inherit from BaseExecutor take a parallelism parameter
from wherever import SomeExecutor, SomeOperator

# Table load jobs are done with parallelism 5
load_tables = SubDagOperator(subdag=DAG("load_tables"), executor=SomeExecutor(parallelism=5))

# Each table load must be it's own job, or must be split into sets of tables of predetermined size, such that num_tables_per_job * parallelism = 5
for table in tables:
    load_table = SomeOperator(task_id=f"load_table_{table}", dag=load_tables)

# Jobs done afterwards are done with higher parallelism
afterwards = SubDagOperator(
    subdag=DAG("afterwards"), executor=SomeExecutor(parallelism=high_parallelism)
)

for job in jobs:
    afterward_job = SomeOperator(task_id=f"job_{job}", dag=afterwards)

# After _all_ table load jobs are complete, start the jobs that should be done afterwards

load_tables > afterwards

这里的次优方面是，对于 DAG 的前半部分，higher_parallelism - 5 将未充分利用集群。

利用作业队列优化实施：

# This is pseudocode, but could be easily adapted to a framework like Celery

# You need two queues
# The table load queue should be initialized with the job items
table_load_queue = Queue(initialize_with_tables)
# The queue for jobs to do afterwards starts empty
afterwards_queue = Queue()

def worker():

    # Work while there's at least one item in either queue
    while not table_load_queue.empty() or not afterwards_queue.empty():
        working_on_table_load = [worker.is_working_table_load for worker in scheduler.active()]

        # Work table loads if we haven't reached capacity, otherwise work the jobs afterwards
        if sum(working_on_table_load) < 5:
            is_working_table_load = True
            task = table_load_queue.dequeue()
        else
            is_working_table_load = False
            task = afterwards_queue.dequeue()

        if task:
            after = work(task)
            if is_working_table_load:

                # After working a table load, create the job to work afterwards
                afterwards_queue.enqueue(after)

# Use all the parallelism available
scheduler.start(worker, num_workers=high_parallelism)

使用这种方法，集群不会被充分利用。