Pythonでメモリ使用量を改善してみる

こんにちはエンジニアののりすけです。

みんなのマーケットではPythonを使用したサービスを動かしています。以前より使用メモリが異常に大きいサービスがあるため、SREチームからなんとかしてほしいとの依頼が来ています。

今回はサンプルコードを使ってPythonのメモリプロファイルを行いながら、省メモリなアプリケーションをどのように実装するか検証したいと思います。

利用するツール

• memory_profiler
• matplotlib

上記のライブラリをpip installでインストールしておきます。

memory_profilerの基本的な使い方

まず使い方を確認します。以下のように確認したい処理にデコレータ@profileを記述します。

from memory_profiler import profile


def large_integer_list():
    return [i for i in range(0,10000000)]


@profile
def main():
    result = sum(large_integer_list())
    print(result)

if __name__ == '__main__':
    main()

上記のコードを実行すると以下のような結果が出力されます。

$ python sample1.py
49999995000000
Filename: sample1.py

Line #    Mem usage    Increment   Line Contents
================================================
     8     12.4 MiB     12.4 MiB   @profile
     9                             def main():
    10     17.0 MiB      4.6 MiB       result = sum(large_integer_list())
    11     17.0 MiB      0.0 MiB       print(result)

左にあるMem usageが該当の行が評価された時点でのメモリ使用量、Incrementが評価された事による使用量の増加を表しています。今回の例の場合、listに入っている数字の合計を算出した時点で4.6MiB増えたことになっています。

この結果を少し考えたいと思います。

もともとの想定では大きなリストを作成し、メモリを大きく使うことを想定していましたが、このプロファイル結果では4.6MiBしか増えていないように見えてしまいます。しかし、よく考えるとこの部分はsum([int]) を評価した結果が+4.6MiBであり、sum関数を処理する上で不要になった情報は削除されていっているように考えられます。

今度は時系列に推移を見たいため、memory_profilerとmatplotlibを使って実行時のメモリ使用量を可視化します。

$ mprof run sample1.py
$ mprof plot

このようにmprofコマンドを利用するとグラフ化することができます。グラフから最大400MiB近くまでメモリを使用したことがわかりました。

次はこの400MiB近くまで使用してしまうこのコードを改善します。

メモリ使用量の改善

改善案としてlarge_integer_list()関数をリスト型ではなく、ジェネレータを返すように変更します。

from memory_profiler import profile

def large_integer_generator():
    # ここをGenerator式に変更
    return (i for i in range(0,10000000))


@profile
def main():
    result = sum(large_integer_generator())
    print(result)

if __name__ == '__main__':
    main()

実行結果

mprof run sample1.py
mprof: Sampling memory every 0.1s
running as a Python program...
49999995000000
Filename: sample1.py

Line #    Mem usage    Increment   Line Contents
================================================
     8     12.8 MiB     12.8 MiB   @profile
     9                             def main():
    10     12.8 MiB      0.0 MiB       result = sum(large_integer_generator())
    11     12.8 MiB      0.0 MiB       print(result)

今度はグラフ上もメモリ使用量が劇的に減っていることがわかります。ジェネレータを使うことで一度に大量のデータを確保することなく一つずつ処理しているため省メモリで処理することができます。

しかし、グラフのx軸を確認すると20秒を超える実行時間になっています。先程のリストを利用した場合は10秒を切る実行時間だったのに対して2倍以上遅い。。。

実行速度の改善

先程のコードではメモリの使用量は改善できましたが、今度は実行速度が問題になってしまいました。原因として考えられるのはジェネレータではyieldするたびに値の評価が行われます。評価回数があまりに多くなってしまったため速度が低下したことが考えられます。

上記の問題点を改善したコードが以下になります。

from memory_profiler import profile
# Iteratorに対する便利な関数がまとまったモジュール
from itertools import (islice, chain)

# iteratorをchunkサイズに分割する関数
def chunks(iterable, chunk_size):
    iterator = iter(iterable)
    chunk = tuple(islice(iterator, chunk_size))
    while chunk:
        yield list(chunk)
        chunk = tuple(islice(iterator, chunk_size))

def large_integer_generator():
    chunk_size = 10000
    # chunk毎にsumを行う
    for chunk in chain(chunks(range(0,10000000), chunk_size)):
        yield sum(chunk)

@profile
def main():
    result = sum(large_integer_generator())
    print(result)

if __name__ == '__main__':
    main()

実行結果

mprof run sample1.py
mprof: Sampling memory every 0.1s
running as a Python program...
49999995000000
Filename: sample1.py

Line #    Mem usage    Increment   Line Contents
================================================
    19     12.6 MiB     12.6 MiB   @profile
    20                             def main():
    21     15.0 MiB      2.4 MiB       result = sum(large_integer_generator())
    22     15.0 MiB      0.0 MiB       print(result)

前回のコードよりchunk分のメモリ使用量は増えていますが、速度が劇的に改善されました。

最後に

今回の検証ではうまくメモリ使用量/実行速度の改善を行うことができました。実際のサービスで同様に改善できる部分があるかは、これから確認しなければいけませんが、今後の実装においても良い検証ができたかと思います。

我々みんなのマーケットテックチームでは「くらしのマーケット」を一緒に作る仲間を募集しています！興味がある方はぜひ気軽に連絡ください （コーポレートサイト https://www.minma.jp/ ） 次回はエンジニアのカーキくんの予定です。