Run Ray on TPU, Part 1: The foundations
GoogleはRay 2.55でGoogle Cloud TPUを正式な第一級アクセラレータとして扱えるようになったと説明しました。GKEとKubeRayを使い、複数ホストTPUスライスをまとめて予約・実行できます。
Rayのタスク、アクター、Ray Train、Ray ServeをTPU上で公式サポートの構成として利用できます。
GKEのRay OperatorとTPU用WebhookがTPUスライスの構成をラベル付けし、Ray Coreのslice_placement_group()がスライス全体を原子的に予約します。
開発者は例えば4x4のトポロジーを宣言し、手作業の配置コードなしでマルチホストTPUジョブを実行できます。
対象はRay 2.55、GKE、Google Cloud TPUを使う開発者であり、Gemini APIやGoogle AI Studioの破壊的変更ではありません。
エージェント学習・推論などの分散PythonワークロードをTPUへ移す際の運用選択肢が広がります。