Linux & Pythonで自分自身のピークのメモリ使用量を取る

結論

いろいろ調べたのですが、これよりマシな方法が見つかりませんでした。
/procファイルシステムを使っているので、サポートしているプラットフォーム(Linux等)でしか動きません。

import os
import re

def peak_memory():
    pid = os.getpid()
    with open(f'/proc/{pid}/status') as f:
        # "VmHWM:    862168 kB" の形式になっている行を抜き出す
        for line in f:
            if not line.startswith('VmHWM:'):
                continue
            return int(re.search('[0-9]+', line)[0])
    raise ValueError('Not Found')

なぜピーク使用量か?

とある関数の内部で、一時的にメモリ使用量がハネ上がる現象が起きました。

その関数をプロファイリングしてメモリ使用量がハネ上がっている箇所を突き止めたいのであればmemory-profilerを使うのが良いかと思います。
今回、どこで増えるかは分かっていて、増えるのは仕方ないが、具体的にどれくらい増えるのかを調べたかったので、ピーク使用量を知る必要がありました。

ピーク使用量を使わないとどうなるか。

現在のメモリ使用量であればpsutilで取れます。このへんを参照。

ですが、ピークではないので、例えばこういうコードのとき

現在の使用メモリを取る

import psutil

def f():
    a = [0] * 20000000

print(psutil.Process().memory_full_info())
f()
print(psutil.Process().memory_full_info())

以下のような結果になります。

あんまり変わらない

pfullmem(rss=12406784, vms=18395136, shared=6369280, text=4096, lib=0, data=6225920, dirty=0, uss=8134656, pss=9319424, swap=0)
pfullmem(rss=12619776, vms=18354176, shared=6385664, text=4096, lib=0, data=6184960, dirty=0, uss=8253440, pss=9438208, swap=0)

f()の中ではメモリをいっぱい使っているのですが、外では使っていないので、外からでは分からないんですね。

fの中でも取ってみる

import psutil

def f():
    a = [0] * 20000000
    print('inner: ', psutil.Process().memory_full_info())

print('before:', psutil.Process().memory_full_info())
f()
print('after: ', psutil.Process().memory_full_info())

この場合、以下のようになります。

真ん中だけいっぱい使っている

before: pfullmem(rss=12476416, vms=18395136, shared=6443008, text=4096, lib=0, data=6225920, dirty=0, uss=8179712, pss=9407488, swap=0)
inner:  pfullmem(rss=172519424, vms=178356224, shared=6520832, text=4096, lib=0, data=166187008, dirty=0, uss=168300544, pss=169528320, swap=0)
after:  pfullmem(rss=12754944, vms=18354176, shared=6520832, text=4096, lib=0, data=6184960, dirty=0, uss=8298496, pss=9526272, swap=0)

メモリを使っていそうなところに、ピンポイントでメモリ使用量を取るコードを入れられるなら、それでいいんですが、そうしたくなかったわけです。

今回の方法ではどうなったか?

ピーク使用量を取る

import os
import re

def peak_memory():
    pid = os.getpid()
    with open(f'/proc/{pid}/status') as f:
        for line in f:
            if not line.startswith('VmHWM:'):
                continue
            return int(re.search('[0-9]+', line)[0])
    raise ValueError('Not Found')

def f():
    a = [0] * 20000000

print('before:', peak_memory(), 'KB')
f()
print('after: ', peak_memory(), 'KB')

結果、こうなりました。

afterで増えてる

before: 9072 KB
after:  165532 KB

めでたしめでたし。

参考文献

Linux プログラムのメモリ消費量を測る

コードについて

本記事に掲載されたコードを、自身・自社が開発されているPythonスクリプトに利用いただくのはご自由にどうぞ。