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 概念說明文件。

• 熱交換：縮短平均復原時間 (MTTR) 的技術。當叢集資源完全用盡，且沒有額外容量 (無論是來自隨需執行個體或現有預留項目) 時，Hotswap 會根據工作負載優先順序進行先占。

  • 如果代管工作負載的節點健康狀態不良，系統會將工作負載重新排定至符合資格的備用節點。如果沒有可用的備用節點，熱交換功能可以優先處理優先順序較低的工作負載，為復原中的工作負載騰出空間。
  • 如果您使用 PriorityClass 設定 Pod，優先順序較高的工作負載會逐出執行中的低優先順序工作負載，以取得其資源。這個淘汰程序稱為「先占」。

用途

請參閱下表，瞭解各用途的最佳做法：

用途	最佳做法	說明
多個工作負載，優先順序各不相同	使用 Kueue 定義佇列，並根據工作負載的重要性指派優先順序。Kueue 可以管理配額，確保特定團隊或專案能存取一定數量的資源。	Kueue 可讓您套用下列設定： 為優先順序較高的 Job 指派較高的 Kueue WorkloadPriority，藉此排定優先順序。 啟用 Kueue 的公平分享佇列，讓所有工作負載最終都能取得資源，即使是低優先順序的工作負載也不例外。 如要測試最佳做法設定，請參閱本文中的 Kueue 範例。
您必須縮短目前的 MTTR。	發生中斷時，使用 Hotswap 在健全的資源中重新排程工作負載，並優先處理高優先順序工作負載，而非低優先順序工作負載。	您可以使用 Hotswap 套用下列設定： 設定 PriorityClasses，定義工作負載的優先順序。 為重要工作負載指派較高的 PriorityClasses。 發生中斷時，自動在健康狀態良好的節點上重新排定工作負載。 如要測試最佳做法設定，請參閱本文中的「熱插拔範例」。
多個 AI 工作負載爭用有限資源	結合 Kueue 和 Hotswap。這兩項功能結合後，可提供強大的系統，在初始排程和執行階段，都會優先處理重要工作負載。	Kueue 和 Hotswap 可讓您套用下列設定： 使用 Kueue 依優先順序管理工作負載的初始排程和許可。 使用熱插拔功能處理工作負載中斷問題，並快速復原。發生中斷時，熱交換功能可縮短高優先順序工作負載的復原時間。 如要測試最佳做法設定，請參閱本文中的 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 設定，其中定義了兩個優先順序類別：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 的排程。

後續步驟

• 瞭解如何在 GKE 中部署 GPU 工作負載。
• 瞭解如何在 GKE 中部署 TPU 工作負載。