Add the ability for backends to add entries to the trace

triton-inference-server · Feb 27, 2023 · 9f30cb2 · 9f30cb2
1 parent f59101d
commit 9f30cb2
Show file tree

Hide file tree

Showing 11 changed files with 81 additions and 10 deletions.
diff --git a/include/triton/core/tritonbackend.h b/include/triton/core/tritonbackend.h
@@ -530,6 +530,15 @@ TRITONBACKEND_RequestOutputBufferProperties(
 TRITONBACKEND_DECLSPEC TRITONSERVER_Error* TRITONBACKEND_RequestRelease(
     TRITONBACKEND_Request* request, uint32_t release_flags);
 
+/// Get the trace associted with a request. The returned trace is owned by the
+/// request, not the caller, and so should not be modified or freed.
+///
+/// \param request The inference request.
+/// \param trace Returns the trace associated with the request.
+/// \return a TRITONSERVER_Error indicating success or failure.
+TRITONBACKEND_DECLSPEC TRITONSERVER_Error* TRITONBACKEND_RequestTrace(
+    TRITONBACKEND_Request* request, TRITONSERVER_InferenceTrace** trace);
+
 ///
 /// TRITONBACKEND_ResponseFactory
 ///

diff --git a/include/triton/core/tritonserver.h b/include/triton/core/tritonserver.h
@@ -707,7 +707,8 @@ typedef enum tritonserver_traceactivity_enum {
   TRITONSERVER_TRACE_REQUEST_END = 6,
   TRITONSERVER_TRACE_TENSOR_QUEUE_INPUT = 7,
   TRITONSERVER_TRACE_TENSOR_BACKEND_INPUT = 8,
-  TRITONSERVER_TRACE_TENSOR_BACKEND_OUTPUT = 9
+  TRITONSERVER_TRACE_TENSOR_BACKEND_OUTPUT = 9,
+  TRITONSERVER_TRACE_CUSTOM_ACTIVITY = 10
 } TRITONSERVER_InferenceTraceActivity;
 
 /// Get the string representation of a trace activity. The returned
@@ -728,7 +729,7 @@ TRITONSERVER_DECLSPEC const char* TRITONSERVER_InferenceTraceActivityString(
 typedef void (*TRITONSERVER_InferenceTraceActivityFn_t)(
     TRITONSERVER_InferenceTrace* trace,
     TRITONSERVER_InferenceTraceActivity activity, uint64_t timestamp_ns,
-    void* userp);
+    void* userp, const char* tag);
 
 /// Type for trace tensor activity callback function. This callback function
 /// is used to report tensor activity occurring for a trace. This function
@@ -810,6 +811,9 @@ TRITONSERVER_DECLSPEC TRITONSERVER_Error* TRITONSERVER_InferenceTraceTensorNew(
     TRITONSERVER_InferenceTraceTensorActivityFn_t tensor_activity_fn,
     TRITONSERVER_InferenceTraceReleaseFn_t release_fn, void* trace_userp);
 
+TRITONSERVER_DECLSPEC TRITONSERVER_Error* TRITONSERVER_ReportActivity(
+    TRITONSERVER_InferenceTrace* trace, const char* tag, uint64_t timestamp);
+
 /// Delete a trace object.
 ///
 /// \param trace The trace object.

diff --git a/src/backend_model.cc b/src/backend_model.cc
@@ -27,6 +27,7 @@
 #include "backend_model.h"
 
 #include <vector>
+
 #include "backend_config.h"
 #include "backend_model_instance.h"
 #include "dynamic_batch_scheduler.h"
@@ -942,6 +943,22 @@ TRITONBACKEND_RequestOutputBufferProperties(
   return nullptr;  // success
 }
 
+
+TRITONBACKEND_DECLSPEC TRITONSERVER_Error*
+TRITONBACKEND_RequestTrace(
+    TRITONBACKEND_Request* request, TRITONSERVER_InferenceTrace** trace)
+{
+#ifdef TRITON_ENABLE_TRACING
+  InferenceRequest* tr = reinterpret_cast<InferenceRequest*>(request);
+  *trace =
+      reinterpret_cast<TRITONSERVER_InferenceTrace*>(tr->TraceProxy()->Trace());
+  return nullptr;
+#else
+  return TRITONSERVER_ErrorNew(
+      TRITONSERVER_ERROR_UNSUPPORTED, "tracing is not supported");
+#endif  // TRITON_ENABLE_TRACING
+}
+
 TRITONAPI_DECLSPEC TRITONSERVER_Error*
 TRITONBACKEND_RequestRelease(
     TRITONBACKEND_Request* request, uint32_t release_flags)

diff --git a/src/dynamic_batch_scheduler.cc b/src/dynamic_batch_scheduler.cc
@@ -189,7 +189,7 @@ DynamicBatchScheduler::Enqueue(std::unique_ptr<InferenceRequest>& request)
   if (request->QueueStartNs() == 0) {
     request->CaptureQueueStartNs();
     INFER_TRACE_ACTIVITY(
-        request->Trace(), TRITONSERVER_TRACE_QUEUE_START,
+        request->TraceProxy(), TRITONSERVER_TRACE_QUEUE_START,
         request->QueueStartNs());
 #ifdef TRITON_ENABLE_TRACING
     request->TraceInputTensors(

diff --git a/src/ensemble_scheduler.cc b/src/ensemble_scheduler.cc
@@ -953,7 +953,7 @@ EnsembleContext::InitStep(
   RETURN_IF_ERROR(irequest->PrepareForInference());
 
 #ifdef TRITON_ENABLE_TRACING
-  auto& parent_trace = request_tracker_->Request()->Trace();
+  auto& parent_trace = request_tracker_->Request()->TraceProxy();
   if (parent_trace != nullptr) {
     irequest->SetTrace(parent_trace->SpawnChildTrace());
     irequest->Trace()->SetModelName(irequest->ModelName());
@@ -1281,7 +1281,7 @@ EnsembleScheduler::Enqueue(std::unique_ptr<InferenceRequest>& request)
   // scheduling process
   request->CaptureQueueStartNs();
   INFER_TRACE_ACTIVITY(
-      request->Trace(), TRITONSERVER_TRACE_QUEUE_START,
+      request->TraceProxy(), TRITONSERVER_TRACE_QUEUE_START,
       request->QueueStartNs());
 #ifdef TRITON_ENABLE_TRACING
   request->TraceInputTensors(

diff --git a/src/infer_request.h b/src/infer_request.h
@@ -322,8 +322,11 @@ class InferenceRequest {
   bool CacheKeyIsSet() const { return cache_key_is_set_; }
 
 #ifdef TRITON_ENABLE_TRACING
-  const std::shared_ptr<InferenceTraceProxy>& Trace() const { return trace_; }
-  std::shared_ptr<InferenceTraceProxy>* MutableTrace() { return &trace_; }
+  const std::shared_ptr<InferenceTraceProxy>& TraceProxy() const
+  {
+    return trace_;
+  }
+  std::shared_ptr<InferenceTraceProxy>* MutableTraceProxy() { return &trace_; }
   void SetTrace(const std::shared_ptr<InferenceTraceProxy>& trace)
   {
     trace_ = trace;

diff --git a/src/infer_trace.h b/src/infer_trace.h
@@ -74,7 +74,18 @@ class InferenceTrace {
     if ((level_ & TRITONSERVER_TRACE_LEVEL_TIMESTAMPS) > 0) {
       activity_fn_(
           reinterpret_cast<TRITONSERVER_InferenceTrace*>(this), activity,
-          timestamp_ns, userp_);
+          timestamp_ns, userp_, nullptr);
+    }
+  }
+
+  void Report(
+      const TRITONSERVER_InferenceTraceActivity activity, uint64_t timestamp_ns,
+      const char* tag)
+  {
+    if ((level_ & TRITONSERVER_TRACE_LEVEL_TIMESTAMPS) > 0) {
+      activity_fn_(
+          reinterpret_cast<TRITONSERVER_InferenceTrace*>(this), activity,
+          timestamp_ns, userp_, tag);
     }
   }
 
@@ -168,6 +179,8 @@ class InferenceTraceProxy {
         memory_type, memory_type_id);
   }
 
+  InferenceTrace* Trace() { return trace_; }
+
   std::shared_ptr<InferenceTraceProxy> SpawnChildTrace();
 
  private:

diff --git a/src/sequence_batch_scheduler.cc b/src/sequence_batch_scheduler.cc
@@ -522,7 +522,7 @@ SequenceBatchScheduler::Enqueue(std::unique_ptr<InferenceRequest>& irequest)
   // scheduling process
   irequest->CaptureQueueStartNs();
   INFER_TRACE_ACTIVITY(
-      irequest->Trace(), TRITONSERVER_TRACE_QUEUE_START,
+      irequest->TraceProxy(), TRITONSERVER_TRACE_QUEUE_START,
       irequest->QueueStartNs());
 
   // Record time at the beginning of the batcher queueing

diff --git a/src/server.cc b/src/server.cc
@@ -520,7 +520,7 @@ InferenceServer::InferAsync(std::unique_ptr<InferenceRequest>& request)
 #ifdef TRITON_ENABLE_STATS
   request->CaptureRequestStartNs();
   INFER_TRACE_ACTIVITY(
-      request->Trace(), TRITONSERVER_TRACE_REQUEST_START,
+      request->TraceProxy(), TRITONSERVER_TRACE_REQUEST_START,
       request->RequestStartNs());
 #endif  // TRITON_ENABLE_STATS
 

diff --git a/src/tritonserver.cc b/src/tritonserver.cc
@@ -26,6 +26,7 @@
 
 #include <string>
 #include <vector>
+
 #include "buffer_attributes.h"
 #include "cuda_utils.h"
 #include "infer_parameter.h"
@@ -894,6 +895,8 @@ TRITONSERVER_InferenceTraceActivityString(
       return "TENSOR_BACKEND_INPUT";
     case TRITONSERVER_TRACE_TENSOR_BACKEND_OUTPUT:
       return "TENSOR_BACKEND_OUTPUT";
+    case TRITONSERVER_TRACE_CUSTOM_ACTIVITY:
+      return "CUSTOM_ACTIVITY";
   }
 
   return "<unknown>";
@@ -1039,6 +1042,20 @@ TRITONSERVER_InferenceTraceModelVersion(
 #endif  // TRITON_ENABLE_TRACING
 }
 
+TRITONAPI_DECLSPEC TRITONSERVER_Error*
+TRITONSERVER_ReportActivity(
+    TRITONSERVER_InferenceTrace* trace, const char* tag, uint64_t timestamp)
+{
+#ifdef TRITON_ENABLE_TRACING
+  tc::InferenceTrace* ltrace = reinterpret_cast<tc::InferenceTrace*>(trace);
+  ltrace->Report(TRITONSERVER_TRACE_CUSTOM_ACTIVITY, timestamp, tag);
+  return nullptr;  // Success
+#else
+  return TRITONSERVER_ErrorNew(
+      TRITONSERVER_ERROR_UNSUPPORTED, "inference tracing not supported");
+#endif  // TRITON_ENABLE_TRACING
+}
+
 //
 // TRITONSERVER_ServerOptions
 //

diff --git a/src/tritonserver_stub.cc b/src/tritonserver_stub.cc
@@ -679,6 +679,10 @@ TRITONBACKEND_RequestOutputBufferProperties()
 {
 }
 TRITONAPI_DECLSPEC void
+TRITONBACKEND_RequestTrace()
+{
+}
+TRITONAPI_DECLSPEC void
 TRITONBACKEND_RequestRelease()
 {
 }
@@ -973,6 +977,10 @@ TRITONAPI_DECLSPEC void
 TRITONBACKEND_BackendAttributeAddPreferredInstanceGroup()
 {
 }
+TRITONAPI_DECLSPEC void
+TRITONSERVER_ReportActivity()
+{
+}
 
 TRITONAPI_DECLSPEC void
 TRITONCACHE_ApiVersion()