Run Ray on TPU, Part 2: Ray AI libraries

Summary

Running Ray on TPUs requires multi-host workloads to remain on one intact TPU slice whose hosts share a high-speed ICI link. This guide presents Ray's AI libraries on top of GKE and Ray Core: Ray Serve for serving, Ray Data's iter_jax_batches() for input pipelines and batch inference, and Ray Train's JaxTrainer for distributed JAX training. Serve uses accelerator_config.topology to defer slice placement to the replica, while iter_jax_batches() returns device-sharded JAX arrays and handles ragged final batches through drop, pad, or raise choices. JaxTrainer takes a training function and slice topology, launches one worker per host, and provides checkpointing, fault-tolerant restarts, and multi-slice scale-out; official TPU images and dashboard metrics round out the workflow.

Takeaways

1. 1

   Without accelerator_config.topology, Serve can place per-chip bundles across separate slices, leaving a multi-host deployment stuck in DEPLOYING because the workers cannot complete their first collective.

2. 2

   JAX should be imported inside train_loop_per_worker because each worker initializes JAX in its own TPU context; JaxTrainer's topology value describes the slice shape, not a chip count.

3. 3

   Official Ray TPU images include the JAX/TPU stack and profiling tools, while the Ray Dashboard shows TPU utilization and memory and can attach to per-worker JAX profiling.