Triton API refactoring (#71)

Extend Triton API C++ wrappers. Remove dali_backend.h.

Signed-off-by: Rafal <[email protected]>

Author: banasraf

qa/L0_multi_input/model_repository/dali_multi_input/config.pbtxt

@@ -39,7 +39,7 @@ input [
 output [
   {
     name: "DALI_unchanged"
-    data_type: TYPE_INT32
+    data_type: TYPE_UINT8
     dims: [ -1 ]
   },
   {

src/CMakeLists.txt

@@ -26,6 +26,7 @@ add_subdirectory(dali_executor)
 add_library(
         triton-dali-backend SHARED
         dali_backend.cc
+        utils/triton.cc
 )
 
 add_library(

src/dali_backend.cc

@@ -20,14 +20,16 @@
 // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 // SOFTWARE.
 
-#include "src/dali_backend.h"
-
 #include <memory>
 
+#include "src/dali_executor/dali_executor.h"
 #include "src/dali_executor/io_buffer.h"
 #include "src/dali_executor/utils/dali.h"
 #include "src/dali_executor/utils/utils.h"
+#include "src/model_provider/model_provider.h"
+#include "src/utils/timing.h"
 #include "src/utils/triton.h"
+#include "triton/backend/backend_common.h"
 #include "triton/backend/backend_model.h"
 #include "triton/backend/backend_model_instance.h"
 
@@ -157,7 +159,7 @@ TRITONSERVER_Error* DaliModel::Create(TRITONBACKEND_Model* triton_model, DaliMod
 }
 
 struct RequestMeta {
-  uint64_t compute_start_ns, compute_end_ns;
+  TimeInterval compute_interval;  // nanoseconds
   int batch_size;
 };
 
@@ -175,20 +177,36 @@ class DaliModelInstance : public ::triton::backend::BackendModelInstance {
     return *dali_model_;
   }
 
-  RequestMeta ProcessRequest(TRITONBACKEND_Response* response, TritonRequest& request) {
+  void Execute(const std::vector<TritonRequest>& requests) {
     DeviceGuard dg(device_id_);
-    RequestMeta ret;
-    auto& outputs_indices = dali_model_->GetOutputOrder();
+    int total_batch_size = 0;
+    TimeInterval batch_compute_interval{};
+    TimeInterval batch_exec_interval{};
+    start_timer_ns(batch_exec_interval);
+    for (size_t i = 0; i < requests.size(); i++) {
+      TimeInterval req_exec_interval{};
+      start_timer_ns(req_exec_interval);
+      auto response = TritonResponse::New(requests[i]);
+      RequestMeta request_meta;
+      TritonError error{};
+      try {
+        request_meta = ProcessRequest(response, requests[i]);
+      } catch (...) { error = ErrorHandler(); }
+
+      if (i == 0) {
+        batch_compute_interval.start = request_meta.compute_interval.start;
+      } else if (i == requests.size() - 1) {
+        batch_compute_interval.end = request_meta.compute_interval.end;
+      }
 
-    auto dali_inputs = GenerateInputs(request);
-    ret.batch_size = dali_inputs[0].meta.shape.num_samples();  // Batch size is expected to be the
-                                                               // same in every input
-    ret.compute_start_ns = detail::capture_time();
-    auto outputs_info = dali_executor_->Run(dali_inputs);
-    ret.compute_end_ns = detail::capture_time();
-    auto dali_outputs = detail::AllocateOutputs(request, response, outputs_info, outputs_indices);
-    dali_executor_->PutOutputs(dali_outputs);
-    return ret;
+      end_timer_ns(req_exec_interval);
+      ReportStats(requests[i], req_exec_interval, request_meta.compute_interval, !error);
+      SendResponse(std::move(response), std::move(error));
+
+      total_batch_size += request_meta.batch_size;
+    }
+    end_timer_ns(batch_exec_interval);
+    ReportBatchStats(total_batch_size, batch_exec_interval, batch_compute_interval);
   }
 
  private:
@@ -201,6 +219,37 @@ class DaliModelInstance : public ::triton::backend::BackendModelInstance {
     dali_executor_ = std::make_unique<DaliExecutor>(std::move(pipeline));
   }
 
+  void ReportStats(TritonRequestView request, TimeInterval exec, TimeInterval compute,
+                   bool success) {
+    LOG_IF_ERROR(TRITONBACKEND_ModelInstanceReportStatistics(triton_model_instance_, request,
+                                                             success, exec.start, compute.start,
+                                                             compute.end, exec.end),
+                 "Failed reporting request statistics.");
+  }
+
+  void ReportBatchStats(uint32_t total_batch_size, TimeInterval exec, TimeInterval compute) {
+    LOG_IF_ERROR(TRITONBACKEND_ModelInstanceReportBatchStatistics(
+                     triton_model_instance_, total_batch_size, exec.start, compute.start,
+                     compute.end, exec.end),
+                 "Failed reporting batch statistics.");
+  }
+
+  /** Run inference for a given \p request and prepare a response. */
+  RequestMeta ProcessRequest(TritonResponseView response, TritonRequestView request) {
+    RequestMeta ret;
+
+    auto dali_inputs = GenerateInputs(request);
+    ret.batch_size = dali_inputs[0].meta.shape.num_samples();  // Batch size is expected to be the
+                                                               // same in every input
+    start_timer_ns(ret.compute_interval);
+    auto outputs_info = dali_executor_->Run(dali_inputs);
+    end_timer_ns(ret.compute_interval);
+    auto dali_outputs = AllocateOutputs(request, response, outputs_info);
+    dali_executor_->PutOutputs(dali_outputs);
+    return ret;
+  }
+
+  /** @brief Generate descriptors of inputs provided by a given request. */
   std::vector<IDescr> GenerateInputs(TritonRequestView request) {
     uint32_t input_cnt = request.InputCount();
     std::vector<IDescr> ret;
@@ -220,6 +269,61 @@ class DaliModelInstance : public ::triton::backend::BackendModelInstance {
     return ret;
   }
 
+  /**
+   * @brief Allocate outputs required by a given request.
+   *
+   * Lifetime of the created buffer is bound to the \p response
+   */
+  std::vector<ODescr> AllocateOutputs(TritonRequestView request, TritonResponseView response,
+                                      const std::vector<OutputInfo>& outputs_info) {
+    uint32_t output_cnt = request.OutputCount();
+    ENFORCE(outputs_info.size() == output_cnt,
+            make_string("Number of outputs in the model configuration (", output_cnt,
+                        ") does not match to the number of outputs from DALI pipeline (",
+                        outputs_info.size(), ")"));
+    const auto& output_indices = dali_model_->GetOutputOrder();
+    std::vector<ODescr> outputs(output_cnt);
+    outputs.reserve(output_cnt);
+    for (uint32_t i = 0; i < output_cnt; ++i) {
+      auto name = request.OutputName(i);
+      int output_idx = output_indices.at(name);
+      IOMeta out_meta{};
+      out_meta.name = name;
+      out_meta.type = outputs_info[output_idx].type;
+      out_meta.shape = outputs_info[output_idx].shape;
+      auto output = response.GetOutput(out_meta);
+      auto buffer = output.AllocateBuffer(outputs_info[output_idx].device, device_id_);
+      outputs[output_idx] = {out_meta, {buffer}};
+    }
+    return outputs;
+  }
+
+  TritonError ErrorHandler() {
+    TritonError error{};
+    try {
+      throw;
+    } catch (TritonError& e) {
+      LOG_MESSAGE(TRITONSERVER_LOG_ERROR, (e.what()));
+      error = std::move(e);
+    } catch (DaliBackendException& e) {
+      LOG_MESSAGE(TRITONSERVER_LOG_ERROR, (e.what()));
+      error = TritonError::Unknown(make_string("DALI Backend error: ", e.what()));
+    } catch (DALIException& e) {
+      LOG_MESSAGE(TRITONSERVER_LOG_ERROR, (e.what()));
+      error = TritonError::Unknown(make_string("DALI error: ", e.what()));
+    } catch (std::runtime_error& e) {
+      LOG_MESSAGE(TRITONSERVER_LOG_ERROR, (e.what()));
+      error = TritonError::Unknown(make_string("Runtime error: ", e.what()));
+    } catch (std::exception& e) {
+      LOG_MESSAGE(TRITONSERVER_LOG_ERROR, (e.what()));
+      error = TritonError::Unknown(make_string("Exception: ", e.what()));
+    } catch (...) {
+      LOG_MESSAGE(TRITONSERVER_LOG_ERROR, ("Unknown error"));
+      error = TritonError::Unknown("Unknown error");
+    }
+    return error;
+  }
+
   std::unique_ptr<DaliExecutor> dali_executor_;
   DaliModel* dali_model_;
 };
@@ -441,76 +545,18 @@ TRITONSERVER_Error* TRITONBACKEND_ModelInstanceFinalize(TRITONBACKEND_ModelInsta
 TRITONSERVER_Error* TRITONBACKEND_ModelInstanceExecute(TRITONBACKEND_ModelInstance* instance,
                                                        TRITONBACKEND_Request** reqs,
                                                        const uint32_t request_count) {
-  DaliModelInstance* instance_state;
+  std::vector<TritonRequest> requests;
+  for (uint32_t idx = 0; idx < request_count; ++idx) {
+    requests.emplace_back(reqs[idx]);
+  }
+  DaliModelInstance* dali_instance;
   RETURN_IF_ERROR(
-      TRITONBACKEND_ModelInstanceState(instance, reinterpret_cast<void**>(&instance_state)));
+      TRITONBACKEND_ModelInstanceState(instance, reinterpret_cast<void**>(&dali_instance)));
   std::vector<TRITONBACKEND_Response*> responses(request_count);
 
-  int total_batch_size = 0;
-  uint64_t exec_start_ns = 0, exec_end_ns = 0, batch_exec_start_ns = 0, batch_exec_end_ns = 0,
-           batch_compute_start_ns = 0, batch_compute_end_ns = 0;
-  batch_exec_start_ns = detail::capture_time();
-  for (size_t i = 0; i < responses.size(); i++) {
-    TritonRequest request(reqs[i]);
-    TRITONSERVER_Error* error = nullptr;  // success
-    exec_start_ns = detail::capture_time();
-    // TODO Do not process requests one by one, but gather all
-    //     into one buffer and process it in DALI all together
-    LOG_IF_ERROR(TRITONBACKEND_ResponseNew(&responses[i], request),
-                 make_string("Failed creating a response, idx: ", i));
-    RequestMeta request_meta;
-
-    try {
-      request_meta = instance_state->ProcessRequest(responses[i], request);
-    } catch (DaliBackendException& e) {
-      LOG_MESSAGE(TRITONSERVER_LOG_ERROR, (e.what()));
-      error = TRITONSERVER_ErrorNew(TRITONSERVER_Error_Code::TRITONSERVER_ERROR_UNKNOWN,
-                                    make_string("DALI Backend error: ", e.what()).c_str());
-    } catch (DALIException& e) {
-      LOG_MESSAGE(TRITONSERVER_LOG_ERROR, (e.what()));
-      error = TRITONSERVER_ErrorNew(TRITONSERVER_Error_Code::TRITONSERVER_ERROR_UNKNOWN,
-                                    make_string("DALI error: ", e.what()).c_str());
-    } catch (std::runtime_error& e) {
-      LOG_MESSAGE(TRITONSERVER_LOG_ERROR, (e.what()));
-      error = TRITONSERVER_ErrorNew(TRITONSERVER_Error_Code::TRITONSERVER_ERROR_UNKNOWN,
-                                    make_string("runtime error: ", e.what()).c_str());
-    } catch (std::exception& e) {
-      LOG_MESSAGE(TRITONSERVER_LOG_ERROR, (e.what()));
-      error = TRITONSERVER_ErrorNew(TRITONSERVER_Error_Code::TRITONSERVER_ERROR_UNKNOWN,
-                                    make_string("exception: ", e.what()).c_str());
-    } catch (...) {
-      LOG_MESSAGE(TRITONSERVER_LOG_ERROR, ("Unknown error"));
-      error = TRITONSERVER_ErrorNew(TRITONSERVER_Error_Code::TRITONSERVER_ERROR_UNKNOWN,
-                                    "Unknown error");
-    }
-
-    exec_end_ns = detail::capture_time();
-    batch_compute_start_ns = batch_compute_start_ns == 0 ?
-                                 request_meta.compute_start_ns :
-                                 batch_compute_start_ns;  // Ternary to please the compiler
-
-    LOG_IF_ERROR(TRITONBACKEND_ModelInstanceReportStatistics(
-                     instance, request, !error, exec_start_ns, request_meta.compute_start_ns,
-                     request_meta.compute_end_ns, exec_end_ns),
-                 make_string("Failed reporting statistics for response idx ", i));
-
-    LOG_IF_ERROR(
-        TRITONBACKEND_ResponseSend(
-            responses[i], TRITONSERVER_ResponseCompleteFlag::TRITONSERVER_RESPONSE_COMPLETE_FINAL,
-            error),
-        make_string("Failed sending response, idx ", i));
-
-    total_batch_size += request_meta.batch_size;
-    batch_exec_end_ns = exec_end_ns;
-    batch_compute_end_ns = request_meta.compute_end_ns;
-  }
-  if (batch_exec_end_ns == 0)
-    batch_exec_end_ns = detail::capture_time();
-
-  LOG_IF_ERROR(TRITONBACKEND_ModelInstanceReportBatchStatistics(
-                   instance, total_batch_size, batch_exec_start_ns, batch_compute_start_ns,
-                   batch_compute_end_ns, batch_exec_end_ns),
-               make_string("Failed reporting batch statistics"));
+  try {
+    dali_instance->Execute(requests);
+  } catch (TritonError& err) { return err.release(); }
 
   return nullptr;
 }