在 AKS 上使用 Ray Serve 为模型提供服务

在本文中,你将使用 Ray Serve on Azure Kubernetes 服务 (AKS) 将微调的极光天气模型部署为持久性 HTTP 终结点。 此方法使用 RayService 资源来实现长期运行的服务,并且不受 Kueue 的准入控制。

Important

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

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

先决条件

设置环境变量。

导航到克隆存储库中的联机服务示例并配置所需的环境变量:

cd <path-to-cloned-repo>/AKS/examples/kueue-and-ray-on-aks/3-workloads/online-serving
export AZURE_STORAGE_ACCOUNT_NAME=$(terraform -chdir=../../1-infrastructure/terraform output -raw storage_account_name)
export AURORA_RUN_ID=<your-aurora-finetune-job-name>
source env.example

注释

AURORA_RUN_ID是你已完成的 Aurora 微调运行中的JOB_NAME。 可以使用以下方法找到它:

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

部署服务

部署 RayService:

./submit-service.sh

该脚本根据服务应用程序代码创建 ConfigMap,通过 envsubst 渲染 RayService 清单,并应用该清单。 KubeRay 操作员创建 Ray 群集并部署 Serve 应用程序。

Tip

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

注释

RayService 不受 Kueue 的准入控制。 服务型工作负载通常是长期运行的,需要专用资源,而不是采用批处理队列机制。

验证部署

持续观察,直到 RayService 显示为 Running:

kubectl -n ray get rayservice ${SERVICE_NAME} -w

预期输出:

NAME           SERVICE STATUS   NUM SERVE ENDPOINTS
aurora-serve   Running          1

测试终结点

将端口转发到 serve 服务并发送测试请求:

kubectl -n ray port-forward svc/${SERVICE_NAME}-serve-svc 8000:8000

运行状况检查:

curl http://localhost:8000${ROUTE_PREFIX}

预期输出:

{"status":"ok","gpu_name":"NVIDIA A100-SXM4-80GB","run_id":"<your-run-id>","adapter":"checkpoints/<your-run-id>/last.safetensors"}

预测请求:

curl -X POST http://localhost:8000${ROUTE_PREFIX} \
  -H 'Content-Type: application/json' \
  -d '{"init_file": "init-2021-01-01-00z.npz", "lead_hours": 6}'

预期输出:

{
  "init_file": "init-2021-01-01-00z.npz",
  "lead_hours": 6,
  "surface_variables": {
    "2t":  {"shape": [1,1,52,100], "mean": 298.76, "min": 272.0, "max": 302.0},
    "10u": {"shape": [1,1,52,100], "mean": -2.21,  "min": -18.0, "max": 15.38},
    "10v": {"shape": [1,1,52,100], "mean": 0.92,   "min": -13.88,"max": 18.88},
    "msl": {"shape": [1,1,52,100], "mean": 101320.07, "min": 100352.0, "max": 101888.0}
  },
  "gpu_name": "NVIDIA A100-SXM4-80GB",
  ...
}

响应包括预测中每个图面变量的每一个变量摘要统计信息(形状、平均值、最小值、最大值)。

配置参考

Variable 默认 Description
AZURE_STORAGE_ACCOUNT_NAME (必填) 模块 1 中的存储帐户
AURORA_RUN_ID (必填) 已完成的 aurora-finetune 运行的 JOB_NAME
AURORA_INPUT_CONTAINER aurora 包含 init 数据的容器
AURORA_ADAPTER_CONTAINER aurora 具有 LoRA 检查点的容器
AURORA_INIT_FILE init-2021-01-01-00z.npz 请求的默认 init 文件
AURORA_LEAD_HOURS 6 默认预测提前期(必须为 6 小时的倍数)
AURORA_LORA_RANK 8 LoRA 排名 (必须匹配训练)
AURORA_REQUIRE_GPU_NAME A100 检查 GPU 名称(留空则跳过)
SERVICE_NAME aurora-serve RayService 资源的名称
ROUTE_PREFIX /aurora 服务终结点的 HTTP 路由前缀
CONFIGMAP_NAME aurora-serve-scripts 保存服务脚本的 ConfigMap 的名称

清理资源

删除 RayService 及其 ConfigMap:

kubectl -n ray delete rayservice ${SERVICE_NAME}
kubectl -n ray delete configmap ${CONFIGMAP_NAME}

后续步骤