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
I want,
@pytest.mark.maxmem(300)
def test_foo():
pass
EDIT:
I tried,
>>> 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.
If you love us? You can donate to us via Paypal or buy me a coffee so we can maintain and grow! Thank you!
Donate Us With