Revision 3

Author: oandreeva-nv

src/tracer.cc

@@ -305,10 +305,7 @@ TraceManager::Trace::~Trace()
     setting_->WriteTrace(streams_);
   } else if (setting_->mode_ == TRACE_MODE_OPENTELEMETRY) {
 #ifndef _WIN32
-    auto root_span = opentelemetry::nostd::get<
-          opentelemetry::nostd::shared_ptr<otel_trace_api::Span>>(
-          this->otel_context_.GetValue(kRootSpan));
-    EndSpanNow(root_span);
+    EndSpan(kRootSpan);
 #else
     LOG_ERROR << "Unsupported trace mode: "
               << TraceManager::InferenceTraceModeString(setting_->mode_);
@@ -340,12 +337,7 @@ TraceManager::Trace::CaptureTimestamp(
           << "{\"name\":\"" << name << "\",\"ns\":" << timestamp_ns << "}]}";
     } else if (setting_->mode_ == TRACE_MODE_OPENTELEMETRY) {
 #ifndef _WIN32
-      auto root_span = opentelemetry::nostd::get<
-          opentelemetry::nostd::shared_ptr<otel_trace_api::Span>>(
-          this->otel_context_.GetValue(kRootSpan));
-      root_span->AddEvent(
-          name, otel_common::SystemTimestamp{
-                    time_offset_ + std::chrono::nanoseconds{timestamp_ns}});
+      AddEvent(kRootSpan, name, timestamp_ns);
 #else
       LOG_ERROR << "Unsupported trace mode: "
                 << TraceManager::InferenceTraceModeString(setting_->mode_);
@@ -379,26 +371,22 @@ TraceManager::Trace::InitTracer(
       std::chrono::duration_cast<std::chrono::nanoseconds>(
           std::chrono::steady_clock::now().time_since_epoch())
           .count();
-  auto otel_start_timestamp = otel_common::SystemTimestamp{
-      time_offset_ + std::chrono::nanoseconds{steady_timestamp_ns}};
-  auto root_span = InitSpan(
-      "InferRequest", otel_start_timestamp, steady_timestamp_ns,
-      true /*is_root_span*/);
+  auto root_span =
+      StartSpan("InferRequest", steady_timestamp_ns, true /*is_root_span*/);
   // Initializing OTel context and storring "InferRequest" span as a root span
   // to keep it alive for the duration of the request.
   otel_context_ = opentelemetry::context::Context({kRootSpan, root_span});
 }
 
 opentelemetry::nostd::shared_ptr<otel_trace_api::Span>
-TraceManager::Trace::InitSpan(
-    std::string name, const otel_common::SystemTimestamp& timestamp_ns,
-    const uint64_t& raw_timestamp_ns, bool is_root_span,
-    std::string parent_span_key)
+TraceManager::Trace::StartSpan(
+    std::string display_name, const uint64_t& raw_timestamp_ns,
+    bool is_root_span, std::string parent_span_key)
 {
-  opentelemetry::nostd::shared_ptr<otel_trace_api::Span> span{nullptr};
   otel_trace_api::StartSpanOptions options;
   options.kind = otel_trace_api::SpanKind::kServer;
-  options.start_system_time = timestamp_ns;
+  options.start_system_time =
+      time_offset_ + std::chrono::nanoseconds{raw_timestamp_ns};
   options.start_steady_time =
       otel_common::SteadyTimestamp{std::chrono::nanoseconds{raw_timestamp_ns}};
 
@@ -410,29 +398,32 @@ TraceManager::Trace::InitSpan(
         otel_context_.GetValue(parent_span_key));
     options.parent = parent_span->GetContext();
   }
-  span = provider_->GetTracer(kTritonTracer)->StartSpan(name, options);
-  return span;
+  return provider_->GetTracer(kTritonTracer)->StartSpan(display_name, options);
 }
 
 void
-TraceManager::Trace::EndSpanNow(
-    opentelemetry::nostd::shared_ptr<otel_trace_api::Span> span)
+TraceManager::Trace::EndSpan(std::string span_key)
 {
-  if (span != nullptr) {
-    auto timestamp_ns = std::chrono::duration_cast<std::chrono::nanoseconds>(
-                            std::chrono::steady_clock::now().time_since_epoch())
-                            .count();
-    EndSpan(span, timestamp_ns);
-  }
+  auto timestamp_ns = std::chrono::duration_cast<std::chrono::nanoseconds>(
+                          std::chrono::steady_clock::now().time_since_epoch())
+                          .count();
+  EndSpan(span_key, timestamp_ns);
 }
 
 
 void
 TraceManager::Trace::EndSpan(
-    opentelemetry::nostd::shared_ptr<otel_trace_api::Span> span,
-    const uint64_t& raw_timestamp_ns)
+    std::string span_key, const uint64_t& raw_timestamp_ns)
 {
-  if (span != nullptr) {
+  if (otel_context_.HasKey(span_key)) {
+    auto span = opentelemetry::nostd::get<
+        opentelemetry::nostd::shared_ptr<otel_trace_api::Span>>(
+        otel_context_.GetValue(span_key));
+
+    if (span == nullptr) {
+      return;
+    }
+
     otel_trace_api::EndSpanOptions end_options;
     end_options.end_steady_time = otel_common::SteadyTimestamp{
         std::chrono::nanoseconds{raw_timestamp_ns}};
@@ -453,106 +444,16 @@ TraceManager::Trace::ReportToOpenTelemetry(
   LOG_TRITONSERVER_ERROR(
       TRITONSERVER_InferenceTraceId(trace, &id), "getting trace id");
   current_span_key = current_span_key + std::to_string(id);
-  switch (activity) {
-    case TRITONSERVER_TRACE_REQUEST_START: {
-      const char* model_name;
-      int64_t model_version;
-      uint64_t parent_id;
-      const char* request_id;
-
-      LOG_TRITONSERVER_ERROR(
-          TRITONSERVER_InferenceTraceModelName(trace, &model_name),
-          "getting model name");
-      LOG_TRITONSERVER_ERROR(
-          TRITONSERVER_InferenceTraceModelVersion(trace, &model_version),
-          "getting model version");
-      LOG_TRITONSERVER_ERROR(
-          TRITONSERVER_InferenceTraceParentId(trace, &parent_id),
-          "getting trace parent id");
-      LOG_TRITONSERVER_ERROR(
-          TRITONSERVER_InferenceTraceRequestId(trace, &request_id),
-          "getting request id");
-
-      // Currently, only 2 types of sub-spans are supported:
-      // request span and compute span. Compute span is a leaf span
-      // and can not be a parent of any sub-span. If parent_id=0,
-      // then current model is either a standalone model, or an ensemble model.
-      // In both of the above cases, the parent of the new request sub-span
-      // is the kRootSpan.
-      std::string parent_span_key = kRootSpan;
-      if (parent_id != 0) {
-        // If parent_id > 0, then this is a child trace, spawned from
-        // the ensamble's main request. For this instance, the parent
-        // span is the ensembles's request span.
-        parent_span_key = kRequestSpan + std::to_string(parent_id);
-      }
-      auto request_span = this->InitSpan(
-          model_name,
-          this->time_offset_ + std::chrono::nanoseconds{timestamp_ns},
-          timestamp_ns, false /*is_root_span*/, parent_span_key);
-      request_span->SetAttribute("triton.model_name", model_name);
-      request_span->SetAttribute("triton.model_version", model_version);
-      request_span->SetAttribute("triton.trace_id", id);
-      request_span->SetAttribute("triton.trace_parent_id", parent_id);
-      request_span->AddEvent(
-          TRITONSERVER_InferenceTraceActivityString(activity),
-          this->time_offset_ + std::chrono::nanoseconds{timestamp_ns});
-      // Adding span to the OpenTelemetry context, to keep it alive and active
-      // until we end it
-      this->otel_context_ =
-          this->otel_context_.SetValue(current_span_key, request_span);
-      break;
-    }
-    case TRITONSERVER_TRACE_COMPUTE_START: {
-      // Any compute span is a child of a request span with the same id.
-      auto compute_span = this->InitSpan(
-          "compute",
-          this->time_offset_ + std::chrono::nanoseconds{timestamp_ns},
-          timestamp_ns, false /*is_root_span*/,
-          kRequestSpan + std::to_string(id));
-      compute_span->AddEvent(
-          TRITONSERVER_InferenceTraceActivityString(activity),
-          this->time_offset_ + std::chrono::nanoseconds{timestamp_ns});
-      // Adding span to the OpenTelemetry context, to keep it alive and active
-      // until we end it
-      this->otel_context_ =
-          this->otel_context_.SetValue(current_span_key, compute_span);
-      break;
-    }
-    case TRITONSERVER_TRACE_QUEUE_START:
-    case TRITONSERVER_TRACE_COMPUTE_INPUT_END:
-    case TRITONSERVER_TRACE_COMPUTE_OUTPUT_START: {
-      if (this->otel_context_.HasKey(current_span_key)) {
-        auto span = opentelemetry::nostd::get<
-            opentelemetry::nostd::shared_ptr<otel_trace_api::Span>>(
-            this->otel_context_.GetValue(current_span_key));
-        span->AddEvent(
-            TRITONSERVER_InferenceTraceActivityString(activity),
-            this->time_offset_ + std::chrono::nanoseconds{timestamp_ns});
-      }
-      break;
-    }
-    case TRITONSERVER_TRACE_COMPUTE_END:
-    case TRITONSERVER_TRACE_REQUEST_END: {
-      if (this->otel_context_.HasKey(current_span_key)) {
-        auto span = opentelemetry::nostd::get<
-            opentelemetry::nostd::shared_ptr<otel_trace_api::Span>>(
-            this->otel_context_.GetValue(current_span_key));
-        span->AddEvent(
-            TRITONSERVER_InferenceTraceActivityString(activity),
-            this->time_offset_ + std::chrono::nanoseconds{timestamp_ns});
-        this->EndSpan(span, timestamp_ns);
-      }
-      break;
-    }
-    case TRITONSERVER_TRACE_TENSOR_QUEUE_INPUT:
-    case TRITONSERVER_TRACE_TENSOR_BACKEND_INPUT:
-    case TRITONSERVER_TRACE_TENSOR_BACKEND_OUTPUT:
-    default: {
-      LOG_ERROR << "Unsupported activity: "
-                << TRITONSERVER_InferenceTraceActivityString(activity);
-      break;
-    }
+
+  MaybeStartSpan(current_span_key, trace, activity, timestamp_ns, id);
+
+  AddEvent(
+      current_span_key, TRITONSERVER_InferenceTraceActivityString(activity),
+      timestamp_ns);
+
+  if (activity == TRITONSERVER_TRACE_REQUEST_END ||
+      activity == TRITONSERVER_TRACE_COMPUTE_END) {
+    EndSpan(current_span_key, timestamp_ns);
   }
 }
 
@@ -590,6 +491,69 @@ TraceManager::Trace::GetSpanNameForActivity(
   return span_name;
 }
 
+void
+TraceManager::Trace::AddEvent(
+    std::string span_key, std::string event, uint64_t timestamp)
+{
+  if (otel_context_.HasKey(span_key)) {
+    auto span = opentelemetry::nostd::get<
+        opentelemetry::nostd::shared_ptr<otel_trace_api::Span>>(
+        otel_context_.GetValue(span_key));
+    span->AddEvent(event, time_offset_ + std::chrono::nanoseconds{timestamp});
+  }
+}
+
+void
+TraceManager::Trace::MaybeStartSpan(
+    std::string span_key, TRITONSERVER_InferenceTrace* trace,
+    TRITONSERVER_InferenceTraceActivity activity, uint64_t timestamp_ns,
+    uint64_t id)
+{
+  if (activity != TRITONSERVER_TRACE_REQUEST_START &&
+      activity != TRITONSERVER_TRACE_COMPUTE_START) {
+    return;
+  }
+
+  // Currently, only 2 types of sub-spans are supported:
+  // request span and compute span. Compute span is a leaf span
+  // and can not be a parent of any sub-span. If parent_id==0,
+  // then current model is either a standalone model, or an ensemble model.
+  // In both cases, the parent of the new request sub-span is the kRootSpan.
+  // If parent_id > 0, then this is a child trace, spawned from
+  // the ensamble's main request. For this instance, the parent
+  // span is the ensembles's request span.
+  uint64_t parent_id;
+  LOG_TRITONSERVER_ERROR(
+      TRITONSERVER_InferenceTraceParentId(trace, &parent_id),
+      "getting trace parent id");
+  std::string parent_span_key =
+      (parent_id != 0) ? kRequestSpan + std::to_string(parent_id) : kRootSpan;
+
+  std::string display_name = "compute";
+  const char* model_name;
+  if (activity == TRITONSERVER_TRACE_REQUEST_START) {
+    LOG_TRITONSERVER_ERROR(
+        TRITONSERVER_InferenceTraceModelName(trace, &model_name),
+        "getting model name");
+    display_name = model_name;
+  }
+
+  auto span = StartSpan(
+      display_name, timestamp_ns, false /*is_root_span*/, parent_span_key);
+
+  if (activity == TRITONSERVER_TRACE_REQUEST_START) {
+    int64_t model_version;
+    LOG_TRITONSERVER_ERROR(
+        TRITONSERVER_InferenceTraceModelVersion(trace, &model_version),
+        "getting model version");
+    span->SetAttribute("triton.model_name", model_name);
+    span->SetAttribute("triton.model_version", model_version);
+    span->SetAttribute("triton.trace_id", id);
+    span->SetAttribute("triton.trace_parent_id", parent_id);
+  }
+
+  otel_context_ = otel_context_.SetValue(span_key, span);
+}
 #endif
 
 void
@@ -650,17 +614,15 @@ TraceManager::TraceActivity(
 
   std::stringstream* ss = nullptr;
   {
-    if (ts->setting_->mode_ == TRACE_MODE_TRITON) {
-      if (ts->streams_.find(id) == ts->streams_.end()) {
-        std::unique_ptr<std::stringstream> stream(new std::stringstream());
-        ss = stream.get();
-        ts->streams_.emplace(id, std::move(stream));
-      } else {
-        ss = ts->streams_[id].get();
-        // If the string stream is not newly created, add "," as there is
-        // already content in the string stream
-        *ss << ",";
-      }
+    if (ts->streams_.find(id) == ts->streams_.end()) {
+      std::unique_ptr<std::stringstream> stream(new std::stringstream());
+      ss = stream.get();
+      ts->streams_.emplace(id, std::move(stream));
+    } else {
+      ss = ts->streams_[id].get();
+      // If the string stream is not newly created, add "," as there is
+      // already content in the string stream
+      *ss << ",";
     }
   }
   // If 'activity' is TRITONSERVER_TRACE_REQUEST_START then collect
@@ -684,26 +646,22 @@ TraceManager::TraceActivity(
         TRITONSERVER_InferenceTraceRequestId(trace, &request_id),
         "getting request id");
 
-    if (ts->setting_->mode_ == TRACE_MODE_TRITON) {
-      *ss << "{\"id\":" << id << ",\"model_name\":\"" << model_name
-          << "\",\"model_version\":" << model_version;
+    *ss << "{\"id\":" << id << ",\"model_name\":\"" << model_name
+        << "\",\"model_version\":" << model_version;
 
-      if (std::string(request_id) != "") {
-        *ss << ",\"request_id\":\"" << request_id << "\"";
-      }
+    if (std::string(request_id) != "") {
+      *ss << ",\"request_id\":\"" << request_id << "\"";
+    }
 
-      if (parent_id != 0) {
-        *ss << ",\"parent_id\":" << parent_id;
-      }
-      *ss << "},";
+    if (parent_id != 0) {
+      *ss << ",\"parent_id\":" << parent_id;
     }
+    *ss << "},";
   }
 
-  if (ts->setting_->mode_ == TRACE_MODE_TRITON) {
-    *ss << "{\"id\":" << id << ",\"timestamps\":["
-        << "{\"name\":\"" << TRITONSERVER_InferenceTraceActivityString(activity)
-        << "\",\"ns\":" << timestamp_ns << "}]}";
-  }
+  *ss << "{\"id\":" << id << ",\"timestamps\":["
+      << "{\"name\":\"" << TRITONSERVER_InferenceTraceActivityString(activity)
+      << "\",\"ns\":" << timestamp_ns << "}]}";
 }
 
 void