How to have pytest place memory limits on tests?

Question

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.

score 8 · Accepted Answer · edited Jan 22 '21 at 08:18

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.

This looks awesome. Could you elaborate on `a bit buggy with setting the limits`? — Roger Dahl, Jan 29 '18 at 04:09
Note that the decorator says `@memory_limit(int(16e8))`, but the try/except block fails for `allocate(int(3e8))`. "A bit buggy" meant that I couldn't find an obvious relation between the two numbers, but for my purposes, an approximate was fine. — , Jan 29 '18 at 18:02
Small but efficient! You could make a pytest plugin out of that, installable through pip / pypi.org — myselfhimself, Sep 11 '20 at 21:27

How to have pytest place memory limits on tests?

1 Answers1