Ajuster le modèle météo Aurora avec Ray sur AKS

Dans cet article, vous soumettez un RayJob qui affine le modèle de fondation météorologique Microsoft Aurora sur des données météorologiques régionales de WeatherBench2 à l’aide de LoRA. La tâche s’exécute sur un seul GPU A100, Kueue gère l’admission, et la tâche écrit le point de contrôle de l’adaptateur et les métriques d’entraînement dans Stockage Blob Azure.

Prérequis

• Infrastructure déployée en suivant Déployer l’infrastructure pour Ray et Kueue sur AKS.
• Les files d’attente Kueue configurées conformément à Configurer les files d’attente Kueue pour les charges de travail Ray sur AKS.
• Au moins un GPU A100 disponible dans le cluster (cet exemple utilise un GPU).
• Données Aurora WeatherBench2 chargées dans le stockage blob (le module d’infrastructure Terraform charge automatiquement ces données).
• envsubst installé (gettext package sur Linux, brew install gettext sur macOS).

Définir des variables d’environnement

Accédez à l’exemple d’optimisation Aurora dans le référentiel cloné et configurez les variables d’environnement requises :

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

Le env.example fichier définit les valeurs par défaut, notamment l’image Ray, le nom de la file d’attente et les paramètres d’entraînement. La JOB_NAME variable est générée automatiquement avec un horodatage.

Soumettre la charge de travail

Soumettez le RayJob au cluster :

./submit.sh

Le script crée un ConfigMap à partir du script d’entraînement Python, affiche le modèle de manifeste via envsubst, et l’applique au cluster. Kueue accepte la tâche lorsque le quota de GPU est disponible dans la file d’attente configurée.

Surveiller la progression

Trouvez et exportez le nom de la tâche si vous êtes dans un nouveau shell :

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

Surveillez l’état de RayJob et l’admission de Kueue :

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

Sortie attendue lorsque le travail se termine :

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

Suivez les journaux du worker :

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

Vérifier les résultats

Vérifiez les artefacts téléversés dans le stockage Blob :

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

Sortie attendue :

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

Téléchargez et inspectez les métriques d’entraînement :

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

Sortie attendue :

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

Le loss_history tableau doit contenir des valeurs finies (et non NaN), confirmant que le chemin d’accès aux données et l’entraînement du modèle fonctionnent correctement.

Référence de configuration

Nettoyer les ressources

Supprimez RayJob et son ConfigMap :

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

Étapes suivantes

• Entraîner un LLM avec Ray
• Exécuter l’inférence par lot avec Ray