Het Aurora-weermodel verfijnen met Ray op AKS

In dit artikel verzendt u een RayJob die het Microsoft Aurora-basismodel op regionale WeatherBench2-gegevens nauwkeurig kan afstemmen met behulp van LoRA. De taak wordt uitgevoerd op één A100 GPU, Kueue bepaalt toegang en de taak schrijft het controlepunt van de adapter en trainingsgegevens naar Azure Blob Storage.

Vereisten

• Infrastructuur geïmplementeerd volgens Infrastructuur voor Ray en Kueue implementeren op AKS.
• Kueue-wachtrijen die zijn geconfigureerd volgens Kueue-wachtrijen configureren voor Ray-workloads op AKS.
• Ten minste één A100 GPU beschikbaar in het cluster (in dit voorbeeld wordt één GPU gebruikt).
• Aurora WeatherBench2-gegevens geüpload naar blobopslag (de Terraform-module voor de infrastructuur uploadt deze gegevens automatisch).
• envsubst geïnstalleerd (gettext pakket in Linux, brew install gettext in macOS).

Omgevingsvariabelen instellen

Navigeer naar het voorbeeld van Aurora in de gekloonde opslagplaats en configureer de vereiste omgevingsvariabelen:

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

Het env.example-bestand stelt standaardinstellingen in, waaronder de Ray-image, de naam van de wachtrij en de trainingsparameters. De JOB_NAME variabele wordt automatisch gegenereerd met een tijdstempel.

Werklast indienen

Verzend de RayJob naar het cluster:

./submit.sh

Het script maakt een ConfigMap op basis van het Python trainingsscript, geeft de manifestsjabloon weer via envsubsten past deze toe op het cluster. Kueue geeft de taak toe wanneer gpu-quotum beschikbaar is in de geconfigureerde wachtrij.

Voortgang bijhouden

Zoek en exporteer de taaknaam als u zich in een nieuwe shell bevindt:

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

Bekijk de RayJob-status en Kueue-toelating:

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

Verwachte uitvoer wanneer de taak is voltooid:

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

Bekijk de workerlogs:

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

Resultaten controleren

Controleer de geüploade artefacten in blobopslagruimte:

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

Verwachte uitvoer:

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

Download en inspecteer de metrische gegevens van de 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

Verwachte uitvoer:

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

De loss_history matrix moet eindige waarden (niet NaN) bevatten, waarbij wordt bevestigd dat het gegevenspad en de modeltraining correct werken.

Configuratiegids

De hulpbronnen opschonen

Verwijder de RayJob en de bijbehorende ConfigMap:

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

Volgende stappen 

• Een LLM trainen met Ray
• Voer batchinferentie uit met Ray