カスタム スケジューラで動的スライスを使用する

このドキュメントでは、 Slice カスタム リソースを直接操作して動的スライスを使用する方法について説明します。スライスを作成し、パーティションの状態をモニタリングして、スライスの健全性を確認できます。

以下の手順を行う前に、動的スライスの コンセプトを理解しておいてください

カスタム スケジューラで動的スライスを使用する理由

複雑なスケジューリング要件がある場合や、動的スライスを既存のスケジューリング インフラストラクチャと統合する場合は、独自のスケジューラを使用して Slice カスタム リソースを管理します。

Slice カスタム リソースを直接管理するのではなく、スケジューラを使用する場合は、GKE で Kueue とトポロジ対応スケジューリング(TAS)との統合が提供されます。詳細については、 Kueue と TAS で動的スライスをスケジュールするをご覧ください。

ワークフローの概要

カスタム スケジューラで動的スライスを使用するには、このドキュメントで次のタスクを行います。

  1. スライス コントローラを有効にする
  2. 増分プロビジョニングでノードプールを作成する
  3. Slice カスタム リソースをワークロード要件 に基づいて作成する。Slice カスタム リソースをクラスタに適用する。
  4. パーティションの状態とスライスの健全性をモニタリングする
  5. 完了したらスライスを削除する。

Slice カスタム リソースのフィールドとステータスの詳細については、 Slice カスタム リソース のリファレンス情報をご覧ください。

要件

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

  • Rapid チャンネルでバージョン 1.35.2-gke.1842000 以降の Standard クラスタを使用する。
  • 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 では、このドキュメントのコマンドを実行できない場合があります。
  • Rapid チャンネルでバージョン 1.35.2-gke.1842000 以降の既存の Standard クラスタがあることを確認します。新しいクラスタを作成するには、リージョン クラスタの作成をご覧ください。
  • リージョンに Ironwood(TPU7x)の十分な割り当てがあることを確認します。
  • マルチスライス ワークロードを実行する場合は、JobSet v0.10.1 以降をインストールします。
  • All Capacity モードで TPU 容量をリクエストする

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

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

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

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

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

    • CLUSTER_NAME: クラスタの名前。
    • LOCATION: 使用可能な TPU 容量があるリージョン
  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 をすべて検出できない場合でも、これらのノードプールをプロビジョニングします。

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

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

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

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 に設定して、1 つのサブブロック内のチップの合計数と一致させます。
  • 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. サブブロックの名前を特定するには、 ブロックのすべてのサブブロックを一覧表示しname: の各エントリの値を reservationSubBlocks の下にコピーします。この値は、このドキュメントのサブブロック名または 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 リンクをすべて有効にするには、次のいずれかの方法で動的スライスを作成します。

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

動的スライスを形成する

ノードプールを作成したら、Slice カスタム リソースを作成して、より大きな動的スライスを形成できます。Pod の代わりに、Slice カスタム リソースを使用して、指定されたトポロジを定義します。このトポロジはスライス コントローラによって有効になります。

ノードとパーティションのステータスを確認する

  1. ノードプールからノード名を取得するには、次のコマンドを実行します。

    kubectl get nodes -l cloud.google.com/gke-nodepool=${NODE_POOL_NAME}
    

    結果は次のようになります。

    NAME                                 STATUS   ROLES    AGE    VERSION
    gke-np-status-update-7b4c890c-0jhp   Ready    <none>   2d1h   v1.35.1-gke.1396002
    gke-np-status-update-7b4c890c-377r   Ready    <none>   2d1h   v1.35.1-gke.1396002
    gke-np-status-update-7b4c890c-gb51   Ready    <none>   2d1h   v1.35.1-gke.1396002
    
  2. ノードのプロビジョニング モデルを確認します。

    kubectl describe node NODE_NAME | grep "cloud.google.com/gke-accelerator-topology-mode"
    

    結果は次のようになります。

    cloud.google.com/gke-accelerator-topology-mode: PROVISION_ONLY
    

    この値は、ワークロード ポリシーの作成時に定義した accelerator-topology-mode=provision_only 設定と一致します。

  3. ノードラベル情報を取得します。

    kubectl describe node NODE_NAME | grep "cloud.google.com/gke-tpu-partition-4x4x4-id"
    

    NODE_NAME は、ノードプール内のノードのいずれかの名前に置き換えます。

    結果は次のようになります。

    cloud.google.com/gke-tpu-partition-4x4x4-id=fba785f80d18552357dcdef6d3d16c27
    

    cloud.google.com/gke-tpu-partition-4x4x4-state アノテーションは、ノードが動的スライスの形成に使用できるかどうかを示します。このラベルは次の値をサポートしています。

    • HEALTHY: ノードは正常で、完全に機能しています。
    • DEGRADED: ノードは破損していますが、動的スライスの形成に使用できます。
    • UNHEALTHY: ノードが誤動作しているため、スライスの形成に使用できません。
    • UNSET: ノードプール内のノードが不足しているため、状態が定義されていません。
    • INCOMPLETE: パーティション内のすべてのノードがプロビジョニングされていません。
  4. ノードに node.gke.io/created-by-mig アノテーションが含まれていることを確認します。

    kubectl describe node NODE_NAME | grep "node.gke.io/created-by-mig"
    

    NODE_NAME は、ノードプール内のノードのいずれかの名前に置き換えます。

    結果は次のようになります。

    node.gke.io/created-by-mig: projects/735972712744/zones/us-central1-ai1a/team/string
    

    出力には node.gke.io/created-by-mig ラベルが含まれています。これにより、GKE コントロール プレーンは Kubernetes ノードを基盤となる Compute Engine リソースにリンクできます。

Slice カスタム リソースを作成する

  1. Slice カスタム リソースを定義します。

    apiVersion: accelerator.gke.io/v1beta1
    kind: Slice
    metadata:
      # Name of the slice resource
      name: SLICE_NAME
    spec:
      # Specify the type of accelerator for this slice
      type: "tpu7x"
      # Define the desired topology for the accelerator slice
      topology: TOPOLOGY
      partitionIds:
        - PARTITION_ID # Example: a9476d1b02bd4f4e75ffffae3bd23c01
        - PARTITION_ID_2
        # ... add more partition IDs as needed
    

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

    • SLICE_NAME: スライスの名前。名前は metadata.name の条件を満たしている必要があります。
    • TOPOLOGY: 動的スライスのトポロジ。トポロジは次の条件を満たしている必要があります。
      • リクエストされたトポロジの各ディメンションは 4 の倍数である必要があります(4A x 4B x 4C など)。
      • トポロジ ディメンション AxBxC の 3 つの値は、減少しない順序(A ≤ B ≤ C)にする必要があります。たとえば、4x4x8 は有効ですが、4x8x4 は無効です。この順序にすることで、スライスを確実に形成し、予期しない動作を回避できます。
      • トポロジ ディメンション A × B × C の 3 つの値の積は 9,216 を超えないようにしてください。
    • PARTITION_ID: スライスを構成する 4x4x4 パーティションを識別する文字列のリスト。 パーティションの数は、チップの合計数に基づいて計算します。各パーティションは 64 個のチップで構成されます。spec.partitionIds リスト内のアイテム数は、計算されたパーティション数((A × B × C) / 64)と完全に一致する必要があります。 partitionIds は次の条件を満たしている必要があります。
      • 各パーティションは予約サブブロックにマッピングする必要があります。
      • 関連するサブブロックはすべて同じ予約に属している必要があります。
      • 関連するブロックはすべて同じ予約内に存在する必要があります。
      • 関連するノードプールのすべてのノードが ready 状態である必要があります。
    • type フィールドの値は tpu7x である必要があります。
    • 必要に応じて、スライス コントローラがスライス 形成中に自動的に再試行できるようにするには、slice.gke.io/retry-on-failure: "true" アノテーション をスライス カスタム リソースに追加します。SliceCreationFailed ステータスが原因でスライスが作成されない場合、スライスが正常に形成されるまでコントローラが再試行します。

    たとえば、4x8x8 スライスを作成するには、4 つの一意のパーティション ID を指定する必要があります。

    apiVersion: accelerator.gke.io/v1beta1
    kind: Slice
    metadata:
        name: test-slice-example
        annotations:
          slice.gke.io/retry-on-failure: "true" # Optional annotation to retry slice formation
    spec:
        type: "tpu7x"
        topology: "4x8x8" # (4*8*8)/64 = 4 partitions
        partitionIds:
            - "p0"
            - "p1"
            - "p2"
            - "p3"
    
  2. Slice カスタム リソースを適用します。

    kubectl apply -f test-slice-example.yaml
    

    この時点で、GKE はスライスの作成を試みます。次のいずれかの問題が発生すると、スライスの作成に失敗し、Slice カスタム リソースのステータスが SliceCreationFailed または FAILED に更新されます。

    • カスタム リソースで選択したノードが存在しない場合、ステータスは SliceCreationFailed になります。
    • カスタム リソースのノードが別のスライスで使用されている場合、ステータスは SliceCreationFailed になります。
    • カスタム リソースのノードが同じ予約ブロックの一部でない場合、ステータスは FAILED になります。
    • ノードが同じ予約にない場合、ステータスは FAILED になります。
    • トポロジがパーティションの数と一致しない場合、ステータスは SliceCreationFailed になります。

    Slice カスタム リソースのステータスの詳細については、スライスのステータスをご覧ください。

Slice カスタム リソースのステータスをモニタリングする

Slice カスタム リソースのステータスを確認するには、次のコマンドを実行します。

kubectl describe slice SLICE_NAME

SLICE_NAME はスライスの名前に置き換えます。

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

Name:         test-slice
Namespace:
Labels:       <none>
Annotations:  <none>
API Version:  accelerator.gke.io/v1beta1
Kind:         Slice
Metadata:
  Creation Timestamp:  2026-01-11T23:45:15Z
  Finalizers:
    accelerator.gke.io/slice-finalizer
  Generation:        1
  Resource Version:  1768175347356335006
  UID:               d0b71e5c-be3f-4788-aead-930c7afec4f2
Spec:
  Partition Ids:
    2c79463990ff67c4e3c2648666bfedfa
    ba898ffcac0ad0946e8ff036d771ee53
    [more partition IDs]
  Topology:  8x16x16
  Type:      tpu7x
Status:
  Conditions:
    Last Transition Time:  2026-01-11T23:45:38Z
    Message:               ""
    
    Reason:                FAILED
    
    Status:                False
    Type:                  Ready
Events:

Slice カスタム リソースのステータスの reason フィールドは、スライスの現在の状態を示します。 Slice カスタム リソースのライフサイクルは次のようになります。

  • SliceNotCreated: コントローラが初期化とリソースのチェックを行います。
    • 前提条件が満たされていない場合、状態は SliceCreationFailed に遷移します。
    • 検証に合格すると、状態は ACTIVATING に遷移します。
  • ACTIVATING: GKE がスライスを形成しています。
    • 成功すると、状態は ACTIVE に遷移します。
    • サブブロックが劣化しているがスライスが使用可能な場合、状態は ACTIVE_DEGRADED に遷移します。
    • 形成に失敗すると、状態は FAILED に遷移します。
  • DEACTIVATING: Slice カスタム リソースが削除された場合、またはアクティブ状態または失敗状態のときに重大な障害が発生した場合、スライスの解体が開始されます。
  • INCOMPLETE: リソースが完全に削除される前の最後のステップ。

Slice カスタム リソースのステータスの詳細については、スライスのステータスをご覧ください。

動的スライスでワークロードを実行する

Slice カスタム リソースが ACTIVE 状態の場合、ワークロードを実行できます。 次のセクションでは、動的スライスを使用するワークロードの例を示します。ワークロードはジョブまたは JobSet として送信されます。

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

次の例は、単一のサブブロック スライスを使用するワークロードを示しています。

  1. 次のサンプル マニフェストを tpu-job-jax-v7x-64.yaml として保存します。

    apiVersion: v1
    kind: Service
    metadata:
    name: headless-svc
    spec:
    clusterIP: None
    selector:
        job-name: tpu-job-jax-v7x-64
    ---
    apiVersion: batch/v1
    kind: Job
    metadata:
    name: tpu-job-jax-v7x-64
    spec:
    backoffLimit: 0
    completions: 16
    parallelism: 16
    completionMode: Indexed
    template:
        metadata:
        annotations:
            cloud.google.com/gke-tpu-slice-topology: 4x4x4
        spec:
        nodeSelector:
            cloud.google.com/gke-tpu-topology: 4x4x4
            cloud.google.com/gke-tpu-accelerator: tpu7x
            cloud.google.com/gke-tpu-slice: test-slice
        subdomain: headless-svc
        restartPolicy: Never
        containers:
        - name: tpu-job-jax
            env:
            - name: TPU_ACCELERATOR_TYPE
              value: tpu7x-128
            image: python:3.12
            securityContext:
            privileged: false
            command:
            - bash
            - -c
            - |
            set -ex
            pip install -U --pre jax jaxlib libtpu requests -i https://us-python.pkg.dev/ml-oss-artifacts-published/jax/simple/ -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
            pip list
            python -c 'import jax; print("Total TPU devices (cores):", jax.device_count())'
            resources:
            requests:
                google.com/tpu: 4
            limits:
                google.com/tpu: 4
    

    このマニフェストの内容:

    • cloud.google.com/gke-tpu-slice-topologycloud.google.com/gke-tpu-topology は、動的スライスのトポロジを定義します。
    • env.value: tpu7x-128 は、TPU アクセラレータ タイプとスライス内のコアの合計数です。 コア数は、トポロジのディメンションにチップあたりのコア数を掛けて計算されます。 たとえば、4x4x4 トポロジの場合、計算は 4 × 4 × 4 × 2 = 128 です。ここで、2tpu7x(Ironwood(TPU7x))のチップあたりのコア数です。 したがって、TPU_ACCELERATOR_TYPEtpu7x-128 です。
  2. tpu-job-jax-v7x-64.yaml マニフェストを適用します。

    kubectl apply -f tpu-job-jax-v7x-64.yaml
    

例 2: JobSet を使用してマルチスライス ノードプールにワークロードをデプロイする

この例では、JobSet を使用してマルチスライス ノードプールにワークロードをデプロイする方法を示します。

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

    kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/v0.10.1/manifests.yaml
    
  2. 次のサンプル マニフェストを tpu-multislice-jax.yaml として保存します。

    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: tpu-multislice-jax
      annotations:
        alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-tpu-slice
    spec:
      failurePolicy:
        maxRestarts: 3
      replicatedJobs:
      - name: slice-job
        replicas: 2
        template:
          spec:
            parallelism: 16
            completions: 16
            backoffLimit: 0
            completionMode: Indexed
            template:
              metadata:
                annotations:
                  # The shape of the slice
                  cloud.google.com/gke-tpu-slice-topology: 4x4x4
              spec:
                hostNetwork: true
                dnsPolicy: ClusterFirstWithHostNet
                nodeSelector:
                  cloud.google.com/gke-tpu-topology: 4x4x4
                  cloud.google.com/gke-tpu-accelerator: tpu7x
                  # IMPORTANT: Do NOT put 'cloud.google.com/gke-tpu-slice' here manually.
                  # The exclusive-topology annotation handles the slice assignment automatically.
                containers:
                - name: jax-worker
                  image: python:3.12
                  securityContext:
                    privileged: true
                  ports:
                  - containerPort: 8471
                  command:
                  - bash
                  - -c
                  - |
                    set -ex
                    pip install -U --pre jax jaxlib libtpu requests -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                    # Verify JobSet injected the specific slice ID for this worker
                    echo "JobSet Index: $JOB_COMPLETION_INDEX"
                    python -c 'import jax; print("Total TPU devices:", jax.device_count())'
                  resources:
                    requests:
                      google.com/tpu: 4
                    limits:
                      google.com/tpu: 4
    
  3. tpu-multislice-jax.yaml マニフェストを適用します。

    kubectl apply -f tpu-multislice-jax.yaml
    

    このマニフェストの内容:

    • replicatedJobsreplicas: 2 フィールドは、JobSet が 2 つの別々のジョブを作成することを示します。各ジョブは 4x4x4 TPU スライスに対応します。
    • alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-tpu-slice アノテーションは、各ジョブが一意の TPU スライスに 割り当てられるようにします。
    • cloud.google.com/gke-tpu-slice-topology: 4x4x4 アノテーションは、各動的スライスのトポロジを定義します。
    • JobSet がスライス割り当てを処理するため、この例では TPU_ACCELERATOR_TYPE 環境変数は明示的に設定されていません。JAX コードは、割り当てられたスライス内の使用可能な TPU デバイスを自動的に検出します。

スライスを削除する

  1. スライスを削除します。

    kubectl patch slice $SLICE_NAME --type json \
      -p='[{"op": "remove", "path": "/metadata/finalizers"}]'
    
  2. スライスが削除されたことを確認します。

    kubectl get slices
    

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

スライス コントローラを無効にするには、クラスタから削除します。

  1. Slice カスタム リソースが空であることを確認します。

    kubectl get slice -A
    
  2. クラスタを更新してスライス コントローラを無効にします。

    gcloud container clusters update ${CLUSTER_NAME} \
        --location=${REGION} \
        --no-enable-slice-controller
    
  3. Slice カスタム リソースを削除します。

    kubectl delete crd slices.accelerator.gke.io
    
  4. Slice カスタム リソースが削除されたことを確認します。

    kubectl get crd | grep slices.accelerator.gke.io
    
  5. スライス コントローラによって追加されたラベルを削除します。次のラベルを削除する必要があります。

    • cloud.google.com/gke-tpu-slice
    • cloud.google.com/gke-tpu-topology
    1. 特定のノードから削除するには、ノード名を更新します。
    export NODE_NAME="gke-tpu-bdac9600-3bdg"
    kubectl label node $NODE_NAME cloud.google.com/gke-tpu-slice- cloud.google.com/gke-tpu-slice-topology-
    
    1. クラスタ内のすべてのノードからこれらのラベルを削除する場合は、次の操作を行います。
    kubectl label nodes --all cloud.google.com/gke-tpu-slice- cloud.google.com/gke-tpu-slice-topology-
    
    1. ノードラベルを確認し、空であることを確認します。
    export NODE_NAME="gke-tpu-bdac9600-3bdg"
    kubectl describe node $NODE_NAME | grep "cloud.google.com/gke-tpu-slice"
    

次のステップ