Ray on TPU vira suporte oficial a partir da versão 2.55

Se você já escala Python com Ray em GPUs, agora dá para rodar o mesmo código em TPU (o acelerador de IA do Google) usando APIs oficiais. É o que anuncia a primeira parte da série "Run Ray on TPU" no Google Developers Blog, escrita por Ivan Nardini e Spencer Peterson.

A mudança concreta: a partir do Ray 2.55, as TPUs do Google Cloud passam a ser um acelerador de primeira classe. Na prática, isso significa que elas entram no pipeline de release do Ray, com imagens pré-construídas oficiais e suporte nas bibliotecas centrais. Antes o caminho era "experimental": você montava seus próprios containers e dependia da comunidade. Para quem constrói no Brasil e usa Google Cloud, é a diferença entre um hack frágil e infra sustentável.

O conceito que muda tudo: slice

Para o Ray, uma TPU é apenas mais um recurso agendável, igual a uma GPU: você pede, o Ray coloca seu trabalho lá. Mas existe uma peculiaridade que vale entender antes de tudo.

Os chips de TPU vêm cabeados num grupo fixo chamado slice: várias máquinas host (VMs) cujos chips compartilham um link dedicado de alta velocidade, o ICI (Inter-Chip Interconnect). Um modelo multi-host precisa cair inteiro em um único slice, senão os workers não se enxergam e o job simplesmente trava.

A analogia da fonte é boa: pense num slice como uma única caixa multi-GPU onde o interconnect rápido (NVLink) só existe dentro da caixa. Espalhe workers entre duas caixas sem cabo entre elas e as operações coletivas (o all-reduce que sincroniza gradientes) nunca terminam. O treino trava. Em GPU você raramente pensa nisso; em TPU é crítico.

Outro termo onipresente é topology, a forma do slice, escrita como 4x4 para um slice de 16 chips. Você pede uma topologia, não uma contagem de chips.

Três camadas, zero código de placement manual

O sistema todo cabe em três camadas. Na esquerda, o código que você já escreve (as bibliotecas do Ray). No meio, o Ray Core, que reserva slices inteiros. Na direita, o GKE (Google Kubernetes Engine) gerenciado, que provisiona o hardware e o rotula para o Ray achar as fronteiras do slice.

O fluxo: o GKE provisiona um slice e rotula seus hosts, o Ray Core lê esses rótulos para reservar o slice de uma vez, e sua chamada de biblioteca fica no topo, declarando só uma topologia. Nenhum código de placement escrito à mão.

GKE e o webhook de TPU

Você roda Ray on TPU via GKE usando o Ray Operator add-on. Uma flag (--enable-ray-operator no Autopilot, ou --addons=RayOperator no Standard) instala duas coisas: o KubeRay, o operador Kubernetes que transforma YAML de RayCluster/RayService/RayJob em clusters rodando (o mesmo que você usaria com GPU), e o Ray TPU webhook, a parte específica de TPU, que carimba cada host com rótulos como ray.io/tpu-slice-name. Esse rótulo é o fio que amarra o sistema inteiro.

No manifesto, você pede TPUs com um nodeSelector para geração e topologia, mais a contagem de chips como recurso. Um slice multi-host adiciona só o campo numOfHosts.

Ray Core: uma função para conhecer

O suporte a TPU no Ray Core mora na API pública ray.util.tpu, e há basicamente uma função que importa: slice_placement_group(). Ela pega aquela garantia de "mantenha meus workers num slice intacto" e vira uma chamada única, reservando o slice de forma atômica (todos os hosts ou nenhum), casando com os rótulos do webhook:

from ray.util.tpu import slice_placement_group

spg = slice_placement_group(topology="4x4", accelerator_version="v6e")
ray.get(spg.placement_group.ready(), timeout=600)

Detalhe importante que a própria fonte faz questão de frisar: você raramente chama slice_placement_group sozinho. As bibliotecas de IA (Data, Train, Serve) chamam por você, então na prática você declara uma topologia e elas cuidam do slice. Só se recorre à função diretamente ao escrever uma carga distribuída customizada.

O trade-off honesto

Vale o alerta que o post traz: a API é pública mas marcada como alpha (@PublicAPI(stability="alpha")). Ou seja, é usável hoje, mas a superfície pode mudar entre releases. Quem for para produção precisa pesar isso, ainda mais em algo tão fundamental quanto o placement.

A parte 2 da série promete cobrir as bibliotecas de IA sobre TPU: serving de LLMs com vLLM, alimentação de slices com Ray Data e treino com JaxTrainer. Há também um exemplo executável no repositório kubernetes-engine-samples que provisiona cluster, serve, data e train num único slice v6e via Terraform, bom ponto de partida para testar antes de recomendar.

Fonte: Google Developers Blog