如何让 pytest 在测试中设置内存限制？

Question

我正在使用 pytest，但我想要一个装饰器，可以设置每次测试的最大内存使用量。与回答的这个问题类似，

@pytest.mark.timeout(300)
def test_foo():
   pass

我想要，

@pytest.mark.maxmem(300)
def test_foo():
   pass

编辑：

我试过，

>>> import os, psutil
>>> import numpy as np
>>> process = psutil.Process(os.getpid())
>>> process.memory_info().rss/1e9
0.01978368
>>> def f():
...     x = np.arange(int(1e9))
... 
>>> process.memory_info().rss/1e9
0.01982464
>>> f()
>>> process.memory_info().rss/1e9
0.019832832

这不会捕获函数中的内存分配。

Answer 1

在学习了如何限制使用的内存并查看当前使用了多少内存后，我编写了一个装饰器，如果内存增量过高，则会出错。设置限制有点麻烦，但对我来说效果很好。

import resource, os, psutil
import numpy


def memory_limit(max_mem):
    def decorator(f):
        def wrapper(*args, **kwargs):
            process = psutil.Process(os.getpid())
            prev_limits = resource.getrlimit(resource.RLIMIT_AS)
            resource.setrlimit(
                resource.RLIMIT_AS, (
                     process.memory_info().rss + max_mem, -1
                )
            )
            result = f(*args, **kwargs)
            resource.setrlimit(resource.RLIMIT_AS, prev_limits)
            return result
        return wrapper
    return decorator
    
    
@memory_limit(int(16e8))
def allocate(N):
    return numpy.arange(N, dtype='u8')
    
a = [allocate(int(1e8)) for i in range(10)]
    
try:
    allocate(int(3e8))
except:
    exit(0)

raise Exception("Should have failed")

至少在我的机器上，代码运行并退出没有错误。