Vertex AI上Python进程因OOM被终止（数据集仅3.5GB）

本地32GB内存机器可正常运行的Python预测任务，迁移至谷歌云Vertex AI的c3-highmem-22实例（176GB内存）后，却因内存不足（OOM）终止进程。任务涉及从BigQuery读取3.5GB数据集到内存，使用本地训练的LSTM模型进行预测。

错误日志 #

2024-04-10 15:27:59.992 BST
[  672.752608] oom-kill:constraint=CONSTRAINT_NONE,nodemask=(null),cpuset=jupyter.service,mems_allowed=0,global_oom,task_memcg=/system.slice/jupyter.service,task=python,pid=3740,uid=1000

2024-04-10 15:28:00.043 BST
[  672.769191] Out of memory: Killed process 3740 (python) total-vm:211774876kB, anon-rss:179645328kB, file-rss:0kB, shmem-rss:0kB, UID:1000 pgtables:354232kB oom_score_adj:0

2024-04-10 15:28:00.043 BST
Apr 10 14:27:59 instance-20240410-151420 kernel: [  672.119856] python invoked oom-killer: gfp_mask=0x100dca(GFP_HIGHUSER_MOVABLE|__GFP_ZERO), order=0, oom_score_adj=0

2024-04-10 15:28:00.093 BST
Apr 10 14:28:00 instance-20240410-151420 kernel: [  672.130423] CPU: 8 PID: 4164 Comm: python Not tainted 5.10.0-28-cloud-amd64 #1 Debian 5.10.209-2

2024-04-10 15:28:00.093 BST
Apr 10 14:28:00 instance-20240410-151420 kernel: [  672.139309] Hardware name: Google Google Compute Engine/Google Compute Engine, BIOS Google 03/29/2024

触发问题的代码片段 #

# Start timing
start_time_predictions = time.time()

# Prepare data as a TensorFlow tensor
data_tensor = tf.constant(data_3d_test, dtype=tf.float32)

# Use the serving signature for prediction.
predictions = model.signatures['serving_default'](lstm_input=data_tensor)['output_0']

# Convert predictions to numpy
predictions_np = predictions.numpy()

# Ensure predictions are shaped as expected
assert predictions_np.shape[1] == 26, "Each title should have 26 predictions"

# End timing
end_time_predictions = time.time()

# Calculate and print the prediction time
predictions_time = end_time_predictions - start_time_predictions
# `format_time` is a function to format the time string
formatted_time = format_time(predictions_time)
print(f"Time taken to generate predictions: {formatted_time}")

已尝试操作与预期 #

• 已提升实例内存至176GB，问题依旧
• 预期模型能像本地环境一样正常生成预测，迁移至Vertex AI是为了实现每周自动运行及云端管理

从日志看，Python进程占用内存接近180GB，远超数据集本身的3.5GB，核心问题在于数据转换、模型预测阶段的内存膨胀或不合理分配，以下是针对性优化方案：

1. 分批处理数据，避免全量加载张量 #

直接用tf.constant会将整个数据集复制为TensorFlow张量，导致内存翻倍。改用tf.data.Dataset分批处理，仅将当前批次数据加载到内存：

import numpy as np
import tensorflow as tf

# 设置合理的批次大小，根据内存调整
batch_size = 1024
# 构建分批数据集
dataset = tf.data.Dataset.from_tensor_slices(data_3d_test).batch(batch_size)

predictions_np = []
start_time_predictions = time.time()

# 逐批次预测
for batch in dataset:
    batch_pred = model.signatures['serving_default'](lstm_input=batch)['output_0']
    predictions_np.append(batch_pred.numpy())

# 合并所有批次结果
predictions_np = np.concatenate(predictions_np, axis=0)
end_time_predictions = time.time()

# 后续断言和计时逻辑保持不变
assert predictions_np.shape[1] == 26, "Each title should have 26 predictions"
formatted_time = format_time(end_time_predictions - start_time_predictions)
print(f"Time taken to generate predictions: {formatted_time}")

2. 调整TensorFlow内存分配策略 #

默认情况下，TensorFlow会尝试占用全部可用内存，在云端环境可能导致内存耗尽。启用按需内存增长或限制内存占用比例：

import tensorflow as tf

# 启用CPU内存按需增长
physical_devices = tf.config.list_physical_devices('CPU')
if physical_devices:
    tf.config.experimental.set_memory_growth(physical_devices[0], True)

# （可选）限制内存占用比例，比如使用80%的实例内存
# tf.config.set_logical_device_configuration(
#     physical_devices[0],
#     [tf.config.LogicalDeviceConfiguration(memory_limit=int(176*1024*0.8))]
# )

3. 检查并优化数据类型 #

确认数据集的数据类型是否合理，比如如果本地用float32而云端意外使用float64，会导致内存翻倍。强制转换为低内存占用的类型：

# 将数据集转换为float32，减少内存占用
data_3d_test = data_3d_test.astype(np.float32)

4. 主动清理内存，释放中间变量 #

在数据转换和预测后，及时删除不需要的变量并触发垃圾回收：

# 转换为张量后删除原numpy数组
del data_3d_test
import gc
gc.collect()

# 预测完成后删除张量变量
del data_tensor, predictions
gc.collect()

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