Kueue と TAS を使用して動的スライスをスケジュールする

このドキュメントでは、Kueue とトポロジ対応スケジューリング(TAS)を使用して、Google Kubernetes Engine(GKE)で TPU ノードプールをプロビジョニングし、動的スライスをスケジュールする方法について説明します。

Slice カスタム リソースを直接操作して、動的スライスを使用することもできます。詳細については、カスタム スケジューラで動的スライスを使用するをご覧ください。

これらの手順に進む前に、動的スライスのコンセプトを理解しておいてください。

要件

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

  • 次のいずれかのバージョンの Rapid チャンネルの Standard クラスタを使用します。
    • 動的スーパースライシング構成(4x4x4 以上のトポロジ)の場合は、バージョン 1.35.2-gke.1842000 以降を使用します。
    • 動的サブスライシング構成(4x4x4 より小さいトポロジ)の場合は、バージョン 1.36.0-gke.3712000 以降を使用します。
  • Ironwood(TPU7x)バージョンを使用します。
  • ノードに Container-Optimized OS イメージを使用します。
  • 増分プロビジョニングを使用するには、All Capacity モードの予約を使用します。All Capacity モードは、TPU Cluster Director によって有効になる機能です。
  • 動的サブスライシングの場合は、ノードに保留中のメンテナンス イベントがあることを確認します。保留中のメンテナンス イベントがないかインスタンスをモニタリングする。ノードに、終了時刻が 2026 年 9 月 18 日から 2026 年 9 月 30 日までの保留中のメンテナンス イベントがある場合は、サブスライシングを使用する前に、それらのノードでホスト メンテナンス イベントを手動でトリガーする必要があります。

始める前に

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

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

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

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

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

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

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

  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
    

Kueue、JobSet、LWS をインストールする

Kueue、JobSet、LWS がすでにインストールされている場合は、このセクションをスキップできます。

Kueue をインストールする

Kueue のドキュメントの手順に沿って操作するか、次のコマンドを実行します。

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

KUEUE_VERSION は、トポロジの要件に基づいて、必要な Kueue バージョンに置き換えます。動的サブスライシングには、Kueue v0.18.2 以降を使用します。動的スーパースライシングには、Kueue v0.16.6 以降を使用します。

JobSet をインストールする

JobSet のドキュメントの手順に沿って操作するか、次のコマンドを実行します。

kubectl apply --server-side -f https://github.com/kubernetes-sigs/jobset/releases/download/JOBSET_VERSION/manifests.yaml

JOBSET_VERSION は、トポロジの要件に基づいて必要な JobSet バージョンに置き換えます。動的サブスライシングの場合は、JobSet v0.12.0 以降を使用します。動的スーパースライシングには、JobSet v0.11.1 以降を使用します。

LWS をインストールする

LeaderWorkerSet(LWS)は、動的サブスライシングにのみ必要です。

LWS のドキュメントの手順に沿って操作するか、次のコマンドを実行します。

kubectl apply --server-side -f https://github.com/kubernetes-sigs/lws/releases/download/LWS_VERSION/manifests.yaml

LWS_VERSION は、必要な LWS バージョンに置き換えます。LWS v0.8.0 以降を使用します。

増分プロビジョニングを使用してノードプールを作成する

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

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

  • TPU の特定のブロック。All Capacity モードの予約で公開されます。ブロック ターゲティングを使用すると、GKE は指定されたブロック内の使用可能なサブブロックにノードプールを作成できます。
  • TPU の特定のサブブロック、または Ironwood(TPU7x)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 フィールドは、1 つのサブブロック内のチップの総数と一致するように、常に 4x4x4 に設定します。
  • 増分プロビジョニング プロセスがトリガーされるように、accelerator-topology-mode フィールドを常に provision_only に設定します。provision_only フィールドが設定されている場合、ノードプールは ICI リンクまたは OCS リンクを形成せずに TPU ノードをプロビジョニングします。

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

すべての容量モードの予約内の特定のサブブロックまたはブロックをターゲットに設定できます。

  • ブロックをターゲットにする: 各ノードプールは、指定されたブロックの容量を使用します。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: 作成したワークロード ポリシーの名前。

この段階では、ノードは作成されますが、ノードの Inter-Chip Interconnect(ICI)リンクはまだアクティブではありません。そのため、これらのノードプールでワークロードを直接実行することはできません。

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

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

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

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

Kueue スライス コントローラをインストールする

  1. Kueue スライス コントローラをインストールするには、次のマニフェストを slice-controller.yaml として保存します。

    # Copyright 2026 Google LLC
    #
    # Licensed under the Apache License, Version 2.0 (the "License");
    # you may not use this file except in compliance with the License.
    # You may obtain a copy of the License at
    #
    #      http://www.apache.org/licenses/LICENSE-2.0
    #
    # Unless required by applicable law or agreed to in writing, software
    # distributed under the License is distributed on an "AS IS" BASIS,
    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    # See the License for the specific language governing permissions and
    # limitations under the License.
    
    apiVersion: v1
    kind: Namespace
    metadata:
      labels:
        app.kubernetes.io/managed-by: kustomize
        app.kubernetes.io/name: slice-controller
        control-plane: controller-manager
      name: slice-controller-system
    ---
    apiVersion: v1
    kind: ServiceAccount
    metadata:
      labels:
        app.kubernetes.io/managed-by: kustomize
        app.kubernetes.io/name: slice-controller
        control-plane: controller-manager
      name: slice-controller-controller-manager
      namespace: slice-controller-system
    ---
    apiVersion: rbac.authorization.k8s.io/v1
    kind: Role
    metadata:
      labels:
        app.kubernetes.io/managed-by: kustomize
        app.kubernetes.io/name: slice-controller
        control-plane: controller-manager
      name: slice-controller-leader-election-role
      namespace: slice-controller-system
    rules:
    - apiGroups:
      - ""
      resources:
      - configmaps
      verbs:
      - get
      - list
      - watch
      - create
      - update
      - patch
      - delete
    - apiGroups:
      - coordination.k8s.io
      resources:
      - leases
      verbs:
      - get
      - list
      - watch
      - create
      - update
      - patch
      - delete
    - apiGroups:
      - ""
      resources:
      - events
      verbs:
      - create
      - patch
    ---
    apiVersion: rbac.authorization.k8s.io/v1
    kind: Role
    metadata:
      labels:
        app.kubernetes.io/managed-by: kustomize
        app.kubernetes.io/name: slice-controller
        control-plane: controller-manager
      name: slice-controller-webhook-secrets-role
      namespace: slice-controller-system
    rules:
    - apiGroups:
      - ""
      resources:
      - secrets
      verbs:
      - get
      - list
      - watch
      - update
    ---
    apiVersion: rbac.authorization.k8s.io/v1
    kind: ClusterRole
    metadata:
      labels:
        control-plane: controller-manager
      name: slice-controller-manager-role
    rules:
    - apiGroups:
      - ""
      resources:
      - nodes
      - pods
      verbs:
      - get
      - list
      - watch
    - apiGroups:
      - accelerator.gke.io
      resources:
      - slices
      verbs:
      - create
      - delete
      - get
      - list
      - patch
      - update
      - watch
    - apiGroups:
      - accelerator.gke.io
      resources:
      - slices/finalizers
      verbs:
      - update
    - apiGroups:
      - accelerator.gke.io
      resources:
      - slices/status
      verbs:
      - get
      - patch
      - update
    - apiGroups:
      - admissionregistration.k8s.io
      resources:
      - mutatingwebhookconfigurations
      verbs:
      - get
      - list
      - update
      - watch
    - apiGroups:
      - batch
      resources:
      - jobs
      verbs:
      - get
      - list
      - patch
      - update
      - watch
    - apiGroups:
      - events.k8s.io
      resources:
      - events
      verbs:
      - create
      - patch
      - update
      - watch
    - apiGroups:
      - jobset.x-k8s.io
      resources:
      - jobsets
      verbs:
      - get
      - list
      - patch
      - update
      - watch
    - apiGroups:
      - kueue.x-k8s.io
      resources:
      - admissionchecks
      verbs:
      - get
      - list
      - watch
    - apiGroups:
      - kueue.x-k8s.io
      resources:
      - admissionchecks/status
      - workloads/status
      verbs:
      - get
      - patch
      - update
    - apiGroups:
      - kueue.x-k8s.io
      resources:
      - workloads
      verbs:
      - create
      - get
      - list
      - patch
      - update
      - watch
    - apiGroups:
      - leaderworkerset.x-k8s.io
      resources:
      - leaderworkersets
      verbs:
      - get
      - list
      - patch
      - update
      - watch
    - apiGroups:
      - ray.io
      resources:
      - rayclusters
      - rayjobs
      verbs:
      - get
      - list
      - patch
      - update
      - watch
    ---
    apiVersion: rbac.authorization.k8s.io/v1
    kind: ClusterRole
    metadata:
      labels:
        control-plane: controller-manager
      name: slice-controller-metrics-auth-role
    rules:
    - apiGroups:
      - authentication.k8s.io
      resources:
      - tokenreviews
      verbs:
      - create
    - apiGroups:
      - authorization.k8s.io
      resources:
      - subjectaccessreviews
      verbs:
      - create
    ---
    apiVersion: rbac.authorization.k8s.io/v1
    kind: ClusterRole
    metadata:
      labels:
        control-plane: controller-manager
      name: slice-controller-metrics-reader
    rules:
    - nonResourceURLs:
      - /metrics
      verbs:
      - get
    ---
    apiVersion: rbac.authorization.k8s.io/v1
    kind: RoleBinding
    metadata:
      labels:
        app.kubernetes.io/managed-by: kustomize
        app.kubernetes.io/name: slice-controller
        control-plane: controller-manager
      name: slice-controller-leader-election-rolebinding
      namespace: slice-controller-system
    roleRef:
      apiGroup: rbac.authorization.k8s.io
      kind: Role
      name: slice-controller-leader-election-role
    subjects:
    - kind: ServiceAccount
      name: slice-controller-controller-manager
      namespace: slice-controller-system
    ---
    apiVersion: rbac.authorization.k8s.io/v1
    kind: RoleBinding
    metadata:
      labels:
        app.kubernetes.io/managed-by: kustomize
        app.kubernetes.io/name: slice-controller
        control-plane: controller-manager
      name: slice-controller-webhook-secrets-rolebinding
      namespace: slice-controller-system
    roleRef:
      apiGroup: rbac.authorization.k8s.io
      kind: Role
      name: slice-controller-webhook-secrets-role
    subjects:
    - kind: ServiceAccount
      name: slice-controller-controller-manager
      namespace: slice-controller-system
    ---
    apiVersion: rbac.authorization.k8s.io/v1
    kind: ClusterRoleBinding
    metadata:
      labels:
        app.kubernetes.io/managed-by: kustomize
        app.kubernetes.io/name: slice-controller
        control-plane: controller-manager
      name: slice-controller-manager-rolebinding
    roleRef:
      apiGroup: rbac.authorization.k8s.io
      kind: ClusterRole
      name: slice-controller-manager-role
    subjects:
    - kind: ServiceAccount
      name: slice-controller-controller-manager
      namespace: slice-controller-system
    ---
    apiVersion: rbac.authorization.k8s.io/v1
    kind: ClusterRoleBinding
    metadata:
      labels:
        control-plane: controller-manager
      name: slice-controller-metrics-auth-rolebinding
    roleRef:
      apiGroup: rbac.authorization.k8s.io
      kind: ClusterRole
      name: slice-controller-metrics-auth-role
    subjects:
    - kind: ServiceAccount
      name: slice-controller-controller-manager
      namespace: slice-controller-system
    ---
    apiVersion: v1
    kind: Secret
    metadata:
      labels:
        control-plane: controller-manager
      name: slice-controller-webhook-server-cert
      namespace: slice-controller-system
    ---
    apiVersion: v1
    kind: Service
    metadata:
      labels:
        app.kubernetes.io/managed-by: kustomize
        app.kubernetes.io/name: slice-controller
        control-plane: controller-manager
      name: slice-controller-controller-manager-metrics-service
      namespace: slice-controller-system
    spec:
      ports:
      - name: https
        port: 8443
        protocol: TCP
        targetPort: 8443
      selector:
        app.kubernetes.io/name: slice-controller
        control-plane: controller-manager
    ---
    apiVersion: v1
    kind: Service
    metadata:
      labels:
        control-plane: controller-manager
      name: slice-controller-webhook-service
      namespace: slice-controller-system
    spec:
      ports:
      - port: 443
        protocol: TCP
        targetPort: 9443
      selector:
        control-plane: controller-manager
    ---
    apiVersion: apps/v1
    kind: Deployment
    metadata:
      labels:
        app.kubernetes.io/managed-by: kustomize
        app.kubernetes.io/name: slice-controller
        control-plane: controller-manager
      name: slice-controller-controller-manager
      namespace: slice-controller-system
    spec:
      replicas: 1
      selector:
        matchLabels:
          app.kubernetes.io/name: slice-controller
          control-plane: controller-manager
      template:
        metadata:
          annotations:
            kubectl.kubernetes.io/default-container: manager
          labels:
            app.kubernetes.io/name: slice-controller
            control-plane: controller-manager
        spec:
          containers:
          - args:
            - --metrics-bind-address=:8443
            - --leader-elect
            - --health-probe-bind-address=:8081
            - --zap-log-level=3
            - --feature-gates=UseRetryMechanismForSliceCreation=true
            - --activation-timeout=60m
            - --kube-api-qps=100
            - --kube-api-burst=200
            - --default-slice-health-node-affinity=HEALTHY_AND_DEGRADED
            - --reuse-slices=true
            - --add-preferred-node-affinity=true
            command:
            - /manager
            image: tpuongke/kueue-slice-controller:v0.10
            livenessProbe:
              httpGet:
                path: /healthz
                port: 8081
              initialDelaySeconds: 15
              periodSeconds: 20
            name: manager
            ports:
            - containerPort: 9443
              name: webhook-server
              protocol: TCP
            readinessProbe:
              httpGet:
                path: /readyz
                port: 8081
              initialDelaySeconds: 5
              periodSeconds: 10
            resources:
              limits:
                cpu: 12
                memory: 32Gi
              requests:
                cpu: 8
                memory: 16Gi
            securityContext:
              allowPrivilegeEscalation: false
              capabilities:
                drop:
                - ALL
            volumeMounts:
            - mountPath: /tmp/k8s-webhook-server/serving-certs
              name: cert
              readOnly: true
          securityContext:
            runAsNonRoot: true
            seccompProfile:
              type: RuntimeDefault
          serviceAccountName: slice-controller-controller-manager
          terminationGracePeriodSeconds: 10
          volumes:
          - name: cert
            secret:
              defaultMode: 420
              secretName: slice-controller-webhook-server-cert
    ---
    apiVersion: admissionregistration.k8s.io/v1
    kind: MutatingWebhookConfiguration
    metadata:
      labels:
        control-plane: controller-manager
      name: slice-controller-mutating-webhook-configuration
    webhooks:
    - admissionReviewVersions:
      - v1
      clientConfig:
        service:
          name: slice-controller-webhook-service
          namespace: slice-controller-system
          path: /mutate-batch-v1-job
      failurePolicy: Fail
      name: mjob.kb.io
      rules:
      - apiGroups:
        - batch
        apiVersions:
        - v1
        operations:
        - CREATE
        resources:
        - jobs
      sideEffects: None
    - admissionReviewVersions:
      - v1
      clientConfig:
        service:
          name: slice-controller-webhook-service
          namespace: slice-controller-system
          path: /mutate-jobset-x-k8s-io-v1alpha2-jobset
      failurePolicy: Fail
      name: mjobset.kb.io
      rules:
      - apiGroups:
        - jobset.x-k8s.io
        apiVersions:
        - v1alpha2
        operations:
        - CREATE
        resources:
        - jobsets
      sideEffects: None
    - admissionReviewVersions:
      - v1
      clientConfig:
        service:
          name: slice-controller-webhook-service
          namespace: slice-controller-system
          path: /mutate-leaderworkerset-x-k8s-io-v1-leaderworkerset
      failurePolicy: Fail
      name: mleaderworkerset.kb.io
      rules:
      - apiGroups:
        - leaderworkerset.x-k8s.io
        apiVersions:
        - v1
        operations:
        - CREATE
        resources:
        - leaderworkersets
      sideEffects: None
    - admissionReviewVersions:
      - v1
      clientConfig:
        service:
          name: slice-controller-webhook-service
          namespace: slice-controller-system
          path: /mutate-ray-io-v1-raycluster
      failurePolicy: Fail
      name: mraycluster.kb.io
      rules:
      - apiGroups:
        - ray.io
        apiVersions:
        - v1
        operations:
        - CREATE
        resources:
        - rayclusters
      sideEffects: None
    - admissionReviewVersions:
      - v1
      clientConfig:
        service:
          name: slice-controller-webhook-service
          namespace: slice-controller-system
          path: /mutate-ray-io-v1-rayjob
      failurePolicy: Fail
      name: mrayjob.kb.io
      rules:
      - apiGroups:
        - ray.io
        apiVersions:
        - v1
        operations:
        - CREATE
        resources:
        - rayjobs
      sideEffects: None
    
  2. slice-controller.yaml マニフェストを適用します。

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

    # Copyright 2026 Google LLC
    #
    # Licensed under the Apache License, Version 2.0 (the "License");
    # you may not use this file except in compliance with the License.
    # You may obtain a copy of the License at
    #
    #      http://www.apache.org/licenses/LICENSE-2.0
    #
    # Unless required by applicable law or agreed to in writing, software
    # distributed under the License is distributed on an "AS IS" BASIS,
    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    # See the License for the specific language governing permissions and
    # limitations under the License.
    
    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
    
  4. 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 ラベル: このレベルは、Pod を特定の VM に割り当て、そのラベルと 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 の割り当て管理機能をバイパスします。

パーティション ヘルス選択を定義する

標準のノードの健全性と readiness に加えて、GKE は cloud.google.com/gke-tpu-partition-[shape]-state ラベルを使用して各パーティション シェイプの特定の状態を公開します([shape] は、2x2x1、2x2x2、2x2x4、2x4x4、4x4x4 などのパーティション ID シェイプと一致します)。このラベルにより、GKE は TPU リンクの状態など、スライスの形成に影響する要因を考慮できます。動的サブスライシング構成(4x4x4 より小さいトポロジ)には、GKE バージョン 1.36.0-gke.3712000 以降が必要です。

パーティション状態ラベルの値は、次のように定義できます。

  • HEALTHY: パーティションは正常で、完全に機能しています。
  • DEGRADED: パーティションのインフラストラクチャが、OCS リンクの劣化などにより、劣化した状態になっています。パーティションは引き続きスライスを形成できますが、全体的なパフォーマンスは正常なパーティションよりも低くなる可能性があります。この状態は、最上位の 4x4x4 トポロジにのみ適用されます。小さなトポロジには、デグレード状態はありません。
  • UNHEALTHY: パーティションが正常ではなく、スライスを形成できません。
  • UNSET: GKE スライス コントローラの初期化が失敗したため、状態が未定義です。
  • INCOMPLETE: パーティション内のすべてのノードがプロビジョニングされているわけではありません。

Kueue Slice Controller Webhook は、ワークロードに特定のパーティションの健全性要件が含まれているかどうかを検証します。優先度が指定されていない場合、Webhook はデフォルトのノード アフィニティを挿入します。

動作は次のとおりです。

  • cloud.google.com/gke-tpu-partition-[shape]-state ラベルをターゲットとする nodeSelector または nodeAffinity が存在する場合、変更されません。
  • このようなラベル構成が存在しない場合、Webhook は次のデフォルトのノード アフィニティを挿入して、使用可能なパーティションのみが使用されるようにします。

    nodeAffinity:
      requiredDuringSchedulingIgnoredDuringExecution:
        nodeSelectorTerms:
        - matchExpressions:
          - key: cloud.google.com/gke-tpu-partition-4x4x4-state
            operator: In
            values:
            - "HEALTHY"
            - "DEGRADED"
    

次のセクションでは、cloud.google.com/gke-tpu-partition-4x4x4-state ラベルが、異なるサブブロックのヘルス構成を指定するように構成されている例を示します。

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

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

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

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

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

    # Copyright 2026 Google LLC
    #
    # Licensed under the Apache License, Version 2.0 (the "License");
    # you may not use this file except in compliance with the License.
    # You may obtain a copy of the License at
    #
    #      http://www.apache.org/licenses/LICENSE-2.0
    #
    # Unless required by applicable law or agreed to in writing, software
    # distributed under the License is distributed on an "AS IS" BASIS,
    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    # See the License for the specific language governing permissions and
    # limitations under the License.
    
    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: big-super-slice
      labels:
        kueue.x-k8s.io/queue-name: lq
    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
                  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 より小さいトポロジ(2x2x1、2x2x2、2x2x4、2x4x4 など)を指定できます。これらの小さなトポロジには、GKE バージョン 1.36.0-gke.3712000 以降が必要です。
      • 動的スーパースライシングの場合: 4x4x4 以上のトポロジを指定できます。動的スーパースライシング構成の場合、リクエストされたトポロジの各ディメンションは 4 の倍数(4A x 4B x 4C など)にする必要があります。
      • トポロジは、AxBxC 形式の 3 次元文字列(4x8x8 など)にする必要があります。
      • ディメンションは、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 に割り当てます。
    • Webhook は、デフォルトのノード アフィニティを挿入して、HEALTHY ノードと DEGRADED ノードが使用されるようにします。
  2. big-super-slice.yaml マニフェストを適用します。

    kubectl apply -f big-super-slice.yaml
    

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

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

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

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

    # Copyright 2026 Google LLC
    #
    # Licensed under the Apache License, Version 2.0 (the "License");
    # you may not use this file except in compliance with the License.
    # You may obtain a copy of the License at
    #
    #      http://www.apache.org/licenses/LICENSE-2.0
    #
    # Unless required by applicable law or agreed to in writing, software
    # distributed under the License is distributed on an "AS IS" BASIS,
    # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    # See the License for the specific language governing permissions and
    # limitations under the License.
    
    apiVersion: jobset.x-k8s.io/v1alpha2
    kind: JobSet
    metadata:
      name: two-super-slices
      labels:
        kueue.x-k8s.io/queue-name: lq
    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 動的スライスを作成し、すべてのスライスがアクティブになるまで待機してからワークロードを許可します。

スライスの再利用によりワークロードのスケジューリングを最適化する

デフォルトでは、Kueue はスライス コントローラに、ワークロードの実行ごとに新しい TPU スライスをプロビジョニングし、ワークロードの完了時に削除するように指示します。この動作により、基盤となるインフラストラクチャでチャーンとオーバーヘッドが発生する可能性があります。

スライスの再利用により、新しいワークロードは、同じトポロジとハードウェア制限を持つ既存のアイドル状態の TPU スライスを再利用できます。スライスの再利用を有効にすると、Kueue はジョブの終了時にスライスを削除するのではなく、アイドル状態としてマークします。同じパーティションをリクエストする新しいジョブは、スライスをすぐに要求して再利用するため、ハードウェアのプロビジョニングの遅延を回避できます。

スライスの再利用を有効にする

スライスの再利用を完全に最適化するには、スライス コントローラと Kueue の両方を連携するように構成する必要があります。

  1. スライス コントローラを構成する: Kueue スライス コントローラのデプロイ引数を更新して、--reuse-slices=true フラグを含めます。ワークロードが既存のスライスに最適にスケジュールされるようにするには、--add-preferred-node-affinity=true フラグも追加します。

  2. Kueue を構成する(推奨): Kueue がコントローラの優先ノード アフィニティの提案を使用できるようにするには、Kueue 構成で TASRespectNodeAffinityPreferred フィーチャー ゲートを有効にする必要があります。

    たとえば、kueue-manager-config ConfigMap を編集してフィーチャー ゲートを追加します。

    apiVersion: v1
    kind: ConfigMap
    metadata:
      name: kueue-manager-config
      namespace: kueue-system
    data:
      controller_manager_config.yaml: |
        apiVersion: config.kueue.x-k8s.io/v1beta1
        kind: Configuration
        featureGates:
          TASRespectNodeAffinityPreferred: true
        # ...
    

    ConfigMap を変更したら、kueue-controller-manager Deployment を再起動して変更を適用します。

    kubectl rollout restart deploy kueue-controller-manager -n kueue-system
    

ワークロードごとにスライスの削除を適用する

特定のワークロードでスライスをクリーンアップし、スライスの再利用を防ぐには、次の操作を行います。

  1. Job または JobSet の定義に次のラベルを適用します。

    kueue.gke.io/tpu-enforce-slice-deletion: "true"
    

    このラベルを適用した後も、ワークロードは既存のアイドル状態のスライスでスケジュールできます。ただし、ワークロードが完了すると、スライスは再利用可能なプールに戻されるのではなく、完全に削除されます。

  2. Kueue 構成の labelKeysToCopy リストにラベルを追加して、このラベルをフレームワーク Job から Kueue Workload オブジェクトに渡すように Kueue を構成します。

    labelKeysToCopy:
    - "kueue.gke.io/tpu-enforce-slice-deletion"
    

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

スライスのステータスを確認し、GKE システム指標でスライス指標をモニタリングできます。

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

ワークロードによってスライスが作成されるか、アクティブに使用されると、スライス コントローラはラベルを付加して、それを使用しているワークロードを識別します。スライス名は動的に生成されるため(再利用が有効になっている場合はワークロード名から切り離されるため)、これらのラベルを使用してワークロードに割り当てられたスライスを見つける必要があります。

  1. ワークロードに割り当てられたスライスを確認するには、次のコマンドを実行します。

    kubectl get slice -l kueue.gke.io/tpu-slice-used-by-namespace=NAMESPACE,kueue.gke.io/tpu-slice-used-by-workload=WORKLOAD_NAME
    

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

    • NAMESPACE: ワークロードの Namespace。
    • WORKLOAD_NAME: Kueue ワークロードの名前。通常は Job または JobSet の名前と同じです。
  2. 前のコマンドの出力からスライス名をコピーします。

  3. スライス名を使用して、動的スライスの詳細なステータスを確認します。

    kubectl describe slice SLICE_NAME
    

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

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

Name:         s-4x4x4-74d3fb0b
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>

スライスの指標をモニタリングする

スライスの状態を公開する次の GKE システム指標をモニタリングできます。

  • kubernetes.io/accelerator/slice/state
  • kubernetes.io/accelerator/partition/state
  • kubernetes.io/accelerator/slice/deformation_durations
  • kubernetes.io/accelerator/slice/formation_durations

指標の詳細については、GKE システム指標をご覧ください。

Kueue Slice コントローラの指標

Kueue Slice Controller によって公開される次の Prometheus 指標をモニタリングして、スライスのライフサイクル、再利用、アドミッション チェックを観察することもできます。

スライスのライフサイクル指標

スライスのライフサイクル指標は、スライスの作成と削除を追跡します。

  • tpu_ac_slice_create_time: スライスの作成にかかった時間を測定するヒストグラム。
  • tpu_ac_slices_created: 作成されたスライスを追跡するカウンタ(is_reusable ラベルを含む)。
  • tpu_ac_slices_status: is_reusable ラベルと in_use ラベルを含む、ステータス別の現在のスライス数を追跡するゲージ。
  • tpu_ac_slice_delete_time: スライスの削除にかかった時間を測定するヒストグラム。
  • tpu_ac_slices_deleted: 削除されたスライスを追跡するカウンタ。
  • tpu_ac_slices_gc: カウンタ トラッキングにより、スライスが自動的にクリーンアップされました。

スライスの再利用指標

スライスの再利用指標は、TPU スライスの再利用の成功と競合を追跡します。

  • tpu_ac_slices_reused: 正常に再利用されたスライスを追跡するカウンタ。
  • tpu_ac_slices_conflict: トポロジまたはシェイプの不一致により削除されたスライスを追跡するカウンタ。

ワークロードとアドミッションの指標

ワークロードとアドミッションの指標は、アドミッション チェックの期間と結果を追跡します。

  • tpu_ac_admission_time: アドミッション チェック プロセスの所要時間を測定するヒストグラム。注: これは、スライス コントローラのアドミッション チェックの時間を測定するものであり、Kueue ワークロードのアドミッション時間全体は含まれません。
  • tpu_ac_workload_processed: target_state ラベルを含む、アドミッション チェックの決定のカウンタ。

一般的な指標

一般的な指標は、運用上の問題とエラーを追跡します。

  • tpu_ac_problem: 運用上の問題やエラーを追跡するカウンタ。この指標には、特定のエラータイプを示す problem ラベルが含まれています。

クリーンアップ

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

  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 カスタム リソースのリファレンス情報をご覧ください。

次のステップ