Ray has become one of the default tools for scaling Python-based machine learning workloads across large fleets of hardware, but for a long time that scaling story worked more smoothly on GPUs than on Google Cloud TPUs. Ray 2.55 changes that by introducing official, first-class support for TPUs, letting developers run distributed workloads on Google's accelerators using the same task-and-actor APIs they already use elsewhere.
The core difficulty TPUs present isn't computational, it's topological. Large training jobs run on a TPU "slice," a set of chips wired together over a high-speed interconnect called ICI, and that interconnect only works if every chip in a job is genuinely part of the same physical slice. Historically, getting Ray to respect that constraint meant manually configuring the scheduler to understand TPU topology, and getting it wrong could silently fragment a job across chips that weren't actually connected, tanking performance in ways that were hard to diagnose after the fact.
Ray's new approach removes that manual step by having the KubeRay Operator on Google Kubernetes Engine automatically provision and label the underlying hardware layout when a TPU slice is created. Ray Core then reads those labels through a new primitive, slice_placement_group, which atomically reserves an entire, co-located slice as a single unit rather than piecing one together from whatever chips happen to be free. Because that atomicity wasn't previously guaranteed, building reliable multi-slice training, deliberately spanning several distinct slices at once, was effectively out of reach; the new primitive gives Ray the foundation to support that pattern going forward.
For developers, the practical upshot is that declaring a hardware topology, something like a "4x4" TPU configuration, is now enough to get a correctly placed job through KubeRay, Ray Train, or Ray Serve, without writing custom placement logic to work around TPU's networking quirks. Ray Train adds alpha support for JAX on TPUs through a new JaxTrainer API, which handles the complex distributed host initialization that JAX workloads normally require, so a developer mainly needs to specify worker count, topology, and accelerator type inside a scaling configuration object rather than wiring up multi-host coordination by hand. PyTorch on TPUs gets similar treatment through Ray Train's existing device-centric mode, which behaves much like ordinary multi-GPU PyTorch training from the developer's point of view.
This work builds on an earlier collaboration between Google and Anyscale aimed at making Ray on TPU feel as close to zero-friction as Ray already is on GPU fleets. Before these changes, engineering teams scaling AI workloads across mixed GPU and TPU hardware were often left managing each accelerator's architectural quirks by hand, even though Ray's core scaling model was meant to abstract exactly that kind of detail away.
The timing lines up with a broader push toward TPU-native tooling across Google's stack, including the recent introduction of eighth-generation TPUs built specifically with agentic and reasoning workloads in mind, and libraries like Tunix that handle the post-training side of that same pipeline. Ray's job sits earlier in that chain, providing the distributed scheduling and orchestration layer that gets a workload placed on the right hardware in the first place. Making that placement step reliable and largely automatic removes one of the more persistent sources of friction for teams trying to run serious JAX or PyTorch workloads on TPUs without a dedicated infrastructure specialist on staff, and sets up the multi-slice scaling patterns that upcoming installments in this series are expected to cover in more depth.