使用自訂排程器進行動態切片

本文說明如何直接與 Slice 自訂資源互動，使用動態切片。您可以建立切片、監控分區狀態，以及驗證切片健康狀態。

按照這些操作說明操作前，請務必瞭解動態切片的概念。

為什麼要搭配自訂排程器使用動態切片？

如果您有複雜的排程需求，或是想將動態切片與現有的排程基礎架構整合，請使用自己的排程器管理 Slice 自訂資源。

如果您偏好使用排程器，而不是直接管理 Slice 自訂資源，GKE 提供 Kueue 和 Topology Aware Scheduling (TAS) 的整合功能。詳情請參閱「使用 Kueue 和 TAS 排程動態切片」。

工作流程總覽

如要搭配自訂排程器使用動態切片，請在本文件中執行下列工作：

  1. 啟用 Slice 控制器。
  2. 建立節點集區並逐步佈建。
  3. 根據工作負載需求建立 Slice 自訂資源。將 Slice 自訂資源套用至叢集。
  4. 監控分區狀態和切片健康狀態。
  5. 完成後，刪除切片。

如要進一步瞭解 Slice 自訂資源的欄位和狀態，請參閱「Slice 自訂資源」參考資訊。

需求條件

如要在 GKE 中使用動態切片，必須符合下列規定：

  • 在下列其中一個版本中，使用搶鮮版通道的 Standard 叢集：
    • 如要設定動態超級切片 (拓撲等於或大於 4x4x4)，請使用 1.35.2-gke.1842000 以上版本。
    • 如要設定動態子切片 (拓撲小於 4x4x4)，請使用 1.36.0-gke.3712000 以上版本。
  • 使用 Ironwood (TPU7x) 版本。
  • 為節點使用 Container-Optimized OS 映像檔。
  • 如要使用增量佈建，請使用「所有容量模式」預留項目。「所有」容量模式是由 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 版本可能不支援執行本文件中的指令。

啟用 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 控制器會啟動該拓撲。
  • 使用 Kueue 和 TAS 安排 GKE 工作負載時程。Kueue 會自動處理 Slice 自訂資源的建立和刪除作業。請避免手動修改 Kueue 建立的 Slice 自訂資源。

使用超級切片或子切片形成動態切片

建立節點集區後，您可以建立 Slice 自訂資源，組成較大的動態超級切片或較小的動態子切片。Slice 自訂資源會定義指定的拓撲，由 Slice 控制器啟動。工作負載接著會在這個動態切片上排定及執行。

動態切片分割區

分割區提供可用於形成工作負載動態切片的拓撲。分割區會顯示每個節點的所有可用拓撲，包括 2x2x1、2x2x2、2x2x4、2x4x4 和 4x4x4。拓撲小於 4x4x4 時，需要 GKE 1.36.0-gke.3712000 以上版本。

大於 4x4x4 的任何切片都沒有分割區標籤，因為這是透過關聯多個 4x4x4 分割區所建立。

增量佈建節點集區中的每個 TPU 節點，都必須具備所有指定的分區 ID 和分區狀態節點標籤。

驗證節點和分割區的狀態

  1. 如要從節點集區取得節點名稱，請執行下列指令：

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

    結果大致如下：

    NAME                                 STATUS   ROLES    AGE    VERSION
gke-np-status-update-7b4c890c-0jhp   Ready    <none>   2d1h   v1.35.1-gke.1396002
gke-np-status-update-7b4c890c-377r   Ready    <none>   2d1h   v1.35.1-gke.1396002
gke-np-status-update-7b4c890c-gb51   Ready    <none>   2d1h   v1.35.1-gke.1396002

  2. 驗證節點的佈建模式：

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

    結果大致如下：

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

  3. 擷取要指定拓撲分割區的節點標籤資訊：

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

    將 NODE_NAME 替換為節點集區中其中一個節點的名稱。

    結果大致如下：

    cloud.google.com/gke-tpu-partition-4x4x4-id=fba785f80d18552357dcdef6d3d16c27
cloud.google.com/gke-tpu-partition-2x4x4-id=e18372d627ac412cb24d5ea8ab8912c9
cloud.google.com/gke-tpu-partition-2x2x4-id=7fbd4e29dc1839217fae41a9dd8211b4
cloud.google.com/gke-tpu-partition-2x2x2-id=a9476d1b02bd4f4e75ffffae3bd23c01
cloud.google.com/gke-tpu-partition-2x2x1-id=0bcfe937d1bb3914a8bdcd94e9f73319

  4. 確認節點包含 node.gke.io/created-by-mig 註解：

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

    將 NODE_NAME 替換為節點集區中其中一個節點的名稱。

    結果大致如下：

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

    輸出內容包含 node.gke.io/created-by-mig 註解，可讓 GKE 控制層將 Kubernetes 節點連結至基礎 Compute Engine 資源。

  5. 擷取要驗證的拓撲分割區狀態的節點標籤資訊：

    kubectl describe node NODE_NAME | grep -E "cloud.google.com/gke-tpu-partition-.*-state"

    結果大致如下：

    cloud.google.com/gke-tpu-partition-4x4x4-state=HEALTHY
cloud.google.com/gke-tpu-partition-2x4x4-state=HEALTHY
cloud.google.com/gke-tpu-partition-2x2x4-state=HEALTHY
cloud.google.com/gke-tpu-partition-2x2x2-state=HEALTHY
cloud.google.com/gke-tpu-partition-2x2x1-state=HEALTHY

    cloud.google.com/gke-tpu-partition-[shape]-state 標籤 (其中 [shape] 對應於分區 ID 的拓撲) 會指出分區是否可用於形成動態切片。這項狀態標籤支援下列值：

    • HEALTHY：分割區運作正常。
    • DEGRADED：分割區受損，但仍可用於動態切片形成。這個狀態只適用於頂層 4x4x4 拓撲。 較小的拓撲不會處於降級狀態。
    • UNHEALTHY：分割區故障，無法用於形成切片。
    • UNSET：由於 GKE 切片控制器初始化失敗，因此狀態未定義。
    • INCOMPLETE：分區內並非所有節點都已佈建。

建立 Slice 自訂資源

視您要建立動態超級切片或動態子切片而定，Slice 自訂資源會略有不同。

  1. 定義 Slice 自訂資源：

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

    更改下列內容：

    • SLICE_NAME：切片名稱。名稱必須符合metadata.name條件，且不得超過 49 個字元。
    • TOPOLOGY：動態切片的拓撲。 拓撲必須符合下列條件：
      • 動態子切片：您可以指定小於 4x4x4 的拓撲，例如 2x2x1、2x2x2、2x2x4 或 2x4x4。這些較小的拓撲需要 GKE 1.36.0-gke.3712000 以上版本。
      • 動態超級切片：您可以指定大於或等於 4x4x4 的拓撲。如要設定動態超級切片，所要求拓撲的每個維度都必須是四的倍數，例如 4A x 4B x 4C。拓撲維度中的三個值 AxBxC 必須依遞增順序排列 (A ≤ B ≤ C)。舉例來說，4x4x8 是有效值，但 4x8x4 不是。 這個順序有助於確保切片形成一致，並避免發生非預期行為。拓撲維度中的三個值 (A × B × C) 相乘後不得超過 9,216。
    • PARTITION_ID：識別組成切片的分割區字串清單。
      • 動態子切片：您必須指定「一個」分割區 ID。
      • 動態超級切片：您必須根據晶片的總數計算分割區數量，每個分割區包含 64 個晶片。spec.partitionIds 清單中的項目數量必須與計算出的分區數量 ((A × B × C) / 64) 完全相符。
      • partitionIds 清單必須符合下列條件：
        • 每個分割區都必須對應至預訂子區塊。
        • 所有相關聯的子區塊都必須屬於同一個預訂。
        • 所有相關聯的子區塊都必須位於同一個預留項目中。
        • 相關聯的節點集區必須處於 ready 狀態。
    • type 欄位的值必須為 tpu7x。
    • 如要讓 Slice 控制器在 Slice 形成期間自動重試，可以選擇將 slice.gke.io/retry-on-failure: "true" 註解新增至 Slice 自訂資源。如果因為 SliceCreationFailed 狀態原因而未建立切片，控制器會重試，直到成功形成切片為止。

    舉例來說，如要建立 4x8x8 切片 (動態超級切片)，您必須提供四個不重複的分割區 ID。

    apiVersion: accelerator.gke.io/v1beta1
kind: Slice
metadata:
    name: test-super-slice-example
    annotations:
      slice.gke.io/retry-on-failure: "true"
spec:
    type: "tpu7x"
    topology: "4x8x8" # (4*8*8)/64 = 4 partitions
    partitionIds:
        - "p0-4x4x4"
        - "p1-4x4x4"
        - "p2-4x4x4"
        - "p3-4x4x4"

    舉例來說，如要建立 2x2x2 區塊 (動態子區塊)，您必須提供一個專屬的磁碟分割區 ID。

    apiVersion: accelerator.gke.io/v1beta1
kind: Slice
metadata:
  name: test-sub-slice-example
  annotations:
      slice.gke.io/retry-on-failure: "true"
spec:
  type: "tpu7x"
  topology: "2x2x2"
  partitionIds:
    - "fba785f80d18552357dcdef6d3d16c27" # Only 1 partitionId for sub-slice

  2. 套用 Slice 自訂資源：

    kubectl apply -f test-slice-example.yaml

    此時，GKE 會嘗試建立切片。如果發生下列任一問題，系統會無法建立 Slice，並將 Slice 自訂資源中的狀態原因更新為 SliceCreationFailed 或 FAILED：

    • 如果自訂資源上選取的節點不存在，狀態原因會是 SliceCreationFailed。
    • 如果自訂資源上的任何節點遭其他切片使用，狀態原因會是 SliceCreationFailed。
    • 如果自訂資源上的節點不屬於同一個預訂子區塊，狀態原因為 FAILED。
    • 如果節點不在同一個預留項目中，狀態原因會是 FAILED。
    • 如果拓撲與分割區數量不符，狀態原因為 SliceCreationFailed。

    如要在狀態為 SliceCreationFailed 時自動重試建立切片，請按照「建立 Slice 自訂資源」一文所述，設定 slice.gke.io/retry-on-failure: "true" 註解。

    如要進一步瞭解 Slice 自訂資源的狀態，請參閱「Slice 狀態」。

監控 Slice 自訂資源的狀態

如要檢查 Slice 自訂資源的狀態，請執行下列指令：

kubectl describe slice SLICE_NAME

將 SLICE_NAME 替換為切片名稱。

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

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

Slice 自訂資源狀態中的 reason 欄位會指出目前的生命週期狀態。可能狀態會因您使用動態子切片或動態超級切片而異。

動態子切片狀態條件

  • SliceNotCreated：控制器會執行初始化作業和資源檢查。
    • 如未符合必要條件，狀態會轉換為 SliceCreationFailed。
    • 如果驗證通過，狀態會轉換為 ACTIVATING。
  • ACTIVATING：GKE 正在形成切片。
    • 如果成功，狀態會轉換為 ACTIVE。
    • 如果子區塊效能降低，但切片仍可使用，狀態會轉換為 ACTIVE_DEGRADED。
    • 如果編組失敗，狀態會轉換為 FAILED。
  • DEACTIVATING：如果刪除 Slice 自訂資源，或在有效或失敗狀態下發生重大錯誤，系統就會開始拆除 Slice。
  • INCOMPLETE：資源完全刪除前的最後一個步驟。

動態超級切片狀態條件

  • SliceNotCreated：尚未建立切片。切片控制器正在初始化，並執行切片形成前的預檢。
  • SliceCreationFailed：建立作業失敗，因為不符合必要條件 (例如缺少必要的 Compute Engine 資源)，或通過前置檢查。如要在這個狀態下自動重試切片形成作業，請設定 slice.gke.io/retry-on-failure: "true" 註解。
  • ACTIVATING：正在形成 (縫合) 切片。
  • ACTIVE：切片已完全形成、健康狀態良好，可執行工作負載。
  • ACTIVE_DEGRADED：切片已形成，但包含降級的立方體，由 ICI 韌性支援。工作負載可以執行，但效能可能會受到影響。動態超級切片可抵禦單一光學電路交換器 (OCS) 故障。如果單一 OCS 裝置故障，所有經過該交換器的光纖連結都會無法使用，導致超級 Pod 中的所有 Cube 運作效能降低。
  • DEACTIVATING：正在拆除切片。
  • FAILED：切片不再可執行工作負載。如果初始形成失敗，或作用中的切片發生重大軟體或硬體故障，就會出現這種狀態。
  • INCOMPLETE：可用的方塊不足，無法啟動超級切片形成程序。

如要進一步瞭解 Slice 自訂資源的狀態，請參閱「Slice 狀態」。

在動態配量上執行工作負載

當 Slice 自訂資源處於 ACTIVE 狀態時，您可以在該資源上執行工作負載。下一節列出使用動態切片的範例工作負載。工作負載會以 Job 或 JobSet 的形式提交。

範例 1：單一工作負載使用單一配量

以下範例顯示使用單一 4x4x4 動態超級切片的作業負載。

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

    在這個資訊清單中：

    • cloud.google.com/gke-tpu-slice-topology 和 cloud.google.com/gke-tpu-topology 定義動態切片的拓撲。
    • env.value: tpu7x-128 是 TPU 加速器類型，也是切片中的核心總數。核心數量的計算方式為將拓撲維度乘以每個晶片的核心數量。舉例來說，如果是 4x4x4 拓撲，計算方式為 4 × 4 × 4 × 2 = 128，其中 2 是 tpu7x (Ironwood (TPU7x)) 的每個晶片核心數。因此 TPU_ACCELERATOR_TYPE 為 tpu7x-128。

  2. 套用 tpu-job-jax-v7x-64.yaml 資訊清單：

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

範例 2：使用 JobSet 在多切片節點集區部署工作負載

這個範例說明如何使用 JobSet，在多切片節點集區中部署工作負載。

  1. 安裝 JobSet：

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

    將 JOBSET_VERSION 替換為 JobSet 的必要版本。如要使用動態子切片，請使用 JobSet v0.12.0 以上版本。如要使用動態超級切片，請使用 JobSet v0.11.1 以上版本。

  2. 將下列範例資訊清單儲存為 tpu-multislice-jax.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: tpu-multislice-jax
  annotations:
    alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-tpu-slice
spec:
  failurePolicy:
    maxRestarts: 3
  replicatedJobs:
  - name: slice-job
    replicas: 2
    template:
      spec:
        parallelism: 16
        completions: 16
        backoffLimit: 0
        completionMode: Indexed
        template:
          metadata:
            annotations:
              # The shape of the slice
              cloud.google.com/gke-tpu-slice-topology: 4x4x4
          spec:
            hostNetwork: true
            dnsPolicy: ClusterFirstWithHostNet
            nodeSelector:
              cloud.google.com/gke-tpu-topology: 4x4x4
              cloud.google.com/gke-tpu-accelerator: tpu7x
              # IMPORTANT: Do NOT put 'cloud.google.com/gke-tpu-slice' here manually.
              # The exclusive-topology annotation handles the slice assignment automatically.
            containers:
            - name: jax-worker
              image: python:3.12
              securityContext:
                privileged: true
              ports:
              - containerPort: 8471
              command:
              - bash
              - -c
              - |
                set -ex
                pip install -U --pre jax jaxlib libtpu requests -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
                # Verify JobSet injected the specific slice ID for this worker
                echo "JobSet Index: $JOB_COMPLETION_INDEX"
                python -c 'import jax; print("Total TPU devices:", jax.device_count())'
              resources:
                requests:
                  google.com/tpu: 4
                limits:
                  google.com/tpu: 4

  3. 套用 tpu-multislice-jax.yaml 資訊清單：

    kubectl apply -f tpu-multislice-jax.yaml

    在這個資訊清單中：

    • replicatedJobs 下方的 replicas: 2 欄位表示 JobSet 會建立兩個不同的 Job，分別對應至一個 4x4x4 TPU 節點。
    • alpha.jobset.sigs.k8s.io/exclusive-topology: cloud.google.com/gke-tpu-slice 註解可確保每個 Job 都會指派給專屬的 TPU 區塊。
    • cloud.google.com/gke-tpu-slice-topology: 4x4x4 註解會定義每個動態配量的拓撲。
    • 本範例未明確設定 TPU_ACCELERATOR_TYPE 環境變數，因為 JobSet 會處理切片指派作業。JAX 程式碼會自動偵測指派的切片中可用的 TPU 裝置。

刪除切片

  1. 刪除切片：

    kubectl patch slice $SLICE_NAME --type json \
  -p='[{"op": "remove", "path": "/metadata/finalizers"}]'

  2. 確認切片已刪除：

    kubectl get slices

升級節點集區

如果升級以遞增佈建方式設定的節點集區，請務必使用特定節點數擴充參數，以免發生容量衝突。

如要設定及執行節點集區升級作業，請按照下列步驟操作：

  1. 如果使用動態超級切片，請先刪除有效切片，再執行手動維護：

    kubectl delete slice SLICE_NAME

    如果使用動態子切片，則不必先刪除活動切片。不過，升級期間仍處於有效狀態的子切片會失敗，並轉換為 FAILED 狀態。

  2. 更新節點集區的升級參數：

    gcloud container node-pools update NODE_POOL_NAME \
    --cluster=CLUSTER_NAME \
    --project=PROJECT_ID \
    --location=LOCATION \
    --max-surge-upgrade=0 \
    --max-unavailable-upgrade=16

    請務必將 --max-surge-upgrade 欄位設為 0 值，以免 GKE 在升級期間嘗試分配額外的 TPU。建議您將 --max-unavailable-upgrade=16 欄位設為同時升級完整的 16 節點子區塊。

  3. 升級節點集區：

    gcloud container clusters upgrade CLUSTER_NAME \
    --node-pool=NODE_POOL_NAME \
    --cluster-version=CLUSTER_VERSION \
    --project=PROJECT_ID \
    --location=LOCATION

硬體維護和故障

如果觸發客戶啟動的維護作業或發生硬體容錯移轉，只有目標節點會受到影響，而非整個節點集區。

如要進行動態超級切片，必須先刪除現有的動態切片，才能執行手動維護作業。如果您使用動態子切片，則不需要先刪除有效切片。如果動態子切片處於啟用狀態時發生維護或硬體故障，系統會按照下列方式處理復原作業：

  1. 自動重塑：當相關聯的主機開始維護時，GKE 會自動重塑作用中的動態切片。
  2. Slice 自訂資源失敗：Slice 自訂資源轉換為 FAILED 狀態。
  3. 排程器觀察：排程器會觀察 Slice 自訂資源的失敗狀態。
  4. 重組程序：排程器會自動嘗試在其他可用的健康節點上重新建立切片設定。

停用 Slice Controller

如要停用切片控制器，請從叢集中移除。

  1. 確認 Slice 自訂資源為空白：

    kubectl get slice -A

  2. 更新叢集以停用切片控制器：

    gcloud container clusters update ${CLUSTER_NAME} \
    --location=${REGION} \
    --no-enable-slice-controller

  3. 刪除 Slice 自訂資源：

    kubectl delete crd slices.accelerator.gke.io

  4. 確認 Slice 自訂資源已刪除：

    kubectl get crd | grep slices.accelerator.gke.io

  5. 移除切片控制器新增的標籤。需要移除這些標籤：

    • cloud.google.com/gke-tpu-slice
    • cloud.google.com/gke-tpu-topology
    1. 如要從特定節點移除，請更新節點名稱
    export NODE_NAME="gke-tpu-bdac9600-3bdg"
kubectl label node $NODE_NAME cloud.google.com/gke-tpu-slice- cloud.google.com/gke-tpu-slice-topology-
    1. 如要從叢集中的每個節點移除這些標籤：
    kubectl label nodes --all cloud.google.com/gke-tpu-slice- cloud.google.com/gke-tpu-slice-topology-
    1. 檢查節點標籤，確認標籤為空白：
    export NODE_NAME="gke-tpu-bdac9600-3bdg"
kubectl describe node $NODE_NAME | grep "cloud.google.com/gke-tpu-slice"

後續步驟