GKE 中断耐受 AI 工作负载设计实践

你已经把一切都做对了。你将庞大的模型训练任务容器化，部署到 Google Kubernetes Engine（GKE），并巧妙地将其调度到 Spot VM 节点池，从而节省高达 90% 的计算成本。

一切都完美运行了 38 个小时。随后，一位高优先级的按需客户需要计算容量，Google Cloud 回收了底层的 Spot VM，你的节点瞬间消失。

无论你是使用可抢占的 Spot VM 来节省成本，还是利用 Dynamic Workload Scheduler（DWS）排队等待稀缺的 GPU 资源，你都是在临时计算资源之上构建系统。硬件最终一定会被收回。要想在未承诺容量上成功运行关键的 AI 工作负载，你的应用架构就必须把故障视为必然事件。

下面是一份在 GKE 上构建可中断工作负载的实用指南。

当 Google Cloud 回收 Spot VM 时，并不会立刻粗暴地切断电源。它会向底层节点发送 ACPI 信号，启动关机流程。Kubernetes 会拦截这个信号，并将其转换为 SIGTERM 信号，直接发送给正在运行的容器。

从收到 SIGTERM 到致命的 SIGKILL 之间，你有一段宽限期（对于非系统 Pod，最长为 15 秒）。

你的应用必须显式监听这个信号。捕获信号后，代码应立即停止接收新的 batch，完成当前循环，将内存中的所有数据刷新到磁盘，然后以状态码 0（成功）退出。

下面是一个使用 Python 捕获该信号的简单示例：

import signal
import sys
import time

def handle_sigterm(signum, frame):
    print("Received SIGTERM. Initiating graceful shutdown...")
    # 1. Stop processing new data
    # 2. Flush memory to persistent storage
    # 3. Save final checkpoint
    print("State saved. Exiting cleanly.")
    sys.exit(0)

# Register the signal handler
signal.signal(signal.SIGTERM, handle_sigterm)

# Your main training loop
print("Starting training loop...")
while True:
    # Train model...
    time.sleep(1)

2. 将 checkpoint 外部化

如果容器终止，其本地文件系统中的所有内容也会随之消失。为了从中断中恢复，你必须定期将进度——包括模型权重、optimizer 状态、epoch 计数器等——保存到外部存储位置。

在 Google Cloud 上，Cloud Storage（GCS）是解决这一问题的常见方案。

频繁保存：确定合理的 checkpoint 间隔，在丢失已完成工作的成本与写入存储带来的开销之间取得平衡。通常可以在每个 epoch 或每几千个 step 后保存一次，但具体频率取决于你的实际需求。

保持同区域：确保 GCS bucket 与 GKE 集群位于同一区域，例如 us-central1，从而尽可能降低延迟，并避免产生出站数据传输费用。

恢复，而不是重启：容器的启动脚本首先应该检查对应的 GCS bucket。如果 bucket 中存在 checkpoint，就加载它，并从完全相同的 step 继续运行。

3. 围绕幂等性进行设计

“幂等性”（Idempotency）是一种听起来很高级的说法，意思是：同一件事执行两次，得到的结果与只执行一次相同。

假设有一个 batch inference 任务，它读取一张图片、处理图片，然后将结果写入数据库。如果 Pod 在写入数据库后的几毫秒内被抢占，但还没来得及把任务标记为已完成，那么重新调度的 Pod 很可能会再次处理这张图片。

如果数据库不加判断地插入新行，你就会得到并非预期的重复数据。

要构建幂等的 pipeline：

根据唯一标识符（例如图片 ID），在数据库中使用 UPSERT（更新或插入）操作。

在消耗昂贵的 GPU 计算周期处理数据之前，先检查对应记录是否已经存在。

4. 为 batch processing 解耦工作队列

如果你要对数千个文件运行大规模的 batch processing 或 inference 任务，不要编写一个遍历静态 CSV 列表的单体 Python 脚本。如果节点在处理到第 5,000 行时终止，管理从哪里重新开始的状态将会是一场噩梦。

应该改为解耦工作负载：

发布任务：将数据集拆分成一条条独立的消息，并将它们推送到 Pub/Sub 等消息代理中。

拉取任务：让 Spot VM 上的 worker Pod 每次从队列中拉取一条消息，或者拉取一小批消息，例如每次 10 条。

确认完成：只有在结果已安全存储之后，才向 Pub/Sub 返回“ACK”（acknowledgment，确认）。

如果 Spot 节点在 inference 过程中被抢占，worker 会在发送 ACK 之前终止。经过短暂的超时后，Pub/Sub 会自动让这条特定消息重新变为可用状态。另一个仍在运行的 worker Pod 会无缝接手并处理它。数据不会丢失，也不需要人工干预。

在 Spot VM 这类临时计算资源上运行任务，不只是一项基础设施选择，更是一项设计选择。通过处理终止信号、积极地将 checkpoint 保存到 GCS、确保操作具备幂等性，以及解耦工作队列，你既能大幅节省成本，又能使用稀缺的 GPU 资源池，同时不牺牲可靠性。

如需采取进一步措施，你可以考虑屏蔽此人和/或举报滥用行为。