How to Run TPU Workloads on Google Cloud with Ray 2.55

Ray 2.55 elevates Google Cloud TPUs to a fully supported accelerator alongside NVIDIA GPUs. As detailed in the foundational release announcement, this update directly integrates TPUs into the Ray release pipelines, eliminating the need to build custom Dockerfiles or manually resolve libtpu dependencies. You can now provision hardware, schedule workloads, and monitor runs using standard Ray primitives and pre-built official images.

This release solves the primary networking constraint of TPU hardware: the Inter-Chip Interconnect (ICI). Rather than manually scripting topology logic to prevent workers from splitting across non-connected hosts, you can rely on Ray’s native placement groups to reserve intact hardware slices atomically.

Official Pre-Built Images

Ray now distributes official TPU container images. These images ship with the necessary jax[tpu] dependencies pre-installed and are tested directly within Ray’s continuous integration pipeline.

To use them, specify the -tpu tag suffix in your environment setup. For example, rayproject/ray:2.55.0-py310-tpu provides Ray 2.55 running on Python 3.10 with native TPU support.

Tag Component	Value	Description
Base	rayproject/ray	The official Docker repository for Ray.
Version	2.55.0	The release version introducing first-class TPU support.
Python	py310	Python version (e.g., 3.10).
Accelerator	-tpu	Bundles the required libtpu and jax[tpu] libraries.

Atomic Slice Reservation

TPU hosts communicate over a high-bandwidth physical network called the Inter-Chip Interconnect. For a multi-host TPU model to function correctly, it must execute on a single “slice”—a fixed physical group of connected host VMs.

If a scheduler splits a workload across multiple slices, the job will hang indefinitely waiting for ICI communication that physically cannot occur. Ray 2.55 introduces the slice_placement_group() primitive in Ray Core to address this.

When you call slice_placement_group(), Ray atomically reserves a complete, interconnected TPU slice before scheduling the job. This ensures that your tasks always land on hardware capable of satisfying the ICI communication requirements. The Ray Train documentation outlines the exact syntax for implementing this primitive in your specific training scripts.

KubeRay and GKE Integration

Deploying on Google Kubernetes Engine (GKE) relies on the KubeRay Operator, which has been updated to automatically provision and label underlying hardware layouts for TPUs.

Instead of writing manual placement code, you declare the desired hardware topology directly in your RayCluster YAML configuration. By specifying a layout string like “4x4”, KubeRay handles the underlying NodePool creation, applies the correct Kubernetes labels, and maps the layout to Ray’s placement groups.

This abstraction significantly reduces the operational overhead of running workloads like production AI inference on GKE. The KubeRay operator natively supports modern hardware revisions, including TPU v6e (requiring GKE version 1.31.2-gke.1115000 or later) and TPU v5p, which provides 4,800 Gb/s of ICI bandwidth per chip.

Library Support

The Ray ecosystem libraries have been updated to utilize these new core primitives natively.

Ray Train

Ray Train now includes a dedicated JaxTrainer. This trainer understands the TPU placement groups and coordinates JAX-based distributed training loops across the reserved hardware slices. It automatically sets the correct environment variables required by JAX for multi-host initialization.

Ray Serve

Ray Serve extends its deployment capabilities to target TPU slices. You can configure deployments to request specific TPU topologies, allowing you to serve large language models on TPUs with the same API used for GPUs.

Observability and Debugging

Google has also made the Ray History Server generally available on GKE. Because TPU workloads often run at scale and fail in complex distributed ways, real-time logging is insufficient. The History Server archives cluster logs, metrics, and event timelines, allowing you to debug placement group failures and performance bottlenecks after a job has completed.

Limitations and Tradeoffs

While Ray 2.55 provides robust foundational support, it lacks advanced cluster management features that you might expect from mature GPU environments.

Specifically, fault tolerance and in-place resizing of TPU clusters are not currently supported. If a single host in a TPU slice fails, the entire slice must typically be recreated, and Ray does not yet automatically migrate the workload to a warm standby slice without job interruption. If your workload requires dynamic scaling, you will need to configure elastic training externally or handle checkpointing and restarts manually until these roadmap items are delivered.

To begin running tasks, configure your KubeRay RayCluster resource with the -tpu image and a declared hardware topology, then submit your standard Ray jobs via the Ray Jobs API.