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

Question

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

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

• 10 个独立的工作节点，因此我可以重新启动它们，以防万一其中任何一个失败

我尝试在下图中解释这一点：

• 如果任何一个工人失败，我可以重新启动它而不影响其他工人。 它仍然使用来自 Group1 的相同线程池，因此强制执行并发限制
• 完成 step1 和 step2 的所有元素后 Group1 将完成
• Step2 没有任何并发措施

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

Answer 1

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

次优实施为 airflow 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)

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