在 AKS 上使用 Ray 运行批处理推理

在本文中,你将提交一个 RayJob,用于使用由 LLM 训练示例生成的 LoRA 适配器运行 vLLM 离线批量推理。 该作业从 Azure Blob 存储 中读取 viggo 测试集划分和 LoRA 适配器,在单个 GPU 上生成预测结果,并上传结果。

Important

AKS 文档和示例中都提到了开源软件。 您部署的软件被排除在 AKS 服务级别协议、有限保修和 Azure 支持之外。 将开源技术与 AKS 一起使用时,请查阅相应社区和项目维护者提供的支持选项来制定计划。

Microsoft 将负责生成我们在 AKS 上部署的开源包。 该责任包括拥有构建、扫描、签名、验证和快速修复流程的完整所有权,并掌控容器镜像中的二进制文件。 如需了解详细信息,请参阅 AKS 漏洞管理AKS 支持范围

先决条件

设置环境变量。

导航到克隆存储库中的批处理推理示例并配置所需的环境变量:

cd <path-to-cloned-repo>/AKS/examples/kueue-and-ray-on-aks/3-workloads/batch-inference
export AZURE_STORAGE_ACCOUNT_NAME=$(terraform -chdir=../../1-infrastructure/terraform output -raw storage_account_name)
source env.example

注释

如果您配置了团队队列(选项 B),请在运行 export QUEUE_NAME=team-a 之前设置 export QUEUE_NAME=team-bsource env.example。 默认值 QUEUE_NAME=default 仅适用于单队列配置(选项 A)。

提交工作负载

提交批量推理 RayJob:

./submit.sh

该脚本根据推理脚本创建 ConfigMap,通过 envsubst 渲染清单模板,并应用该清单。 当已配置的队列中有 1 个 GPU 可用时,Kueue 会准入该作业。

Tip

运行 ./submit.sh --dry-run 以验证呈现的清单,而不将其应用到群集。

监视进度

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

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

查看 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
batch-inference-xxxxxxxxxx   SUCCEEDED    Complete            2026-01-01T00:00:00Z   2026-01-01T00:09:00Z   9m
NAME                                      QUEUE     RESERVED IN     ADMITTED   FINISHED   AGE
rayjob-batch-inference-xxxxxxxxxx-xxxxx   default   cluster-queue   True       True       9m

尾部工作器日志:

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

验证结果

列出 Blob 存储中的预测文件:

az storage blob list -c llm-pipeline --prefix "inference/${JOB_NAME}/" \
  --account-name ${AZURE_STORAGE_ACCOUNT_NAME} --auth-mode login -o table

预期输出:

Name                                                 Blob Type    Blob Tier    Length    Content Type
---------------------------------------------------  -----------  -----------  --------  ------------------------
inference/<job-name>/metrics.json                    BlockBlob    Hot          128       application/octet-stream
inference/<job-name>/predictions.jsonl               BlockBlob    Hot          411777    application/octet-stream

下载并检查预测:

az storage blob download -c llm-pipeline \
  -n "inference/${JOB_NAME}/predictions.jsonl" \
  --account-name ${AZURE_STORAGE_ACCOUNT_NAME} --auth-mode login \
  --file /tmp/predictions.jsonl
head -1 /tmp/predictions.jsonl | python3 -m json.tool

预期输出:

{
    "input": "I'm wondering, have you played any games you found genuinely shocking?",
    "expected": "request(specifier[shocking])",
    "generated": "request_explanation"
}

配置参考

Variable 默认 Description
AZURE_STORAGE_ACCOUNT_NAME (必填) 模块 1 中的存储帐户
JOB_NAME batch-inference-<timestamp> 唯一的 RayJob 名称
QUEUE_NAME default Kueue LocalQueue 名称
LLM_DATA_CONTAINER llm-pipeline 具有 viggo 测试数据的 Blob 容器
LLM_LORA_CONTAINER llm-pipeline 带有 LoRA 适配器的 Blob 容器
CONFIGMAP_NAME batch-inference-scripts 保存推理脚本的 ConfigMap 的名称

清理资源

删除 RayJob 及其 ConfigMap:

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

若要拆解所有基础结构,请参阅 “部署基础结构 - 清理资源”。

后续步骤