限制Python程序使用的RAM

Question

我想限制一个 Python 程序使用的 RAM 为一半，这样当所有 RAM 都被占用时程序就不会完全冻结。我正在使用下面的代码实现，但它并不起作用，我的笔记本电脑仍然会冻结：

import sys
import resource

def memory_limit():
    rsrc = resource.RLIMIT_DATA
    soft, hard = resource.getrlimit(rsrc)
    soft /= 2
    resource.setrlimit(rsrc, (soft, hard))

if __name__ == '__main__':
    memory_limit() # Limitates maximun memory usage to half
    try:
        main()
    except MemoryError:
        sys.stderr.write('MAXIMUM MEMORY EXCEEDED')
        sys.exit(-1)

我正在使用从main函数调用的其他函数。

我做错了什么？

提前感谢您。

附注：我已经搜索了此问题并找到了我所放置的代码，但仍然无法工作...

- Ulises CT

---

也许你想在 Python 脚本之外使用 ulimit 或 prlimit 命令。或者设置一个内存限制的 cgroup 并在其中运行脚本。我不确定自我限制是否是最好的主意——如果尝试检查或执行限制的代码需要在过程中分配内存，会发生什么？ - twalberg

你尝试过 soft /= 100 或者 soft //= 2 吗？ - Gribouillis

如果我使用它，它会说退出太快 @Gribouillis - Ulises CT

为什么你的所有RAM都被使用了？你正在将大量数据加载到内存中吗？你尝试过使用生成器吗？ - roymustang86

@roymustang86 这是一个大学作业，我们必须使用深度优先搜索策略生成一棵树，但由于任务说明的原因，它会陷入无限循环。 - Ulises CT

Answer 1

我做了一些研究，发现了一个从Linux系统获取内存的函数，链接在这里：Determine free RAM in Python。我稍微修改了一下，将内存硬限制设置为可用空闲内存的一半。

import resource
import sys

def memory_limit():
    """Limit max memory usage to half."""
    soft, hard = resource.getrlimit(resource.RLIMIT_AS)
    # Convert KiB to bytes, and divide in two to half
    resource.setrlimit(resource.RLIMIT_AS, (get_memory() * 1024 / 2, hard))

def get_memory():
    with open('/proc/meminfo', 'r') as mem:
        free_memory = 0
        for i in mem:
            sline = i.split()
            if str(sline[0]) in ('MemFree:', 'Buffers:', 'Cached:'):
                free_memory += int(sline[1])
    return free_memory  # KiB

if __name__ == '__main__':
    memory_limit_half()
    try:
        main()
    except MemoryError:
        sys.stderr.write('\n\nERROR: Memory Exception\n')
        sys.exit(1)

- Ulises CT

Answer 2

我修改了@Ulises CT的答案，因为我认为改变原始函数太多不是很好，所以我将其转换为一个装饰器。希望这有帮助。

import resource
import platform
import sys

def memory_limit(percentage: float):
    """
    只在linux操作系统起作用
    """
    if platform.system() != "Linux":
        print('Only works on linux!')
        return
    soft, hard = resource.getrlimit(resource.RLIMIT_AS)
    resource.setrlimit(resource.RLIMIT_AS, (get_memory() * 1024 * percentage, hard))

def get_memory():
    with open('/proc/meminfo', 'r') as mem:
        free_memory = 0
        for i in mem:
            sline = i.split()
            if str(sline[0]) in ('MemFree:', 'Buffers:', 'Cached:'):
                free_memory += int(sline[1])
    return free_memory

def memory(percentage=0.8):
    def decorator(function):
        def wrapper(*args, **kwargs):
            memory_limit(percentage)
            try:
                return function(*args, **kwargs)
            except MemoryError:
                mem = get_memory() / 1024 /1024
                print('Remain: %.2f GB' % mem)
                sys.stderr.write('\n\nERROR: Memory Exception\n')
                sys.exit(1)
        return wrapper
    return decorator

@memory(percentage=0.8)
def main():
    print('My memory is limited to 80%.')

- qiuyuan

---

在内存修饰器中，你应该return function(*args, **kwargs)。 - jaksco

有没有其他方法使这个装饰器在所有操作系统上都能工作？ - shivang patel

Answer 3

由于 https://git.kernel.org/pub/scm/linux/kernel/git/torvalds/linux.git/commit/?id=34e431b0ae398fc54ea69ff85ec700722c9da773 的原因，建议使用以下代码：

if str(sline[0]) == 'MemAvailable:':
            free_memory = int(sline[1])
            break

答案代码告诉我，在一台拥有1TB RAM的机器上，可用内存为8GB。

- The_Pingu

Answer 4

一个更简单的解决方案是：

import psutil
import resource

# Calculate the maximum memory limit (80% of available memory)
virtual_memory = psutil.virtual_memory()
available_memory = virtual_memory.available
memory_limit = int(available_memory * 0.8)

# Set the memory limit
resource.setrlimit(resource.RLIMIT_AS, (memory_limit, memory_limit))

- kiran.gilvaz