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しておくのが基本ですが、どうしても駄目なときは削除も試してみましょう。
追記
プロセスの開始方法に"spawn"を指定することでも可能だとコメントでご指摘をいただきました。
spawn
親プロセスは新たに python インタープリタープロセスを開始します。子プロセスはプロセスオブジェクトの run() メソッドの実行に必要なリソースのみ継承します。特に、親プロセスからの不要なファイル記述子とハンドルは継承されません。この方式を使用したプロセスの開始は fork や forkserver に比べ遅くなります。Unix と Windows で利用可能。Windows でのデフォルト。
multiprocessing.get_context('spawn')とすると、multiprocessingモジュールと同じAPIを持つオブジェクトが返り、これからPoolを作ることで解決できます。これを利用しても良さそうです。
