我不會用 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 里面。