上文https://www.jianshu.com/p/9736a1bc4ac6
简要叙述了如何使用gunicorn+flask将深度学习模型服务化。gunicorn的workers通常具有相同的启动参数，这对于单纯运算在CPU上的进程实例来说没什么问题，因为每个进程的功能都是一模一样的。
而在深度学习服务场景，需要CPU与GPU协同运算，在具有多个显卡的服务器中，要将计算进程均匀绑定不同的GPU，最大化利用显卡资源。在gunicorn+flask+mxnet的架构中，要实现这个需求，需要解决两个问题
1）如何为gunicorn的每个worker分配一个独立的特征（ID）。通过gunicorn server hooks实现。
2）获取显卡数量。参见https://www.jianshu.com/p/6753ffee711d
server hooks分配worker_id
编写gunicorn.conf.py配置文件，除了基本配置参数外，添加下面的钩子函数。
def on_starting(server):
server._worker_id_overload = set()
def nworkers_changed(server, new_value, old_value):
server._worker_id_current_workers = new_value
def _next_worker_id(server):
if server._worker_id_overload:
return server._worker_id_overload.pop()
in_use = set(w._worker_id for w in server.WORKERS.values() if w.alive)
free = set(range(0, server._worker_id_current_workers)) - in_use
return free.pop()
def on_reload(server):
server._worker_id_overload = set(range(0, server.cfg.workers))
def pre_fork(server, worker):
worker._worker_id = _next_worker_id(server)
def post_fork(server, worker):
os.environ['IMAGE_APP_WORKER_ID'] = str(worker._worker_id)
在pre_fork()钩子中，为每个worker分配了一个唯一的worker_id；在post_fork()钩子中，将该id写入环境变量，worker启动后读取环境变量即可；
启动脚本根据worker_id和显卡数量，分配显卡id
flask启动入口 server.py
import sys
import logging
import traceback
import os
import nvgpu
from flask import Flask
from flask_restful import Api
from wdpredictor import wdplib
from wdpredictor import flask_resource
app = Flask(__name__)
api = Api(app)
worker_id = int(os.environ.get('IMAGE_APP_WORKER_ID',0))
len(nvgpu.available_gpus())))
image_predictors = wdplib.InitImagePredictors(worker_id % len(nvgpu.available_gpus()))
api.add_resource(flask_resource.ImageRemotePrediction, '/image', resource_class_args=(image_predictors, ))
if __name__ == '__main__':
app.run('0.0.0.0', debug=True)
flask启动代码将worker_id 对显卡数量取模，实现了mxnet实例平均绑定至不同显卡的需求。
网友评论