Ray AI ワークロードに GKE を使用する理由: 移植性、スケーラビリティ、管理性、費用

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

生成 AI と大規模言語モデル（LLM）の革命により、モデルのサイズが大きくなり、コンピューティング インフラストラクチャに対する需要も増大しています。こうした進歩を自社アプリケーションに組み入れようとしている組織では、スケジューリングのオーバーヘッドを最小限に抑える分散型コンピューティング ソリューションの必要性がますます高まっています。スケーラブルな生成 AI ソリューションのニーズが高まるなか、AI ワークロードのスケーリングと分散を目的に設計されたオープンソースの Python フレームワークである Ray の人気が高まっています。

仮想マシン（VM）上の従来型の Ray のデプロイには、スケーラビリティ、リソースの効率性、インフラストラクチャの管理性の点で制限があります。一つの代替案は、Kubernetes のパワーと柔軟性を活用するとともに、Ray のデプロイと管理を簡素化するオープンソースの Kubernetes オペレーターである KubeRay を使用して、Google Kubernetes Engine（GKE）に Ray をデプロイすることです。

「Ray on GKE のおかげで、当社の AI 担当者は、基盤となるプラットフォームの複雑さの把握や管理に頭を悩ませることなく、アプリケーションに必要なインフラストラクチャ リソース、柔軟性、スケーラビリティを簡単にオーケストレートできます。」 - Instadeep、インフラストラクチャ責任者 Nacef Labidi 氏

このブログ投稿では、GKE で Ray を実行することによる数多くの利点（スケーラビリティ、費用対効果、フォールト トレランス、独立性、移植性など）と、その利用開始方法に関するリソースについて説明します。

簡単なスケーリングとノードの自動プロビジョニング

VM 上では、Ray のスケーラビリティは本質的にクラスタ内の VM の数によって制限されます。特定のクラウド用に構成された自動スケーリングとノード プロビジョニング（例はこちら）には、マシンタイプとネットワーク構成に関する詳細な知識が必要になります。対照的に、Kubernetes はコンテナ、Pod、VM をスケジューリング ユニットとして使用してインフラストラクチャ リソースをオーケストレートしますが、Ray はアプリケーション内でデータ並列プロセスを分散し、スケジューリングにアクターとタスクを使用します。

KubeRay では、クラウドに依存しない自動スケーリングが導入されており、workerGroupSpec 内で最小および最大のレプリカを定義できるようになります。この構成に基づいて、Ray オートスケーラーはタスクの要求に応じてさらに多くの Kubernetes Pod をスケジューリングします。また、GKE Autopilot の運用モードを選択するか、GKE Standard で NAP を有効にすると、ノードのプロビジョニングが自動的に行われるため、手動構成が不要になります。

効率性の向上と起動レイテンシの改善

GKE では、確約利用割引、GPU の新しい料金モデルと予約など、割引ベースの節約が可能です。さらに、GKE を使用すると、YAML 構成を介した Spot ノードなどの費用削減手段を簡単に利用できます。

GKE では、確約利用割引、Autopilot モードでの GPU の新しい料金モデルと予約など、割引ベースの節約が可能です。さらに、GKE を使用すると、YAML 構成を介した Spot ノードなどの費用削減手段を簡単に利用できます。

リソースの使用量を最適化し、迅速な回復、より高速なイテレーション、弾力性を確保するには、起動レイテンシを低く抑える必要があります。GKE のイメージ ストリーミングを使用すると、イメージ全体のダウンロードの完了を待つことなく、Artifact Registry から対象となるコンテナ イメージを初期化し、実行できます。テストでは「ray-ml」コンテナ イメージのコンテナが「ContainerCreating」状態から「Running」状態に移行するまでに 8.82 秒かかることが実証されましたが、イメージ ストリーミングなしの場合は 5 分 17 秒でした。これは 35 倍もの速さです。Autopilot クラスタではイメージ ストリーミングが自動的に有効になります。また、Standard クラスタでもイメージ ストリーミングが利用可能です。

フォールト トレランスと独立性を実現する自動化されたインフラストラクチャ管理

VM 上で Ray クラスタを管理すると、詳細な VM 構成を通じてフォールト トレランスと独立性を制御できます。ただし、これには Kubernetes が提供するような移植可能で自動化された自己回復機能がありません。

Kubernetes は、明確かつ宣言的で、べき等の、望ましい状態の構成で表現される反復可能な自動化に優れています。また、自動の自己回復機能を提供しており、Kuberay 2.0 以降では、ヘッドノードがダウンしたときに Ray クラスタがクラッシュしなくてすむように拡張されています。

事実、Ray Serve のドキュメントでは、本番環境のワークロードに Kubernetes を推奨しています。その際は、RayService カスタム リソースを使用してヘルスチェック、ステータス レポート、障害復旧、アップグレードを自動的に処理します。

GKE では、宣言型 YAML ベースのアプローチにより、デプロイと管理が簡素化されるだけでなく、セキュリティと独立性のプロビジョニングも行えます。これは、Kubernetes の RBAC を Google Cloud の Identity and Access Management（IAM）と統合することで実現され、管理者は各 Ray クラスタに付与される権限を細かく調整できるようになります。たとえば、データの取り込みまたはモデルの保存のために Google Cloud Storage バケットへのアクセスを必要とする Ray クラスタには、そのバケットに対する読み取りと書き込みのみにアクションを制限する特定のロールを割り当てることができます。これを構成するには、Kubernetes サービス アカウント（KSA）を Ray クラスタ「workerGroupSpec」の Pod テンプレートの一部として指定し、Workload Identity アノテーションを使用して適切な権限を持つ Google サービス アカウントを KSA に関連付けます。

Kubernetes 名前空間を使用した簡単なマルチチーム共有

Ray には、初期状態では Ray クラスタ間のセキュリティの分離がありません。Kubernetes を使用すると、名前空間を利用してチームごとに Ray クラスタを作成し、Kubernetes のロールベース アクセス制御（RBAC）、リソースの割り当て、ネットワーク ポリシーを使用できます。これにより、名前空間ベースの信頼性の境界が形成され、複数のチームがそれぞれ、より大規模な Kubernetes 共有クラスタで Ray クラスタを管理できるようになります。

柔軟性と移植性

Kubernetes はデータや AI 以外にも使用できます。Kubernetes は汎用プラットフォームとして、クラウドとオンプレミス間で移植可能であり、充実したエコシステムを備えています。

Kubernetes を使用すると、同じインフラストラクチャ上で Ray ワークロードと非 Ray ワークロードを混在させることができます。これにより、中央のプラットフォーム チームは、インフラストラクチャとリソースの管理を GKE に任せながら、単一の共通コンピューティング レイヤを管理できるようになります。お客様自身の個人的な SRE になぞらえて考えてみてください。

GKE で Kuberay を使ってみる

まとめると、GKE 上で Ray を実行することで、クラウドの移植性を確保しながら、本番環境のワークロードのスケーラビリティ、費用対効果、フォールト トレランス、独立性を簡単に実現できます。変化する需要に迅速に適応する柔軟性を手にできるため、進化し続ける生成 AI 環境における先進的な組織にとって理想的な選択肢となります。

GKE で Kuberay の利用を開始するには、こちらの手順をご参照ください。このリポジトリには、GPU および TPU で Kuberay を実行するための Terraform テンプレートのほか、トレーニングとサービングの例が含まれています。また、GKE での AI / ML ページでその他のチュートリアルやコードサンプルを見つけることもできます。

-シニア エンジニアリング マネージャー Alex Bulankou

-AI / ML プロダクト マネージャー Winston Chiang