Run Ray on TPU, Part 1: The foundations
- Google Developers Blog
Search
Run Ray on TPU, Part 1: The foundations
JULY 20, 2026
Ivan Nardini
AI Developer Relations
Share
TL;DR : If you already scale Python with Ray on GPUs, your code can now run on TPU (Tensor Processing Unit, Google's AI accelerator chip) with fully supported official APIs you already know. The task-and-actor model, a JaxTrainer, the same Ray Serve deployment just pointed at TPUs orchestrated by Google Kubernetes Engine (GKE).<br>As of Ray 2.55, Google Cloud TPUs are a first-class accelerator in Ray . This means that TPUs are now in Ray's release pipelines with official pre-built images and support across the core libraries, instead of the old "experimental" path where you built your own containers and leaned on community help. In this "Run Ray on TPU" series, you will learn how a TPU slice is just another accelerator Ray schedules onto (Part 1), then walk through each library (Part 2).<br>Ray and TPU: A quick introduction<br>Ray is a distributed-computing framework: you write Python, and Ray runs it across a cluster as tasks (stateless functions) and actors (stateful workers). To Ray, a TPU is just another schedulable resource, the way a GPU is. You ask for it, Ray places your work on it.<br>But there is one thing to keep in mind, and then we move on.<br>TPU chips are wired together into a fixed group called a slice : several host machines (VMs) whose chips share a dedicated high-speed link called the ICI (Inter-Chip Interconnect). A multi-host model has to land on one whole slice , or its workers can't reach each other and the job just hangs.<br>If you think in GPUs, picture a slice as a single multi-GPU box where the fast interconnect (NVLink) only exists inside the box. Split your workers across two boxes with no cable between them and the collective operations, the all-reduce steps that synchronize gradients, never finish. Training just hangs. A TPU slice behaves the same way: the ICI is that cable, and it only reaches the chips of one slice.
That's the whole reason "Ray on TPU" needs anything special. Something has to guarantee all your workers land on one intact slice. On GPUs you barely think about it; on TPUs it's crucial, and Ray and GKE (Google Kubernetes Engine, Google's managed Kubernetes) handle it for you.<br>One more word you'll see everywhere is topology : it's the shape of a slice, written like 4x4 for a 16-chip slice. You ask for a topology, not a chip count.<br>Once you understand slicing and topology with TPU, the existing Ray stack and your development process remains unchanged and it runs on TPU slices that GKE provisions. The diagram below is the whole system in one picture: on the left, the code you write (the Ray libraries you already use); in the middle, the Ray Core layer that reserves whole slices; on the right, the GKE managed layer that provisions the hardware and labels it so Ray can find slice boundaries.
GKE provisions a slice and labels its hosts, Ray Core reads those labels to reserve the whole slice at once, and your library call sits on top, declaring a topology and nothing more. No hand-written placement code anywhere. The rest of this part walks the bottom two layers, GKE then Ray Core while Part 2 covers the Ray AI libraries.<br>How GKE orchestrates Ray on TPU<br>You run Ray on TPU through GKE using the Ray Operator add-on .
# Autopilot (fully managed nodes)<br>gcloud container clusters create-auto CLUSTER \<br>--enable-ray-operator --location=LOCATION
# or Standard (you manage node pools)<br>gcloud container clusters create CLUSTER \<br>--addons=RayOperator --location=LOCATION
Shell
Copied
That single flag installs two things that matter for TPU. The first, KubeRay , is the Kubernetes operator that turns RayCluster, RayService, and RayJob YAML into running Ray clusters; it's the same KubeRay you'd use with GPUs. The second is the TPU-specific part: the Ray TPU webhook , which stamps every TPU host with labels like ray.io/tpu-slice-name so Ray can tell which machines are wired into the same slice. That label is the thread the whole system pulls on.<br>From there, you ask for TPUs in a manifest the same way you'd ask for any node, with a nodeSelector for the generation and topology and the chip count as a resource. A multi-host slice adds one field, numOfHosts.
# inside a RayCluster workerGroupSpec<br>nodeSelector:<br>cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice # the TPU generation<br>cloud.google.com/gke-tpu-topology: "4x4" # the slice shape<br># ... and request chips via the google.com/tpu resource limit<br>numOfHosts: 4 # multi-host: how many host VMs make up this slice
Shell
Copied
GKE provisions the slice, the webhook labels it, Ray reads the labels. You write Python. Once the add-on is up, you'll see the KubeRay operator pod running, and applying that manifest brings up a head pod plus one worker pod per host in the slice. The cluster step of the get-started example provisions all of this...