在 GKE Standard 中部署 TPU 工作负载

本页面为您了解如何在 Google Kubernetes Engine (GKE) 中使用 TPU 加速机器学习 (ML) 工作负载提供了基础知识。TPU 专为矩阵乘法处理而设计，例如大规模深度学习模型训练。TPU 经过优化，可处理庞大的数据集和复杂的机器学习模型，因此具有出色的性能，可为机器学习工作负载提供更高的性价比和能效。在本指南中，您将了解如何使用 Cloud TPU 加速器部署机器学习工作负载、为 TPU 配置配额、为运行 TPU 的节点池配置升级，以及监控 TPU 工作负载指标。

本教程适用于希望使用 Kubernetes 容器编排功能通过 TPU 管理大规模模型训练、调优和推理工作负载的机器学习 (ML) 工程师、平台管理员和运维人员。如需详细了解我们在 Google Cloud 内容中提及的常见角色和示例任务，请参阅常见的 GKE 用户角色和任务。

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

• Cloud TPU
• Cloud TPU 系统架构
• GKE 中的 TPU

准备工作

在开始之前，请确保您已执行以下任务：

• 启用 Google Kubernetes Engine API。
启用 Google Kubernetes Engine API
• 如果您要使用 Google Cloud CLI 执行此任务，请安装并初始化 gcloud CLI。 如果您之前安装了 gcloud CLI，请运行 gcloud components update 以获取最新版本。

规划 TPU 配置

根据您的模型及其所需的内存量规划 TPU 配置。在使用本指南在 TPU 上部署工作负载之前，请先完成规划 TPU 配置中的规划步骤。

确保您拥有 TPU 配额

按需虚拟机或 Spot 虚拟机的配额

如果您要创建具有按需或 Spot 虚拟机的 TPU 切片节点池，则必须在您要使用的区域中拥有足够的 TPU 配额。

创建使用 TPU 预留的 TPU 切片节点池不需要任何 TPU 配额。¹ 对于预留的 TPU，您可以放心地跳过此部分。

在 GKE 中创建按需或 Spot TPU 切片节点池需要 Compute Engine API 配额。Compute Engine API 配额 (compute.googleapis.com) 与 Cloud TPU API 配额 (tpu.googleapis.com) 不同，使用 Cloud TPU API 创建 TPU 时需要后者。

如需查看 TPU 的 Compute Engine API 配额的限制和当前用量，请按照以下步骤操作：

1. 前往 Google Cloud 控制台中的配额页面：

   转到“配额”

2. 在 filter_list 过滤条件框中，执行以下操作：

   1. 请使用下表根据 TPU 版本和机器类型选择并复制配额的属性。例如，如果您计划创建其机器类型以 ct5lp- 开头的按需 TPU v5e 节点，请输入 Name: TPU v5 Lite PodSlice chips。

      TPU 版本，机器类型以下列字符串开头	按需实例的配额的属性和名称	Spot2 实例的配额的属性和名称
      TPU v3， ct3-	Dimensions (e.g. location): tpu_family:CT3	不适用
      TPU v3， ct3p-	Dimensions (e.g. location): tpu_family:CT3P	不适用
      TPU v4， ct4p-	Name: TPU v4 PodSlice chips	Name: Preemptible TPU v4 PodSlice chips
      TPU v5e， ct5lp-	Name: TPU v5 Lite PodSlice chips	Name: Preemptible TPU v5 Lite Podslice chips
      TPU v5p， ct5p-	Name: TPU v5p chips	Name: Preemptible TPU v5p chips
      TPU Trillium， ct6e-	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 预留中的芯片数量。例如，为机器类型以 ct5lp- 开头的 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 节点数上限。

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

确保预留可用性

创建预留 TPU 切片节点池（使用预留的 TPU 切片节点池）不需要任何 TPU 配额。但是，在创建节点池时，预留必须具有足够的可用或未使用的 TPU 芯片。

如需查看项目中存在哪些预留以及 TPU 预留中可用的 TPU 芯片数，请查看预留列表。

在 GKE 中预配 TPU 的选项

借助 GKE，您可以通过在工作负载清单中使用 Kubernetes nodeSelector 或通过创建具有 TPU 的 Standard 模式节点池，直接在各个工作负载中使用 TPU。

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

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

创建集群

在具有可用 TPU 的区域中创建 Standard 模式的 GKE 集群。

gcloud container clusters create CLUSTER_NAME \
  --location LOCATION \
  --cluster-version VERSION

替换以下内容：

• CLUSTER_NAME：新集群的名称。
• LOCATION：具有可用 TPU 容量的区域。
• VERSION：GKE 版本，必须支持您要使用的机器类型。请注意，默认 GKE 版本可能无法为您的目标 TPU 提供可用性。如需了解 TPU 机器类型可用的最低 GKE 版本，请参阅 GKE 中的 TPU 可用性。

创建节点池

您可以创建单主机或多主机 TPU 切片节点池。

创建单主机 TPU 切片节点池

您可以使用 Google Cloud CLI、Terraform 或 Google Cloud 控制台创建单主机 TPU 切片节点池。

gcloud container node-pools create NODE_POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONES \
    --machine-type=MACHINE_TYPE \
    [--sandbox=type=gvisor]

resource "google_container_node_pool" "NODE_POOL_RESOURCE_NAME" {
  provider           = google
  project            = PROJECT_ID
  cluster            = CLUSTER_NAME
  name               = POOL_NAME
  location           = CLUSTER_LOCATION
  node_locations     = [NODE_ZONES]

  node_config {
    machine_type = MACHINE_TYPE
    reservation_affinity {
      consume_reservation_type = "SPECIFIC_RESERVATION"
      key = "compute.googleapis.com/reservation-name"
      values = [RESERVATION_LABEL_VALUES]
    }
    spot = true
    flex_start = false
  }
}

创建多主机 TPU 切片节点池

您可以使用 Google Cloud CLI、Terraform 或 Google Cloud 控制台创建多主机 TPU 切片节点池。

gcloud container node-pools create POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME \
    --node-locations=NODE_ZONES \
    --machine-type=MACHINE_TYPE \
    --tpu-topology=TPU_TOPOLOGY \
    [--num-nodes=NUM_NODES] \
    [--spot \]
    [--flex-start \]
    [--enable-autoscaling \
      --max-nodes MAX_NODES]
    [--reservation-affinity=specific \
    --reservation=RESERVATION_NAME] \
    [--node-labels cloud.google.com/gke-nodepool-group-name=COLLECTION_NAME,cloud.google.com/gke-workload-type=HIGH_AVAILABILITY]
    [--placement-type=COMPACT]

GKE 如何处理容量问题

如果 GKE 由于可用 TPU 容量不足而无法创建 TPU 切片节点池，GKE 会返回一条错误消息，指示由于容量不足而无法创建 TPU 切片节点。

如果您要创建单主机 TPU 切片节点池，则错误消息将类似于以下内容：

2 nodes cannot be created due to lack of capacity. The missing nodes will be
created asynchronously once capacity is available. You can either wait for the
nodes to be up, or delete the node pool and try re-creating it again later.

如果您要创建多主机 TPU 切片节点池，则错误消息将类似于以下内容：

The nodes (managed by ...) cannot be created now due to lack of capacity. They
will be created asynchronously once capacity is available. You can either wait
for the nodes to be up, or delete the node pool and try re-creating it again
later.

TPU 预配请求可能会在队列中长时间保留，并在队列中时一直保持“正在预配”状态。

具备可用容量后，GKE 会创建尚未创建的其余节点。

如果您需要更快地需要容量，请考虑使用 Spot 虚拟机，但请注意，Spot 虚拟机消耗的配额与按需实例不同。

您可以通过删除 TPU 切片节点池来删除已加入队列的 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 记录。

   使用 Kuberay 部署 TPU 多主机资源时，GKE 提供可部署的 webhook 作为实验性 Terraform 模板的一部分，用于运行 Ray on GKE。如需了解如何使用 TPU 运行 Ray on GKE，请参阅实验性 TPU 用户指南。变更 webhook 会将这些环境变量注入请求 google.com/tpu 属性的 Ray 集群和多主机 cloud.google.com/gke-tpu-topology 节点选择器。

   • 在工作负载清单中，添加 Kubernetes 节点选择器，以确保 GKE 根据您定义的 TPU 机器类型和 TPU 拓扑调度 TPU 工作负载：

       nodeSelector:
         cloud.google.com/gke-tpu-accelerator: TPU_ACCELERATOR
         cloud.google.com/gke-tpu-topology: TPU_TOPOLOGY
       

     替换以下内容：

     • TPU_ACCELERATOR：TPU 加速器的名称。
     • TPU_TOPOLOGY：TPU 切片的物理拓扑。拓扑格式取决于 TPU 版本。如需了解详情，请参阅规划 GKE 中的 TPU。

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

以下部分举例说明了如何运行使用 TPU 执行简单计算的 Job。

示例 1：运行显示 TPU 切片节点池中可用 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.
           securityContext:
             privileged: true
           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
   

示例 2：运行显示 TPU 切片中可用 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.
       securityContext:
         privileged: true
       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
     

使用加速器（GPU 和 TPU）升级节点池

GKE 会自动升级 Standard 集群，包括节点池。如果您希望节点更快地升级到更高版本，也可以手动升级节点池。如需控制集群升级的方式，请使用发布渠道、维护窗口和排除项以及 发布顺序。

您还可以为节点池配置节点升级策略，例如超额配置升级、蓝绿升级或短期升级。通过配置这些策略，您可以确保节点池的升级方式能够在速度和环境中断之间实现最佳平衡。对于多主机 TPU 切片节点池，GKE 会在单一步骤中以原子方式重新创建整个节点池，而不是使用配置的节点升级策略。如需了解详情，请参阅 GKE 中的 TPU 相关术语中的“原子性”定义。

使用节点升级策略会暂时要求 GKE 预配更多资源，具体取决于配置。如果 Google Cloud对节点池的资源具有容量限制（例如，您尝试创建更多具有 GPU 或 TPU 的节点时出现资源可用性错误），请参阅在资源受限的环境中升级。

清理

为避免因本指南中使用的资源而导致您的 Google Cloud 账号产生费用，请考虑删除不再具有调度工作负载的 TPU 切片节点池。如果工作负载运行必须正常终止，请使用 kubectl drain 清理工作负载，然后再删除节点。

1. 删除 TPU 分片节点池：

   gcloud container node-pools delete POOL_NAME \
       --location=LOCATION \
       --cluster=CLUSTER_NAME
   

   替换以下内容：

   • POOL_NAME：节点池的名称。
   • CLUSTER_NAME：集群的名称。
   • LOCATION：集群的计算位置。

配置其他设置

以下各部分介绍了可应用于 TPU 工作负载的其他配置。

管理集合调度

在 TPU Trillium 中，您可以使用集合调度对 TPU 切片节点进行分组。对这些 TPU 切片节点进行分组后，您可以更轻松地调整副本数量，以满足工作负载需求。 Google Cloud 会控制软件更新，以确保集合中的切片充足，始终可用于处理流量。

TPU Trillium 支持为运行推理工作负载的单主机和多主机节点池提供集合调度。以下内容介绍了集合调度行为如何取决于您使用的 TPU 切片类型：

• 多主机 TPU 切片：GKE 会将多主机 TPU 切片分组以形成集合。每个 GKE 节点池都是此集合中的副本。如需定义集合，请创建多主机 TPU 切片并为集合分配唯一名称。如需向集合中添加更多 TPU 切片，请创建另一个具有相同集合名称和工作负载类型的多主机 TPU 切片节点池。
• 单主机 TPU 切片：GKE 会将整个单主机 TPU 切片节点池视为一个集合。如需向集合中添加更多 TPU 切片，您可以调整单主机 TPU 切片节点池的大小。

如需管理集合，请根据您使用的节点池类型执行以下任一操作。

管理多主机 TPU 切片节点池中的集合调度

您可以使用以下任务来管理多主机 TPU 切片节点池。

• 如需检查多主机 TPU 切片池是否属于集合，请运行以下命令：

  gcloud container node-pools describe NODE_POOL_NAME \
      --location LOCATION \
      --cluster CLUSTER_NAME \
      --format="json" | jq -r \
      '"nodepool-group-name: \(.config.labels["cloud.google.com/gke-nodepool-group-name"] // "")\ngke-workload-type: \(.config.labels["cloud.google.com/gke-workload-type"] // "")"'
  

  输出类似于以下内容：

  nodepool-group-name: <code><var>NODE_POOL_COLLECTION_NAME</var></code>
  gke-workload-type: HIGH_AVAILABILITY
  

  如果多主机 TPU 切片池属于集合，则输出具有以下标签：

  • cloud.google.com/gke-workload-type: HIGH_AVAILABILITY
  • cloud.google.com/gke-nodepool-group-name: <code><var>COLLECTION_NAME</var></code>

• 如需获取集群中的集合列表，请运行以下命令：

  #!/bin/bash
  
  # Replace with your cluster name, project, and location
  CLUSTER_NAME=CLUSTER_NAME
  PROJECT=PROJECT_ID
  LOCATION=LOCATION
  
  declare -A collection_names
  
  node_pools=$(gcloud container node-pools list --cluster "$CLUSTER_NAME" --project "$PROJECT" --location "$LOCATION" --format="value(name)")
  
  # Iterate over each node pool
  for pool in $node_pools; do
      # Describe the node pool and extract labels using jq
      collection_name=$(gcloud container node-pools describe "$pool" \
          --cluster "$CLUSTER_NAME" \
          --project "$PROJECT" \
          --location "$LOCATION" \
          --format="json" | jq -r '.config.labels["cloud.google.com/gke-nodepool-group-name"]')
  
      # Add the collection name to the associative array if it's not empty
      if [[ -n "$collection_name" ]]; then
          collection_names["$collection_name"]=1
      fi
  done
  
  # Print the unique node pool collection names
  echo "Unique cloud.google.com/gke-nodepool-group-name values:"
  for name in "${!collection_names[@]}"; do
      echo "$name"
  done
  

  输出类似于以下内容：

  Unique cloud.google.com/gke-nodepool-group-name values: {COLLECTION_NAME_1}, {COLLECTION_NAME_2}, {COLLECTION_NAME_3}
  

• 如需获取属于某个集合的节点池的列表，请运行以下命令：

  #!/bin/bash
  
  TARGET_COLLECTION_NAME=COLLECTION_NAME
  CLUSTER_NAME=CLUSTER_NAME
  PROJECT=PROJECT_ID
  LOCATION=LOCATION
  
  matching_node_pools=()
  
  # Get the list of all node pools in the cluster
  node_pools=$(gcloud container node-pools list --cluster "$CLUSTER_NAME" --project "$PROJECT" --location "$LOCATION" --format="value(name)")
  
  # Iterate over each node pool
  for pool in $node_pools; do
      # Get the value of the cloud.google.com/gke-nodepool-group-name label
      collection_name=$(gcloud container node-pools describe "$pool" \
          --cluster "$CLUSTER_NAME" \
          --project "$PROJECT" \
          --location "$LOCATION" \
          --format="json" | jq -r '.config.labels["cloud.google.com/gke-nodepool-group-name"]')
  
      # Check if the group name matches the target value
      if [[ "$collection_name" == "$TARGET_COLLECTION_NAME" ]]; then
          matching_node_pools+=("$pool")
      fi
  done
  
  # Print the list of matching node pools
  echo "Node pools with collection name '$TARGET_COLLECTION_NAME':"
  for pool in "${matching_node_pools[@]}"; do
      echo "$pool"
  done
  

  输出类似于以下内容：

  Node pools with collection name 'COLLECTION_NAME':
  {NODE_POOL_NAME_1}
  {NODE_POOL_NAME_2}
  {NODE_POOL_NAME_3}
  

• 如需扩容集合，请创建另一个多主机 TPU 切片节点池，并添加 cloud.google.com/gke-workload-type 和 cloud.google.com/gke-nodepool-group-name。在 cloud.google.com/gke-nodepool-group-name 中使用相同的集合名称，并运行相同的工作负载类型。如果在集群上启用了节点自动预配，GKE 会根据工作负载需求自动创建节点池。

• 如需缩容集合，请删除节点池。

• 如需删除集合，请移除所有关联的节点池。您可以删除节点池或删除集群。删除集群会移除其中的所有集合。

管理单主机 TPU 切片节点池中的集合调度

您可以使用以下任务来管理单主机 TPU 切片节点池。

• 如需检查单主机 TPU 切片池是否启用了集合调度，请运行以下命令：

  gcloud container node-pools describe NODE_POOL_NAME \
      --cluster CLUSTER_NAME \
      --project PROJECT_NAME \
      --location LOCATION \
      --format="json" | jq -r '.config.labels["cloud.google.com/gke-workload-type"]'
  

  输出类似于以下内容：

  gke-workload-type: HIGH_AVAILABILITY
  

  如果单主机 TPU 切片池是集合的一部分，则输出具有 cloud.google.com/gke-workload-type: HIGH_AVAILABILITY 标签。

• 如需扩容集合，请使用节点自动预配手动或自动调整节点池的大小。

• 如需缩容集合，请删除节点池。

• 如需删除集合，请移除所有关联的节点池。您可以删除节点池或删除集群。删除集群会移除其中的所有集合。

使用多切片

您可以在多切片中汇总较小的时间片，以处理较大训练工作负载。如需了解详情，请参阅 GKE 中的多切片 TPU。

迁移 TPU 预留

如果您有现有的 TPU 预留，则必须先将 TPU 预留迁移到新的基于 Compute Engine 的预留系统。您还可以创建基于 Compute Engine 的预留系统，而无需迁移。如需了解如何迁移 TPU 预留，请参阅 TPU 预留。

启用日志记录

GKE 节点（包括 TPU 虚拟机）上运行的容器发出的日志会由 GKE 日志记录代理收集，发送到 Logging，并且显示在 Logging 中。

使用 GKE 节点自动预配功能

您可以将 GKE 配置为自动创建和删除节点池，以满足 TPU 工作负载的资源需求。如需了解详情，请参阅配置 Cloud TPU。

使用自定义计算类预配 TPU

您还可以将 GKE 配置为在创建新节点的扩缩操作期间使用自定义计算类请求 TPU。

您可以在自定义计算类规范中指定 TPU 配置选项。当 GKE 工作负载使用该自定义计算类时，GKE 会尝试预配在扩容时使用指定配置的 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 节点。

导致 TPU 切片节点健康状况不佳的条件包括：

• 任何具有通用节点条件的 TPU 切片节点。
• 任何不可分配 TPU 数量大于零的 TPU 切片节点。
• TPU 切片中的任何已停止（由于抢占）或终止的虚拟机实例。
• 节点维护：如果多主机 TPU 切片节点池中的任何 TPU 切片节点因主机维护而关闭，GKE 会重新创建整个 TPU 切片节点池。

您可以在操作历史记录中查看修复状态（包括失败原因）。 如果失败是由配额不足引起的，请与您的Google Cloud 客户代表联系，以增加相应的配额。

为 TPU 切片节点配置正常终止

在控制平面运行 1.29.1-gke.1425000 或更高版本的 GKE 集群中，TPU 切片节点支持 SIGTERM 信号，用于提醒节点即将关停。在 TPU 节点中，即将关停通知的可配置时间最长为五分钟。

如需将 GKE 配置为在此通知时间范围内正常终止工作负载，请按照管理 GKE 节点中断情况（适用于 GPU 和 TPU）中的步骤操作。

在不使用特权模式的情况下运行容器

在 GKE 1.28 版或更高版本的节点中运行的容器不需要启用特权模式来访问 TPU。GKE 1.28 版及更低版本中的节点需要特权模式。

如果您的 TPU 切片节点运行的版本低于 1.28，请参阅以下部分：

在 TPU 切片中的虚拟机上运行的容器需要访问更高的锁定内存限制，以便驱动程序可以通过直接内存访问 (DMA) 与 TPU 芯片进行通信。如需实现这一目标，您必须配置更高的 ulimit。如果您想缩小容器的权限范围，请完成以下步骤：

1. 修改 securityContext 以包含以下字段：

   securityContext:
     capabilities:
       add: ["SYS_RESOURCE"]
   

2. 在将工作负载设置为使用 TPU 资源之前，请在容器内运行以下命令以增加 ulimit：

   ulimit -l 68719476736
   

对于 TPU v5e，1.27.4-gke.900 版及更高版本的集群可以在无特权模式下运行容器。

可观测性和指标

信息中心

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 系统指标。

已知问题

• 在新的 TPU 分片节点报告可用 TPU 之前，集群自动扩缩器可能会错误地计算这些节点的容量。然后，集群自动扩缩器可能会执行额外的扩容，因此创建超出需要的节点。集群自动扩缩器会在常规缩容操作后缩容额外的节点（如果不需要）。
• 集群自动扩缩器会取消对等待状态超过 10 小时的 TPU 切片节点池的扩容操作。集群自动扩缩器稍后将重试此类扩容操作。对于不使用预留的客户，此行为可能会降低 TPU 可获取性。
• 容忍 TPU 污点的非 TPU 工作负载可能会阻止缩容在排空 TPU 切片节点池期间重新创建的节点池。
• 内存带宽利用率指标不适用于 v5e TPU。

后续步骤

• 详细了解如何使用 TPU 在 GKE 上设置 Ray
• 利用 GKE 在 Cloud TPU 上构建大规模的机器学习服务
• 在 TPU 上使用 KubeRay 提供大语言模型
• 排查 GKE 中的 TPU 问题
• 了解如何使用 GKE Sandbox 对 TPU 工作负载进行沙盒处理