Implantar infraestrutura para Ray e Kueue no AKS

Neste artigo, você implanta a infraestrutura necessária para executar cargas de trabalho de IA do Ray com o Kueue no AKS (Serviço de Kubernetes do Azure). O módulo Terraform provisiona um cluster do AKS com pools de nós de GPU, instala o operador KubeRay e o controlador Kueue por meio do Helm, cria Armazenamento de Blobs do Azure com conjuntos de dados pré-preparados e configura a identidade da carga de trabalho para acesso seguro.

Importante

O software de código aberto é mencionado em toda a documentação e amostras do AKS. O software que você implanta está excluído dos contratos de nível de serviço do AKS, garantia limitada e suporte do Azure. Ao usar tecnologia de código aberto junto com o AKS, consulte as opções de suporte disponíveis nas comunidades e mantenedores de projetos respectivos para desenvolver um plano.

A Microsoft assume a responsabilidade por criar os pacotes de código aberto que implantamos no AKS. Essa responsabilidade inclui ter propriedade completa do processo de criação, verificação, sinalização, validação e hotfix, junto com o controle sobre os binários em imagens de contêiner. Para obter mais informações, confira Gerenciamento de vulnerabilidades para o AKS e Cobertura de suporte do AKS.

Pré-requisitos

  • Uma assinatura do Azure. Se você não tiver uma, crie uma conta gratuita.
  • CLI do Azure versão 2.70 ou posterior, autenticada com az login. Se você tiver várias assinaturas, defina a correta com az account set --subscription <subscription-id>.
  • Terraform versão 1.6 ou posterior.
  • kubectl versão 1.28 ou posterior. Instalar com az aks install-cli.
  • Python versão 3.10 ou posterior (necessária para a geração de dados do Aurora).
  • Cota de GPU na região de Azure de destino. A configuração padrão provisiona um nó Standard_ND96amsr_A100_v4 (8 GPUs A100 de 80 GB).

Clonar o repositório

Clone o repositório Azure/AKS e navegue até o módulo de infraestrutura:

git clone https://github.com/Azure/AKS
cd AKS/examples/kueue-and-ray-on-aks/1-infrastructure

Instalar dependências de Python

A etapa de carregamento de dados do Aurora requer Python pacotes. Instale-os antes de implantar:

pip install -r terraform/requirements-generator.txt

Verificar cota de GPU

Verifique se sua assinatura tem cota para o SKU de GPU em sua região de destino. Substitua eastus2 por sua região preferencial:

LOCATION=<your-azure-region>
az vm list-usage --location "$LOCATION" --output table | grep -i "NDAMSv4_A100"

Observação

Se você não tiver cota de GPU ou quiser validar a infraestrutura sem GPUs, implante com gpu_enabled=false:

terraform apply -var="subscription_id=<your-subscription-id>" -var="gpu_enabled=false"

O caminho somente CPU é útil para validar a infraestrutura e a configuração da fila. As cargas de trabalho que solicitam GPUs permanecem no estado Pendente.

Implantação com Terraform

Inicialize e aplique o módulo Terraform:

cd terraform
terraform init
terraform apply -var="subscription_id=<your-subscription-id>"

O módulo Terraform cria:

  • Um cluster AKS com emissor OIDC e identidade de carga de trabalho habilitados.
  • Um pool de nós do sistema e um pool de nós de GPU (quando gpu_enabled=true)
  • Operador KubeRay e controlador Kueue (versões Helm do MCR).
  • DaemonSet de monitoramento de GPU (quando gpu_enabled=true)
  • Conta de armazenamento do Armazenamento de Blobs do Azure com contêineres aurora e llm-pipeline
  • Identidade gerenciada atribuída a um usuário com a função Colaborador de Dados do Blob de Armazenamento
  • Credencial de identidade federada para o ray-workload ServiceAccount
  • Conjuntos de dados pré-preparados (dados do Aurora WeatherBench2 e conjunto de dados viggo NLG)

Para obter a configuração completa do módulo Terraform, consulte o diretório 1-infrastructure no repositório.

Conectar-se ao cluster

Recupere as credenciais do seu novo cluster do AKS:

eval "$(terraform output -raw get_credentials_command)"

Verificar a implantação

Confirme se os operadores estão em execução e se os nós estão disponíveis:

  1. Verifique o operador KubeRay:

    kubectl -n kuberay-system get pods
    

    Resultado esperado:

    NAME                                READY   STATUS    RESTARTS   AGE
    kuberay-operator-xxxxxxxxxx-xxxxx   1/1     Running   0          5m
    
  2. Verifique o controlador Kueue:

    kubectl -n kueue-system get pods
    

    Resultado esperado:

    NAME                                            READY   STATUS    RESTARTS   AGE
    kueue-controller-manager-xxxxxxxxxx-xxxxx       1/1     Running   0          5m
    
  3. Verifique os nós:

    kubectl get nodes
    

    Saída esperada (com gpu_enabled=true):

    NAME                              STATUS   ROLES    AGE   VERSION
    aks-gpupool-xxxxxxxx-vmssxxxxxx   Ready    <none>   5m    v1.35.x
    aks-system-xxxxxxxx-vmssxxxxxx    Ready    <none>   10m   v1.35.x
    aks-system-xxxxxxxx-vmssxxxxxx    Ready    <none>   10m   v1.35.x
    
  4. Verificar o monitoramento de GPU (somente quando gpu_enabled=true):

    kubectl -n gpu-monitoring get daemonsets
    

    Resultado esperado:

    NAME                    DESIRED   CURRENT   READY   UP-TO-DATE   AVAILABLE   NODE SELECTOR   AGE
    gpu-monitoring-a100     1         1         1       1            1           <none>          5m
    

Limpar os recursos

Para excluir todos os recursos Azure que você criou usando este módulo:

terraform destroy -var="subscription_id=<your-subscription-id>"

Próxima etapa