开源项目 | 五分钟搭建BERT服务，实现1000+QPS

作者丨刘欣

单位丨香侬科技算法架构负责人

研究方向丨NLP工程化、算法平台架构

深度学习模型在训练和测试时，通常使用小批量（mini-batch）的方式将样本组装在一起，这样能充分利用 GPU 的并行计算特性，加快运算速度。 

但在将使用了深度学习模型的服务部署上线时，由于用户请求通常是离散和单次的，若采取传统的循环服务器或多线程服务器，在短时间内有大量请求时，会造成 GPU 计算资源闲置，用户等待时间线性变长。 

基于此，我们开发了 service-streamer，它是一个中间件，将服务请求排队组成一个完整的 batch，再送进 GPU 运算。这样可以牺牲最小的时延（默认最大 0.1s），提升整体性能，极大优化 GPU 利用率。

Github开源链接：

https://github.com/ShannonAI/service-streamer

功能特色

安装步骤

可通过 pip 安装，要求 Python>=3.5：

pip installservice_streamer

五分钟搭建BERT服务

为了演示 API 使用方法，service-streamer 提供了一个完整的教程和示例代码。如何在五分钟搭建起基于 BERT 模型的完形填空服务，每秒处理 1000+ 请求。点击文章左下角“阅读原文”查看完整代码。
1. 首先我们定义一个完型填空模型（bert_model.py），其 predict 方法接受批量的句子，并给出每个句子中 [MASK] 位置的预测结果。

class TextInfillingModel(object);       ...batch=［"twinkle twinkle [MASK] star",            "Happy birthday to [MASK]",            'the answer to life, the [MASK], and everything']model=TextaInfillingModel()outputs=model.predict(batch)print(outputs)#['little', 'you', 'universe' ]

2. 然后使用 Flask 将模型封装成 web 服务 flask_example.py。这时候你的 web 服务每秒钟只能完成 12 句请求。

model=TextInfillingModel()@app.route("/naive", methods=["POST"])def naive_predict( ):      inputs = request.form.getlist("s")      outputs = model.predict(inputs)      return jsonify(outputs)app.run(port=5005)@app.route("/naive", methods=["POST"])
def naive_predict( ):
      inputs = request.form.getlist("s")
      outputs = model.predict(inputs)
      return jsonify(outputs)

app.run(port=5005)

3. 下面我们通过 service_streamer 封装你的模型函数，三行代码使 BERT 服务的预测速度达到每秒 200+ 句（16 倍 QPS）。

from service_streamer import ThreadStreamer streamer= ThreadedStreamer (model.predict,batch_size=64, max_latency=0.1)@app.route("/stream", methods=["POST"])def stream_predict( ):     inputs = request.form.getlist("s")    outputs = streamer.predict(inputs)    return isonify(outputs)app.run(port=5005, debug=False)import ThreadStreamer streamer= ThreadedStreamer (model.predict,batch_size=64, max_latency=0.1)

@app.route("/stream", methods=["POST"])
def stream_predict( ):
     inputs = request.form.getlist("s")
    outputs = streamer.predict(inputs)
    return isonify(outputs)

app.run(port=5005, debug=False)

4. 最后，我们利用 Streamer 封装模型，启动多个 GPU worker，充分利用多卡性能实现每秒 1000+ 句（80 倍 QPS）。

import multiprocessingfrom service_streamer import ManagedModel, Streamermultiprocessing.set_start_method("spawn", force=True)class ManagedBertModel(ManagedModel):     def init_model(self):           self.model = TextInfillingModel( )      def predict(self, batch):            return self.model.predict(batch)streamer =Streamer(ManagedBertModel, batch_size=64, max_latency=0.1, worker_num = 8, cuda_devices=(0,1,2,3))app.run(port=5005, debug=False)
from service_streamer import ManagedModel, Streamer
multiprocessing.set_start_method("spawn", force=True)

class ManagedBertModel(ManagedModel):

     def init_model(self):
           self.model = TextInfillingModel( )

      def predict(self, batch):
            return self.model.predict(batch)

streamer =Streamer(ManagedBertModel, batch_size=64, max_latency=0.1, 
worker_num = 8, cuda_devices=(0,1,2,3))

app.run(port=5005, debug=False)

运行 flask_multigpu_example.py 这样即可启动 8 个 GPU worker，平均分配在 4 张卡上。

更多指南

除了上面的 5 分钟教程，service-streamer 还提供了： 

API介绍

快速入门
通常深度学习的 inference 按 batch 输入会比较快。

outputs = model.predict(batch_inputs)

用 service_streamer 中间件封装 predict 函数，将 request 排队成一个完整的 batch，再送进 GPU。牺牲一定的时延（默认最大 0.1s），提升整体性能，极大提高 GPU 利用率。

from service_streamer import ThreadedStreamer# 用Streamer封装batch_predict函数streamer = ThreadedStreamer(model.predict, batch_size=64, max_latency=0.1)# 用Streamer异步调用predict函数outputs = streamer.predict(batch_inouts)import ThreadedStreamer

# 用Streamer封装batch_predict函数
streamer = ThreadedStreamer(model.predict, batch_size=64, max_latency=0.1)

# 用Streamer异步调用predict函数
outputs = streamer.predict(batch_inouts)

然后你的 web server 需要开启多线程（或协程）即可。 
短短几行代码，通常可以实现数十（batch_size/batch_per_request）倍的加速。
分布式GPU worker 
上面的例子是在 web server 进程中，开启子线程作为 GPU worker 进行 batch predict，用线程间队列进行通信和排队。 
实际项目中 web server 的性能（QPS）远高于 GPU 模型的性能，所以我们支持一个 web server 搭配多个 GPU worker 进程。

import multiprocessing; multiprocessing.set_start_method("spawn", force=True)from service_streamer import Streamer# spawn出4个gpu worker进程streamer = Streamer（model.predict, 64, 0.1, worker_num=4）outputs = streamer.redict(batch)
multiprocessing.set_start_method("spawn", force=True)
from service_streamer import Streamer

# spawn出4个gpu worker进程
streamer = Streamer（model.predict, 64, 0.1, worker_num=4）
outputs = streamer.redict(batch)

Streamer 默认采用 spawn 子进程运行 gpu worker，利用进程间队列进行通信和排队，将大量的请求分配到多个 worker 中处理，再将模型 batch predict 的结果传回到对应的 web server，并且返回到对应的 http response。

上面这种方式定义简单，但是主进程初始化模型，多占了一份显存，并且模型只能运行在同一块 GPU 上，所以我们提供了 ManageModel 类，方便模型 lazy 初始化和迁移，以支持多 GPU。

from service_streamer import ManagedModelclass ManagedBertModel(ManagedModel):   def init_model(self):       self.model = Model( )  def predict(self, batch):      return self.model.predict(batch) # spawn出4个gpu worker进程，平均分数在0/1/2/3号GPU上streamer = Streamer(ManagedBertModel, 64, 0.1, worker_num=4,cuda_devices=(0,1,2,3))outputs = streamer.predict(batch)import ManagedModel

class ManagedBertModel(ManagedModel):

   def init_model(self):
       self.model = Model( )

  def predict(self, batch):
      return self.model.predict(batch) 

# spawn出4个gpu worker进程，平均分数在0/1/2/3号GPU上
streamer = Streamer(ManagedBertModel, 64, 0.1, worker_num=4,cuda_devices=(0,1,2,3))
outputs = streamer.predict(batch)

分布式web server 
有时候，你的 web server 中需要进行一些 CPU 密集型计算，比如图像、文本预处理，再分配到 GPU worker 进入模型。CPU 资源往往会成为性能瓶颈，于是我们也提供了多 web server 搭配（单个或多个）GPU worker 的模式。
使用 RedisStreamer 指定所有 web server 和 GPU worker 共用的 redis broker 地址。

# 默认参数可以省略，使用localhost:6379streamer = RedisStreamer(redis_broker="172.22.22.22:6379")
streamer = RedisStreamer
(redis_broker="172.22.22.22:6379")

然后跟任意 python web server 的部署一样，用 gunicorn 或 uwsgi 实现反向代理和负载均衡。

cd examplegunicorn -c redis_streamer_gunicorn.py flask_example:app

这样每个请求会负载均衡到每个 web server 中进行 CPU 预处理，然后均匀的分布到 GPU worke 中进行模型 predict。

Future API

如果你使用过任意 concurrent 库，应该对 future 不陌生。当你的使用场景不是 web service，又想使用 service_streamer 进行排队或者分布式 GPU 计算，可以直接使用 Future API。

from service_streamer import ThreadedStreamerstreamer = ThreadedStreamer(model.predict, 64, 0.1)xs ={}for i in range(200):    future = streamer.submit(["Happy birthday to [MASK]",                                            "Today is my lucky [MASK]"])     xs.append(future)# 先拿到所有future对象，再等待异步返回for future in xs:     outputs = future.result()     print(outputs)import ThreadedStreamer
streamer = ThreadedStreamer(model.predict, 64, 0.1)

xs ={}
for i in range(200):
    future = streamer.submit(["Happy birthday to [MASK]", 
                                           "Today is my lucky [MASK]"])
     xs.append(future)

# 先拿到所有future对象，再等待异步返回
for future in xs:
     outputs = future.result()
     print(outputs)

基准测试

如何做基准测试
我们使用 wrk 来使做基准测试。 
环境

单个GPU进程

# start flask threaded serverpython example/flask_example.py# benchmark naive api without service_streamer./wrk -t 4 -c 128 -d 20s --timeout=10s -s scripts/streamer.lua http://127.0.0.1:5005/naive# benchmark stream api with service_streamer./wrk -t 4 -c 128 -d 20s --timeout=10s -s scripts/streamer.lua http://127.0.0.1:5005/naive
python example/flask_example.py

# benchmark naive api without service_streamer
./wrk -t 4 -c 128 -d 20s --timeout=10s -s scripts/streamer.lua http://127.0.0.1:5005/naive

# benchmark stream api with service_streamer
./wrk -t 4 -c 128 -d 20s --timeout=10s -s scripts/streamer.lua http://127.0.0.1:5005/naive

多个GPU进程 
这里对比单 web server 进程的情况下，多 GPU worker 的性能，验证通过和负载均衡机制的性能损耗。Flask 多线程 server 已经成为性能瓶颈，故采用 gevent server。

利用Future API使用多个GPU 
为了规避 web server 的性能瓶颈，我们使用底层 Future API 本地测试多 GPU worker 的 benchmark。

可以看出 service_streamer 的性能跟 GPUworker 数量及乎成线性关系，其中进程间通信的效率略高于 redis 通信。

点击以下标题查看更多往期内容： 

• KDD Cup 2019 AutoML Track冠军团队技术分享

• 神经网络架构搜索（NAS）综述 | 附资料推荐
  

• 小米拍照黑科技：基于NAS的图像超分辨率算法
  

• 深度解读：小米AI实验室最新成果FairNAS

• 自动机器学习（AutoML）最新综述

• NAS-FPN：基于自动架构搜索的特征金字塔网络

#投 稿 通 道#

 让你的论文被更多人看到 

如何才能让更多的优质内容以更短路径到达读者群体，缩短读者寻找优质内容的成本呢？答案就是：你不认识的人。

总有一些你不认识的人，知道你想知道的东西。PaperWeekly 或许可以成为一座桥梁，促使不同背景、不同方向的学者和学术灵感相互碰撞，迸发出更多的可能性。

PaperWeekly 鼓励高校实验室或个人，在我们的平台上分享各类优质内容，可以是最新论文解读，也可以是学习心得或技术干货。我们的目的只有一个，让知识真正流动起来。

? 来稿标准：

• 稿件确系个人原创作品，来稿需注明作者个人信息（姓名+学校/工作单位+学历/职位+研究方向） 

• 如果文章并非首发，请在投稿时提醒并附上所有已发布链接 

• PaperWeekly 默认每篇文章都是首发，均会添加“原创”标志

? 投稿邮箱：

• 投稿邮箱：hr@paperweekly.site 

• 所有文章配图，请单独在附件中发送 

• 请留下即时联系方式（微信或手机），以便我们在编辑发布时和作者沟通

?

现在，在「知乎」也能找到我们了

进入知乎首页搜索「PaperWeekly」

点击「关注」订阅我们的专栏吧

关于PaperWeekly

PaperWeekly 是一个推荐、解读、讨论、报道人工智能前沿论文成果的学术平台。如果你研究或从事 AI 领域，欢迎在公众号后台点击「交流群」，小助手将把你带入 PaperWeekly 的交流群里。

▽ 点击 | 阅读原文 | 访问项目主页