Pythonの使用メモリを空きメモリまでに制限する

Deep LearningでメモリとCPUをぶん回している時、調子に乗って多重起動してOSをハングアップさせてしまう事が多いので、起動時に使用メモリを制限しようと思いました。

Stack Overflowに良いサンプルがあったので、利用させて頂きました。 /proc/meminfoはKB基準で、setrlimitはBytes基準なので1024を掛ける。 soft limitとhard limitの仕組みについては公式リファレンスを参照すると良いかと思う。

import sys, resource

def memory_limit():
    free_memory = get_memory()
    soft = free_memory * 1024 * 0.8
    hard = free_memory * 1024
    resource.setrlimit(resource.RLIMIT_AS, (soft, hard))

def get_memory():
    with open('/proc/meminfo', 'r') as mem:
        free_memory = 0
        for i in mem:
            sline = i.split()
            if str(sline[0]) in ('MemFree:', 'Buffers:', 'Cached:'):
                free_memory += int(sline[1])
    return free_memory

if __name__ == '__main__':
    memory_limit() # Limitates maximun memory usage to half
    try:
        main()
    except MemoryError:
        sys.stderr.write('\n\nERROR: Memory Exception\n')
        sys.exit(1)

上記では未使用メモリの80%を越えたらエラーが出るようになっている。 この設定をしてからOSハングアップ地獄から解放されて助かっている。 再度調べるのが面倒なのでメモ。