使用 RayJob 和 Kueue 的优先级调度

本指南展示了如何将 使用 Ray Data 微调 PyTorch Lightning 文本分类器 示例作为 RayJob 运行，并利用 Kueue 进行优先级调度和配额管理。

什么是 Kueue？

Kueue 是一个 Kubernetes 原生的作业排队系统，用于管理配额以及作业如何消耗这些配额。Kueue 决定何时：

• 使一个任务等待

• 要允许一个作业开始，这意味着 Kubernetes 会创建 Pod。

• 要抢占一个作业，意味着 Kubernetes 会删除活跃的 Pod。

Kueue 对一些 KubeRay API 提供了原生支持。具体来说，您可以使用 Kueue 来管理 RayJob 和 RayCluster 消耗的资源。请参阅 Kueue 文档 以了解更多信息。

步骤 0：在 GKE 上创建一个 Kubernetes 集群（可选）

如果你已经有一个带有GPU的Kubernetes集群，你可以跳过这一步。否则，请按照 在GKE上使用GPU启动KubeRay的Kubernetes集群 来设置一个Kubernetes集群。

步骤 1：安装 KubeRay 操作员

按照 部署 KubeRay 操作员 来从 Helm 仓库安装最新稳定的 KubeRay 操作员。如果你正确地为 GPU 节点池设置了污点，那么 KubeRay 操作员 Pod 必须在 CPU 节点上。

步骤 2: 安装 Kueue

VERSION=v0.6.0
kubectl apply --server-side -f https://github.com/kubernetes-sigs/kueue/releases/download/$VERSION/manifests.yaml

有关安装 Kueue 的更多详细信息，请参阅 Kueue 安装。

步骤 3：使用优先级调度配置 Kueue

要理解本教程，理解以下 Kueue 概念非常重要：

• ResourceFlavor

• ClusterQueue

• LocalQueue

• 工作负载优先级类

# kueue-resources.yaml
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: "default-flavor"
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ClusterQueue
metadata:
  name: "cluster-queue"
spec:
  preemption:
    withinClusterQueue: LowerPriority
  namespaceSelector: {} # Match all namespaces.
  resourceGroups:
  - coveredResources: ["cpu", "memory", "nvidia.com/gpu"]
    flavors:
    - name: "default-flavor"
      resources:
      - name: "cpu"
        nominalQuota: 2
      - name: "memory"
        nominalQuota: 8G
      - name: "nvidia.com/gpu" # ClusterQueue only has quota for a single GPU.
        nominalQuota: 1
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: LocalQueue
metadata:
  namespace: "default"
  name: "user-queue"
spec:
  clusterQueue: "cluster-queue"
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: WorkloadPriorityClass
metadata:
  name: prod-priority
value: 1000
description: "Priority class for prod jobs"
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: WorkloadPriorityClass
metadata:
  name: dev-priority
value: 100
description: "Priority class for development jobs"

YAML 清单配置：

• ResourceFlavor

  • ResourceFlavor default-flavor 是一个空的 ResourceFlavor，因为 Kubernetes 集群中的计算资源是同质的。换句话说，用户可以请求 1 个 GPU 而不必考虑它是 NVIDIA A100 还是 T4 GPU。

• ClusterQueue

  • ClusterQueue cluster-queue 只有一个 ResourceFlavor default-flavor，配额为 2 个 CPU、8G 内存和 1 个 GPU。它完全匹配 1 个 RayJob 自定义资源的需求。因此，一次只能运行 1 个 RayJob。

  • ClusterQueue cluster-queue 有一个抢占策略 withinClusterQueue: LowerPriority。此策略允许待处理的 RayJob，如果其在 ClusterQueue 的名义配额内无法容纳，则可以抢占 ClusterQueue 中优先级较低的活跃 RayJob 自定义资源。

• LocalQueue

  • LocalQueue user-queue 是 default 命名空间中的一个命名空间对象，属于一个 ClusterQueue。典型的做法是为组织的租户、团队或用户分配一个命名空间。用户将作业提交到 LocalQueue，而不是直接提交到 ClusterQueue。

• 工作负载优先级类

  • 工作负载优先级类 prod-priority 的值高于工作负载优先级类 dev-priority。这意味着具有 prod-priority 优先级类的 RayJob 自定义资源优先于具有 dev-priority 优先级类的 RayJob 自定义资源。

创建 Kueue 资源：

kubectl apply -f kueue-resources.yaml

步骤 4：部署 RayJob

下载执行 微调 PyTorch Lightning 文本分类器 中所有步骤的 RayJob。源代码 也在 KubeRay 仓库中。

curl -LO https://raw.githubusercontent.com/ray-project/kuberay/master/ray-operator/config/samples/pytorch-text-classifier/ray-job.pytorch-distributed-training.yaml

在创建 RayJob 之前，请使用以下内容修改 RayJob 元数据：

metadata:
  generateName: dev-pytorch-text-classifier-
  labels:
    kueue.x-k8s.io/queue-name: user-queue
    kueue.x-k8s.io/priority-class: dev-priority

• kueue.x-k8s.io/queue-name: user-queue: 如前一步所述，用户将作业提交到 LocalQueue，而不是直接提交到 ClusterQueue。

• kueue.x-k8s.io/priority-class: dev-priority: 使用 dev-priority WorkloadPriorityClass 分配 RayJob。

• 一个修改后的名称，以表明此任务用于开发。

此外，通过查看 Ray 头 Pod 请求的资源，注意此 RayJob 所需的资源：

resources:
  limits:
    memory: "8G"
    nvidia.com/gpu: "1"
  requests:
    cpu: "2"
    memory: "8G"
    nvidia.com/gpu: "1"

现在部署 RayJob：

$ kubectl create -f ray-job.pytorch-distributed-training.yaml
rayjob.ray.io/dev-pytorch-text-classifier-r6d4p created

验证 RayCluster 和提交者 Kubernetes Job 是否正在运行：

$ kubectl get pod
NAME                                                      READY   STATUS    RESTARTS   AGE
dev-pytorch-text-classifier-r6d4p-4nczg                   1/1     Running   0          4s  # Submitter Kubernetes Job
torch-text-classifier-r6d4p-raycluster-br45j-head-8bbwt   1/1     Running   0          34s # Ray head Pod

在确认作业已成功完成后，删除 RayJob。

$ kubectl get rayjobs.ray.io dev-pytorch-text-classifier-r6d4p -o jsonpath='{.status.jobStatus}'
SUCCEEDED
$ kubectl get rayjobs.ray.io dev-pytorch-text-classifier-r6d4p -o jsonpath='{.status.jobDeploymentStatus}'
Complete
$ kubectl delete rayjob dev-pytorch-text-classifier-r6d4p
rayjob.ray.io "dev-pytorch-text-classifier-r6d4p" deleted

步骤 5：排队多个 RayJob 资源

创建3个RayJob自定义资源，以查看Kueue如何与KubeRay交互以实现作业排队。

$ kubectl create -f ray-job.pytorch-distributed-training.yaml
rayjob.ray.io/dev-pytorch-text-classifier-8vg2c created
$ kubectl create -f ray-job.pytorch-distributed-training.yaml
rayjob.ray.io/dev-pytorch-text-classifier-n5k89 created
$ kubectl create -f ray-job.pytorch-distributed-training.yaml
rayjob.ray.io/dev-pytorch-text-classifier-ftcs9 created

由于每个 RayJob 请求 1 个 GPU，而 ClusterQueue 的配额仅为 1 个 GPU，Kueue 会自动暂停新的 RayJob 资源，直到 GPU 配额可用。

你也可以检查 ClusterQueue 以查看可用和已使用的配额：

$ kubectl get clusterqueue
NAME            COHORT   PENDING WORKLOADS
cluster-queue            2
$ kubectl get clusterqueue cluster-queue -o yaml
apiVersion: kueue.x-k8s.io/v1beta1
kind: ClusterQueue
...
...
...
status:
  admittedWorkloads: 1  # Workloads admitted by queue.
  flavorsReservation:
  - name: default-flavor
    resources:
    - borrowed: "0"
      name: cpu
      total: "8"
    - borrowed: "0"
      name: memory
      total: 19531250Ki
    - borrowed: "0"
      name: nvidia.com/gpu
      total: "2"
  flavorsUsage:
  - name: default-flavor
    resources:
    - borrowed: "0"
      name: cpu
      total: "8"
    - borrowed: "0"
      name: memory
      total: 19531250Ki
    - borrowed: "0"
      name: nvidia.com/gpu
      total: "2"
  pendingWorkloads: 2   # Queued workloads waiting for quotas.
  reservingWorkloads: 1 # Running workloads that are using quotas.

步骤 6：部署一个优先级更高的 RayJob

此时，有多个 RayJob 自定义资源排队，但只有足够的配额来运行一个 RayJob。现在，您可以创建一个优先级更高的新 RayJob 来抢占已排队的 RayJob 资源。使用以下方式修改 RayJob：

metadata:
  generateName: prod-pytorch-text-classifier-
  labels:
    kueue.x-k8s.io/queue-name: user-queue
    kueue.x-k8s.io/priority-class: prod-priority

• kueue.x-k8s.io/queue-name: user-queue: 如前一步所述，用户将作业提交到 LocalQueue，而不是直接提交到 ClusterQueue。

• kueue.x-k8s.io/priority-class: dev-priority: 使用 prod-priority WorkloadPriorityClass 分配 RayJob。

• 一个修改后的名称，以表明此任务用于生产环境。

创建新的 RayJob：

$ kubectl create -f ray-job.pytorch-distributed-training.yaml
rayjob.ray.io/prod-pytorch-text-classifier-gkp9b created

请注意，当配额不足以同时满足两者时，优先级较高的作业会抢占优先级较低的作业：

$ kubectl get pods
NAME                                                      READY   STATUS    RESTARTS   AGE
prod-pytorch-text-classifier-gkp9b-r9k5r                  1/1     Running   0          5s
torch-text-classifier-gkp9b-raycluster-s2f65-head-hfvht   1/1     Running   0          35s