最佳化 GKE AI/機器學習工作負載優先順序

本文說明相關工具和最佳做法，協助您在 Google Kubernetes Engine (GKE) 中，盡量提高異質 AI/ML 工作負載的資源使用率，並減少停機時間，特別是在預留容量不足或使用隨選資源時。異質工作負載是指在同一個 GKE 叢集中同時執行的不同類型 AI/ML 工作負載。舉例來說，您可能會執行對延遲時間敏感的線上推論服務，同時執行一系列可中斷的批次訓練工作。

本指南為平台管理員和營運人員，以及資料和 AI 專家提供建議。

AI/機器學習工作負載優先順序的優點

異質工作負載的優先順序各不相同，且共用有限的容量和資源。本頁的最佳做法說明如何設定 GKE 和開放原始碼工具，以享有下列優點：

  • 盡量減少高優先順序工作負載的停機時間。
  • 快速執行高優先順序工作負載。
  • 盡量減少資源消耗。

背景

GKE 支援下列開放原始碼工具，可協助您盡可能提高資源使用率。

  • Kueue：Kubernetes 原生工作負載佇列系統，專為批次、AI 和高效能運算工作負載設計。Kueue 可擴充管理其他工作負載類型，例如自訂資源定義 (如 leaderworkerset) 定義的類型。Kueue 可管理 Kubernetes 叢集中的配額，以及工作負載耗用配額的方式。Kueue 會決定工作負載的等待時間、工作負載的啟動時間 (例如建立 Pod)，以及工作負載所屬 Pod 的搶占時間。

    如要進一步瞭解 Kueue，請參閱 Kueue 概念說明文件。

  • 熱插拔：這項技術可縮短平均復原時間 (MTTR)。當叢集資源用盡，且沒有額外容量 (來自隨選執行個體或現有預留資源) 時，Hotswap 會根據工作負載優先順序進行搶占。

    • 如果代管工作負載的節點狀況不佳，系統會將工作負載重新排程到符合資格的備用節點。如果沒有可用的備用節點，Hotswap 可以搶占優先順序較低的工作負載，為正在復原的工作負載騰出空間。
    • 如果使用 PriorityClass 設定 Pod，優先順序較高的工作負載會驅逐優先順序較低的工作負載，以取得資源。這項移除程序稱為搶占。

用途

請參閱下表，瞭解各個使用案例的最佳做法：

用途 最佳做法 說明
多個工作負載，優先順序各不相同 使用 Kueue 定義佇列，並根據工作負載的重要性指派優先順序。Kueue 可以管理配額，確保特定團隊或專案能存取一定數量的資源。

Kueue 可讓您套用下列設定：

  • 指派較高的 Kueue WorkloadPriority，優先處理高優先順序工作。
  • 啟用 Kueue 的公平分享佇列，確保所有工作負載最終都能取得資源，即使是低優先順序的工作負載也不例外。

如要測試最佳做法設定，請參閱本文中的 Kueue 範例。
你必須縮短目前的 MTTR。 發生中斷時，使用 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 叢集佇列。

熱插拔

以下範例顯示定義兩個優先順序類別 (low-priority-job 和 high-priority-job) 的 Hotswap 設定。這項 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 會在恢復運作的節點集區中，重新排程低優先順序工作負載。

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。

後續步驟