GKE 動的スライスを使用して TPU を構成する

このドキュメントでは、Google Kubernetes Engine（GKE）で動的スライスを使用する方法について説明します。動的スライスを使用すると、プロビジョニングされた TPU サブブロックをさまざまなトポロジに構成できます。この機能により、ノードプールを再作成する必要性が軽減され、障害発生時の自動復旧が可能になることでフォールト トレランスが強化され、リソース使用率が最適化されます。

このドキュメントは、TPU の使用率を最適化し、プロビジョニング時間を短縮し、大規模なトレーニングと推論のワークロードのフォールト トレランスを向上させたい AI/ML エンジニアとプラットフォーム管理者を対象としています。

このドキュメントを読む前に、次をよく理解しておいてください。

動的スライスとは

動的スライスを使用すると、TPU のプロビジョニングを切り離すことができるため、Cloud TPU の容量を柔軟に管理できます。動的スライスには、次のプロセスが含まれます。

  1. リソースは、サブブロックと呼ばれる小さな単位でプロビジョニングします。サブブロックは、Ironwood（TPU7x）容量の基本的な論理ビルディング ユニットです。Ironwood（TPU7x）の場合、相互接続された TPU チップの 4x4x4 トポロジを持つ TPU VM の 16 ノード グループを表します。TPU All Capacity モードと動的スライシングのコンテキストでは、ノードプールはサブブロックに直接マッピングされます。
  2. 動的スライスでは、これらのサブブロックを結合して大きなスライスにします。

動的スライスのメリット

動的スライスを使用すると、次のことが可能になります。

  • プロビジョニング時間を短縮する: サブブロックを個別にプロビジョニングすると、単一の障害の影響が最小限に抑えられるため、全体的なプロビジョニングが高速化されます。
  • 復旧時間を短縮する: TPU チップの障害が発生した場合、障害の最小単位はサブブロックです。動的スライスは、障害のあるサブブロックを分離するため、ワークロードを大規模なスライス全体を再プロビジョニングするよりも迅速に正常なサブブロックに再スケジュールできます。
  • 容量の再形成: ワークロード要件が多様な場合、動的スライスなしではトポロジの変更のためにノードプールを削除して再作成する必要がありますが、動的スライスを使用するとその必要がなくなります。代わりに、指定されたシェイプに合わせてプロビジョニングされたノードプールを動的に再構成できます。

動的スライシングの主な要素

動的スライスでは、次の重要なコンセプトが導入されています。

  • ノードプールの増分プロビジョニング: 動的スライスでは、ノードプールのフォールト トレラントなプロビジョニング モデルである増分プロビジョニングが使用されます。このモデルは、すべての TPU 容量を 16 ノードの TPU VM グループのノードプールに変換します。
  • スライス コントローラ: GKE コントロール プレーン内で実行され、動的スライスを管理する Kubernetes カスタム リソース コントローラ。スライス コントローラは、動的スライスを表す Slice カスタム リソースのライフサイクルを管理します。スライス コントローラは、スライスの作成、継続的なモニタリング、削除を処理します。スケジューラを使用すると、スケジューラは Slice カスタム リソースの作成と削除を指示します。
  • スライス カスタム リソース: リクエストされた TPU トポロジに基づいて、これらのサブブロックを動的に結合します。このプロセスでは、OCS ネットワークの動的再構成を利用して TPU ノードプールを接続し、パフォーマンスの最適化を実現します。Slice カスタム リソースのステータス フィールドを調べることで、動的スライスの形成の進行状況や健全性を確認できます。

要件

GKE で動的スライスを使用するには、次の要件を満たす必要があります。

  • バージョン 1.35.0-gke.274500 以降の Standard クラスタ（Rapid チャンネル）を使用します。
  • Ironwood（TPU7x）バージョンを使用します。
  • ノードに Container-Optimized OS イメージを使用します。
  • 増分プロビジョニングを使用するには、All Capacity モードの予約を使用します。All Capacity モードは、TPU Cluster Director によって有効になる機能です。

始める前に

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

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

制限事項

  • 単一のスライスは、予約内の同じ TPU ブロック内のサブブロックを使用する必要があります。TPU ブロック間でサブブロックを使用するには、TPU マルチスライスを使用します。
  • 動的スライスは、4x4x4 より小さいトポロジをサポートしていません。

Kueue を使用して GKE で動的スライスを使用する

このセクションでは、GKE で動的スライスを使用するワークフローについて説明します。

  1. All Capacity モードの予約のトポロジと健全性ステータスを表示する。
  2. クラスタでスライス コントローラを有効にします。
  3. TPU ノードプールを作成します。
  4. Slice カスタム リソースを作成するように Kueue を構成します。
  5. Kueue を使用して動的スライスでワークロードを実行する。
  6. クリーンアップします。

スライス コントローラを有効にする

動的スライスを使用するには、クラスタでスライス コントローラを有効にします。

  1. クラスタを更新します。

    gcloud container clusters update CLUSTER_NAME \
    --location=LOCATION \
    --enable-slice-controller

    次のように置き換えます。

  2. kubectl コマンドを使用してクラスタと通信できるように、認証情報を取得します。

    gcloud config set container/cluster CLUSTER_NAME
gcloud container clusters get-credentials CLUSTER_NAME \
    --location=LOCATION

  3. 次のコマンドの出力で、slices.accelerator.gke.io 値が存在することを確認します。

    kubectl get crd slices.accelerator.gke.io

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

    slices.accelerator.gke.io                2026-01-09T23:58:02Z

増分プロビジョニングでノードプールを作成する

このセクションでは、増分プロビジョニングを使用して TPU ノードプールを作成する方法について説明します。GKE は、すべての TPU 容量を 16 ノードの TPU VM グループ（サブブロック）のノードプールに変換します。GKE は、ホストマシンの正常な部分にノードを配置し、異常なマシンを修復しながら段階的にプロビジョニングすることで、16 個の正常な VM をすべて検出できない場合でも、これらのノードプールをプロビジョニングします。

ノードプールは次のいずれかに属するようにターゲット設定できます。

  • All Capacity モードの予約で公開される特定の TPU ブロック。ブロック ターゲティングを使用すると、GKE は指定されたブロック内の使用可能なサブブロックにノードプールを作成できます。
  • よりきめ細かい制御を行うための、TPU の特定のサブブロックまたは特定の 16 ノードの TPU VM グループ。

ワークロード ポリシーの作成

Ironwood（TPU7x）で TPU スライス ノードプールを作成するには、まず accelerator-topology-mode フィールドが provision_only に設定されたワークロード ポリシーを作成する必要があります。この設定により、増分プロビジョニング プロセスがトリガーされます。

ワークロード ポリシーを作成します。

gcloud compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \
        --project=PROJECT_ID \
        --region=REGION  \
        --type=HIGH_THROUGHPUT \
        --accelerator-topology=4x4x4 \
        --accelerator-topology-mode=provision_only

次のように置き換えます。

  • WORKLOAD_POLICY_NAME: ワークロード ポリシーの名前。
  • PROJECT_ID: 実際の Google Cloud プロジェクト ID。
  • REGION: ワークロード ポリシーのリージョン。

このコマンドでは、次の操作を行います。

  • accelerator-topology フィールドは常に 4x4x4 に設定して、単一のサブブロック内のチップの総数と一致させます。
  • 増分プロビジョニング プロセスがトリガーされるように、accelerator-topology-mode フィールドは常に provision_only に設定してください。provision_only フィールドが設定されている場合、ノードプールは ICI リンクまたは OCS リンクを形成せずに TPU ノードをプロビジョニングします。

ノードプールがブロックまたはサブブロックに属するようにターゲットを設定する

All Capacity モードの予約内の特定のサブブロックまたはブロックをターゲットにできます。

  • ブロックをターゲットにする: 各ノードプールは、指定されたブロックの容量を使用します。GKE は、そのブロック内の使用可能なサブブロックにノードプールを配置します。使用するブロックのサブブロックの数だけノードプールを作成する必要があります。

  • サブブロックをターゲットにする: 各ノードプールは、特定の利用可能なサブブロックにマッピングされます。サブブロック ターゲティングを使用する場合、GKE は少なくとも 1 つの VM が正常であればノードプールを作成します。増分プロビジョニングにより、すべてのノードが指定されたサブブロック内に配置されます。

ブロック

  1. 予約のブロックの名前と、ブロックで使用可能なサブブロックの数を取得するには、All Capacity モードの予約のトポロジと健全性ステータスを表示するドキュメントの手順を完了します。

    1. すべての予約ブロックを一覧表示し、name: フィールドの値をコピーして、ブロックの名前を特定します。この値は、このドキュメント内のブロックまたは BLOCK_NAME の名前です。

    2. 予約ブロックを記述し、reservationSubBlockCount フィールドの値を確認して、作成するノードプールの数を決定します。この値は、使用可能なサブブロックの数です。たとえば、reservationSubBlockCount: 4 値は、ブロックに 4 つのサブブロックが使用可能であることを示します。この場合、4 つの個別のノードプールを作成する必要があります。

  2. 予約パスを設定します。

    export RESERVATION_PATH="projects/PROJECT_ID/reservations/RESERVATION_NAME/reservationBlocks/BLOCK_NAME"

    次のように置き換えます。

    • RESERVATION_NAME: TPU 予約の名前。
    • BLOCK_NAME: ブロックの名前。

  3. 前の手順で特定したサブブロックごとにノードプールを作成します。たとえば、カウントが 4 の場合は、このコマンドを 4 回実行します。ノードプールごとに一意の名前を使用します。

    gcloud container node-pools create NODE_POOL_NAME \
      --cluster=CLUSTER_NAME \
      --node-locations=ZONE \
      --machine-type=tpu7x-standard-4t \
      --num-nodes=16 \
      --placement-policy=WORKLOAD_POLICY_NAME \
      --reservation-affinity=specific \
      --reservation=${RESERVATION_PATH}

    次のように置き換えます。

    • NODE_POOL_NAME: 新しいノードプールの名前。
    • CLUSTER_NAME: GKE クラスタの名前。
    • WORKLOAD_POLICY_NAME: 作成したワークロード ポリシーの名前。
    • ZONE: ノードプールのゾーン（us-central1-a など）。

サブブロック

  1. ブロックの名前と使用可能なサブブロックの ID を取得するには、All Capacity モードの予約のトポロジと健全性ステータスを表示するドキュメントで次の手順を完了します。

    1. ブロックの名前を確認するには、すべての予約ブロックを一覧表示し、name: フィールドの値をコピーします。この値は、このドキュメントのブロックまたは BLOCK_NAME の名前です。

    2. サブブロックの名前を特定するには、ブロックのすべてのサブブロックを一覧表示し、reservationSubBlocks の各エントリの name: フィールドの値をコピーします。この値は、このドキュメントのサブブロックまたは SUBBLOCK_NAME の名前です。

  2. 予約パスを設定します。

    export RESERVATION_PATH="projects/PROJECT_ID/reservations/RESERVATION_NAME/reservationBlocks/BLOCK_NAME/reservationSubBlocks/SUBBLOCK_NAME"

    次のように置き換えます。

    • RESERVATION_NAME: TPU 予約の名前。
    • BLOCK_NAME: ブロックの名前。
    • SUBBLOCK_NAME: サブブロックの名前。

  3. ノードプールを作成します。

    gcloud container node-pools create NODE_POOL_NAME \
        --project=PROJECT_ID \
        --cluster=CLUSTER_NAME \
        --node-locations=ZONE \
        --machine-type=tpu7x-standard-4t \
        --num-nodes=16 \
        --placement-policy=WORKLOAD_POLICY_NAME \
        --reservation-affinity=specific \
        --reservation=${RESERVATION_PATH}

    次のように置き換えます。

    • NODE_POOL_NAME: 新しいノードプールの一意の名前（sub-block-pool-1 など）。
    • PROJECT_ID: 実際の Google Cloud プロジェクト ID。
    • CLUSTER_NAME: GKE クラスタの名前。
    • ZONE: ノードプールのゾーン（us-central2-b など）。
    • WORKLOAD_POLICY_NAME: 作成したワークロード ポリシーの名前。

この段階では、ノードは作成されますが、チップ間インターコネクト（ICI）リンクはまだアクティブではありません。そのため、これらのノードプールでワークロードを直接実行することはできません。

必要な ICI リンクをすべて有効にしてスライスを形成し、ワークロードのスケジュール設定を許可するには、次のいずれかの方法で動的スライスを作成します。

  • スライス カスタム リソースを作成します。Pod の代わりに、Slice カスタム リソースを使用して指定されたトポロジを定義します。このトポロジは、スライス コントローラによって有効になります。
  • Kueue と TAS を使用して GKE ワークロードをスケジュールします。Kueue は、Slice カスタム リソースの作成と削除を自動的に処理します。Kueue によって作成された Slice カスタム リソースを手動で変更しないでください。

Kueue と TAS を使用して動的スライスを作成する

このセクションでは、Kueue と TAS を使用して GKE ワークロードをスケジュールします。

動的スライス用の JobSet リソースと Kueue リソースをインストールする

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

    helm install jobset oci://registry.k8s.io/jobset/charts/jobset \
        --version 0.10.1 \
        --namespace jobset-system \
        --create-namespace \
        --set controller.resources.requests.cpu=4 \
        --set controller.resources.requests.memory=16Gi

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

    helm install kueue oci://registry.k8s.io/kueue/charts/kueue \
        --version 0.16.1 \
        --namespace kueue-system \
        --create-namespace \
        --wait \
        --set controllerManager.replicas=3 \
        --set controllerManager.manager.resources.requests.cpu=16 \
        --set controllerManager.manager.resources.requests.memory=64Gi

  3. Kueue スライス コントローラをインストールします。

    kubectl apply -f https://gist.githubusercontent.com/mwysokin/cd90010d0d375b3bf57c536905692547/raw/506c36dd070f4ac222ba8a5e58ba28bbfcfa8ed3/kueue-slice-controller-v0.8.0-130.yaml

  4. 動的スライス用に Kueue を構成するには、次のマニフェストを dynamic-slice-topology.yaml として保存します。

    apiVersion: kueue.x-k8s.io/v1beta1
kind: Topology
metadata:
  name: superslice-topology
spec:
  levels:
  # Label to identify the physical block a sub-block belongs to.
  # Only sub-blocks from the same block can form a slice.
  - nodeLabel: cloud.google.com/gce-topology-block
  # Label to identify individual TPU sub-blocks (4x4x4 topology).
  - nodeLabel: cloud.google.com/gke-tpu-partition-4x4x4-id
  # Standard Kubernetes label for individual nodes.
  # Required to assign Pods to specific VMs.
  - nodeLabel: kubernetes.io/hostname
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ResourceFlavor
metadata:
  name: superslice-rf
spec:
  nodeLabels:
    cloud.google.com/gke-tpu-accelerator: tpu7x
  topologyName: superslice-topology
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: AdmissionCheck
metadata:
  name: superslice-ac
spec:
  controllerName: accelerator.gke.io/slice
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: ClusterQueue
metadata:
  name: cq
spec:
  namespaceSelector: {}
  admissionChecks:
  - superslice-ac
  resourceGroups:
  - coveredResources:
    - google.com/tpu
    flavors:
    - name: superslice-rf
      resources:
      - name: google.com/tpu
        nominalQuota: "999999"  # modeling unlimited quota
---
apiVersion: kueue.x-k8s.io/v1beta1
kind: LocalQueue
metadata:
  name: lq
  namespace: default
spec:
  clusterQueue: cq

  5. dynamic-slice-topology.yaml マニフェストを適用します。

    kubectl apply -f dynamic-slice-topology.yaml

    このマニフェストでは、次のリソースを定義して、動的スライス用に Kueue を構成します。

    • Ironwood（TPU7x）動的スライス トポロジ（superslice-topology）: トポロジは、Kueue が動的スライス ワークロードをスケジュールするときに考慮するレベルを定義します。これらのレベルは次のとおりです。
      • cloud.google.com/gce-topology-block ラベル: 同じブロックのサブブロックのみがスライスを形成できるため、どのサブブロックがどのブロックに属するかを理解するには、このレベルが必要です。
      • cloud.google.com/gke-tpu-partition-4x4x4-id ラベル: このレベルは、個々の Ironwood（TPU7x）サブブロック（4x4x4 トポロジ）を表します。
      • kubernetes.io/hostname ラベル: 特定の VM に Pod を割り当て、ラベルと taint を確認するには、このレベルが必要です。
    • Ironwood（TPU7x）SuperSlice ResourceFlavor（superslice-rf）: Ironwood（TPU7x）サブブロックのリソース フレーバーには、Ironwood（TPU7x）マシンとノードを照合するための cloud.google.com/gke-tpu-accelerator: tpu7x ラベルが含まれています。
    • SuperSlice AdmissionCheck（superslice-ac）: このアドミッション チェックは、GKE スライス コントローラがスライスがアクティブになったことを確認するまで、ワークロードをスケジュールしないように Kueue に指示します。アドミッション チェックが最初に定義され、動的スライス ワークロードを処理する ClusterQueue に追加されます。
    • ClusterQueue（cq）と LocalQueue（lq）: これらのフィールドは google.com/tpu リソースを管理します。cq ClusterQueue には superslice-ac アドミッション チェックが含まれています。google.com/tpu の nominalQuota は、次の 2 つの方法で構成できます。
      • 特定の割り当て: フェア シェアリングと割り当て管理のために、既存の容量に合わせて nominalQuota を設定します。
      • 無制限の割り当て: nominalQuota を "999999" などの非常に高い値に設定して、無制限の割り当てをモデル化します。TAS と動的スライスに焦点を当てるため、この構成では Kueue の割り当て管理機能がバイパスされます。

サブブロックの健全性の選択を定義する

標準のノードの健全性と準備状況に加えて、GKE は cloud.google.com/gke-tpu-partition-4x4x4-state ラベルを使用して各サブブロックの特定の状態を公開します。このラベルを使用すると、GKE は TPU リンクの状態など、スライス形成に影響する要因を考慮できます。

cloud.google.com/gke-tpu-partition-4x4x4-state ラベルの値は次のように定義できます。

  • HEALTHY: インフラストラクチャは正常です。
  • DEGRADED: OCS リンクの劣化などにより、サブブロックのインフラストラクチャが劣化状態です。サブブロックはスライスを形成できますが、健全なサブブロックと比較して全体的なパフォーマンスが低下する可能性があります。パフォーマンスの低下を許容できる場合は、例 3 に示すように、ノード アフィニティを使用して DEGRADED サブブロックを使用するようにワークロードを構成できます。
  • UNHEALTHY: サブブロックが異常で、スライスを形成できません。

Kueue スライス コントローラ Webhook は、ワークロードに特定の下位ブロックの健全性要件が含まれているかどうかを検証します。設定が指定されていない場合、Webhook はデフォルトのノード アフィニティを挿入します。

動作は次のとおりです。

  • cloud.google.com/gke-tpu-partition-4x4x4-state ラベルをターゲットとする nodeSelector または nodeAffinity が存在する場合、それらは変更されません。

  • このようなラベル構成が存在しない場合、Webhook は次のデフォルトのノード アフィニティを挿入して、使用可能なサブブロックのみが使用されるようにします。

    nodeSelector:
  cloud.google.com/gke-tpu-partition-4x4x4-state: "HEALTHY"

次のセクションでは、cloud.google.com/gke-tpu-partition-4x4x4-state ラベルが構成され、さまざまなサブブロックの健全性構成を指定する例を示します。

Kueue を使用して動的スライスでテスト ワークロードを実行する

このセクションでは、Kueue と TAS を使用して動的スライスにワークロードをデプロイする方法について説明します。このサンプルには、動的スライス ワークロードと複数のスライスで構成されるワークロードを作成する方法を示す 3 つの例が含まれています。ワークロードは JobSet として送信されます。

例 1: 単一のワークロードが単一の動的スライスを使用する

次の例は、12 個のサブブロックで構成される 4x12x16 トポロジのスライスを使用してワークロードを作成する方法を示しています。Pod の数は、(4 * 12 * 16) / ノードあたり 4 チップ = 192 Pod として計算されました。

  1. 次のマニフェストを big-super-slice.yaml として保存します。

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: big-super-slice
  labels:
    kueue.x-k8s.io/queue-name: lq
  annotations:
spec:
  replicatedJobs:
    - name: job-jax
      replicas: 1
      template:
        spec:
          parallelism: 192  # pods per slice calculation: 4*12*16 / 4 = 192
          completions: 192
          backoffLimit: 10
          template:
            metadata:
              annotations:
                cloud.google.com/gke-tpu-slice-topology: 4x12x16
            spec:
              tolerations:
                - key: "google.com/tpu"
                  operator: "Equal"
                  value: "present"
                  effect: "NoSchedule"
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu7x
                cloud.google.com/gke-tpu-partition-4x4x4-state: "HEALTHY"
              containers:
                - name: jax
                  image: python:latest
                  command:
                    - bash
                    - -c
                    - |
                      printenv
                      pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                      python -c 'import jax; print("Global device count:", jax.device_count(), "Local device count:", jax.local_device_count())'
                  resources:
                    limits:
                      google.com/tpu: 4
              restartPolicy: Never

    このマニフェストでは、次のアノテーションによって、Kueue にスライスの特性とトポロジが通知され、次の構成が行われます。

    • cloud.google.com/gke-tpu-slice-topology: 動的スライス トポロジとして "4x12x16" を指定します。tpu7x アクセラレータ トポロジの要件には、次のルールが含まれます。
      • 最小のトポロジは 4x4x4 です。
      • トポロジは、AxBxC 形式の 3 次元文字列にする必要があります。例: 4x8x8
      • 各ディメンション（A、B、C）は 4 の倍数にする必要があります。
      • ディメンションは非減少順（A <= B <= C）で並べ替える必要があります。たとえば、4x8x4 は無効です。4x4x8 にする必要があります。
      • 寸法（ABC）の積は 9,216 を超えないようにしてください。
      • サポートされている最大のスライス トポロジには、最大 32 個のサブブロックを含めることができます。たとえば、32 個のサブブロックを含む 8x16x16、30 個のサブブロックを含む 8x12x20、27 個のサブブロックを含む 12x12x12 は、許容範囲内です。
    • cloud.google.com/gke-tpu-accelerator: tpu7x: Ironwood（TPU7x）を実行する VM に Pod をスケジュールします。
    • kueue.x-k8s.io/queue-name: JobSet を Kueue LocalQueue に割り当てます。

  2. big-super-slice.yaml マニフェストを適用します。

    kubectl apply -f big-super-slice.yaml

    マニフェストを適用すると、Kueue は big-super-slice という名前の JobSet を作成します。Kueue は、4x12x16 トポロジを使用して単一の動的スライスを形成しようとします。スライスがアクティブになると、Kueue はワークロードを許可し、192 個の Pod がノードでスケジュールされて、ワークロードを実行する動的スライスが形成されます。

例 2: 複数のレプリカを持つワークロード

次の例は、それぞれ 4 つのサブブロックで構成される 2 つの動的スライスを使用するワークロードを作成する方法を示しています。

  1. 次のマニフェストを two-super-slices.yaml として保存します。

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
name: two-super-slices
labels:
    kueue.x-k8s.io/queue-name: lq
annotations:
spec:
replicatedJobs:
    - name: job-jax
    replicas: 2
    template:
        spec:
        parallelism: 64  # Pods per slice calculation: (4*8*8) / 4 = 64
        completions: 64
        backoffLimit: 10
        template:
            metadata:
            annotations:
                cloud.google.com/gke-tpu-slice-topology: 4x8x8
            spec:
            tolerations:
                - key: "google.com/tpu"
                operator: "Equal"
                value: "present"
                effect: "NoSchedule"
            nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu7x
                cloud.google.com/gke-tpu-partition-4x4x4-state: "HEALTHY"
            containers:
                - name: jax
                image: python:latest
                command:
                    - bash
                    - -c
                    - |
                    printenv
                    pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                    python -c 'import jax; print("Global device count:", jax.device_count(), "Local device count:", jax.local_device_count())'
                resources:
                    limits:
                    google.com/tpu: 4
            restartPolicy: Never

  2. two-super-slices.yaml マニフェストを適用します。

    kubectl apply -f two-super-slices.yaml

このマニフェストでは、replicatedJobs フィールドで replicas: 2 を設定します。マニフェストを適用すると、Kueue は 4x8x8 トポロジを使用して 2 つの個別のスライスを形成しようとします。Kueue は、jobset.spec.replicatedJobs[].replicas で定義されたレプリカごとに動的スライスを作成します。n レプリカが指定されている場合、Kueue はワークロード用に n 個の動的スライスを作成し、すべてのスライスがアクティブになるまで待機してからワークロードを承認します。

例 3: 単一の動的スライスと NodeAffinity を使用するワークロード

Kueue 0.15 以降、Kueue は TAS ノード選択の NodeAffinity をサポートしています。この機能を使用すると、HEALTHY ノードと DEGRADED ノードの両方を動的スライスの一部にすることができます。次の例は、単一の動的スライスと NodeAffinity を使用してワークロードを構成する方法を示しています。

  1. 次のマニフェストを slice-8x8x8-na.yaml として保存します。

    apiVersion: jobset.x-k8s.io/v1alpha2
kind: JobSet
metadata:
  name: slice-8x8x8-na
  labels:
    kueue.x-k8s.io/queue-name: lq
spec:
  replicatedJobs:
    - name: rj1
      replicas: 1
      template:
        spec:
          parallelism: 128
          completions: 128
          backoffLimit: 10
          template:
            metadata:
              annotations:
                cloud.google.com/gke-tpu-slice-topology: 8x8x8
            spec:
              tolerations:
                - key: "google.com/tpu"
                  operator: "Equal"
                  value: "present"
                  effect: "NoSchedule"
              nodeSelector:
                cloud.google.com/gke-tpu-accelerator: tpu7x
              affinity:
                nodeAffinity:
                  requiredDuringSchedulingIgnoredDuringExecution:
                    nodeSelectorTerms:
                      - matchExpressions:
                          - key: cloud.google.com/gke-tpu-partition-4x4x4-state
                            operator: In
                            values:
                              - "HEALTHY"
                              - "DEGRADED"
              containers:
                - name: jax
                  image: python:latest
                  command:
                    - bash
                    - -c
                    - |
                      printenv
                      pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                      python -c 'import jax; print("Global device count:", jax.device_count(), "Local device count:", jax.local_device_count())'
                  resources:
                    limits:
                      google.com/tpu: 4
              restartPolicy: Never

  2. slice-8x8x8-na.yaml マニフェストを適用します。

    kubectl apply -f slice-8x8x8-na.yaml

    マニフェストを適用すると、Kueue は slice-8x8x8-na という名前の JobSet を作成します。次に、Kueue は 8x8x8 トポロジを持つ単一の動的スライスを形成しようとします。これにより、指定された NodeAffinity により、HEALTHY ノードと DEGRADED ノードの両方を含めることができます。スライスがアクティブになると、Kueue はワークロードを承認し、128 個の Pod が動的スライスを形成するノードにスケジュールされます。

スライスのステータスをモニタリングする

動的スライスのステータスを確認するには、次のコマンドを実行します。

kubectl describe slice SLICE_NAME

SLICE_NAME は、スライスの名前に置き換えます。通常、スライス名は JobSet 名とレプリカ インデックスから派生します。例 1 の場合、Kueue によって作成されたスライスの名前は default-jobset-big-super-slice-yyyyy-job-jax-0 のようになります。

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

Name:         test-slice
Namespace:
Labels:       <none>
Annotations:  <none>
API Version:  accelerator.gke.io/v1beta1
Kind:         Slice
Metadata:
  Creation Timestamp:  2026-02-12T23:44:28Z
  Finalizers:
    accelerator.gke.io/slice-finalizer
  Generation:        1
  Resource Version:  1770939905695871008
  UID:               6dbbfe14-4486-4462-864d-e078d0ca8b5b
Spec:
  Partition Ids:
    5eae6a4f59d59cf30a9bf49de618eb2b
  Topology:  4x4x4
  Type:      tpu7x
Status:
  Conditions:
    Last Transition Time:  2026-02-12T23:45:05Z
    Message:
    Reason:                ACTIVE
    Status:                True
    Type:                  Ready
    Last Transition Time:  2026-02-12T23:45:05Z
    Message:               NodeLabelingCompleted
    Reason:                NodeLabelIsAdded
    Status:                True
    Type:                  NodeLabeled
Events:                    <none>

スライス名は、基盤となる Compute Engine リソースの命名規則との互換性を確保するため、次のルールに従います。

  • テンプレート: {namespace}-jobset-{jobset.metadata.name}-kueueHash[5-character]-{jobset.spec.replicatedJobs[].name}-sliceIndex。
  • 長さ: 名前が 54 文字以下である。コントローラは、ハイフンと 8 文字のクラスタ ハッシュを追加して、63 文字の制限がある Compute Engine リソース名を作成します。
  • 形式: 名前が正規表現 ^[a-z]([-a-z0-9]*[a-z0-9])?$ と一致します。名前には次の特性があります。
    • 先頭は小文字でなければなりません。
    • 小文字、数字、ハイフン（-）のみが含まれています。
    • 末尾は英小文字または数字にします（末尾をハイフンにすることはできません）。

クリーンアップ

予期しない請求が発生しないようにするには、ノードプールを削除する前にスライスを削除してください。

  1. JobSet を削除します。このアクションにより、Kueue が関連する Slice カスタム リソースを削除します。

    kubectl delete jobset JOBSET_NAME

    JOBSET_NAME は、JobSet の名前に置き換えます（big-super-slice など）。

  2. TPU ノードプールを削除します。

    gcloud container node-pools delete NODE_POOL_NAME \
    --cluster=CLUSTER_NAME \
    --location=LOCATION

（省略可）独自のスケジューラで動的スライスを使用する

このドキュメントでは、Kueue と TAS の使用に焦点を当てます。ただし、独自のカスタム スケジューラを使用して動的スライスを管理することもできます。別のスケジューラを使用する場合は、Slice カスタム リソースのリファレンス情報を参照してください。

次のステップ