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 を見つけられない場合でも、これらのノードプールをプロビジョニングします。

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

  • All Capacity モードの予約で公開される特定の TPU ブロック。ブロック ターゲティングを使用すると、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 フィールドは常に 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. サブブロックの名前を特定するには、ブロックのすべてのサブブロックを一覧表示し、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 の割り当て管理機能がバイパスされます。

パーティションの健全性の選択を定義する

標準のノードの健全性と準備状況に加えて、GKE は cloud.google.com/gke-tpu-partition-[shape]-state ラベル（[shape] はパーティション ID の形状（2x2x1、2x2x2、2x2x4、2x4x4、4x4x4 など）と一致）を使用して、各パーティション形状の特定の状態を公開します。このラベルを使用すると、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 アクセラレータ トポロジの要件には、次のルールが含まれます。
      • 動的サブスライスの場合: 2x2x1、2x2x2、2x2x4、2x4x4 など、4x4x4 より小さいトポロジを指定できます。これらの小規模なトポロジには、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. スライス コントローラを構成する: --reuse-slices=true フラグを含めるように Kueue スライス コントローラのデプロイ引数を更新します。ワークロードが既存のスライスに最適にスケジュールされるように、--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. このラベルをフレームワーク Job から Kueue Workload オブジェクトに渡すように Kueue を構成するには、Kueue 構成の labelKeysToCopy リストにラベルを追加します。

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

次のステップ