ブログ

Airflow で異なる DAG 間の依存関係を設定する方法

2019.12.15

Tech

こんにちは。ソフトウェアエンジニアの冨田です。
FLYWHEEL Advent Calendar 2019 の15日目は、Apache Airflow で ExternalTaskSensor を使い異なる DAG に属するタスク間に依存関係を設定する方法を紹介します。
最初にAirflowの紹介は簡単にするものの、Airflow 使用したことがある読者を想定しています。

Apache Airflow について

フライウィールが開発しているデータプラットフォームではデータパイプラインの管理に Apache Airflow を採用しています。
Airflow は Python で定義されたワークフローの管理やスケジューリングを行うためのプラットフォームで、GCP では Cloud Composer というマネージドのサービスが提供されています。
Airflow ではワークフロー内のタスクの依存関係を DAG (Directed Acyclic Graph, 有向非巡回グラフ) と呼ばれる閉路を持たない有向グラフで定義することができます。定義した DAG は定期的に実行するようスケジュールを設定したり、Web View から実行状況を確認、必要であれば再実行するなどの処理を行うことができます。
例えば下記のようなDAGを実装すると、タスクCはタスクAとタスクBが成功した場合にのみ実行されるようになります。

DAG 間の依存関係

一つの DAG の肥大化を防ぐために、私のプロジェクトでは DAG を役割ごとに複数個定義していました。日々追加される商品データを処理するための DAG、ユーザーのウェブサイト上での行動ログを処理するための DAG、受注データを処理するための DAG、といったイメージです。
しかし、新機能の開発を進めていく上で、複数の DAG 間に依存関係を設定したいケースが出てくるようになりました。状況の一例を図で示すと次のようになります。DAG-1 のタスクCの終了を待ってから、DAG-2 のタスクEを実行したい、つまり赤い点線のような依存関係を設定したい状況です。

依存関係は通常は同じ DAG 内のタスクに対して設定するのですが、このような DAG を超えた依存関係はどのように設定できるのでしょうか。

ExternalTaskSensor の活用

異なる DAG に存在するタスクに依存関係を設定するには、ExternalTaskSensor を活用することができます。ExternalTaskSensor では  external_dag_id  と  external_task_id  を指定することができ、指定されたタスクのステータスを定期的1に確認します。指定されたタスクの状態が  allowed_status  に含まれるいずれかの状態に遷移した時、ExternalTaskSensor 自体も success ステータスとなります。
各パラメータの細かい定義はドキュメントを参照していただきたいのですが、具体的な使用例は下記のようになります。先ほど紹介した DAG-1 を午前10時、DAG-2 を午前10時1分に実行する場合は次のような ExternalTaskSensor を定義すれば良いことになります。なお、動作検証等は Airflow Version 1.10.6 で行っています。

# dag1.py
import datetime
from airflow.models import DAG
from airflow.operators.dummy_operator import DummyOperator
DAG_ID = "dag_1"
DEFAULT_ARGS = {'start_date': datetime.datetime(2019, 12, 15)}
with DAG(DAG_ID, default_args=DEFAULT_ARGS,
     schedule_interval="0 10 * * *") as dag:
  task_a = DummyOperator(task_id='task_a')
  task_b = DummyOperator(task_id='task_b')
  task_c = DummyOperator(task_id='task_c')
  [task_a, task_b] >> task_c
# dag2.py
import datetime
from airflow.models import DAG
from airflow.operators.dummy_operator import DummyOperator
from airflow.sensors.external_task_sensor import ExternalTaskSensor
DAG_ID = "dag_2"
DEFAULT_ARGS = {'start_date': datetime.datetime(2019, 12, 15)}
with DAG(DAG_ID, default_args=DEFAULT_ARGS,
     schedule_interval="01 10 * * *") as dag:
  task_d = DummyOperator(task_id='task_d')
  task_e = DummyOperator(task_id='task_e')
  task_c_sensor = ExternalTaskSensor(
    task_id='task_c_sensor',
    external_dag_id='dag_1',
    external_task_id='task_c',
    allowed_states=['success'],
    execution_delta=datetime.timedelta(minutes=1))
[task_d, task_c_sensor] >> task_e
実際にこの2つの DAG を設定すると、タスクCの終了を待ってからタスクEが実行される様子が確認できます。
ExternalTaskSensor の注意点
いままで ExternalTaskSensor を使うと異なる DAG に対して依存関係を設定できることを見てきました。
 個人的に ExternalTaskSensor を使う上で注意した方が良いと感じるのは、依存関係が Web View で明示的に可視化されない点です。通常、同一 DAG 内の依存関係は Airflow の Graph View で確認できるのですが、ExternalTaskSensor の依存関係はここでは可視化されません。
 ExternalTaskSensor で複雑な依存関係を設定しすぎると、依存関係に気づかないで DAG を更新してしまい思わぬ事故が発生する可能性もあるので、ExternalTaskSensor の使用は必要最低限に止めるべきでしょう。
おまけ: XCOM で値も受け渡す
依存関係だけでなく、異なる DAG で計算した値を受け取って使用したい場面もあるかもしれません。先ほどの例でいうと、タスクCからタスクEに何か値を渡したいとします。
 このような状況では XCOM を使用することができます。XCOM は cross-communication を意味し異なるタスク間でメッセージをやり取りするために用意された手段です。
 公式ドキュメントに詳しく説明されているように、タスクは key-value ペアを XCOM に push したり、別のタスクが push した値 pull することができます。
 BashOperator の場合、 xcom_push=True  というオプションを指定すると、 キーが “return_value” で値が標準出力の最後の行という key-value ペアを push してくれます。例えば下記のようにすれば、 “hello world!” という文字列を XCOM に push できます。
task_c = BashOperator(task_id='task_c',
             bash_command='sleep 120 && echo "hello world!"',
             xcom_push=True)
次にタスクEではこの値を pull します。pull は TaskInstance オブジェクトから行うことができ、マクロを使うと  {{ ti.xcom_pull(…) }}  といった形で使用できます。
 タスクEで先ほどタスクCから push した値を pull する場合は下記のようになります。
task_e = BashOperator(
  task_id='task_e',
  bash_command=
  'echo "recieved: {{ ti.xcom_pull(dag_id="dag_1", task_ids="task_c", include_prior_dates=True) }}"'
)
 include_prior_dates=True  となっているのがポイントで、デフォルトの値である  include_prior_dates=False  の場合、同じ  execution_date  を持つ DAG からしか値を pull することができません。
 今回の例のように  execution_date  が古い DAG (過去にトリガーされた DAG) から値を pull したい場合は、 include_prior_dates=True  とすると直近に実行された DAG に含まれるタスクから値を pull してきてくれます。
まとめ
今回は ExternalTaskSensor や XCOM を使用して異なる DAG 間で依存関係の設定や値の受け渡しを行う方法を紹介しました。
 基本的には一つの DAG 内で依存関係が完結することが望ましいとは思いますが、どうしても必要な場合は ExternalTaskSensor の使用も検討してみてはいかがでしょうか。
 フライウィールではこのような技術を使い、データプラットフォームを一緒に開発してくれるエンジニアを募集しています。我こそは！と思う方はぜひオフィスに遊びに来てみてください。
 
Notes
  1. デフォルトの設定では BaseSensorOperator の poke_interval の初期値である60秒ごと ↩
 
         
一覧に戻る