如何实现一个 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 里面。