GKE AI/ML ワークロードの優先順位付けに関するベスト プラクティス

Autopilot Standard

このドキュメントでは、Google Kubernetes Engine（GKE）で異種 AI/ML ワークロードのリソース使用率を最大化し、ダウンタイムを最小限に抑えるためのツールとベスト プラクティスについて説明します。特に、予約またはオンデマンド リソースで容量が不足している場合に役立ちます。異種ワークロードとは、同じ GKE クラスタで同時に実行されるさまざまなタイプの AI/ML ワークロードを指します。 たとえば、割り込み可能な一連のバッチトレーニング ジョブと並行して、レイテンシの影響を受けやすいオンライン推論サービスを実行できます。

このガイドでは、プラットフォーム管理者とオペレーター、データおよび AI スペシャリスト向けの推奨事項について説明します。GKE のすべてのベスト プラクティスの概要については、 GKE のベスト プラクティスをご覧ください。

AI/ML ワークロードの優先順位付けのメリット

異種ワークロードには異なる優先順位があり、限られた容量とリソースを共有します。このページのベスト プラクティスでは、GKE とオープンソース ツールを構成して、次のメリットを得る方法について説明します。

• 優先度の高いワークロードのダウンタイムを最小限に抑えます。
• 優先度の高いワークロードを迅速に実行します。
• リソース消費を最適化します。

背景

GKE は、リソース使用率を最適化するために次のオープンソース ツールをサポートしています。

• Kueue: バッチ、AI、ハイ パフォーマンス コンピューティング ワークロード向けに設計された Kubernetes ネイティブのワークロード キューイング システム。Kueue は、leaderworkerset などのカスタム リソース定義で定義されたワークロード タイプなど、他のワークロード タイプを管理するように拡張できます。Kueue は、Kubernetes クラスタ内の割り当てとワークロードによる割り当ての使用方法を管理します。Kueue は、ワークロードの待機、ワークロードの開始（Pod の作成など）、ワークロードに属する Pod のプリエンプトのタイミングを決定します。

  Kueue の詳細については、Kueue コンセプトのドキュメントをご覧ください。

• Hotswap: 平均復旧時間（MTTR）を短縮する手法。Hotswap を使用すると、クラスタ リソースが完全に使用され、オンデマンド インスタンスまたは既存の予約から追加の容量を使用できない場合に、ワークロードの優先度に基づいてプリエンプションを行うことができます。

  • ワークロードをホストするノードが異常になると、ワークロードは使用可能なスペアノードに再スケジュールされます。スペアノードがない場合、Hotswap は優先度の低いワークロードをプリエンプトして、復旧中のワークロードのスペースを確保できます。
  • Pod を PriorityClass で構成すると、優先度の高いワークロードが実行中の優先度の低いワークロードを排除してリソースを取得します。この排除プロセスはプリエンプションと呼ばれます。

ユースケース

次の表を使用して、各ユースケースのベスト プラクティスを確認してください。

ユースケース	効果的な手法	説明
優先度の異なる複数のワークロード	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 スライスのノードラベルを指定する tpu-v6e-slice という名前の ResourceFlavor を定義します。
• TPU リソースの割り当てを管理する tpu-training-cq という名前の ClusterQueue を定義します。
• default Namespace のワークロードが tpu-training-cq クラスタキューを使用できるようにする default-queue という名前の LocalQueue を定義します。

Hotswap

次の例は、2 つの優先度クラス 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 は、インフラストラクチャが復旧する前に、優先度の低いワークロードをプリエンプトして、優先度の高いワークロードを再スケジュールします。優先度の低いワークロードは失敗したままになります。このプロセスにより、ワークロードのアイドル時間が大幅に短縮されます。
• インフラストラクチャが復旧すると、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 ワークロードをデプロイする 方法を確認する。