Delay time-point of response complete timestamp in GPRC and SageMaker…

… endpoint
triton-inference-server · Mar 7, 2024 · d905aeb · d905aeb
1 parent 1151a4d
commit d905aeb
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 12 deletions.
diff --git a/src/grpc/infer_handler.cc b/src/grpc/infer_handler.cc
@@ -992,11 +992,6 @@ ModelInferHandler::InferResponseComplete(
     state->context_->EraseInflightState(state);
   }
 
-#ifdef TRITON_ENABLE_TRACING
-  state->trace_timestamps_.emplace_back(std::make_pair(
-      "INFER_RESPONSE_COMPLETE", TraceManager::CaptureTimestamp()));
-#endif  // TRITON_ENABLE_TRACING
-
   // If gRPC Stream is cancelled then no need of forming and returning
   // a response.
   if (state->IsGrpcContextCancelled()) {
@@ -1064,6 +1059,11 @@ ModelInferHandler::InferResponseComplete(
     return;
   }
 
+#ifdef TRITON_ENABLE_TRACING
+  state->trace_timestamps_.emplace_back(std::make_pair(
+      "INFER_RESPONSE_COMPLETE", TraceManager::CaptureTimestamp()));
+#endif  // TRITON_ENABLE_TRACING
+
 #ifdef TRITON_ENABLE_TRACING
   state->trace_timestamps_.emplace_back(
       std::make_pair("GRPC_SEND_START", TraceManager::CaptureTimestamp()));

diff --git a/src/sagemaker_server.cc b/src/sagemaker_server.cc
@@ -394,13 +394,6 @@ SagemakerAPIServer::SagemakeInferRequestClass::InferResponseComplete(
     err = infer_request->FinalizeResponse(response);
   }
 
-#ifdef TRITON_ENABLE_TRACING
-  if (infer_request->trace_ != nullptr) {
-    infer_request->trace_->CaptureTimestamp(
-        "INFER_RESPONSE_COMPLETE", TraceManager::CaptureTimestamp());
-  }
-#endif  // TRITON_ENABLE_TRACING
-
 
   LOG_TRITONSERVER_ERROR(
       TRITONSERVER_InferenceResponseDelete(response),
@@ -426,6 +419,12 @@ SagemakerAPIServer::SagemakeInferRequestClass::InferResponseComplete(
   if ((flags & TRITONSERVER_RESPONSE_COMPLETE_FINAL) == 0) {
     return;
   }
+#ifdef TRITON_ENABLE_TRACING
+  if (infer_request->trace_ != nullptr) {
+    infer_request->trace_->CaptureTimestamp(
+        "INFER_RESPONSE_COMPLETE", TraceManager::CaptureTimestamp());
+  }
+#endif  // TRITON_ENABLE_TRACING
   evthr_defer(infer_request->thread_, ReplyCallback, infer_request);
 }