🦇

Pythonのウェブフレームワークで実行時間のプロファイリングを取得し可視化する

2022/12/15に公開

概要

こんにちは、UbieでMachine Learning Engineerをしているyagです。最近では業務で機械学習に触れる機会が少なくなり、日々BigQueryやSQLと戯れております。

さて、昨日のUbie Advent CalendarにてUbieがGoとNode.jsの会社になるという記事が界隈で話題になりましたが、機械学習周りの実装は依然としてPython一択という状況で、例外として生き残りそうです。現在でもPythonで書かれた機械学習のWeb APIエンドポイントが幾つも存在しており、私も開発に携わっているWeb APIの改善やメンテナンスをしております。今回はそのレポジトリで起きた話と、問題の特定のために行ったプロファイリングを紹介したいと思います。やることは本記事のタイトルにある通りです。

特定のリクエストがタイムアウトする問題

ある朝起きて仕事をしていると、Sentryから1通の通知がslackに届きました。調べてみると、自分の管理するAPIが何かの理由で処理を完了できずgunicornのタイムアウトが発生していました。リトライにも失敗しているようで、なんやかんやあって(ここは中略)よくよく調べてみると、そのリクエストにおいてあるリストの値が多いときに、処理に莫大な時間がかかってしまうという問題がありました。よくある組合せ爆発的なヤツですね。

幸いにもコードから計算に時間がかかっている部分の箇所を特定することは容易でしたが、「推測するな、計測せよ」という格言のもと、パフォーマンス修正が正しく行われたか、また修正前後の比較のためにも、プロファイリングをきちんと取ることにしました。

Flaskのプロファイラ

今回のAPIはFlaskをウェブフレームワークに利用していたので、Flaskの例を紹介します。FastAPIなどにも似たようなプロファイラを利用するツールやパッケージは存在します。

基本的な使い方

さて、FlaskはwerkzeugというWSGI準拠のウェブフレームワークを利用しており、これにはプロファイラが付属しています。そのため、特にパッケージをインストールする必要はなく、これを読み込んでFlaskのappにセットアップすれば完了です。ちなみに裏側はPythonのプロファイラであるcProfileです。

from werkzeug.middleware.profiler import ProfilerMiddleware
app.config['PROFILE'] = True
app.wsgi_app = ProfilerMiddleware(app.wsgi_app)

サンプルとして単にHTMLを返すAPIだけだと一瞬で終わってしまいプロファイリングの意味がないので、わかりやすくするために乱数を生成したりsleepを実行したりしてみます。

@app.route("/")
def hello_world():
    _ = [random.random() for _ in range(100000)]
    time.sleep(5)

上記のエンドポイントにアクセスすると、以下のような結果が標準出力に出力されます。

--------------------------------------------------------------------------------
PATH: '/'
         100609 function calls (100593 primitive calls) in 5.018 seconds

   Ordered by: internal time, call count

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        1    5.002    5.002    5.002    5.002 {built-in method time.sleep}
        1    0.008    0.008    0.011    0.011 /Users/yag_ays/dev/playground/20221214flask/app.py:14(<listcomp>)
   100000    0.003    0.000    0.003    0.000 {method 'random' of '_random.Random' objects}
[...]

この結果を見ると、built-in methodのtime.sleepで5秒経過していたり、randomを100,000回実行していることが直感的に分かりますね。

実行時間の内訳を可視化する

上記の結果だけでもボトルネックの特定は十分可能ですが、可視化することによってより直感的に全体のプロファイルの規模や相対的な処理時間が理解できます。プロファイリング結果をcProfileのフォーマットでダンプして、snakevizで可視化してみましょう。

まずProfilerMiddlewareのパラメータで、出力ディレクトリとそのファイル名のフォーマットを指定します。

app.wsgi_app = ProfilerMiddleware(
  app.wsgi_app,
  restrictions=[30],
  nprofile_dir="./",
  filename_format="{method}-{path}-{time:.0f}-{elapsed:.0f}ms.prof",
)

そして、snakevizを使って可視化してみます(インストール方法は割愛します)。

$ snakeviz your-profile.prof

このように可視化できました。sleepが長すぎて乱数生成の部分がほぼ見えませんが、かすかに右下部に伸びている部分が乱数生成の処理時間に対応しています。こうした小さい領域も、画面下部の行をクリックすることでズームすることができるので、ドリルダウンして詳細を見るなどインタラクティブな分析に便利です。

まとめ

今回は、エンドポイントの実行が遅い問題に対して、プロファイラを利用して実行時間を分析し可視化する一連の手順を紹介しました。デバッグや高速化に使える技術ですので、ローカル開発時にフラグを付けると一連のプロファイリングを実行するようにしておくと便利そうです。

参考

Discussion