Run Ray on TPU, Part 1: The foundations

TL;DR: If you already scale Python with Ray on GPUs, your code can now run on TPU (Tensor Processing Unit, Google’s AI accelerator chip) with fully supported official APIs you already know. The task-and-actor model, a JaxTrainer, the same Ray Serve deployment just pointed at TPUs orchestrated by Google Kubernetes Engine (GKE).

As of Ray 2.55, Google Cloud TPUs are a first-class accelerator in Ray. This means that TPUs are now in Ray’s release pipelines with official pre-built images and support across the core libraries, instead of the old “experimental” path where you built your own containers and leaned on community help. In this “Run Ray on TPU” series, you will learn how a TPU slice is just another accelerator Ray schedules onto (Part 1), then walk through each library (Part 2).

Ray and TPU: A quick introduction

Ray is a distributed-computing framework: you write Python, and Ray runs it across a cluster as tasks (stateless functions) and actors (stateful workers). To Ray, a TPU is just another schedulable resource, the way a GPU is. You ask for it, Ray places your work on it.

But there is one thing to keep in mind, and then we move on.

TPU chips are wired together into a fixed group called a slice: several host machines (VMs) whose chips share a dedicated high-speed link called the ICI (Inter-Chip Interconnect). A multi-host model has to land on one whole slice, or its workers can’t reach each other and the job just hangs.

If you think in GPUs, picture a slice as a single multi-GPU box where the fast interconnect (NVLink) only exists inside the box. Split your workers across two boxes with no cable between them and the collective operations, the all-reduce steps that synchronize gradients, never finish. Training just hangs. A TPU slice behaves the same way: the ICI is that cable, and it only reaches the chips of one slice.

fig1 (2)

That’s the whole reason “Ray on TPU” needs anything special. Something has to guarantee all your workers land on one intact slice. On GPUs you barely think about it; on TPUs it’s crucial, and Ray and GKE (Google Kubernetes Engine, Google’s managed Kubernetes) handle it for you.

One more word you’ll see everywhere is topology: it’s the shape of a slice, written like 4×4 for a 16-chip slice. You ask for a topology, not a chip count.

Once you understand slicing and topology with TPU, the existing Ray stack and your development process remains unchanged and it runs on TPU slices that GKE provisions. The diagram below is the whole system in one picture: on the left, the code you write (the Ray libraries you already use); in the middle, the Ray Core layer that reserves whole slices; on the right, the GKE managed layer that provisions the hardware and labels it so Ray can find slice boundaries.

fig2 (2)

GKE provisions a slice and labels its hosts, Ray Core reads those labels to reserve the whole slice at once, and your library call sits on top, declaring a topology and nothing more. No hand-written placement code anywhere. The rest of this part walks the bottom two layers, GKE then Ray Core while Part 2 covers the Ray AI libraries.

How GKE orchestrates Ray on TPU

You run Ray on TPU through GKE using the Ray Operator add-on.

# Autopilot (fully managed nodes)
gcloud container clusters create-auto CLUSTER \
  --enable-ray-operator --location=LOCATION

# or Standard (you manage node pools)
gcloud container clusters create CLUSTER \
  --addons=RayOperator --location=LOCATION

Shell

That single flag installs two things that matter for TPU. The first, KubeRay, is the Kubernetes operator that turns RayCluster, RayService, and RayJob YAML into running Ray clusters; it’s the same KubeRay you’d use with GPUs. The second is the TPU-specific part: the Ray TPU webhook, which stamps every TPU host with labels like ray.io/tpu-slice-name so Ray can tell which machines are wired into the same slice. That label is the thread the whole system pulls on.

From there, you ask for TPUs in a manifest the same way you’d ask for any node, with a nodeSelector for the generation and topology and the chip count as a resource. A multi-host slice adds one field, numOfHosts.

# inside a RayCluster workerGroupSpec
nodeSelector:
  cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice   # the TPU generation
  cloud.google.com/gke-tpu-topology: "4x4"              # the slice shape
# ... and request chips via the google.com/tpu resource limit
numOfHosts: 4   # multi-host: how many host VMs make up this slice

Shell

GKE provisions the slice, the webhook labels it, Ray reads the labels. You write Python. Once the add-on is up, you’ll see the KubeRay operator pod running, and applying that manifest brings up a head pod plus one worker pod per host in the slice. The cluster step of the get-started example provisions all of this with Terraform.

What actually keeps your workers together is a Ray Core primitive sitting just above this layer, the slice placement group, and that’s where the rest of the guide starts.

Ray Core on TPU

Ray Core is the base layer, the task-and-actor engine and scheduler everything else sits on. Its TPU support lives in the public ray.util.tpu API, and there’s really one function to know: slice_placement_group(). It takes that “keep my workers on one intact slice” guarantee from earlier and turns it into a single call, reserving a whole slice atomically (all hosts or none) by matching on the webhook labels.

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): ...

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)
]

Python

It is important to highlight that you rarely call slice_placement_group yourself. The Ray AI libraries (Data, Train, Serve) call it for you, so in practice you declare a topology and they handle the slice. You’d only reach for slice_placement_group() directly when you’re writing a custom distributed workload that isn’t Train, Serve, or Data. One caveat worth knowing: the API is public but marked alpha (@PublicAPI(stability="alpha")), so it’s usable today but the surface can still shift between releases.

That’s the foundation. Next are the libraries.

You now have the whole mental model: a slice has to stay intact, GKE provisions and labels it, and Ray Core reserves it as a unit so you never hand-write placement code. Everything you actually build sits on top of that and reuses it.

In Part 2, we will explore how you can use Ray AI libraries on TPU for serving LLMs with vLLM, feeding slices with Ray Data and training with JaxTrainer.