AI21 usa fila Kueue para gerenciar frota de 10.000 GPUs no Google Cloud

O problema: negociar GPUs por mensagens

A AI21 Labs opera um cluster GKE compartilhado no Google Cloud com cerca de 10.000 GPUs, usado por múltiplos times (pesquisadores, MLOps, engenharia e alinhamento). A empresa mantém um único pool compartilhado de propósito, em vez de um cluster por time: o Kubernetes facilita criar muitos clusters, mas dificulta compartilhar capacidade entre eles, e o pool único é o que mantém a utilização alta - o que, por sua vez, pressiona o agendamento interno.

Antes da mudança, a coordenação acontecia no canal interno #gpu-resources: quem precisava de um tipo de GPU totalmente ocupado pedia no canal e esperava que alguém liberasse capacidade. Isso funcionava enquanto havia folga, mas quando a utilização ficava perto de 100% - meta para uma frota reservada cara - cada pedido virava uma conversa de evição de soma zero. Os problemas observados foram tempo de engenharia desperdiçado em triagem e negociação, filas injustas, capacidade ociosa em alguns pontos enquanto outros times passavam fome, e dependência de "comandos mágicos" ad hoc.

Os quatro critérios e a avaliação de ferramentas

A equipe definiu quatro pilares para um sistema ideal de gerenciamento de GPUs: justiça (decisões sistemáticas conforme uso da conta, não primeiro-que-chega), ordem (fila refletindo prioridades de negócio, não só timestamp), eficiência (dispensar devs de deletar manualmente workloads ociosos) e integridade técnica (jobs multi-pod do tipo gang só iniciam quando todos os recursos estão prontos).

Os recursos nativos do Kubernetes ficaram aquém: PriorityClass opera no nível de pod, sem noção de gang, podendo matar um job inteiro para liberar uma GPU; ResourceQuota não permite empréstimo entre times quando há capacidade ociosa; LimitRange é guardrail, não agendador; e Jobs/Indexed Jobs não têm semântica all-or-nothing - 3 de 4 pods agendados continuam rodando e queimando GPUs.

A AI21 avaliou então três agendadores de batch - Apache YuniKorn, Volcano e Kueue - e escolheu o Kueue, um sistema de fila de jobs open source nativo do Kubernetes, pela simplicidade e integração, apesar de o Volcano ter mais funcionalidades na época.

Implementação e limites da v0.10

O Kueue admite, prioriza e preempta workloads automaticamente, com ClusterQueues, Workload Priority Classes e semântica all-or-nothing integrada aos Indexed Jobs. Na primeira versão (v0.10), a empresa estruturou três camadas de capacidade: garantida (quota por time, nunca preemptada), preemptível (pool compartilhado com fair sharing, que empresta capacidade ociosa e devolve na retomada) e on-demand (provisionamento dinâmico quando tudo está cheio). Do lado do desenvolvedor, três rótulos - gasto (baixo/alto), preemptibilidade e prioridade - mais single-node ou multi-node determinavam a fila de destino.

Duas falhas surgiram. A primeira foi a lacuna de justiça: à época não existia Admission Fair Sharing (AFS), impedindo conciliar fair scheduling por chips entre tenants com a regra de não preempção entre times. A AI21 levou a demanda ao time do Kueue, que lançou a AFS: ela preempte workloads de times que famintizam outros e reordena a fila priorizando times de menor uso histórico.

A segunda foi fragmentação: cada máquina tem 8 GPUs, e uma carga que exige 8 em um único nó não cabe se a capacidade estiver espalhada como 1+1+4+2 em quatro nós, embora a matemática de quotas indicasse 8 GPUs livres. A solução foi Topology Aware Scheduling (TAS): com uma Topology customizada, o Kueue recusa admission que não cabe em um único nó e agenda com o algoritmo LeastFreeCapacity, empacotando primeiro nos nós mais utilizados.

Redesign (v0.15) e resultados

Com AFS e TAS, e a pedidos novos dos stakeholders - preempção por prioridade dentro do pool single-node preemptível e uma nova faixa multi-node best-effort de prioridade baixa -, a equipe redesenhou para a v0.15. O pool preemptível single-node passou a se dividir por prioridade, e a interface do desenvolvedor permaneceu igual, mudando apenas o roteamento.

Hoje o canal #gpu-resources está arquivado e todos os workloads da AI21 são direcionados automaticamente às filas corretas. Os números reportados: intervenções manuais semanais de 20 para 0; jobs "zumbis" e parcialmente alocados eliminados; fragmentação de 15% para 8%; horas semanais de negociação em canais reduzidas a zero; e tempo médio de fome de um hero job de 72 para 12 horas. O Google Cloud publicou o caso destacando 83% de redução no time-to-start de workloads de alta prioridade.

Lições

As conclusões da AI21: agendar no nível de workload, não de pod, pois treinamento distribuído exige admission all-or-nothing; uma fila só vira agendador quando entende a topologia do cluster; e manter a frota propositalmente perto de 100% de utilização é seguro quando existe uma faixa preemptível que empresta capacidade ociosa e a devolve na retomada.