PLATFORMS Signal 174
Run Ray on TPU, Part 2: Ray AI libraries
↗
The critical detail is the `topology` field in Ray Serve—omit it and a multi-host model's workers can scatter across slices with no ICI between them, leaving the deployment stuck in DEPLOYING indefinitely without a crash. Ray Data's `iter_jax_batches()` and Ray Train's JaxTrainer similarly abstract slice-aware placement, removing the need to hand-code for TPU slice shapes.
Written by elseif from the cluster below · every claim links back to a sourceThe three things worth knowing
Ray Serve requires a `topology` field (e.g., "4x4") to gang-schedule multi-host models on a single TPU slice; without it, per-chip bundles can split across slices and hang forever.
Ray Data's `iter_jax_batches()` delivers device-sharded JAX arrays directly, bypassing host-side copy bottlenecks and handling ragged final batches with explicit drop, pad, or raise options.
JaxTrainer brings Ray Train's checkpointing, fault tolerance, and multi-slice scale-out to JAX workloads without requiring manual topology accounting in user code.
THE CLUSTER