Perfezionare il modello meteorologico Aurora con Ray in AKS

In questo articolo si invia un RayJob che ottimizza il modello di base meteo Microsoft Aurora sui dati WeatherBench2 regionali usando LoRA. Il processo viene eseguito su una singola GPU A100, Kueue controlla l'ammissione e il processo scrive i checkpoint dell'adattatore e le metriche di training in archiviazione BLOB di Azure.

Prerequisiti

• Infrastruttura distribuita seguendo Distribuire l'infrastruttura per Ray e Kueue in AKS.
• Code Kueue configurate seguendo Configurare le code Kueue per i carichi di lavoro Ray in AKS.
• Almeno una GPU A100 disponibile nel cluster (questo esempio usa una GPU).
• Dati di Aurora WeatherBench2 caricati nell'archiviazione BLOB (il modulo Terraform dell'infrastruttura carica automaticamente tali dati).
• envsubst installato (gettext pacchetto su Linux, brew install gettext su macOS).

Impostare le variabili di ambiente

Passare all'esempio di ottimizzazione di Aurora nel repository clonato e configurare le variabili di ambiente necessarie:

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

Il env.example file imposta le impostazioni predefinite, tra cui l'immagine Ray, il nome della coda e i parametri di training. La JOB_NAME variabile viene generata automaticamente con un timestamp.

Invia il carico di lavoro

Invia il RayJob al cluster:

./submit.sh

Lo script crea un oggetto ConfigMap dallo script di training Python, esegue il rendering del modello di manifesto tramite envsubste lo applica al cluster. Kueue accetta il processo quando la quota GPU è disponibile nella coda configurata.

Monitorare lo stato di avanzamento

Trovare ed esportare il nome del processo se si è in una nuova shell:

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

Controlla lo stato di RayJob e lo stato di ammissione di Kueue:

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

Output previsto al termine del processo:

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

Eseguire la parte finale dei log del ruolo di lavoro:

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

Verificare i risultati

Controllare gli artefatti caricati nell'archivio BLOB:

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

Output previsto:

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

Scaricare ed esaminare le metriche di training:

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

Output previsto:

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

La loss_history matrice deve contenere valori finiti (non NaN), confermando che il percorso dati e il training del modello funzionano correttamente.

Informazioni di riferimento sulla configurazione

Pulire le risorse

Eliminare RayJob e il relativo ConfigMap:

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

Passaggi successivi

• Eseguire il training di un LLM con Ray
• Eseguire l'inferenza batch con Ray