Add request statistic reporting for decoupled mode (triton-inference-server#163)

Tabrizian · web-flow · commit 73fdcdad66f7 · 2022-06-06T12:41:55.000-04:00
diff --git a/src/pb_metric_reporter.cc b/src/pb_metric_reporter.cc
@@ -36,7 +36,8 @@ PbMetricReporter::PbMetricReporter(
     std::shared_ptr<std::vector<TRITONBACKEND_Response*>> responses)
     : instance_(instance), requests_(requests), request_count_(request_count),
       responses_(responses), total_batch_size_(0), exec_start_ns_(0),
-      compute_start_ns_(0), compute_end_ns_(0), exec_end_ns_(0)
+      compute_start_ns_(0), compute_end_ns_(0), exec_end_ns_(0),
+      success_status_(true)
 {
 }
 
@@ -51,21 +52,28 @@ PbMetricReporter::~PbMetricReporter()
     // request object. We use the execution start/end time for
     // compute also so that the entire execution time is associated
     // with the inference computation.
-    LOG_IF_ERROR(
-        TRITONBACKEND_ModelInstanceReportStatistics(
-            instance_, request,
-            ((*responses_)[r] != nullptr) /* success */, exec_start_ns_,
-            compute_start_ns_, compute_end_ns_, exec_end_ns_),
-        "failed reporting request statistics");
+    if (responses_) {
+      LOG_IF_ERROR(
+          TRITONBACKEND_ModelInstanceReportStatistics(
+              instance_, request, ((*responses_)[r] != nullptr) /* success */,
+              exec_start_ns_, compute_start_ns_, compute_end_ns_, exec_end_ns_),
+          "failed reporting request statistics");
+    } else {
+      LOG_IF_ERROR(
+          TRITONBACKEND_ModelInstanceReportStatistics(
+              instance_, request, success_status_, exec_start_ns_,
+              compute_start_ns_, compute_end_ns_, exec_end_ns_),
+          "failed reporting request statistics");
+    }
   }
 
   // Report the entire batch statistics. This backend does not support
   // batching so the total batch size is always 1.
   if (total_batch_size_ != 0) {
     LOG_IF_ERROR(
         TRITONBACKEND_ModelInstanceReportBatchStatistics(
-            instance_, total_batch_size_, exec_start_ns_,
-            compute_start_ns_, compute_end_ns_, exec_end_ns_),
+            instance_, total_batch_size_, exec_start_ns_, compute_start_ns_,
+            compute_end_ns_, exec_end_ns_),
         "failed reporting batch request statistics");
   }
 }
@@ -100,4 +108,10 @@ PbMetricReporter::SetExecEndNs(const uint64_t exec_end_ns)
   exec_end_ns_ = exec_end_ns;
 }
 
+void
+PbMetricReporter::SetSuccessStatus(const bool success_status)
+{
+  success_status_ = success_status;
+}
+
 }}}  // namespace triton::backend::python
diff --git a/src/pb_metric_reporter.h b/src/pb_metric_reporter.h
@@ -42,6 +42,7 @@ class PbMetricReporter {
   uint64_t compute_start_ns_;
   uint64_t compute_end_ns_;
   uint64_t exec_end_ns_;
+  bool success_status_;
 
  public:
   PbMetricReporter(
@@ -54,5 +55,6 @@ class PbMetricReporter {
   void SetComputeStartNs(const uint64_t compute_start_ns);
   void SetComputeEndNs(const uint64_t compute_end_ns);
   void SetExecEndNs(const uint64_t exec_end_ns);
+  void SetSuccessStatus(const bool success_status);
 };
 }}};  // namespace triton::backend::python
diff --git a/src/python_be.cc b/src/python_be.cc
@@ -926,7 +926,8 @@ ModelInstanceState::ResponseSendDecoupled(
 TRITONSERVER_Error*
 ModelInstanceState::ProcessRequestsDecoupled(
     TRITONBACKEND_Request** requests, const uint32_t request_count,
-    std::vector<std::unique_ptr<InferRequest>>& pb_inference_requests)
+    std::vector<std::unique_ptr<InferRequest>>& pb_inference_requests,
+    PbMetricReporter& reporter)
 {
   NVTX_RANGE(nvtx_, "ProcessRequests " + Name());
   closed_requests_ = {};
@@ -954,6 +955,10 @@ ModelInstanceState::ProcessRequestsDecoupled(
       requests, request_count, pb_inference_requests, request_batch,
       responses));
 
+  uint64_t compute_start_ns = 0;
+  SET_TIMESTAMP(compute_start_ns);
+  reporter.SetComputeStartNs(compute_start_ns);
+
   std::unique_ptr<IPCMessage> ipc_message;
   RETURN_IF_EXCEPTION(
       ipc_message =
@@ -971,6 +976,12 @@ ModelInstanceState::ProcessRequestsDecoupled(
 
   AllocatedSharedMemory<ResponseBatch> response_batch =
       Stub()->ShmPool()->Load<ResponseBatch>(received_message_->Args());
+
+  uint64_t compute_end_ns = 0;
+  SET_TIMESTAMP(compute_end_ns);
+  reporter.SetComputeEndNs(compute_end_ns);
+  reporter.SetBatchStatistics(request_count);
+
   if (response_batch.data_->has_error) {
     if (response_batch.data_->is_error_set) {
       auto error = PbString::LoadFromSharedMemory(
@@ -1819,10 +1830,24 @@ TRITONBACKEND_ModelInstanceExecute(
     }
   } else {
     std::vector<std::unique_ptr<InferRequest>> infer_requests;
+
+    uint64_t exec_start_ns = 0;
+    SET_TIMESTAMP(exec_start_ns);
+
+    PbMetricReporter reporter(
+        instance_state->TritonModelInstance(), requests, request_count,
+        nullptr);
+    reporter.SetExecStartNs(exec_start_ns);
+
     error = instance_state->ProcessRequestsDecoupled(
-        requests, request_count, infer_requests);
+        requests, request_count, infer_requests, reporter);
+
+    uint64_t exec_end_ns = 0;
+    SET_TIMESTAMP(exec_end_ns);
+    reporter.SetExecEndNs(exec_end_ns);
 
     if (error != nullptr) {
+      reporter.SetSuccessStatus(false);
       for (uint32_t r = 0; r < request_count; ++r) {
         TRITONBACKEND_Request* request = requests[r];
         if (!instance_state->ExistsInClosedRequests(
diff --git a/src/python_be.h b/src/python_be.h
@@ -320,7 +320,8 @@ class ModelInstanceState : public BackendModelInstance {
   // Process all the requests in the decoupled mode.
   TRITONSERVER_Error* ProcessRequestsDecoupled(
       TRITONBACKEND_Request** requests, const uint32_t request_count,
-      std::vector<std::unique_ptr<InferRequest>>& pb_infer_requests);
+      std::vector<std::unique_ptr<InferRequest>>& pb_infer_requests,
+      PbMetricReporter& pb_metric_reporter);
 
   bool ExistsInClosedRequests(intptr_t closed_request);
 
