Ajustar el modelo meteorológico Aurora con Ray en AKS

En este artículo, envías un RayJob que ajusta con precisión el modelo fundacional meteorológico Microsoft Aurora con datos regionales de WeatherBench2 mediante LoRA. La tarea se ejecuta en una única GPU A100, Kueue controla la admisión y la tarea escribe el punto de control del adaptador y las métricas de entrenamiento en Azure Blob Storage.

Prerequisites

• Infraestructura implementada siguiendo Implementar la infraestructura para Ray y Kueue en AKS.
• Las colas de Kueue se configuran siguiendo Configurar colas de Kueue para cargas de trabajo de Ray en AKS.
• Al menos una GPU A100 disponible en el clúster (en este ejemplo se usa una GPU).
• Datos de Aurora WeatherBench2 cargados en el almacenamiento de blobs (el módulo de infraestructura de Terraform carga estos datos automáticamente).
• envsubst instalado (gettext paquete en Linux, brew install gettext en macOS).

Establecimiento de variables de entorno

Diríjase al ejemplo de ajuste fino de Aurora en el repositorio clonado y configure las variables de entorno necesarias:

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

El env.example archivo establece los valores predeterminados, como la imagen de Ray, el nombre de la cola y los parámetros de entrenamiento. La JOB_NAME variable se genera automáticamente con una marca de tiempo.

Enviar la carga de trabajo

Envíe el RayJob al clúster:

./submit.sh

El script crea un ConfigMap a partir del script de entrenamiento en Python, genera el manifiesto a partir de la plantilla mediante envsubst y lo aplica al clúster. Kueue admite el trabajo cuando hay cuota de GPU disponible en la cola configurada.

Supervisión de progreso

Busca y exporta el nombre del trabajo si te encuentras en un nuevo shell:

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

Vea el estado de RayJob y la admisión de Kueue:

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

Salida esperada cuando se completa el trabajo:

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

Consulta los registros del trabajador:

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

Comprobación de los resultados

Compruebe los artefactos cargados en Blob Storage:

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

Descargue e inspeccione las métricas de entrenamiento:

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

La loss_history matriz debe contener valores finitos (no NaN), lo que confirma que la ruta de acceso de datos y el entrenamiento del modelo funcionan correctamente.

Referencia de configuración

Limpieza de recursos

Elimine RayJob y su ConfigMap:

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

Pasos siguientes

• Entrena un LLM con Ray
• Ejecución de la inferencia por lotes con Ray