如何限制python进程的内存使用量

一个进程如果疯狂的使用内存,那么就会抢占服务器的资源导致其他任务无法正常执行,那么有没有什么办法可以限制进程所能使用的内存总量呢?本文以python为例,向你展示如何限制进程的内存。

1. 使用resource模块

python的resource模块提供了限制内存使用的功能,你可以设置进程所能申请的最大内存,当内存申请量超过限制时,进行会被kill掉,下面是一个简单的示例

import resource
import time
import psutil
​
p = psutil.Process()
print(p.pid)
​
def limit_memory(maxsize):
    soft, hard = resource.getrlimit(resource.RLIMIT_AS)
    resource.setrlimit(resource.RLIMIT_AS, (maxsize, hard))
​
limit_memory(1024*1024*180)   # 限制180M ,可申请内存,对应虚拟内存
​
lst = []
while True:
    lst.append("a"*1000000)
    time.sleep(0.5)

执行脚本后,程序会输出自己的进行pid,使用命令 top -p pid, 就可以查看到这个进程的资源使用情况,重点关注VIRT 的数值,这个是虚拟内存,resource模块限制的正是这个数值,这个数值在显示时默认是字节,按一下e 键,就可以切换单位,切换到m更容易观察,当VIRT接近180M或者略微超出时,进程就被kill了。

这里有必要解释一下,为什么VIRT 小于所限制的数值时进程会被kill。我在程序里使用列表不停的append字符串,python里的列表与C++ STL 里的vector有一点像,列表会先申请一部分内存,当这部分内存快耗尽时,它会依据算法申请一块更大的内存备用,当VIRT 接近180M时,已经申请的内存不够用了,于是python再次申请,结果一下子超过了180M,终端来没有来得及显示最终的内存量,进程就被kill了。

resource.setrlimit 可以限制python进行对资源的使用,resource.RLIMIT_AS 是 进程可能占用的地址空间的最大区域,单位是字节,soft 和 hard 的数值,你不必去关心,在linux下,默认都是-1,其实getrlimit方法可以不调用,直接将hard替换成-1就行。

2. 使用cgroup

2.1 cgroup

Cgroup 是 Linux kernel 的一项功能:它是在一个系统中运行的层级制进程组,你可对其进行资源分配(如 CPU 时间、系统内存、网络带宽或者这些资源的组合),在使用前,你需要安装

 yum install -y libcgroup libcgroup-tools

2.2 创建cgroup

 cd /sys/fs/cgroup/memory/
mkdir memory_limit 
cd memory_limit
mkdir 30m_limit
cd 30m_limit
echo 30m > memory.limit_in_bytes   # 限制程序最多只能使用30M内存

/sys/fs/cgroup/memory/ 可以理解为限制内存使用的根目录,我创建memory_limit的目的,是为了在memory_limit里创建多个cgroup,不然都在/sys/fs/cgroup/memory/ 下创建会很乱。

2.3 验证内存被限制

创建一个python脚本 memory_limit.py

import resource
import time
import psutil
​
p = psutil.Process()
print(p.pid)
​
lst = []
while True:
    lst.append("a"*100000)
    time.sleep(0.3)
    print(p.memory_info())

执行脚本

cgexec -g memory:memory_limit/30m_limit python memory_limit.py

cat 30m_limit/cgroup.procs 可以查看进程的pid 28636

使用命令 top -p 28636 可以查看该进程使用的内存,当RES 略微超出30m时,进程被kill。为什么会超出一点点呢,在程序里,一次申请"a"*100000 这么多内存,刚好超过了30m,系统判断内存超出限制后才会kill进程。

3. 父进程监控方式

使用resource模块只能控制进程使用的虚拟内存,使用cgroup虽然可以限制物理内存,但必须借助cgroup,而且cgroup也并非万能,假设你对进程A进行了内存资源限制,进程A在运行期间产生了子进程B,子进程B又产生了自己的子进程C,那么这时,进程ABC三个进程作为整体来遵守内存资源限制。当内存资源达到所设置的上限时,cgroup不会将这三个进程都kill,而是kill掉消耗资源最多的那一个,这样就可能会产生孤立进程,而且并没有什么设置方法可以让cgroup如你所期望的那样在这种情况下kill所有进程。

鉴于此,我自己试验了一种方式,用父进程来监控子进程的内存资源使用情况。假设我有一个脚本limit.py

import time
import psutil
from subprocess import Popen
​
​
child = Popen(['python', '/root/test/child1.py'],start_new_session=True)
print("child1: " + str(child.pid))
​
child = Popen(['python', '/root/test/child1.py'],start_new_session=True)
print("child1_2: " + str(child.pid))
​
lst = []
​
while True:
    lst.append("a"*1000000)
    time.sleep(0.5)

limit.py脚本会启动两个子进程child1.py, child1.py 会启动子进程child2.py

import time
from subprocess import Popen
​
child = Popen(['python', '/root/test/child2.py'],start_new_session=True)
​
print("child2: " + str(child.pid))
​
while True:
    time.sleep(5)

child2.py 又会启动child3.py

import time
​
lst = []
while True:
    lst.append("a"*1000000)
    time.sleep(2)

直接正常启动limit脚本,会生成6个子进程,两个child1, 两个child2, 两个child3,使用cgroup启动limit,这7个进程会作为一个整体来遵守资源限制条件,当资源使用超出限制时,cgroup会杀掉消耗资源最多的那个进程,而非全部。

如果,我用一个监控脚本来启动limit.py, 那么我就可以监控limit进程及其子进程的内存使用情况,当内存超出限制时,kill所有进程,为了实现这一目标,我需要做到以下几点:

  1. 获取limit进程及其子进程的内存使用总和
  2. 内存使用超出限制后,kill所有子进程
  3. 当limit正常退出后,监控进程应当正常退出
  4. 避免limit成为僵尸进程

对子进程的操作,可以使用psutil这个第三方库,示例代码如下

import psutil
import os
import time
import errno
import signal
from subprocess import Popen
​
​
pid = os.getpid()
print("parent pid: " + str(pid))
p = psutil.Process(pid)
​
def wait_child(signum, frame):
    try:
        while True:
            childpid, status = os.waitpid(-1, os.WNOHANG)
            if childpid == 0:
                print('没有立即可用的子进程')
                break
            exitcode = status >> 8
            print(f'子进程{childpid}退出,状态码是{exitcode}')
    except OSError as e:
        if e.errno == errno.ECHILD:
            print("没有需要等待wait的子进程")
        else:
            raise
​
child = Popen(['python', '/root/test/limit.py'],start_new_session=True)
print("child pid: " + str(child.pid))
​
signal.signal(signal.SIGCHLD, wait_child)  # wait 子进程,防止僵尸进程
while True:
    state = child.poll()
    if state is not None:     # 子进程结束
        print("子进程结束")
        break
    try: 
        # 获取所有子进程的内存使用总量
        child_process = psutil.Process(child.pid)
        memory_sum = 0
        for child_pro in child_process.children(recursive=True): 
            info = child_pro.memory_full_info()
            memory = info.uss / (1024*1024)     # 物理内存
            memory_sum += memory
​
        info = child_process.memory_full_info()
        memory = info.uss / (1024*1024)
        memory_sum += memory
      
        print(memory_sum)
        if memory_sum > 200:
            print("内存超出限制,子进程被kill")
            # 内存使用超出限制后,kill所有子进程
            for child_pro in child_process.children(recursive=True):
                print("kill  ", child_pro.pid)
                child_pro.kill()
​
            child_process.kill()
    except:
        pass
​
 
    time.sleep(5)

4. 使用hook的方式监控内存使用

方法1最简单,但只能限制虚拟内存,方法2可以限制物理内存,但要借助cgroup来实现,方法3是纯python实现,虽然没有修改被监控的脚本,但是需要借助额外的监控脚本来实现内存使用的控制。那么有没有其他的方法,可以不主动的使用额外的监控脚本,也不用修改被监控内存的脚本,也能达到内存监控使用的目的呢?答案是借助sitecustomize.py

sitecustomize.py是一个特殊的脚本,通常放在site-packages目录下,你也可以通过PYTHONPATH环境变量将它所在的目录加入到sys.path中,当你用python执行一个脚本时,python会先加载并执行这个脚本,你可以这样理解它,在程序正式执行前,会先执行sitecustomize.py,这个时间有多提前呢,在执行sitecustomize.py时,sys.argv还没有被初始化。这这个脚本里,我启动一个线程,来监控当前进程及其子进程使用的内存总量,当内存使用总量超出限制时终结进程。

import threading
import psutil
import os
import time
import errno
import signal
import sys
​
​
def wait_child(signum, frame):
    try:
        while True:
            childpid, status = os.waitpid(-1, os.WNOHANG)
            if childpid == 0:
                print('没有立即可用的子进程')
                break
            exitcode = status >> 8
            print(f'子进程{childpid}退出,状态码是{exitcode}')
    except OSError as e:
        if e.errno == errno.ECHILD:
            print("没有需要等待wait的子进程")
        else:
            raise
​
​
def limit_memory():
    pid = os.getpid()
    p = psutil.Process(pid)
   
    while True:
​
        memory_sum = 0
        for child_pro in p.children(recursive=True): 
            info = child_pro.memory_full_info()
            memory = info.uss / (1024*1024)
            memory_sum += memory
​
        info = p.memory_full_info()
        memory = info.uss / (1024*1024)
        memory_sum += memory
    
        print(memory_sum, os.getpid())
​
        if memory_sum > 100:
            for child_pro in p.children(recursive=True):
                print("kill  ", child_pro.pid)
                child_pro.kill()
        
            p.kill() 
        time.sleep(3)
​
signal.signal(signal.SIGCHLD, wait_child)
t = threading.Thread(target=limit_memory)
t.start()

使用sitecustomize.py,就不再需要方法3里的监控脚本了,不知不觉中,就对python进程的内存使用进行控制,这对于执行python脚本的用户来说,完全无感知。

但这种方法也有自己的缺陷,由于sitecustomize.py 是自动默认加载的,因此所有的python进程都将被限制内存使用,包括进程启动后创建的子进程,这着实不是一个完美的方案。那么有没有什么办法可以做到只限制那些希望被限制内存使用的脚本呢?

我想通过sys.argv来判断这个脚本是否需要被限制内存,如果希望被限制内存,那么在使用python命令执行脚本时加上一个参数,类似这样

python limit.py limit_memory

如果sys.argv中有limit_memory 这个字符串,那么就对内存进行限制,否则正常执行,想法是美好的,现实却很残酷,由于sitecustomize.py的执行时间过于靠前,以至于sys模块的argv还没有被初始化,在sitecustomize.py脚本里,根本拿不到执行程序的参数。经过一番google,终于还是找到了办法,可以获取启动命令中的参数

def get_python_interpreter_arguments():
    argc = ctypes.c_int()
    argv = ctypes.POINTER(ctypes.c_wchar_p if sys.version_info >= (3, ) else ctypes.c_char_p)()
    ctypes.pythonapi.Py_GetArgcArgv(ctypes.byref(argc), ctypes.byref(argv))
    arguments = list()
    for i in range(argc.value):
        arguments.append(argv[i])
​
    return arguments
​
argv = get_python_interpreter_arguments()
​
​
​
if 'limit_memory' in argv:
    signal.signal(signal.SIGCHLD, wait_child)
    t = threading.Thread(target=limit_memory)
    t.start()

get_python_interpreter_arguments 函数可以获取进程启动时的命令参数,这样,就可以通过这些参数来决定是否需要对进程使用的内存进行监控

扫描关注, 与我技术互动

QQ交流群: 211426309

加入知识星球, 每天收获更多精彩内容

分享日常研究的python技术和遇到的问题及解决方案