服务流处理器

助力您的深度学习应用Web服务。 中文README

什么是服务流处理器？ • 亮点 • 安装 • 5分钟开发BERT服务 • API • 基准测试 • 常见问题解答 •

• 由ShannonAI制作 • :globe_with_meridians: http://www.shannonai.com/

什么是服务流处理器？

在深度学习模型中，通常会将数据样本收集为一个“小批量”进行处理。这样可以充分利用GPU的并行计算能力。然而，用户对Web服务的请求通常是离散的。如果使用传统的循环服务器或线程服务器，GPU每次只能处理一个请求，导致空闲时间增加，并且在有并发请求时，延迟会线性增长。

ServiceStreamer是一款用于机器学习应用Web服务的中间件。它能够将用户的请求排队并采样成小批量，通过提高GPU利用率来显著提升系统的整体性能。

亮点

• :hatching_chick: 易于使用：只需少量改动，即可使模型加速十倍。
• :zap: 处理速度快：降低机器学习模型在线推理的延迟。
• :octopus: 扩展性强：易于应用于多GPU场景，以处理海量请求。
• :crossed_swords: 适用性广：可与任何Web框架和/或深度学习框架配合使用。

安装

使用pip安装ServiceStream，需要Python >= 3.5：

pip install service_streamer 

5分钟开发BERT服务

我们提供了一个分步教程，帮助您在5分钟内将BERT部署上线。该服务每秒可处理1400个句子。

“文本补全”是自然语言处理中的一个任务：给定一个随机缺失了若干单词的句子，模型需要根据上下文预测这些被删除的单词。

“BERT”在过去两年中备受关注，在许多NLP任务上都取得了最先进的成果。BERT利用“掩码语言模型（MLM）”作为其预训练目标之一。MLM会随机遮盖输入中的部分标记，目标是根据上下文预测被遮盖单词的原始词汇ID。MLM与文本补全任务有相似之处，因此将BERT引入文本补全任务是非常自然的选择。

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 = TextInfillingModel()
   outputs = model.predict(batch)
   print(outputs)
   # ['little', 'you', 'universe']
   

   注意：请先下载预训练的BERT模型。

2. 其次，利用Flask将预测接口封装成Web服务。flask_example.py

   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)
   

   请运行flask_example.py，您将得到一个基础的Web服务器。

   curl -X POST http://localhost:5005/naive -d 's=Happy birthday to [MASK].' 
   ["you"]
   

   此时，您的Web服务器每秒只能处理12个请求。更多细节请参见基准测试。

3. 第三，通过service_streamer封装模型函数。仅需三行代码，BERT服务的预测速度就能达到每秒200+个句子（快了16倍）。

   from service_streamer import ThreadedStreamer
   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 jsonify(outputs)
   
   app.run(port=5005, debug=False)
   

   运行flask_example.py，然后使用wrk测试性能。

   wrk -t 2 -c 128 -d 20s --timeout=10s -s benchmark.lua http://127.0.0.1:5005/stream
   ...
   Requests/sec:    200.31
   

4. 最后，通过Streamer封装模型，并在多个GPU上启动服务工作进程。Streamer进一步提升了推理速度，使服务达到每秒1000+个句子（快了80倍）。

   from service_streamer import ManagedModel, Streamer
   
   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)
   

   可以启动8个GPU工作进程，并均匀分配到4个GPU上。

API

快速入门

一般来说，利用并行计算可以加快推理速度。

outputs = model.predict(batch_inputs)

ServiceStreamer是用于机器学习应用Web服务的中间件。它会将用户的请求排队并调度成小批量，然后转发到GPU工作进程中。ServiceStreamer会在一定程度上牺牲延迟（默认最大为0.1秒），并通过提高GPU利用率来提升整体性能。

from service_streamer import ThreadedStreamer

# 使用Streamer封装batch_predict函数

streamer = ThreadedStreamer(model.predict, batch_size=64, max_latency=0.1)

# 将 model.predict 替换为 streamer.predict

outputs = streamer.predict(batch_inputs)

在多线程（或协程）环境下启动 Web 服务器。只需添加几行代码，您的服务器通常就能提升 10 倍（batch_size/batch_per_request）的处理速度。

分布式 GPU 工作进程

实际应用中，Web 服务器的性能（QPS）远高于 GPU 模型本身。我们还支持一个 Web 服务器搭配多个 GPU 工作进程。

from service_streamer import Streamer

# 使用 spawn 启动 4 个 GPU 工作进程
streamer = Streamer(model.predict, 64, 0.1, worker_num=4)
outputs = streamer.predict(batch)

Streamer 默认使用 spawn 子进程来运行 GPU 工作进程。它通过进程间队列进行通信和任务排队，能够将大量请求分发到多个工作进程进行处理。

随后，模型的预测结果会批量返回给对应的 Web 服务器，并进一步转发到相应的 HTTP 响应中。

+-----------------------------------------------------------------------------+
| NVIDIA-SMI 390.116                Driver Version: 390.116                   |
|-------------------------------+----------------------+----------------------+
...
+-----------------------------------------------------------------------------+
| Processes:                                                       GPU Memory |
|  GPU       PID   Type   Process name                             Usage      |
|=============================================================================|
|    0      7574      C   /home/liuxin/nlp/venv/bin/python            1889MiB |
|    1      7575      C   /home/liuxin/nlp/venv/bin/python            1889MiB |
|    2      7576      C   /home/liuxin/nlp/venv/bin/python            1889MiB |
|    3      7577      C   /home/liuxin/nlp/venv/bin/python            1889MiB |
+-----------------------------------------------------------------------------+

上述方法定义简单，但主进程需要初始化模型，这会额外占用一部分内存。而且模型只能运行在同一个 GPU 上。因此，我们提供了 ManagedModel 类，以方便模型的延迟初始化和迁移，同时支持多 GPU 运行。

from service_streamer import ManagedModel

class ManagedBertModel(ManagedModel):

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

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


# 使用 spawn 启动 4 个 GPU 工作进程，并均匀分配到 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 服务器

一些 CPU 密集型计算，例如图像和文本预处理，通常需要先在 Web 服务器中完成。预处理后的数据再传递给 GPU 工作进程进行预测。然而，在实际应用中，CPU 资源往往会成为性能瓶颈。为此，我们也提供了多台 Web 服务器与单个或多个 GPU 工作进程相匹配的模式。

使用 RedisStream 为所有 Web 服务器和 GPU 工作进程指定一个唯一的 Redis 地址。

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

我们利用 gunicorn 或 uwsgi 实现反向代理和负载均衡。

cd example
gunicorn -c redis_streamer_gunicorn.py flask_example:app

每个请求都会被负载均衡到不同的 Web 服务器进行 CPU 预处理，然后再均匀分配到 GPU 工作进程进行模型预测。

Future API

如果您使用过任何并发库，应该对 future 不陌生。如果希望将 service_streamer 用于请求排队或分布式 GPU 计算，且应用场景并非 Web 服务，可以直接使用 Future API。

from service_streamer import ThreadedStreamer
streamer = ThreadedStreamer(model.predict, 64, 0.1)

xs = []
for i in range 200:
    future = streamer.submit([["How", "are", "you", "?"], ["Fine", "."], ["Thank", "you", "."]])
    xs.append(future)


# 获取所有 Future 对象实例并等待异步响应。
for future in xs:
    outputs = future.result()
    print(outputs)

基准测试

基准测试

我们使用 wrk 进行基准测试。

测试示例和脚本可在 example 中找到。

环境

• GPU：Titan Xp
• CUDA：9.0
• PyTorch：1.1

单 GPU 进程

# 启动 Flask 多线程服务器
python example/flask_example.py

# 测试无 service_streamer 的原生 API
wrk -t 4 -c 128 -d 20s --timeout=10s -s benchmark.lua http://127.0.0.1:5005/naive
# 测试带有 service_streamer 的流式 API
wrk -t 4 -c 128 -d 20s --timeout=10s -s benchmark.lua http://127.0.0.1:5005/stream

	Naive	ThreaedStreamer	Streamer	RedisStreamer
qps	12.78	207.59	321.70	372.45
延迟	8440ms	603.35ms	392.66ms	340.74ms

多 GPU 进程

与单个 Web 服务器进程相比，多 GPU 工作进程的通信和负载均衡机制会导致一定的性能损失。由于多线程 Flask 服务器已成为性能瓶颈，我们采用了 gevent 服务器。请参阅 flask_multigpu_example.py。

wrk -t 8 -c 512 -d 20s --timeout=10s -s benchmark.lua http://127.0.0.1:5005/stream

gpu_worker_num	Naive	ThreadedStreamer	Streamer	RedisStreamer
1	11.62	211.02	362.69	365.80
2	N/A	N/A	488.40	609.63
4	N/A	N/A	494.20	1034.57

• Threaded Streamer 由于 Python GIL 的限制，多工作进程并无意义。我们仅使用单 GPU 工作进程进行对比研究。

• Streamer 在超过 2 个 GPU 工作进程时，性能提升不再线性。此时 CPU 利用率达到 100%，瓶颈在于 CPU，Flask 的性能问题也成为阻碍。

使用 Future API 启动多 GPU 进程

我们采用 Future API 在本地进行多 GPU 基准测试，以减少 Web 服务器对性能的影响。请参考 future_example.py 中的代码示例。

gpu_worker_num	批量处理	ThreadedStreamer	Streamer	RedisStreamer
1	422.883	401.01	399.26	384.79
2	无	无	742.16	714.781
4	无	无	1400.12	1356.47

可以看出，service_streamer 的性能几乎与 GPU 工作进程的数量呈线性关系。在 service_streamer 中，进程间通信比使用 Redis 更高效。

常见问题解答

问： 使用从 allennlp 训练的模型，在推理时将 Streamer 的 worker_num=4，为什么会出现 16 核 CPU 被打满，且速度反而比 worker_num=1 的 Streamer 还慢的情况？

答： 对于多进程推理，如果模型在处理数据时使用了多线程的 NumPy 操作，可能会导致较高的 CPU 开销，从而使得多核并行计算的速度反而不如单核。这类问题在使用 alennlp、spacy 等第三方库时较为常见。可以通过设置 numpy threads 相关环境变量来解决。

import os
os.environ["MKL_NUM_THREADS"] = "1"  # export MKL_NUM_THREADS=1 
os.environ["NUMEXPR_NUM_THREADS"] = "1"  # export NUMEXPR_NUM_THREADS=1 
os.environ["OMP_NUM_THREADS"] = "1"  # export OMP_NUM_THREADS=1
import numpy

请注意，必须在 import numpy 之前设置这些环境变量。

问： 使用 RedisStreamer 时，如果只有一个 Redis 代理，但有多个模型，输入批次的结构可能不同。如何处理这种情况？

答： 在初始化工作进程和流式处理器时指定前缀，每个流式处理器将使用唯一的通道。

初始化工作进程的示例：

from service_streamer import run_redis_workers_forever
from bert_model import ManagedBertModel

if __name__ == "__main__":
    from multiprocessing import freeze_support
    freeze_support()
    run_redis_workers_forever(ManagedBertModel, 64, prefix='channel_1')
    run_redis_workers_forever(ManagedBertModel, 64, prefix='channel_2')

使用流式处理器获取结果的示例：

from service_streamer import RedisStreamer

streamer_1 = RedisStreaemr(prefix='channel_1')
streamer_2 = RedisStreaemr(prefix='channel_2')

# 预测
output_1 = streamer_1.predict(batch)
output_2 = streamer_2.predict(batch)

Service Streamer 快速上手指南

Service Streamer 是一款专为深度学习 Web 服务设计的中间件。它通过将离散的用户请求自动聚合成 mini-batch，充分利用 GPU 的并行计算能力，从而显著提升模型推理的吞吐量（QPS）并降低延迟。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• 操作系统：Linux / macOS / Windows
• Python 版本：>= 3.5
• 硬件依赖：推荐使用配备 NVIDIA GPU 的环境以发挥最大性能（CPU 环境亦可运行）
• 前置框架：已安装任意深度学习框架（如 PyTorch, TensorFlow）及 Web 框架（如 Flask, FastAPI）

安装步骤

使用 pip 进行安装。国内开发者建议使用清华源或阿里源以加速下载：

# 使用默认源
pip install service_streamer

# 或使用国内镜像源（推荐）
pip install service_streamer -i https://pypi.tuna.tsinghua.edu.cn/simple

基本使用

只需几行代码即可将现有的模型预测函数升级为支持高并发的流式服务。以下以 Flask + BERT 文本填充任务为例。

1. 定义模型

首先定义一个标准的模型类，包含 predict 方法，该方法接收一个列表（batch）并返回结果。

class TextInfillingModel(object):
    def __init__(self):
        # 初始化你的深度学习模型，例如加载 BERT 权重
        pass

    def predict(self, batch):
        # batch 是一个字符串列表，例如 ["Hello [MASK].", "World [MASK]."]
        # 执行模型推理并返回结果列表
        results = [] 
        # ... 模型推理逻辑 ...
        return results

model = TextInfillingModel()

2. 集成 Service Streamer

引入 ThreadedStreamer，将模型的 predict 方法包装起来。设置合适的 batch_size（批大小）和 max_latency（最大等待延迟，单位秒）。

from service_streamer import ThreadedStreamer

# 初始化 Streamer
# batch_size: 每次推理的最大样本数
# max_latency: 请求在队列中等待的最大时间（秒），超时即强制推理
streamer = ThreadedStreamer(model.predict, batch_size=64, max_latency=0.1)

3. 部署 Web 服务

在 Web 框架（如 Flask）中，用 streamer.predict 替换原有的 model.predict。

from flask import Flask, request, jsonify

app = Flask(__name__)

@app.route("/predict", methods=["POST"])
def stream_predict():
    # 获取请求中的多个句子
    inputs = request.form.getlist("s")
    
    # 使用 streamer 进行预测（自动 batching）
    outputs = streamer.predict(inputs)
    
    return jsonify(outputs)

if __name__ == "__main__":
    # 生产环境建议关闭 debug 模式
    app.run(port=5005, debug=False)

进阶：多 GPU 分布式部署

若需利用多张 GPU 进一步提升性能，可使用 Streamer 配合 ManagedModel 启动多个工作进程：

from service_streamer import ManagedModel, Streamer

class ManagedBertModel(ManagedModel):
    def init_model(self):
        # 在每个子进程中独立初始化模型
        self.model = TextInfillingModel()

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

# 启动 8 个工作进程，均匀分布在 4 张 GPU 上
streamer = Streamer(
    ManagedBertModel, 
    batch_size=64, 
    max_latency=0.1, 
    worker_num=8, 
    cuda_devices=(0, 1, 2, 3)
)

通过以上步骤，您的服务即可实现高吞吐量的模型推理。