pytest-monitorとpytest-memrayを用いてソフトウェアのボトルネックを発見する
本記事を閲覧いただきありがとうございます。シンプルフォーム社のエンジニアの駒井 (@rindybell) です。本記事は SimpleForm Advent Calendar 2023 の２３日目の記事となっています。
本記事では、ソフトウェアのボトルネック発見に有用な、pytest-monitor と pytest-memray というライブラリを紹介します。
1. 背景
プロダクトマネージャといったステークホルダがソフトウェアを評価するとき、「何らかの検索をできるか」とか「期待するデータが表示されるか」といった、機能面をしばしば重視します。ただし、運用コストに影響するため、メモリやＣＰＵの効率性といった非機能の観点も、本来は重要な指標です。
しかしながら、ソフトウェアにおいて支障をきたす何らかのボトルネックがあるとき、そのボトルネックを同定することはしばしば容易ではありません。効率的なソフトウェアを構築するためにも、メモリやＣＰＵの効率性を機械的に評価できることは非常に重要です。
そこで、pytest-monitor や pytest-memray という、単体テストのパフォーマンスを評価できるライブラリについて紹介します。pytest-memray につきましては、本家のmemrayの pytest 向けライブラリのようです。
2. 環境構築と実験設定
同アドベントカレンダーの別記事グラフデータベースのためのテスト基盤を構築し、開発サイクルを高速化してみたにて作成した検証用のリポジトリにて実験します。
今回は、上記のリポジトリ上のソフトウェアに対して、ボトルネックを人工的に追加し、各種ライブラリが本箇所を同定できるかどうかを検証します。まず、以下のような断片をボトルネックとしてソースコードに埋め込みます。埋め込み先ですが、src.graphs.router.read_graphsとしています。テストコードの観点ではtests.test_router.test_getが対応します。
import numpy as np
x = np.random.random((10000, 10000)).astype(dtype=np.float64)
x.dot(x)
次に、各種ライブラリをインストールします。
$ pip install pytest-memray
$ pip install pytest-monitor
3. 実験
pytest-monitor
pytest-monitor は単体テストのパフォーマンスを参照できるライブラリです。pytest-monitor をインストールしている状態で、単体テスト（pytest）を実行するときに--dbオプションを有効にすることで pytest-monitor を使用できます。コマンド例を次に示します。ここで、メトリクスの出力先をファイルパスの形式で指定できますが、次の例ではpytest-monitor.dbとしています。
$ python -m pytest --db pytest-monitor.db
上記コマンドを実行後、pytest-monitor.dbファイルが生成されます。続いて、メトリクスを確認するために sqlite3 コマンドを実行します。
$ sqlite3 pytest-monitor.db
SQLite version 3.37.2 2022-01-06 13:25:41
Enter ".help" for usage hints.
sqlite>
sqlite3 にて、出力形式を変更します。
sqlite> .mode column
sqlite>
どのようなメトリクスがあるかの説明については公式のドキュメントを参照いただきたいですが、ここでは、ＣＰＵやメモリについての指標を参照します。
sqlite> select ITEM_PATH,ITEM,TOTAL_TIME,CPU_USAGE,MEM_USAGE from TEST_METRICS order by MEM_USAGE desc;
ITEM_PATH ITEM TOTAL_TIME CPU_USAGE MEM_USAGE
----------------- ------------------ ------------------ ----------------- -----------
tests.test_router test_get 9.27683258056641 11.9297183644165 1573.53125
tests.test_router test_post_vertices 0.078852653503418 1.01455051219723 52.18359375
tests.test_router test_post_edges 0.098275899887085 1.0175434680822 41.125
test_daos test_create_edge 0.0947296619415283 0.844508450261103 39.16796875
test_daos test_create_vertex 0.0755016803741455 0.927131683071422 39.16015625
test_daos test_create_edge 0.0837421417236328 1.07472770755048 39.16015625
test_daos test_read_graphs 0.0671150684356689 0.893987019584254 39.15625
（省略）
上記に示す通り、tests.test_router の test_get においてＣＰＵやメモリの観点が異常値を示していることが分かります。この結果から、pytest-monitor はソースコードのボトルネックを同定することに役立ちそうです。
pytest-memray
次に pytest-memray について説明します。pytest-memray は pytest 向けのメモリプロファイラです。pytest 時に、--memrayオプションを追加することで利用できます。
$ python -m pytest --memray tests
（中略）
📦 Total memory allocated: 1.6GiB
📏 Total allocations: 44489
📊 Histogram of allocation sizes: |▁ █ ▁ |
🥇 Biggest allocating functions:
- read_graphs:/home/rindybell/workspace/simpleform/simpleform/PROJECTS/2023_ac_example_graphdb_application/src/graphs/router.py:46 -> 795.4MiB
- read_graphs:/home/rindybell/workspace/simpleform/simpleform/PROJECTS/2023_ac_example_graphdb_application/src/graphs/router.py:45 -> 762.9MiB
- _call_with_frames_removed:<frozen importlib._bootstrap>:219 -> 1.1MiB
- _call_with_frames_removed:<frozen importlib._bootstrap>:219 -> 372.7KiB
上記の出力には、割り当てられたメモリの総量が記されています。例えば router.py の 45 行目や 46 行目において、おおよそ 800MB のメモリが割り当てられていることが分かります。埋め込みした断片は、サイズが(10000,10000)の行列を定義しているわけですが、それぞれの成分を 64bit で表しているため、800MB のメモリが割り当てられていることについては納得感があります。
Tips: テストに対してメモリの制約
pytest-memray は、テストコードにおいて利用されるメモリについて制限を課すことが可能です。制限を課す時、@pytest.mark.limit_memoryというデコレータを付与します。試しに、24MB をメモリの上限として設定します。
@pytest.mark.asyncio
@pytest.mark.parametrize("expected, end_point, params", data_for_test_get)
@pytest.mark.limit_memory("24 MB")
async def test_get(self, client, gremlin_service, expected, end_point, params):
response = await client.get(end_point, json=params)
assert expected == response.status_code
上記の設定を追加し、再度テストを実行します。
$ python -m pytest --memray tests
（中略）
=========================================================================================== short test summary info ============================================================================================
MEMORY PROBLEMS tests/test_router.py::TestRouter::test_get[200-/v1/graphs-params0]
================================================================================== 1 failed, 10 passed, 60 warnings in 28.98s ==================================================================================
すると、メモリの上限を超えた場合に、テストが failed として扱われていることが分かります。以上より、pytest-memray も効率性を考慮したソフトウェア開発に貢献してくれそうです。
4. おわりに
本記事では、pytest-monitor や pytest-memray といったライブラリの使用例を紹介しました。いかがでしたでしょうか。普段の開発では、メモリの効率性を意識しながらソースコードを記述することが難しい場合があり、本ライブラリのような非機能面のメトリクスを参照できるライブラリは非常に重要と考えています。本記事が少しでも読者の皆様のお力になれたら幸いです。
Appendix
本記事をお読みいただきましてありがとうございました。明日以降の弊社のアドベントカレンダーもどうぞよろしくお願いいたします。また、Open Data Labという団体にて、大規模言語モデル入門・輪読会 #2 の開催しておりますので、そちらも併せてご確認いただければ幸いです。
Discussion