如何限制堆大小？

Question

我有时候写Python程序，但很难在执行之前判断它会用多少内存。因此，有时候我会运行一个Python程序，它试图分配大量的内存，这样就会导致系统的内核频繁地进行内存交换，进而影响其他正在运行的程序的性能。

因此，我想限制Python程序使用的内存大小。当达到这个限制时，程序可以直接崩溃。请问有什么好的方法可以做到这一点呢？

如果有影响的话，我的很多代码是用Cython写的，所以需要考虑在那里的内存分配。我并不一定要用纯Python的解决方案（不需要跨平台），所以在Linux上能用的任何方法都可以。

Answer 1

下面的代码为指定的最大常驻内存大小分配内存。

import resource

def set_memory_limit(memory_kilobytes):
    # ru_maxrss: peak memory usage (bytes on OS X, kilobytes on Linux)
    usage_kilobytes = lambda: resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
    rlimit_increment = 1024 * 1024
    resource.setrlimit(resource.RLIMIT_DATA, (rlimit_increment, resource.RLIM_INFINITY))

    memory_hog = []

    while usage_kilobytes() < memory_kilobytes:
        try:
            for x in range(100):
                memory_hog.append('x' * 400)
        except MemoryError as err:
            rlimit = resource.getrlimit(resource.RLIMIT_DATA)[0] + rlimit_increment
            resource.setrlimit(resource.RLIMIT_DATA, (rlimit, resource.RLIM_INFINITY))

set_memory_limit(50 * 1024)  # 50 mb

在Linux机器上测试过。

Answer 2

看看这个ulimit。它可以用来设置资源的限制。可能还需要一些合适的内核设置。

Answer 3

你可以看看 resource.setrlimit() 这个函数。它只在Unix系统上有效，但看起来可能正是你需要的，因为你可以用resource.RLIMIT_DATA参数为你的进程和它的子进程设置一个最大堆大小。

补充：这里有个例子：

import resource

rsrc = resource.RLIMIT_DATA
soft, hard = resource.getrlimit(rsrc)
print 'Soft limit starts as  :', soft

resource.setrlimit(rsrc, (1024, hard)) #limit to one kilobyte

soft, hard = resource.getrlimit(rsrc)
print 'Soft limit changed to :', soft

我不太确定你具体的使用场景是什么，但你可能需要用resource.RLIMIT_STACK来限制栈的大小。如果超过这个限制，系统会给你的进程发送一个SIGSEGV信号。要处理这个信号，你需要使用一个备用信号栈，具体可以参考setrlimit在Linux上的 手册。不过我不确定Python是否实现了sigaltstack，所以如果你想从这个限制中恢复，可能会有点困难。