Edit

Fine-tune the Aurora weather model with Ray on AKS

In this article, you submit a RayJob that fine-tunes the Microsoft Aurora weather foundation model on regional WeatherBench2 data using LoRA. The job runs on a single A100 GPU, Kueue controls admission, and the job writes the adapter checkpoint and training metrics to Azure Blob Storage.

Important

Open-source software is mentioned throughout AKS documentation and samples. Software that you deploy is excluded from AKS service-level agreements, limited warranty, and Azure support. As you use open-source technology alongside AKS, consult the support options available from the respective communities and project maintainers to develop a plan.

Microsoft takes responsibility for building the open-source packages that we deploy on AKS. That responsibility includes having complete ownership of the build, scan, sign, validate, and hotfix process, along with control over the binaries in container images. For more information, see Vulnerability management for AKS and AKS support coverage.

Prerequisites

Set environment variables

Navigate to the Aurora fine-tune example in the cloned repository and configure the required environment variables:

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

The env.example file sets defaults including the Ray image, queue name, and training parameters. The JOB_NAME variable is autogenerated with a timestamp.

Note

If you configured team queues (Option B), set export QUEUE_NAME=team-a or export QUEUE_NAME=team-b before running source env.example. The default QUEUE_NAME=default only works with the single-queue configuration (Option A).

Submit the workload

Submit the RayJob to the cluster:

./submit.sh

The script creates a ConfigMap from the Python training script, renders the manifest template via envsubst, and applies it to the cluster. Kueue admits the job when GPU quota is available in the configured queue.

Tip

Run ./submit.sh --dry-run to validate the rendered manifest without applying it to the cluster.

Monitor progress

Find and export the job name if you're in a new shell:

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

Watch the RayJob status and Kueue admission:

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

Expected output when the job completes:

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

Tail the worker logs:

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

Verify results

Check the uploaded artifacts in blob storage:

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

Expected output:

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 and inspect the training metrics:

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

Expected output:

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

The loss_history array should contain finite values (not NaN), confirming the data path and model training are working correctly.

Configuration reference

Variable Default Description
AZURE_STORAGE_ACCOUNT_NAME (required) Storage account from Module 1
AURORA_INPUT_CONTAINER aurora Container for input data
AURORA_OUTPUT_CONTAINER aurora Container for checkpoints
AURORA_INIT_FILE init-2021-01-01-00z.npz Init NPZ file name
AURORA_TRUTH_FILE truth-2021-01-01-06z.npz Truth NPZ file name
AURORA_MAX_STEPS 1 Training steps
AURORA_LORA_RANK 8 LoRA rank
AURORA_LEAD_HOURS 6 Forecast lead time (must be 6h multiple)
AURORA_REQUIRE_GPU_NAME A100 GPU name substring guard
QUEUE_NAME default Kueue LocalQueue name
CONFIGMAP_NAME aurora-finetune-scripts Name for the ConfigMap holding the training script

Clean up resources

Delete the RayJob and its ConfigMap:

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

Next steps

Serve the fine-tuned Aurora model with Ray Serve

  • Last updated on