Run Ray on TPU, Part 1: The foundations

Summary

Ray 2.55 makes Google Cloud TPUs a first-class accelerator in Ray, with official APIs, pre-built images, and core-library support rather than the earlier experimental path. The central constraint is TPU slicing: a multi-host workload must land on one intact slice because its hosts communicate through the slice’s ICI (Inter-Chip Interconnect). On GKE, the Ray Operator installs KubeRay and a TPU webhook; the webhook labels hosts, while Ray Core uses those labels to reserve a whole slice atomically through the public slice_placement_group() API, with topology and accelerator version as inputs. The guide shows node selection by TPU generation and topology, a multi-host numOfHosts field, and a RayCluster manifest that brings up one worker per host. Ray AI libraries such as Data, Train, and Serve generally invoke this placement primitive themselves, while direct use targets custom distributed workloads; the API is public but alpha.

Takeaways

1. 1

   A TPU slice is requested by topology, such as 4x4, rather than only by chip count; its hosts share the ICI required for multi-host communication.

2. 2

   The Ray TPU webhook labels hosts with slice information such as ray.io/tpu-slice-name, allowing Ray Core to reserve all hosts in one slice atomically.

3. 3

   Ray Data, Train, and Serve generally call slice placement internally; direct use of slice_placement_group() is mainly for custom distributed workloads, and the public API is marked alpha.