modular otel_trace

Signed-off-by: Zhang Haotong <[email protected]>

Author: Zhang Haotong

tensorrt_llm/executor/result.py

@@ -11,9 +11,7 @@
 import torch
 import torch.nn.functional as F
 
-from tensorrt_llm.llmapi.otel_tracing import (
-    SpanAttributes, SpanKind, extract_trace_context, global_otlp_tracer,
-    insufficient_request_metrics_warning)
+from tensorrt_llm.llmapi import tracing
 
 from .._utils import nvtx_range_debug
 from ..bindings import executor as tllm
@@ -400,23 +398,23 @@ def do_tracing(
         output: CompletionOutput,
         req_perf_metrics_dict: Optional[dict[str, float]] = None,
     ):
-        if not global_otlp_tracer():
+        if not tracing.global_otlp_tracer():
             return
 
         metrics_dict = self.metrics_dict
         if not metrics_dict or not req_perf_metrics_dict:
             # Insufficient request metrics available; trace generation aborted.
-            insufficient_request_metrics_warning()
+            tracing.insufficient_request_metrics_warning()
             return
 
-        trace_context = extract_trace_context(self.trace_headers)
+        trace_context = tracing.extract_trace_context(self.trace_headers)
         sampling_params = self.sampling_params
 
         # TODO: Add request arrival time
         arrival_time = time.time() - metrics_dict.get(MetricNames.E2E, -1)
-        with global_otlp_tracer().start_as_current_span(
+        with tracing.global_otlp_tracer().start_as_current_span(
                 "llm_request",
-                kind=SpanKind.SERVER,
+                kind=tracing.SpanKind.SERVER,
                 context=trace_context,
                 start_time=int(arrival_time * 1e9),
         ) as span:
@@ -428,38 +426,41 @@ def safe_set_attr(span, attr, value):
             e2e_time = metrics_dict.get(MetricNames.E2E, -1)
             safe_set_attr(
                 span,
-                SpanAttributes.GEN_AI_REQUEST_TEMPERATURE,
+                tracing.SpanAttributes.GEN_AI_REQUEST_TEMPERATURE,
                 sampling_params.temperature,
             )
-            safe_set_attr(span, SpanAttributes.GEN_AI_REQUEST_TOP_P,
+            safe_set_attr(span, tracing.SpanAttributes.GEN_AI_REQUEST_TOP_P,
                           sampling_params.top_p)
             safe_set_attr(
                 span,
-                SpanAttributes.GEN_AI_REQUEST_MAX_TOKENS,
+                tracing.SpanAttributes.GEN_AI_REQUEST_MAX_TOKENS,
                 sampling_params.max_tokens,
             )
-            safe_set_attr(span, SpanAttributes.GEN_AI_REQUEST_N,
+            safe_set_attr(span, tracing.SpanAttributes.GEN_AI_REQUEST_N,
                           sampling_params.n)
             # TODO: Add prompt info in result base
             safe_set_attr(
                 span,
-                SpanAttributes.GEN_AI_USAGE_PROMPT_TOKENS,
+                tracing.SpanAttributes.GEN_AI_USAGE_PROMPT_TOKENS,
                 getattr(self.postproc_params.postproc_args, "num_prompt_tokens",
                         None) if self.postproc_params
                 and self.postproc_params.postproc_args else None,
             )
-            safe_set_attr(span, SpanAttributes.GEN_AI_USAGE_COMPLETION_TOKENS,
+            safe_set_attr(span,
+                          tracing.SpanAttributes.GEN_AI_USAGE_COMPLETION_TOKENS,
                           output.length)
             safe_set_attr(
                 span,
-                SpanAttributes.GEN_AI_LATENCY_TIME_TO_FIRST_TOKEN,
+                tracing.SpanAttributes.GEN_AI_LATENCY_TIME_TO_FIRST_TOKEN,
                 metrics_dict.get(MetricNames.TTFT, -1),
             )
-            safe_set_attr(span, SpanAttributes.GEN_AI_LATENCY_E2E, e2e_time)
-            safe_set_attr(span, SpanAttributes.GEN_AI_REQUEST_ID, self.id)
+            safe_set_attr(span, tracing.SpanAttributes.GEN_AI_LATENCY_E2E,
+                          e2e_time)
+            safe_set_attr(span, tracing.SpanAttributes.GEN_AI_REQUEST_ID,
+                          self.id)
             safe_set_attr(
                 span,
-                SpanAttributes.GEN_AI_LATENCY_TIME_IN_QUEUE,
+                tracing.SpanAttributes.GEN_AI_LATENCY_TIME_IN_QUEUE,
                 metrics_dict.get(MetricNames.REQUEST_QUEUE_TIME, -1),
             )
 

tensorrt_llm/llmapi/llm.py

@@ -16,7 +16,7 @@
 from tensorrt_llm.inputs.data import TextPrompt
 from tensorrt_llm.inputs.multimodal import MultimodalParams
 from tensorrt_llm.inputs.registry import DefaultInputProcessor
-from tensorrt_llm.llmapi.otel_tracing import init_tracer
+from tensorrt_llm.llmapi import tracing
 
 from .._utils import nvtx_range_debug
 from ..bindings import executor as tllm
@@ -213,7 +213,7 @@ def __init__(self,
 
         try:
             if self.args.otlp_traces_endpoint:
-                init_tracer("trt.llm", self.args.otlp_traces_endpoint)
+                tracing.init_tracer("trt.llm", self.args.otlp_traces_endpoint)
                 logger.info(
                     f"Initialized OTLP tracer successfully, endpoint: {self.args.otlp_traces_endpoint}"
                 )

tensorrt_llm/llmapi/otel_tracing.py renamed to tensorrt_llm/llmapi/tracing.py

@@ -1,12 +1,27 @@
 # Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
 
+__all__ = [
+    'SpanAttributes',
+    'contains_trace_headers',
+    'extract_trace_context',
+    'extract_trace_headers',
+    'get_span_exporter',
+    'global_otlp_tracer',
+    'init_tracer',
+    'insufficient_request_metrics_warning',
+    'is_otel_available',
+    'is_tracing_enabled',
+    'log_tracing_disabled_warning',
+]
+
 import os
 from collections.abc import Mapping
 from typing import Optional
 
 from tensorrt_llm._utils import run_once
 from tensorrt_llm.logger import logger
 
+# Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0rc1/vllm/tracing.py#L11
 TRACE_HEADERS = ["traceparent", "tracestate"]
 
 _global_tracer_ = None
@@ -19,7 +34,7 @@
         OTEL_EXPORTER_OTLP_TRACES_PROTOCOL
     from opentelemetry.sdk.trace import TracerProvider
     from opentelemetry.sdk.trace.export import BatchSpanProcessor
-    from opentelemetry.trace import SpanKind, Tracer, set_tracer_provider
+    from opentelemetry.trace import Tracer, set_tracer_provider
     from opentelemetry.trace.propagation.tracecontext import \
         TraceContextTextMapPropagator
 

tensorrt_llm/serve/openai_server.py

@@ -26,12 +26,9 @@
 from tensorrt_llm.inputs import prompt_inputs
 from tensorrt_llm.inputs.utils import ConversationMessage, apply_chat_template
 from tensorrt_llm.llmapi import DisaggregatedParams as LlmDisaggregatedParams
+from tensorrt_llm.llmapi import tracing
 from tensorrt_llm.llmapi.disagg_utils import MetadataServerConfig, ServerRole
 from tensorrt_llm.llmapi.llm import RequestOutput
-from tensorrt_llm.llmapi.otel_tracing import (contains_trace_headers,
-                                              extract_trace_headers,
-                                              is_tracing_enabled,
-                                              log_tracing_disabled_warning)
 from tensorrt_llm.logger import logger
 from tensorrt_llm.metrics.collector import MetricsCollector
 from tensorrt_llm.serve.chat_utils import (check_multiple_response,
@@ -529,8 +526,8 @@ async def _get_trace_headers(
         self,
         headers: Headers,
     ) -> Optional[Mapping[str, str]]:
-        if is_tracing_enabled():
-            return extract_trace_headers(headers)
-        if contains_trace_headers(headers):
-            log_tracing_disabled_warning()
+        if tracing.is_tracing_enabled():
+            return tracing.extract_trace_headers(headers)
+        if tracing.contains_trace_headers(headers):
+            tracing.log_tracing_disabled_warning()
         return None