How to Install / Upgrade: Run Ray on TPU, Part 1: The foundations

Ray 2.55 introduces official first-class support for Google Cloud TPUs, allowing developers to run distributed Python workloads on Google accelerators using Ray task-and-actor APIs. To satisfy networking requirements for multi-host TPU slices over Inter-Chip Interconnect (ICI), the KubeRay Operator on GKE automatically provisions and labels the hardware layout. Ray Core then uses these labels through the slice_placement_group() primitive to manage hardware placement.

To perform the install or upgrade, update your deployment to Ray 2.55 on Google Kubernetes Engine (GKE) running the KubeRay Operator. Configure your Python workloads using the familiar Ray task-and-actor APIs. Incorporate the slice_placement_group() primitive from Ray Core into your code to leverage the labeled TPU hardware layout provisioned by the operator.

A primary gotcha is the strict networking requirement that multi-host TPU slices must remain grouped over their Inter-Chip Interconnect (ICI). To verify setup, confirm that the KubeRay Operator on GKE is automatically provisioning and labeling the underlying hardware layout. Finally, ensure that Ray Core reads these labels correctly via slice_placement_group() when scheduling distributed Python tasks across the Google Cloud TPU slices.