使用 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 集群，则可以跳过此步骤。否则，请按照 使用 GPU 为 KubeRay 启动 Google Cloud GKE 集群 在 GKE 上设置 Kubernetes 集群。

步骤 1：安装 KubeRay Operator#

请按照 部署 KubeRay Operator 从 Helm 仓库安装最新的稳定版 KubeRay Operator。如果您正确设置了 GPU 节点池的污点，KubeRay Operator Pod 必须位于 CPU 节点上。

步骤 2：安装 Kueue#

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

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

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

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

# 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 manifest 配置了：

  • 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。此策略允许不符合其 ClusterQueue 名义配额的待处理 RayJob 抢占 ClusterQueue 中具有较低优先级的活动 RayJob 自定义资源。

  • LocalQueue

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

  • WorkloadPriorityClass

    • WorkloadPriorityClass prod-priority 的值高于 WorkloadPriorityClass 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：为 RayJob 分配 dev-priority WorkloadPriorityClass。

  • 修改名称以表明此任务用于开发。

同时，通过查看 Ray head 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: prod-priority：为 RayJob 分配 prod-priority WorkloadPriorityClass。

  • 修改名称以表明此任务用于生产。

创建新的 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