TPU で Ray を実行する — Part 2: Ray AI ライブラリRun Ray on TPU, Part 2: Ray AI libraries

Google's developer blog has continued its tutorial series on pairing Ray, the open-source distributed computing framework, with Tensor Processing Units (TPUs), Google's custom machine learning accelerators. This second part moves beyond basic cluster setup to focus on the Ray AI libraries — chiefly Ray Train and Ray Data — and how they can drive large, distributed training and data-processing jobs on TPU hardware. The topic matters because it connects a widely used, Python-native orchestration layer with accelerators that many teams find harder to program than GPUs, potentially lowering the barrier to scaling machine learning workloads on Google Cloud.

Ray originated at UC Berkeley's RISELab and is now maintained largely by Anyscale. Its design separates a low-level engine, Ray Core, from a set of higher-level libraries aimed at machine learning. Ray Train handles distributed model training, Ray Data manages streaming ingestion and preprocessing, Ray Tune covers hyperparameter tuning, and Ray Serve addresses deployment. Part 1 of this series, as its framing implies, dealt with getting Ray running on TPUs at the Core level; this installment builds on that foundation to show the libraries most developers actually use day to day.

The article's technical core is the coordination of multi-host TPU workloads. A TPU slice or pod comprises multiple host VMs, each attached to several TPU chips, with hosts linked by a high-speed inter-chip interconnect. Training a single model across such a slice follows an SPMD pattern, in which the same program runs on every host and the chips exchange gradients through collective communication. Ray's tasks and actors provide the scheduling and process-management primitives to launch and supervise these per-host workers, and the guide explains how TPU capacity is surfaced as schedulable resources inside a Ray cluster so that jobs land on the intended hardware.

For training, Ray Train typically wraps a user-defined function and replicates it across a group of workers. On TPUs this generally means each worker runs a JAX or PyTorch/XLA process that participates in the collective operations spanning the slice. Ray handles placement, fault signaling, and lifecycle management, while the underlying framework performs the actual accelerator computation. Ray Data complements this by building input pipelines that run preprocessing on CPU resources and stream batches to the accelerators, an arrangement intended to keep expensive TPU chips busy rather than idle