zoukankan      html  css  js  c++  java
  • 我不会用 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 里面。

  • 相关阅读:
    UVA 11235 Frequent Values ---RMQ
    UVA 12266 Stock prices --优先队列
    HDU 1896 Stones --优先队列+搜索
    POJ 1442 Black Box -优先队列
    POJ 2263 Heavy Cargo 多种解法
    POJ 3250 Bad Hair Day --单调栈(单调队列?)
    FZU1894 志愿者选拔 --单调队列
    POJ 2823 Sliding Window 再探单调队列
    UVA 11992 Fast Matrix Operations (二维线段树)
    两道相似KMP题
  • 原文地址:https://www.cnblogs.com/zzk0/p/15496171.html
Copyright © 2011-2022 走看看