From Experimental to Official
If a team already scales Python workloads with Ray on GPUs, that same code can now run on Google Cloud TPUs using fully supported, official APIs. As of Ray 2.55, TPUs are a first-class accelerator in Ray, complete with pre-built images and support baked into the core release pipeline, replacing what used to be an experimental path that required building custom containers and leaning on community troubleshooting.
Why TPUs Needed Special Handling
Ray treats a resource like a GPU as something schedulable: ask for it, and Ray places work there. TPUs complicate that picture slightly, because TPU chips are wired together into a fixed group called a slice, made up of several host machines whose chips share a dedicated high-speed connection known as the Inter-Chip Interconnect, or ICI. A multi-host model has to land entirely on one slice, or its workers can't reach each other and the job simply hangs.
The rough GPU analogy is a single multi-GPU box where the fast NVLink interconnect only exists inside that box. Split workers across two boxes with no cable between them, and the collective operations that synchronize gradients during training never finish. A fragmented TPU slice fails the same way.
Automating What Used to Be Manual
Previously, getting this right meant manually configuring Ray to understand a cluster's specific hardware topology, and getting it wrong meant jobs landing on fragmented, unconnected slices with severe performance consequences. The new setup changes that. Running on GKE through the Ray Operator add-on, the KubeRay Operator automatically provisions and labels the underlying hardware layout, and Ray Core uses those labels through a slice_placement_group() primitive to atomically reserve an entire slice as one unit.
That single change unlocks something that wasn't reliably possible before: true multi-slice training that intentionally spans multiple distinct slices, since Ray previously couldn't guarantee single-slice atomicity in the first place.
Familiar APIs, New Hardware
Support now extends across Ray's core libraries. Ray Train offers a JaxTrainer for JAX workloads on multi-host TPUs, handling the complex distributed host initialization automatically so a developer only needs to declare hardware requirements, worker count, topology, accelerator type, inside a simple scaling configuration object. Ray Serve deployments work the same way, just pointed at TPUs instead of GPUs, without rewriting core orchestration logic.
This first installment lays the groundwork: a TPU slice is just another resource Ray schedules onto, provided the platform respects the topology constraints underneath it. Part two of the series is expected to walk through each supporting library in more hands-on detail, covering training, serving, and elastic workload patterns for teams ready to move existing Ray pipelines onto TPU hardware.