在 GKE 中以全容量模式部署 TPU

标准

本文档介绍了如何在 GKE 中使用 TPU 全容量模式功能来部署和管理工作负载。全容量模式预留可让您更好地控制 TPU 资源，从而在预留容量内以精细的控制方式放置工作负载。

本文档适用于希望使用 Kubernetes 容器编排功能来精细控制 TPU 部署的机器学习 (ML) 工程师、平台管理员和运维人员。

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

什么是 TPU 全容量模式？

借助 TPU Cluster Director 启用的 TPU 全容量模式，您可以完全控制预留的 TPU 容量。TPU Cluster Director 是一项管理服务，可让您基于预留控制 TPU。

与之前的托管式模式不同，在托管式模式下， Google Cloud 会预留一部分容量来处理硬件故障，而在全容量模式下，您可以访问预留的全部 TPU 资源。此模式可让您全面了解硬件状态，但同时也会将管理节点故障和计划内维护的责任转移给您。

如需详细了解全容量模式的主要功能，请参阅 TPU Cluster Director 概览中的全容量模式。

与 GKE 中的“所有容量”模式相关的术语

下表列出了 Ironwood (TPU7x) 版本中按块、子块和立方体大小定义的术语和等效项。立方体是一种由互连的 TPU 芯片组成的 4x4x4 拓扑，仅适用于 3 元组 ({A}x{B}x{C}) 中的拓扑。

TPU 资源	核心数	芯片数	主机数	立方体
1 个芯片	2	1	-	-
1 位主持人	8	4	1	-
1 个子区块	128	64	16	1
一个块包含 144 个子块	18432	9216	2304	144

如需详细了解块中允许的拓扑，请参阅选择拓扑。

准备工作

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

启用 Google Kubernetes Engine API。

启用 Google Kubernetes Engine API

如果您要使用 Google Cloud CLI 执行此任务，请安装并初始化 gcloud CLI。如果您之前安装了 gcloud CLI，请通过运行 gcloud components update 命令来获取最新版本。较早版本的 gcloud CLI 可能不支持运行本文档中的命令。
注意：对于现有 gcloud CLI 安装，请务必设置 compute/region 属性。如果您主要使用可用区级集群，请改为设置 compute/zone。通过设置默认位置，您可以避免 gcloud CLI 中出现如下错误：One of [--zone, --region] must be supplied: Please specify location。如果集群的位置与您设置的默认位置不同，您可能需要在某些命令中指定位置。

确保您已有 1.34.0-gke.2201000 版或更高版本的 Standard 集群。如需创建新集群，请参阅创建区域级集群。
确保您在要使用的区域中拥有足够的 TPU 配额。
安装 JobSet v0.2.3 或更高版本。

限制

GKE 中的 TPU 全容量模式仅支持 Ironwood (TPU7x) 版本。

在 GKE 中使用 TPU 全容量模式

本部分介绍了在 GKE 中使用 TPU 全容量模式的工作流程。

在全容量模式预留中创建节点池

GKE 中的“所有容量”模式允许您通过以下方式创建节点池：

一种节点池，其中 GKE 会在您的“所有容量”TPU 预留中选择块或子块。
一种节点池，用于在 TPU 全容量模式预留中指定特定区块或子区块。

GKE 会在 TPU 全容量预留中选择块或子块

在此模式下，GKE 会在 TPU 全容量预留中选择节点池的放置位置。此流程与使用其他 TPU 预配选项（例如按需或 Spot 虚拟机）创建节点池类似。

如需创建节点池，请将 gcloud container node-pools create 命令与 --reservation 标志结合使用。将 TPU 预留的完整资源名称指定为 --reservation 标志的值。

如需查看节点池创建命令的示例，请参阅手动创建节点池。

定位预留中的块或子块

借助 TPU 全容量模式，您可以将 TPU 预留中的特定块或子块指定为并行工作负载的目标。此功能适用于需要 TPU 芯片紧密相邻以最大限度缩短延迟时间的工作负载。

查看可用的块、子块和主机配置。完成查看“所有容量”模式预留的拓扑和健康状态文档中的步骤。
创建工作负载政策：

注意：您无需为每个节点池创建新的工作负载政策。每个项目、每个区域和每个拓扑的工作负载政策都是唯一的。如果您已有符合要求的工作负载政策，可以跳过此步骤，并在下一步中重复使用现有政策。如需查看工作负载政策列表，请使用 gcloud compute resource-policies list --filter="region:REGION" 命令。
```
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) 拓扑，请参阅规划 GKE 中的 TPU。
- PROJECT_ID：您的 Google Cloud 项目 ID。
- REGION：工作负载政策的区域。工作负载政策是一种区域级资源，可在共享相同拓扑的节点池中重复使用。
如需创建节点池并指定预留的特定块或子块，请使用 --reservation 标志指定预留中目标块或子块的完整资源名称。

如需定位预留中的特定块，请使用以下命令：
```
gcloud container node-pools create NODE_POOL_NAME \
    --cluster=CLUSTER_NAME \
    --machine-type=tpu7x-standard-4t \
    --placement-policy=WORKLOAD_POLICY_NAME \
    --zone=ZONE \
    --reservation=project/PROJECT/reservation/RESERVATION_NAME/reservationBlocks/BLOCK_NAME
```
如需定位到块中的特定子块，请使用以下命令：
```
gcloud container node-pools create NODE_POOL_NAME \
    --cluster=CLUSTER_NAME \
    --machine-type=tpu7x-standard-4t \
    --placement-policy=WORKLOAD_POLICY_NAME \
    --zone=ZONE \
    --reservation=project/PROJECT/reservation/RESERVATION_NAME/reservationBlocks/BLOCK_NAME/reservationSubBlocks/SUB_BLOCK_NAME
```
替换以下内容：
- NODE_POOL_NAME：新节点池的名称。
- CLUSTER_NAME：GKE 集群的名称。
- WORKLOAD_POLICY_NAME：您创建的工作负载政策的名称。
- ZONE：节点池的可用区，例如 us-central1-a。
- PROJECT：您的 Google Cloud 项目 ID。
- RESERVATION_NAME：TPU 预留的名称。
- BLOCK_NAME：预留中的特定块。
- SUB_BLOCK_NAME：预留中的特定子块。
在上述命令中，您

安排工作负载

创建包含以“所有容量”模式运行的 TPU 虚拟机的节点池后，您可以像部署任何其他 TPU 节点池一样部署工作负载。对于 TPU 全容量模式，在调度工作负载方面，与使用标准 SLO 支持的预留的节点池相比，没有其他区别。

如需详细了解使用 TPU 的工作负载并查看相关示例，请参阅在 TPU 切片节点上运行工作负载和运行多切片工作负载。

管理节点故障

TPU 全容量模式预留是无保留预留。无保留是指您会获得完整的 TPU 容量，包括Google Cloud 通常在受管容量模式下保留用于故障转移的部分。

在 TPU 全容量模式下，如果虚拟机因硬件故障等问题而发生故障，系统会 Google Cloud 尝试在同一主机上恢复虚拟机（就地修复）。因此，我们建议您保持备用容量，以便在基础架构发生故障时重新安排工作负载。

节点故障和恢复

当全容量模式预留中的节点发生故障时，会发生以下事件：

Google Cloud会为失败的 Compute Engine 虚拟机实例启动修复事件。此过程会尝试将虚拟机恢复为 RUNNING 状态，并将 GKE 节点状态恢复为 READY。
TPU 虚拟机进入修复状态，在该节点上运行的任何工作负载可能会失败，具体取决于其故障切换政策。即使节点池中的一个或多个虚拟机出现故障，节点池的状态也不会更改为 ERROR。

如需监控节点的健康状况，请按以下步骤操作：

列出节点池中的节点：
```
kubectl get nodes
```
失败的节点的状态为 NotReady。
监控 Compute Engine 节点的状态：

对于全容量模式下的 TPU 虚拟机，请使用 gcloud compute instances describe。此命令还提供拓扑的物理状态，以查找主机、子块和块详细信息。
```
gcloud compute instances describe VM_NAME \
        --format="table[box,title=VM-Position](resourceStatus.physical_host_topology:label=location)" \
        --zone=ZONE
```
替换以下内容：
- VM_NAME：TPU 虚拟机实例的名称。
- ZONE：虚拟机的可用区，例如 us-central1-a。
如需详细了解如何检索全容量模式容量的拓扑和健康信息，请参阅查看全容量模式预留的拓扑和健康状态。

管理维护

为了应对潜在的中断并帮助确保工作负载保持弹性，您可以使用 GKE 维护政策来管理各个节点的维护。如需了解详情，请参阅维护窗口和维护排除项。

GKE 不支持在全容量模式预留中对 TPU 虚拟机进行分组维护。如需在子块、块或预留级层执行分组维护，请使用 Compute Engine API。如需了解详情，请参阅使用全容量模式的 TPU 管理维护事件。

清理

为避免您的 Google Cloud 账号产生不必要的费用，请删除不再具有调度工作负载的 TPU 节点池。如果必须正常终止正在运行的工作负载，请使用 kubectl drain 命令清理工作负载，然后再删除节点池。

删除 TPU 节点池：
```
gcloud container node-pools delete NODE_POOL_NAME \
    --location=LOCATION \
    --cluster=CLUSTER_NAME
```
替换以下内容：
- NODE_POOL_NAME：节点池的名称。
- CLUSTER_NAME：集群的名称。
- LOCATION：集群的计算位置。