一级Python工程师的总结——关于python中显存回收的问题

技术背景

笔者在执行一个Jax的任务中，又发现了一个奇怪的问题，就是明明只分配了很小的矩阵空间，但是在多次的任务执行之后，显存突然就爆了。而且此时已经按照Jax的官方说明配置了 XLA_PYTHON_CLIENT_PREALLOCATE 这个参数为 false ，也就是不进行显存的预分配（默认会分配90%的显存空间以供使用）。然后在网上找到了一些类似的问题，比如参考链接中的1、2、3、4，都是在一些操作后发现未释放显存，这里提供一个实例问题和处理的思路，如果有更好的方案欢迎大家在评论区留言。

问题复现

在未执行任何GPU的任务时，我们可以看到此时 nvidia-smi 的输出如下：

Tue Dec 14 16:14:32 2021       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.42.01    Driver Version: 470.42.01    CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Quadro RTX 4000     On   | 00000000:03:00.0  On |                  N/A |
| 30%   43C    P8    20W / 125W |   1260MiB /  7979MiB |     10%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  Quadro RTX 4000     On   | 00000000:A6:00.0 Off |                  N/A |
| 30%   34C    P8     7W / 125W |     10MiB /  7982MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      1673      G   /usr/lib/xorg/Xorg                110MiB |
|    0   N/A  N/A      3015      G   /usr/lib/xorg/Xorg                661MiB |
|    0   N/A  N/A      3251      G   /usr/bin/gnome-shell              132MiB |
|    0   N/A  N/A   1142734      G   ...AAAAAAAAA= --shared-files       64MiB |
|    0   N/A  N/A   1337710      G   ...AAAAAAAAA= --shared-files       80MiB |
|    0   N/A  N/A   1371509      G   ...369783.log --shared-files       63MiB |
|    0   N/A  N/A   1506625      G   ...AAAAAAAAA= --shared-files       89MiB |
|    1   N/A  N/A      1673      G   /usr/lib/xorg/Xorg                  4MiB |
|    1   N/A  N/A      3015      G   /usr/lib/xorg/Xorg                  4MiB |
+-----------------------------------------------------------------------------+

此时启动一个ipython的终端窗口，执行如下的Jax任务：

In [1]: import numpy as np

In [2]: import os
   ...: os.environ['CUDA_VISIBLE_DEVICES']='1'
   ...: os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"

In [3]: from jax import numpy as jnp

In [4]: a = np.ones(1000000)

In [5]: b = jnp.array(a)

此时再次查看 nvidia-smi 的结果如下：

Tue Dec 14 16:18:26 2021       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.42.01    Driver Version: 470.42.01    CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Quadro RTX 4000     On   | 00000000:03:00.0  On |                  N/A |
| 30%   42C    P8    20W / 125W |   1238MiB /  7979MiB |     10%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  Quadro RTX 4000     On   | 00000000:A6:00.0 Off |                  N/A |
| 30%   36C    P0    35W / 125W |    114MiB /  7982MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      1673      G   /usr/lib/xorg/Xorg                110MiB |
|    0   N/A  N/A      3015      G   /usr/lib/xorg/Xorg                661MiB |
|    0   N/A  N/A      3251      G   /usr/bin/gnome-shell              129MiB |
|    0   N/A  N/A   1142734      G   ...AAAAAAAAA= --shared-files       44MiB |
|    0   N/A  N/A   1337710      G   ...AAAAAAAAA= --shared-files       80MiB |
|    0   N/A  N/A   1371509      G   ...369783.log --shared-files       63MiB |
|    0   N/A  N/A   1506625      G   ...AAAAAAAAA= --shared-files       89MiB |
|    1   N/A  N/A      1673      G   /usr/lib/xorg/Xorg                  4MiB |
|    1   N/A  N/A      3015      G   /usr/lib/xorg/Xorg                  4MiB |
|    1   N/A  N/A   1743467      C   /usr/local/bin/python             101MiB |
+-----------------------------------------------------------------------------+

此时的结果还是比较符合我们的预期的，这个python的进程占用了101MB的空间。但是此时如果我们在ipython中把这个对象删除了：

In [6]: del b

In [7]: b
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<ipython-input-7-89e6c98d9288> in <module>
----> 1 b

NameError: name 'b' is not defined

然后再次查看 nvidia-smi 的结果：

Tue Dec 14 16:21:12 2021       
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 470.42.01    Driver Version: 470.42.01    CUDA Version: 11.4     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Quadro RTX 4000     On   | 00000000:03:00.0  On |                  N/A |
| 30%   42C    P5    21W / 125W |   1231MiB /  7979MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
|   1  Quadro RTX 4000     On   | 00000000:A6:00.0 Off |                  N/A |
| 30%   35C    P8     7W / 125W |    114MiB /  7982MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+
                                                                               
+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A      1673      G   /usr/lib/xorg/Xorg                110MiB |
|    0   N/A  N/A      3015      G   /usr/lib/xorg/Xorg                662MiB |
|    0   N/A  N/A      3251      G   /usr/bin/gnome-shell              111MiB |
|    0   N/A  N/A   1142734      G   ...AAAAAAAAA= --shared-files       55MiB |
|    0   N/A  N/A   1337710      G   ...AAAAAAAAA= --shared-files       80MiB |
|    0   N/A  N/A   1371509      G   ...369783.log --shared-files       63MiB |
|    0   N/A  N/A   1506625      G   ...AAAAAAAAA= --shared-files       89MiB |
|    1   N/A  N/A      1673      G   /usr/lib/xorg/Xorg                  4MiB |
|    1   N/A  N/A      3015      G   /usr/lib/xorg/Xorg                  4MiB |
|    1   N/A  N/A   1743467      C   /usr/local/bin/python             101MiB |
+-----------------------------------------------------------------------------+

此时我们可以看到，虽然已经把对象给删除了，在python的程序中已然找不到这个对象，但是在显存中的数据并未被消除。而且如果一直不消除，这块显存就会一直占用在那里，直到python进程（此时作为该进程的一个守护进程）的结束。

解决思路

暂时还不清楚这个问题发生的机制，在一些特定场景下出现僵尸进程的问题似乎跟我复现的这个场景也有所不同。只是考虑到在python的进程结束之后，这一块的显存还是被成功释放了的，因此我考虑直接用进程的方法来解决这个显存分配和清空的方法，以下是一个基于进程实现的案例：

import os
os.environ['CUDA_VISIBLE_DEVICES']='1'
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"

import time
from multiprocessing import Pool
import numpy as np
from jax import numpy as jnp

a = np.ones(1000000)

def f(a):
    b = jnp.array(a)
    time.sleep(2)
    print('Array b has been deleted!')
    return True

with Pool(1) as p:
    res = p.map(f, [(a,)])

print ('Is jax array deleted successfully?\t{}'.format(res))
time.sleep(6)

在这个程序中，我们把要执行的相关任务，包含GPU矩阵的转化与分配，都放到了一个函数中，然后通过multiprocessing开启一个子进程，来执行这个任务，并且在任务中甚至不需要手动执行 del 这个删除的操作。这么一来，我们既可以实现对象的即时销毁，也通过进程控制的机制确保在显存中占用的位置被清空。如果进程执行中存在一些问题，还可以通过 terminate 的操作来直接杀死进程，同样也可以确保显存占用不会发生堆积的情况。程序的执行结果如下：

Array b has been deleted!
Is jax array deleted successfully?      [True]

在程序执行的过程中我们也可以看到，在nvidia-smi中的显存占用，仅仅持续了2秒，也就是我们在函数内部设置的进程sleep参数。而在之后6秒的sleep时间中，这一块内存占用是被清空了的，这也就达到了我们最初的目的。当然，最重要的是，我们依然可以从函数中获取到返回值，这就确保后面有需要存储或者使用到的参数不被同步的销毁。需要注意的是，在同等条件下，如果不使用子进程来执行这个函数，而是直接使用 res=f(a) 的形式来执行，作为临时变量的b最终依然存在于显存之中，这是一个非常可怕的事情。

总结概要

在使用一些python的GPU模块，或者写CUDA时，有时会发现显存被无端占用的场景，即时执行了cudaFree()或者python的del操作，也无法消除这一块的显存占用。最终我们发现，可以通过额外开启一个子进程的方法来封装相关的操作，通过对进程的存活控制来实现对GPU显存占用的控制，有可能是一个临时规避问题的思路。

版权声明

本文首发链接为： 关于python中显存回收的问题 - DECHIN - 博客园