Implementación de la infraestructura para Ray y Kueue en AKS

En este artículo, implementará la infraestructura necesaria para ejecutar cargas de trabajo de Ray AI con Kueue en Azure Kubernetes Service (AKS). El módulo Terraform aprovisiona un clúster de AKS con grupos de nodos de GPU, instala el operador KubeRay y el controlador kueue a través de Helm, crea Azure Blob Storage con conjuntos de datos preconfigurados y configura la identidad de carga de trabajo para el acceso seguro.

Importante

El software de código abierto se menciona en toda la documentación y ejemplos de AKS. El software que implemente se excluye de los contratos de nivel de servicio de AKS, la garantía limitada y el soporte técnico de Azure. A medida que usa la tecnología de código abierto junto con AKS, consulte las opciones de soporte técnico disponibles en las comunidades y los mantenedores de proyectos respectivos para desarrollar un plan.

Microsoft asume la responsabilidad de crear los paquetes de código abierto que implementamos en AKS. Esa responsabilidad incluye ser plenamente responsable del proceso de compilación, escaneo, firma, validación y corrección rápida, junto con el control de los binarios en las imágenes de contenedor. Para obtener más información, vea Administración de vulnerabilidades para AKS y Cobertura del soporte técnico de AKS.

Prerequisites

  • Una suscripción a Azure. Si no tiene una, cree una cuenta gratuita.
  • CLI de Azure versión 2.70 o posterior, autenticado con az login. Si tiene varias suscripciones, establezca la correcta con az account set --subscription <subscription-id>.
  • Terraform versión 1.6 o posterior.
  • kubectl versión 1.28 o posterior. Instale con az aks install-cli.
  • Python versión 3.10 o posterior (necesaria para la generación de datos Aurora).
  • Cuota de GPU en la región de Azure de destino. La configuración predeterminada aprovisiona un Standard_ND96amsr_A100_v4 nodo (8×A100 80 GB de GPU).

Clonación del repositorio

Clone el repositorio de Azure/AKS y vaya al módulo de infraestructura:

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

Instalación de dependencias de Python

El paso de carga de datos de Aurora requiere los paquetes de Python. Instálelos antes de implementar:

pip install -r terraform/requirements-generator.txt

Comprobación de la cuota de GPU

Compruebe que la suscripción tiene cuota para la SKU de GPU en la región de destino. Reemplace por eastus2 su región preferida:

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

Note

Si no tiene cuota de GPU o quiere validar la infraestructura sin GPU, implemente con gpu_enabled=false:

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

La ruta solo con CPU resulta útil para validar la infraestructura y la configuración de la cola. Las cargas de trabajo que solicitan GPUs permanecen pendientes.

Implementación con Terraform

Inicialice y aplique el módulo de Terraform:

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

El módulo de Terraform crea:

  • Un clúster de AKS con emisor OIDC e identidad de carga de trabajo habilitados
  • Un grupo de nodos del sistema y un grupo de nodos de GPU (cuando gpu_enabled=true)
  • Operador KubeRay y controlador Kueue (lanzamientos de Helm desde MCR)
  • DaemonSet de supervisión de GPU (cuando gpu_enabled=true)
  • Cuenta de Azure Blob Storage con contenedores aurora y llm-pipeline
  • Identidad administrada asignada por el usuario con el rol Colaborador de datos de Storage Blob
  • Credencial de identidad federada para ServiceAccount ray-workload
  • Conjuntos de datos preconfigurados (datos Aurora WeatherBench2 y conjunto de datos NLG viggo)

Para obtener la configuración completa del módulo de Terraform, consulte el directorio 1-infrastructure del repositorio.

Conectarse al clúster

Recupere las credenciales del nuevo clúster de AKS:

eval "$(terraform output -raw get_credentials_command)"

Comprobación de la implementación

Confirme que los operadores están en ejecución y que los nodos están disponibles:

  1. Compruebe el 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. Compruebe el 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. Compruebe los nodos:

    kubectl get nodes
    

    Salida esperada (con 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 la monitorización de la GPU (solo cuando 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
    

Limpieza de recursos

Para eliminar todos los recursos de Azure creados mediante este módulo:

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

Paso siguiente