时间戳

首页 > 解决方案 > SLURM Python脚本在循环中累积内存

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

问题描述

我在 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 是罪魁祸首。

我添加了（关闭 pyplot 窗口）：

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

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

标签： pythonmemory-leaksout-of-memoryslurm

解决方案

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

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

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

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

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

加法

现在您已经将问题范围缩小到pcolormesh，下一步是阅读有关 matplotlib 和使用内存的文档或教程。pcolormesh此外，网络搜索可pcolormesh memory leak找到有关此的特定提示。

最简单的尝试是添加调用以ax.cla()清除轴，如本例所示。

您可以从 pyplot 切换到 matplotlib 的面向对象接口，如果有任何全局状态，它不会保留太多。相反，我认为pyplot保留figand ax，在这种情况下释放变量不足以释放它们的对象。

假设您的数据适合矩形网格，显然imshow 比 pcolormesh 使用更少的内存和时间。

注意#1741 之类的问题建议pcolormesh只创建一次，然后在每次循环迭代中设置其数据——你可以做mesh = plt.pcolormesh(...)一次，然后mesh.set_array(np.log10(den))在每次迭代中进行类似的操作吗？它还建议调用cla().

推荐阅读