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

步骤 1：安装 KubeRay Operator#

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

步骤 2：安装 Kueue#

VERSION=v0.8.2
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 清单配置了：

  • 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。

  • 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。

  • 修改名称以指示此作业用于开发。

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

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：为 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