データ分析

Airflow DAG の改良による Cloud Composer の最適化

2023年2月2日

Google Cloud Japan Team

※この投稿は米国時間 2023 年 1 月 21 日に、Google Cloud blog に投稿されたものの抄訳です。

データパイプラインのホスティング、オーケストレーション、管理は、どの企業にとっても複雑なプロセスです。Google Cloud が提供している Cloud Composer というフルマネージドのワークフローオーケストレーションサービスを利用すれば、企業はクラウドとオンプレミスデータセンターにまたがるワークフローの作成、スケジューリング、モニタリング、管理に対応できます。Cloud Composer は広く利用されている Apache Airflow のオープンソースプロジェクトを基に構築されており、Python プログラミング言語を使用して動作します。Apache Airflow では、ユーザーがタスクの有向非巡回グラフ（DAG）を作成できるようになっています。この DAG は、特定の間隔で実行されるようにスケジュールすることも、外部イベントによってトリガーすることも可能です。

このガイドには、Apache Airflow DAG を作成する際に一般的に適用できる、作業項目のチェックリストが掲載されています。これらのチェック項目は、Google Cloud とオープンソースコミュニティが判断したベストプラクティスに沿ったものとなっています。一連の高パフォーマンスの DAG によって Cloud Composer の動作が最適化され、標準化された作成方法によって、デベロッパーが数百個、あるいは数千個の DAG でも管理できるようになります。チェック項目のそれぞれが、Cloud Composer 環境と開発プロセスにメリットをもたらします。

はじめに

ファイル名を標準化します。作成した DAG ファイルのコレクションを他のデベロッパーが容易に参照できるようにするためです。
a. 例: team_project_workflow_version.py
DAG は決定的でなければなりません。
a. 特定の入力によって常に同じ出力が生成される必要があります。
DAG はべき等でなければなりません。
a. DAG を何度トリガーしても、毎回同じ効果 / 結果が得られなければなりません。
タスクはアトミック、かつ、べき等でなければなりません。
a. タスクごとに、他のオペレーションとは独立して再実行できる 1 つのオペレーションを処理するようにします。タスクがアトミックな場合、そのタスクの一部が成功したことは、タスク全体が成功したことを意味します。
可能な限り単純な DAG にします。
a. タスク間の依存関係が少ない単純な DAG にすると、オーバーヘッドが少なくなるため、スケジューリングのパフォーマンスが向上する傾向があります。一般的に、多数の依存関係がある深くネストされたツリー構造よりも線形構造（例: A -> B -> C）にしたほうが、効率的な DAG になります。
Python の docstring 規則に従って、各関数のファイルの上部にドキュメントを記述してください。
a. Python の docstring 規則は、他のデベロッパーやプラットフォームエンジニアが Airflow DAG を理解するために役立ちます。
b. 関数と同様に BashOperator のドキュメントも作成するようにしてください。DAG で bash スクリプトを参照している場合、スクリプトの目的を記したドキュメントがないと、このスクリプトに詳しくないデベロッパーにはトラブルシューティングが困難です。

DAG の作成を標準化する

7.default_args にオーナーを追加します。
a. デベロッパーのメールアドレス / ID を使用するか、配布リスト / チーム名を使用するかを決定します。

dag = DAG() ではなく with DAG() as dag を使用します。
a. すべてのオペレーターまたはタスクグループに DAG オブジェクトを渡す必要がなくなるようにします。
DAG ID 内にバージョンを設定します。
a. DAG 内のコードを変更するたびにバージョンを更新します。
b. こうすると、削除されたタスクログが UI に表示されなくなることや、ステータスのないタスクが古い DAG 実行に対して生成されること、DAG 変更時の一般的な混乱を防ぐことができます。
c. Airflow オープンソースには、将来的にバージョニングが実装される予定です。

10.DAG にタグを追加します。
a. デベロッパーがタグをフィルタして Airflow UI をナビゲートできるようにします。
b. 組織、チーム、プロジェクト、アプリケーションなどを基準に DAG をグループ化します。

DAG の説明を追加します。
a. 他のデベロッパーが自分の DAG の内容を理解できるようにします。
作成時には DAG を一時停止します。
a. こうすると、誤って DAG が実行されて Cloud Composer 環境の負荷が増すという事態を回避できます。
catchup=False に設定して、自動キャッチアップによる Cloud Composer 環境の過負荷を避けます。
DAG が完了せずに Cloud Composer 環境のリソースが保持されることや、再試行時に競合が引き起こされることのないよう、 dagrun_timeout を設定します。
インスタンス化で DAG に引数を渡し、すべてのタスクにデフォルトで同じ start_date が設定されるようにします。
DAG では静的な start_date を使用します。
a. 動的な start_date を使用した場合、誤った開始日が導き出され、失敗したタスクインスタンスやスキップされた DAG 実行を消去するときにエラーが発生する可能性があります。
retries を、DAG レベルで適用される default_arg として設定します。これよりも細かいレベルで設定するのは、特定のタスクで必要な場合に限定してください。
1. 適切な再試行回数は 1～4 回です。再試行回数が多すぎると、Cloud Composer 環境に不要な負荷がかかります。

読み込んでいます...

18. 以下のコールバック関数のそれぞれで行う処理を定義します（メールの送信、コンテキストのロギング、Slack チャネルへのメッセージ送信など）。DAG によっては、何もしなくても問題がない場合があります。
a. success
b. failure
c. sla_miss
d. retry

例:

読み込んでいます...

from airflow import DAG
from airflow.operators.python_operator import PythonOperator
default_args = {
   'owner': 'me',
   'retries': 2, # 最大再試行回数は 2～4 回にすること
   'retry_delay': timedelta(minutes=5)
}
def on_success_callback(context):
   # DAG 内のタスクが正常に完了した場合
   print(f"Task {context['task_instance_key_str']} succeeded!")
def sla_miss_callback(context):
   # DAG 内のタスクが SLA を満たさない場合
   print(f"Task {context['task_instance_key_str']} missed its SLA!")
def on_retry_callback(context):
   # DAG 内のタスクが再試行を行う場合
   print(f"Task {context['task_instance_key_str']} retrying...")
def on_failure_callback(context):
   # DAG 内のタスクが失敗した場合
   print(f"Task {context['task_instance_key_str']} failed!")
# DAG を作成し、コールバックを設定する
with DAG(
   'dag_id_v1_0_0',
   default_args=default_args,
   description='This is a detailed description of the DAG',
   start_date=datetime(2022, 1, 1),
   dagrun_timeout=timedelta(minutes=10),
 
   tags=['example', 'versioned_dag_id'],
   is_paused_upon_creation= True,
   catchup= False,
   schedule_interval=None,
   on_success_callback=on_success_callback, # 正常に完了した場合の処理
   sla_miss_callback=on_sla_miss_callback, # SLA を満たさない場合の処理
   on_retry_callback=on_retry_callback, # 再試行を行う場合の処理
   on_failure_callback=on_failure_callback # 失敗した場合の処理
) as dag:
   def example_task(**kwargs):
       # DAG に含めるタスクの例
       print(f"Running example task with context: {kwargs}")
   # タスクを作成して DAG に追加する
   task = PythonOperator(
       task_id="example_task",
       python_callable=example_task,
       provide_context=True,
   )

19. タスクグループを使用してタスクを整理します。

例:

読み込んでいます...

20.オペレーターの変数名と task_id の引数を一致させる必要があります。

読み込んでいます...

21. タスクに SLA を追加します。
a. SLA をタスクレベルに設定し、タスクの実行時間が想定より長い場合は通知を受け取るようにします。

読み込んでいます...

Composer 環境の負荷を軽減する

22. Python 関数ではなく Jinja テンプレート / マクロを使用します。
a. Airflow のテンプレートフィールドを使用して、環境変数の値と Jinja テンプレートに含まれる値を DAG に取り込むことができます。こうすると、DAG をべき等にできるとともに（つまり、何度呼び出しても結果が変わらないこと）、スケジューラのハートビート中に不要な関数が実行されることを防止できます。
b. The Airflow エンジンはデフォルトで、すべてのテンプレート内でアクセスできる変数を渡します。

以下の例ではベストプラクティスに反し、Python 関数 datetime に基づく変数を定義しています。

読み込んでいます...

DAG ファイルに上記のコードが含まれている場合、スケジューラがハードビートを送信するたびに、これらの関数が実行されますが、パフォーマンスが低い可能性があります。さらに重要なことに、これではべき等の DAG にはなりません。datetime.today() は DAG 実行の日付ではなく現在の日付を基準とするため、過去の日付に対して失敗した DAG を再実行することはできません。

これよりも適切な実装方法は、次のように Airflow 変数を使用することです。

読み込んでいます...

23. 独自の Airflow 変数を作成することは避けます。
a. 独自に作成した変数はメタデータデータベースに格納されるため、変数を取得するにはデータベース接続が必要になります。こうすると、Cloud Composer 環境のパフォーマンスに影響を与えかねません。そのため、環境関数または Google Cloud シークレットを使用してください。

24.すべての DAG をまったく同じスケジュールで実行しないようにします（可能な限り、ワークロードを分散して実行します）。
a. スケジュール間隔には、Airflow のマクロや time_deltas ではなく cron 式を優先して使用するようにします。こうすると、より厳密なスケジュールを指定できるとともに、1 日の中でワークロードを分散して実行しやすくなり、Cloud Composer 環境への負荷が軽減されます。
b. 特定の cron スケジュール式を生成するには、Crontab.guru が役立ちます。例についてはこちらでご確認ください。

例:

読み込んでいます...

25. 少量のデータの場合を除き、XComs を使用しないようにします。
a. 使用すると、保存量もデータベースへの接続数も増えます。
b. どうしても必要な場合は、値として JSON 辞書を使用します（辞書内の多数の値に対して 1 つの接続）

26. DAG / Google Cloud Storage のパスに不要なオブジェクトを追加しないようにします。
a. 追加しなければならない場合は、GCS のパスに、Airflow スケジューラが解析する必要のない .airflowignore ファイルを追加します（SQL、プラグインなどの場合）。

27. タスクの実行タイムアウトを設定します。

例:

読み込んでいます...

28. 可能な場合は、センサーではなく遅延可能な演算子を使用します。
a. 遅延可能な演算子は、待機の必要があることを把握すると自身を一時停止してワーカーを解放し、再開対象のジョブをトリガーに引き渡します。したがって、一時停止（遅延）中の遅延可能な演算子によってワーカースロットが占有されることはないため、アイドル状態の演算子やセンサーで浪費されるクラスタリソースの数 / 量が少なくなります。

例:

読み込んでいます...

29. センサーを使用するときは、必ず mode、poke_interval、timeout を定義します。
a. センサーを使用するには、Airflow ワーカーが実行されている必要があります。
b. n 秒間隔（つまり、poke_interval の値が 60 未満）でチェックするセンサーには、mode=poke を使用します。mode=poke が設定されたセンサーは、n 秒間隔で継続的にポーリングし、Airflow ワーカーリソースを保持します。
c. n 分間隔（つまり、poke_interval の値が 60 以上）でチェックするセンサーには、mode=reschedule を使用します。mode=reschedule が設定されたセンサーは、次のチェックまでの n 分間、Airflow ワーカーリソースを解放します。

例:

読み込んでいます...

30. Cloud Composer 環境への負荷を最小限にするために、処理を外部サービス（BigQuery、Dataproc、Cloud Functions など）にオフロードします。
a. 一般的に、これらのサービスではそれぞれに固有の Airflow オペレーターを使用します。

31. サブ DAG を使用してはなりません。
a. サブ DAG は、DAG 内に再利用可能なタスクのグループを作成するための機能として古いバージョンの Airflow で使用されていました。しかし、パフォーマンスと機能上の問題の原因となっていたため、Airflow 2.0 でサブ DAG は非推奨となりました。

32. DAG 間の依存関係には Pub/Sub を使用します。
a. マルチクラスタ / DAG 間の依存関係の例をご確認ください。

33. DAG が迅速に読み込まれるようにします。
a. 不要な「最上位レベル」の Python コードは使用しないようにします。DAG 外部からの多数のインポート、変数、関数を使用する DAG では、Airflow スケジューラの解析に時間がかかるため、Cloud Composer / Airflow のパフォーマンスとスケーラビリティが低下します。
b. インポートと関数を DAG 内に移動すると、解析時間を（数秒単位にまで）短縮できます。
c. 開発された DAG によって DAG 解析時間が過剰に増加していないことを確認します。

例:

読み込んでいます...

開発とテストを効率化する

34. 「セルフチェック」を実装します（センサーまたは遅延可能な演算子を使用）。
a. タスクが想定どおりに機能していることを確認するために、DAG にチェックを追加できます。たとえば、タスクがデータを BigQuery パーティションに push する場合、その次のタスクに、パーティションが生成されたこと、およびデータが正しいことを確認するためのチェックを追加します。

例:

読み込んでいます...

# ------------------------------------------------------------
  # ソースデータを変換して、パーティション分割テーブルに転送する
  # ------------------------------------------------------------
  create_or_replace_partitioned_table_job = BigQueryInsertJobOperator(
      task_id="create_or_replace_comments_partitioned_query_job",
      configuration={
          "query": {
	"query": 'sql/create_or_replace_comments_partitioned.sql',
	"useLegacySql": False,
          }
      },
      location="US",
  )
  create_or_replace_partitioned_table_job_error = dummy_operator.DummyOperator(
      task_id="create_or_replace_partitioned_table_job_error",
      trigger_rule="one_failed",
  )
  create_or_replace_partitioned_table_job_ok = dummy_operator.DummyOperator(
      task_id="create_or_replace_partitioned_table_job_ok",      trigger_rule="one_success"
  )
  # ------------------------------------------------------------
  # 今日の分のパーティションが comments_partitioned 内に存在するかどうか確認する
  # ------------------------------------------------------------
table_partition_sensor = BigQueryTablePartitionExistenceSensor(
      project_id="{{ project_id }}",
      task_id="bq_check_table_partition",
      dataset_id="{{ dataset }}",
      table_id="comments_partitioned",
      partition_id="{{ ds_nodash }}",
      mode="reschedule"
      poke_interval=60,
      timeout=60 * 5
  )
  create_or_replace_partitioned_table_job >> [
      create_or_replace_partitioned_table_job_error,
      create_or_replace_partitioned_table_job_ok,
  ]
  create_or_replace_partitioned_table_job_ok >> table_partition_sensor

35. 同様のタスク / タスクグループ / DAG を Python コードで動的に生成する機会を探します。
a. こうすると、DAG の開発プロセスを簡素化および標準化できます。

例:

読み込んでいます...

36. DAG の単体テストを実装します。

例:

読み込んでいます...

37. Composer ローカル開発 CLI ツールを使用して、ローカル開発を行います。
a. Composer ローカル開発 CLI ツールは、Airflow 環境をローカルで実行することにより、Cloud Composer 2 向けの Apache Airflow DAG 開発を合理化します。このローカルの Airflow 環境では、特定の Cloud Composer バージョンのイメージが使用されます。

38. 可能な場合は Cloud Composer ステージング環境を保持して、本番環境にデプロイする前に完全な DAG 実行のテストを十分に行います。
a. 変数（Google Cloud Storage オペレーションの出力パス、構成を読み取るために使用するデータベースなど）を変更するには、DAG をパラメータ化します。DAG 内に値をハードコードして、環境に応じて手動で変更しないでください。

39. Pylint や Flake8 などの Python lint チェックツールを使用して、標準化されたコードの lint チェックを行います。

40. Black や YAPF などの Python フォーマットツールを使用して、標準化されたコードのフォーマットをチェックします。

次のステップ

このブログ記事では、Google Cloud Composer で使用する Airflow DAG を開発する際のベストプラクティスを網羅したチェックリストをご紹介しました。デベロッパーはこれらのベストプラクティスに沿って、Cloud Composer の動作を最適化し、よく整理されて管理しやすい DAG を開発できます。

Cloud Composer の詳細については、以下の関連するブログ投稿とドキュメントのページをご覧ください。

-戦略的クラウドエンジニア Christian Yarros

投稿先