AI21 revoluciona gestão de GPUs com Kueue e acelera cargas de trabalho em 83%

A gestão eficiente de recursos computacionais é um gargalo para o avanço da Inteligência Artificial, especialmente com a alta demanda por unidades de processamento gráfico (GPUs). A AI21, empresa de IA generativa, detalhou como superou esse desafio ao adotar o Kueue, um sistema de filas e agendamento nativo do Kubernetes.

Desenvolvido pelo grupo `kubernetes-sigs` (responsável por projetos que estendem o Kubernetes), o Kueue atua como um gerenciador de filas para cargas de trabalho de IA e Machine Learning. Ele decide quando um trabalho pode ser admitido, priorizado e, se necessário, preempcionado. Para a AI21, que opera uma frota de 10.000 GPUs em um cluster compartilhado do Google Cloud, a solução otimizou o uso dos recursos, substituindo a negociação manual por canais de mensagem por um sistema automatizado que garante justiça e eficiência no uso das GPUs.

A implementação do Kueue pela AI21 mostra uma evolução notável na gestão de GPUs. Inicialmente, a empresa usava um sistema v0.10 baseado em cotas e preempção, mas enfrentava desafios de justiça e fragmentação de recursos. O feedback da AI21 foi crucial para o projeto Kueue, levando à incorporação de duas funcionalidades essenciais: o Admission Fair Sharing (AFS) e o Topology Aware Scheduling (TAS).

Com o AFS, o Kueue passou a reequilibrar a ordem das filas com base no histórico de uso de chips por equipe, não apenas por prioridade ou timestamp. O TAS resolveu o problema da fragmentação, garantindo que os pods só fossem admitidos se houvesse um nó com capacidade física suficiente para eles. Essas melhorias levaram à versão v0.15 da implementação interna da AI21, automatizando a alocação e otimizando a utilização das GPUs de maneira sem precedentes.

A experiência da AI21 demonstra a importância da infraestrutura robusta para acelerar o desenvolvimento de IA. O uso de ferramentas como o Kueue reduz drasticamente o tempo de espera para cargas de trabalho críticas, como a AI21 obteve uma redução de 83%. Isso significa ciclos de pesquisa mais rápidos, menor tempo de inatividade para os engenheiros e um retorno sobre o investimento (ROI) significativamente maior em frotas de GPU, que são caras.

A automação da gestão de recursos libera equipes de pesquisa e desenvolvimento de tarefas operacionais, permitindo que se concentrem na inovação e na criação de modelos de IA mais sofisticados. A colaboração entre usuários e desenvolvedores de projetos open-source, como o Kueue, também garante que a ferramenta evolua para atender às necessidades reais do mercado.

Repositório oficial: kubernetes-sigs/kueue