近年来，在深度学习模型的部署中，涌现了许多新技术，从最早的tf/pytorch直接部署，到onnx，到trnsorrt(trt)，模型的运行效率越来越高。但即便是使用trt方式部署模型，也依旧无法避免巨大的工作量。在使用c++部署trt时，代码量较大，还需要预处理困难。而使用python部署trt，受限于python的GIL问题，需要同时部署多个模型，又失去了trt节省显存的优点。同时，不论使用何种方式，都无法避免面对组batch等问题，整个系统的工程实现会变得非常复杂。
在这样的情况下，triton应运而生一站式的帮助我们解决了几乎所有的工程问题，作为一款强大的模型执行引擎，triton支持几乎所有主流的网络模型，对于pytorch，tensorflow，onnx，trt都有良好的支持。

本文将从trt/onnx模型的部署介绍triton的部署流程。

triton项目的组件

triton server
triton的服务端可以直接使用docker方式部署，从triton的镜像列表中选择tag为xx.xx-py3的镜像即可。
triton client
triton的客户端可以从pip直接安装，但为了测试方便，可以直接使用包含测试端的容器环境，tag为xx.xx-py3-sdk的镜像中包含了客户端。

triton server

部署

triton部署模型可以参考文档1和文档2，但是对于onnx和trt模型，由于模型内已经包含了输入和输出的信息，因此triton可以自动生成配置文件，部署会变得非常简单。

按照triton的教程，我们创建三层目录结构，之后直接把onnx或trt模型拷贝进去即可。

onnx的默认模型名称为model.onnx，而trt的默认模型名称为model.plan

mkdir -p model_repository/your_model_name/1
# onnx
cp your_model_name.onnx model_repository/your_model_name/1/model.onnx
# trt
cp your_model_name.trt model_repository/your_model_name/1/model.plan

此时你的目录结构如下

<model-repository-path>/
  <your_model_name>/
    1/
      model.plan

准备好模型文件的目录结构之后，我们启动triton服务，并使用--strict-model-config=false要求他自动生成模型文件

docker run --rm --gpus all \
-p8000:8000 -p8001:8001 -p8002:8002 \
-v model_repository:/models \
nvcr.io/nvidia/tritonserver:21.05-py3 \
tritonserver --strict-model-config=false --model-repository=/models

如果模型正常启动，你将会看到类似下面的的输出

+----------------------+---------+--------+
| Model                | Version | Status |
+----------------------+---------+--------+
| <model_name>         | <v>     | READY  |
| ..                   | .       | ..     |
| ..                   | .       | ..     |
+----------------------+---------+--------+
...
...
...
I1002 21:58:57.891440 62 grpc_server.cc:3914] Started GRPCInferenceService at 0.0.0.0:8001
I1002 21:58:57.893177 62 http_server.cc:2717] Started HTTPService at 0.0.0.0:8000
I1002 21:58:57.935518 62 http_server.cc:2736] Started Metrics Service at 0.0.0.0:8002

config

接下来我们需要修改triton为我们生成的配置文件，使用如下命令可以获得当前的配置

1	curl localhost:8000/v2/models/<your_model_name>/config

得到json输出之后，我们需要收到修改为pbconfig的格式。可以参考这篇教程，基本格式如下。

platform: "your_model_name"
max_batch_size: 8
input [
  {
    name: "input0"
    data_type: TYPE_FP32
    dims: [ 16 ]
  }
]
output [
  {
    name: "output0"
    data_type: TYPE_FP32
    dims: [ 4 ]
  }
]

修改之后，保存为model_repository/your_model_name/config.pbtxt，保存后你的模型目录结构如下：

<model-repository-path>/
  <your_model_name>/
    config.pbtxt
    1/
      model.plan

这时我们可以关闭triton，去掉--strict-model-config=false选项后重启服务，修改好的配置文件就会生效了。

triton client

安装

此处以python client为例

triton python client的安装非常简单，使用如下方式即可，如果你只使用http或grpc方式调用，可以将tritonclient[all]替换为tritonclient[http]或tritonclient[grpc]

1 2	pip install nvidia-pyindex pip install tritonclient[all]

如果你安装完之后报错ModuleNotFoundError: No module named 'tritonclient'不妨尝试一下

1 2	python -m pip install nvidia-pyindex python -m pip install tritonclient[all]

这样可以确保你的依赖被安装到你使用的python环境中。

使用

根据我的测试，grpc方式调用triton的性能远远高于http方式，在极端情况下甚至可以达到十倍的性能差距，因此如果没有特殊需要，我不建议使用http方式调用。

客户端示例可以看此处，推荐从simple_grpc_infer_client.py这个例子入手

shared memory

对于本机内部调用，triton还支持使用shm方式进行数据共享，减小通信开销。更进一步的，你还可以使用cuda shm方式直接共享显存，在数据发送端就将显存设置好，直接调用triton进行推理。但需要注意的是cuda shm会消耗较多显存，需要用户自己去控制，否则会导致显存用尽，在我自己的测试中。

如果有这部分需要可以参考Triton shared memory

异步模式(async mode)

triton client还支持异步调用，可以查看simple_grpc_async_infer_client.py。

将triton客户端封装为协程

triton的客户端仅支持异步模式而不支持协程，需要我们手动进行封装

一次原始的异步调用形式如下，当执行结束，callback函数会被调用

triton_client.async_infer(
        model_name=model_name,
        inputs=inputs,
        callback=callback,
        outputs=outputs)

因此我们需要在callback中通知协程。需要注意的是callback会在另个一个triton client的线程中被调用，因此需要使用loop.call_soon_threadsafe方法为future设置结果。

import asyncio
from functools import partial

def callback(loop, future, result, error):
    if error:
        loop.call_soon_threadsafe(future.set_exception, error)
    else:
        loop.call_soon_threadsafe(future.set_result, result)

def await_infer(triton_client, model_name, inputs, outputs)
    loop = asyncio.get_running_loop()
    future = loop.create_future()

    triton_client.async_infer(
        model_name=model_name,
        inputs=inputs,
        callback=partial(callback, loop, future),
        outputs=outputs)
    return future

包装好后我们就可以使用await关键字去调用await_infer了

async_result = await await_infer(
                triton_client = service.triton_client, 
                model_name = model_name, 
                inputs = inputs,
                outputs = outputs
            )

Maple.Link

Nvidia Triton 使用教程