Por que executar Ray em TPUs?

Se você já usa Ray para escalar workloads Python em GPUs, agora pode rodar o mesmo código em TPUs do Google Cloud com suporte oficial a partir do Ray 2.55. Nada de containers experimentais ou hacks da comunidade—o Ray agora trata TPUs como aceleradores de primeira classe, com imagens pré-construídas e suporte completo nas bibliotecas principais.

Mas tem um detalhe: TPUs são fisicamente conectadas em grupos fixos chamados slices. Cada slice consiste em várias VMs host cujos chips são conectados via um link de alta velocidade chamado ICI (Inter-Chip Interconnect). Se seus workers não estiverem no mesmo slice, eles não conseguem se comunicar e o treinamento trava.

Neste tutorial, você vai aprender a configurar Ray em TPU usando GKE e usar a API slice_placement_group para garantir que seus workers fiquem em um slice intacto.

Ray on TPU architecture diagram showing GKE and Ray Core layers Software Concept Art

Configurando GKE com Ray Operator

Primeiro, crie um cluster GKE com o add-on Ray Operator. Você pode escolher Autopilot (totalmente gerenciado) ou Standard (você gerencia os node pools).

Cluster Autopilot

gcloud container clusters create-auto CLUSTER \
  --enable-ray-operator --location=LOCATION

Cluster Standard

gcloud container clusters create CLUSTER \
  --addons=RayOperator --location=LOCATION &&

gcloud container node-pools create v6e-16-slice \
  --cluster=CLUSTER \
  --location=LOCATION \
  --machine-type=ct6e-standard-4t \
  --tpu-topology=4x4 \
  --num-nodes=4

O Ray Operator instala dois componentes principais:

  • KubeRay: Transforma YAMLs de RayCluster, RayService e RayJob em clusters Ray rodando.
  • Ray TPU webhook: Rotula cada host TPU com ray.io/tpu-slice-name para que o Ray identifique quais máquinas fazem parte do mesmo slice.

Agora, solicite TPUs no seu manifest RayCluster usando nodeSelector e limites de recursos:

# dentro de workerGroupSpec do RayCluster
nodeSelector:
  cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice # geração da TPU
  cloud.google.com/gke-tpu-topology: "4x4" # formato do slice
# ... solicite chips via o limite de recurso google.com/tpu
numOfHosts: 4 # multi-host: número de VMs host no slice

Aplique o manifest, e o GKE provisiona o slice, o webhook rotula, e o Ray está pronto para agendar seu trabalho.

Python code for slice_placement_group in a terminal Developer Related Image

Usando slice_placement_group para Reserva Atômica

O suporte a TPU no Ray Core está na API pública ray.util.tpu. A função chave é slice_placement_group(), que reserva um slice inteiro atomicamente (todos os hosts ou nenhum) combinando com os rótulos do webhook.

Aqui está um exemplo completo:

from ray.util.tpu import slice_placement_group
from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy

# Reserva um slice v6e 4x4 inteiro (16 chips em 4 hosts), atomicamente
spg = slice_placement_group(topology="4x4", accelerator_version="v6e")
ray.get(spg.placement_group.ready(), timeout=600)

@ray.remote(resources={"TPU": 4})
def worker(rank, world):
    # Seu código de treinamento ou inferência distribuída aqui
    pass

tasks = [
    worker.options(
        scheduling_strategy=PlacementGroupSchedulingStrategy(
            placement_group=spg.placement_group
        )
    ).remote(rank=i, world=spg.num_hosts)
    for i in range(spg.num_hosts)
]

# Aguarda todos os tasks completarem
ray.get(tasks)

Importante: Na prática, você raramente chama slice_placement_group diretamente. As bibliotecas Ray AI (Data, Train, Serve) chamam isso para você. Você só precisa disso ao escrever workloads distribuídos personalizados.

Limitações e Cuidados

  • A API é marcada como @PublicAPI(stability="alpha"), então pode mudar entre releases.
  • Slices multi-host exigem planejamento cuidadoso; se seus workers não caírem no mesmo slice, seu job vai travar.
  • A topologia da TPU deve corresponder ao formato físico do slice (ex: 4x4 para 16 chips).

TPU slice with multiple host VMs connected via ICI Technical Structure Concept

Conclusão e Próximos Passos

Agora você tem um modelo mental sólido: um slice de TPU deve permanecer intacto, o GKE provisiona e rotula, e o Ray Core reserva como uma unidade—então você nunca escreve código de placement manualmente.

Para aprofundar, confira a Parte 2 desta série onde exploramos o uso das bibliotecas Ray AI em TPU para servir LLMs com vLLM, alimentar slices com Ray Data e treinar com JaxTrainer.

Também não perca estas leituras relacionadas:

Bons builds!

Este conteúdo foi elaborado com o auxílio de ferramentas de IA, com base em fontes confiáveis, e revisado pela nossa equipe editorial antes da publicação. Não substitui o aconselhamento de um profissional especializado.