{"slug": "run-ray-on-tpu-part-1-the-foundations", "title": "Run Ray on TPU, Part 1: The foundations", "summary": "Ray 2.55 makes Google Cloud TPUs a first-class accelerator, with official pre-built images and support across core libraries. The distributed-computing framework now treats TPU slices as schedulable resources, using GKE's Ray Operator add-on and a TPU webhook to label hosts so Ray reserves entire slices. This eliminates the need for hand-written placement code, enabling existing Ray code for GPUs to run on TPUs unchanged.", "body_md": "*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).*\n\nAs 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).\n\nRay 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.\n\nBut there is one thing to keep in mind, and then we move on.\n\n*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.*\n\nIf 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.\n\nThat'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.\n\nOne more word you'll see everywhere is **topology**: it's the shape of a slice, written like 4x4 for a 16-chip slice. You ask for a topology, not a chip count.\n\nOnce 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.\n\nGKE 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.\n\nYou run Ray on TPU through GKE using the **Ray Operator add-on**.\n\n```\n# Autopilot (fully managed nodes)\ngcloud container clusters create-auto CLUSTER \\\n  --enable-ray-operator --location=LOCATION\n\n# or Standard (you manage node pools)\ngcloud container clusters create CLUSTER \\\n  --addons=RayOperator --location=LOCATION\n```\n\nThat 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.\n\nFrom 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.\n\n```\n# inside a RayCluster workerGroupSpec\nnodeSelector:\n  cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice   # the TPU generation\n  cloud.google.com/gke-tpu-topology: \"4x4\"              # the slice shape\n# ... and request chips via the google.com/tpu resource limit\nnumOfHosts: 4   # multi-host: how many host VMs make up this slice\n```\n\nGKE 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](https://github.com/GoogleCloudPlatform/kubernetes-engine-samples/tree/main/ai-ml/gke-ray/tpu/get-started/cluster) of the get-started example provisions all of this with Terraform.\n\nWhat 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.\n\nRay 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()`\n\n. 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.\n\n``` python\nfrom ray.util.tpu import slice_placement_group\nfrom ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy\n\n# Reserve one whole v6e 4x4 slice (16 chips across 4 hosts), atomically\nspg = slice_placement_group(topology=\"4x4\", accelerator_version=\"v6e\")\nray.get(spg.placement_group.ready(), timeout=600)\n\n@ray.remote(resources={\"TPU\": 4})\ndef worker(rank, world): ...\n\ntasks = [\n    worker.options(\n        scheduling_strategy=PlacementGroupSchedulingStrategy(\n            placement_group=spg.placement_group)\n    ).remote(rank=i, world=spg.num_hosts)\n    for i in range(spg.num_hosts)\n]\n```\n\nIt is important to highlight that **you rarely call** `slice_placement_group`\n\n**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()`\n\ndirectly 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\")`\n\n), so it's usable today but the surface can still shift between releases.\n\nYou 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.\n\nIn 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.\n\n`kubernetes-engine-samples`\n\n, cluster, serve, data, and train on a single v6e slice.For now, thanks for reading! And if you have any additional questions or feedback, feel free to reach out on socials ([LinkedIn](https://www.linkedin.com/in/ivan-nardini), [X](https://x.com/ivnardini)).\n\nHappy building!", "url": "https://wpnews.pro/news/run-ray-on-tpu-part-1-the-foundations", "canonical_source": "https://developers.googleblog.com/run-ray-on-tpu-part-1-the-foundations/", "published_at": "2026-07-20 16:25:07.085116+00:00", "updated_at": "2026-07-20 16:25:09.631167+00:00", "lang": "en", "topics": ["ai-infrastructure", "ai-chips"], "entities": ["Google Cloud", "Ray", "TPU", "GKE", "KubeRay", "Ray Operator"], "alternates": {"html": "https://wpnews.pro/news/run-ray-on-tpu-part-1-the-foundations", "markdown": "https://wpnews.pro/news/run-ray-on-tpu-part-1-the-foundations.md", "text": "https://wpnews.pro/news/run-ray-on-tpu-part-1-the-foundations.txt", "jsonld": "https://wpnews.pro/news/run-ray-on-tpu-part-1-the-foundations.jsonld"}}