Ajuste o modelo meteorológico Aurora com Ray no AKS

Neste artigo, você enviará um RayJob que ajusta o modelo meteorológico básico do Microsoft Aurora em dados regionais do WeatherBench2 usando LoRA. O trabalho é executado em uma única GPU A100, o Kueue controla a admissão e o trabalho grava o ponto de verificação do adaptador e as métricas de treinamento no Armazenamento de Blobs do Azure.

Pré-requisitos

• Infraestrutura implantada seguindo Implantar infraestrutura para Ray e Kueue no AKS.
• Filas do Kueue configuradas seguindo 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).
• Dados do Aurora WeatherBench2 enviados para o armazenamento de blob (o módulo de infraestrutura do Terraform faz o upload desses dados automaticamente).
• envsubst instalado (pacote gettext no Linux, brew install gettext no macOS).

Definir variáveis de ambiente

Navegue até o 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 arquivo env.example define valores padrão, incluindo a imagem do Ray, o nome da fila e os parâmetros de treinamento. A variável JOB_NAME é gerada automaticamente com um carimbo de data e hora.

Enviar a carga de trabalho

Submeta o RayJob ao cluster:

./submit.sh

O script cria um ConfigMap do script de treinamento Python, renderiza o modelo de manifesto por meio envsubste o aplica ao cluster. O Kueue admite o trabalho quando a cota da GPU estiver disponível na fila configurada.

Monitorar o progresso

Encontre e exporte o nome do trabalho se estiver em um novo shell:

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

Assista ao status do RayJob e à admissão de 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

Acompanhe os logs do worker:

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

Verificar os resultados

Verifique os artefatos enviados para o armazenamento de blobs:

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

Resultado esperado:

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

Baixe e inspecione as métricas de treinamento:

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

Resultado esperado:

{
    "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",
    ...
}

A loss_history matriz deve conter valores finitos (não NaN), confirmando que o caminho de dados e o treinamento do modelo estão funcionando corretamente.

Referência de configuração

Limpar os recursos

Exclua o RayJob e seu ConfigMap:

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

Próximas Etapas 

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