uWSGI+FlaskでDeepLearningをAPI化するとスタックする問題に対して

こんにちは、けんご(@N30nnnn)です。

uWSGI+Flaskと PytorchやKerasなどのDeepLearningフレームワークを組み合わせると正常に動かない(レスポンスが返ってこない)パターンがあります。

レスポンスが正常に返ってこないパターンの再現リポジトリを作りました。 リポジトリに再現コードや、再現までのコマンドを記しています。

これを元に簡単に記事を書きます。

github.com

問題

通常、機械学習アプリケーションを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の事例で確認されてます。

原因と対策

どうやらこの問題は、uWSGIによってForkされたWorkerからマルチスレッドでグローバルのモデルを参照していることが問題の様で、推定ハンドラ内でリクエストのたびに都度モデルをロードする場合は発現しません。しかし、リクエストのたびにロードするのはレスポンスタイムを考えるとナンセンスです。それならばアプリ起動時に予め各プロセスのためにモデルをロードしておきましょう。

uWSGIはデフォルトではPreforkモデルです。 すなわち、アプリのロードが1回のみ行われて、各WorkerにはForkが配られます。 一方で各Workerで別個にアプリをロードするモードは lazy-apps と呼ばれています。

起動方法は

$ uwsgi --lazy-apps

または

iniファイル内で

[uwsgi]
lazy-apps = true

です。

ちなみに、似たような設定で lazy モードがありますが、過剰に設定を書き換えてしまう様で、非推奨の様です。