2017-09-20 12 views
1

pytestを使用していますが、テストごとに最大メモリ使用量を設定できるデコレータが必要です。答えたこのquestionテストでpytestの場所のメモリ制限を得るには?

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

と同様に、私が欲しい、

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

はEDIT:

私が試した、

メモリ割り当てをキャッチしていない
>>> 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 

関数内にあります。

答えて

1

limit the memory usedの方法を学んだあと、memory is currently usedの量を見て、メモリのインクリメントが高すぎるとエラーが出るデコレータを書きました。それは限界を設定すると少しバギーですが、それは私のために十分に動作します。

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

少なくとも私のマシンでは、コードは実行され、エラーなしで終了します。

+0

これは素晴らしいですね。あなたは限界を設定して '少しバギーに精緻化できますか? –

+1

デコレータは '@memory_limit(int(16e8))'と言っていますが、 'allocate(int(3e8))'のtry/exceptブロックは失敗します。 「ちょっとしたバグ」は、2つの数字の間に明白な関係があるとは思えませんでしたが、私の目的のためにはおおよそのものでした。 –