Add request id to trace (#169)

Author: Tabrizian

include/triton/core/tritonserver.h

@@ -91,7 +91,7 @@ struct TRITONSERVER_MetricFamily;
 ///   }
 ///
 #define TRITONSERVER_API_VERSION_MAJOR 1
-#define TRITONSERVER_API_VERSION_MINOR 18
+#define TRITONSERVER_API_VERSION_MINOR 19
 
 /// Get the TRITONBACKEND API version supported by the Triton shared
 /// library. This value can be compared against the
@@ -857,6 +857,17 @@ TRITONSERVER_DECLSPEC TRITONSERVER_Error*
 TRITONSERVER_InferenceTraceModelVersion(
     TRITONSERVER_InferenceTrace* trace, int64_t* model_version);
 
+/// Get the request id 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 request_id Returns the version of the model associated
+/// with the trace.
+/// \return a TRITONSERVER_Error indicating success or failure.
+TRITONSERVER_DECLSPEC TRITONSERVER_Error* TRITONSERVER_InferenceTraceRequestId(
+    TRITONSERVER_InferenceTrace* trace, const char** request_id);
+
 /// TRITONSERVER_InferenceRequest
 ///
 /// Object representing an inference request. The inference request

src/ensemble_scheduler.cc

@@ -957,6 +957,7 @@ EnsembleContext::InitStep(
   if (parent_trace != nullptr) {
     irequest->SetTrace(parent_trace->SpawnChildTrace());
     irequest->Trace()->SetModelName(irequest->ModelName());
+    irequest->Trace()->SetRequestId(irequest->Id());
     irequest->Trace()->SetModelVersion(irequest->ActualModelVersion());
   }
 #endif

src/infer_trace.h

@@ -61,9 +61,11 @@ class InferenceTrace {
 
   const std::string& ModelName() const { return model_name_; }
   int64_t ModelVersion() const { return model_version_; }
+  const std::string& RequestId() const { return request_id_; }
 
   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; }
 
   // Report trace activity.
   void Report(
@@ -117,6 +119,7 @@ class InferenceTrace {
 
   std::string model_name_;
   int64_t model_version_;
+  std::string request_id_;
 
   // Maintain next id statically so that trace id is unique even
   // across traces
@@ -137,8 +140,10 @@ class InferenceTraceProxy {
   int64_t Id() const { return trace_->Id(); }
   int64_t ParentId() const { return trace_->ParentId(); }
   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 SetRequestId(const std::string& n) { trace_->SetRequestId(n); }
   void SetModelVersion(int64_t v) { trace_->SetModelVersion(v); }
 
   void Report(

src/tritonserver.cc

@@ -998,6 +998,20 @@ TRITONSERVER_InferenceTraceModelName(
 #endif  // TRITON_ENABLE_TRACING
 }
 
+TRITONAPI_DECLSPEC TRITONSERVER_Error*
+TRITONSERVER_InferenceTraceRequestId(
+    TRITONSERVER_InferenceTrace* trace, const char** request_id)
+{
+#ifdef TRITON_ENABLE_TRACING
+  tc::InferenceTrace* ltrace = reinterpret_cast<tc::InferenceTrace*>(trace);
+  *request_id = ltrace->RequestId().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)
@@ -2907,6 +2921,7 @@ TRITONSERVER_ServerInferAsync(
     tc::InferenceTrace* ltrace = reinterpret_cast<tc::InferenceTrace*>(trace);
     ltrace->SetModelName(lrequest->ModelName());
     ltrace->SetModelVersion(lrequest->ActualModelVersion());
+    ltrace->SetRequestId(lrequest->Id());
 
     lrequest->SetTrace(std::make_shared<tc::InferenceTraceProxy>(ltrace));
 #else

src/tritonserver_stub.cc

@@ -174,6 +174,10 @@ TRITONSERVER_InferenceTraceModelVersion()
 {
 }
 TRITONAPI_DECLSPEC void
+TRITONSERVER_InferenceTraceRequestId()
+{
+}
+TRITONAPI_DECLSPEC void
 TRITONSERVER_InferenceRequestNew()
 {
 }