[Python] pytestで処理の遅い箇所を可視化して特定する

プロジェクトのテスト数が増えてくるとCI回すのも遅くなるし、ローカルでこまめにテスト回すぞっていうときもストレスは増えていきます。テストが１万件もあればしょうがないかな、という気にもなりますが、時に「あれ、なんかやたら遅いテストケースがあるぞ？」ということもあります。テストが遅い原因がプロダクトコードにあるのであればより問題です。
pythonのテストフレームワークはpytestを使っているので、以下pytestの場合はこんなことして遅い処理を探し出しているという備忘録。

主にやったことは以下の2つ。

① pytestの--durationsオプションで処理時間の大きいテストケースを出力する
② pytest-profilingでプロファイルして処理の遅い箇所を特定

テストコード

本記事で使用したテストコードです。1～10000000の整数を順に突っ込んだリストを4パターンの方法で作成して最後にassertで比較しています。
setup()のなかでスリープを読んでいるのはテスト処理時間をアウトプットする場合にわかりやすくするために入れてるだけで処理的には意味はありません。

ソースコードはココにおいてあります。

test_sample.py

import time
import pytest
import sys


COUNT = 10000000

@pytest.fixture(scope='module')
def expected():
  # assertで比較する期待値の作成
  return [i for i in range(1, COUNT)]

@pytest.fixture(scope='function', autouse=True)
def setup():
  # 処理的には意味のないスリーブ
  time.sleep(0.1)
  yield
  time.sleep(0.2)

def test_1_1(expected):
  actual = [i for i in range(1, COUNT)]
  assert expected == actual


def test_1_2(expected):
  actual = list({i for i in range(1, COUNT)})
  assert expected == actual


def test_1_3(expected):
  actual = []
  for i in range(1, COUNT):
    actual.append(i)
  assert expected == actual


def test_1_4(expected):
  actual = []
  for i in range(1, COUNT):
    # actual = actual + [i] とすると死ぬほど時間かかる
    actual += [i]
  assert expected == actual

どのテストが遅いか割り出す

pytestの実行時に--durations=Nというオプションをつけると、最も遅いテスト＋前後の処理(setup/teardown)N個を実行結果に表示してくれます。
N=0とするとすべての結果を出力します。

--durations=0を指定してpytestを実行した結果が以下です。なお、デフォルトでは0.01s以下の結果は非表示になりますが、-vvオプションをつけるとすべての結果が表示されます。

=========== slowest test durations =========== 
2.13s call     tests/test_sample.py::test_1_3
1.25s call     tests/test_sample.py::test_1_4
1.08s call     tests/test_sample.py::test_1_2
0.81s call     tests/test_sample.py::test_1_1
0.66s setup    tests/test_sample.py::test_1_1
0.20s teardown tests/test_sample.py::test_1_2
0.20s teardown tests/test_sample.py::test_1_4
0.20s teardown tests/test_sample.py::test_1_3
0.20s teardown tests/test_sample.py::test_1_1
0.10s setup    tests/test_sample.py::test_1_4
0.10s setup    tests/test_sample.py::test_1_2
0.10s setup    tests/test_sample.py::test_1_3
============= 4 passed in 7.40s =============

上位4つは4つのテストケースの実行結果になっています。シンプルに内包表記で書くのが速いですね。
続いてtest_1_1()の前処理(setup)に0.66秒程かかっています。これはexpected()の処理時間＋setup()内のスリープ0.1秒が加わったものだと思われます。
続いて各テストの後処理(teardown)が4つ、test_1_1()以外のsetup処理が続きます。これはsetup()に記述したスリープ時間そのままですね。

テスト内のどの処理が遅いか割り出す

どのテストが遅いかまで特定した後は速度改善のためにテスト内のどの処理が遅いかを特定する必要があります。
ネットの英知によればpytestとプロファイラを同時に実行する方法が出てきますが、プロファイルの出力結果をpstatsなどを使って別途分析してやる必要があります。
この辺を簡単にできないか調べていたところ辿り着いたのがpytest-profilingです。

pytest-profiling

pipでインストールできます。

pip install pytest-profiling

またプロファイルの結果をSVGとして出力することが出来ますが、Graphvizが必要なので別途インストールしておいてください。

使い方は簡単でpytest実行時に--profile-svgオプションをつけるだけです。

pytest tests/test_sample.py --profile-svg

--profile-svgオプションをつけて実行するとpytest-profilingは以下の処理を行います。

1. pytest及びcProfileの実行
2. プロファイル結果をpstatsで解析
3. gprof2dotで解析結果をグラフ(dot file)に変換
4. dotコマンド(Graphviz)でグラフをSVG画像に変換

なお--profile-svgオプションの代わりに--profileをつけると、上記の3番の処理まで行うようです。

一連の処理が終わるとprofディレクトリが作成され、その下にpstats解析情報のprofファイルとSVG画像が作成されます。
profディレクトリ内には以下のようなファイルが作成されます。

prof
├ combined.prof
├ combined.svg
├ test_1_1.prof
├ test_1_2.prof
├ test_1_3.prof
└ test_1_4.prof

test_1_1.prof～test_1_4.profは各テストケースごとのプロファイル解析情報です。ファイル名にはテストの名称がつけられますが、日本語などはアンダースコアに置き換えられて作成されます。
combined.profは実行したすべてのテストに対するプロファイル解析情報です。combined.svgはそれをグラフに変換してさらにSVG画像に変換したもので、以下のような画像になります。

gprof2dotのREADMEにグラフの見方が解説されています。
グラフの各ノードは以下のような情報が記載されています。

+------------------------------+
|        function name         |
| total time % ( self time % ) |
|         total calls          |
+------------------------------+

• function name … 関数名
• total time % … 全体の実行時間に対する、この関数及びすべての子関数の実行時間の割合
• self time % … 全体の処理時間に対する、この関数のみの実行時間の割合
• total calls … この関数が呼び出された総会数(再帰呼び出し含む)

また呼び出し元の関数と呼び出された関数（親と子と表現される）をつなぐエッジには以下のような情報が記載されています。

           total time %
              calls
parent --------------------> children

• total time % … 子から親に対して転送された実行時間の割合
• calls … 親が子を読んだ回数

最上位の親ノードからtotal time % の大きい順にたどっていけば、簡単に処理の遅い箇所を特定することが出来そうです。

課題：windowsだとうまく動かない

最初Windowsで動かしていましたが、profファイルは生成されるのにSVGファイルが作成されませんでした。よくよく見るとSupported OSにWindowsが含まれていないようです。

ただprofファイルは生成されるので、profファイルからSVGまでの変換を手動で行うことは可能でした。
gprof2dotはpytest-profilingをインストールすると合わせてインストールされると思います。

gprof2dot -f pstats prof/combined.prof > prof/tmp
dot -Tsvg -o prof/combined.svg prof/tmp

ちなみにVisualStudioCodeのターミナルから実行しましたがdotコマンドの実行にエラーが出ました。既定のシェルをpowershellにしているとうまくいかないようです。「既定のシェルの選択」からcmdに変えるとうまくいきました。

最後に

pytestのconfigurationファイルに以下のような記載をしてプロジェクトのルートディレクトリ配下においておけばpytest実行時にpytest-profilingが実行されます。

pytest.ini

[pytest]
testpaths = tests
python_files = test_*.py
addopts = -vv --durations=0 --profile-svg
pytest_plugins = ['pytest_profiling']

ただテストケースが多いと解析情報も煩雑になると思うので、まずは--durationsオプションでテストの遅い箇所を特定、その後該当するテストにのみpytest-profilingを実行して処理の遅い箇所を特定する、としたほうがやりやすい気はします。configurationファイルに書くと毎度実行されてしまいますしね。