Ajuste o modelo meteorológico da Aurora com o Ray no AKS

Neste artigo, submetes um RayJob que afina o modelo de fundação meteorológica do Microsoft Aurora em dados regionais do WeatherBench2 usando LoRA. A tarefa é executada numa única GPU A100, Kueue controla a admissão, e a tarefa grava o checkpoint do adaptador e as métricas de treino no Armazenamento de Blobs do Azure.

Pré-requisitos

• Infraestrutura implementada seguindo Implementar a infraestrutura para Ray e Kueue no AKS.
• Filas do Kueue configuradas de acordo com Configurar filas do Kueue para cargas de trabalho do Ray no AKS.
• Pelo menos uma GPU A100 disponível no cluster (este exemplo usa uma GPU).
• Os dados do Aurora WeatherBench2 são carregados para armazenamento em blob (o módulo Terraform da infraestrutura carrega estes dados automaticamente).
• envsubst instalado (gettext pacote para Linux, brew install gettext no macOS).

Definir variáveis de ambiente

Navegue até ao exemplo de ajuste fino do Aurora no repositório clonado e configure as variáveis de ambiente necessárias:

cd <path-to-cloned-repo>/AKS/examples/kueue-and-ray-on-aks/3-workloads/aurora-finetune
export AZURE_STORAGE_ACCOUNT_NAME=$(terraform -chdir=../../1-infrastructure/terraform output -raw storage_account_name)
source env.example

O env.example ficheiro define os valores predefinidos incluindo a imagem Ray, o nome da fila e os parâmetros de treino. A JOB_NAME variável é gerada automaticamente com um carimbo temporal.

Enviar a carga de trabalho

Submeter o RayJob ao cluster:

./submit.sh

O script cria um ConfigMap a partir do script de treino em Python, renderiza o modelo de manifesto via envsubst, e aplica-o ao cluster. Kueue aceita a tarefa quando a quota de GPUs está disponível na fila configurada.

Monitorizar progresso

Localize e exporte o nome da tarefa caso esteja numa nova shell:

export JOB_NAME=$(kubectl -n ray get rayjob --no-headers -o custom-columns=":metadata.name" | grep aurora-finetune)

Monitorize o estado do RayJob e a admissão do Kueue:

kubectl -n ray get rayjob ${JOB_NAME} -w
kubectl -n ray get workload -w

Saída esperada quando o trabalho for concluído:

NAME                         JOB STATUS   DEPLOYMENT STATUS   START TIME             END TIME               AGE
aurora-finetune-xxxxxxxxxx   SUCCEEDED    Complete            2026-01-01T00:00:00Z   2026-01-01T00:08:00Z   8m

NAME                                      QUEUE     RESERVED IN     ADMITTED   FINISHED   AGE
rayjob-aurora-finetune-xxxxxxxxxx-xxxxx   default   cluster-queue   True       True       8m

Seguir os registos dos trabalhadores:

kubectl -n ray logs -l ray.io/cluster=$(kubectl -n ray get rayjob ${JOB_NAME} -o jsonpath='{.status.rayClusterName}') -f

Verificar resultados

Verifique os artefactos carregados no armazenamento de blobs:

az storage blob list -c aurora --prefix checkpoints/${JOB_NAME}/ \
  --account-name ${AZURE_STORAGE_ACCOUNT_NAME} --auth-mode login -o table

Produção esperada:

Name                                                    Blob Type    Blob Tier    Length    Content Type
------------------------------------------------------  -----------  -----------  --------  ------------------------
checkpoints/<job-name>/last.safetensors                 BlockBlob    Hot          11432328  application/octet-stream
checkpoints/<job-name>/train-metrics.json               BlockBlob    Hot          985       application/json

Descarregue e inspecione as métricas de formação:

az storage blob download -c aurora \
  -n checkpoints/${JOB_NAME}/train-metrics.json \
  --account-name ${AZURE_STORAGE_ACCOUNT_NAME} --auth-mode login \
  --file /tmp/train-metrics.json
cat /tmp/train-metrics.json | python3 -m json.tool

Produção esperada:

{
    "final_loss": 24051.97,
    "initial_loss": 24183.10,
    "loss_history": [24183.10],
    "loss_improvement": 131.13,
    "max_steps": 1,
    "trainable_parameters": 2850816,
    "gpu_name": "NVIDIA A100-SXM4-80GB",
    ...
}

O loss_history array deve conter valores finitos (não NaN), confirmando que o caminho dos dados e o treino do modelo estão a funcionar corretamente.

Referência de configuração

Limpeza de recursos

Apague o RayJob e o seu ConfigMap:

kubectl -n ray delete rayjob ${JOB_NAME}
kubectl -n ray delete configmap ${CONFIGMAP_NAME}

Passos seguintes

• Treine um LLM com o Ray
• Executar inferência em lote com o Ray