Python multiprocessing结合scipy periodogram内存泄漏问题排查

最小复现代码 #

#!/usr/bin/env python3
import numpy as np
from scipy import signal
from memory_profiler import profile
import gc

fs = 30 # Sampling rate

def genSpec(args):
    time = np.arange(60*30) / fs # 1 minute at 30 hz
    freq = 65
    x = np.sin(2*np.pi*freq*time)
    w = fs*10 # Window length
    spec = [signal.periodogram(x[i:i+w], fs=fs, nfft=int(fs/.001)) for i in range(len(x)-w+1)]
    return spec

def genNoSpec(args):
    spec = genSpec(args)
    return None

def genSpecNP(args):
    spec = genSpec(args)
    return np.array(spec)

def doMultiprocess(numItems, func):
    from multiprocessing import Pool
    with Pool() as p:
        list(p.imap_unordered(func, range(numItems))) # Do nothing with it

@profile
def gimmeData(numItems):
    print('Generating without spec')
    doMultiprocess(numItems, genNoSpec)
    print('Generating spec wrapped in np.array')
    doMultiprocess(numItems, genSpecNP)
    print('Generating with spec')
    doMultiprocess(numItems, genSpec)
    print('Generation completed')
    gc.collect()
    print('Garbage collected')

gimmeData(16)
input(f'Press [Enter] to exit')

内存占用统计结果 #

运行上述代码后，memory_profiler输出的内存统计如下：

Line #    Mem usage    Increment  Occurences   Line Contents
============================================================
    30     76.5 MiB     76.5 MiB           1   @profile
    31                                         def gimmeData(numItems):
    32     76.5 MiB      0.0 MiB           1       print('Generating without spec')
    33     77.4 MiB      0.9 MiB           1       doMultiprocess(numItems, genNoSpec)
    34     77.4 MiB      0.0 MiB           1       print('Generating spec wrapped in np.array')
    35     78.4 MiB      1.0 MiB           1       doMultiprocess(numItems, genSpecNP)
    36     78.4 MiB      0.0 MiB           1       print('Generating with spec')
    37   5389.3 MiB   5310.9 MiB           1       doMultiprocess(numItems, genSpec)
    38   5389.3 MiB      0.0 MiB           1       print('Generation completed')
    39   5384.7 MiB     -4.7 MiB           1       gc.collect()
    40   5384.7 MiB      0.0 MiB           1       print('Garbage collected')

观测规律 #

从运行结果可总结出以下规律：

• 若scipy.signal.periodogram生成的频谱计算结果从子进程传回父进程（对应genSpec函数场景），即使结果传回后立即离开作用域，也无法被垃圾回收；
• 若计算结果不传回父进程（对应genNoSpec函数场景），子进程内的内存可被正常回收；
• 若先将计算结果封装为numpy数组再传回父进程（对应genSpecNP函数场景），内存也可被正常回收。

初步推测与测试环境 #

基于第三点现象，初步推测numpy的数据转换过程清除了periodogram调用生成的、导致内存无法正常回收的关联对象，但暂未定位到具体根因，需要明确该问题的产生原因与修复方案。
已完成复现的测试环境如下：

• 初始测试基于Anaconda环境，Python版本为3.7.6，scipy版本为1.6.2，已在Ubuntu 18.04.6、Gentoo系统下复现该问题；
• 后续在Arch Linux系统、Python 3.10.5、scipy 1.8.1环境下也观测到完全相同的现象。

编辑记录：已对示例做最小化调整，将直接调用signal.spectrogram的逻辑替换为封装periodogram结果为numpy数组的逻辑

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