使用 KAI Scheduler 为 RayClusters 进行抢占式调度、队列优先级和 GPU 共享#

本指南演示了如何使用 KAI Scheduler 为 RayClusters 设置具有配额、抢占式调度和 GPU 共享的分层队列。

KAI Scheduler#

KAI Scheduler 是一个高性能、可扩展的 Kubernetes 调度器，专为 AI/ML 工作负载而构建。KAI 旨在大规模协调 GPU 集群，优化 GPU 分配，并支持完整的 AI 生命周期——从交互式开发到大规模分布式训练和推理。一些主要功能包括：

  • 装箱和分散调度：通过使用装箱最小化碎片来优化节点使用，或通过使用分散调度来提高弹性和负载均衡。

  • GPU 共享：允许 KAI 将来自不同团队的多个 Ray 工作负载合并到同一 GPU 上，使您的组织能够将更多工作负载部署到现有硬件上，并减少 GPU 空闲时间。

  • 工作负载自动伸缩：在 min/max 范围内缩放 Ray 副本或工作程序，同时尊重抢占式约束。

  • 集群自动伸缩：兼容动态云基础设施（包括像 Karpenter 这样的自动伸缩器）。

  • 工作负载优先级：在队列中有效地区分 Ray 工作负载的优先级。

  • 分层队列和公平性：具有配额、超额权重、限制的两级队列，并使用 DRF 等实现队列之间的公平资源分配。有关更多详细信息和关键功能，请参阅文档。

核心组件#

  1. PodGroups：PodGroups 是用于调度的原子单元，代表一个或多个相互依赖的 Pod，调度器将它们作为一个单元执行，也称为抢占式调度。它们对于分布式工作负载至关重要。KAI Scheduler 包括一个 PodGrouper，可自动处理抢占式调度。

PodGrouper 的工作原理

RayCluster "distributed-training":
├── Head Pod: 1 GPU
└── Worker Group: 4 × 0.5 GPU = 2 GPUs
Total Group Requirement: 3 GPUs

PodGrouper schedules all 5 pods (1 head + 4 workers) together or none at all.

  1. 队列：队列通过以下方式强制执行资源分配的公平性：

  • 配额：为队列保证的基线资源量。调度器首先分配配额以确保公平性。

  • 队列优先级：确定队列在获得超出其配额的资源时的顺序。调度器首先为高优先级队列服务。

  • 超额权重：控制调度器如何在同一优先级级别内的队列之间分配剩余资源。具有较高权重的队列将获得更大比例的额外资源。

  • 限制：定义队列可以消耗的最大资源量。

您可以为具有多个团队的组织（例如，拥有多个团队的部门）安排分层队列。

先决条件#

  • 具有 GPU 节点的 Kubernetes 集群

  • NVIDIA GPU Operator

  • 已配置 kubectl 以访问您的集群

  • 安装启用了 GPU 共享的 KAI Scheduler。从 KAI Scheduler releases 中选择所需的版本，并替换以下命令中的 <KAI_SCHEDULER_VERSION>。建议选择 v0.10.0 或更高版本。

# Install KAI Scheduler
helm upgrade -i kai-scheduler oci://ghcr.io/nvidia/kai-scheduler/kai-scheduler -n kai-scheduler --create-namespace --version <KAI_SCHEDULER_VERSION> --set "global.gpuSharing=true"

步骤 1：安装 KubeRay Operator 并将 KAI Scheduler 设置为批处理调度器#

遵循 KubeRay Operator 的官方安装文档，并添加以下配置以启用 KAI Scheduler 集成。

--set batchScheduler.name=kai-scheduler

步骤 2：创建 KAI Scheduler 队列#

为 department-1 及其子团队 team-a 创建基本的队列结构。出于演示目的，此示例不强制执行任何配额、超额权重或限制。您可以根据需要配置这些参数。

apiVersion: scheduling.run.ai/v2
kind: Queue
metadata:
  name: department-1
spec:
  #priority: 100 (optional)
  resources:
    cpu:
      quota: -1
      limit: -1
      overQuotaWeight: 1
    gpu:
      quota: -1
      limit: -1
      overQuotaWeight: 1
    memory:
      quota: -1
      limit: -1
      overQuotaWeight: 1
---
apiVersion: scheduling.run.ai/v2
kind: Queue
metadata:
  name: team-a
spec:
  #priority: 200 (optional)
  parentQueue: department-1
  resources:
    cpu:
      quota: -1
      limit: -1
      overQuotaWeight: 1
    gpu:
      quota: -1
      limit: -1
      overQuotaWeight: 1
    memory:
      quota: -1
      limit: -1
      overQuotaWeight: 1

注意：为使此演示更易于理解，我们在下一步的 RayCluster 示例中组合了这些队列定义。您可以使用单个合并的 YAML 文件，一次性应用队列和工作负载。

步骤 3：使用 KAI Scheduler 进行抢占式调度#

关键模式是将队列标签添加到您的 RayCluster。KubeRay 存储库中的这是一个基本示例。

metadata:
  name: raycluster-sample
  labels:
    kai.scheduler/queue: team-a    # This is the essential configuration.

应用此带有队列的 RayCluster。

curl -LO https://raw.githubusercontent.com/ray-project/kuberay/master/ray-operator/config/samples/ray-cluster.kai-scheduler.yaml

kubectl apply -f ray-cluster.kai-scheduler.yaml

#Verify queues are created
kubectl get queues
# NAME           PRIORITY   PARENT         CHILDREN     DISPLAYNAME
# department-1                             ["team-a"]   
# team-a                    department-1                

# Watch the pods get scheduled
kubectl get pods -w
# NAME                                    READY   STATUS              RESTARTS   AGE
# kuberay-operator-7d86f4f46b-dq22x       1/1     Running             0          50s
# raycluster-sample-head-rvrkz            0/1     ContainerCreating   0          13s
# raycluster-sample-worker-worker-mlvtz   0/1     Init:0/1            0          13s
# raycluster-sample-worker-worker-rcb54   0/1     Init:0/1            0          13s
# raycluster-sample-worker-worker-mlvtz   0/1     Init:0/1            0          40s
# raycluster-sample-worker-worker-rcb54   0/1     Init:0/1            0          41s
# raycluster-sample-head-rvrkz            0/1     Running             0          42s
# raycluster-sample-head-rvrkz            1/1     Running             0          54s
# raycluster-sample-worker-worker-rcb54   0/1     PodInitializing     0          59s
# raycluster-sample-worker-worker-mlvtz   0/1     PodInitializing     0          59s
# raycluster-sample-worker-worker-rcb54   0/1     Running             0          60s
# raycluster-sample-worker-worker-mlvtz   0/1     Running             0          60s
# raycluster-sample-worker-worker-rcb54   1/1     Running             0          71s
# raycluster-sample-worker-worker-mlvtz   1/1     Running             0          71s

为工作负载设置优先级#

在 Kubernetes 中，为工作负载分配不同的优先级可确保高效的资源管理，最大限度地减少服务中断，并支持更好的伸缩。通过为工作负载设置优先级，KAI Scheduler 会根据分配的优先级来调度作业。当没有足够的资源供工作负载使用时，调度器可以抢占低优先级的工作负载，为高优先级的工作负载释放资源。此方法可确保调度器在资源分配中始终优先处理任务关键型服务。

KAI Scheduler 部署带有多项预定义优先级类别

  • train (50) - 用于可抢占的训练工作负载

  • build-preemptible (75) - 用于可抢占的构建/交互式工作负载

  • build (100) - 用于构建/交互式工作负载（不可抢占）

  • inference (125) - 用于推理工作负载（不可抢占）

您可以提交具有特定优先级的前缀的相同工作负载。将前面的示例修改为 build 类工作负载。

  labels:
    kai.scheduler/queue: team-a    # This is the essential configuration.
    priorityClassName: build       # Here you can specify the priority class in metadata.labels (optional)

有关更多信息，请参阅文档。

步骤 4：提交具有 GPU 共享的 Ray 工作程序#

此示例创建了两个工作程序，它们在 RayCluster 中共享一个 GPU，每个工作程序占用 0.5 的 GPU（通过时间切片）。请参阅YAML 文件。

curl -LO https://raw.githubusercontent.com/ray-project/kuberay/master/ray-operator/config/samples/ray-cluster.kai-gpu-sharing.yaml

kubectl apply -f ray-cluster.kai-gpu-sharing.yaml

# Watch the pods get scheduled
kubectl get pods -w
# NAME                                          READY   STATUS    RESTARTS   AGE
# kuberay-operator-7d86f4f46b-dq22x             1/1     Running   0          4m9s
# raycluster-half-gpu-head-9rtxf                0/1     Running   0          4s
# raycluster-half-gpu-shared-gpu-worker-5l7cn   0/1     Pending   0          4s
# raycluster-half-gpu-shared-gpu-worker-98tzh   0/1     Pending   0          4s
# ... (skip for brevity)
# raycluster-half-gpu-shared-gpu-worker-5l7cn   0/1     Init:0/1   0          6s
# raycluster-half-gpu-shared-gpu-worker-5l7cn   0/1     Init:0/1   0          7s
# raycluster-half-gpu-shared-gpu-worker-98tzh   0/1     Init:0/1   0          8s
# raycluster-half-gpu-head-9rtxf                1/1     Running    0          19s
# raycluster-half-gpu-shared-gpu-worker-5l7cn   0/1     PodInitializing   0          19s
# raycluster-half-gpu-shared-gpu-worker-98tzh   0/1     PodInitializing   0          19s
# raycluster-half-gpu-shared-gpu-worker-5l7cn   0/1     Running           0          20s
# raycluster-half-gpu-shared-gpu-worker-98tzh   0/1     Running           0          20s
# raycluster-half-gpu-shared-gpu-worker-5l7cn   1/1     Running           0          31s
# raycluster-half-gpu-shared-gpu-worker-98tzh   1/1     Running           0          31s

注意：此示例中的 GPU 共享（带时间切片）仅发生在 Kubernetes 层，允许多个 Pod 共享单个 GPU 设备。调度器不强制执行内存隔离，因此应用程序必须自行管理其使用情况以防止干扰。有关其他 GPU 共享方法（例如 MPS），请参阅 KAI 的文档。

验证 GPU 共享是否正常工作#

要确认 GPU 共享是否正常工作，请使用以下命令：

# 1. Check GPU fraction annotations and shared GPU groups
kubectl get pods -l ray.io/cluster=raycluster-half-gpu -o custom-columns="NAME:.metadata.name,NODE:.spec.nodeName,GPU-FRACTION:.metadata.annotations.gpu-fraction,GPU-GROUP:.metadata.labels.runai-gpu-group"

您应该在同一个节点上看到两个工作程序 Pod，其 GPU-FRACTION: 0.5 和相同的 GPU-GROUP ID。

NAME                                          NODE               GPU-FRACTION   GPU-GROUP
raycluster-half-gpu-head                      ip-xxx-xx-xx-xxx   <none>         <none>
raycluster-half-gpu-shared-gpu-worker-67tvw   ip-xxx-xx-xx-xxx   0.5            3e456911-a6ea-4b1a-8f55-e90fba89ad76
raycluster-half-gpu-shared-gpu-worker-v5tpp   ip-xxx-xx-xx-xxx   0.5            3e456911-a6ea-4b1a-8f55-e90fba89ad76

这表明两个工作程序具有相同的 NVIDIA_VISIBLE_DEVICES（相同的物理 GPU）和 GPU-FRACTION: 0.50。

故障排除#

检查是否缺少队列标签#

如果 Pod 仍然处于 Pending 状态，最常见的问题是缺少队列标签。

检查 Operator 日志中是否有 KAI Scheduler 错误，并查找类似以下的错误消息：

"Queue label missing from RayCluster; pods will remain pending"

解决方案：确保您的 RayCluster 具有集群中存在的队列标签。

metadata:
  labels:
    kai.scheduler/queue: default  # Add this label