在 GKE Autopilot 上部署 TPU 工作负载

本页面介绍了如何通过在 Google Kubernetes Engine (GKE) Autopilot 集群中使用 Cloud TPU 加速器 (TPU) 来加速机器学习 (ML) 工作负载。该指导可帮助您为机器学习应用框架选择正确的库，将 TPU 工作负载设置为在 GKE 上以最佳方式运行，并在部署后监控工作负载。

本页面适用于希望在 TPU 上准备和运行机器学习工作负载的平台管理员和运维人员、数据和 AI 专家以及应用开发者。如需详细了解我们在 Google Cloud 内容中提及的常见角色、职责和示例任务，请参阅常见的 GKE 用户角色和任务。

在阅读本页面之前，请确保您熟悉以下资源：

• Cloud TPU 简介
• Cloud TPU 系统架构
• 关于 GKE 中的 TPU

Autopilot 中 TPU 的工作原理

如需在 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 以获取最新版本。

• 确保您拥有运行 GKE 1.32.3-gke.1927000 版或更高版本的 Autopilot 集群。如需查看相关说明，请参阅创建 Autopilot 集群。
• 如需使用预留的 TPU，请确保您已有特定容量预留。如需查看相关说明，请参阅使用预留的可用区级资源。

确保 TPU 和其他 GKE 资源的配额

在 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. 在 filter_list 过滤条件框中，执行以下操作：

   1. 请使用下表，根据 TPU 版本和 cloud.google.com/gke-tpu-accelerator 节点选择器中的值选择并复制配额的属性。例如，如果您计划创建在 cloud.google.com/gke-tpu-accelerator 节点选择器中其值为 tpu-v5-lite-podslice 的按需 TPU v5e 节点，请输入 Name: TPU v5 Lite PodSlice chips。

      TPU 版本， cloud.google.com/gke-tpu-accelerator	按需实例的配额的属性和名称	Spot2 实例的配额的属性和名称
      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

   2. 选择维度（例如位置）属性，然后输入 region:，后跟您计划在其中创建 GKE 中 TPU 的区域的名称。例如，如果您计划在可用区 us-west4-a 中创建 TPU 切片节点，请输入 region:us-west4。TPU 配额是区域性的，因此同一区域内的所有可用区都使用同一 TPU 配额。

如果没有配额与您输入的过滤条件匹配，则说明项目尚未获得所需区域的任何指定配额，您必须申请 TPU 配额调整。

创建 TPU 预留后，相应配额的限制值和当前使用值都会增加 TPU 预留中的芯片数量。例如，为在 cloud.google.com/gke-tpu-accelerator 节点选择器中其值为 tpu-v5-lite-podslice 的 16 个 TPU v5e 芯片创建预留后，相关区域中 TPU v5 Lite PodSlice chips 配额的限制和当前用量都会增加 16。

其他 GKE 资源的配额

您可能需要在 GKE 创建资源的区域增加以下与 GKE 相关的配额。

• Persistent Disk SSD (GB) 配额：默认情况下，每个 Kubernetes 节点的启动磁盘需要 100GB。因此，此配额应至少设置为您预计创建的 GKE 节点数上限和 100GB 的乘积（节点数 * 100GB）。
• 使用中的 IP 地址配额：每个 Kubernetes 节点使用一个 IP 地址。因此，此配额应至少设置为您预计创建的 GKE 节点数上限。
• 确保 max-pods-per-node 与子网范围一致：每个 Kubernetes 节点都对 Pod 使用次要 IP 范围。例如，max-pods-per-node 为 32 时需要 64 个 IP 地址，这相当于每个节点一个 /26 子网。请注意，此范围不应与任何其他集群共享。为避免耗尽 IP 地址范围，请使用 --max-pods-per-node 标志来限制允许在节点上调度的 Pod 数量。max-pods-per-node 的配额应至少设置为您预计创建的 GKE 节点数上限。

如需申请增加配额，请参阅申请配额调整。

在 GKE 中预配 TPU 的选项

借助 GKE Autopilot，您可以通过使用 Kubernetes nodeSelector 直接在各个工作负载中使用 TPU。

或者，您也可以使用自定义计算类来请求 TPU。借助自定义计算类，平台管理员可以定义节点配置的层次结构，以便 GKE 在节点扩缩决策期间确定优先级，从而使工作负载在您选择的硬件上运行。

如需查看相关说明，请参阅使用自定义计算类集中预配 TPU 部分。

准备 TPU 应用

TPU 工作负载具有以下准备要求。

1. JAX、PyTorch 和 TensorFlow 等框架使用 libtpu 共享库访问 TPU 虚拟机。libtpu 包括 XLA 编译器、TPU 运行时软件和 TPU 驱动程序。每个 PyTorch 和 JAX 版本都需要特定的 libtpu.so 版本。如需使用 GKE 中的 TPU，请确保您使用以下版本：

   TPU 类型	libtpu.so版本
   TPU Trillium (v6e) tpu-v6e-slice	推荐的 jax[tpu] 版本：v0.4.9 或更高版本 推荐的 torchxla[tpuvm] 版本：v2.1.0 或更高版本。
   TPU v5e tpu-v5-lite-podslice	推荐的 jax[tpu] 版本：v0.4.9 或更高版本 推荐的 torchxla[tpuvm] 版本：v2.1.0 或更高版本。
   TPU v5p tpu-v5p-slice	推荐的 jax[tpu] 版本：0.4.19 或更高版本。 推荐的 torchxla[tpuvm] 版本：建议使用 2023 年 10 月 23 日的夜间版本 build。
   TPU v4 tpu-v4-podslice	推荐的 jax[tpu]：v0.4.4 或更高版本 推荐的 torchxla[tpuvm]：v2.0.0 或更高版本
   TPU v3 tpu-v3-slice tpu-v3-device	推荐的 jax[tpu]：v0.4.4 或更高版本 推荐的 torchxla[tpuvm]：v2.0.0 或更高版本

2. 为请求 TPU 资源的容器设置以下环境变量：
   • TPU_WORKER_ID：每个 Pod 的唯一整数。此 ID 表示 TPU 切片中的唯一工作器 ID。此字段支持的值范围从 0 到 Pod 数量减去 1。
   • TPU_WORKER_HOSTNAMES：需要在切片内相互通信的 TPU 虚拟机主机名或 IP 地址的英文逗号分隔列表。切片中的每个 TPU 虚拟机都应该有一个对应的主机名或 IP 地址。IP 地址或主机名列表按 TPU_WORKER_ID 排序且为零索引。

   在 completionMode: Indexed, subdomain, parallelism > 1 情况下创建 Job 并请求 google.com/tpu 属性时，GKE 会使用变更 webhook 自动注入这些环境变量。GKE 会添加一个无头 Service，以便为支持该 Service 的 Pod 添加 DNS 记录。

完成工作负载准备后，您可以运行使用 TPU 的 Job。

在工作负载中请求 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 切片上运行，则还必须创建一个按名称选择工作负载的无头 Service。该无头 Service 通过更新 Kubernetes DNS 配置以指向工作负载中的 Pod，可让多主机切片中不同节点上的 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: python:3.10
           ports:
           - containerPort: 8471 # Default port using which TPU VMs communicate
           - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
           command:
           - bash
           - -c
           - |
             pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
             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 支持的值。
   • TOPOLOGY：切片中 TPU 芯片的排列方式，例如 2x2x4。必须是所选 TPU 类型支持的拓扑。
   • NUMBER_OF_CHIPS：容器要使用的 TPU 芯片数量。对于 limits 和 requests，值必须相同。
   • MEMORY_SIZE：TPU 使用的最大内存量。内存限制取决于您使用的 TPU 版本和拓扑。如需了解详情，请参阅加速器的最小值和最大值。
   • ** 可选 runtimeClassname: gvisor*：此设置可让您在 GKE Sandbox 中运行此 Pod。如需使用，请取消对这一行的注释。GKE Sandbox 支持 TPU v4 及更高版本。如需了解详情，请参阅 GKE Sandbox。

2. 部署作业：

   kubectl create -f tpu-autopilot.yaml
   

   创建此作业时，GKE 会自动执行以下操作：

   1. 预配节点以运行 Pod。这些节点可以是单主机切片或多主机切片，具体取决于您指定的 TPU 类型、拓扑和资源请求。
   2. 向 Pod 添加污点并向节点添加容忍设置，以防止任何其他工作负载与 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 切片

单主机 TPU 切片节点中的集合调度适用于 1.31.2-gke.1088000 版及更高版本的 Autopilot 集群。如需创建单主机 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 并部署工作负载，请完成以下步骤：

1. 将以下清单保存为 tpu-compute-class.yaml：

   apiVersion: cloud.google.com/v1
   kind: ComputeClass
   metadata:
     name: tpu-class
   spec:
     priorities:
     - tpu:
         type: tpu-v5-lite-podslice
         count: 4
         topology: 2x4
     - spot: true
       tpu:
         type: tpu-v5-lite-podslice
         count: 4
         topology: 2x4
     - flexStart:
         enabled: true
       tpu:
         type: tpu-v6e-slice
         count: 4
         topology: 2x4
     nodePoolAutoCreation:
       enabled: true
   

2. 部署计算类：

   kubectl apply -f tpu-compute-class.yaml
   

   如需详细了解自定义计算类和 TPU，请参阅 TPU 配置。

3. 将以下清单保存为 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: python:3.10
           ports:
           - containerPort: 8471 # Default port using which TPU VMs communicate
           - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
           command:
           - bash
           - -c
           - |
             pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
             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，值必须相同，且等于所选自定义计算类中 tpu.count 字段的值。
   • MEMORY_SIZE：TPU 使用的最大内存量。内存限制取决于您使用的 TPU 版本和拓扑。如需了解详情，请参阅加速器的最小值和最大值。
   • NUMBER_OF_CHIPS：容器要使用的 TPU 芯片数量。对于 limits 和 requests，值必须相同。

4. 部署作业：

   kubectl create -f tpu-job.yaml
   

   创建此作业时，GKE 会自动执行以下操作：

   • 预配节点以运行 Pod。这些节点可以是单主机切片或多主机切片，具体取决于您指定的 TPU 类型、拓扑和资源请求。 根据最高优先级的 TPU 资源的可用性，GKE 可能会回退到较低优先级，以最大限度地提高可获取性。
   • 向 Pod 添加污点并向节点添加容忍设置，以防止任何其他工作负载与 TPU 工作负载在相同节点上运行。

   如需了解详情，请参阅自定义计算类简介。

5. 完成本部分后，您可以删除所创建的资源以避免继续计费：

   kubectl delete -f tpu-job.yaml
   

示例：显示多主机切片中的 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
           cloud.google.com/gke-tpu-topology: 2x2x4
         containers:
         - name: tpu-job
           image: python:3.10
           ports:
           - containerPort: 8471 # Default port using which TPU VMs communicate
           - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
           command:
           - bash
           - -c
           - |
             pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
             python -c 'import jax; print("TPU cores:", jax.device_count())'
           resources:
             requests:
               cpu: 10
               memory: 407Gi
               google.com/tpu: 4
             limits:
               cpu: 10
               memory: 407Gi
               google.com/tpu: 4

2. 部署清单：

   kubectl create -f available-chips-multihost.yaml
   

   GKE 会运行具有四个虚拟机的 TPU v4 切片（多主机 TPU 切片）。该切片具有 16 个互连 TPU 芯片。

3. 验证作业是否创建了四个 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
       cloud.google.com/gke-tpu-topology: 2x4
     containers:
     - name: tpu-job
       image: python:3.10
       ports:
       - containerPort: 8431 # Port to export TPU runtime metrics, if supported.
       command:
       - bash
       - -c
       - |
         pip install 'jax[tpu]' -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
         python -c 'import jax; print("Total TPU chips:", jax.device_count())'
       resources:
         requests:
           google.com/tpu: 8
         limits:
           google.com/tpu: 8

2. 部署清单：

   kubectl create -f available-chips-singlehost.yaml
   

   GKE 会预配具有八个单主机 TPU 切片的节点，这些切片使用 TPU v5e。每个 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 节点池状态信息中心：

前往 GKE TPU 节点池状态

此信息中心可让您全面了解多主机 TPU 节点池的健康状况。如需了解详情，请参阅监控 TPU 节点和节点池的健康状况指标。

在Google Cloud 控制台的 Kubernetes 集群页面中，可观测性标签页还会在加速器 > TPU 标题下显示 TPU 可观测性指标，例如 TPU 使用情况。如需了解详情，请参阅查看可观测性指标。

仅当您在 GKE 集群中启用了系统指标时，系统才会填充 TPU 信息中心。

运行时指标

在 GKE 1.27.4-gke.900 版或更高版本中，使用 JAX 0.4.14 版或更高版本并指定 containerPort: 8431 的 TPU 工作负载会以 GKE 系统指标的形式导出 TPU 利用率指标。Cloud Monitoring 提供以下指标来监控 TPU 工作负载的运行时性能：

• 工作周期：TensorCore 在 TPU 芯片上活跃处理的时间占过去的采样周期（60 秒）的百分比。百分比越大，表示 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 节点和节点池的健康状况指标

如果训练作业出错或以失败告终，您可以检查与底层基础架构相关的指标，以确定中断是否是由底层节点或节点池的问题引起的。

节点状态

在 GKE 1.32.1-gke.1357001 版或更高版本中，以下 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}]))

节点池状态

以下 k8s_node_pool 受监控资源的 GKE 系统指标可显示 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 示例显示了指定集群的 7 天平均中断时间 (MTBI)：

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 切片中的虚拟机会将 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 虚拟机）上运行的容器发出的日志会由 GKE 日志记录代理收集，发送到 Logging，并且显示在 Logging 中。

针对 Autopilot 中的 TPU 工作负载的建议

以下建议可能会提高 TPU 工作负载的效率：

• 在长达七天宽限期内使用延长运行时间的 Pod，然后 GKE 终止 Pod 进行缩容或节点升级。您可以将维护窗口和排除项与延长运行时间的 Pod 搭配使用，以进一步延迟自动节点升级。
• 使用容量预留可确保工作负载接收请求的 TPU，而无需排入队列来实现可用性。

如需了解如何设置 GKE 中的 Cloud TPU，请参阅以下 Google Cloud 资源：

• 规划 GKE 中的 TPU，以开始设置 TPU
• 在 GKE Autopilot 中部署 TPU 工作负载
• 在 GKE Standard 中部署 TPU 工作负载
• 了解将 Cloud TPU 用于机器学习任务的最佳做法。
• 视频：利用 GKE 在 Cloud TPU 上构建大规模的机器学习服务。
• 在 TPU 上使用 KubeRay 提供大语言模型
• 了解如何使用 GKE Sandbox 对 GPU 工作负载进行沙盒处理