AI Platform のステータス通知
この記事は TensorFlow Advent Calendar 24 日目の記事です。
TensorFlow の面白ネタを用意することに失敗したので、TensorFlow で書いたコードの実行環境としてよく使われる AI Platform という Google Cloud Platform のサービスの話をしようと思います。
AI Platform とは
いわゆるジョブシステム的なやつで、任意のコードを submit すると、そいつをクラウド上で実行してくれるというものです。
昔は TensorFlow 用に最適化されたランタイム上でコードを実行していたのですが、最近では Docker イメージを渡してコンテナを実行することもできるようになり、TensorFlow ユーザー以外も安心して使えるようになったりしました。
TensorFlow の分散学習用のクラスタを作る機能なんかもあるので、気軽に分散学習を試してみたい方にもおすすめです。
今回の話
そんな AI Platform なんですが、ジョブの終了を良い感じに通知してくれないという欠点がありました。
自分からジョブのステータスを見に行けばもちろんわかるのですが、終了したら AI Platform 側から通知して、Slack に投稿するなり他のシステムと繋ぐなりしたいですよね？
Stackoverflow でも同様の質問があるのですが、残念ながら自分で気合い入れて実装してくれというのが公式な回答のようです。
というわけで、仕方なく実装してみました。
アーキテクチャ
こんな感じの構成です:
- AI Platform の必要なログを拾って、Stackdriver Logging 経由で Pub/Sub に流し込む
- Pub/Sub から Cloud Run にそのまま push でログを送って、Cloud Run 上でメッセージ内容からジョブステータスを判定
- 成功や失敗に対応するログを発見したら、それを適切に加工して、また Pub/Sub に流して完了
ということをします。
ジョブのステータスの情報が Pub/Sub に流れさえすれば、そこから push で別のところに通知を送るのはどうにでもなるはずです。
コード
そんなわけで、ここにコードを公開しておきます:
構成は簡単なんですが、権限周りなどの設定項目がけっこう多いので、練習がてら Terraform のモジュールにしてみました。
Cloud Run の部分はあらかじめ Docker イメージを Google Container Registry に公開しておく必要があったので gcr.io/sfujiwara/ai-platform-notification というイメージを用意しました。
ひとまず public にしてありますが真面目にメンテするかどうかは保証できないので、このリポジトリは参考程度に留めておいて、ちゃんと使いたい人はフォークするなどして自分の GCP プロジェクトでなんやかんやするのをおすすめします。
まとめ
簡単そうに見えて権限周りとかわりと面倒で、正直作るのくっそだるかったです。
AI Platform を使っている TensorFlow ユーザーで同じ悩みを抱えている人はけっこういそうなので、何かの参考になれば幸いです。
あと Google はこのへんの機能を公式に作ってくれ〜。
Discussion