Dostrojenie modelu pogodowego Aurora z Rayem w usłudze AKS

W tym artykule przesyłasz zadanie RayJob, które dostraja bazowy model pogodowy Microsoft Aurora na podstawie regionalnych danych WeatherBench2 przy użyciu techniki LoRA. Zadanie działa na pojedynczym GPU A100, Kueue kontroluje dopuszczanie zadań, a zadanie zapisuje punkt kontrolny adaptera i metryki trenowania do usługi Azure Blob Storage.

Prerequisites

• Infrastruktura wdrożona zgodnie z wdrażaniem infrastruktury dla Ray i Kueue na platformie AKS.
• Kolejki Kueue skonfigurowane zgodnie z instrukcjami w artykule Konfigurowanie kolejek Kueue dla obciążeń Ray w usłudze AKS.
• Co najmniej jeden procesor GPU A100 dostępny w klastrze (w tym przykładzie jest używany jeden procesor GPU).
• Dane Aurora WeatherBench2 zostały przesłane do magazynu Blob Storage (moduł Terraform dla infrastruktury automatycznie przesyła te dane).
• envsubst zainstalowano (gettext pakiet w systemie Linux, brew install gettext w systemie macOS).

Ustawianie zmiennych środowiskowych

Przejdź do przykładu aurora fine-tune w sklonowanym repozytorium i skonfiguruj wymagane zmienne środowiskowe:

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

Plik env.example ustawia wartości domyślne, w tym obraz Ray, nazwę kolejki i parametry trenowania. Zmienna JOB_NAME jest generowana automatycznie za pomocą znacznika czasu.

Prześlij zadanie

Prześlij RayJob do klastra:

./submit.sh

Skrypt tworzy obiekt ConfigMap na podstawie skryptu treningowego w Pythonie, renderuje szablon manifestu za pomocą envsubst i wdraża go w klastrze. Kueue dopuszcza zadanie, gdy w skonfigurowanej kolejce jest dostępny limit GPU.

Monitorowanie postępu

Znajdź i wyeksportuj nazwę zadania, jeśli pracujesz w nowej powłoce:

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

Obserwuj stan RayJob i dopuszczenie przez Kueue:

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

Oczekiwane dane wyjściowe po zakończeniu zadania:

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

Ujmij dzienniki procesów roboczych:

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

Weryfikowanie wyników

Sprawdź przesłane artefakty w magazynie obiektów blob:

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

Oczekiwane dane wyjściowe:

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

Pobierz i sprawdź metryki trenowania:

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

Oczekiwane dane wyjściowe:

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

Tablica loss_history powinna zawierać wartości skończone (a nie NaN), potwierdzając, że ścieżka danych i trenowanie modelu działają poprawnie.

Dokumentacja konfiguracji

Uprzątnij zasoby

Usuń RayJob i jego ConfigMap:

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

Następne kroki

• Trenuj model LLM za pomocą Ray
• Uruchamianie inferencji wsadowej za pomocą Ray