优化 GKE AI/机器学习工作负载优先级

本文档介绍了相关工具和最佳实践，旨在最大限度地提高 Google Kubernetes Engine (GKE) 中异构 AI/ML 工作负载的资源利用率并最大限度地缩短其停机时间，尤其是在预留中或通过按需资源没有容量时。异构工作负载是指在同一 GKE 集群中同时运行的不同类型的 AI/机器学习工作负载。例如，您可能会运行对延迟时间敏感的在线推理服务，同时运行一系列可中断的批量训练作业。

本指南为平台管理员和运维人员以及数据和 AI 专家提供了一些建议。

AI/机器学习工作负载优先级的优势

异构工作负载具有不同的优先级，并且共享有限的容量和资源。本页中的最佳实践介绍了如何配置 GKE 和开源工具，以帮助您获得以下好处：

  • 最大限度地缩短高优先级工作负载的停机时间。
  • 快速执行高优先级工作负载。
  • 优化资源消耗。

背景

GKE 支持以下开源工具来优化资源利用率。

  • Kueue：一个 Kubernetes 原生工作负载排队系统，专为批处理、AI 和高性能计算工作负载而设计。Kueue 可以扩展，以管理其他工作负载类型，例如由自定义资源定义（如 leaderworkerset）定义的工作负载类型。Kueue 可管理 Kubernetes 集群中的配额以及工作负载使用配额的方式。Kueue 会决定工作负载的等待时间、工作负载的启动时间（例如，通过创建 Pod）以及属于工作负载的 Pod 的抢占时间。

    如需详细了解 Kueue，请参阅 Kueue 概念文档。

  • 热插拔：一种可缩短平均恢复时间 (MTTR) 的技术。当集群资源完全利用且没有额外的容量（无论是来自按需实例还是现有预留）时，热交换可根据工作负载优先级实现抢占。

    • 当托管工作负载的节点运行状况不佳时，系统会在符合条件的备用节点上重新调度工作负载。如果没有可用的备用节点，热交换可以抢占低优先级工作负载，以便为正在恢复的工作负载腾出空间。
    • 如果您使用 PriorityClass 配置 Pod，则配置了较高优先级的工作负载会逐出正在运行的低优先级工作负载，以获取其资源。此逐出过程称为抢占。

使用场景

请参阅下表，了解每种应用场景的最佳实践：

使用场景 最佳做法 说明
具有不同优先级的多个工作负载 使用 Kueue 定义队列并根据工作负载的重要性为其分配优先级。Kueue 可以管理配额，以便特定团队或项目能够访问一定数量的资源。

借助 Kueue，您可以应用以下配置：

  • 通过为高优先级作业分配较高的 Kueue WorkloadPriority，优先处理这些作业。
  • 启用 Kueue 的公平共享排队功能，以便所有工作负载（即使是低优先级的工作负载）最终都能获得资源。

如需测试最佳实践配置，请参阅本文档中的 Kueue 示例。
您必须缩短当前的平均修复时间。 使用 Hotswap 在发生中断时重新安排健康资源中的工作负载，并抢占低优先级工作负载以支持高优先级工作负载。

通过热插拔，您可以应用以下配置：

  • 配置 PriorityClasses 以定义工作负载的优先级。
  • 为关键工作负载分配更高的 PriorityClasses。
  • 在发生中断时，自动将工作负载重新调度到运行状况良好的节点。

如需测试最佳实践配置，请参阅本文档中的热插拔示例。
多个 AI 工作负载争夺有限的资源 将 Kueue 和 Hotswap 结合使用。这种组合提供了一个强大的系统，可在初始调度期间和运行时优先处理关键工作负载。

Kueue 和 Hotswap 可让您应用以下配置：

  • 使用 Kueue 根据优先级管理工作负载的初始调度和准入。
  • 使用 Hotswap 处理工作负载中断并实现快速恢复。热交换有助于缩短高优先级工作负载在发生中断时的恢复时间。

如需测试最佳实践配置，请参阅本文档中的 Kueue 和热交换示例。

最佳实践实施示例

以下示例演示了如何实现 Kueue 和 Hotswap，以及如何将它们结合起来以实现上一部分中所述的最佳实践。

Kueue

以下示例清单展示了 Kueue 配置：

  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ResourceFlavor
  metadata:
    name: tpu-v6e-slice
  spec:
    nodeLabels:
      cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
  ---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ClusterQueue
  metadata:
    name: tpu-training-cq
  spec:
    resourceGroups:
    - flavors:
      - name: tpu-v6e-slice
        resources:
        - name: google.com/tpu
          nominalQuota: 32
    queueingStrategy: BestEffortFIFO
    preemption:
      reclaimWithinCohort: Never
      reclaimOutOfCohort:
        enable: true
        reclaimMoreThanNominalQuota: false
  ---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: LocalQueue
  metadata:
    name: default-queue
    namespace: default
  spec:
    clusterQueue: tpu-training-cq

此清单执行以下操作：

  • 定义了一个名为 tpu-v6e-slice 的 ResourceFlavor，用于指定 TPU v6e 切片的节点标签。
  • 定义了一个名为 tpu-training-cq 的 ClusterQueue，用于管理 TPU 资源的配额。
  • 定义一个名为 default-queue 的 LocalQueue，允许 default 命名空间中的工作负载使用 tpu-training-cq 集群队列。

热交换

以下示例展示了一个定义了两个优先级类（low-priority-job 和 high-priority-job）的热插拔配置。此热插拔配置会创建高优先级的 JobSet 工作负载并使用 MaxText。

  apiVersion: scheduling.k8s.io/v1
  kind: PriorityClass
  metadata:
    name: low-priority-job
  value: 1000000
  globalDefault: false
  description: "This priority class should be used for low priority pods only."
  ---
  apiVersion: scheduling.k8s.io/v1
  kind: PriorityClass
  metadata:
    name: high-priority-job
  value: 2000000
  globalDefault: false
  description: "This priority class should be used for critical pods only."
  ---
  apiVersion: jobset.x-k8s.io/v1alpha2
  kind: JobSet
  metadata:
    name: high-jax-trillium
    annotations:
      alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
  spec:
    failurePolicy:
      maxRestarts: 10
      restartStrategy: BlockingRecreate
    replicatedJobs:
    - name: slice
      replicas: 2
      template:
        spec:
          backoffLimit: 0
          completions: 4
          parallelism: 4
          template:
            spec:
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                cloud.google.com/gke-tpu-topology: 4x4
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              priorityClassName: high-priority-job
              containers:
              - name: jax-program
                image: <IMAGE LOCATION>
                command:
                -   python3
                -   MaxText/train.py
                -   MaxText/configs/base.yml
                -   model_name=llama2-7b
                -   run_name=<UNIQUE RUN NAME>
                -   steps=300
                -   base_output_directory=gs://<OUTPUT BUCKET>
                -   dataset_path=gs://max-datasets-rogue
                -   max_target_length=4096
                -   dataset_type=synthetic
                -   enable_checkpointing=False
                resources:
                  limits:
                    google.com/tpu: 4

根据此配置，Hotswap 会执行以下操作：

  • 如果基础设施故障中断了高优先级工作负载，JobSet 会重新启动该工作负载。热交换会抢占低优先级工作负载，以便在基础架构恢复之前重新调度高优先级工作负载。低优先级工作负载仍处于失败状态。此流程可显著减少工作负载的空闲时间。
  • 当基础设施恢复后，Hotswap 会在恢复的节点池中重新调度低优先级工作负载。

Kueue 和 Hotswap

当您在资源有限的复杂环境中运行时，可将 Kueue 与 Hotswap 结合使用。这种组合提供了一个强大的系统，可在初始调度期间和运行时优先处理关键工作负载。

以下示例展示了 Kueue 和 Hotswap 的组合配置。此示例使用 MaxText：

  apiVersion: scheduling.k8s.io/v1
  kind: PriorityClass
  metadata:
    name: low-priority-job
  value: 1000000
  globalDefault: false
  description: "This priority class should be used for low priority pods only."
  ---
  apiVersion: scheduling.k8s.io/v1
  kind: PriorityClass
  metadata:
    name: high-priority-job
  value: 2000000
  globalDefault: false
  description: "This priority class should be used for critical pods only."
  ---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ResourceFlavor
  metadata:
    name: tpu-v6e-slice
  spec:
    nodeLabels:
      cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
  ---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ClusterQueue
  metadata:
    name: tpu-training-cq
  spec:
    resourceGroups:
    - flavors:
      - name: tpu-v6e-slice
        resources:
        - name: google.com/tpu
          nominalQuota: 32
    queueingStrategy: BestEffortFIFO
    preemption:
      reclaimWithinCohort: Never
      reclaimOutOfCohort:
        enable: true
        reclaimMoreThanNominalQuota: false
  ---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: LocalQueue
  metadata:
    name: default-queue
    namespace: default
  spec:
    clusterQueue: tpu-training-cq
  ---
  apiVersion: jobset.x-k8s.io/v1alpha2
  kind: JobSet
  metadata:
    name: low-jax-trillium
    annotations:
      kueue.x-k8s.io/queue-name: default-queue
      alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
  spec:
    failurePolicy:
      maxRestarts: 10
      restartStrategy: BlockingRecreate
    replicatedJobs:
    - name: slice
      replicas: 2
      template:
        spec:
          backoffLimit: 0
          completions: 4
          parallelism: 4
          template:
            metadata:
              labels:
                kueue.x-k8s.io/managed-by: kueue
                kueue.x-k8s.io/priority-class: low-priority-job
            spec:
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                cloud.google.com/gke-tpu-topology: 4x4
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              priorityClassName: low-priority-job
              containers:
              - name: jax-program
                image: <IMAGE LOCATION>
                command:
                - python3
                - MaxText/train.py
                - MaxText/configs/base.yml
                - model_name=llama2-7b
                - run_name=low-priority-run
                - steps=30000
                - base_output_directory=gs://<OUTPUT BUCKET>
                - dataset_path=gs://max-datasets-rogue
                - max_target_length=4096
                - dataset_type=synthetic
                - enable_checkpointing=False
                resources:
                  limits:
                    google.com/tpu: 4
  ---
  apiVersion: jobset.x-k8s.io/v1alpha2
  kind: JobSet
  metadata:
    name: high-jax-trillium
    annotations:
      kueue.x-k8s.io/queue-name: default-queue
      alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-nodepool
  spec:
    failurePolicy:
      maxRestarts: 10
      restartStrategy: BlockingRecreate
    replicatedJobs:
    - name: slice
      replicas: 2
      template:
        spec:
          backoffLimit: 0
          completions: 4
          parallelism: 4
          template:
            metadata:
              labels:
                kueue.x-k8s.io/managed-by: kueue
                kueue.x-k8s.io/priority-class: high-priority-job
            spec:
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
                cloud.google.com/gke-tpu-topology: 4x4
              hostNetwork: true
              dnsPolicy: ClusterFirstWithHostNet
              priorityClassName: high-priority-job
              containers:
              - name: jax-program
                image: <IMAGE LOCATION>
                command:
                - python3
                - MaxText/train.py
                - MaxText/configs/base.yml
                - model_name=llama2-7b
                - run_name=high-priority-run
                - steps=300
                - base_output_directory=gs://<OUTPUT BUCKET>
                - dataset_path=gs://max-datasets-rogue
                - max_target_length=4096
                - dataset_type=synthetic
                - enable_checkpointing=False
                resources:
                  limits:
                    google.com/tpu: 4

根据此配置，Kueue 与 Hotswap 结合使用，并执行以下操作：

  • Kueue 会根据 low-jax-trillium 和 high-jax-trillium JobSet 的已定义优先级和可用资源，管理它们进入集群队列的准入。
  • 如果 high-jax-trillium JobSet 因基础设施故障而中断，Hotswap 会抢占 low-jax-trillium JobSet 以重新调度高优先级 JobSet。
  • 热交换可确保高优先级 JobSet 快速重启，从而最大限度地缩短其空闲时间。
  • 当基础设施恢复时，Hotswap 会在恢复的节点池中重新调度低优先级 JobSet。

后续步骤