Triton Inference Server 框架实践

Triton Inference Server 框架实践

最近，我参与了一个紧急项目，负责将项目中的一个模块利用 TritonServer 框架进行加速，主要通过其 dynamic_batch 特性提升单卡吞吐量。经过一段时间的努力，我在今晚值班时整理了相关经验，与大家分享。

背景介绍

该模块是 embedding 模块，用于将文本表示成浮点数向量，输入为 List[str]，输出为 List[List[float]]。服务部署了两种模型，通过输入字段决定使用哪个模型推理。目前线上版本使用 FastAPI 封装，提供 HTTP 接口服务。

由于时间紧迫，且服务协议不能改变，还需要对模型输出进行后处理，而我对 torch 转 onnx 模型不够熟练，因此先直接使用 python backend，后续再研究如何转 tensorrt 或 onnx 进行模型加速。

使用方法

基本用法

首先明确模型的输入输出，然后输出 config.pbtxt 文件。该模型较为简单，配置如下：

name: "model_name"
backend: "python"
input [
  {
    name: "texts"
    data_type: TYPE_STRING
    dims: [ -1 ]
  }
]
output [
  {
    name: "prompt_tokens"
    data_type: TYPE_INT32
    dims: [ 1 ]
  },
  {
    name: "total_tokens"
    data_type: TYPE_INT32
    dims: [ 1 ]
  },
  {
    name: "embedding"
    data_type: TYPE_FP64
    dims: [ -1, 768 ]
  }
]
dynamic_batching {
  preferred_batch_size: [ 4, 8 ]
  max_queue_delay_microseconds: 1000
}

需要注意的是，shape 默认包含第一维度 batch_size。虽然输入 texts 的 shape 是 -1，但输入应为二维，例如：

texts =