python の tracemalloc / psutil でメモリ使用量を確認

やりたいこと

python プログラムでメモリを消費している箇所を特定する。
ここでは、tracemalloc と psutil でを使用してメモリ使用量を確認する。

tracemalloc

tracemalloc で snapshot をとると、snapshot 間でのメモリ増減の差分を取得することができる。
また、現在のメモリ使用量、ピーク時のメモリ使用量を取得することもできる。
なお、tracemalloc によるメモリ使用量を確認するには、tracemalloc.start() を実行しておく必要がある。

psutil

psutil では、rss、vms 等のメモリ使用量を取得することができる。
インストール方法は以下の通り。

pip install psutil

プログラム例

test1.py

import os
import psutil
import tracemalloc


def main():
    tracemalloc.start()  # tracemalloc の開始

    snapshot_before = tracemalloc.take_snapshot()  # メモリ使用前に snapshot を取得

    for i in range(0, 100):
        a = [1] * 1000
        b = [2] * 2000
        c = [3] * 3000
        d = [4] * 4000
        e = [5] * 5000

    snapshot_after = tracemalloc.take_snapshot()  # メモリ使用後に snapshot を取得

    # snapshot の差分で増加が大きい 3 行を表示
    all_stats = snapshot_after.compare_to(snapshot_before, 'lineno')
    inc_stats = [stat for stat in all_stats if stat.size_diff > 0]
    print("\nsnapshot diff")
    for stat in inc_stats[:3]:
        print(stat)

    # 現在、ピークのメモリ使用量表示
    cur, peak = tracemalloc.get_traced_memory()
    print("\ncurrent, peak memory")
    print(f"current: {cur}")
    print(f"peak:    {peak}")

    # psutil によるメモリ使用量表示
    process = psutil.Process(os.getpid())
    mem_info = process.memory_info()
    print("\nmemory")
    print(mem_info)

    return 0


if __name__ == '__main__':
    res = main()
    exit(res)

実行結果例

$ python test1.py

snapshot diff
test1.py:16: size=39.1 KiB (+39.1 KiB), count=1 (+1), average=39.1 KiB
test1.py:15: size=31.2 KiB (+31.2 KiB), count=1 (+1), average=31.2 KiB
test1.py:14: size=23.4 KiB (+23.4 KiB), count=1 (+1), average=23.4 KiB

current, peak memory
current: 123176
peak:    160672

memory
pmem(rss=14548992, vms=18509824, shared=7602176, text=4096, lib=0, data=8527872, dirty=0)

snapshot の diff ではメモリ増加の上位 16-14行目は、それぞれ e = [5] * 5000、d = [4] * 4000、c = [3] * 3000 となっており、メモリ使用量が大きい行を特定できている。