在 GKE Autopilot 上部署 TPU 工作負載

本頁說明如何在 Google Kubernetes Engine (GKE) Autopilot 叢集中使用 Cloud TPU 加速器 (TPU)，加快機器學習 (ML) 工作負載的速度。這份指南可協助您為機器學習應用程式架構選取正確的程式庫、設定 TPU 工作負載，在 GKE 上以最佳狀態執行，以及在部署後監控工作負載。

本頁面適用於平台管理員和營運人員、資料和 AI 專家，以及想在 TPU 上準備及執行 ML 工作負載的應用程式開發人員。如要進一步瞭解我們在 Google Cloud 內容中提及的常見角色、職責和範例工作，請參閱「常見的 GKE 使用者角色和工作」。

閱讀本頁面之前，請務必先熟悉下列資源：

TPU 在 Autopilot 中的運作方式

如要在 Autopilot 工作負載中使用 TPU，請在工作負載資訊清單中指定下列項目：

  • spec.nodeSelector 欄位中的 TPU 版本。
  • spec.nodeSelector 欄位中的 TPU 拓撲。拓撲必須支援指定的 TPU 版本。
  • spec.containers.resources.requests 和 spec.containers.resources.limits 欄位中的 TPU 晶片數量。

部署工作負載時，GKE 會佈建具有所要求 TPU 設定的節點，並在節點上排定 Pod。GKE 會將每個工作負載放在自己的節點上，讓每個 Pod 都能存取節點的完整資源，並將中斷風險降到最低。

Autopilot 中的 TPU 支援下列功能：

  1. Spot Pod
  2. 特定運算資源預留項目
  3. 延長執行時間的 Pod
  4. 彈性啟動

規劃 TPU 設定

使用本指南部署 TPU 工作負載前，請根據模型和所需記憶體量規劃 TPU 設定。詳情請參閱「規劃 TPU 設定」。

定價

如需定價資訊，請參閱「Autopilot 定價」。

事前準備

開始之前，請確認您已完成下列工作：

  • 啟用 Google Kubernetes Engine API。
    • 啟用 Google Kubernetes Engine API
  • 如要使用 Google Cloud CLI 執行這項工作，請安裝並初始化 gcloud CLI。如果您先前已安裝 gcloud CLI，請執行 gcloud components update 指令，取得最新版本。較舊的 gcloud CLI 版本可能不支援執行本文中的指令。
  • 確認您擁有執行 GKE 1.32.3-gke.1927000 以上版本的 Autopilot 叢集。如需操作說明，請參閱「建立 Autopilot 叢集」。
  • 如要使用預留的 TPU，請確認您已預留特定容量。如需操作說明，請參閱「使用預訂」。

確認 TPU 和其他 GKE 資源的配額

下列各節說明如何確保在 GKE 中使用 TPU 時有足夠的配額。

如要建立 TPU 節點，您必須有可用的 TPU 配額，除非您使用現有的容量預留項目。如果您使用的是預留 TPU，請略過本節。

在 GKE 中建立 TPU 節點時，需要 Compute Engine API 配額 (compute.googleapis.com)，而非 Cloud TPU API 配額 (tpu.googleapis.com)。一般 Autopilot Pod 和 Spot Pod 的配額名稱不同。

如要查看 TPU 的 Compute Engine API 配額上限和目前用量，請按照下列步驟操作：

  1. 前往 Google Cloud 控制台的「配額」頁面：

    前往配額頁面

  2. 在「篩選器」方塊中，執行下列操作：

    1. 請根據 TPU 版本和 cloud.google.com/gke-tpu-accelerator 節點選取器中的值，使用下表選取並複製配額的屬性。舉例來說，如果您打算建立隨選 TPU v5e 節點，且 cloud.google.com/gke-tpu-accelerator 節點選取器中的值為 tpu-v5-lite-podslice，請輸入 Name: TPU v5 Lite PodSlice chips。

      TPU 版本 cloud.google.com/gke-tpu-accelerator 隨選執行個體配額的屬性和名稱 Spot 2 執行個體的配額屬性和名稱
      TPU v3、
      tpu-v3-device 
      Dimensions (e.g. location):
      tpu_family:CT3
      		 不適用
      TPU v3、
      tpu-v3-slice 
      Dimensions (e.g. location):
      tpu_family:CT3P
      		 不適用
      TPU v4、
      tpu-v4-podslice 
      Name:
      TPU v4 PodSlice chips
      		 
      Name:
      Preemptible TPU v4 PodSlice chips
      TPU v5e、
      tpu-v5-lite-podslice 
      Name:
      TPU v5 Lite PodSlice chips
      		 
      Name:
      Preemptible TPU v5 Lite Podslice
      chips
      TPU v5p、
      tpu-v5p-slice 
      Name:
      TPU v5p chips
      		 
      Name:
      Preemptible TPU v5p chips
      TPU Trillium、
      tpu-v6e-slice 
      Dimensions (e.g. location):
      tpu_family:CT6E
      		 
      Name:
      Preemptible TPU slices v6e
      Ironwood (TPU7x) (預覽版)、
      tpu7x 
      Dimensions (e.g. location):
      tpu_family:tpu7x
      		 
      Name:
      Preemptible TPU slices tpu7x

    2. 選取「Dimensions (e.g. locations)」(維度，例如位置) 屬性，然後輸入 region:，接著輸入您打算在 GKE 中建立 TPU 的區域名稱。舉例來說，如果您打算在 us-west4-a 區域建立 TPU 節點，請輸入 region:us-west4。TPU 配額屬於區域性配額，因此同一區域內的所有可用區都會消耗相同的 TPU 配額。

如果沒有任何配額符合您輸入的篩選條件，表示專案在您需要的區域中未獲得任何指定配額，因此您必須申請調整 TPU 配額。

建立 TPU 預訂時，對應配額的上限和目前用量值都會增加，增幅為 TPU 預訂中的晶片數量。舉例來說，如果為 16 個 TPU v5e 晶片建立預訂，且cloud.google.com/gke-tpu-accelerator節點選取器中的值為 tpu-v5-lite-podslice，則相關區域中 TPU v5 Lite PodSlice chips 配額的「限制」和「目前用量」都會增加 16。

其他 GKE 資源的配額

您可能需要在 GKE 建立資源的區域中，增加下列 GKE 相關配額。

  • SSD 永久磁碟 (GB) 配額：每個 Kubernetes 節點的開機磁碟預設需要 100 GB。因此，這項配額至少應設為您預期建立的 GKE 節點數量上限與 100 GB 的乘積 (節點 * 100 GB)。
  • 使用中 IP 位址配額：每個 Kubernetes 節點會占用一個 IP 位址。 因此，這項配額至少應設為您預期建立的 GKE 節點數量上限。
  • 確認 max-pods-per-node 與子網路範圍一致：每個 Kubernetes 節點都會使用 Pod 的次要 IP 範圍。舉例來說，32 個 Pod max-pods-per-node 需要 64 個 IP 位址，因此每個節點需要 /26 子網路。請注意，這個範圍不得與任何其他叢集共用。如要避免 IP 位址範圍用盡，請使用 --max-pods-per-node 標記，限制可在節點上排程的 Pod 數量。max-pods-per-node 的配額應至少設為您預計建立的 GKE 節點數量上限。

如要申請提高配額，請參閱「要求調整配額」。

準備 TPU 應用程式

TPU 工作負載必須符合下列準備條件。

  1. JAX、PyTorch 和 TensorFlow 等架構會使用 libtpu 共用程式庫存取 TPU VM。libtpu 包含 XLA 編譯器、TPU 執行階段軟體和 TPU 驅動程式。每個版本的 PyTorch 和 JAX 都需要特定版本的 libtpu.so。為避免套件版本衝突，建議使用 JAX AI 映像檔。如要在 GKE 中使用 TPU，請務必使用下列版本：tpu7x
    TPU 類型 「libtpu.so」版本
    Ironwood (TPU7x) (預先發布版)
    TPU Trillium (v6e)
    tpu-v6e-slice
    TPU v5e
    tpu-v5-lite-podslice
    TPU v5p
    tpu-v5p-slice
    TPU v4
    tpu-v4-podslice
    TPU v3
    tpu-v3-slice
    tpu-v3-device
  2. 在工作負載資訊清單中新增 Kubernetes 節點選取器，確保 GKE 會在您定義的 TPU 機器類型和 TPU 拓撲上排定 TPU 工作負載：

      nodeSelector:
    cloud.google.com/gke-tpu-accelerator: TPU_ACCELERATOR
    cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
    cloud.google.com/placement-policy-name: WORKLOAD_POLICY # Required only for Ironwood (TPU7x)

    更改下列內容：

    • TPU_ACCELERATOR：TPU 加速器的名稱。例如，使用 tpu7x-standard-4t。
    • TPU_TOPOLOGY：TPU 配量的實體拓撲。拓撲格式取決於 TPU 版本。例如，使用 2x2x2。詳情請參閱「在 GKE 中規劃 TPU」。
    • WORKLOAD_POLICY：您要用來放置 TPU Pod 的工作負載政策名稱。只有 Ironwood (TPU7x) 需要這個節點選取器。

完成工作負載準備後，即可執行使用 TPU 的工作。

在 GKE 中佈建 TPU 的選項

如要在 GKE 中佈建 TPU，您可以選擇下列設定選項：
  • 工作負載要求：您可以在 spec.nodeSelector 欄位中指定 TPU 版本和拓撲，並在 spec.containers.resources 區段中指定 TPU 晶片數量。部署工作負載時，GKE 會自動佈建具有正確 TPU 設定的節點，並將每個工作負載放在專屬節點上，確保工作負載能完整存取節點資源。如需操作說明，請參閱「在工作負載中要求使用 TPU」。

    • 使用自訂運算類別集中佈建 TPU

    以下各節說明如何建立自訂 ComputeClass，然後建立使用 ComputeClass 中定義 TPU 的 Job。

    建立自訂 ComputeClass

    如要建立符合 TPU 規則的自訂 ComputeClass，步驟會因您使用 Ironwood (TPU7x) 或較早的 TPU 版本而異。

    Ironwood (TPU7x)

    1. 建立工作負載政策。只有在建立多主機節點集區時，才需要執行這個步驟，這取決於您選擇的拓撲。如果您使用單一主機節點集區，請略過這個步驟。

      gcloud compute resource-policies create workload-policy WORKLOAD_POLICY_NAME \
    --type=HIGH_THROUGHPUT \
    --accelerator-topology=TPU_TOPOLOGY \
    --project=PROJECT_ID \
    --region=REGION

      更改下列內容：

      • WORKLOAD_POLICY_NAME：工作負載政策的名稱。
      • TPU_TOPOLOGY：TPU Ironwood (TPU7x) 拓撲。例如，使用 2x2x2。如要進一步瞭解所有支援的 Ironwood (TPU7x) 拓撲，請參閱拓撲部分。
      • PROJECT_ID：您的 Google Cloud 專案 ID
      • REGION：工作負載政策的區域。工作負載政策是區域資源，可用於節點集區。

    2. 將下列資訊清單儲存為 tpu-compute-class.yaml：

      apiVersion: cloud.google.com/v1
kind: ComputeClass
metadata:
  name: tpu-class
spec:
  priorities:
    - tpu:
        type: tpu7x
        topology: TPU_TOPOLOGY
        count: 4
      placement:
        policyName: WORKLOAD_POLICY_NAME
  nodePoolAutoCreation:
    enabled: true

    3. (選用) 您可以消耗特定預留項目或子區塊。舉例來說，您可以在 ComputeClass 資訊清單中新增下列 specs：

        reservations:
    affinity: Specific
    specific:
      - name: RESERVATION_NAME
        reservationBlock:
          name: RESERVATION_BLOCK_NAME
          reservationSubBlock:
            name: RESERVATION_SUB_BLOCK_NAME

      更改下列內容：

      • RESERVATION_NAME：Compute Engine 容量預留項目的名稱。
      • RESERVATION_BLOCK_NAME：Compute Engine 容量預留區塊的名稱。
      • RESERVATION_SUB_BLOCK_NAME：Compute Engine 容量預留子區塊的名稱。

      詳情請參閱「耗用可用區保留資源」。

    其他 TPU 版本

    如要使用為 TPU 設定的自訂 ComputeClass，佈建 v3、v4、v5p、v5e 或 v6e (Trillium) TPU，請完成下列步驟：

    1. 將下列資訊清單儲存為 tpu-compute-class.yaml：

      apiVersion: cloud.google.com/v1
kind: ComputeClass
metadata:
  name: tpu-class
spec:
  priorities:
  - tpu:
      type: TPU_TYPE
      count: NUMBER_OF_CHIPS
      topology: TOPOLOGY
  - spot: true
    tpu:
      type: {"<var>"}}TPU_TYPE
      count: NUMBER_OF_CHIPS
      topology: TOPOLOGY
  - flexStart:
      enabled: true
    tpu:
      type: {"<var>"}}TPU_TYPE
      count: NUMBER_OF_CHIPS
      topology: TOPOLOGY
  nodePoolAutoCreation:
    enabled: true

      更改下列內容：

      • TPU_TYPE：要使用的 TPU 類型，例如 tpu-v4-podslice。必須是 GKE 支援的值。
      • TOPOLOGY：配量中的 TPU 晶片排列方式，例如 2x2x4。必須是所選 TPU 類型支援的拓撲。
      • NUMBER_OF_CHIPS：供容器使用的 TPU 晶片數量。limits 和 requests 的值必須相同。

    2. 部署 ComputeClass：

      kubectl apply -f tpu-compute-class.yaml

      如要進一步瞭解自訂 ComputeClass 和 TPU，請參閱 TPU 設定。

    建立會耗用 TPU 的 Job

    1. 將下列資訊清單儲存為 tpu-job.yaml：

      apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/compute-class: tpu-class
      containers:
      - name: tpu-job
        image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS
          limits:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS

      更改下列內容：

      • NUMBER_OF_CHIPS：供容器使用的 TPU 晶片數量。limits 和 requests 的值必須相同，且等於所選自訂 ComputeClass 中的 CHIP_COUNT 值。
      • MEMORY_SIZE：TPU 使用的記憶體量上限。記憶體限制取決於您使用的 TPU 版本和拓撲。詳情請參閱「加速器的最低和最高值」。
      • NUMBER_OF_CHIPS：供容器使用的 TPU 晶片數量。limits 和 requests 的值必須相同。

    2. 部署 Job：

      kubectl create -f tpu-job.yaml

      建立這項 Job 時，GKE 會自動執行下列操作：

      • 佈建節點來執行 Pod。視您指定的 TPU 類型、拓撲和資源要求而定，這些節點可以是單一主機配量或多主機配量。視最高優先順序的 TPU 資源可用性而定，GKE 可能會改用較低的優先順序，盡量取得資源。
      • 在 Pod 中新增 taint，並在節點中新增容許規定，防止其他工作負載與 TPU 工作負載在相同節點上執行。

      詳情請參閱「關於自訂 ComputeClass」。

    3. 完成本節後，您可以刪除已建立的資源，避免繼續計費：

      kubectl delete -f tpu-job.yaml

    在工作負載中要求 TPU

    本節說明如何建立在 Autopilot 中要求 TPU 的 Job。在任何需要 TPU 的工作負載中，您都必須指定下列項目：

    • TPU 版本和拓撲的節點選取器
    • 工作負載中容器的 TPU 晶片數量

    如要查看支援的 TPU 版本、拓撲，以及配量中對應的 TPU 晶片和節點數量，請參閱「選擇 TPU 版本」。

    工作負載中 TPU 要求的注意事項

    Pod 中只有一個容器可以使用 TPU。容器要求的 TPU 晶片數量必須等於配量中節點所連結的 TPU 晶片數量。舉例來說，如果您要求使用 2x4 拓撲的 TPU v5e (tpu-v5-lite-podslice)，可以要求下列任一項目：

    • 4 個晶片，這會建立兩個多主機節點，每個節點有 4 個 TPU 晶片
    • 8 晶片，建立具有 8 個 TPU 晶片的單一主機節點

    為盡量提高成本效益，最佳做法是務必使用所要求切片中的所有 TPU。如果您要求兩個節點的多主機配量，每個節點有 4 個 TPU 晶片，則應部署在兩個節點上執行的工作負載，並耗用配量中的所有 8 個 TPU 晶片。

    建立要求 TPU 的工作負載

    下列步驟會建立要求 TPU 的 Job。如果您有在多主機 TPU 配量上執行的工作負載，也必須建立無週邊服務，並依名稱選取工作負載。這個無頭服務可讓多主機切片中不同節點上的 Pod 相互通訊，方法是更新 Kubernetes DNS 設定，指向工作負載中的 Pod。

    1. 將下列資訊清單儲存為 tpu-autopilot.yaml：

      apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-job
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-job
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      # Optional: Run in GKE Sandbox
      # runtimeClassName: gvisor
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: TPU_TYPE
        cloud.google.com/gke-tpu-topology: TOPOLOGY
      containers:
      - name: tpu-job
        image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          python -c 'import jax; print("TPU cores:", jax.device_count())'
        resources:
          requests:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS
          limits:
            cpu: 10
            memory: MEMORY_SIZE
            google.com/tpu: NUMBER_OF_CHIPS

      更改下列內容：

      • TPU_TYPE：要使用的 TPU 類型，例如 tpu-v4-podslice。必須是 GKE 支援的值。
      • ：配量中的 TPU 晶片排列方式，例如 2x2x4。TOPOLOGY必須是所選 TPU 類型支援的拓撲。
      • NUMBER_OF_CHIPS：供容器使用的 TPU 晶片數量。limits 和 requests 的值必須相同。
      • MEMORY_SIZE：TPU 使用的記憶體上限。記憶體限制取決於您使用的 TPU 版本和拓撲。詳情請參閱「加速器的最低和最高值」。

      此外，您也可以選擇變更下列欄位：

      • image：要使用的 JAX AI 圖片。在範例資訊清單中，這個欄位設為最新的 JAX AI 映像檔。如要設定其他版本，請參閱目前的 JAX AI 圖片清單。
      • runtimeClassname: gvisor：這項設定可讓您在 GKE Sandbox 中執行這個 Pod。如要使用，請取消註解這一行。GKE Sandbox 支援 TPU 第 4 版以上。詳情請參閱「GKE Sandbox」。

    2. 部署 Job：

      kubectl create -f tpu-autopilot.yaml

      建立這項 Job 時，GKE 會自動執行下列操作：

      1. 佈建節點來執行 Pod。視您指定的 TPU 類型、拓撲和資源要求而定，這些節點可以是單一主機配量或多主機配量。
      2. 在 Pod 中新增 taint，並在節點中新增容許規定，防止其他工作負載與 TPU 工作負載在相同節點上執行。

    3. 完成本節後，您可以刪除建立的工作負載，避免系統繼續計費：

      kubectl delete -f tpu-autopilot.yaml

    建立要求 TPU 和集合排程的工作負載

    在 TPU Trillium 中，您可以使用集合排程，將 TPU 節點切片分組。將這些 TPU 節點分組後，就能更輕鬆地調整副本數量，以滿足工作負載需求。 Google Cloud 會控管軟體更新，確保集合內有足夠的切片，隨時可處理流量。

    TPU Trillium 支援單一主機和多主機節點集區的集合排程，可執行推論工作負載。以下說明收集排程行為如何取決於所用 TPU 配量類型：

    • 多主機 TPU 配量：GKE 會將多主機 TPU 配量分組，形成集合。每個 GKE 節點集區都是這個集合中的副本。如要定義集合，請建立多主機 TPU 配量，並為集合指派專屬名稱。如要將更多 TPU 節點加入集合，請建立另一個多主機 TPU 節點集區，並使用相同的集合名稱和工作負載類型。
    • 單一主機 TPU 配量：GKE 會將整個單一主機 TPU 配量節點集區視為集合。如要將更多 TPU 配量新增至集合，可以調整單一主機 TPU 配量節點集區的大小。

    如要瞭解收藏內容排程的限制，請參閱「收藏內容排程的運作方式」

    使用多主機 TPU 配量

    在多主機 TPU 節點中排定收集作業，適用於 1.31.2-gke.1537000 以上版本的 Autopilot 叢集。只有 1.31.2-gke.1115000 以上版本支援具有 2x4 拓撲的多主機 TPU 節點。如要建立多主機 TPU 節點，並將其分組為集合，請在工作負載規格中新增下列 Kubernetes 標籤：

    • cloud.google.com/gke-nodepool-group-name：每個集合在叢集層級都應有專屬名稱。cloud.google.com/gke-nodepool-group-name 標籤中的值必須符合叢集標籤規定。

    • cloud.google.com/gke-workload-type: HIGH_AVAILABILITY

      舉例來說，下列程式碼區塊會定義具有多主機 TPU 切片的集合：

        nodeSelector:
    cloud.google.com/gke-nodepool-group-name: ${COLLECTION_NAME}
    cloud.google.com/gke-workload-type: HIGH_AVAILABILITY
    cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
    cloud.google.com/gke-tpu-topology: 4x4
...

    使用單一主機 TPU 配量

    在 1.31.2-gke.1088000 以上版本的 Autopilot 叢集中，單一主機 TPU 節點可進行收集排程。如要建立單一主機 TPU 節點，並將其分組為集合，請在工作負載規格中新增 cloud.google.com/gke-workload-type:HIGH_AVAILABILITY 標籤。

    舉例來說，下列程式碼區塊會定義具有單一主機 TPU 切片的集合：

      nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v6e-slice
    cloud.google.com/gke-tpu-topology: 2x2
    cloud.google.com/gke-workload-type: HIGH_AVAILABILITY
  ...

    使用自訂運算類別部署集合

    如要進一步瞭解如何使用自訂運算類別，部署要求 TPU 工作負載和集合排程的工作負載，請參閱「TPU 多主機集合」和「定義 TPU SLO 的工作負載類型」。

    範例：顯示多主機配量中的 TPU 晶片總數

    下列工作負載會傳回多主機 TPU 節點中所有節點的 TPU 晶片數量。如要建立多主機切片，工作負載必須具備下列參數：

    • TPU 版本：TPU v4
    • 拓撲：2x2x4

    這個版本和拓撲選取項目會產生多主機配量。

    1. 將下列資訊清單儲存為 available-chips-multihost.yaml： 
      apiVersion: v1
kind: Service
metadata:
  name: headless-svc
spec:
  clusterIP: None
  selector:
    job-name: tpu-available-chips
---
apiVersion: batch/v1
kind: Job
metadata:
  name: tpu-available-chips
spec:
  backoffLimit: 0
  completions: 4
  parallelism: 4
  completionMode: Indexed
  template:
    spec:
      subdomain: headless-svc
      restartPolicy: Never
      nodeSelector:
        cloud.google.com/gke-tpu-accelerator: tpu-v4-podslice # Node selector to target TPU v4 slice nodes.
        cloud.google.com/gke-tpu-topology: 2x2x4 # Specifies the physical topology for the TPU slice.
      containers:
      - name: tpu-job
        image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
        ports:
        - containerPort: 8471 # Default port using which TPU VMs communicate
        - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
        command:
        - bash
        - -c
        - |
          python -c 'import jax; print("TPU cores:", jax.device_count())' # Python command to count available TPU chips.
        resources:
          requests:
            cpu: 10
            memory: 407Gi
            google.com/tpu: 4 # Request 4 TPU chips for this workload.
          limits:
            cpu: 10
            memory: 407Gi
            google.com/tpu: 4 # Limit to 4 TPU chips for this workload.
    2. 部署資訊清單： 
      kubectl create -f available-chips-multihost.yaml

      GKE 會執行含有四個 VM 的 TPU v4 配量 (多主機 TPU 配量)。配量有 16 個互連的 TPU 晶片。

    3. 確認 Job 已建立四個 Pod： 
      kubectl get pods

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

      NAME                       READY   STATUS      RESTARTS   AGE
tpu-job-podslice-0-5cd8r   0/1     Completed   0          97s
tpu-job-podslice-1-lqqxt   0/1     Completed   0          97s
tpu-job-podslice-2-f6kwh   0/1     Completed   0          97s
tpu-job-podslice-3-m8b5c   0/1     Completed   0          97s
    4. 取得其中一個 Pod 的記錄： 
      kubectl logs POD_NAME

      將 POD_NAME 替換為其中一個已建立的 Pod 名稱。例如：tpu-job-podslice-0-5cd8r。

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

      TPU cores: 16
    5. 選用：移除工作負載： 
      kubectl delete -f available-chips-multihost.yaml

    範例：顯示單一節點中的 TPU 晶片

    下列工作負載是靜態 Pod，會顯示附加至特定節點的 TPU 晶片數量。如要建立單一主機節點，工作負載必須具備下列參數：

    • TPU 版本：TPU v5e
    • 拓撲：2x4

    這個版本和拓撲選取項目會產生單一主機配量。

    1. 將下列資訊清單儲存為 available-chips-singlehost.yaml： 
      apiVersion: v1
kind: Pod
metadata:
  name: tpu-job-jax-v5
spec:
  restartPolicy: Never
  nodeSelector:
    cloud.google.com/gke-tpu-accelerator: tpu-v5-lite-podslice # Node selector to target TPU v5e slice nodes.
    cloud.google.com/gke-tpu-topology: 2x4 # Specify the physical topology for the TPU slice.
  containers:
  - name: tpu-job
    image: us-docker.pkg.dev/cloud-tpu-images/jax-ai-image/tpu:latest
    ports:
    - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
    command:
    - bash
    - -c
    - |
      python -c 'import jax; print("Total TPU chips:", jax.device_count())'
    resources:
      requests:
        google.com/tpu: 8 # Request 8 TPU chips for this container.
      limits:
        google.com/tpu: 8 # Limit to 8 TPU chips for this container.
    2. 部署資訊清單： 
      kubectl create -f available-chips-singlehost.yaml

      GKE 會佈建節點，其中包含八個使用 TPU v5e 的單一主機 TPU 配量。每個 TPU 節點有八個 TPU 晶片 (單一主機 TPU 配量)。

    3. 取得 Pod 的記錄： 
      kubectl logs tpu-job-jax-v5

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

      Total TPU chips: 8
    4. 選用：移除工作負載： 
        kubectl delete -f available-chips-singlehost.yaml

    觀察及監控 TPU

    資訊主頁

    Google Cloud 控制台中的節點集區可觀測性功能現已全面開放使用。 如要查看 GKE 上 TPU 多主機節點集區的狀態，請前往 Cloud Monitoring 提供的「GKE TPU Node Pool Status」(GKE TPU 節點集區狀態) 資訊主頁：

    前往 GKE TPU 節點集區狀態

    這個資訊主頁可提供多主機 TPU 節點集區的健康狀態深入分析。 詳情請參閱「監控 TPU 節點和節點集區的健康狀態指標」。

    在Google Cloud 控制台的「Kubernetes Clusters」頁面中，「Observability」分頁標籤也會顯示「Accelerators」>「TPU」標題下方的 TPU 可觀測性指標，例如 TPU 使用量。詳情請參閱「查看可觀測性指標」。

    只有在 GKE 叢集中啟用系統指標時，TPU 資訊主頁才會顯示資料。

    執行階段指標

    在 GKE 1.27.4-gke.900 以上版本中，如果 TPU 工作負載使用 JAX 0.4.14 以上版本，並指定 containerPort: 8431，就會將 TPU 使用率指標匯出為 GKE 系統指標。您可以在 Cloud Monitoring 中使用下列指標，監控 TPU 工作負載的執行階段效能：

    • 任務週期：在過去取樣期間 (60 秒) 內，TensorCore 在 TPU 晶片上主動處理作業的時間百分比。百分比越高，表示 TPU 使用率越高。
    • 已用記憶體：分配的加速器記憶體量，以位元組為單位。每 60 秒取樣一次。
    • 記憶體總量：加速器記憶體總量，以位元組為單位。每 60 秒取樣一次。

    這些指標位於 Kubernetes 節點 (k8s_node) 和 Kubernetes 容器 (k8s_container) 架構中。

    Kubernetes 容器：

    • kubernetes.io/container/accelerator/duty_cycle
    • kubernetes.io/container/accelerator/memory_used
    • kubernetes.io/container/accelerator/memory_total

    Kubernetes 節點：

    • kubernetes.io/node/accelerator/duty_cycle
    • kubernetes.io/node/accelerator/memory_used
    • kubernetes.io/node/accelerator/memory_total

    監控 TPU 節點和節點集區的健康狀態指標

    如果訓練工作發生錯誤或終止失敗，您可以檢查基礎架構相關指標，判斷中斷是否是由基礎節點或節點集區的問題所導致。

    節點狀態

    在 1.32.1-gke.1357001 以上版本的 GKE 中，下列 GKE 系統指標會公開 GKE 節點的狀態：

    • kubernetes.io/node/status_condition

    condition 欄位會回報節點的狀況，例如 Ready、DiskPressure 和 MemoryPressure。status 欄位會顯示回報的狀況狀態，可以是 True、False 或 Unknown。這是具有 k8s_node 受控資源類型的指標。

    這項 PromQL 查詢會顯示特定節點是否為 Ready：

    kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    node_name="NODE_NAME",
    condition="Ready",
    status="True"}

    如要協助排解叢集問題，您可能需要查看出現其他狀況的節點：

    kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    condition!="Ready",
    status="True"}

    您可能想特別查看未Ready的節點：

    kubernetes_io:node_status_condition{
    monitored_resource="k8s_node",
    cluster_name="CLUSTER_NAME",
    condition="Ready",
    status="False"}

    如果沒有資料，節點就已準備就緒。系統每 60 秒會取樣一次狀態條件。

    您可以使用下列查詢，瞭解整個車隊的節點狀態：

    avg by (condition,status)(
  avg_over_time(
    kubernetes_io:node_status_condition{monitored_resource="k8s_node"}[${__interval}]))

    節點集區狀態

    下列 GKE 系統指標適用於 k8s_node_pool 受監控的資源，可顯示 GKE 節點集區的狀態：

    • kubernetes.io/node_pool/status

    這項指標只會針對多主機 TPU 節點集區回報。

    status 欄位會回報節點集區的狀態，例如 Provisioning、Running、Error、Reconciling 或 Stopping。狀態會在 GKE API 作業完成後更新。

    如要確認特定節點集區是否處於 Running 狀態，請使用下列 PromQL 查詢：

    kubernetes_io:node_pool_status{
    monitored_resource="k8s_node_pool",
    cluster_name="CLUSTER_NAME",
    node_pool_name="NODE_POOL_NAME",
    status="Running"}

    如要監控專案中節點集區的數量 (依狀態分組)，請使用下列 PromQL 查詢：

    count by (status)(
  count_over_time(
    kubernetes_io:node_pool_status{monitored_resource="k8s_node_pool"}[${__interval}]))

    節點集區可用性

    下列 GKE 系統指標會顯示多主機 TPU 節點集區是否可用：

    • kubernetes.io/node_pool/multi_host/available

    如果節點集區中的所有節點均可使用，指標值為 True，否則為 False。系統每 60 秒會對指標取樣一次。

    如要檢查專案中是否提供多主機 TPU 節點集區，請使用下列 PromQL 查詢：

    avg by (node_pool_name)(
  avg_over_time(
    kubernetes_io:node_pool_multi_host_available{
      monitored_resource="k8s_node_pool",
      cluster_name="CLUSTER_NAME"}[${__interval}]))

    節點中斷次數

    下列 GKE 系統指標會回報自上次取樣以來，GKE 節點的中斷次數 (指標每 60 秒取樣一次)：

    • kubernetes.io/node/interruption_count

    interruption_type (例如 TerminationEvent、MaintenanceEvent 或 PreemptionEvent) 和 interruption_reason (例如 HostError、Eviction 或 AutoRepair) 欄位可協助提供節點中斷的原因。

    如要查看專案中叢集 TPU 節點的中斷情形和原因，請使用下列 PromQL 查詢：

      sum by (interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_interruption_count{monitored_resource="k8s_node"}[${__interval}]))

    如要只查看主機維護事件，請更新查詢，篩選 interruption_reason 的 HW/SW Maintenance 值。使用下列 PromQL 查詢：

      sum by (interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_interruption_count{monitored_resource="k8s_node", interruption_reason="HW/SW Maintenance"}[${__interval}]))

    如要查看依節點集區匯總的中斷次數，請使用下列 PromQL 查詢：

      sum by (node_pool_name,interruption_type,interruption_reason)(
    sum_over_time(
      kubernetes_io:node_pool_interruption_count{monitored_resource="k8s_node_pool", interruption_reason="HW/SW Maintenance", node_pool_name=NODE_POOL_NAME }[${__interval}]))

    節點集區復原時間 (TTR)

    下列 GKE 系統指標會回報 GKE 多主機 TPU 節點集區的復原期長度分布：

    • kubernetes.io/node_pool/accelerator/times_to_recover

    這項指標中記錄的每個樣本，都代表節點集區從停機期間復原的單一事件。

    這項指標有助於追蹤多主機 TPU 節點集區的復原時間，以及服務中斷的間隔時間。

    您可以使用下列 PromQL 查詢，計算叢集過去 7 天的平均修復時間 (MTTR)：

    sum(sum_over_time(
  kubernetes_io:node_pool_accelerator_times_to_recover_sum{
    monitored_resource="k8s_node_pool", cluster_name="CLUSTER_NAME"}[7d]))
/
sum(sum_over_time(
  kubernetes_io:node_pool_accelerator_times_to_recover_count{
    monitored_resource="k8s_node_pool",cluster_name="CLUSTER_NAME"}[7d]))

    節點集區中斷間隔時間 (TBI)

    節點集區中斷情形間隔時間會衡量基礎架構在發生中斷情形前運作的時間長度。 計算方式是取一段時間內的平均值，分子是基礎架構的總運作時間，分母是基礎架構的總中斷次數。

    下列 PromQL 範例顯示指定叢集的中斷情形平均間隔時間 (MTBI) (7 天)：

    sum(count_over_time(
  kubernetes_io:node_memory_total_bytes{
    monitored_resource="k8s_node", node_name=~"gke-tpu.*|gk3-tpu.*", cluster_name="CLUSTER_NAME"}[7d]))
/
sum(sum_over_time(
  kubernetes_io:node_interruption_count{
    monitored_resource="k8s_node", node_name=~"gke-tpu.*|gk3-tpu.*", cluster_name="CLUSTER_NAME"}[7d]))

    主機指標

    在 GKE 1.28.1-gke.1066000 以上版本中，TPU 節點中的 VM 會將 TPU 使用率指標匯出為 GKE 系統指標。您可以使用 Cloud Monitoring 監控下列 TPU 主機效能指標：

    • TensorCore 使用率：目前 TensorCore 的用量占比。TensorCore 值等於矩陣乘法單元 (MXU) 加上向量單元的總和。 TensorCore 使用率值是將過去取樣期間 (60 秒) 執行的 TensorCore 作業數量，除以相同期間支援的 TensorCore 作業數量。 值越大代表使用率越高。
    • 記憶體頻寬使用率：目前使用的加速器記憶體頻寬百分比。將取樣期間 (60 秒) 內使用的記憶體頻寬，除以相同取樣期間內支援的最大頻寬，即可得出這項指標。

    這些指標位於 Kubernetes 節點 (k8s_node) 和 Kubernetes 容器 (k8s_container) 架構中。

    Kubernetes 容器：

    • kubernetes.io/container/accelerator/tensorcore_utilization
    • kubernetes.io/container/accelerator/memory_bandwidth_utilization

    Kubernetes 節點：

    • kubernetes.io/node/accelerator/tensorcore_utilization
    • kubernetes.io/node/accelerator/memory_bandwidth_utilization

    詳情請參閱 Kubernetes 指標和 GKE 系統指標。

    記錄

    GKE 節點 (包括 TPU VM) 上執行的容器所產生的記錄檔，會由 GKE 記錄代理程式收集，並傳送至 Logging，然後顯示在 Logging 中。

    Autopilot 中 TPU 工作負載的最佳化建議

    下列建議可能有助於提升 TPU 工作負載的效率：

    • 使用延長執行時間的 Pod，在 GKE 終止 Pod 之前，最多可享有七天的緩衝期，以進行縮減或節點升級。您可以搭配維護期間和排除項目，以及執行時間較長的 Pod，進一步延後節點自動升級作業。
    • 使用容量預留功能，確保工作負載能取得要求的 TPU，不必排隊等待。

    如要瞭解如何在 GKE 中設定 Cloud TPU，請參閱下列 Google Cloud 資源：