The Ray 2.55 release introduces official, first-class support for Google Cloud TPUs, enabling developers to execute distributed Python workloads on Google's accelerators. This integration utilizes familiar Ray task-and-actor APIs. The KubeRay Operator on Google Kubernetes Engine (GKE) manages the provisioning and labeling of underlying hardware layouts to meet the strict networking requirements of multi-host TPU "slices" over their Inter-Chip Interconnect (ICI).
Key Points
- Ray 2.55 introduces official, first-class support for Google Cloud TPUs.
- Developers can run distributed Python workloads on Google Cloud TPUs using existing Ray task-and-actor APIs.
- The KubeRay Operator on GKE automatically provisions and labels hardware for multi-host TPU slices.
- Ray Core uses its
slice_placement_group()primitive to reserve complete TPU slices atomically. - Jobs can be deployed through KubeRay, Ray Train, or Ray Serve by declaring a hardware topology like "4x4" without custom placement code.
- Ray 2.55 includes TPU support in Ray's release pipelines with official pre-built images.
- A TPU "slice" consists of several host machines (VMs) with chips sharing a dedicated high-speed Inter-Chip Interconnect (ICI).
Context
According to the Google Developers Blog, Ray 2.55 integrates Google Cloud TPUs as a first-class accelerator within the Ray ecosystem. This means TPUs are now part of Ray's release pipelines, offering official pre-built images and support across core libraries, moving beyond previous experimental paths that required custom container builds and community assistance. The blog notes that a TPU is treated as another schedulable resource by Ray, similar to a GPU.
Why It Matters
This update provides builders with a direct path to leverage Google Cloud TPUs for distributed Python workloads using established Ray frameworks. The integration simplifies deployment by automating complex hardware placement, potentially reducing development effort and enabling access to specialized AI acceleration.
What To Do
- Note that Ray 2.55 is the version introducing this official support.
- Test the
slice_placement_group()primitive in Ray Core for atomic slice reservation. - Explore deploying jobs via KubeRay, Ray Train, or Ray Serve by specifying hardware topology.
- Review the Google Developers Blog series "Run Ray on TPU" for detailed guidance on GKE and Ray Core layers.
Keep Exploring
/atlas/gemini-family