---
title: Fine-tune the Aurora weather model with Ray on Azure Kubernetes Service (AKS)
description: Learn how to fine-tune the Microsoft Aurora weather foundation model using Ray and Kueue on AKS with LoRA adapters.
ms.topic: how-to
ms.service: azure-kubernetes-service
ms.date: 06/26/2026
author: feiskyer
ms.author: peni
ms.custom: 'stateful-workloads'
---

# Fine-tune the Aurora weather model with Ray on AKS

In this article, you submit a RayJob that fine-tunes the [Microsoft Aurora](https://github.com/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.

[!INCLUDE [open source disclaimer](./includes/open-source-disclaimer.md)]

## Prerequisites

* Infrastructure deployed following [Deploy infrastructure for Ray and Kueue on AKS](./deploy-ray-infrastructure.md).
* Kueue queues configured following [Configure Kueue queues for Ray workloads on AKS](./configure-ray-kueue-queues.md).
* At least one A100 GPU available in the cluster (this example uses one GPU).
* Aurora WeatherBench2 data uploaded to blob storage (the infrastructure Terraform module uploads this data automatically).
* `envsubst` installed (`gettext` package on Linux, `brew install gettext` on macOS).

## Set environment variables

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

```bash
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:

```bash
./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:

```bash
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:

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

Expected output when the job completes:

```output
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
```

```output
NAME                                      QUEUE     RESERVED IN     ADMITTED   FINISHED   AGE
rayjob-aurora-finetune-xxxxxxxxxx-xxxxx   default   cluster-queue   True       True       8m
```

Tail the worker logs:

```bash
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:

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

Expected output:

```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:

```bash
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:

```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:

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

## Next steps

> [!div class="nextstepaction"]
> [Serve the fine-tuned Aurora model with Ray Serve](./ray-online-serving.md)

* [Train an LLM with Ray](./ray-train-llm.md)
* [Run batch inference with Ray](./ray-batch-inference.md)
