GKE上TPU出现内存泄漏导致OOM或Unavailable: Socket Closed错误如何解决

问题根因 #

你遇到的是TensorFlow 2.3版本TPUStrategy的已知未修复缺陷：循环创建、训练多个模型时，TPU的高带宽内存(HBM)中残留的旧计算图、缓存函数句柄、废弃张量不会被CPU侧的tf.keras.backend.clear_session()、del model等操作释放，内存持续占满后触发OOM，进而导致TPU进程重启抛出Socket Closed类连接错误。 #

可行解决方案 #

1. 优先升级TensorFlow版本 #

该内存泄漏问题已在TensorFlow 2.6及后续版本修复，直接使用适配Cloud TPU的TF2.6+官方运行镜像，无需大幅修改代码即可解决绝大多数此类泄漏问题。

2. 手动重置TPU状态（适配无法升级TF版本的场景） #

如果受环境限制只能使用TF2.3，每次完成单个模型训练、且已将需要的指标/参数保存到CPU侧存储后，调用以下接口硬重置TPU：

# 重置TPU内存，清空所有残留资源
tf.tpu.experimental.initialize_tpu_system(tpu)
# 重置后重新创建TPUStrategy，不要复用之前的实例
tpu_strategy = tf.distribute.TPUStrategy(tpu)

该操作会清空TPU上的所有缓存数据，完全释放占用的HBM内存，从根源避免残留堆积。

3. 优化现有代码结构 #

你当前的代码存在两处会加剧泄漏的写法，可直接修改：

• 不要在create_network函数内重复创建TPUStrategy实例，直接复用全局初始化的tpu_strategy即可，重复初始化会残留多余的连接句柄
• 推理步骤y_pred = model(x_test)会生成额外的TF函数缓存，建议替换为model.evaluate(x_test, y_test)，或用strategy.run封装推理逻辑，不要在CPU侧持有TPU返回的张量的强引用

4. 大规模超参搜索场景适配方案 #

因为你使用的是GKE调度的可抢占式TPU，可调整调度逻辑：每个超参数实验单独申请一个TPU节点，实验结束后直接释放TPU节点，完全避免单节点多实验的资源残留问题，调度开销低且稳定性更高。 #

内容的提问来源于stack exchange，提问作者Lexi31415