Ray 2.55でTPUが「標準の計算資源」に GKE上のAI基盤はどう変わるか｜2026年7月23日版
Google CloudのTPUを使う分散AI処理が、GPU向けに書いたRayの運用モデルへ大きく近づきました。Ray 2.55ではTPUが正式にサポートされ、タスクやActor、学習、推論サービスを共通のAPIで扱えます。
重要なのは、単に「TPUでもRayが動く」ことではありません。複数ホストにまたがるTPUの物理構成を、GKEとKubeRayが把握し、まとまりを崩さずに配置できるようになったことが今回の核心です。
- Ray 2.55でGoogle Cloud TPUが正式サポート対象になった
- GKEのRay OperatorがKubeRayとTPU用Webhookを導入する
- TPUスライスの構成をラベルで識別し、分散処理に必要な配置を維持する
- GPU向けRay基盤を持つ開発チームは、既存の設計を生かしてTPUを比較検討しやすくなる
何が変わったのか
Google Developers Blogの発表によると、Ray 2.55からGoogle Cloud TPUは「first-class accelerator」として扱われます。公式のビルド済みイメージが提供され、Rayのリリース工程でも継続的に検証される位置付けです。
従来は、独自コンテナの構築やコミュニティーベースの対応に頼る場面がありました。正式サポートへの移行により、開発チームはTPU専用の実験的な経路ではなく、Rayの標準機能を中心に構成できます。
Rayは、Pythonの関数を分散実行する「タスク」と、状態を保持するワーカーである「Actor」を基本単位にします。今回の対応では、この考え方を保ったままTPUを割り当てられます。
対応範囲
主な利用経路は次の通りです。
- Ray Core：タスクやActorにTPU資源を割り当てる
- Ray Train：
JaxTrainerなどを使って分散学習を実行する - Ray Serve：TPU上で推論サービスを展開する
- GKE／KubeRay：クラスタ作成、Pod配置、ジョブ運用を担う
Rayのアクセラレーター資料でも、Google TPUは「Rayチームがテストし、サポートする」対象として掲載されています。AMD GPUやIntel Gaudiなど、実験的またはコミュニティーサポートとされる一部のアクセラレーターとは扱いが異なります。
難所はTPUスライスを崩さない配置
TPUの運用では、チップ数だけを数えて空いているマシンへ順番に処理を置けばよいとは限りません。
複数ホスト型のTPU Podスライスは、複数のチップを高速なInter-Chip Interconnect（ICI）で接続します。分散学習のワーカーを正しく動かすには、同じスライスに属するホストを一体として扱い、要求したトポロジーを保つ必要があります。
ここがポイント: Ray 2.55のTPU対応は、TPUを単なる個数として数えるだけでなく、「どのスライスに属し、どのように接続されているか」をスケジューリングへ反映します。
GKEとWebhookの役割
GKEでRay Operatorアドオンを有効にすると、KubeRayに加えてTPU用のWebhookが導入されます。WebhookはTPUホストへ、次のような情報を示すラベルを付けます。
ray.io/tpu-slice-name：所属するTPUスライスray.io/tpu-worker-id：スライス内のワーカー番号ray.io/tpu-topology：チップの物理トポロジーray.io/accelerator-type：TPUの世代や種類ray.io/tpu-pod-type：Podの構成
Rayはこれらを手掛かりに、同じスライスへ置く必要があるワーカーを識別します。開発者が毎回、物理配線の詳細をアプリケーションコードへ埋め込む仕組みではありません。
マルチホストは一体で増減する
RayのKubeRay向けTPU資料では、複数ホスト型のワーカーグループについて、レプリカを個別Podではなくスライス単位で扱う構成が説明されています。
GKEもマルチホスト型TPUを個々のノードではなく、ノードプール単位でまとめてスケールします。これにより、一部のホストだけが起動して分散ジョブを開始できない、といった不整合を避けます。
TPU対応の本体は、APIの追加よりもこの配置制御にあります。 大規模な学習では、計算資源が存在するだけでなく、必要な接続構成のまま確保されることが性能と安定性を左右するためです。
開発者のコードと設定はどうなるか
Ray Coreでは、タスクやActorが必要とするTPU数をリソースとして宣言できます。たとえば、@ray.remote(resources={"TPU": 1})という指定で、1基分のTPU資源を要求します。
ただし、TPUは小数単位で共有できません。Rayの資料では、Google TPUはfractional resourceをサポートしないと明記されています。GPUで使われることがある「0.25基ずつ複数タスクへ割り当てる」といった設計を、そのまま移植することはできません。
Kubernetes側で必要な指定
GKE上のRayClusterでは、ワーカーPodに対して主に次の情報を指定します。
google.com/tpuのrequestsとlimits- TPUの世代を示すnode selector
- 必要なTPUトポロジー
- マルチホスト構成で作るホスト数
requestsとlimitsには同じチップ数を指定します。また、TPUの世代によって利用可能なリージョンや必要なGKEバージョンが異なるため、アプリケーションだけでなくクラスタ側の条件確認が欠かせません。
たとえばTPUクラスタ構築ガイドでは、TPU v6eにGKE 1.31.2-gke.1115000以降が必要とされています。導入時は、このような世代別要件と利用枠を先に確認する必要があります。
GPUからTPUへ移る際の現実的な判断
Rayの抽象化によって運用方法はそろいやすくなりましたが、GPU用の処理が無条件で同じ性能やコストで動くわけではありません。
移行しやすくなる部分
すでにRayを使うチームには、次の利点があります。
- タスクとActorを中心とした分散設計を維持できる
- KubeRayの
RayCluster、RayJob、RayServiceを引き続き使える - 学習と推論を別々の独自基盤に分けず、Rayの各ライブラリで管理できる
- ハードウェアの種類を変えても、ジョブ投入やサービス化の考え方を共通化しやすい
これは、GPUをすぐTPUへ全面移行すべきという意味ではありません。同じRay基盤上で、GPUとTPUを実測比較する入口が整ったと捉えるのが適切です。
個別検証が残る部分
一方、次の項目はワークロードごとの確認が必要です。
- JAXなど、TPUに適したフレームワークを利用できるか
- モデルの演算や入力形状がTPUの特性に合うか
- コンパイル時間を含めた実効レイテンシー
- TPUスライスを確保できるリージョンと割り当て上限
- データ読み込みがアクセラレーターの処理速度に追いつくか
- 障害時の再実行、チェックポイント、サービス切り替え手順
特にオンライン推論では、演算時間だけでなく起動時間やトラフィック変動への追従も評価対象です。学習で高い利用率を出せても、小さなリクエストが断続的に届くサービスでは別の結果になる可能性があります。
日本の開発チームが見るべきポイント
国内のAI開発者やクラウド運用担当者にとって、今回の更新は「TPUを採用するか」より、「比較検証を既存基盤へ組み込めるか」が重要です。
開発者
PythonとRayのタスク／Actorモデルを維持しながらTPUを試せます。ただし、リソース指定が共通化されても、モデルコードやフレームワークの互換性は別問題です。まず小規模な学習またはバッチ推論で、出力の一致と処理時間を測るのが現実的です。
クラウド運用担当者
GKEのバージョン、Ray 2.55系のイメージ、KubeRay、TPU用Webhookを一組として管理する必要があります。TPUスライスのラベルが期待通り付与されているかも、監視項目へ加えるべきです。
技術責任者
GPUとTPUの比較では、アクセラレーター単価だけを見ないことが重要です。クラスタの待機時間、コンパイル、データ転送、運用工数、障害復旧まで含めたジョブ単位の費用と所要時間で判断する必要があります。
継続ウォッチ
今回の発表は「Part 1」とされており、続編ではRayの各ライブラリを使った具体的な運用が扱われる予定です。次に確認したいのは以下の点です。
- Ray Trainでのマルチスライス学習と障害復旧の実例
- Ray Serveでのオートスケール、起動時間、レイテンシー
- Ray Dataを含む入力パイプラインの構成
- GPUとTPUを同じクラスタ設計で使い分ける際の運用例
- TPU世代ごとの対応範囲と公式イメージの更新頻度
今日のまとめ
Ray 2.55では、Google Cloud TPUが正式なアクセラレーターとしてRayの標準運用へ入りました。GKEとKubeRayはTPUスライスの所属やトポロジーをラベルで把握し、複数ホストを必要な構成のまま配置します。
これにより、GPU向けにRayを利用してきた開発チームは、タスク、Actor、学習、推論サービスという既存の枠組みを保ちながらTPUを検証しやすくなります。
次の実務的な一歩は全面移行ではありません。代表的なモデルを1つ選び、処理時間、起動時間、費用、復旧性をGPUと同じ条件で測れる検証環境を作ることです。その比較で初めて、TPUの正式サポートが自社のAI基盤にどれだけ効くかを判断できます。