Microsoft abre o código do TauGrid, plataforma Kubernetes para cargas de IA em GPU

Times de plataforma que rodam IA no Kubernetes raramente executam apenas uma coisa. Eles mantêm um sistema de filas, um runtime distribuído, verificações de saúde das GPUs, dashboards e uma camada de scripts de submissão costurando tudo isso. A equipe de engenharia do Azure Kubernetes Service acaba de abrir o código do TauGrid, que colapsa todo esse trabalho de integração em uma única instalação via Helm.

O que é o TauGrid

O TauGrid é uma plataforma auto-hospedada para executar cargas de IA no Kubernetes. Ele combina cinco componentes que normalmente seriam integrados à mão: a CLI tau, o enfileiramento e admissão de workloads via Kueue, a orquestração de clusters Ray via KubeRay, o monitoramento de saúde das GPUs em nível de nó e a observabilidade de cluster e workloads. O código é escrito majoritariamente em Go.

A divisão de responsabilidades é o ponto central do design. Times de plataforma cuidam de workspaces, filas, perfis de computação, armazenamento, identidade e observabilidade. Pesquisadores trabalham a partir de um repositório e da CLI, submetendo cargas sem precisar configurar o Kubernetes diretamente.

Como um job percorre o sistema

Uma carga é descrita em um arquivo tau.yaml. O exemplo de treinamento em GPU publicado pela Microsoft roda um job PyTorch em uma única A100. Ao executar tau run, o TauGrid resolve a política da plataforma, renderiza um Job do Kubernetes ou um RayJob do KubeRay e o submete pelo Kueue. As seis etapas documentadas são: submissão, enfileiramento, execução, monitoramento, recuperação e evidência.

Os registros de evidência capturam metadados, configuração, logs, métricas e checkpoints — o que torna cada execução reproduzível e auditável. Quando vários times compartilham um cluster, seus jobs caem em uma ClusterQueue compartilhada do Kueue, que admite cada um por cota e prioridade e o coloca em GPUs saudáveis.

Instalação e limites atuais

A instalação é um chart Helm direto do Microsoft Container Registry:

helm install taugrid \
  oci://mcr.microsoft.com/aks/ai-runtime/helm/taugrid \
  --version 0.4.2 \
  --namespace tau-system \
  --create-namespace

Os pré-requisitos são um cluster Kubernetes 1.30+ com nós de GPU, kubectl e Helm 3.0 ou superior. Dois detalhes operacionais importam para quem avalia fora do Azure: o TauGrid não envia telemetria à Microsoft por padrão, mas a observabilidade via Azure Data Explorer ainda é específica do Azure. A intenção declarada é suportar Kubernetes em nuvem e on-premises sem dependência do Azure — e as contribuições nesse sentido estão abertas.