准备维护事件

TPU 虚拟机是挂接了 TPU 硬件的 Compute Engine 虚拟机实例。Compute Engine 虚拟机受 Compute Engine 虚拟机维护事件影响。每个 TPU 都连接到 Compute Engine 虚拟机，因此使用更多 TPU（例如在 TPU 切片中）会提高某个虚拟机遇到维护事件的可能性。

本文档介绍了在 Cloud TPU 上处理长时间运行的训练作业的维护事件的方法。如需了解如何在 Google Kubernetes Engine (GKE) 中处理 TPU 的维护事件，请参阅管理 GPU 和 TPU 的 GKE 节点中断。

查看即将进行的维护通知

您可以查看即将进行的主机维护事件的通知。通过监控实例即将开始的维护窗口，您可以主动准备工作负载以处理即将进行的维护，同时将中断降至最低。如需了解详情，请参阅查看维护通知。

使用检查点从维护事件中快速恢复

检查点是维护事件短时间恢复的关键，应经常保存。最好是大约每小时保存一次检查点。检查点通常不足以由于维护事件或其他训练中断而丢失大量训练进度。

检查点通常是指训练中使用的所有已保存的参数（例如模型权重）。保存检查点所需的时间介于数秒到数分钟之间。

虽然 TPU 可以自动从大多数维护事件中恢复，并且训练作业会继续进行，无需人工干预，但可能存在极端情况，即作业不会重启并自动继续。发生这种情况时，您需要删除并重新创建 TPU 资源，然后从已保存的检查点重启训练作业。如需了解如何检测自动恢复故障并从中恢复，请参阅检测 TPU 故障并从中恢复。

保存和加载各个机器学习框架的检查点的机制各不相同。受支持的 Cloud TPU 模型通常内置了检查点。如需详细了解检查点，请参阅 TensorFlow 2.x、PyTorch 或 JAX/flax。

使用自动检查点

您可以使用自动检查点功能来保留训练进度，方法是将代码配置为在发生维护事件时保存非预定检查点。如需详细了解自动检查点，请参阅 Cloud TPU 自动检查点。

重试训练脚本

训练脚本可能会因中断事件而停止。您可以使用 bash 脚本不断重试训练脚本，直到训练完成。例如：

while ! gcloud compute tpus tpu-vm ssh ${TPU_NAME} --command "python3 TRAINING_COMMAND"; do sleep 1; done

每次重试都应从最新的检查点继续执行，因此您应始终将重试脚本与检查点结合使用。

可用于生产用途的训练流水线应使用资源管理系统，例如 Google Kubernetes Engine (GKE)。如需详细了解如何将 Google Kubernetes Engine 与 TPU 虚拟机搭配使用，请参阅部署 TPU 工作负载。

检测 TPU 故障并从中恢复

如果 TPU 没有从维护事件中恢复，则您可以使用恢复脚本检测 TPU 状态并删除和重新创建 TPU。如需查看恢复脚本示例，请参阅 retry.sh。如果运行训练脚本的进程崩溃，您可以修改恢复脚本以重试运行训练脚本。

如需了解如何手动删除和重新创建 TPU，请参阅管理 TPU 资源。

使用集合调度

Cloud TPU 具有集合调度的概念，提供两种类型的集合，客户可以使用这些集合来支持训练或服务和推理工作负载。当您使用此功能部署 Cloud TPU 实例时， Google Cloud 会应用最适合应用的唯一维护时间表。您可以预期每种集合类型具有以下行为：

  • 训练（默认）：此集合类型对典型的训练工作负载很有用，在这种工作负载中，您需要尽可能缩短所有实例的停机时间，并尽可能减少意外中断，以便在维护事件期间快速恢复服务。训练集合类型可为一组实例并行调度和执行维护事件。

  • 服务（可使用 --workload-type=AVAILABILITY_OPTIMIZED）：此集合类型对大多数服务或推理工作负载都很有用，在这种工作负载中，您需要在部分实例（副本）上尽可能缩短停机时间，以确保服务连续性，即使在维护事件期间也是如此。服务集合类型可为一组实例分阶段安排和执行维护事件。仅 TPU v6e 支持指定服务集合。

如需详细了解集合调度，请参阅为推理工作负载调度 TPU 集合。

检测维护事件

您可以使用以下 gcloud compute tpus tpu-vm describe 命令检测 TPU 上是否发生了维护事件以及发生时间：

$ gcloud compute tpus tpu-vm describe tpu-name --zone=zone  | grep 'health'

此命令的输出显示 TPU 的当前状态以及最近一次维护事件的说明。输出应如下所示：

health: HEALTHY
healthDescription: The TPU had a maintenance event at 2022-01-26T03:44:36.265703305Z

查看维护事件日志

您可以在系统事件审核日志中查看有关 TPU 的维护事件的历史日志。

  1. 在 Google Cloud 控制台导航菜单中，前往 Logs Explorer 页面：

    转到日志浏览器

  2. 使用以下搜索查询查看已终止或已重启的所有 TPU：

    "tpu.nodes.terminate" OR "tpu.nodes.restart"

    结果会显示搜索时间范围内 TPU 工作器的所有中断和修复日志。日志包括：

    • 事件的日期和时间
    • 事件的类型
    • 对于“终止”事件，为 protoPayload.metadata.terminateReason 字段中的终止原因

后续步骤