feat: otel trace impl

Author: 皓聪

requirements-dev.txt

@@ -31,3 +31,9 @@ ruff==0.9.4
 lm_eval[api]==0.4.8
 docstring_parser
 genai-perf==0.0.13
+typing_extensions
+# observation.tracing
+opentelemetry-sdk
+opentelemetry-api
+opentelemetry-exporter-otlp
+opentelemetry-semantic-conventions-ai

requirements.txt

@@ -59,3 +59,9 @@ ninja
 etcd3
 blake3
 llguidance==0.7.29
+typing_extensions
+# observation.tracing
+opentelemetry-sdk
+opentelemetry-api
+opentelemetry-exporter-otlp
+opentelemetry-semantic-conventions-ai

tensorrt_llm/commands/serve.py

@@ -84,6 +84,7 @@ def get_llm_args(model: str,
                  num_postprocess_workers: int = 0,
                  trust_remote_code: bool = False,
                  reasoning_parser: Optional[str] = None,
+                 otlp_traces_endpoint: Optional[str] = None,
                  **llm_args_extra_dict: Any):
 
     if gpus_per_node is None:
@@ -125,6 +126,7 @@ def get_llm_args(model: str,
         "num_postprocess_workers": num_postprocess_workers,
         "postprocess_tokenizer_dir": tokenizer or model,
         "reasoning_parser": reasoning_parser,
+        "otlp_traces_endpoint": otlp_traces_endpoint,
     }
 
     return llm_args, llm_args_extra_dict
@@ -249,6 +251,13 @@ def launch_server(host: str,
     default=None,
     help="Server role. Specify this value only if running in disaggregated mode."
 )
+@click.option(
+    "--otlp_traces_endpoint",
+    type=str,
+    default=None,
+    help=
+    "Target URL to which OpenTelemetry traces will be sent."
+)
 def serve(model: str, tokenizer: Optional[str], host: str, port: int,
           log_level: str, backend: str, max_beam_width: int,
           max_batch_size: int, max_num_tokens: int, max_seq_len: int,
@@ -258,7 +267,8 @@ def serve(model: str, tokenizer: Optional[str], host: str, port: int,
           num_postprocess_workers: int, trust_remote_code: bool,
           extra_llm_api_options: Optional[str], reasoning_parser: Optional[str],
           metadata_server_config_file: Optional[str],
-          server_role: Optional[str]):
+          server_role: Optional[str],
+          otlp_traces_endpoint: Optional[str]):
     """Running an OpenAI API compatible server
 
     MODEL: model name | HF checkpoint path | TensorRT engine path
@@ -281,7 +291,8 @@ def serve(model: str, tokenizer: Optional[str], host: str, port: int,
         free_gpu_memory_fraction=kv_cache_free_gpu_memory_fraction,
         num_postprocess_workers=num_postprocess_workers,
         trust_remote_code=trust_remote_code,
-        reasoning_parser=reasoning_parser)
+        reasoning_parser=reasoning_parser,
+        otlp_traces_endpoint=otlp_traces_endpoint)
 
     llm_args_extra_dict = {}
     if extra_llm_api_options is not None:

tensorrt_llm/executor/executor.py

@@ -1,4 +1,5 @@
 import atexit
+from collections.abc import Mapping
 import faulthandler
 import multiprocessing
 import platform
@@ -122,6 +123,7 @@ def generate_async(
             mrope_config: Optional[dict] = None,
             kv_cache_retention_config: Optional[KvCacheRetentionConfig] = None,
             disaggregated_params: Optional[DisaggregatedParams] = None,
+            trace_headers: Optional[Mapping[str, str]] = None,
             postproc_params: Optional[PostprocParams] = None,
             return_perf_metrics: Optional[bool] = False) -> GenerationResult:
         """Generate output for the given prompt token ids in the asynchronous mode.
@@ -149,7 +151,8 @@ def generate_async(
                 mrope_config=mrope_config,
                 kv_cache_retention_config=kv_cache_retention_config,
                 disaggregated_params=disaggregated_params,
-                return_perf_metrics=return_perf_metrics))
+                return_perf_metrics=return_perf_metrics,
+                trace_headers=trace_headers))
         return result
 
     def generate(

tensorrt_llm/executor/request.py

@@ -1,3 +1,4 @@
+from collections.abc import Mapping
 import os
 from dataclasses import dataclass
 from typing import List, Optional, Union
@@ -87,6 +88,7 @@ def __init__(
         mrope_config: Optional[dict] = None,
         kv_cache_retention_config: Optional[KvCacheRetentionConfig] = None,
         disaggregated_params: Optional[DisaggregatedParams] = None,
+        trace_headers: Optional[Mapping[str, str]] = None,
         postproc_params: Optional[PostprocParams] = None,
         return_perf_metrics: Optional[bool] = False,
     ):
@@ -114,6 +116,7 @@ def __init__(
         self.id: Optional[int] = None
         self.disaggregated_params = disaggregated_params
         self.return_perf_metrics = return_perf_metrics
+        self.trace_headers = trace_headers
 
     def set_id(self, id):
         assert self.id is None, f"Request ID is already set: {self.id}"

tensorrt_llm/executor/result.py

@@ -1,4 +1,5 @@
 import asyncio
+from collections.abc import Mapping
 import json
 import weakref
 from dataclasses import dataclass, field
@@ -7,6 +8,7 @@
                     Optional, TypeAlias, Union)
 from weakref import WeakMethod
 
+from tensorrt_llm.llmapi.otel_tracing import SpanAttributes, SpanKind, extract_trace_context, global_otlp_tracer
 import torch
 import torch.nn.functional as F
 
@@ -160,6 +162,7 @@ def __init__(self,
         self.decoding_iter = 0
         self._done = False
         self.metrics_dict = {}
+        self.trace_headers = None
 
         if has_event_loop():
             self.aqueue = AsyncQueue()
@@ -284,6 +287,7 @@ def _handle_sequence(self,
                 raise ValueError(
                     f"Unknown finish reason: {finish_reasons[src_idx]}")
             self.record_stats(output, req_perf_metrics_dict)
+            self.do_tracing(output, req_perf_metrics_dict,req_perf_metrics_dict)
 
     @nvtx_range_debug("handle_response",
                       color="red",
@@ -375,6 +379,70 @@ def record_stats(self,
             metrics_stats.update(processed_metrics_stat)
         self.metrics_dict = metrics_stats
 
+    def do_tracing(
+        self,
+        output: CompletionOutput,
+        req_perf_metrics_dict: Optional[dict[str, float]] = None,
+    ):
+        if not global_otlp_tracer():
+            return
+
+        metrics_dict = self.metrics_dict
+        if not metrics_dict:
+            # Insufficient request metrics available; trace generation aborted.
+            return
+
+        trace_context = extract_trace_context(self.trace_headers)
+        sampling_params = self.sampling_params
+        with global_otlp_tracer().start_as_current_span(
+            "llm_request",
+            kind=SpanKind.SERVER,
+            context=trace_context,
+            start_time=int(
+                req_perf_metrics_dict.get(RequestEventTiming.ARRIVAL_TIME, 0)
+            ),
+        ) as span:
+
+            def safe_set_attr(span, attr, value):
+                if value is not None:
+                    span.set_attribute(attr, value)
+
+            e2e_time = metrics_dict.get(SupportedMetricNames.E2E, -1)
+            safe_set_attr(
+                span,
+                SpanAttributes.GEN_AI_REQUEST_TEMPERATURE,
+                sampling_params.temperature,
+            )
+            safe_set_attr(
+                span, SpanAttributes.GEN_AI_REQUEST_TOP_P, sampling_params.top_p
+            )
+            safe_set_attr(
+                span,
+                SpanAttributes.GEN_AI_REQUEST_MAX_TOKENS,
+                sampling_params.max_tokens,
+            )
+            safe_set_attr(span, SpanAttributes.GEN_AI_REQUEST_N, sampling_params.n)
+            safe_set_attr(
+                span,
+                SpanAttributes.GEN_AI_USAGE_PROMPT_TOKENS,
+                self.postproc_params.postproc_args.num_prompt_tokens,
+            )
+            safe_set_attr(
+                span, SpanAttributes.GEN_AI_USAGE_COMPLETION_TOKENS, output.length
+            )
+            safe_set_attr(
+                span,
+                SpanAttributes.GEN_AI_LATENCY_TIME_TO_FIRST_TOKEN,
+                metrics_dict.get(SupportedMetricNames.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,
+                SpanAttributes.GEN_AI_LATENCY_TIME_IN_QUEUE,
+                metrics_dict.get(SupportedMetricNames.REQUEST_QUEUE_TIME, -1),
+            )
+
 
 class DetokenizedGenerationResultBase(GenerationResultBase):
     ''' The base class for the generation result with detokenization support. '''
@@ -462,6 +530,7 @@ def __init__(
         self.disaggregated_params = disaggregated_params
         # minimal sampling params needed for logprob calculation
         self._logprob_params = logprob_params
+        self.trace_headers = generation_request.trace_headers
 
         # for aborting the request
         self._executor: Optional[weakref.ReferenceType[

tensorrt_llm/llmapi/llm.py

@@ -1,4 +1,5 @@
 import atexit
+from collections.abc import Mapping
 import json
 import os
 import shutil
@@ -14,6 +15,7 @@
 
 from tensorrt_llm.inputs.data import TextPrompt
 from tensorrt_llm.inputs.registry import DefaultInputProcessor
+from tensorrt_llm.llmapi.otel_tracing import init_tracer
 
 from .._utils import nvtx_range_debug
 from ..bindings import executor as tllm
@@ -207,6 +209,13 @@ def __init__(self,
                 self.mpi_session.shutdown()
             raise
 
+        try:
+            if self.args.otlp_traces_endpoint:
+                init_tracer("trt.llm",self.args.otlp_traces_endpoint)
+                logger.info(f"Initialize otlp tracer success, endpont: {self.args.otlp_traces_endpoint}")
+        except Exception as e:
+            logger.error(f"Failed to initialize otlp tracer: {e}")
+
         exception_handler.register(self, 'shutdown')
         atexit.register(LLM._shutdown_wrapper, weakref.ref(self))
 
@@ -305,6 +314,7 @@ def generate_async(
         streaming: bool = False,
         kv_cache_retention_config: Optional[KvCacheRetentionConfig] = None,
         disaggregated_params: Optional[DisaggregatedParams] = None,
+        trace_headers: Optional[Mapping[str, str]] = None,
         _postproc_params: Optional[PostprocParams] = None,
     ) -> RequestOutput:
         """Generate output for the given prompt in the asynchronous mode.
@@ -413,6 +423,7 @@ def generate_async(
             mrope_config=mrope_config,
             kv_cache_retention_config=kv_cache_retention_config,
             disaggregated_params=disaggregated_params,
+            trace_headers=trace_headers,
             postproc_params=_postproc_params,
             return_perf_metrics=self.args.return_perf_metrics,
         )

tensorrt_llm/llmapi/llm_args.py

@@ -1001,6 +1001,11 @@ class BaseLlmArgs(BaseModel):
         json_schema_extra={"type": "Optional[MpiSession]"},
         exclude=True,
         alias="_mpi_session")
+    
+    otlp_traces_endpoint: Optional[str] = Field(
+        default=None,
+        description="Target URL to which OpenTelemetry traces will be sent.",
+        alias="otlp_traces_endpoint")
 
     backend: Optional[str] = Field(
         default=None,