Google Kubernetes Engine で Ray を始める

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

急速に進化する分散型コンピューティングの分野では、効率的でスケーラブルなフレームワークが求められています。Ray.io は、クラスタ内の複数のノードで Python アプリケーションを簡単にスケールアップするためのオープンソース フレームワークです。Ray は分散型の並列アプリケーション、特にディープ ラーニングのアプリケーションを構築するためのシンプルな API を提供します。

Google Kubernetes Engine（GKE）は、コンテナ化アプリケーションを簡単にデプロイおよび管理できるようにするマネージド コンテナ オーケストレーション サービスです。GKE は基盤となるインフラストラクチャを抽象化する、スケーラブルで柔軟なプラットフォームを提供します。

KubeRay を使用すると、Ray を Kubernetes にデプロイできます。Ray が提供する Pythonic の見事に統合されたエクスペリエンスと、GKE マネージド Kubernetes によるエンタープライズ向けの信頼性やスケールが得られます。これらを組み合わせることで、分散型アプリケーションの構築とデプロイ、管理の際のスケーラビリティやフォールト トレラント、および利便性が向上します。

このブログ投稿では、GKE 上で Ray を簡単に始めるためのソリューションとして、テンプレートを共有します。ソリューションのコンポーネントについて説明し、Stable Diffusion に Ray Serve を使用した推論の例を紹介します。

ソリューションの概要

このソリューションのテンプレートでは、KubeRay を使用します。これは、ワークロードをプロビジョニングするオペレーターとして Kubernetes 上の Ray クラスタを管理する OSS ソリューションです。README ファイルにある段階的な手順に沿って、ご利用を開始してください。このソリューションには、プラットフォーム レベルとユーザーレベルの 2 つのリソース グループが含まれています。

プラットフォーム レベルのリソースは、システム管理者が開発環境ごとに 1 回デプロイすることが想定されています。ここには、すべてのユーザーが共有する一般的なインフラストラクチャと GCP サービスの統合が含まれます。

• GKE クラスタとノードプール。main.tf ファイル内で構成を変更できます。このモジュールは、GPU に必要な Nvidia ドライバなどの GPU ノードプールを備えた GKE クラスタをデプロイします。これらは別のマシンタイプと交換できます。

• Kubernetes システムの名前空間とサービス アカウント、および必要な IAM ポリシー バインディング。これによりプラットフォーム管理者は、Ray クラスタ リソースに対し、きめ細かなユーザー アクセス制御と割り当てポリシーを提供できるようになります。

• KubeRay オペレーター。オペレーターには、KubeRay リソース内の変更を監視し、KubeRay クラスタの状態を調整する責任があります。  

• ロギング。「logging_config」セクションでは、システム コンポーネントおよびワークロードからのログを Cloud logging に書き込めるようにします。

• モニタリング。「monitoring_config」セクションでは、マネージド Prometheus の統合が可能になります。これにより、デプロイが自動的にシステムレベルの指標を収集し、マネージド指標サービスへの書き込みができるようになります。

• Workload identity。これにより、ワークロードは Google IAM サービス アカウントを使用して他の GCP サービスで認証できるようになります。

ユーザーレベルのリソースは、開発環境で各ユーザーが 1 回デプロイすることが想定されています。

• KubeRay クラスタ。これは、ワークロードに使用する実際の Ray クラスタです。Workload Identity プールと、GCP サービスにきめ細かなアクセスを提供する IAM バインドされたサービス アカウントを使用するよう構成されています。Ray クラスタの設定は、kuberay-values.yaml ファイルを編集してカスタマイズできます。

• ロギング。このソリューションは、各 KubeRay ワーカーノードとともにデプロイされたサイドカー コンテナを追加します。これは、fluentbit を使用して Ray ログをヘッドノードから Cloud Logging に転送するものです。fluentbit-config ファイルを編集すると、ロギング コンテナがログをフィルタおよびフラッシュする方法を変更できます。

• モニタリング。このモジュールは、ユーザーの Ray クラスタから指標を収集し、データポイントを Google マネージド Prometheus にアップロードする PodMonitoring リソースを提供します。任意で Grafana ダッシュボードをインストールし、ウェブブラウザからアクセスできます。

• JupyterHub サーバー。このモジュールを使用して、ユーザーの名前空間に JupyterHub ノートブック サーバーをインストールすることで、ユーザーは Ray クラスタを直接操作できるようになります。

Ray クラスタでワークロードを実行する

Ray Serve で提供された例を実行し、Stable Diffusion をデプロイしてみましょう。この例は、こちらの Ray Serve ドキュメントから独自に抜粋したものです。Jupyter ノートブックにある例を開くには、ブラウザの proxy-public の外部 IP（IP を取得する手順）に移動してください。次に、[ファイル] -> [URL から開く] をクリックし、ノートブックの未加工 URL を入力すると開きます。

ノートブックは Ray クラスタと同じ Kubernetes クラスタ内で実行されるため、クラスタ内部のサービス エンドポイントを使用して Ray クラスタと直接通信できます。したがって、Ray クラスタを公共のインターネット トラフィックに公開する必要はありません。本番環境ワークロードでは、GCP アカウント認証情報を使用してエンドポイントを保護する必要があります。Google Cloud Identity Aware Proxy（IAP）を使用すると、Ray クラスタなどのユーザー リソースへのきめ細かなアクセス制御が可能になり、GCP リソースを不必要に公開することなく保護できます。GKE クラスタで IAP を有効にする方法については、こちらで詳細なチュートリアルをご覧いただけます。

このノートブックには、事前トレーニングされたモデルをライブ エンドポイントにデプロイするためのコードが含まれています。最後のセルは、作成されたサービス エンドポイントを呼び出します。

このノートブックを実行すると、かわいい猫のユニークな写真が含まれたファイルが生成されます。こちらが一例です。

さあ、これで、画像生成用の大規模モデルを GKE にデプロイできました。

ロギングとモニタリング

前述したように、このソリューションにより自動的なロギングとモニタリングが可能になります。これらのログを探しましょう。

Cloud コンソールで、[ロギング] -> [ログ エクスプローラ] と進みます。クエリのテキスト ボックスに次のように入力します。

ここで、クラスタから転送された Ray ログが表示されるはずです。

モニタリング指標を表示させるには、Cloud コンソールの Metrics Explorer に移動します。「Target」のメニューで「Prometheus Target」、「Ray」の順に選択します。表示させたい指標を選択します（例: prometheus/ray_component_cpu_percentage/gauge）。

このデプロイには、Grafana デプロイも付属しています。このガイドに沿ってデプロイを開くと、Ray クラスタの指標を表示できます。

まとめ

Ray と GKE とを組み合わせることで、分散型アプリケーションを構築、デプロイ、管理するためのシンプルかつ効果的なソリューションが得られます。Ray のシンプルさは、データやモデルの開発者にとって魅力的な選択肢となり、また GKE のスケーラビリティと信頼性は、エンタープライズ プラットフォームにとっての事実上の選択肢となっています。このブログ投稿で紹介したソリューション テンプレートは、GKE に Ray をデプロイするために推奨される、KubeRay をすぐに開始できる便利な方法を提供しています。

Kubernetes および GKE での Ray の構築に関してご不明な点があれば、直接 ray-on-gke@google.com にお問い合わせいただくか、GitHub にコメントしてください。GKE を使用した AI Platforms の構築について、詳しくはユーザーガイドをご覧ください。