Add the ability for backends to add entries to the trace

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
 ///

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
@@ -810,6 +811,15 @@ TRITONSERVER_DECLSPEC TRITONSERVER_Error* TRITONSERVER_InferenceTraceTensorNew(
     TRITONSERVER_InferenceTraceTensorActivityFn_t tensor_activity_fn,
     TRITONSERVER_InferenceTraceReleaseFn_t release_fn, void* trace_userp);
 
+/// Report a trace activity.
+///
+/// \param trace The trace object.
+/// \param timestamp The timestamp associated with the trace activity.
+/// \return a TRITONSERVER_Error indicating success or failure.
+TRITONSERVER_DECLSPEC TRITONSERVER_Error*
+TRITONSERVER_InferenceTraceReportActivity(
+    TRITONSERVER_InferenceTrace* trace, uint64_t timestamp);
+
 /// Delete a trace object.
 ///
 /// \param trace The trace object.
@@ -847,6 +857,34 @@ TRITONSERVER_DECLSPEC TRITONSERVER_Error* TRITONSERVER_InferenceTraceParentId(
 TRITONSERVER_DECLSPEC TRITONSERVER_Error* TRITONSERVER_InferenceTraceModelName(
     TRITONSERVER_InferenceTrace* trace, const char** model_name);
 
+/// Set the name associated with a trace.
+///
+/// \param trace The trace.
+/// \param name The name associated with the trace.
+/// \return a TRITONSERVER_Error indicating success or failure.
+TRITONSERVER_DECLSPEC TRITONSERVER_Error*
+TRITONSERVER_InferenceTraceSetModelName(
+    TRITONSERVER_InferenceTrace* trace, const char* name);
+
+/// Get the name associated with a trace. The caller does
+/// not own the returned string and must not modify or delete it. The
+/// lifetime of the returned string extends only as long as 'trace'.
+///
+/// \param trace The trace.
+/// \param name Returns the name associated with the trace.
+/// \return a TRITONSERVER_Error indicating success or failure.
+TRITONSERVER_DECLSPEC TRITONSERVER_Error* TRITONSERVER_InferenceTraceName(
+    TRITONSERVER_InferenceTrace* trace, const char** name);
+
+/// Set the name associated with a trace.
+///
+/// \param trace The trace.
+/// \param trace_name The anme associated with a trace.
+/// \return a TRITONSERVER_Error indicating success or failure.
+TRITONSERVER_DECLSPEC TRITONSERVER_Error*
+TRITONSERVER_InferenceTraceSetName(
+    TRITONSERVER_InferenceTrace* trace, const char* trace_name);
+
 /// Get the version of the model associated with a trace.
 ///
 /// \param trace The trace.

src/backend_model.cc

@@ -942,6 +942,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)

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(

src/ensemble_scheduler.cc

@@ -953,12 +953,12 @@ 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());
-    irequest->Trace()->SetRequestId(irequest->Id());
-    irequest->Trace()->SetModelVersion(irequest->ActualModelVersion());
+    irequest->TraceProxy()->SetModelName(irequest->ModelName());
+    irequest->TraceProxy()->SetRequestId(irequest->Id());
+    irequest->TraceProxy()->SetModelVersion(irequest->ActualModelVersion());
   }
 #endif
 
@@ -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(

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;

src/infer_trace.h

@@ -66,6 +66,8 @@ class InferenceTrace {
   void SetModelName(const std::string& n) { model_name_ = n; }
   void SetModelVersion(int64_t v) { model_version_ = v; }
   void SetRequestId(const std::string& request_id) { request_id_ = request_id; }
+  void SetTraceName(const std::string& trace_name) { trace_name_ = trace_name; }
+  const std::string& TraceName() { return trace_name_; }
 
   // Report trace activity.
   void Report(
@@ -120,6 +122,7 @@ class InferenceTrace {
   std::string model_name_;
   int64_t model_version_;
   std::string request_id_;
+  std::string trace_name_;
 
   // Maintain next id statically so that trace id is unique even
   // across traces
@@ -139,10 +142,12 @@ class InferenceTraceProxy {
   ~InferenceTraceProxy() { trace_->Release(); }
   int64_t Id() const { return trace_->Id(); }
   int64_t ParentId() const { return trace_->ParentId(); }
+  const std::string& TraceName() { return trace_->TraceName(); }
   const std::string& ModelName() const { return trace_->ModelName(); }
   const std::string& RequestId() const { return trace_->RequestId(); }
   int64_t ModelVersion() const { return trace_->ModelVersion(); }
   void SetModelName(const std::string& n) { trace_->SetModelName(n); }
+  void SetTraceName(const std::string& n) { trace_->SetTraceName(n); }
   void SetRequestId(const std::string& n) { trace_->SetRequestId(n); }
   void SetModelVersion(int64_t v) { trace_->SetModelVersion(v); }
 
@@ -168,6 +173,8 @@ class InferenceTraceProxy {
         memory_type, memory_type_id);
   }
 
+  InferenceTrace* Trace() { return trace_; }
+
   std::shared_ptr<InferenceTraceProxy> SpawnChildTrace();
 
  private:

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

src/server.cc

@@ -522,7 +522,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
 

src/tritonserver.cc

@@ -894,6 +894,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>";
@@ -1025,6 +1027,34 @@ TRITONSERVER_InferenceTraceRequestId(
 #endif  // TRITON_ENABLE_TRACING
 }
 
+TRITONAPI_DECLSPEC TRITONSERVER_Error*
+TRITONSERVER_InferenceTraceName(
+    TRITONSERVER_InferenceTrace* trace, const char** trace_name)
+{
+#ifdef TRITON_ENABLE_TRACING
+  tc::InferenceTrace* ltrace = reinterpret_cast<tc::InferenceTrace*>(trace);
+  *trace_name = ltrace->TraceName().c_str();
+  return nullptr;  // Success
+#else
+  return TRITONSERVER_ErrorNew(
+      TRITONSERVER_ERROR_UNSUPPORTED, "inference tracing not supported");
+#endif  // TRITON_ENABLE_TRACING
+}
+
+TRITONAPI_DECLSPEC TRITONSERVER_Error*
+TRITONSERVER_InferenceTraceSetName(
+    TRITONSERVER_InferenceTrace* trace, const char* trace_name)
+{
+#ifdef TRITON_ENABLE_TRACING
+  tc::InferenceTrace* ltrace = reinterpret_cast<tc::InferenceTrace*>(trace);
+  ltrace->SetTraceName(trace_name);
+  return nullptr;  // Success
+#else
+  return TRITONSERVER_ErrorNew(
+      TRITONSERVER_ERROR_UNSUPPORTED, "inference tracing not supported");
+#endif  // TRITON_ENABLE_TRACING
+}
+
 TRITONAPI_DECLSPEC TRITONSERVER_Error*
 TRITONSERVER_InferenceTraceModelVersion(
     TRITONSERVER_InferenceTrace* trace, int64_t* model_version)
@@ -1039,6 +1069,20 @@ TRITONSERVER_InferenceTraceModelVersion(
 #endif  // TRITON_ENABLE_TRACING
 }
 
+TRITONAPI_DECLSPEC TRITONSERVER_Error*
+TRITONSERVER_InferenceTraceReportActivity(
+    TRITONSERVER_InferenceTrace* trace, uint64_t timestamp)
+{
+#ifdef TRITON_ENABLE_TRACING
+  tc::InferenceTrace* ltrace = reinterpret_cast<tc::InferenceTrace*>(trace);
+  ltrace->Report(TRITONSERVER_TRACE_CUSTOM_ACTIVITY, timestamp);
+  return nullptr;  // Success
+#else
+  return TRITONSERVER_ErrorNew(
+      TRITONSERVER_ERROR_UNSUPPORTED, "inference tracing not supported");
+#endif  // TRITON_ENABLE_TRACING
+}
+
 //
 // TRITONSERVER_ServerOptions
 //

src/tritonserver_stub.cc

@@ -178,6 +178,14 @@ TRITONSERVER_InferenceTraceRequestId()
 {
 }
 TRITONAPI_DECLSPEC void
+TRITONSERVER_InferenceTraceName()
+{
+}
+TRITONAPI_DECLSPEC void
+TRITONSERVER_InferenceTraceSetName()
+{
+}
+TRITONAPI_DECLSPEC void
 TRITONSERVER_InferenceRequestNew()
 {
 }
@@ -679,6 +687,10 @@ TRITONBACKEND_RequestOutputBufferProperties()
 {
 }
 TRITONAPI_DECLSPEC void
+TRITONBACKEND_RequestTrace()
+{
+}
+TRITONAPI_DECLSPEC void
 TRITONBACKEND_RequestRelease()
 {
 }
@@ -973,6 +985,10 @@ TRITONAPI_DECLSPEC void
 TRITONBACKEND_BackendAttributeAddPreferredInstanceGroup()
 {
 }
+TRITONAPI_DECLSPEC void
+TRITONSERVER_InferenceTraceReportActivity()
+{
+}
 
 TRITONAPI_DECLSPEC void
 TRITONCACHE_ApiVersion()