Agentic AI
Ray adds first-class TPU support: scale distributed Python across Google's AI chips
Ray 2.55 now treats Google Cloud TPUs as a native accelerator alongside GPUs, letting Python developers scale task-and-actor workloads on TPU slices through Kubernetes with the same APIs they already know.
Google Developers Blog (AI)
Python developers who use Ray to distribute work across GPU clusters can now apply the same approach to Google's Tensor Processing Units (TPUs). Starting with Ray 2.55, TPUs have graduated from experimental status to a fully supported, first-class accelerator within Ray's core platform.
What changed#
Ray now includes official pre-built container images and native scheduling for TPUs as part of its regular release pipeline, phasing out the previous workaround of building custom containers and relying on community support. The shift signals that TPU workloads will receive the same integration depth and stability guarantees as GPU deployments.
The TPU slice constraint#
The crucial difference between scheduling on TPUs versus GPUs comes down to how the hardware is wired. TPU chips are physically bundled into fixed groups called slices—multiple host machines whose processors connect via a dedicated high-speed link (the Inter-Chip Interconnect, or ICI). Unlike a GPU cluster where you can scatter workers across separate machines, a distributed job on TPUs must land entirely within a single slice. Split your workers across two slices and collective operations—the gradient-synchronization steps that distributed training depends on—will hang indefinitely.
Ray and Google Kubernetes Engine (GKE) handle slice allocation automatically. A webhook stamps each TPU host with labels marking its slice membership, and Ray's scheduler reads those labels to reserve an entire slice as an atomic unit, guaranteeing all workers stay connected.
How it works in practice#
You declare a TPU topology (written as a shape like 4x4 for a 16-chip slice) and request it in a Kubernetes manifest using a nodeSelector and the Ray Operator add-on. Behind the scenes:
- GKE provisions the physical slice and applies slice-identity labels to its hosts
- Ray Core reads those labels and reserves the whole slice with a
slice_placement_group()primitive - Ray AI libraries (Data, Train, Serve) call that reservation automatically when you specify a topology
In most cases, you never invoke slice_placement_group() directly—the higher-level libraries do it for you.
