feat: otel trace impl

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

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

@@ -63,6 +63,11 @@ meson
 ninja
 etcd3
 blake3
+# observation.tracing
+opentelemetry-sdk
+opentelemetry-api
+opentelemetry-exporter-otlp
+opentelemetry-semantic-conventions-ai
 soundfile
 triton==3.3.1; platform_machine == "x86_64"
 tiktoken

tensorrt_llm/_utils.py

@@ -28,7 +28,8 @@
 from enum import EnumMeta
 from functools import lru_cache, partial, wraps
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Sequence, Union
+from typing import Any, Callable, Dict, List, Optional, Sequence, Union
+from typing_extensions import ParamSpec
 
 import numpy as np
 import nvtx
@@ -1127,3 +1128,16 @@ def set_prometheus_multiproc_dir() -> object:
         os.environ["PROMETHEUS_MULTIPROC_DIR"] = prometheus_multiproc_dir.name
     logger.info(
         f"PROMETHEUS_MULTIPROC_DIR: {os.environ['PROMETHEUS_MULTIPROC_DIR']}")
+
+
+P = ParamSpec("P")
+
+# From: https://stackoverflow.com/a/4104188/2749989
+def run_once(f: Callable[P, None]) -> Callable[P, None]:
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> None:
+        if not wrapper.has_run:  # type: ignore[attr-defined]
+            wrapper.has_run = True  # type: ignore[attr-defined]
+            return f(*args, **kwargs)
+
+    wrapper.has_run = False  # type: ignore[attr-defined]
+    return wrapper

tensorrt_llm/commands/serve.py

@@ -86,6 +86,7 @@ def get_llm_args(model: str,
                  trust_remote_code: bool = False,
                  reasoning_parser: Optional[str] = None,
                  fail_fast_on_attention_window_too_large: bool = False,
+                 otlp_traces_endpoint: Optional[str] = None,
                  **llm_args_extra_dict: Any):
 
     if gpus_per_node is None:
@@ -148,6 +149,7 @@ def get_llm_args(model: str,
         reasoning_parser,
         "fail_fast_on_attention_window_too_large":
         fail_fast_on_attention_window_too_large,
+        "otlp_traces_endpoint": otlp_traces_endpoint,
     }
 
     return llm_args, llm_args_extra_dict
@@ -285,6 +287,13 @@ def launch_server(host: str,
     help=
     "Exit with runtime error when attention window is too large to fit even a single sequence in the KV cache."
 )
+@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,
@@ -295,7 +304,8 @@ def serve(model: str, tokenizer: Optional[str], host: str, port: int,
           extra_llm_api_options: Optional[str], reasoning_parser: Optional[str],
           metadata_server_config_file: Optional[str],
           server_role: Optional[str],
-          fail_fast_on_attention_window_too_large: bool):
+          fail_fast_on_attention_window_too_large: bool,
+          otlp_traces_endpoint: Optional[str]):
     """Running an OpenAI API compatible server
 
     MODEL: model name | HF checkpoint path | TensorRT engine path
@@ -321,7 +331,8 @@ def serve(model: str, tokenizer: Optional[str], host: str, port: int,
         trust_remote_code=trust_remote_code,
         reasoning_parser=reasoning_parser,
         fail_fast_on_attention_window_too_large=
-        fail_fast_on_attention_window_too_large)
+        fail_fast_on_attention_window_too_large,
+        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
@@ -119,6 +120,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,
         multimodal_params: Optional[MultimodalParams] = None,
         scheduling_params: Optional[SchedulingParams] = None,
@@ -144,6 +146,7 @@ def generate_async(
             streaming=streaming,
             kv_cache_retention_config=kv_cache_retention_config,
             disaggregated_params=disaggregated_params,
+            trace_headers=trace_headers,
             multimodal_params=multimodal_params,
             scheduling_params=scheduling_params)
         result = self.submit(request)

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
@@ -94,6 +95,7 @@ def __init__(
         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,
         multimodal_params: Optional[MultimodalParams] = None,
         scheduling_params: Optional[SchedulingParams] = None,
@@ -121,6 +123,7 @@ def __init__(
         self.kv_cache_retention_config = kv_cache_retention_config
         self.id: Optional[int] = None
         self.disaggregated_params = disaggregated_params
+        self.trace_headers = trace_headers
         self.scheduling_params = scheduling_params
 
     def set_id(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()
@@ -288,6 +291,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",
@@ -388,6 +392,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. '''
@@ -498,6 +566,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
@@ -15,6 +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 .._utils import nvtx_range_debug
 from ..bindings import executor as tllm
@@ -209,6 +211,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))
 
@@ -313,6 +322,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,
         scheduling_params: Optional[SchedulingParams] = None,
     ) -> RequestOutput:
@@ -431,6 +441,7 @@ def generate_async(
             streaming=streaming,
             kv_cache_retention_config=kv_cache_retention_config,
             disaggregated_params=disaggregated_params,
+            trace_headers=trace_headers,
             postproc_params=_postproc_params,
             multimodal_params=multimodal_params,
             scheduling_params=scheduling_params,

tensorrt_llm/llmapi/llm_args.py

@@ -1310,6 +1310,11 @@ class BaseLlmArgs(StrictBaseModel):
         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,