高内存notebook环境下multiprocessing.Pool初始化过慢问题求解

父进程内存占用较高时，启动multiprocessing.Pool的耗时会明显增加，初步推测该现象与fork机制相关。

此前曾尝试调用get_context("spawn")创建由空Python进程构成的进程池，但该方案未能生效，主要原因有两点：

• spawn模式下无法传入无法被导入的ad-hoc即席自定义函数，而我基于notebook开展科研计算工作，这类随写随用的临时函数是实验迭代的必需组件
• 通过htop观测发现，创建Pool(40)时内存占用并未出现40倍的暴涨，说明系统存在写时复制（copy-on-write）懒加载机制，这种前提下进程池初始化耗时过高的现象十分反常。

核心需求：寻找可行方案，能够在高内存占用的notebook环境中，快速初始化无需访问父进程大块内存数据的进程池。

问题复现代码 #

import time
import numpy as np
from multiprocessing import Pool, get_context

without_x = get_context("spawn")

t0 = time.time()
with Pool(40) as pool:
  print(f'无数据时fork模式进程池初始化耗时 {time.time() - t0:.2f}s')

t0 = time.time()
with without_x.Pool(40) as pool:
  print(f'无数据时spawn模式进程池初始化耗时  {time.time() - t0:.2f}s')

# 分配4GB内存
x = np.zeros(1024 ** 3, dtype=np.uint32)
x[:]=1

t0 = time.time()
with Pool(40) as pool:
  print(f'加载4GB数据后fork模式进程池初始化耗时 {time.time() - t0:.2f}s')

t0 = time.time()
with without_x.Pool(40) as pool:
  print(f'加载4GB数据后spawn模式进程池初始化耗时 {time.time() - t0:.2f}s')

初始测试运行结果 #

无数据时fork模式进程池初始化耗时 0.33s
无数据时spawn模式进程池初始化耗时  0.21s
加载4GB数据后fork模式进程池初始化耗时 1.71s
加载4GB数据后spawn模式进程池初始化耗时 1.93s

@juanpa.arrivillaga 提出的使用__name__ == "__main__"保护大内存分配逻辑的建议未达预期效果（测试前已重启进程），对应测试代码如下：

if __name__ == "__main__":
# if True:
  print('正在初始化4GB内存数据')
  x = np.zeros(1024 ** 3, dtype=np.uint32)
  x[:]=1

t0 = time.time()
with Pool(40) as pool:
  print(f'加载4GB数据后fork模式进程池初始化耗时 {time.time() - t0:.2f}s')

t0 = time.time()
with without_x.Pool(40) as pool:
  print(f'加载4GB数据后spawn模式进程池初始化耗时 {time.time() - t0:.2f}s')

该方案测试运行结果 #

正在初始化4GB内存数据
加载4GB数据后fork模式进程池初始化耗时 3.48s
加载4GB数据后spawn模式进程池初始化耗时 3.60s

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