Implementar infraestrutura para Ray e Kueue no AKS

Neste artigo, implementa-se a infraestrutura necessária para executar cargas de trabalho de IA do Ray com Kueue no Azure Kubernetes Service (AKS). O módulo Terraform prevê um cluster AKS com pools de nós GPU, instala o operador KubeRay e o controlador Kueue via 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 é excluído dos contratos de nível de serviço do AKS, da garantia limitada e do suporte do Azure. Ao usar a tecnologia de código aberto ao lado do AKS, consulte as opções de suporte disponíveis nas respetivas comunidades e mantenedores do projeto para desenvolver um plano.

A Microsoft assume a responsabilidade pela criação dos pacotes de código aberto que implantamos no AKS. Essa responsabilidade inclui ter a propriedade completa do processo de compilação, digitalização, assinatura, validação e correção rápida, juntamente com o controlo dos binários nas imagens de contentor. Para obter mais informações, consulte Gestão de vulnerabilidades para AKS e Cobertura de suporte AKS.

Pré-requisitos

  • Uma assinatura do Azure. Se não tiver uma, crie uma conta gratuita.
  • CLI do Azure versão 2.70 ou posterior, autenticado com az login. Se tiveres várias subscrições, define a correta com az account set --subscription <subscription-id>.
  • Terraform versão 1.6 ou posterior.
  • Kubectl versão 1.28 ou posterior. Instale com az aks install-cli.
  • Python versão 3.10 ou posterior (necessária para geração de dados Aurora).
  • Quota de GPU na sua região Azure de destino. A configuração predefinida aprovisiona 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é ao módulo de infraestrutura:

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

Instalar dependências em Python

A etapa de upload de dados do Aurora requer pacotes em Python. Instale-os antes de implementar:

pip install -r terraform/requirements-generator.txt

Verifique a quota de GPU

Verifica se a tua subscrição tem quota para o SKU da GPU na tua região de destino. Substitua eastus2 pela sua região preferida:

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

Nota

Se não tiver quota de GPUs ou quiser validar a infraestrutura sem GPUs, implemente com gpu_enabled=false:

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

A via exclusivamente de CPU é útil para validar a infraestrutura e a configuração da fila. As cargas de trabalho que solicitam GPUs permanecem pendentes.

Implementar 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 ativada
  • Um pool de nós do sistema e um pool de nós de GPU (quando gpu_enabled=true)
  • Operador KubeRay e controlador Kueue (Lançamentos do Helm da MCR)
  • DaemonSet de monitorização da GPU (quando gpu_enabled=true)
  • Conta do Armazenamento de Blobs do Azure com contentores aurora e llm-pipeline
  • Identidade gerida atribuída pelo utilizador com papel de Contribuidor de Dados do Blob de Armazenamento
  • Credencial de identidade federada para a ray-workload ServiceAccount
  • Conjuntos de dados pré-preparados (dados Aurora WeatherBench2 e conjunto de dados viggo NLG)

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

Ligar ao cluster

Recupere as credenciais do seu novo cluster AKS:

eval "$(terraform output -raw get_credentials_command)"

Verificar a implantação

Confirme que os operadores estão a funcionar e que os nós estão disponíveis:

  1. Verifique o operador KubeRay:

    kubectl -n kuberay-system get pods
    

    Produção esperada:

    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
    

    Produção esperada:

    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. Verifique a monitorização da GPU (apenas quando gpu_enabled=true):

    kubectl -n gpu-monitoring get daemonsets
    

    Produção esperada:

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

Limpeza de recursos

Para eliminar todos os recursos do Azure que criaste usando este módulo:

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

Passo seguinte