並列計算におけるメモリリークの回避（python）

はじめに

強引な解決策なので悪しからず。メモリリークの原因特定をしたいならば一番下のリンクが参考になる。

症状

強化学習を行っている際に、次のようなコードを用いてデータ生成を並列化した。時間が経つと子プロセスのメモリ使用量がどんどん増え、最終的にはメモリ上限に達し、計算できなくなった（ちなみにメモリは24GB）。ループの中身のメモリリークを調べたが原因が特定できなかった。なお、生成するデータ数には上限を設けており、このデータでメモリがあふれることはないものとする。

mp.py

import multiprocessing as mp
import time
N_PROCESS = 10

# この関数がメモリリークを起こす
def do_something(queue):
    while True:
        # データ生成

if __name__ == "__main__":
    queue = mp.Queue(maxsize=10000)
    for _ in range(N_PROCESS):
        p = mp.Process(target=do_something, args=(queue))
        p.start()
    while True:
        time.sleep(10)
        # queueの読み込み、学習など

対策

子プロセスで行っていた無限ループを一定回数のループに変更し、終了した子プロセスを再度立ち上げることで同じ内容を実現。子プロセスが走っている間はメモリが増えるが、終了時に子プロセスのメモリはリセットされるため、際限のない増大は抑えられた。

mp.py

import multiprocessing as mp
import time
N_PROCESS = 10
N_LOOP = 100

def do_something(queue):
    for _ in N_LOOP: # 無限ループを避ける
        # データ生成

if __name__ == "__main__":
    queue = mp.Queue(maxsize=10000)
    processes = []
    for _ in range(N_PROCESS):
        p = mp.Process(target=do_something, args=(queue))
        p.start()
        processes.append(p)
    while True:
        time.sleep(10)
        # 終了したプロセスは再度立ち上げる
        for i in range(N_PROCESS):
            if not processes[i].is_alive():
                processes[i].terminate()
                processes[i] = mp.Process(target=do_something, args=(queue))
                processes[i].start()
        # queueの読み込み、学習など

参考：メモリリークの調査で使ったもの

• tracememlock（公式やqiita記事参照）
• object graphs（マニュアル参照）