What you need to know from Part 1
Running Ray on TPU comes down to one caveat. TPU chips are wired into fixed groups called slices, and a multi-host model has to land on one whole slice or its workers can't reach each other and the job hangs. Google Kubernetes Engine with the Ray Operator add-on provisions slices and labels their hosts, and a Ray Core primitive called slice_placement_group() reserves a whole slice at once. You declare a topology (the slice shape, like 4x4 for 16 chips) and the libraries handle the placement for you.
With Core handling placement underneath, the libraries all follow the same pattern: declare a topology, let Core reserve the slice. What changes per library is only what you declare it on. The libraries in the order most teams adopt them are Ray Serve, Ray Data, and Ray Train.
Ray Serve: vLLM-backed LLM serving on a TPU slice
Ray Serve gives you autoscaling, load-balancing, and multi-model composition, and on TPU it serves LLMs through vLLM. A multi-host model gets gang-scheduled onto a single slice using a single accelerator_config.topology field. The deployment is a standard Ray Serve config, with the topology parameter being the only TPU-specific addition. Once the model is on the slice, vLLM handles the inference the same way it does on GPUs, with the only difference being that the slice placement is enforced by Ray rather than assumed by the runtime.
Ray Data: feeding the slice without becoming the bottleneck
A fast accelerator is only as useful as the data you can keep flowing into it, and TPUs are fast enough that a naive loader becomes the bottleneck. The iter_jax_batches() function solves this. It hands you batches that are already JAX arrays and already device-sharded, so a training input pipeline or a large batch-inference job pulls straight from a Ray Data pipeline with no host-side NumPy-to-JAX copy stalling the step.
The pattern works as the input side of a JaxTrainer job, and it's just as useful on its own for offline batch inference over a big dataset on a TPU slice. The function landed recently in Ray, and the data step of the get-started example uses it for dataset prep and batch inference.
JaxTrainer: training on a TPU slice without writing slice code
Training used to be the confusing part of Ray on TPU, because of topology and having to account for the slice shape in your code. JaxTrainer addresses that. It brings Ray Train's training loop (checkpointing, fault tolerance, multi-slice scale-out) to JAX, Google's array and autodiff library and the native framework for TPU. You hand it a training function and a slice shape, and Ray launches one worker per host, wires them into a single mesh, and runs your function on each.
The single ScalingConfig is the only TPU-specific parameter. Everything else (checkpointing, fault tolerance, the training loop itself) is the same Ray Train API you would use on GPUs. The win is that the slice shape is a configuration parameter, not a constraint that propagates through your code. The job that took 200 lines of topology-aware JAX to express on the old Ray versions is now a 30-line training function with one config line for the slice shape.
The official images
As part of the first-class accelerator support, Ray now publishes official rayproject/ray:*-tpu images with the JAX/TPU stack (jax[tpu], flax, optax, orbax-checkpoint) and profiling tooling already installed. You base your image on the tagged -tpu one and skip the work of assembling a working TPU environment by hand. For teams that have spent days fighting a TPU environment, that single change is worth more than any of the library updates.
The wrap
In this two-part guide, the journey was: how Ray runs on TPU, then how to run AI workloads on it. Part 1 covered the one caveat (whole slice, all-or-nothing) and the two layers underneath (GKE through the Ray Operator add-on, Ray Core through slice_placement_group()). Part 2 put the AI libraries on top: Ray Serve for vLLM-backed inference, Ray Data for JAX-native batch feeding, JaxTrainer for distributed training. The same Ray you already use on GPUs, now on TPU.