使用 Cloud TPU Pod 管理维护事件

概览

TPU 节点和 TPU 虚拟机是挂接了 TPU 硬件的 Compute Engine 虚拟机实例。Compute Engine 虚拟机容易受到 Compute Engine 虚拟机维护事件的影响。每个 TPU 都连接到一个 Compute Engine 虚拟机，因此使用更多的 TPU（例如，在 TPU Pod 中）会增加您的某个虚拟机遇到维护事件的可能性。

本文档介绍了在 Cloud TPU 上处理长时间运行训练作业所需的维护事件方法。

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

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

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

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

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

检测维护事件

您可以使用以下 gcloud describe 命令检测 TPU 上是否发生维护事件以及何时发生：

TPU 虚拟机

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

TPU 节点

$ gcloud compute tpus 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 上的维护事件的历史日志。

在 Google Cloud 控制台导航菜单中，点击 Compute Engine > 虚拟机实例，然后搜索，例如：

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

在搜索时间范围内，系统会显示任何中断和修复的 TPU 工作器。日志将在 protoPayload.metadata.terminateReason 中显示事件的日期和时间、事件类型；对于“terminate”事件，则会显示终止原因。

处理维护事件

您可以使用多种方法来减轻维护事件中断。

  1. 定期保存检查点

    在理想情况下，当发生“中断事件”时，训练会从最新的检查点继续。

  2. 训练脚本重试

    训练脚本可能会因“中断事件”而停止。您可以使用 bash 脚本不断重试训练脚本，直到训练完成。每次重试都应从最新的检查点继续执行，因此重试脚本应始终与检查点结合使用。

    可正式投入使用的训练流水线应使用 Google Kubernetes Engine (GKE) 等资源管理系统。如需详细了解如何将 Google Kubernetes Engine 与 TPU 虚拟机架构配合使用，请参阅部署 TPU 工作负载。如需详细了解如何将 Google Kubernetes Engine 与 TPU 节点架构配合使用，请参阅在 Google Kubernetes Engine 上运行 TPU 应用。否则，您可以实现 bash 脚本以不断重试训练脚本，直到完成为止。例如：

    使用 TPU 节点：

    （从您的虚拟机）bash while ! python3 [training command]; do sleep 1; done

    使用 TPU 虚拟机：

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

    （请注意，您需要从 Cloud Shell 或终端运行 TPU 虚拟机命令，而不是从 TPU 虚拟机运行）。

  3. 检测 TPU 故障并从中恢复

    如果 TPU 无法从维护事件恢复，您可以使用恢复脚本检测 TPU 状态，删除并重新创建 TPU。如需查看此脚本的示例，请点击此处。 如需详细了解如何手动删除并重新创建 TPU，请参阅管理 TPU。

    创建或重新创建 TPU 虚拟机时，您可以使用 --metadata startup-script 参数指定启动脚本。每当创建 TPU 虚拟机时，都会运行启动脚本。如需了解详情，请参阅运行标准安装脚本。