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

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

按照這些操作說明進行之前,請務必瞭解動態切片的概念

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

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

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

工作流程總覽

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

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

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

需求條件

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

  • 在 Rapid 搶鮮版管道中使用 1.35.2-gke.1842000 以上版本的 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 版本可能不支援執行本文件中的指令。

啟用 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 自訂資源。

形成動態切片

建立節點集區後,您可以建立 Slice 自訂資源,組成更大的動態 Slice。您可以使用 Slice 自訂資源定義指定拓撲,然後由 Slice 控制器啟動,而非使用 Pod。

驗證節點和分割區的狀態

  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
    

    這個值與您建立工作負載政策時定義的 accelerator-topology-mode=provision_only 設定相符。

  3. 擷取節點標籤資訊:

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

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

    結果大致如下:

    cloud.google.com/gke-tpu-partition-4x4x4-id=fba785f80d18552357dcdef6d3d16c27
    

    cloud.google.com/gke-tpu-partition-4x4x4-state 註解會指出節點是否可用於形成動態切片。這個標籤支援下列值:

    • HEALTHY:節點健康狀態良好,且功能正常運作。
    • DEGRADED:節點受損,但仍可用於形成動態切片。
    • UNHEALTHY:節點故障,無法用於形成切片。
    • UNSET:節點集區中的節點數量不足,因此狀態未定義。
    • INCOMPLETE:分區內並非所有節點都已佈建。
  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 資源。

建立 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 條件。
    • TOPOLOGY:動態切片的拓撲。拓撲必須符合下列條件:
      • 所要求拓撲的每個維度都必須是四的倍數,例如 4A x 4B x 4C
      • 拓撲維度中的三個值 AxBxC 必須依遞增順序排列 (A ≤ B ≤ C)。舉例來說,4x4x8 是有效的,但 4x8x4 無效。這個順序有助於確保切片形成一致,避免發生非預期行為。
      • 拓撲維度中的三個值相乘後,A × B × C 不得超過 9,216。
    • PARTITION_ID:識別構成切片的 4x4x4 分區字串清單。根據晶片總數計算分區數量,每個分區包含 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-slice-example
        annotations:
          slice.gke.io/retry-on-failure: "true" # Optional annotation to retry slice formation
    spec:
        type: "tpu7x"
        topology: "4x8x8" # (4*8*8)/64 = 4 partitions
        partitionIds:
            - "p0"
            - "p1"
            - "p2"
            - "p3"
    
  2. 套用 Slice 自訂資源:

    kubectl apply -f test-slice-example.yaml
    

    此時,GKE 會嘗試建立切片。如果發生下列任一問題,切片建立作業就會失敗,且 Slice 自訂資源中的狀態原因會更新為 SliceCreationFailedFAILED

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

    如要進一步瞭解 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 欄位會指出 Slice 的目前狀態。Slice 自訂資源的生命週期如下:

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

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

在動態切片上執行工作負載

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

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

以下範例顯示使用單一子區塊切片的作業負載。

  1. 將下列範例資訊清單儲存為 tpu-job-jax-v7x-64.yaml

    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-topologycloud.google.com/gke-tpu-topology 會定義動態配量的拓撲。
    • env.value: tpu7x-128 是 TPU 加速器類型,也是切片中的核心總數。核心數量的計算方式為將拓撲維度乘以每個晶片的核心數量。舉例來說,如果是 4x4x4 拓撲,計算方式為 4 × 4 × 4 × 2 = 128,其中 2tpu7x (Ironwood (TPU7x)) 的每個晶片核心數。因此 TPU_ACCELERATOR_TYPEtpu7x-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/v0.10.1/manifests.yaml
    
  2. 將下列範例資訊清單儲存為 tpu-multislice-jax.yaml

    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
    

停用 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"
    

後續步驟