Python多进程处理数据内存耗尽问题排查与优化求助

先揪出内存泄漏的核心原因 #

• Queue的隐形开销：就算限制了队列大小，要是主进程没及时取结果，或者worker崩溃导致队列堆积数据，内存会持续上涨。而且multiprocessing的Queue传输大对象时依赖pickle序列化，过程中会生成临时大对象，叠加后内存占用远超理论值。
• 子进程内存未及时释放：如果worker处理完任务未正常退出，或者反复创建新进程而非复用，闲置进程会一直占用内存。
• 主进程一次性加载全量数据：2000组1×4000的数组看似仅几十MB，但主进程一次性加载所有数据后，加上附加元数据、缓存，再结合写时复制机制（修改数据后会占用真实内存），总内存占用会大幅飙升。

直接可用的优化方案 #

1. 替换Queue为轻量通信方式 #

• 用multiprocessing.Pipe代替Queue：Pipe是点对点通信，比带锁和后台线程的Queue内存开销小很多，适合一对一的worker通信场景。
• 优先用进程池的回调机制：如果能用multiprocessing.Pool，直接用Pool.imap_unordered或Pool.apply_async的回调函数处理结果，无需手动维护Queue，避免数据堆积。

2. 延迟加载数据，让worker按需读取 #

• 主进程仅存储数据路径或索引，worker自行按需加载单组数据，处理完成后立即释放该组数据的内存。示例代码：

  def worker(data_idx):
      # 按需加载单组数据，仅在worker内存中临时存储
      spectrum = load_spectrum_from_file(data_idx)
      result = process_spectrum(spectrum)
      save_result(result)
      # 显式删除大变量，强制释放内存
      del spectrum, result
      return None
  
  if __name__ == "__main__":
      from multiprocessing import Pool
      # 进程数不超过Slurm分配的CPU核心数
      with Pool(processes=4) as pool:
          pool.map(worker, range(2000))
  

3. 优化数据传输，减少大对象传递 #

• 用高效序列化方式：比如pickle的protocol=5（支持大对象，序列化更快更省内存）。
• 避免传输完整数组：worker处理完直接将结果保存到文件，主进程仅接收处理完成的信号；若无需完整结果，只传输峰值、均值等统计值。

4. 严格控制子进程的数量与生命周期 #

• 进程数不超过Slurm分配的CPU核心数，过多进程只会增加内存开销和上下文切换成本。
• 用multiprocessing.Pool复用进程，避免频繁创建销毁进程产生的内存碎片。
• 在worker函数末尾强制触发垃圾回收：添加import gc; gc.collect()，及时释放未被自动回收的内存。

5. 调整Slurm作业配置 #

• 合理设置内存参数：根据进程数和单进程预估内存占用配置--mem或--mem-per-cpu，比如8个进程、单进程占20GB，可设置--mem=160G，避免因Slurm调度问题导致内存分配不足。
• 启用OOM监控：添加--mail-type=OUT_OF_MEMORY，OOM时会收到包含内存快照的邮件，便于定位内存占用异常的进程。

6. 优化数据结构 #

• 用numpy数组存储光谱数据：Python list每个元素都是独立对象，内存开销远大于numpy数组。
• 降低数据精度：若业务允许，将float64类型转为float32，直接减少一半内存占用。

排查验证方法 #

• 用tracemalloc监控内存变化，定位内存增长节点：

  import tracemalloc
  tracemalloc.start()
  
  # 执行部分处理逻辑
  snapshot = tracemalloc.take_snapshot()
  top_stats = snapshot.statistics('lineno')
  for stat in top_stats[:10]:
      print(stat)
  

• 用psutil打印worker内存占用，确认单进程内存是否超标：

  import psutil
  def worker(data_idx):
      proc = psutil.Process()
      print(f"Worker {data_idx} 初始内存: {proc.memory_info().rss / 1024**2:.2f} MB")
      # 处理数据逻辑
      print(f"Worker {data_idx} 处理后内存: {proc.memory_info().rss / 1024**2:.2f} MB")
  

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