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

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

按照這些操作說明操作之前,請務必瞭解動態配量的概念。

為什麼要搭配使用動態配量和自訂排程器?

如有複雜的排程需求,或想將動態配量功能與現有排程基礎架構整合,請使用自己的排程器管理 Slice 自訂資源。

如果您偏好使用排程器,而不是直接管理 Slice 自訂資源,GKE 提供與 Kueue 和拓撲感知排程 (TAS) 的整合功能。詳情請參閱「使用 Kueue 和 TAS 排定動態配量時間」。

工作流程總覽

如要搭配使用動態配量和自訂排程器,請執行本文中的下列工作:

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

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

需求條件

如要在 GKE 使用動態配量,必須符合下列條件:

  • 在下列其中一個版本中,使用快速管道的標準叢集:
    • 如要設定動態超級切片 (拓撲等於或大於 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 版本可能不支援執行本文中的指令。

啟用切片控制器

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

  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 的特定子區塊,或特定 16 節點的 TPU VM 群組,可提供更精細的控制。

建立工作負載政策

如要使用 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 自訂資源,形成較大的動態超級切片或較小的動態子切片。配量自訂資源會定義指定的拓撲,配量控制器會啟動該拓撲。工作負載接著會在這個動態切片上排程及執行。

動態切片分割區

分區提供可用來為工作負載建立動態配量的拓撲。每個節點都會顯示一個分割區,其中包含所有可用拓撲,包括 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 的拓撲。如要設定動態超級切片,所要求拓撲的每個維度都必須是 4 的倍數,例如 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.gke.io/retry-on-failure: "true" 註解新增至切片自訂資源。如果因 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 自訂資源中的狀態原因更新為 SliceCreationFailed 或 FAILED:

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

    如要在狀態為 SliceCreationFailed 時自動重試建立 Slice,請按照「建立 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 裝置故障,所有經過該交換器的光纖連結都會無法使用,導致 Superpod 內的所有 Cube 運作效能降低。
  • DEACTIVATING:正在拆除切片。
  • FAILED:切片不再可執行工作負載。如果初始建立失敗,或運作中的配量發生重大軟硬體故障,就會進入這個狀態。
  • INCOMPLETE:可用的立方體不足,無法啟動超級切片形成。

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

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

當 Slice 自訂資源處於 ACTIVE 狀態時,您可以在該資源上執行工作負載。 以下各節列舉使用動態配量的工作負載。工作負載會以工作或 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 0.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"
    

後續步驟