GKE AI/ML ワークロードの優先度設定を最適化する

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

このガイドでは、プラットフォーム管理者とオペレーター、データおよび AI スペシャリスト向けの推奨事項について説明します。

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

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

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

背景

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

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

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

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

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

ユースケース

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

ユースケース ベスト プラクティス 説明
優先度の異なる複数のワークロード Kueue を使用してキューを定義し、重要度に基づいてワークロードに優先度を割り当てます。Kueue は割り当てを管理して、特定のチームやプロジェクトが一定量のリソースにアクセスできるようにします。

Kueue では、次の構成を適用できます。

  • 優先度の高いジョブに高い Kueue WorkloadPriority を割り当てて、優先度を上げます。
  • Kueue のフェアシェア キューイングを有効にして、優先度の低いワークロードも含め、すべてのワークロードが最終的にリソースを受け取れるようにします。

ベスト プラクティスの構成をテストするには、このドキュメントの Kueue の例をご覧ください。
現在の MTTR を短縮する必要があります。 Hotswap を使用すると、中断が発生したときに正常なリソースでワークロードを再スケジュールし、優先度の高いワークロードのために優先度の低いワークロードをプリエンプトできます。

ホットスワップでは、次の構成を適用できます。

  • PriorityClasses を構成して、ワークロードの優先度レベルを定義します。
  • 重要なワークロードに高い PriorityClasses を割り当てます。
  • 中断が発生したときに、正常なノードでワークロードを自動的に再スケジュールします。

ベスト プラクティス構成をテストするには、このドキュメントのホットスワップの例をご覧ください。
限られたリソースを奪い合う複数の AI ワークロード Kueue と Hotswap を組み合わせます。この組み合わせにより、初期スケジューリングとランタイムの両方で重要なワークロードを優先する堅牢なシステムが実現します。

Kueue と Hotswap を使用すると、次の構成を適用できます。

  • Kueue を使用して、優先度に基づいてワークロードの初期スケジューリングとアドミッションを管理します。
  • 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 を定義します。

ホットスワップ

次の例は、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 とホットスワップ

リソースが限られた複雑な環境で運用する場合は、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 を再スケジュールします。

次のステップ