SLURM Python脚本在循环中累积内存

Question

我在 HPC 的 SLURM 调度程序上运行一个简单的 python 脚本。 它读入一个数据集（大约 6GB）并绘制和保存部分数据的图像。这些数据文件有好几个，所以我使用循环进行迭代，直到完成每个文件的数据绘制。

但是，由于某种原因，每个循环中的内存使用量都会增加。我已经使用 getsizeof() 映射了我的变量，但它们似乎不会随着迭代而改变。所以我不确定这个内存“泄漏”可能来自哪里。

这是我的脚本：

import os, psutil
import sdf_helper as sh
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.ticker as plticker
plt.rcParams['figure.figsize'] = [6, 4]
plt.rcParams['figure.dpi'] = 120 # 200 e.g. is really fine, but slower
from sys import getsizeof


for i in range(5,372):
    plt.clf()   
    fig, ax = plt.subplots()
    #dd gets data using the epoch specific SDF file reader sh.getdata
    dd = sh.getdata(i,'/dfs6/pub/user');
    #extract density data as 2D array
    den = dd.Derived_Number_Density_electron.data.T;
    nmin = np.min(dd.Derived_Number_Density_electron.data[np.nonzero(dd.Derived_Number_Density_electron.data)])
    #extract grid points as 2D array
    xy = dd.Derived_Number_Density_electron.grid.data
    #extract single number time
    time = dd.Header.get('time')
    #free up memory from dd
    dd = None
    #plotting
    plt.pcolormesh(xy[0], xy[1],np.log10(den), vmin = 20, vmax = 30)
    cbar = plt.colorbar()
    cbar.set_label('Density in log10($m^{-3}$)')
    plt.title("time:   %1.3e s \n Min e- density:   %1.2e $m^{-3}$" %(time,nmin))
    ax.set_facecolor('black')
    plt.savefig('D00%i.png'%i, bbox_inches='tight')
    print("dd: ", getsizeof(dd))
    print("den: ",getsizeof(den))
    print("nmin: ",getsizeof(nmin))
    print("xy: ",getsizeof(xy))
    print("time: ",getsizeof(time))
    print("fig: ",getsizeof(fig))
    print("ax: ",getsizeof(ax))
    process = psutil.Process(os.getpid())
    print(process.memory_info().rss)

输出

Reading file /dfs6/pub/user/0005.sdf
dd:  16
den:  112
nmin:  32
xy:  56
time:  24
fig:  48
ax:  48
8991707136

Reading file /dfs6/pub/user0006.sdf
dd:  16
den:  112
nmin:  32
xy:  56
time:  24
fig:  48
ax:  48
13814497280

Reading file /dfs6/pub/user/0007.sdf
dd:  16
den:  112
nmin:  32
xy:  56
time:  24
fig:  48
ax:  48
18648313856

SLURM 输入

#!/bin/bash

#SBATCH -p free
#SBATCH --job-name=epochpyd1
#SBATCH --nodes=1
#SBATCH --ntasks=1
#SBATCH --mem-per-cpu=20000


#SBATCH --mail-type=begin,end
#SBATCH --mail-user=**

module purge
module load python/3.8.0

python3 -u /data/homezvol0/user/CNTDensity.py > density.out

SLURM 输出

/data/homezvol0/user/CNTDensity.py:21: RuntimeWarning: divide by zero encountered in log10
  plt.pcolormesh(xy[0], xy[1],np.log10(den), vmin = 20, vmax = 30)
/export/spool/slurm/slurmd.spool/job1910549/slurm_script: line 16:  8004 Killed                  python3 -u /data/homezvol0/user/CNTDensity.py > density.out
slurmstepd: error: Detected 1 oom-kill event(s) in step 1910549.batch cgroup. Some of your processes may have been killed by the cgroup out-of-memory handler.

据我所知，一切似乎都在工作。不确定什么会占用超过 20GB 的内存。

编辑 所以我开始从下往上注释掉循环的各个部分。现在很明显 pcolormesh 是罪魁祸首。

我已添加 (Closing pyplot windows)：

fig.clear()
plt.clf()
plt.close('all')
fig = None
ax = None
del fig
del ax

到最后，无论如何记忆都在不断攀升。我对正在发生的事情完全不知所措。

Answer 1

您在正确的轨道上，已经使每次迭代的内存累积量可见。调试的下一步是考虑内存可能在哪里积累的假设以及测试这些假设的方法。

每次迭代后保留内存的是den 之类的变量。您可以通过dd = None 清除这些变量，或通过del dd 删除它们，或将循环体的部分移动到子例程中来排除这些假设（从而缩小问题范围），这样一些变量就会消失当这些子程序返回时。 （并且分解出子例程也可以使这些部分更可重用和更容易测试。）这种技术将排除问题的一些可能原因，但我不希望这些变量分配在迭代中累积内存，如果代码在每次迭代中将数据添加到 dict 或 list 就会出现这种情况。

另一个假设是，matplotlib 中的状态累积不会被plt.clf() 清除，或者sdf_helper 中的状态累积。我对这些库知之甚少，无法提供直接的见解，但他们的文档应该说明如何清除状态。即使不知道如何清除它们的状态，我们也可以想办法检验这些假设。例如。注释掉 plt 调用或至少是数据密集型调用，然后查看内存是否仍在累积。

你可能会想出比我更多的假设。首先头脑风暴假设是一种很好的方法，因为其中一个可能是明显的最佳候选者，或者其中一个可能比其他的更容易测试。

请注意，累积内存可能有多种原因，在这种情况下，修复一个原因会减少内存累积，但不会修复它。由于您正在测量内存积累，因此您将能够检测到这一点。在许多调试情况下，我们看不到多个原因对问题的增量贡献，例如 flakey 结果，因此另一种技术是删除可能导致问题的所有内容，然后一次添加一个。

添加

既然您已将问题范围缩小到 pcolormesh，下一步是阅读 docs 或有关 matplotlib 和 pcolormesh 如何使用内存的教程。此外，通过网络搜索 pcolormesh memory leak 可以找到关于此的具体提示。

最简单的尝试是添加对ax.cla() 的调用以清除坐标区，如this example。

您可以从 pyplot 切换到 matplotlib 的面向对象接口，该接口不会保留任何全局状态。相比之下，我认为pyplot 保留了fig 和ax，在这种情况下释放变量不足以释放它们的对象。

显然是imshowuses less memory and time than pcolormesh, assuming your data fits on a rectangular grid。

注意Issue like #1741，它建议只创建一次pcolormesh，然后在每次循环迭代中设置它的数据——你能做一次mesh = plt.pcolormesh(...)，然后在每次迭代中像mesh.set_array(np.log10(den))一样吗？它还建议致电cla()。