分散 AI と ML の未来に向けて Ray と Kubernetes をともに進化させる

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

Ray は、Google Cloud のデベロッパーの間で人気のある OSS コンピューティング エンジンで、CPU、GPU、TPU にわたる複雑な分散 AI ワークロードを処理します。同様に、プラットフォーム エンジニアは、Kubernetes、特に Google Kubernetes Engine の強力で信頼性の高いインフラストラクチャ オーケストレーションに長い間信頼を寄せてきました。今年初め、Google は Anyscale とのパートナーシップを発表し、Ray と Kubernetes の優れた機能を組み合わせて、最も要求の厳しい AI ワークロードに対応する分散オペレーティング システムを構築しました。今回は、Ray と Kubernetes で共同開発したオープンソースの機能強化についてご紹介します。

Ray と Kubernetes のラベルベースのスケジューリング

Ray の主なメリットの一つは、柔軟なプリミティブ セットです。これにより、デベロッパーは基盤となるハードウェアを直接意識することなく、分散アプリケーションを記述できます。しかし、Ray の仮想リソースの既存のサポートでは十分にカバーされないユースケースもあります。

スケジューリングの柔軟性を高め、Ray アプリケーションの自動スケーリングを Ray と Kubernetes のスケジューラがより適切に実行できるようにするため、ラベルセレクタを Ray に導入します。Ray のラベルセレクタは、Kubernetes のラベルとセレクタに大きく影響を受けており、両方のシステム間で使い慣れたエクスペリエンスとスムーズな統合を提供することを目的としています。Ray Label Selector API は Ray v2.49 以降で利用可能で、分散タスクとアクターのスケジューリングの柔軟性が向上します。

新しい Label Selector API により、Ray はデベロッパーが以下のようなことを直接行えるようにします。

• Ray クラスタ内のノードにラベルを割り当てる（例: gpu-family=L4, market-type=spot, region=us-west-1）。

• タスク、アクター、プレースメント グループを起動する際に、実行するゾーン、リージョン、アクセラレータ タイプを宣言する。

• カスタムラベルを使用して、トポロジと高度なスケジューリング ポリシーを定義する。

GKE で分散アプリケーションをスケジューリングするには、Ray と Kubernetes のラベルセレクタを組み合わせて使用することで、アプリケーションと基盤となるインフラストラクチャを完全に制御できます。また、この組み合わせを GKE のカスタム コンピューティング クラスと併用して、特定の GPU タイプが利用できない場合のフォールバック動作を定義することもできます。具体的な例を見てみましょう。

以下は、利用可能な容量に応じてさまざまな GPU タイプで実行できる Ray リモートタスクの例です。Ray v2.49 以降では、プライマリ GPU タイプまたはマーケット タイプが利用できない場合に、フォールバック動作で GPU をバインドするアクセラレータ タイプを定義できるようになりました。この例では、リモートタスクは L4 GPU を使用したスポット容量をターゲットにしていますが、オンデマンドへのフォールバックも可能です。

GKE では、カスタム コンピューティング クラスを使用して同じフォールバック ロジックを結合し、Ray クラスタの基盤となるインフラストラクチャが同じフォールバック動作と一致するようにできます。

Ray ラベルセレクタの使用を開始するには、Ray のドキュメントをご覧ください。

Ray と Kubernetes でのアクセラレータ サポートの強化

今年初め、Google は新しい Ray Serve LLM API を使用して、A3 High および A3 Mega マシン インスタンスで GKE 上に DeepSeek-R1 などの大規模モデルをデプロイする機能を実証しました。GKE v1.33 と KubeRay v1.4 以降では、動的リソース割り当て（DRA）を使用してハードウェア アクセラレータを柔軟にスケジューリングして共有できるため、Ray で次世代の AI アクセラレータを使用できます。具体的には、NVIDIA GB200 NVL72 ラックスケール アーキテクチャを利用する A4X シリーズのマシンに、DRA を使用して Ray クラスタをデプロイできるようになりました。A4X での Ray で DRA を使用するには、A4X 上に AI 向けに最適化された GKE クラスタを作成し、NVL72 ラックを表す ComputeDomain リソースを定義します。

次に、Ray ワーカーの Pod テンプレートでクレームを指定します。

DRA と Ray を組み合わせることで、最も要求の厳しい Ray ワークロードで最適な GPU パフォーマンスを実現するために、Ray ワーカー グループが同じ GB200 NVL72 ラックに正しくスケジューリングされます。

また、Anyscale とのパートナーシップにより、Ray でよりネイティブな TPU エクスペリエンスを実現し、JAX などのフレームワークとのエコシステム統合を強化します。Ray Train では、Ray v2.49 以降で JAXTrainer API が導入され、JAX を使用した TPU でのモデル トレーニングが効率化されました。Ray でのこれらの TPU の改善について詳しくは、「A more native experience for Cloud TPUs with Ray on GKE」をご覧ください。

Kubernetes の書き込み可能な cgroup を使用した Ray ネイティブのリソース分離

書き込み可能な cgroup を使用すると、コンテナのルートプロセスは、特権機能を必要とすることなく、同じコンテナ内にネストされた cgroup を作成できます。この機能は、同じコンテナ内でユーザーコードと並行して複数のコントロール プレーン プロセスを実行する Ray にとって特に重要です。最も負荷の高いワークロードでも、Ray はコンテナ リソースの合計の一部をシステム クリティカルなタスク用に動的に予約できるため、Ray クラスタの信頼性が大幅に向上します。

GKE v1.34.X-gke.X 以降では、次のアノテーションを追加することで、Ray クラスタの書き込み可能な cgroup を有効にできます。

書き込み可能な cgroup を使用して Ray のリソース分離を有効にするには、ray start で次のフラグを設定します。

この機能は、セキュリティを損なうことなくスタック全体の信頼性を向上させるために、Ray と Kubernetes を進化させている例の一つです。

近日中に、タスクおよびアクターごとのリソース制限と要件のサポートも導入する予定です。これは、Ray で長らく要望が寄せられていた機能です。さらに、この機能をアップストリームするために、オープンソースの Kubernetes コミュニティと協力しています。

Pod のインプレース サイズ変更による Ray の垂直自動スケーリング

Kubernetes v1.33 で Pod のインプレース サイズ変更が導入されたことで、Kubernetes で実行する際の Ray の垂直スケーリング機能を統合する初期段階に入りました。初期のベンチマークでは、Pod を水平スケーリングの前に垂直スケーリングすることで、ワークロードの効率が 30% 向上することが示されています。

Pod のインプレース サイズ変更により、ワークロードの効率が次のように向上します。

• タスク / アクターのスケールアップの高速化: インプレース サイズ変更により、Ray ワーカーは利用可能なリソースを数秒でスケールアップできます。新しいノードをプロビジョニングするのに数分かかる場合があることを考えると、大幅に改善されています。この機能により、新しい Ray タスクのスケジューリング時間が大幅に短縮されます。

• ビンパッキングとリソース使用率の向上: Pod のインプレース サイズ変更により、Ray ワーカーを Kubernetes ノードに効率的にビンパッキングできます。新しい Ray ワーカーがスケールアップすると、利用可能なノード容量の小さな部分を予約し、残りの容量を他のワークロードのために解放できます。

• 信頼性の向上と障害の減少: メモリのインプレース スケーリングにより、メモリ不足（OOM）エラーを大幅に削減できます。失敗したジョブを再起動する必要がないため、この機能によりワークロード全体の効率と安定性が向上します。 

Ray + Kubernetes = AI 向けの分散 OS

Google は、Anyscale とのパートナーシップから生まれた最近の共同イノベーションを紹介できることを嬉しく思います。Ray と Kubernetes は、その強力な相乗効果により、最新の AI / ML 向けの分散オペレーティング システムとしての地位を確立しています。Google は、継続的なパートナーシップがオープンソースの Ray と Kubernetes のエコシステム内のイノベーションを加速し、最終的には分散 AI / ML の未来を推進すると考えています。

これらのアップデートにより、Ray が GKE でシームレスに動作するようになるための大きな一歩を踏み出しました。ご利用方法は以下のとおりです。

1. 容量をリクエストする: Dynamic Workload Scheduler Flex Start を使用して、TPU と GPU をすぐに利用できます。これにより、7 日未満で実行されるジョブのコンピューティングにアクセスできます。

2. Ray on GKE を使ってみる

3. TPU で JaxTrainer を試す

-Google、スタッフ ソフトウェア エンジニア Andrew Sy Kim
-Anyscale、スタッフ ソフトウェア エンジニア、Edward Oakes 氏