シグナルを使って statistical profiler をつくろう

Pythonの標準ライブラリにはプロファイラがあるが、これは deterministic profiler という種類のプロファイラだ。関数が呼び出されたときと関数から抜けるときに登録したコールバック関数を呼び出す仕組みが用意されており、そのコールバック関数で時間を測ることで、各関数の実行時間を測ることができる。

このプロファイラは、短時間の実行で正確な時間を計測できるので、マイクロプロファイリングには適している。しかし、プロファイルを実行するとプログラムの実行速度がかなり低下してしまうという問題もある。

deterministic profiler と別の方式に、 statistical profiler というものがある。
これは、何らかの周期でプログラムのどこを実行しているのかをサンプリングするものだ。実行時間の長い関数ほど何度もサンプリングされるので、サンプル数が実行時間に比例する。
この方式の利点は、サンプリング周期が十分に長ければ、アプリケーションの速度低下を最低限に抑えられることだ。

さて、「何らかの周期」だけど、これは setitimer が便利だ。実時間、CPU時間(USER)、CPU時間(USER+SYSTEM)の3種類の周期でシグナルを発行してくれる。これを使って簡単な statistical profiler を作ってみよう。

import sys
import signal
from collections import defaultdict

_counter = defaultdict(int)
_cum_counter = defaultdict(int)

def _sig_handler(signal, frame):
    code = frame.f_code
    _counter[code] += 1
    _cum_counter[code] += 1
    s = set()
    s.add(code)
    f = frame.f_back
    while f is not None:
        code = f.f_code
        if code not in s:
            _cum_counter[f.f_code] += 1
            s.add(code)
        f = f.f_back

def install(target='cpu', interval=0.01):
    if target == 'cpu':
        ttype = signal.ITIMER_PROF
        sigtype = signal.SIGPROF
    elif target == 'user':
        ttype = signal.ITIMER_VIRTUAL
        sigtype = signal.SIGVTALRM
    elif target == 'real':
        ttype = signal.ITIMER_REAL
        sigtype = signal.SIGALRM

    signal.signal(sigtype, _sig_handler)
    signal.setitimer(ttype, interval, interval)

    import atexit
    atexit.register(_show_profile)

def _print_profile(counter, total):
    items = counter.items()
    items.sort(key=lambda x: x[1], reverse=True)
    for code, count in items:
        print "%20s:%-4d %20s %f" % (code.co_filename, code.co_firstlineno,
                                     code.co_name, 100.0*count/total)

def _show_profile():
    total = sum(_counter.itervalues())
    print "== time =="
    _print_profile(_counter, total)
    print "== cumulative time =="
    _print_profile(_cum_counter, total)

def test1():
    n=0
    while n < 10**5:
        n+=1

def test2():
    n=0
    while n < 10**6:
        n+=1

def loop():
    write=sys.stdout.write
    from time import time
    while True:
        test1()
        test2()

if __name__ == '__main__':
    install()
    loop()

このスクリプトをしばらく実行したあとに Ctrl-C で止めると、次のような出力が得られる。

== time ==
         sprofile.py:58                  test2 90.737012
         sprofile.py:53                  test1 9.262988
== cumulative time ==
         sprofile.py:1                <module> 100.000000
         sprofile.py:63                   loop 100.000000
         sprofile.py:58                  test2 90.737012
         sprofile.py:53                  test1 9.262988

USR1 などのシグナルを受けたらプロファイル結果を表示するなどのカスタマイズをすれば、デーモン型のプログラムのプロファイリングが簡単にできるはずだ。

あと、 frame オブジェクトには実行中の行数も入っているので、関数の中のどの部分が重いのかが判らない場合はそれもサンプリングするといいだろう。