はじめに
Amazon Managed Workflows for Apache Airflow (MWAA) は、AirflowのAWSマネージドサービスであり、簡単にAirflowの環境を構築できます。今回はAirflowでのSensorのよくある課題と、その解決方法としてDeferrable Operatorsを利用する方法を記載します。
Sensorの概要
Airflowではワークフロー（AirflowではDAGと呼ばれる）をPythonで記載できます。DAGでは、実施したいアクションをOperatorを使って記述します。そのOperatorの中の一つに、何かのイベントを待つために使うSensorというOperatorがあります。
例えば、S3にファイルが置かれたらGlueジョブを実行する、というDAGを作成したいときは、Sensorの一つであるS3KeySensorを使うことで実現できます。
Sensorのよくある課題
このSensorですが、デフォルトではワーカースロットを占有してしまいます。WorkerはTaskを実行するコンピューティング環境であり、例えばMWAAの環境クラスmw1.smallでは同時に実行できるTaskはデフォルトで5つとなっています。
同時に実行できるTaskが5つなので、例えば5つのSensorがS3からのファイルを待っている状態だと、他のデータ処理などのTaskが実行できなくなります。
もう少し具体的に言うと、Taskはキューで管理されており、最大同時実行数に達するとそれ以降のTaskはキューに入ったままとなります。
待ちが発生するのはSensor以外のOperatorでも同じく起きますが、Sensorは「アクションを待つ」という性質上、実行が長くなりがちで、この課題が起きる可能性も高いと言えます。
Sensorの実行モードでrescheduleを指定すればワーカースロットを使わずに待つこともできますが、時間で繰り返し実行されリソース使用量が多くなることなどから、今回は取り扱いません。Deferrable Operatorsとの比較は公式ドキュメントをご覧ください。
Deferrable Operatorsを使った解決方法
この課題を解決する方法として、Deferrable Operatorsを利用する方法があります。
これを使うと、待機中はWorkerを利用せず、Triggerという別のプロセスで実施してくれます。つまり先ほどの例だと、5つのSensorを実行していても、それらはTriggerで実行されるため、Workerでは別の処理が実行できるということです。
このDeferrable Operatorsは、Airflow v2.2.0以降で利用でき、MWAAだとv2.7.2以降で利用できます。
具体的な使い方
対応しているOperatorであれば、以下のようにdeferrable=Trueを指定するだけで利用できます。簡単ですね。
wait_for_source_data = S3KeySensor (
task_id="WaitForSourceData",
bucket_name="source_bucket_name",
bucket_key = "object_key",
aws_conn_id="aws_default",
deferrable=True
)
実際に使ってみる
今回は、S3KeySensorを、以下2通りの方法で利用してみます。
- 通常のSensorとして利用
- Deferrable Operatorsとして利用
具体的には、まずS3KeySensorでS3の指定した場所に指定のファイルが置かれるのを待ちます。その後、DummyOperatorを実行（何も起きないダミーのOperator）して終了です。
このDAGを、同時実行数を超えた回数実行して、挙動を見てみます。
上記はAirflow UIの画面ですが、ダークモードにしているため背景が黒くなっています。ダークモードは、Airflow v2.10以降で利用できます。
実行する環境
- 東京リージョン
- MWAA v2.10.1
- 環境クラス mw1.micro
- デフォルトの最大同時実行タスクは3つ
- 実行ロールには、適切なIAMポリシーを付与済み
- S3KeySensorで見に行くS3バケットの参照権限など
実行するソース
以下のソースのDAGを実行します。
deferrable=Trueを指定するとDeferrable Operatorsとして実行されるため、通常のSensorとして実行したい場合はコメントアウトします。
from airflow import DAG
from airflow.providers.amazon.aws.sensors.s3 import S3KeySensor
from airflow.operators.dummy import DummyOperator
from datetime import datetime, timedelta
default_args = {
'depends_on_past': False,
'email_on_failure': False,
'email_on_retry': False,
'retries': 1,
'retry_delay': timedelta(minutes=5)
}
# DAGの定義
dag = DAG(
's3_key_sensor_example',
default_args=default_args,
description='A simple DAG with S3KeySensor and a dummy task',
schedule_interval=None,
start_date=datetime(2023, 5, 20),
catchup=False
)
# S3KeySensorタスク
check_s3_file = S3KeySensor(
task_id="WaitForSourceData",
bucket_key="s3://<bucket-name>/data/test.txt",
dag=dag,
deferrable=True, # ★★★ Deferrable 有効化 ★★★
poke_interval=10,
timeout=600
)
# ダミータスク
dummy_task = DummyOperator(
task_id='dummy_task',
dag=dag
)
# タスクの依存関係を設定
check_s3_file >> dummy_task
① 通常のSensorとして利用
まずは、通常のSensorとして実行してみます。
デフォルトの最大同時実行タスクは3つであるため、わざとDAGを5回実行してみます。
最初の3回の実行状態は、緑色のrunningになっています。これはWorkerでSensorが実行されている状態です。
4，5回目の実行状態は、茶色のscheduluedになっています。つまり、まだ実行されていない状態です。
通常のSensorでは、Workerを使っているので4, 5回目のDAG（の中のTask）は実行されませんでした。
② Deferrable Operatorsとして利用
Deferrable Operatorsを有効化して、DAGを5回実行してみます。
5つ、全ての状態が紫色のdeferredとなっていることが分かります。TriggerでSensorが実行されているので、Workerは使用していません。
Deferrable Operatorsを使うと、Workerを使用せず5つのDAG（の中のTask）が全て実行されました。
おわりに
Deferrable Operatorsを有効/無効にした際、最大同時実行タスク以上の数のTaskを実行したらどうなるのかを見てみました。
基本的にSensorを使う場合は、有効化してよいのではないかなと思います。Sensorで待機していたから他のTaskが実行できない、というのはあまり望ましくない動きですよね。
今回利用したS3KeySensorは既にオプションがあったので簡単に有効化できましたが、例えばGlueのジョブ完了を待つGlueJobSensorなどにはそのオプションがないんですよね。
classairflow.providers.amazon.aws.sensors.glue.GlueJobSensor(*, job_name, run_id, verbose=False, aws_conn_id='aws_default', **kwargs)
一方、以前の記事で紹介したAWS Batchのジョブ完了を待つBatchSensorにはオプションがあるようです。
classairflow.providers.amazon.aws.sensors.batch.BatchSensor(*, job_id, aws_conn_id='aws_default', region_name=None, deferrable=conf.getboolean('operators', 'default_deferrable', fallback=False), poke_interval=30, max_retries=4200, **kwargs)
このように、Deferrable Operatorsを使いたい場合はドキュメントで確認することをお勧めします。
Glueジョブを実行するGlueJobOperatorにはdeferrableオプションが存在します。これを使うことで、そもそもSensorは使わずに、非同期実行+ワーカースロット解放が実現できそうです。
参考にした記事