Add the ability for backends to add entries to the trace

include/triton/core/tritonbackend.h

@@ -1525,7 +1525,7 @@ TRITONBACKEND_BackendAttributeSetParallelModelInstanceLoading(
 ///
 /// \param batcher User-defined placeholder for backend to store and
 /// retrieve information about the batching strategy for this
-/// model.RITONBACKEND_ISPEC return a TRITONSERVER_Error indicating success or
+/// model. TRITONBACKEND_ISPEC returns a TRITONSERVER_Error indicating success or
 /// failure. \param model The backend model for which Triton is forming a batch.
 /// \return a TRITONSERVER_Error indicating success or failure.
 TRITONBACKEND_ISPEC TRITONSERVER_Error* TRITONBACKEND_ModelBatcherInitialize(

include/triton/core/tritonserver.h

@@ -712,7 +712,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
@@ -785,6 +786,16 @@ TRITONSERVER_DECLSPEC struct TRITONSERVER_Error* TRITONSERVER_InferenceTraceNew(
     TRITONSERVER_InferenceTraceActivityFn_t activity_fn,
     TRITONSERVER_InferenceTraceReleaseFn_t release_fn, void* trace_userp);
 
+/// 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 trace_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** trace_name);
+
 /// Create a new inference trace object. The caller takes ownership of
 /// the TRITONSERVER_InferenceTrace object and must call
 /// TRITONSERVER_InferenceTraceDelete to release the object.
@@ -818,6 +829,17 @@ TRITONSERVER_InferenceTraceTensorNew(
     TRITONSERVER_InferenceTraceTensorActivityFn_t tensor_activity_fn,
     TRITONSERVER_InferenceTraceReleaseFn_t release_fn, void* trace_userp);
 
+/// Report a trace activity. All the traces reported using this API will be
+/// using TRITONSERVER_TRACE_CUSTOM_ACTIVITY type.
+///
+/// \param trace The trace object.
+/// \param timestamp The timestamp associated with the trace activity.
+/// \param name The trace activity name.
+/// \return a TRITONSERVER_Error indicating success or failure.
+TRITONSERVER_DECLSPEC TRITONSERVER_Error*
+TRITONSERVER_InferenceTraceReportActivity(
+    TRITONSERVER_InferenceTrace* trace, uint64_t timestamp, const char* name);
+
 /// Delete a trace object.
 ///
 /// \param trace The trace object.

src/backend_model.cc

@@ -1191,35 +1191,31 @@ TRITONBACKEND_RequestOutputBufferProperties(
   return nullptr;  // success
 }
 
-TRITONAPI_DECLSPEC TRITONSERVER_Error*
-TRITONBACKEND_RequestRelease(
-    TRITONBACKEND_Request* request, uint32_t release_flags)
-{
-  InferenceRequest* tr = reinterpret_cast<InferenceRequest*>(request);
-  std::unique_ptr<InferenceRequest> ur(tr);
-  InferenceRequest::Release(std::move(ur), release_flags);
-  return nullptr;  // success
-}
-
-TRITONAPI_DECLSPEC TRITONSERVER_Error*
+TRITONBACKEND_DECLSPEC TRITONSERVER_Error*
 TRITONBACKEND_RequestTrace(
     TRITONBACKEND_Request* request, TRITONSERVER_InferenceTrace** trace)
 {
 #ifdef TRITON_ENABLE_TRACING
   InferenceRequest* tr = reinterpret_cast<InferenceRequest*>(request);
-  if (tr->TraceProxy() != nullptr) {
-    *trace = reinterpret_cast<TRITONSERVER_InferenceTrace*>(
-        tr->TraceProxy()->Trace());
-  } else {
-    *trace = nullptr;
-  }
-  return nullptr;  // success
+  *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)
+{
+  InferenceRequest* tr = reinterpret_cast<InferenceRequest*>(request);
+  std::unique_ptr<InferenceRequest> ur(tr);
+  InferenceRequest::Release(std::move(ur), release_flags);
+  return nullptr;  // success
+}
+
 ///
 /// TRITONBACKEND_State
 ///

src/infer_request.h

@@ -327,7 +327,7 @@ class InferenceRequest {
   {
     return trace_;
   }
-  std::shared_ptr<InferenceTraceProxy>* MutableTrace() { return &trace_; }
+  std::shared_ptr<InferenceTraceProxy>* MutableTraceProxy() { return &trace_; }
   void SetTrace(const std::shared_ptr<InferenceTraceProxy>& trace)
   {
     trace_ = trace;

src/infer_trace.h

@@ -67,6 +67,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(
@@ -121,6 +123,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
@@ -141,10 +144,12 @@ class InferenceTraceProxy {
   InferenceTrace* Trace() { return trace_; }
   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); }
 
@@ -170,6 +175,8 @@ class InferenceTraceProxy {
         memory_type, memory_type_id);
   }
 
+  InferenceTrace* Trace() { return trace_; }
+
   std::shared_ptr<InferenceTraceProxy> SpawnChildTrace();
 
  private:

src/tritonserver.cc

@@ -915,6 +915,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>";
@@ -1046,6 +1048,20 @@ 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_InferenceTraceModelVersion(
     TRITONSERVER_InferenceTrace* trace, int64_t* model_version)
@@ -1073,6 +1089,19 @@ TRITONSERVER_InferenceTraceSpawnChildTrace(
   } else {
     *child_trace = nullptr;
   }
+  return nullptr;
+#else
+  return TRITONSERVER_ErrorNew(
+      TRITONSERVER_ERROR_UNSUPPORTED, "inference tracing not supported");
+#endif  // TRITON_ENABLE_TRACING
+}
+
+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(

src/tritonserver_stub.cc

@@ -182,6 +182,10 @@ TRITONSERVER_InferenceTraceSpawnChildTrace()
 {
 }
 TRITONAPI_DECLSPEC void
+TRITONSERVER_InferenceTraceName()
+{
+}
+TRITONAPI_DECLSPEC void
 TRITONSERVER_InferenceRequestNew()
 {
 }
@@ -691,6 +695,10 @@ TRITONBACKEND_RequestOutputBufferProperties()
 {
 }
 TRITONAPI_DECLSPEC void
+TRITONBACKEND_RequestTrace()
+{
+}
+TRITONAPI_DECLSPEC void
 TRITONBACKEND_RequestRelease()
 {
 }
@@ -1022,6 +1030,10 @@ TRITONAPI_DECLSPEC void
 TRITONBACKEND_BackendAttributeAddPreferredInstanceGroup()
 {
 }
+TRITONAPI_DECLSPEC void
+TRITONSERVER_InferenceTraceReportActivity()
+{
+}
 
 TRITONAPI_DECLSPEC void
 TRITONBACKEND_BackendAttributeSetParallelModelInstanceLoading()