ZOZO Advent Calendar 2025

Airflowで子DAGの完了を待機するポーリング制御の実装

airflow

Last updated at 2025-12-03Posted at 2025-12-03

概要

業務で1000以上のBigQueryテーブルに対してストアドを並列実行する機会があり、その際に実装した親子DAGのポーリング制御ロジックを紹介します。
今回はairflow dags trigger で子DAGを起動したため、親DAGで子DAGの完了を待機するポーリング制御を実装しています。

実装

def wait_for_worker_dags(**kwargs):
    """
    子DAGの完了を待機（ポーリング）
    """
    from airflow.models import DagRun
    from airflow.utils.state import State

    # 前のタスクでトリガーしたrun_idのリストを取得
    dag_run_ids = kwargs['ti'].xcom_pull(
        task_ids='trigger_worker_dags',
        key='dag_run_ids'
    )

    if not dag_run_ids:
        print("No DAG runs to wait for")
        return

    # 並列実行数の上限
    MAX_ACTIVE_DAG_RUNS = 13

    # 待機時間の設定
    first_wait = 900       # 初回15分
    subsequent_wait = 3600  # 以降1時間

    attempt = 0
    completed_runs = set()  # 完了済みのrun_idを保持

    print(f"Waiting for {len(dag_run_ids)} DAG runs to complete...")

    while True:
        active_count = 0

        # 各run_idの状態をチェック
        for run_id in dag_run_ids:
            # 既に完了済みならスキップ
            if run_id in completed_runs:
                continue

            # DagRunの状態を取得
            dag_runs = DagRun.find(
                dag_id='worker_dag',
                run_id=run_id
            )

            if not dag_runs:
                print(f"DAG run not found: {run_id}")
                continue

            state = dag_runs[0].state

            # 完了（成功または失敗）をチェック
            if state in [State.SUCCESS, State.FAILED]:
                completed_runs.add(run_id)
                print(f"DAG run completed: {run_id}, state: {state}")
            else:
                active_count += 1

        # 全て完了したか確認
        if len(completed_runs) == len(dag_run_ids):
            print(f"All {len(dag_run_ids)} DAG runs completed")
            break

        # 進捗状況を表示
        print(f"Progress: {len(completed_runs)}/{len(dag_run_ids)} completed, {active_count} active")

        # アクティブなDAG数が上限に達している場合
        if active_count >= MAX_ACTIVE_DAG_RUNS:
            print(f"Active DAG runs ({active_count}) reached limit ({MAX_ACTIVE_DAG_RUNS})")

        # 待機時間を段階的に調整
        wait_time = first_wait if attempt == 0 else subsequent_wait
        print(f"Waiting {wait_time} seconds before next check...")
        time.sleep(wait_time)
        attempt += 1

ポイントは以下の通りです。

FAILEDも完了扱い → 無限ループ回避
待機時間を段階調整できるようにする → 初回15分、以降1時間にしてリソース消費を抑制

この実装により1000件以上のテーブルに対して、BigQueryのスロット消費を抑制しつつ効率的に並列処理を実行できました。
ユースケースに応じてDagRunが見つからない場合の対策やタイムアウト対策を追加しても良さそうだなと思います。

以上です。
どなたかのお役に立てれば幸いです。

Airflowで子DAGの完了を待機するポーリング制御の実装

概要

実装

関連資料