在 AKS 上使用 Ray 微调 Aurora 天气模型

在本文中，你将提交一个 RayJob，使用 LoRA 基于区域 WeatherBench2 数据对 Microsoft Aurora 天气基础模型进行微调。 作业在单个 A100 GPU 上运行，Kueue 负责准入控制，并将适配器检查点和训练指标写入 Azure Blob 存储。

先决条件

• 按照 在 AKS 上为 Ray 和 Kueue 部署基础结构进行部署的基础结构。
• 按照 为 AKS 上的 Ray 工作负载配置 Kueue 队列 中所述配置的 Kueue 队列。
• 群集中至少有一个 A100 GPU 可用（此示例使用一个 GPU）。
• Aurora WeatherBench2 数据已上传到 Blob 存储中（基础设施 Terraform 模块会自动上传这些数据）。
• envsubst 已安装（Linux 上为 gettext 软件包，macOS 上为 brew install gettext）。

设置环境变量。

前往已克隆的代码库中的 Aurora 微调示例，并配置所需的环境变量：

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

文件 env.example 设置默认值，包括 Ray 图像、队列名称和训练参数。 变量 JOB_NAME 是使用时间戳自动生成的。

提交工作负载

将 RayJob 提交到群集：

./submit.sh

该脚本根据 Python 训练脚本创建 ConfigMap，通过 envsubst 渲染清单模板，并将其应用到集群。 当配置的队列中有可用的 GPU 配额时，Kueue 会准入该作业。

监视进度

如果你当前处于新的 shell 中，请先查找并导出作业名称：

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

查看 RayJob 状态和 Kueue 准入情况：

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

作业完成时的预期输出：

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

尾部工作器日志：

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

验证结果

检查 Blob 存储中上传的项目：

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

预期输出：

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

下载并检查训练指标：

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

预期输出：

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

数组 loss_history 应包含有限值（而不是 NaN），确认数据路径和模型训练正常工作。

配置参考

清理资源

删除 RayJob 及其 ConfigMap：

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

后续步骤

• 使用 Ray 训练 LLM
• 使用 Ray 运行批处理推理