GPU アクセラレーションを備えた GKE Autopilot で TensorFlow を実行する

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

ML と AI への関心がますます高まる中で、「AI / ML ワークロードの最適な実行環境は何だろう？」と疑問に思うことはないでしょうか。

この疑問にお応えするべく開発されたのが、GPU に対応した Google Kubernetes Engine（GKE）の Autopilot オペレーション モードです。Autopilot を使用すると、インフラストラクチャの管理がすべて自動化されるため、お客様は推論、トレーニング、その他の GPU タスクといった AI / ML ワークロードの実行に専念できます。Pod または Job の定義とコンテナを提供し、それを Autopilot でスケジュールするだけで、あとは適切な GPU が自動的にプロビジョニングされてワークロードが実行されます。課金されるのは Job が実行されている間だけで、Job が完了する（または job を手動で終了させる）と課金は直ちに停止します。クリーンアップも自動的に行われます。

この投稿では、AI / ML ワークロードの作成、実行、破棄までのデモをお見せします。使用するワークロードは、NVIDIA T4 上で動作する Tensorflow 対応の Jupyter ノートブックです。このノートブックを使用して、さまざまな AI / ML トレーニングのサンプルを実行できます。Jupyter ノートブックは AI / ML について学んだり、試してみたりするのに最適です。今回は、何回かの実行の間で作業内容を保存できるよう、永続ディスクをマウントします。

こちらからデモの動画をご覧いただくこともできます。

セットアップ

まず、GKE Autopilot クラスタを作成します。GPU はすべてのリージョンで使用できるわけではないため、必要な GPU を使用できるリージョンを選択します（今回は NVIDIA T4 を使用します）。GPU が使用可能なリージョンは、Autopilot の料金表に示されています。

クラスタを作成します。

インストール

これで、GPU アクセラレーションを備えた Tensorflow 対応の Jupyter ノートブックをデプロイできます。

以下の StatefulSet の定義は、TensorFlow 環境にある Jupyter ノートブックを提供する tensorflow/tensorflow:latest-gpu-jupyter コンテナのインスタンスを作成し、NVIDIA T4 GPU をプロビジョニングして、PersistentVolume を /tf/saved パスにマウントします。これにより、作業内容を保存でき、作業再開時に同じ内容が維持されるようになります。このインスタンスは Spot で実行されるため、費用を 60～91% 削減できます（プリエンプトされた場合も作業内容は保存されます）。

これは長期にわたって使用できる正当な Jupyter ノートブックです。

デスクトップからこのノートブックに接続できるよう、ロードバランサも必要です。

これらの構成を次のようにデプロイします。

コマンドの終了を待つ間、次のコマンドを実行することで、クラスタにおけるイベントを監視してクラスタが正常に機能するかを確認できます（関連するイベントを示すため、出力は一部省略しています）。

Kubernetes と Autopilot の仕組みにより、最初に FailedScheduling が発生しますが、その理由は、コードをデプロイする時点では Pod を処理できるリソースが存在しないためです。しかしその後、TriggeredScaleUp が発生し、Autopilot がリソースを自動的に追加したことが示されます。そして、Pod にリソースが割り当てられると、最後に Scheduled が発生します。GPU ノードは通常の CPU ノードよりもプロビジョニングに要する時間が若干長く、このコンテナは起動するまで少し時間がかかります。今回の例では、Pod のスケジュールから実行まで 5 分ほどかかりました。

ノートブックの使用

次に接続してみましょう。まず、ロードバランサの外部 IP を取得します。

ロードバランサを参照します。

次のように Kubernetes で exec を指定して、提案されたコマンドを実行できます。

ログインするには、トークン（この例では e54a0e8129ca3918db604f5c79e8a9712aa08570e62d2715）をコピーして入力ボックスに貼り付け、[Log In] をクリックします。

注: このステップを省略する場合は、独自のトークンを構成に設定できます。その場合、env の行のコメントを解除し、独自のトークンを定義してください。

2 つのフォルダが表示されます。一方のフォルダにはサンプルが含まれていて、もう一方の「saved」は今回永続ディスクからマウントしたフォルダです。セッション間で状態を維持するために操作は「saved」フォルダから行い、作業を開始する前に「tensorflow-tutorials」ディレクトリを「saved」ディレクトリに移動することをおすすめします。下図の UI を参考にして、フォルダの移動や独自のノートブックのアップロードを行ってください。

では、付属サンプルをいくつか実行してみましょう。

classification.ipynb サンプル

overfit_and_underfit.ipynb サンプル

Tensorflow ドキュメントにある例のように、独自のプロジェクトをアップロードすることができます。ドキュメントからノートブックをダウンロードし、それを Jupyter の saved/ フォルダにアップロードして実行するだけです。

GPU アクセラレーションを利用した Tensorflow basics.ipynb チュートリアル

以上で、NVIDIA T4 で実行される再利用可能な TensorFlow Jupyter ノートブックが作成されました。これはただのおもちゃなどではなく、PersistentVolume に接続されているので作業内容が（StatefulSet が削除されたり、Pod が停止したりした場合でも）保存されます。Spot コンピューティングを使用しているため、費用が抑えられます。すべては 2 つの YAML ファイルからプロビジョニングされ、基になるコンピューティング ハードウェアについて考える必要がありません。素晴らしいですね！

モニタリングとトラブルシューティング

「The kernel appears to have died. It will restart automatically.」のようなメッセージが表示された場合は、まずログの最後の部分を見ます。

私が経験したよくある問題は、2 つのノートブックを実行しようとしたときに、GPU のメモリが枯渇することでした（ログには CUDA_ERROR_OUT_OF_MEMORY と記録されます）。簡単な解決策は、使用中のノートブック以外のすべてのものをシャットダウンすることです。

次のコマンドで、GPU 使用率を監視できます。

なんらかの理由でセットアップをやり直す必要がある場合、Pod を削除しさえすれば、Kubernetes によって Pod が再作成されます。この処理は Autopilot では非常に迅速に行われ、クラスタ内の GPU 対応ノードのリソースがアイドル状態になるのは短時間だけです。

次のステップ

シェル環境に入って任意のコードを実行する（ノートブック UI は使用しない）には、次のコマンドを使用します。その際、永続化するデータを /tf/saved/ に保存することを忘れないでください。

他のチュートリアルを試してみたい場合は、TensorFlow チュートリアルと Keras をご覧ください。

私の場合、Keras リポジトリのクローンを永続ボリュームに作成し、Keras のすべてのチュートリアルを自分のノートブックでも使用できるようにしました。

ご自身のノートブックに Pandas のような追加の Python モジュールが必要な場合は、同じ方法でセットアップできます。ただし、より耐久性のあるセットアップを作成するには、上述したコマンドをさらに充実させた独自の Dockerfile が必要になることでしょう（今後の投稿でそのようなレシピを共有していただける場合はお知らせください）。

いくつかのサンプルを実行した際の出力を以下に示します。

Keras timeseries/ipynb/timeseries_weather_forecasting.ipynb サンプルの出力

Keras generative/ipynb/text_generation_with_miniature_gpt.ipynb サンプルでの、あるエポックのランダム イテレーション

クリーンアップ

GPU は安価なリソースではないため、作業が終了したら必ずリソースを削除してください。次のように、StatefulSet と Service を削除してクリーンアップします。

繰り返しになりますが、Autopilot の長所は、Kubernetes リソース（今回の例では StatefulSet と LoadBalancer）を削除すると、それに関連する課金が行われなくなる点です。

残るは永続ディスクだけですが、これは残しておいてもかまいません（残しておくと、上述の StatefulSet を再作成した場合にこの永続ディスクが再びアタッチされ、作業内容が保存されます）。または、不要であれば、次のコマンドで削除できます。

クラスタも、これ以上必要ない場合は削除できます。

次のステップ

ご覧いただいたように、Autopilot で GPU ワークロードを実行することは非常に簡単です。

必要な GPU リソースを含む Kubernetes ワークロードを定義しさえすれば、残りのものは自動的に管理されます。作業が終了したら、オブジェクトを削除すれば課金は直ちに停止します。ノードのクリーンアップについて心配する必要はありません。

独自の GKE クラスタを使ってみるには、https://console.cloud.google.com/kubernetes に進んでください。Google Cloud のご利用が初めての方は、$300 分の無料トライアルを活用していただけます。

- Google Kubernetes Engine、プロダクト マネージャー William Denniss