uWSGI+FlaskでDeepLearningをAPI化するとスタックする問題に対して
こんにちは、けんご(@N30nnnn)です。
uWSGI+Flaskと PytorchやKerasなどのDeepLearningフレームワークを組み合わせると正常に動かない(レスポンスが返ってこない)パターンがあります。
レスポンスが正常に返ってこないパターンの再現リポジトリを作りました。 リポジトリに再現コードや、再現までのコマンドを記しています。
これを元に簡単に記事を書きます。
問題
通常、機械学習アプリケーションをAPI化する場合、グローバル変数にモデルを事前ロードし、推定ハンドラでそれを参照する形を取ると思います。
# https://github.com/keng000/uwsgi-flask-ml/blob/master/myapp/api/wsgi.py import torch from flask import Flask from myapp.ml.model import Net from myapp.path import PathManager app = Flask(__name__) # Model pre-loading net = Net() model_path = PathManager.MODELS / "model.pth" net.load_state_dict(torch.load(model_path, map_location=lambda storage, loc: storage)) net.eval() @app.route("/estimate", methods=["GET"]) def estimate(): dummy = torch.rand(1, 3, 32, 32) net(dummy) # <- here, the process will stuck return "OK", 200
この形をとって、 app
をuWSGIのデフォルト設定でロードすると、推定処理(コメント付与箇所)のところでスタックします。
自分の場合では、 Pytorch 0.4.0
まで動作していたところ、 1.1.0
で突然スタックする様になりました。
この問題はPytorchとKerasの事例で確認されてます。
- nginx - Pytorch model prediction in production with uwsgi - Stack Overflow
- Flask + Tensorflow + Keras で モデルからのpredictが動作しない問題 - DEV Community 👩💻👨💻
- Tensorflow backend - bug in model._make_predict_function(...) · Issue #2397 · keras-team/keras · GitHub
原因と対策
どうやらこの問題は、uWSGIによってForkされたWorkerからマルチスレッドでグローバルのモデルを参照していることが問題の様で、推定ハンドラ内でリクエストのたびに都度モデルをロードする場合は発現しません。しかし、リクエストのたびにロードするのはレスポンスタイムを考えるとナンセンスです。それならばアプリ起動時に予め各プロセスのためにモデルをロードしておきましょう。
uWSGIはデフォルトではPreforkモデルです。
すなわち、アプリのロードが1回のみ行われて、各WorkerにはForkが配られます。
一方で各Workerで別個にアプリをロードするモードは lazy-apps
と呼ばれています。
起動方法は
$ uwsgi --lazy-apps
または
iniファイル内で
[uwsgi] lazy-apps = true
です。
ちなみに、似たような設定で lazy
モードがありますが、過剰に設定を書き換えてしまう様で、非推奨の様です。