Why Run Ray on TPUs?

If you're already using Ray to scale Python workloads on GPUs, you can now run the same code on Google Cloud TPUs with official support starting from Ray 2.55. No more experimental containers or community hacks—Ray now treats TPUs as first-class accelerators, with pre-built images and full support across core libraries.

But there's a catch: TPUs are physically wired into fixed groups called slices. Each slice consists of several host VMs whose chips are connected via a high-speed link called ICI (Inter-Chip Interconnect). If your workers aren't placed on the same slice, they can't communicate, and your training will hang.

In this tutorial, you'll learn how to set up Ray on TPU using GKE and use Ray's slice_placement_group API to ensure your workers stay on one intact slice.

Ray on TPU architecture diagram showing GKE and Ray Core layers IT Technology Image

Setting Up GKE with Ray Operator

First, create a GKE cluster with the Ray Operator add-on. You can choose Autopilot (fully managed) or Standard (you manage node pools).

Autopilot Cluster

gcloud container clusters create-auto CLUSTER \
  --enable-ray-operator --location=LOCATION

Standard Cluster

gcloud container clusters create CLUSTER \
  --addons=RayOperator --location=LOCATION &&

gcloud container node-pools create v6e-16-slice \
  --cluster=CLUSTER \
  --location=LOCATION \
  --machine-type=ct6e-standard-4t \
  --tpu-topology=4x4 \
  --num-nodes=4

The Ray Operator installs two key components:

  • KubeRay: Turns RayCluster, RayService, and RayJob YAML into running Ray clusters.
  • Ray TPU webhook: Labels each TPU host with ray.io/tpu-slice-name so Ray can identify which machines are part of the same slice.

Now, request TPUs in your RayCluster manifest using a nodeSelector and resource limits:

# inside a RayCluster workerGroupSpec
nodeSelector:
  cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice # TPU generation
  cloud.google.com/gke-tpu-topology: "4x4" # slice shape
# ... request chips via the google.com/tpu resource limit
numOfHosts: 4 # multi-host: number of host VMs in the slice

Apply the manifest, and GKE will provision the slice, the webhook labels it, and Ray is ready to schedule your work.

Python code for slice_placement_group in a terminal Software Concept Art

Using slice_placement_group for Atomic Reservation

Ray Core's TPU support lives in the public ray.util.tpu API. The key function is slice_placement_group(), which reserves a whole slice atomically (all hosts or none) by matching the webhook labels.

Here's a complete example:

from ray.util.tpu import slice_placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy

# Reserve one whole v6e 4x4 slice (16 chips across 4 hosts), atomically
spg = slice_placement_group(topology="4x4", accelerator_version="v6e")
ray.get(spg.placement_group.ready(), timeout=600)

@ray.remote(resources={"TPU": 4})
def worker(rank, world):
    # Your distributed training or inference code here
    pass

tasks = [
    worker.options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=spg.placement_group
        )
    ).remote(rank=i, world=spg.num_hosts)
    for i in range(spg.num_hosts)
]

# Wait for all tasks to complete
ray.get(tasks)

Important: In practice, you rarely call slice_placement_group directly. Ray AI libraries (Data, Train, Serve) call it for you. You only need it when writing custom distributed workloads.

Caveats and Limitations

  • The API is marked as @PublicAPI(stability="alpha"), so it may change between releases.
  • Multi-host slices require careful planning; if your workers don't land on the same slice, your job will hang.
  • TPU topology must match the physical slice shape (e.g., 4x4 for 16 chips).

TPU slice with multiple host VMs connected via ICI Programming Illustration

Conclusion and Next Steps

You now have a solid mental model: a TPU slice must stay intact, GKE provisions and labels it, and Ray Core reserves it as a unit—so you never write placement code by hand.

To go deeper, check out Part 2 of this series where we explore using Ray AI libraries on TPU for serving LLMs with vLLM, feeding slices with Ray Data, and training with JaxTrainer.

Also, don't miss these related reads:

Happy building!

This content was drafted using AI tools based on reliable sources, and has been reviewed by our editorial team before publication. It is not intended to replace professional advice.