トポロジ対応スケジューリング（TAS）で GKE ワークロードをスケジュールする

AI / ML ワークロードでは、Pod 間通信が頻繁に行われます。この要件により、Pod 間のネットワーク帯域幅は、ワークロードの実行時間と費用に直接影響します。この帯域幅は、クラスタ内の仮想マシン（VM）インスタンスの配置によって異なります。

このドキュメントでは、パフォーマンスと信頼性の両方を考慮して、Google Kubernetes Engine（GKE）クラスタで大規模な AI または ML ワークロードのスケジューリングを最適化する方法について説明します。具体的には、低レイテンシ通信にトポロジ認識スケジューリング（TAS）を使用するようにクラスタを構成します。このアプローチにより、通信オーバーヘッドを最小限に抑え、ワークロードのパフォーマンスを最大限に高めることができます。

トポロジ認識スケジューリング（TAS）とは

TAS を使用すると、 大規模言語モデル（LLM）のトレーニングの効率を大幅に向上させることができます。TAS は、ワーカーが特定のランク順で通信する必要がある勾配集約時の通信オーバーヘッドを最小限に抑えるために、ワーカーをネットワーク トポロジに戦略的に配置します。TAS は、順次通信するワーカー間のネットワーク ホップを最小限に抑えることで、ネットワーク競合を減らし、帯域幅の使用率を最適化し、収束を高速化してトレーニング時間を短縮します。LLM モデルがますます大規模化する中、TAS は分散トレーニングのパフォーマンスとスケーラビリティを最大化するために不可欠です。

TAS は、予約によって取得できる密度の高い容量で最適に機能します。Flex Start VM または Spot VM の場合、容量が近くに割り当てられる可能性が低いため、このシナリオでは TAS が適切に機能しない可能性があります。

始める前に

作業を始める前に、次のタスクが完了していることを確認してください。

  • Google Kubernetes Engine API を有効にする。
    • Google Kubernetes Engine API の有効化
  • このタスクに Google Cloud CLI を使用する場合は、gcloud CLI をインストールして初期化する。gcloud CLI をインストール済みの場合は、gcloud components update コマンドを実行して最新のバージョンを取得します。以前のバージョンの gcloud CLI では、このドキュメントのコマンドを実行できない場合があります。

  • クラスタに接続するには、次のコマンドを実行します。

    gcloud container clusters get-credentials CLUSTER_NAME

    CLUSTER_NAME は、使用するクラスタの名前に置き換えます。

GKE クラスタを準備する

TAS でワークロードを実行するように GKE クラスタを準備する手順は次のとおりです。

  1. TAS を有効にして Kueue をインストールする

  2. GKE クラスタのトポロジを表示する

  3. Kueue を構成する

TAS を有効にして Kueue をインストールする

TAS は、割り当てとジョブによる割り当ての使用方法を管理する Kubernetes ネイティブ システムである Kueueで 使用することをおすすめします。TAS には Kueue バージョン 0.10.0 以降が必要であり、明示的に有効にする必要があります。

Kueue をインストールして TAS を有効にするには、次のいずれかのオプションを選択します。

Kueue マニフェスト

  1. Kueue をインストールします。

    kubectl apply --server-side -f https://github.com/kubernetes-sigs/kueue/releases/download/v0.10.0/manifests.yaml

  2. Kueue で TAS を有効にします。

    kubectl -n kueue-system patch deployment kueue-controller-manager \
    --type json -p='[{"op": "add", "path": "/spec/template/spec/containers/0/args/-", "value": "--feature-gates=TopologyAwareScheduling=true"}]'

Helm チャート

Helm チャートを使用して、TAS を有効にして Kueue をインストールします。

helm install kueue oci://us-central1-docker.pkg.dev/k8s-staging-images/charts/kueue \
    --version="v0.10.0" \
    --create-namespace \
    --namespace=kueue-system \
    --set="controllerManager.featureGates[0].name=TopologyAwareScheduling,controllerManager.featureGates[0].enabled=true"

Kueue をインストールしたら、次のセクションで説明するように、管理するインフラストラクチャを認識するように構成する必要があります。

GKE クラスタのトポロジを表示する

Spot VM としてプロビジョニングされた A4X、A4、A3 Ultra、A3 Mega、A3 High（GPU 数 8） ノードのトポロジを表示する前に、 コンパクト プレースメントを GKE ノードに定義して、TAS の物理トポロジを公開する必要があります。それ以外の場合はエラーが発生します。

特定のノードプール内の GKE クラスタノードのトポロジを表示するには、次のコマンドを実行します。

kubectl get nodes -l cloud.google.com/gke-nodepool=NODE_POOL_NAME \
    -ocustom-columns='NAME:.metadata.name,BLOCK:.metadata.labels.cloud\.google\.com/gce-topology-block,SUBBLOCK:.metadata.labels.cloud\.google\.com/gce-topology-subblock,HOST:.metadata.labels.cloud\.google\.com/gce-topology-host' | sort -k2,4

NODE_POOL_NAME は、ノードプールの名前に置き換えます。

出力の VM 上の GKE ノードの物理トポロジを把握するには、次のノードラベルを参照してください。

  • cloud.google.com/gce-topology-block: VM が配置されている予約済みブロックの組織固有の ID。

  • cloud.google.com/gce-topology-subblock: VM が配置されているサブブロックの組織固有の ID。

  • cloud.google.com/gce-topology-host: VM が配置されているホストの ID。

  • kubernetes.io/hostname: Kubernetes ノードのホスト名。通常、このホスト名は GKE ノード名でもあります。

2 つの VM が共有するラベル値が多いほど、VM は物理的に近い場所に配置されます。これらの用語の詳細については、 用語をご覧ください。

Kueue を構成する

Kueue をインストールしたら、管理するインフラストラクチャを指定するように Kueue を構成する必要があります。通常、Kueue には、 ClusterQueue リソース 割り当て定義が必要です。これは、静的インフラストラクチャまたは クラスタの自動スケーリングが有効な動的インフラストラクチャのいずれかです。ClusterQueue は、ワークロードがリクエストするリソースが ClusterQueue で定義されたリソースプール以下の場合にのみ、ワークロードを受け入れます。このセクションで説明するように Kueue を構成すると、Kueue は TAS を使用して次のようにワークロードを受け入れます。

  • TAS ワークロード: Kueue は、物理 インフラストラクチャのトポロジと現在の使用状況の両方を確認します。

  • TAS 以外のワークロード: Kueue は、物理 インフラストラクチャのトポロジを確認しません。Kueue は構成で定義された割り当て全体を管理し、ノード割り当ては kube-scheduler に任せます。

ClusterQueue リソース割り当て定義を Kueue に提供する方法については、次の例をご覧ください。

  • 非常に高い割り当て: Kueue は、リクエストされたリソースに基づいてワークロードの受け入れを停止することはほとんどありません。TAS の定義に基づいて、Kueue はインフラストラクチャ トポロジに基づいてワークロードを受け入れる場合と受け入れない場合があります。詳細については、非常に高いリソース割り当てをご覧ください。

  • 現実的な割り当て: Kueue は、 ワークロードがリクエストするリソースがこのリソース割り当ての上限内にある場合にのみ、ワークロードを受け入れます。TAS の定義に基づいて、Kueue はワークロードを受け入れる前にインフラストラクチャ トポロジを確認します。詳細については、 現実的なリソース割り当てをご覧ください。

以降のセクションのリソース割り当てへの参照はすべて、ClusterQueue リソース割り当てを指します。

非常に高いリソース割り当て

次の例では、非常に高いリソース割り当てを使用しているため、Kueue は使用可能なリソース割り当てに基づいてワークロードを停止しません。代わりに、Kueue は使用可能なノードのトポロジ情報を使用して、トポロジをワークロードの要件に一致させようとします。

次のリソース割り当て定義を使用する手順は次のとおりです。

  1. 任意のファイル エディタを開きます。次に、次の割り当て定義を kueue-tas-config-very-high-quota.yaml という名前の YAML ファイルに含めます。

      apiVersion: kueue.x-k8s.io/v1alpha1
  kind: Topology
  metadata:
    name: "gke-default"
  spec:
    levels:
    - nodeLabel: "cloud.google.com/gce-topology-block"
    - nodeLabel: "cloud.google.com/gce-topology-subblock"
    - nodeLabel: "cloud.google.com/gce-topology-host"
    - nodeLabel: "kubernetes.io/hostname"
---
  kind: ResourceFlavor
  apiVersion: kueue.x-k8s.io/v1beta1
  metadata:
    name: "tas-flavor"
  spec:
    nodeLabels:
      cloud.google.com/gke-nodepool: "NODE_POOL_NAME"
    topologyName: "gke-default"
    tolerations:
    - key: "nvidia.com/gpu"
      operator: "Exists"
      effect: NoSchedule
---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ClusterQueue
  metadata:
    name: "tas-cluster-queue"
  spec:
    namespaceSelector: {}
    resourceGroups:
    - coveredResources: ["nvidia.com/gpu"]
      flavors:
      - name: "tas-flavor"
        resources:
        - name: "nvidia.com/gpu"
          nominalQuota: 10000000
---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: LocalQueue
  metadata:
    namespace: "default"
    name: "tas-user-queue"
  spec:
    clusterQueue: "tas-cluster-queue"

    NODE_POOL_NAME は、ノードプールの名前に置き換えます。

  2. Kueue ジョブ キューイング システムのリソース割り当て構成を作成して適用します。

    kubectl create -f kueue-tas-config-very-high-quota.yaml

現実的なリソース割り当て

前の例では、GPU リソースのみを構成しました。ただし、Kueue は Kubernetes 互換のすべてのリソースを管理できます。

次の例では、CPU、メモリ、GPU など、より現実的なリソース割り当てを定義します。これは 100 台の a3-ultragpu-8g マシン用です。1 台のマシンには、224 個の vCPU、2,944 GB のメモリ、8 個の GPU が搭載されています。

次のリソース割り当て定義を使用する手順は次のとおりです。

  1. 任意のファイル エディタを開きます。次に、次の割り当て定義を kueue-tas-config-real-quota.yaml という名前の YAML ファイルに含めます。

      apiVersion: kueue.x-k8s.io/v1alpha1
  kind: Topology
  metadata:
    name: "gke-default"
  spec:
    levels:
    - nodeLabel: "cloud.google.com/gce-topology-block"
    - nodeLabel: "cloud.google.com/gce-topology-subblock"
    - nodeLabel: "cloud.google.com/gce-topology-host"
    - nodeLabel: "kubernetes.io/hostname"
---
  kind: ResourceFlavor
  apiVersion: kueue.x-k8s.io/v1beta1
  metadata:
    name: "tas-flavor"
  spec:
    nodeLabels:
      cloud.google.com/gke-nodepool: "NODE_POOL_NAME"
    topologyName: "gke-default"
    tolerations:
    - key: "nvidia.com/gpu"
      operator: "Exists"
      effect: NoSchedule
---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: ClusterQueue
  metadata:
    name: "tas-cluster-queue"
  spec:
    namespaceSelector: {} # match all
    resourceGroups:
    - coveredResources: ["cpu", "memory", "nvidia.com/gpu"]
      flavors:
      - name: "tas-flavor"
        resources:
        # numbers below represent quota of 100 a3-ultragpu-8g machines
        - name: "cpu"
          nominalQuota: 22400
        - name: "memory"
          nominalQuota: 294400Gi
        - name: "nvidia.com/gpu"
          nominalQuota: 800
---
  apiVersion: kueue.x-k8s.io/v1beta1
  kind: LocalQueue
  metadata:
    namespace: "default"
    name: "tas-user-queue"
  spec:
    clusterQueue: "tas-cluster-queue"

    NODE_POOL_NAME は、ノードプールの名前に置き換えます。

  2. Kueue ジョブ キューイング システムのリソース割り当て構成を作成して適用します。

    kubectl create -f kueue-tas-config-real-quota.yaml

    出力は次のようになります。

    topology.kueue.x-k8s.io/gke-default created
resourceflavor.kueue.x-k8s.io/tas-flavor created
clusterqueue.kueue.x-k8s.io/tas-cluster-queue created
localqueue.kueue.x-k8s.io/tas-user-queue created

Kueue を使用して TAS でワークロードをスケジュールする

次のシナリオでは、トポロジ リクエスト タイプとトポロジ リクエスト レベルを使用して、一般的なワークロードとインフラストラクチャの組み合わせを管理するように Kueue と TAS に指示する方法を示します。

  • 使用可能なトポロジ リクエスト タイプ（優先または必須）は次のとおりです。

    • kueue.x-k8s.io/podset-preferred-topology: Kueue は、特定のトポロジ レベル内でワークロード全体をスケジュールすることを優先しますが、このトポロジ レベルに収まらないワークロードも受け入れます。単一のトポロジ レベルに収まる可能性のあるワークロードの場合、Kueue はそのワークロードをそのトポロジ レベルの複数のインスタンスにスケジュールする可能性があります。

    • kueue.x-k8s.io/podset-required-topology: ワークロード全体が選択したトポロジ レベルに収まるまで、Kueue はこのワークロードの受け入れを試行し続けます。

  • 使用可能なトポロジ リクエスト レベルは次のとおりです。Job の実行を優先または 必須とする物理インフラストラクチャをより具体的に指定できます。

    • cloud.google.com/gce-topology-block

    • cloud.google.com/gce-topology-subblock

    • cloud.google.com/gce-topology-host

    • kubernetes.io/hostname

これらの値を使用してワークロードをスケジュールするには、次の Job YAML ファイルを使用します。

apiVersion: batch/v1
kind: Job
metadata:
  generateName: JOB_NAME
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
spec:
  parallelism: NUMBER_OF_REPLICAS
  completions: NUMBER_OF_REPLICAS
  completionMode: Indexed
  template:
    metadata:
      annotations:
        ANNOTATIONS_STRING
    spec:
      containers:
      - name: dummy-job
        image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
        args: ["60s"]
        resources:
          requests:
            nvidia.com/gpu: "1"
          limits:
            nvidia.com/gpu: "1"
      restartPolicy: Never

次の変数を置き換えます。

  • JOB_NAME: Job の名前。

  • NUMBER_OF_REPLICAS: 並行して実行されている Pod の数。

  • ANNOTATIONS_STRING: 次の表をご覧ください。

    リクエストされたトポロジ タイプとレベル 説明 ANNOTATIONS_STRING
    ホスト名 内で実行することを優先 （推奨） この構成では、容量が断片化されていても、ワークロードのリソース要件を満たすのに十分なリソースが使用可能であれば、ワークロードを受け入れます。Kueue は、Pod をできるだけコンパクトにスケジュールします。 kueue.x-k8s.io/podset-preferred-topology: "kubernetes.io/hostname"
    ホスト 内で実行することが必須

    この構成では、ワークロードのリソース要件を満たすのに十分なリソースを持つホストが使用可能な場合にのみ、ワークロードを受け入れます。

    これは、ホストごとに複数の VM がある場合（小規模なマシンタイプなど）や、1 つのノードで複数の Pod を実行できる場合に便利です。このような場合、ワークロードが受け入れられると、単一のホストで実行されます。

    		 kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-host"
    ホスト 内で実行することを優先 この構成では、容量が断片化されていても、ワークロードのリソース要件を満たすのに十分なリソースが使用可能であれば、ワークロードを受け入れます。Kueue は、ホスト内で Pod をスケジュールしようとし、必要に応じて追加のホストを使用します。 kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-host"
    サブブロック 内で実行することが必須 この構成では、ワークロードのリソース要件を満たすのに十分なリソースを持つサブブロックが使用可能な場合にのみ、ワークロードを受け入れます。 kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-subblock"
    サブブロック 内で実行することを優先 この構成では、容量が断片化されていても、ワークロードのリソース要件を満たすのに十分なリソースが使用可能であれば、ワークロードを受け入れます。Kueue は、サブブロック内で Pod をスケジュールしようとし、必要に応じて追加のサブブロックを使用します。この場合、Kueue は、要件を満たすのに十分な容量を持つサブブロックと比較して、断片化されていても使用可能な容量が多いサブブロックを上位にランク付けします。 kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-subblock"
    ブロック 内で実行することが必須 この構成では、ブロック内で使用可能なリソースがワークロードのリソース要件を満たす場合にのみ、ワークロードを受け入れます。受け入れられた場合、Kueue はワークロードをスケジュールするサブブロックとホストの数を最小限に抑えます。これにより、使用可能な容量が断片化される可能性があります。 kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
    ブロック 内で実行することを優先 この構成では、容量が断片化されていても、ワークロードのリソース要件を満たすのに十分なリソースが使用可能であれば、ワークロードを受け入れます。Kueue は、ブロック内で Pod をスケジュールしようとし、必要に応じて追加のブロックを使用します。 kueue.x-k8s.io/podset-preferred-topology: "cloud.google.com/gce-topology-block"

Kueue を使用して TAS で PodGroup を使用してワークロードをスケジュールする

PodGroup を使用する場合は、PodGroup 内のすべての Pod に対して 3 つの追加フィールドを指定する必要があります。

使用する ML フレームワークに応じて、PodGroup のリーダーに GPU が必要な場合と不要な場合があります。Kueue の制限により、これらのケースは別々に処理する必要があります。次の例では、1 つのリーダーと 2 つのワーカーで構成される 3 つの Pod の PodGroup を作成する方法を示します。

ケース 1: リーダーもワーカーであり、GPU が必要

リーダーがワーカーの 1 つであり、GPU も必要な場合、リーダーは PodGroup 内の任意の数を持つことができます。簡単にするため、次の例ではリーダーのインデックスは 0 です。

apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-leader-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group"
    kueue.x-k8s.io/pod-group-pod-index: "0"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  containers:
  - name: leader
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"
  restartPolicy: Never
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-1-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group"
    kueue.x-k8s.io/pod-group-pod-index: "1"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-2-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group"
    kueue.x-k8s.io/pod-group-pod-index: "2"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"

ケース 2: リーダーはワーカーではなく、GPU は不要

Kueue の制限により、リーダーがワーカーの 1 つでない場合、Kueue が PodSet を作成する方法により、リーダーは PodGroup の最後のインデックスを持つ必要があります。リーダーが最後のインデックスを持たず、最初のワーカーが最初のインデックスを使用しない場合、Kueue はランク割り当てを適用しません。

次の例をご覧ください。

---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-leader-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group2"
    kueue.x-k8s.io/pod-group-pod-index: "2"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  containers:
  - name: leader
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        cpu: "1"
      limits:
        cpu: "1"
  restartPolicy: Never
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-0-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group2"
    kueue.x-k8s.io/pod-group-pod-index: "0"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"
---
apiVersion: v1
kind: Pod
metadata:
  generateName: tas-podgroup-worker-1-
  labels:
    kueue.x-k8s.io/queue-name: tas-user-queue
    kueue.x-k8s.io/pod-group-name: "tas-podgroup-example-group2"
    kueue.x-k8s.io/pod-group-pod-index: "1"
  annotations:
    kueue.x-k8s.io/pod-group-total-count: "3"
    kueue.x-k8s.io/podset-required-topology: "cloud.google.com/gce-topology-block"
spec:
  restartPolicy: Never
  containers:
  - name: worker
    image: gcr.io/k8s-staging-perf-tests/sleep:v0.1.0
    args: ["600s"]
    resources:
      requests:
        nvidia.com/gpu: "1"
      limits:
        nvidia.com/gpu: "1"

次のステップ