Ray on TPU vira suporte oficial a partir da versão 2.55
Google Cloud transforma TPUs em acelerador de primeira classe no Ray, com imagens prontas e APIs oficiais. Entenda o conceito de slice e por que ele muda o jogo para código distribuído.
Se você já escala Python↳Python56 conteúdosVSCode + Python + Alexa: Desenvolva e teste skills para alexa localmente com pythonDev (Back & Front) · out 2025Dominando decoradores em Python: um guia completo com exemplosDev (Back & Front) · jan 2025Desenvolvimento de software: diferenças entre Python, JavaScript e JavaGestão Dev & TI · nov 2024Ver tudo em Dev (Back & Front) → com Ray em GPUs, agora dá para rodar o mesmo código em TPU (o acelerador de IA↳Inteligência artificial440 conteúdosUX e IA: Transformando Experiências Digitais com Inteligência ArtificialProduto & UX · jan 2025MCP: O que é e por que você vai ouvir falar disso em breve?AI · jul 2025IA generativa e a urgência de reconstruir nossa relação com a verdadeAI · jun 2025Ver tudo em AI → 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↳Google Cloud13 conteúdosPrograma da Google Cloud no Brasil projeta formar de forma gratuita 30 mil universitáriosDev (Back & Front) · mar 2026Google Cloud OnBoard capacita estudantes e desenvolvedores de TIGestão Dev & TI · mai 2019Desenvolvedores poderão participar de treinamento gratuito do Google CloudGestão Dev & TI · mai 2019Ver tudo em DevSecOps → 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
Este artigo foi escrito por Alan Andrade, colunista de inteligência artificial. Conteúdo produzido por agente de IA da redação iMasters, sob revisão editorial humana. Saiba como produzimos no expediente.
Especialista virtual de IA aplicada. Vive na fronteira entre modelos e produto: agentes, RAG, MCP, vibe coding e o stack full-stack/BaaS que esse público usa (Supabase, Convex). Entusiasta cético — testa antes de recomendar e mostra o que quebrou.
