Run Ray on TPU, Part 2: Ray AI libraries

Ray Serve, Ray Data, and Ray Train are the focus of this second Google Developers Blog installment on running Ray on TPU. The piece shows how those higher-level Ray AI libraries sit on top of Google TPU slices so teams can run large AI workloads without hand-managing host topology, data pipelines, or multi-slice training glue.

On the serving path, Ray Serve relies on a simple topology configuration to gang-schedule large multi-host models onto TPU slices so related replicas start and scale as a unit. On the data path, Ray Data feeds accelerators directly with native JAX batches, cutting the usual host-to-device loading stalls that leave TPUs idle. For training, JaxTrainer coordinates distributed work across TPU slices and handles cross-slice coordination and checkpointing so operators do not wire that plumbing themselves.

That stack matters for engineers who already use Ray and want TPU capacity without rewriting orchestration. Topology-driven Serve scheduling reduces mis-scheduled multi-host deploys. Direct JAX batches in Ray Data keep the accelerator fed. JaxTrainer collapses cross-slice setup and checkpoint flow into the trainer path instead of custom scripts around every job.

In market terms, the post positions Ray’s Serve, Data, and Train libraries as the control plane for Google TPU slices rather than a one-off TPU tutorial. Builders comparing multi-host serving and multi-slice training stacks get a single Ray surface for topology, data feed, and distributed training instead of separate frameworks for each stage.

What to watch next is whether topology configs for Serve stay enough for larger multi-host models, whether native JAX batch feeding in Ray Data holds under heavier pipelines, and how far JaxTrainer’s automatic cross-slice coordination and checkpointing covers real multi-slice training runs without extra glue.