Ray集成NVIDIA KAI调度器实现组调度与工作负载优先级

在Ray中使用NVIDIA KAI调度器启用组调度和工作负载优先级

NVIDIA KAI调度器现已与KubeRay原生集成，将某中心Run:ai高需求和高规模环境中使用的相同调度引擎直接引入Ray集群。这意味着您现在可以在Ray环境中利用组调度、工作负载自动扩缩容、工作负载优先级、分层队列等特性。这些功能通过协调作业启动、高效共享GPU和优先处理工作负载，使您的基础设施更加智能。

这对Ray用户意味着什么

组调度：无部分启动
分布式Ray工作负载需要其所有工作节点和参与者同时启动，或者根本不启动。KAI确保它们作为协调的组启动，防止浪费性的部分分配，避免训练或推理管道停滞。

工作负载和集群自动扩缩容
对于离线批量推理等工作负载，当集群资源可用或队列允许超配额使用时，Ray集群可以扩展。随着需求减少，它们也可以缩减，提供与资源可用性和工作负载需求相匹配的弹性计算，无需手动干预。

工作负载优先级：不同类型工作负载的平滑共存
如果资源有限，高优先级推理作业可以自动抢占低优先级批量训练，保持应用程序响应性，无需手动干预。

带优先级的分层队列：动态资源共享
为不同项目团队创建具有明确优先级的队列，以便在容量可用时，较高优先级队列可以从其他团队借用空闲资源。

技术设置

本示例假设以下环境：

具有一个NVIDIA A10G GPU的Kubernetes集群
已安装某中心GPU Operator
已部署NVIDIA KAI调度器
KubeRay Operator夜间镜像或通过Helm chart配置使用KAI调度器

步骤1：设置KAI调度器队列

在提交Ray工作负载之前，必须为KAI调度器定义队列。KAI调度器支持分层队列，使团队和部门能够组织成多级结构，具有对资源分布的细粒度控制。

在此示例中，将创建一个简单的两级层次结构，其中包含一个名为department-1的顶级父队列和一个名为team-a的子队列。

apiVersion: scheduling.run.ai/v2
kind: Queue
metadata:
  name: department-1
spec:
  resources:
    cpu:
      quota: -1
      limit: -1
      overQuotaWeight: 1
    gpu:
      quota: -1
      limit: -1
      overQuotaWeight: 1
    memory:
      quota: -1
      limit: -1
      overQuotaWeight: 1

apiVersion: v1
kind: Secret
metadata:
  name: hf-token
type: Opaque
stringData:
  hf_token: $HF_TOKEN

应用工作负载：

kubectl apply -f ray-service.kai-scheduler.llm-serve.yaml

观察抢占行为

应用后，您会注意到KAI调度器抢占训练作业，为推理工作负载腾出空间，因为两者竞争相同的GPU，但推理工作负载具有更高的优先级。

与部署的模型交互

在端口转发以访问Ray仪表板或推理端点之前，列出可用服务以确保定位正确的服务：

端口转发Ray仪表板：

kubectl port-forward svc/ray-kai-scheduler-serve-llm-xxxxx-head-svc 8265:8265

然后打开http://127.0.0.1:8265查看Ray仪表板并确认部署处于活动状态。

接下来，端口转发LLM服务端点：

kubectl port-forward svc/ray-kai-scheduler-serve-llm-serve-svc 8000:8000

最后，使用OpenAI兼容的API调用查询模型：

curl http://localhost:8000/v1/completions -H "Content-Type: application/json" -d '{
        "model": "qwen2.5-7b-instruct",
        "prompt": "San Francisco is a",
        "max_tokens": 13,
        "temperature": 0}'

总结

本文探讨了KAI调度器如何提升Ray工作负载的调度能力，包括组调度和分层队列。演示了如何有效优先处理训练和推理工作负载，当资源有限时，推理工作负载能够抢占训练作业。

虽然此演示为方便起见使用了简单的开放权重模型和Hugging Face，但某中心Run:ai Model Streamer是一个生产级选项，通过将模型权重直接从S3或其他高带宽存储流式传输到GPU内存，减少模型启动时间。它还与vLLM原生集成，并与Ray开箱即用。
更多精彩内容请关注我的个人公众号公众号（办公AI智能小助手）或者我的个人博客 https://blog.qife122.com/
对网络安全、黑客技术感兴趣的朋友可以关注我的安全公众号（网络安全技术点滴分享）