Source: Google Developers Blog — 2026-07-20
Summary
Ray 2.55 ships official, first-class support for Google Cloud TPUs, letting developers scale distributed Python workloads on TPUs with the same task/actor APIs they already use for GPUs. The core new primitive is slice_placement_group(), which atomically reserves an entire TPU slice — all hosts or none — by declaring a hardware topology (e.g. "4x4" on v6e chips) instead of writing custom placement logic. The KubeRay Operator on GKE provisions that topology via a TPU-specific webhook that labels TPU hosts so Ray can identify which machines belong to the same physical slice.
Key Takeaways
slice_placement_group(topology="4x4", accelerator_version="v6e")reserves a whole TPU slice atomically — partial allocations aren't possible, which matters because TPU pods only work as complete, ICI-connected units.- KubeRay is the same Kubernetes operator already used for Ray-on-GPU clusters (RayCluster/RayService/RayJob CRDs); TPU support is additive via a dedicated Ray TPU webhook, not a separate stack.
- Ray Train and Ray Serve can consume the same slice topology declaration, so training and serving jobs share one placement abstraction instead of hand-rolled scheduling code per workload type.
- Positions Ray as a portable orchestration layer across GPU and TPU fleets — the same Python task/actor code targets either accelerator by changing the topology spec, not the application logic.
Reel Script
Hook If you've ever tried to run a distributed training job across a full TPU pod, you know the pain: half the pod comes up, half doesn't, and now you're debugging a hung job instead of training a model. Google and Ray just shipped a fix for exactly that failure mode.
Core Concept
TPUs don't work like GPUs where you can grab however many you want. A TPU "slice" is a fixed block of chips wired together with a high-speed interconnect, and it only works if every chip in that block comes online together. Ray 2.55 adds a primitive called slice_placement_group that treats the whole slice as one atomic unit — Ray either reserves the entire block of hardware, or none of it, so your job never starts running on half a slice and hanging. Underneath, KubeRay — the Kubernetes operator that already manages Ray clusters on GPUs — got a TPU-specific webhook that labels each machine so Ray knows which hosts belong to the same physical slice. That's the mechanism: instead of you writing custom scheduling code to coordinate multi-host TPU jobs, you just declare the shape you need.
Hands-On
The actual code is almost suspiciously simple: slice_placement_group(topology="4x4", accelerator_version="v6e"). That one line asks for a 4x4 arrangement of v6e TPU chips as a single reserved unit, and Ray Train or Ray Serve can schedule directly against it. No custom bin-packing logic, no manual host coordination — the topology string is the whole placement spec. That's the artifact worth screenshotting: one function call replacing what used to be cluster-specific glue code.
Takeaway My take: this matters less as a TPU story and more as a portability story — the same Ray code that runs your GPU pipeline can now target TPUs by swapping a topology string, which lowers the switching cost for anyone Google is trying to pull off Nvidia hardware. If your team runs Ray on GKE, this is worth a look before your next TPU procurement conversation. Follow for more on the infra plumbing behind the AI compute shift.