pytest如何在测试中设置内存限制? [英] How to have pytest place memory limits on tests?

问题描述

我正在使用pytest，但我想拥有一个装饰器，该装饰器可以设置每个测试的最大内存使用量.类似于与此问题的答案，

I am using pytest, but I would like to have a decorator that could set a maximum memory usage per test. Similar to this question which was answered with,

@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

哪个没有捕获函数中的内存分配.

Which doesn't catch the memory allocation in the function.

推荐答案

学习了如何 ，我写了一个装饰器，如果内存增量过高，则会出错.设置限制有点麻烦，但是对我来说效果很好.

After learning how to limit the memory used and seeing how much memory is currently used, I wrote a decorator that errors out if the memory increment is too high. It's a bit buggy with setting the limits, but it works well enough for me.

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")

至少在我的机器上，代码可以正常运行和退出.

At least on my machine, code runs and exits without an error.

这篇关于pytest如何在测试中设置内存限制?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！