使用 GKE 動態切片設定 TPU

本文說明如何在 Google Kubernetes Engine (GKE) 中使用動態切片。動態切片功能可讓您將佈建的 TPU 子區塊設定為不同拓撲。這項功能可減少重新建立節點集區的需求,並在發生故障時自動復原,提高容錯能力,同時最佳化資源用量。

本文適用於 AI/ML 工程師和平台管理員,可協助他們最佳化 TPU 使用率、縮短佈建時間,以及提升大規模訓練和推論工作負載的容錯能力。

閱讀本文前,請先熟悉下列項目:

什麼是動態切片?

動態切片可讓您解除 TPU 佈建作業的耦合,彈性管理 Cloud TPU 容量。動態切片包含下列程序:

  1. 您會以較小的單位 (稱為「子區塊」) 佈建資源。子區塊是 Ironwood (TPU7x) 容量的基本邏輯建構單元。如果是 Ironwood (TPU7x),則代表 16 個節點的 TPU VM 群組,並採用 4x4x4 拓撲互連 TPU 晶片。在 TPU 全容量模式和動態切片環境中,節點集會直接對應至子區塊。
  2. 動態切片功能會將這些子區塊縫合在一起,形成較大的切片。

動態切片的優點

動態切片可協助您達成下列目標:

  • 縮短佈建時間:個別佈建子區塊可加快整體佈建速度,因為這樣可將單一失敗事件的影響降到最低。
  • 縮短復原時間:如果發生 TPU 晶片故障,最小的故障單位是子區塊。動態切片會隔離故障的子模塊,因此工作負載可以比重新佈建整個大型切片更快地在健康的子模塊上重新排程。
  • 重塑容量:如果工作負載需求多樣,您不需要刪除並重新建立節點集區來變更拓撲,這是在沒有動態切片的情況下必須執行的操作。您可以動態重新設定佈建的節點集區,以符合指定的形狀。

動態切片的重要元素

動態切片功能會導入下列重要概念:

  • 節點集區的增量佈建:動態切片使用增量佈建,這是節點集區的容錯佈建模式。這個模型會將所有 TPU 容量轉換為 TPU VM 的 16 節點群組節點集區。
  • 切片控制器:在 GKE 控制層中執行的 Kubernetes 自訂資源控制器,可管理動態切片。切片控制器會管理 Slice 自訂資源的生命週期,代表動態切片。切片控制器負責建立、持續監控及刪除切片。使用排程器時,排程器會導向建立及刪除 Slice 自訂資源。
  • 配量自訂資源:根據要求的 TPU 拓撲,動態將這些子區塊縫合在一起。這項程序會動態重新設定 OCS 網路,以連線至 TPU 節點集區,確保效能達到最佳狀態。您可以檢查 Slice 自訂資源的狀態欄位,瞭解動態切片形成進度或健康狀態。

需求條件

如要在 GKE 中使用動態切片,必須符合下列條件:

  • 使用快速通道中 1.35.0-gke.274500 以上版本的 Standard 叢集。
  • 使用 Ironwood (TPU7x) 版本。
  • 為節點使用 Container-Optimized OS 映像檔。
  • 如要使用增量佈建,請預留所有容量模式。所有容量模式都是由 TPU Cluster Director 啟用的功能。

事前準備

開始之前,請確認您已完成下列工作:

  • 啟用 Google Kubernetes Engine API。
  • 啟用 Google Kubernetes Engine API
  • 如要使用 Google Cloud CLI 執行這項工作,請安裝初始化 gcloud CLI。如果您先前已安裝 gcloud CLI,請執行 gcloud components update 指令,取得最新版本。較舊的 gcloud CLI 版本可能不支援執行本文件中的指令。

限制

  • 單一切片必須使用預訂項目下同一 TPU 區塊內的子區塊。如要在 TPU 區塊之間使用子區塊,請使用 TPU 多重切片
  • 動態切片不支援小於 4x4x4 的拓撲。

在 GKE 中使用 Kueue 進行動態切片

本節說明在 GKE 中使用動態切片的流程。

  1. 查看所有容量模式預訂的拓撲和健康狀態
  2. 在叢集中啟用 Slice 控制器
  3. 建立 TPU 節點集區
  4. 設定 Kueue,建立 Slice 自訂資源
  5. 使用 Kueue 在動態切片上執行工作負載
  6. 清除所用資源

啟用 Slice 控制器

如要使用動態切片,請在叢集中啟用切片控制器。

  1. 更新叢集:

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

    更改下列內容:

  2. 取得憑證,以便使用 kubectl 指令與叢集通訊:

    gcloud config set container/cluster CLUSTER_NAME
    gcloud container clusters get-credentials CLUSTER_NAME \
        --location=LOCATION
    
  3. 在下列指令的輸出內容中,確認 slices.accelerator.gke.io 值是否存在:

    kubectl get crd slices.accelerator.gke.io
    

    輸出結果會與下列內容相似:

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

建立節點集區並啟用增量佈建

本節說明如何使用增量佈建建立 TPU 節點集區。GKE 會將所有 TPU 容量轉換為 16 個節點的 TPU VM 群組或子區塊節點集區。即使 GKE 無法找到所有 16 個健康狀態良好的 VM,仍會佈建這些節點集區,方法是將節點放在主機的健康狀態良好部分,並在修復健康狀態不良的機器時,逐步佈建這些機器。

您可以將節點集區指定為下列任一項目:

  • 特定 TPU 區塊,在「所有容量」模式預留中公開。區塊目標設定可讓 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,以符合單一子區塊內的晶片總數。
  • 請務必將 accelerator-topology-mode 欄位設為 provision_only,確保系統觸發增量佈建程序。設定 provision_only 欄位後,節點集區會佈建 TPU 節點,但不形成 ICI 或 OCS 連結。

將節點集區指定為屬於區塊或子區塊

您可以在「所有容量」模式預訂中指定子區塊或區塊。

  • 指定區塊:每個節點集區都會使用指定區塊的容量。GKE 會將節點集區放在該區塊的可用子區塊中。您必須建立與要使用的區塊中子區塊數量相同的節點集區。
  • 指定子區塊:每個節點集區都會對應至特定可用子區塊。使用子區塊目標時,只要至少有一個 VM 正常運作,GKE 就會建立節點集區。增量佈建可確保所有節點都放置在指定的子區塊內。

封鎖

  1. 如要擷取預留區塊中的區塊名稱,以及區塊中可用子區塊的數量,請在「查看所有容量模式預留的拓撲和健康狀態」文件中完成下列步驟:

    1. 如要找出封鎖名稱,請列出所有預訂封鎖,然後複製 name: 欄位中的值。這個值是這個文件中的區塊或 BLOCK_NAME 名稱。

    2. 如要判斷要建立多少節點集區,請說明預訂區塊,並找出 reservationSubBlockCount 欄位中的值。這個值代表可用的子區塊數量。舉例來說,reservationSubBlockCount: 4 值表示該區塊有四個子區塊可用,因此您需要建立四個不同的節點集區。

  2. 設定預訂路徑:

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

    更改下列內容:

    • RESERVATION_NAME:TPU 預留項目的名稱。
    • BLOCK_NAME:區塊名稱。
  3. 為上一步驟中識別的每個子區塊建立節點集區。 舉例來說,如果計數為 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,請在「查看所有容量模式預訂的拓撲和健康狀態」文件中完成下列步驟:

    1. 如要找出區塊名稱,請列出所有預訂區塊,然後複製 name: 欄位中的值。這個值是文件中的區塊名稱或 BLOCK_NAME

    2. 如要找出子區塊的名稱,請列出區塊的所有子區塊,然後複製 reservationSubBlocks 底下每個項目的 name: 欄位值。這個值是本文件中子區塊或 SUBBLOCK_NAME 的名稱。

  2. 設定預訂路徑:

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

    更改下列內容:

    • RESERVATION_NAME:TPU 預留項目的名稱。
    • BLOCK_NAME:區塊名稱。
    • SUBBLOCK_NAME:子區塊的名稱。
  3. 建立節點集區:

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

    更改下列內容:

    • NODE_POOL_NAME:新節點集區的專屬名稱,例如 sub-block-pool-1
    • PROJECT_ID:您的 Google Cloud 專案 ID。
    • CLUSTER_NAME:GKE 叢集名稱。
    • ZONE:節點集區的可用區,例如 us-central2-b
    • WORKLOAD_POLICY_NAME:您建立的工作負載政策名稱。

此時節點已建立,但晶片間互連 (ICI) 連結尚未啟用。因此,您無法直接在這些節點集區上執行工作負載。

如要啟用所有必要的 ICI 連結來形成切片,並允許排定工作負載,請使用下列其中一種方法建立動態切片:

  • 建立 Slice 自訂資源。您可以使用 Slice 自訂資源定義指定拓撲,而非 Pod,Slice 控制器會啟動該拓撲。
  • 使用 KueueTAS 安排 GKE 工作負載時程。Kueue 會自動處理 Slice 自訂資源的建立和刪除作業。請避免手動修改 Kueue 建立的 Slice 自訂資源。

使用 Kueue 和 TAS 建立動態切片

在本節中,您將使用 Kueue 和 TAS 排程 GKE 工作負載。

安裝 JobSet 和 Kueue 資源,進行動態切片

  1. 安裝 JobSet:

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

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

    kubectl apply -f https://gist.githubusercontent.com/mwysokin/cd90010d0d375b3bf57c536905692547/raw/506c36dd070f4ac222ba8a5e58ba28bbfcfa8ed3/kueue-slice-controller-v0.8.0-130.yaml
    
  4. 如要設定 Kueue 進行動態切片,請將下列資訊清單儲存為 dynamic-slice-topology.yaml

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

    kubectl apply -f dynamic-slice-topology.yaml
    

    在這個資訊清單中,您將定義下列資源,為動態切片設定 Kueue:

    • Ironwood (TPU7x) 動態切片拓撲 (superslice-topology):拓撲會定義 Kueue 排程動態切片工作負載時考量的層級。這些等級如下:
      • cloud.google.com/gce-topology-block 標籤:這個層級是必要項目,因為只有來自相同區塊的子區塊才能形成切片,因此需要瞭解哪些子區塊屬於哪些區塊。
      • cloud.google.com/gke-tpu-partition-4x4x4-id 標籤:這個層級代表個別的 Ironwood (TPU7x) 子區塊 (4x4x4 拓撲)。
      • kubernetes.io/hostname 標籤:您必須具備這個層級,才能將 Pod 指派給特定 VM,並觀察其標籤和汙點。
    • Ironwood (TPU7x) SuperSlice ResourceFlavor (superslice-rf):Ironwood (TPU7x) 子區塊的資源風味包含 cloud.google.com/gke-tpu-accelerator: tpu7x 標籤,可與 Ironwood (TPU7x) 機器的節點相符。
    • SuperSlice AdmissionCheck (superslice-ac):這項許可檢查會告知 Kueue,在 GKE 切片控制器確認切片已啟用前,請勿排定工作負載。首先定義准入檢查,然後新增至處理動態切片工作負載的 ClusterQueue
    • ClusterQueue (cq) 和 LocalQueue (lq):這些欄位會管理google.com/tpu資源。cq ClusterQueue 包含 superslice-ac 准入檢查。google.com/tpunominalQuota 可以透過下列兩種方式設定:
      • 特定配額:設定 nominalQuota 以符合現有容量, 確保公平分享和配額管理。
      • 無限制配額:將 nominalQuota 設為極高的值 (例如 "999999"),模擬無限制配額。為著重說明 TAS 和動態切片,這項設定會略過 Kueue 的配額管理功能。

定義子區塊健康狀態選取範圍

除了標準節點健康狀態和就緒狀態之外,GKE 還會使用 cloud.google.com/gke-tpu-partition-4x4x4-state 標籤,顯示每個子區塊的特定狀態。這個標籤可讓 GKE 考量影響切片形成的因素,例如 TPU 連結的狀態。

您可以按照下列方式定義 cloud.google.com/gke-tpu-partition-4x4x4-state 標籤的值:

  • HEALTHY:基礎架構健康狀態良好。
  • DEGRADED:子區塊的基礎架構處於降級狀態,例如 OCS 連結降級。子區塊仍可形成切片,但整體效能可能不如健全的子區塊。如果可以容忍效能可能降低,您可以設定工作負載,使用節點親和性來使用 DEGRADED 子區塊,如範例 3 所示。
  • UNHEALTHY:子區塊不健全,無法形成切片。

Kueue Slice Controller 網路鉤子會驗證工作負載是否包含特定子區塊健康狀態需求。如未指明偏好設定,Webhook 會插入預設節點親和性。

行為如下:

  • 如果存在以 cloud.google.com/gke-tpu-partition-4x4x4-state 標籤為目標的 nodeSelectornodeAffinity,則不會變更。
  • 如果沒有這類標籤設定,Webhook 會插入下列預設節點親和性,確保只使用可用的子區塊:

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

下一節的範例說明如何設定 cloud.google.com/gke-tpu-partition-4x4x4-state 標籤,指定不同的子區塊健康狀態設定。

使用 Kueue 透過動態切片執行測試工作負載

本節說明如何使用 Kueue 和 TAS,在動態切片上部署工作負載。其中包含三個範例,說明如何建立動態切片工作負載,以及由多個切片組成的工作負載。工作負載會以 JobSet 形式提交。

範例 1:單一工作負載使用單一動態切片

以下範例說明如何使用具有 4x12x16 拓撲的切片建立工作負載,該拓撲由 12 個子區塊組成。Pod 數量的計算方式如下:(4 * 12 * 16) / 每個節點 4 個晶片 = 192 個 Pod。

  1. 將下列資訊清單儲存為 big-super-slice.yaml

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

    在這個資訊清單中,下列註解會告知 Kueue 切片特徵和拓撲,以設定下列項目:

    • cloud.google.com/gke-tpu-slice-topology:將 "4x12x16" 指定為動態切片拓撲。tpu7x加速器拓撲的相關規定 包括下列規則:
      • 拓撲下限為 4x4x4
      • 拓撲必須是三維字串,格式為 AxBxC。例如:4x8x8
      • 每個維度 (A、B 和 C) 都必須是 4 的倍數。
      • 尺寸必須以非遞減順序排序:A <= B <= C。舉例來說,4x8x4 無效,應為 4x4x8
      • 尺寸 (ABC) 的乘積不得超過 9,216。
      • 支援的最大切片拓撲最多可包含 32 個子區塊。舉例來說,8x16x16 有 32 個子區塊、8x12x20 有 30 個子區塊,或 12x12x12 有 27 個子區塊,都符合可接受的限制。
    • cloud.google.com/gke-tpu-accelerator: tpu7x:在執行 Ironwood (TPU7x) 的 VM 上排定 Pod。
    • kueue.x-k8s.io/queue-name:將 JobSet 指派給 Kueue LocalQueue。
  2. 套用 big-super-slice.yaml 資訊清單:

    kubectl apply -f big-super-slice.yaml
    

    套用資訊清單後,Kueue 會建立名為 big-super-sliceJobSet。 Kueue 接著會嘗試使用 4x12x16 拓撲形成單一動態切片。切片啟用後,Kueue 會允許工作負載,並在節點上排定 192 個 Pod,形成動態切片來執行工作負載。

範例 2:工作負載有多個副本

以下範例說明如何建立工作負載,其中使用兩個動態切片,每個切片由四個子區塊組成。

  1. 將下列資訊清單儲存為 two-super-slices.yaml

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

    kubectl apply -f two-super-slices.yaml
    

在這個資訊清單中,您會在 replicatedJobs 欄位中設定 replicas: 2。 套用資訊清單後,Kueue 會嘗試以 4x8x8 拓撲形成兩個獨立的切片。Kueue 會為 jobset.spec.replicatedJobs[].replicas 中定義的每個副本建立動態切片。如果指定 n 副本,Kueue 會為工作負載建立 n 個動態切片,並等待所有切片都處於啟用狀態,再允許工作負載。

範例 3:具有單一動態切片和 NodeAffinity 的工作負載

自 Kueue 0.15 起,Kueue 支援 NodeAffinity,可選取 TAS 節點。這項功能可讓 HEALTHYDEGRADED 節點成為動態切片的一部分。以下範例說明如何設定具有單一動態切片和 NodeAffinity 的工作負載:

  1. 將下列資訊清單儲存為 slice-8x8x8-na.yaml

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

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

    套用資訊清單後,Kueue 會建立名為 slice-8x8x8-naJobSet。 Kueue 接著會嘗試形成具有 8x8x8 拓撲的單一動態切片,由於指定了 NodeAffinity,因此可同時納入 HEALTHYDEGRADED 節點。切片啟用後,Kueue 會允許工作負載,並在構成動態切片的節點上排定 128 個 Pod 的時程。

監控切片狀態

如要檢查動態切片的狀態,請執行下列指令:

kubectl describe slice SLICE_NAME

SLICE_NAME 替換為切片名稱。切片名稱通常衍生自 JobSet 名稱和副本索引。以「範例 1」來說,Kueue 建立的切片名稱會類似於 default-jobset-big-super-slice-yyyyy-job-jax-0

輸出結果會與下列內容相似:

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

為確保與基礎 Compute Engine 資源命名慣例相容,切片名稱須遵守下列規則:

  • 範本: {namespace}-jobset-{jobset.metadata.name}-kueueHash[5-character]-{jobset.spec.replicatedJobs[].name}-sliceIndex
  • 長度:名稱不得超過 54 個字元。控制器會附加連字號和 8 個半形字元的叢集雜湊值,建立 Compute Engine 資源名稱,長度上限為 63 個半形字元。
  • 格式:名稱符合規則運算式 ^[a-z]([-a-z0-9]*[a-z0-9])?$。名稱具有下列特徵:
    • 開頭為小寫英文字母。
    • 只能包含小寫英文字母、數字和連字號 (-)。
    • 結尾為小寫英文字母或數字 (不得為連字號)。

清除所用資源

為避免產生非預期的費用,請先刪除節點集區的切片,再刪除節點集區。

  1. 刪除 JobSet。這項操作會觸發 Kueue 刪除相關聯的 Slice 自訂資源。

    kubectl delete jobset JOBSET_NAME
    

    JOBSET_NAME 替換為 JobSet 的名稱,例如 big-super-slice

  2. 刪除 TPU 節點集區:

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

(選用) 使用自己的排程器進行動態切片

本文著重於使用 Kueue 和 TAS。不過,您也可以使用自己的自訂排程器管理動態切片。如要使用其他排程器,請參閱 Slice 自訂資源參考資訊。

後續步驟