静かなる名辞

pythonとプログラミングのこと



multiprocessing.Poolがやたらメモリを消費するときの対策

概要

 multiprocessing.Poolは原理的にプロセスをforkさせるので、メインプロセスに大きなデータが残っているとそれが丸々コピーされてメモリ領域を食います。

 グローバル関数限定ですが、initializerを使って必要ないデータを消すことができます。また、Poolを作るタイミングを工夫することでそもそも大きいデータが子プロセスに引き継がれないようにすることができます。

前提状況の説明

 以下のようなコードです。

import subprocess
from multiprocessing import Pool

import numpy as np

a = np.arange(10**7)

def f():
    subprocess.run("ps -aux | grep [m]emory_test", shell=True)

p = Pool(1)
p.apply(f)
p.close()
p.terminate()
print(a.shape)

 見るからにメモリをドカ食いしそうな10**7のnumpy配列を確保しています。実行すると、以下のようになります。

username      7407 44.0  1.2 543952 100056 pts/0   Sl+  20:25   0:00 python memory_test.py
username      7411  0.0  1.1 347344 94640 pts/0    S+   20:25   0:00 python memory_test.py
(10000000,)

 もし子プロセスで走らせたいのがaを使わない処理なら、無駄に大容量のメモリを食っていることになります。

対策1:initializerで消す

実験

 以下のようなコードを書いてみます。

import subprocess
from multiprocessing import Pool

import numpy as np

a = np.arange(10**7)

def f():
    subprocess.run("ps -aux | grep [m]emory_test", shell=True)

def initializer():
    del globals()["a"]

p = Pool(1, initializer=initializer)
p.apply(f)
p.close()
p.terminate()
print(a.shape)  # ちゃんといることの確認
username      7427  0.0  1.2 543948 100112 pts/0   Sl+  20:26   0:00 python memory_test.py
username      7431  0.0  0.2 269212 16548 pts/0    S+   20:26   0:00 python memory_test.py
(10000000,)

 だいぶ改善しました。

本末転倒というか・・・

 まあ、見ての通りエレガントな方法ではありません。また、globals()は書き換えられてもlocals()は書き換えられないので、ローカル変数には効きません。

 そこで2番目の対策を考えます。

対策2:早めにPoolを作る

説明

 上のコードでaが作られる以前にPoolを作れば、その時点でforkするのでメモリどか食い現象は回避できます。

 こんな感じですね。

import subprocess
from multiprocessing import Pool

import numpy as np

def f():
    subprocess.run("ps -aux | grep [m]emory_test", shell=True)

p = Pool(1)
a = np.arange(10**7)

p.apply(f)
p.close()
p.terminate()
print(a.shape)
username      7525  0.0  1.2 543952 100116 pts/0   Sl+  20:31   0:00 python memory_test.py
username      7529  0.0  0.2 269216 16512 pts/0    S+   20:31   0:00 python memory_test.py
(10000000,)

 initializerで消すのと同等の効果がありますが、こちらだとローカル変数でも大丈夫です。また、グローバル変数をdelする方法だと、initializerが走るまでの一瞬の間は無駄なデータがメモリを消費する訳で、そういう面でもこちらの方が有利だと思います。

まとめ

 早めに(重いデータがメモリに読み込まれる前に)forkしておくのが基本ですが、どうしても駄目なときは削除も試してみましょう。