我不会用 Triton 系列：如何实现一个 backend
如何实现一个 backend
这篇文章主要讲如何实现一个 Triton Backend，以 Pytorch Backend 为例子。
Backend API
我们需要实现两个类来存储状态以及七个 Backend API。
- ModelState
- ModelInstanceState
- TRITONBACKEND_Initialize
- TRITONBACKEND_Finalize
- TRITONBACKEND_ModelInitialize
- TRITONBACKEND_ModelFinalize
- TRITONBACKEND_ModelInstanceInitialize
- TRITONBACKEND_ModelInstanceFinalize
- TRITONBACKEND_ModelInstanceExecute
ModelState 和 ModelInstanceState 这两个类可以绑定到 Triton 提供的指针上，你可以在 ModelState 和 ModelInstanceState 里面存储任何你想要的状态，然后将它绑定到 Triton 提供的指针上。这两个类并非必须的，它的作用相当于存储。在阅读了 Pytorch Backend 之后，会发现如果要写新的 Backend，七个 Backend API 并不需要做任何更改，只需要修改 ModelState 和 ModelInstanceState 即可。这两个类里面只需要做几个事情：模型配置文件检验、处理请求。
API 调用的时机和次数
简单概括就是：
- 动态链接库加载的时候，执行 TRITONBACKEND_Initialize
- 当属于一个 Backend 的模型都被删除了，且 Triton 开启了热更新，它会卸载动态链接库，执行 TRITONBACKEND_Finalize
- 其他的方法看名字就好了，不然就说了很多废话。比如 “在模型初始化的时候，调用模型初始化”
- 一个 Backend 对应多个 Model，Backend 只调用一次，Model 调用次数和仓库中模型数量一样多
- 一个 Model 对应多个 ModelInstance，根据模型的配置文件，调用 “模型实例” 的初始化方法。
Pytorch Backend 例子
地址：https://github.com/triton-inference-server/pytorch_backend/blob/main/src/libtorch.cc
我们以 Pytorch Backend 为学习例子，看看应该如何实现。
ModelState
一个 ModelState 和一个 TRITONBACKEND_Model 相关联，这个类主要提供一些模型配置检查、参数校验、模型实例共用的属性和方法。比如，加载模型的方法是所有模型实例初始化的时候需要的。
ModelInstanceState
一个 ModelInstanceState 和一个 TRITONBACKEND_ModelInstance 相关联，多个 ModelInstanceState 共享一个 ModelState。这个类主要提供一些处理请求、前向传播执行的方法。
令人颇感困惑的是，Pytorch 将模型输入输出配置的检查放到了这个函数里面，而 Tensorflow 的 backend 实现中，将模型输入输出的检查放到 ModelState。从抽象的分层来看，我认为模型配置的检查应该放到实例化之前，这样就可以避免每次初始化 “模型实例” 的时候都检查一次。Pytorch 这么做的原因是，设计了一个 ModelInstanceState 相关的内部状态 input_index_map_，这个状态的初始化依赖于模型的配置。
TRITONBACKEND_Initialize
Pytorch Backend 里面没有什么需要特别处理的东西，就是模板代码就好了，打印 backend 名字和版本之类的。
TRITONBACKEND_Finalize
没有提供实现。卸载动态链接库，直接移除就好了，没有需要清理的东西。
TRITONBACKEND_ModelInitialize
调用 Create 方法创建一个 ModelState，使用 TRITONBACKEND_ModelSetState 将 ModelState 绑定到传进来的 TRITONBACKEND_Model 上面。
TRITONBACKEND_ModelFinalize
前面绑定的是一个指针，所以要在这里删除指针。
TRITONBACKEND_ModelInstanceInitialize
“模型实例” 初始化和 TRITONBACKEND_ModelInitialize 的逻辑基本一致。不过需要使用多几个 API，这个方法传进来只有模型实例，我们可以从实例拿到绑定的 Model，再从 Model 拿出 ModelState，然后调用 “模型实例” 的 Create 方法进行初始化，最后同样调用 API 绑定到 ModelInstance。
TRITONBACKEND_ModelInstanceFinalize
前面绑定的是一个指针，所以要在这里删除指针。
TRITONBACKEND_ModelInstanceExecute
这个 API 的输入是 “模型实例” 和 “请求”，这里从 “模型实例” 中取出 ModelInstanceState，然后调用处理请求的方法即可。
实现细节
模型配置文件检验
在 Pytorch 的实现中，将模型配置文件的检验放到了 “模型实例” 初始化的时候，因为它设计了一些 “模型实例” 相关的状态，并且需要使用到模型配置文件。于是它一边进行模型配置文件的检验，一边初始化 “模型实例” 相关的状态。
在 OneFlow 的实现中，计划将模型配置文件的校验放到模型初始化 TRITONBACKEND_ModelInitialize 里面，而不是 “模型实例” 初始化的时候。不过，这取决于后面的 OneFlow C++ API 是如何实现的。
那么 Pytorch Backend 是如何实现的呢？
整个检验过程是：进行模型配置文件的检验，之后设置 “模型实例” 相关的状态。一边分析输入输出的名字是否符合规则，一边将输入输出的名字映射到 id。这么设计的 主要原因 是 forward 接口需要用户按照一定顺序将 tensor 输入，并没有提供一个 map 结构的输入。写成这样的 次要原因 是提高效率，一次 parse 就好。
关键数据结构：
triton::common::TritonJson::Value // 其实就是 JSONObject
关键函数调用:
GetBooleanSequenceControlProperties // 获取 sequence control 属性
GetTypedSequenceControlProperties // 同上
model_state_->ModelConfig().MemberAsArray("input", &ios) // 获取一个 JSON 的成员并作为 JSONArray
ios.ArraySize() // JSONArray 数组大小
ios.IndexAsObject(i, &io) // 获取 JSONArray 中的一个元素
io.MemberAsString("name", &io_name) // 获取 JSONArray 中的一个元素并作为 String
TRITON_ENABLE_GPU
如果 Triton 开启了 GPU，那么需要做一些特别的处理。比如在同步 CudaStream。
处理请求
TRITONBACKEND_ModelInstanceExecute 拿到的是一个二维指针，即一个请求的数组，这一批请求需要一起做处理。在进行前向传播之前，我们需要将输出收集起来。Triton 提供了一些工具类帮助我们去做收集，估计下面这个收集的方法是一个异步的方法，这样可以提高性能，不过需要我们显式使用同步操作。下面的 input_buffer 是一个指针，可能指向 Host，也可能指向 Device，不管这个指针指向哪里，后面使用 libtorch 的方法，创建一个 tensor，这样我们就获取了输入了。
需要注意的是：分配内存需要调用 Triton 的方法，然后用 torch 创建 tensor。不管 Tensor 所属的内存是 CPU 还是 GPU 的，都是由 Triton 来管理。
collector->ProcessTensor(
input_name, input_buffer, batchn_byte_size, memory_type,
memory_type_id);
关键数据结构
BackendMemory // 内存的抽象
BackendMemory::AllocationType // 分配的类型: CPU 或者 GPU
关键 API 调用
TRITONBACKEND_RequestInputByIndex(requests[i], 0 /* index */, &input); // 获取绑定在 request 上的输入
TRITONBACKEND_InputProperties(input, nullptr, nullptr, &shape, nullptr, nullptr, nullptr); // 获取输入的属性
TRITONBACKEND_RequestInputCount(requests[0], &input_count) // 获取输入的个数
const int64_t batchn_byte_size = GetByteSize(input_datatype, batchn_shape); // 输入字节大小
响应请求
在前处传播完了之后，就可以获取到输出的 Tensor 了，我们只需要从 Tensor 中取出数据指针就可以了，然后调用 Triton 提供的工具，帮我们将数据拷贝到指定的 repsonse 上面。
responder.ProcessTensor(
name, output_dtype, batchn_shape, output_buffer,
(device_.type() == torch::kCPU) ? TRITONSERVER_MEMORY_CPU
: TRITONSERVER_MEMORY_GPU,
(device_.type() == torch::kCPU) ? 0 : device_.index());
关键 API 调用
auto err = TRITONBACKEND_ResponseNew(&response, requests[i]); // 根据 request 创建 response
报告统计数据
时间戳宏
为了方便获取时间戳，Triton 提供了一个宏函数，方便调用。
#define SET_TIMESTAMP(TS_NS) \
{ \
TS_NS = std::chrono::duration_cast<std::chrono::nanoseconds>( \
std::chrono::steady_clock::now().time_since_epoch()) \
.count(); \
}
需要哪些时间戳
uint64_t exec_start_ns = 0; // 开始执行，从请求中取出数据开始
uint64_t compute_start_ns = 0; // 推理开始
uint64_t compute_end_ns = 0; // 推理结束
uint64_t exec_end_ns = 0; // 结束执行，在报告统计数据之前。
统计数据
TRITONBACKEND_ModelInstanceReportStatistics // 一条请求
TRITONBACKEND_ModelInstanceReportBatchStatistics // 一个 Batch
总结
简单梳理了一下，其实就几个事情：
- 模型配置文件检测，验证输入输出的写法是否正确
- 请求和响应，从请求中取出输入，将输出写到响应
- 内存管理的方法，Triton 管理内存，Pytorch 则接收或输出一个指针
- 统计数据，获取几个时间戳，调用 Triton API 来设置
深度学习框架在上面的过程中，只负责了一小部分，从 Triton 拿到指针，变成一个框架可以处理的 Tensor，然后进行推理，获取输出，最后将输出变成一个指针，返回给 Triton。于是，Triton 拿到指针之后，写到 response 里面。
