Cloud Composer と GKE を活用して機械学習のワークフローを構築する

AUTHOR :   ギックス

GKE を効率的に使うために

弊社の機械学習基盤では、ワークフロー管理ツールとして Cloud Composer (Airflow) を利用しており、機械学習タスクは別の Google Kubernetes Engine (以下、 GKE) クラスタ上で実行する構成を取っています。

GKE では複数の node pool を定義できるため、予め用途に応じた複数の node pool を用意しておくことで、タスクに応じた環境の割当てを容易に実現することができます。(現在は β版の提供に留まっていますが、GKE 側で利用リソースに応じて自動的に node pool の作成・削除を行ってくれる node-autoprovisioning も非常に有用です)

例えば、以下のような活用が考えられます。

  • 機械学習モデルそれぞれに適した環境の割り当て(CPU重視、メモリ重視、GPU 利用有無、など)
  • データセットのサイズに応じたスペックの変更
  • preemptible インスタンスを利用して安価にハイスペックなマシンを利用

このように GKE を使用することによって柔軟な処理基盤を構築することができますが、 Cloud Composer と組み合わせる場合にはいくつか留意しておくべき点があります。

Kubernetes Executor

Airflow には Executor という概念があり、DAG 内で定義された各タスクをどのような方式で実行するか、制御方法を指定することができます。この Executor の 1 つとして Kubernetes Executor というものがあり、実行ノードの指定やボリュームのマウントなどの Kubernetes の機能を享受しながら、各タスクを別の Pod で独立に実行することができます。

しかし、最大の問題点として Cloud Composer は Celery Executor 以外の Executor を2020年2月現在時点ではサポートしていません。よって、この方法はそもそも利用できないということになります。

参考までに、Airflow で現在利用できるExecutor の一覧はこちらから確認できます。

KubernetesPodOperator (GKEPodOperator)

Kubernetes の Pod 上でタスクを実行する別の方法として、 KubernetesPodOperator があります。(GKE 専用の Operator として GKEPodOperator が別に存在していますが、機能は基本的に同一です)こちらは単一のタスクを Kubernetes の Pod 上で実行する Operator なので、 Celery Executor からでも利用できる点が前述の Kubernetes Executor と大きく異なります。

実際に GKEPodOperator を利用して Cloud Composer とは別の GKE クラスタで処理を行うワークフローを作成すると以下のようになります。ここでは詳しく紹介しませんが、 Kubernetes のマニフェスト (yaml) で書いていた内容が Python のコードに変わっただけなので、 Torelation や Volume, NodeSelector, Resource Requests/Limits などのパラメータも指定することが可能です。詳細はGCP の公式ページなどをご参照ください。


これで万事解決!と言いたいところでしたが、現実はそう甘くはありませんでした。

KubernetesPodOperator で実行されるタスクはタスクごとに独立した Pod 上で実行されますが、あくまで Pod であり Kubernetes の Job ではありません。そのため、リトライや並列実行などの機能は一切利用できず、リトライ管理や並列実行制御などを全て Cloud Composer 側で行う必要があります。もちろん Airflow が提供するタスクや DAG のリトライ機能は利用できますが、全てのジョブの実行状態を Cloud Composer 側で持つ場合、 DAG の並列数や同時実行タスク数の上限にもより気を配る必要性が出てきます。

また、 Pod の起動に失敗したタスクはエラーとされてしまうため、リソースが確保できるまでスケジューリングを待機する Kubernetes の仕組みとはあまり相性がよくありません。タスクを実行する GKE のリソースが逼迫している場合、タイムアウト時間を非常に長くとるか、リソースが空くまで監視し続ける必要があります。

以上の理由から、弊社の環境では別の策を考えることになりました。

BashOperator と kubectl コマンドの利用

Airflow 公式では Kubernetes の Job としてタスクを実行できる Operator を提供していないため、今回は愚直に BashOperator で kubectl を叩く方式を採用しました。

最適な実現方法ではないかもしれませんが、メンテナンスコストや実装の手間を考えると、お手軽に利用できるのが BashOperator の利点です。以下に Composer から別の GKE クラスタに対して Job の実行と終了待ちを行う Bash スクリプトの例を記載しておきます。この例では全てのパラメータを固定値としていますが、実用上は xcom や Operator 実行時のパラメータを埋め込むことで、より柔軟に処理内容を変更することができます。(ここでは取り上げませんが、 Airflow Sensor を使うことで、より柔軟な処理待ちの機構を実現できます。)

Cloud Composer の今後への期待

Cloud Composer は Airflow へのバージョン追従が遅い点やランニングコストの面でまだまだ課題を感じる部分もありますが、マネージドサービスとして気軽に利用できるのが良いところだと感じています。

Kubernetes との連携を考えると Argo や Flyte など別のワークフローエンジンも候補に挙がるかとは思いますが、簡単に使い始められることや Dataflow など他の GCP 製品との連携が容易なことが Cloud Composer の利点だと思います。興味を持たれた方はぜひ試してみてください。


Masaaki Hirotsu
MLOps Div. 所属 / Kaggle Master
機械学習・データ分析基盤の構築に関わる事例や、クラウドを活用したアーキテクチャについて発信していきます。

SERVICE