GKE AI/机器学习工作负载优先级划分最佳实践

Autopilot Standard

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

本指南为平台管理员和运维人员以及数据和 AI 专家提供了建议。如需全面了解所有 GKE 最佳实践，请参阅 GKE 最佳实践。

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

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

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

背景

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

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

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

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

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

使用场景

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

使用场景	最佳实践	说明
具有不同优先级的多个工作负载	使用 Kueue 定义队列，并根据工作负载的重要性为其分配优先级。Kueue 可以管理配额，以便某些团队或项目可以访问一定数量的资源。	Kueue 可让您应用以下配置： 通过为高优先级作业分配较高的 Kueue WorkloadPriority 来确定其优先级。 启用 Kueue 的公平共享排队，以便所有工作负载（即使是低优先级工作负载）最终都能获得资源。 如需测试最佳实践配置，请参阅本文档中的 Kueue 示例。
您必须缩短当前的 MTTR。	当发生中断时，使用 Hotswap 在运行状况良好的资源中重新调度工作负载，并抢占低优先级工作负载以支持高优先级工作负载。	Hotswap 可让您应用以下配置： 配置 PriorityClasses 以定义工作负载的优先级。 为关键工作负载分配较高的 PriorityClasses。 当发生中断时，在运行状况良好的节点上自动 重新调度工作负载。 如需测试最佳实践配置，请参阅本文档中的 Hotswap 示例。
多个 AI 工作负载争用有限的资源	结合使用 Kueue 和 Hotswap。这种组合提供了一个强大的系统 可在初始调度和运行时确定关键工作负载的优先级。	Kueue 和 Hotswap 可让您应用以下配置： 使用 Kueue 根据优先级管理 工作负载的初始调度和准入。 使用 Hotswap 处理工作负载 中断并实现快速恢复。当发生中断时，Hotswap 有助于缩短高优先级工作负载的恢复时间。 如需测试最佳实践配置，请参阅本文档中的 Kueue 和 Hotswap 示例。

最佳实践实现示例

以下示例演示了如何实现 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 集群队列。

Hotswap

以下示例展示了一个 Hotswap 配置，该配置定义了两个优先级类：low-priority-job 和 high-priority-job。此 Hotswap 配置会创建一个高优先级 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 会抢占低优先级工作负载，以便在基础架构恢复之前重新调度高优先级工作负载。低优先级工作负载将保持失败状态。此过程可显著缩短工作负载的闲置时间。
• 当基础架构恢复时，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。
• Hotswap 可确保高优先级 JobSet 快速重启，从而最大限度地缩短其闲置时间。
• 当基础架构恢复时，Hotswap 会在已恢复的节点池中重新调度低优先级 JobSet。

后续步骤

• 了解如何在 GKE 中部署 GPU 工作负载。
• 了解如何在 GKE 中部署 TPU 工作负载。