fix

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

Author: Zhang Haotong

tensorrt_llm/executor/result.py

@@ -316,7 +316,7 @@ def _handle_response(self,
             else:
                 self._outputs[0]._postprocess_result = response.res
             if response.metrics:
-                self.metrics_dict = response.metrics
+                self.metrics_dict.update(response.metrics)
 
             if response.error:
                 if self._background_error_handler is not None and (
@@ -391,7 +391,7 @@ def record_stats(self,
             stats, len(output.token_ids), self.sampling_params.n > 1)
         if processed_metrics_stat:
             metrics_stats.update(processed_metrics_stat)
-        self.metrics_dict = metrics_stats
+        self.metrics_dict.update(metrics_stats)
 
     def do_tracing(
         self,
@@ -410,20 +410,29 @@ def do_tracing(
         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)
+        # Since arrival_time and other timing metrics are based on different time origins,
+        # we need to apply corrections to align them with absolute timestamps
+        time_correction = 0
+        arrival_timestamp = metrics_dict.get(MetricNames.ARRIVAL_TIMESTAMP, 0)
+        arrival_time = req_perf_metrics_dict.get(
+            RequestEventTiming.ARRIVAL_TIME, 0)
+        if arrival_timestamp > 0:
+            time_correction = arrival_timestamp - arrival_time
+        else:
+            time_correction = time.time() - metrics_dict.get(
+                MetricNames.E2E, -1) - arrival_time
+
         with tracing.global_otlp_tracer().start_as_current_span(
                 "llm_request",
                 kind=tracing.SpanKind.SERVER,
                 context=trace_context,
-                start_time=int(arrival_time * 1e9),
+                start_time=int((arrival_time + time_correction) * 1e9),
         ) as span:
 
             def safe_set_attr(span, attr, value):
                 if value is not None:
                     span.set_attribute(attr, value)
 
-            e2e_time = metrics_dict.get(MetricNames.E2E, -1)
             safe_set_attr(span,
                           tracing.SpanAttributes.GEN_AI_REQUEST_TEMPERATURE,
                           sampling_params.temperature)
@@ -451,14 +460,36 @@ def safe_set_attr(span, attr, value):
                 span, tracing.SpanAttributes.GEN_AI_LATENCY_TIME_TO_FIRST_TOKEN,
                 metrics_dict.get(MetricNames.TTFT, -1))
             safe_set_attr(span, tracing.SpanAttributes.GEN_AI_LATENCY_E2E,
-                          e2e_time)
+                          metrics_dict.get(MetricNames.E2E, -1))
             safe_set_attr(span,
                           tracing.SpanAttributes.GEN_AI_LATENCY_TIME_IN_QUEUE,
                           metrics_dict.get(MetricNames.REQUEST_QUEUE_TIME, -1))
             safe_set_attr(
                 span, tracing.SpanAttributes.GEN_AI_RESPONSE_FINISH_REASONS,
                 json.dumps([output.finish_reason])
                 if output.finish_reason else None)
+            safe_set_attr(
+                span,
+                tracing.SpanAttributes.GEN_AI_LATENCY_KV_CACHE_TRANSFER_TIME,
+                req_perf_metrics_dict.get(
+                    RequestEventTiming.KV_CACHE_TRANSFER_END, 0.0) -
+                req_perf_metrics_dict.get(
+                    RequestEventTiming.KV_CACHE_TRANSFER_START, 0.0))
+
+            if req_perf_metrics_dict.get(
+                    RequestEventTiming.KV_CACHE_TRANSFER_START,
+                    0) and req_perf_metrics_dict.get(
+                        RequestEventTiming.KV_CACHE_TRANSFER_END, 0):
+                tracing.add_event(
+                    tracing.SpanEvents.KV_CACHE_TRANSFER_START,
+                    timestamp=int((req_perf_metrics_dict.get(
+                        RequestEventTiming.KV_CACHE_TRANSFER_START, 0.0) +
+                                   time_correction) * 1e9))
+                tracing.add_event(
+                    tracing.SpanEvents.KV_CACHE_TRANSFER_END,
+                    timestamp=int((req_perf_metrics_dict.get(
+                        RequestEventTiming.KV_CACHE_TRANSFER_END, 0.0) +
+                                   time_correction) * 1e9))
 
 
 class DetokenizedGenerationResultBase(GenerationResultBase):

tensorrt_llm/executor/worker.py

@@ -1062,7 +1062,15 @@ def _get_metrics_dict(
                 req_perf_metrics.timing_metrics.first_scheduled_time.
                 total_seconds(),
                 RequestEventTiming.LAST_TOKEN_TIME:
-                req_perf_metrics.timing_metrics.last_token_time.total_seconds()
+                req_perf_metrics.timing_metrics.last_token_time.total_seconds(),
+                RequestEventTiming.KV_CACHE_TRANSFER_START:
+                req_perf_metrics.timing_metrics.kv_cache_transfer_start.
+                total_seconds(),
+                RequestEventTiming.KV_CACHE_TRANSFER_END:
+                req_perf_metrics.timing_metrics.kv_cache_transfer_end.
+                total_seconds(),
+                RequestEventTiming.KV_CACHE_SIZE:
+                req_perf_metrics.timing_metrics.kv_cache_size,
             }
     return metrics_dict
 

tensorrt_llm/llmapi/llm.py

@@ -17,6 +17,7 @@
 from tensorrt_llm.inputs.multimodal import MultimodalParams
 from tensorrt_llm.inputs.registry import DefaultInputProcessor
 from tensorrt_llm.llmapi import tracing
+from tensorrt_llm.metrics.enums import MetricNames
 
 from .._utils import nvtx_range_debug
 from ..bindings import executor as tllm
@@ -449,6 +450,10 @@ def generate_async(
             scheduling_params=scheduling_params,
         )
 
+        if sampling_params.return_perf_metrics:
+            result.metrics_dict.update(
+                {MetricNames.ARRIVAL_TIMESTAMP: time.time()})
+
         return RequestOutput._from_generation_result(result, prompt,
                                                      self.tokenizer)
 

tensorrt_llm/llmapi/tracing.py

@@ -2,10 +2,10 @@
 
 __all__ = [
     'SpanAttributes', 'SpanKind', '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',
-    'set_global_otlp_tracer'
+    'extract_trace_context', 'get_span_exporter', 'global_otlp_tracer',
+    'init_tracer', 'insufficient_request_metrics_warning', 'is_otel_available',
+    'is_tracing_enabled', 'log_tracing_disabled_warning',
+    'set_global_otlp_tracer', 'extract_trace_headers'
 ]
 
 import functools
@@ -98,16 +98,23 @@ def extract_trace_context(
         return None
 
 
-def extract_trace_headers(headers: Mapping[str, str]) -> Mapping[str, str]:
-    # Return only recognized trace headers with normalized lowercase keys
-    lower_map = {k.lower(): v for k, v in headers.items()}
-    return {h: lower_map[h] for h in TRACE_HEADERS if h in lower_map}
+def extract_trace_headers(
+        headers: Mapping[str, str]) -> Optional[Mapping[str, str]]:
+    if is_tracing_enabled():
+        # Return only recognized trace headers with normalized lowercase keys
+        lower_map = {k.lower(): v for k, v in headers.items()}
+        return {h: lower_map[h] for h in TRACE_HEADERS if h in lower_map}
+    if contains_trace_headers(headers):
+        log_tracing_disabled_warning()
+    return None
 
 
 def inject_trace_headers(headers: Mapping[str, str]) -> Mapping[str, str]:
-    trace_headers = extract_trace_headers(headers) if not headers else {}
-    TraceContextTextMapPropagator().inject(trace_headers)
-    return trace_headers
+    if is_tracing_enabled():
+        trace_headers = extract_trace_headers(headers) if not headers else {}
+        TraceContextTextMapPropagator().inject(trace_headers)
+        return trace_headers
+    return None
 
 
 def global_otlp_tracer() -> Tracer:
@@ -138,9 +145,17 @@ class SpanAttributes:
     GEN_AI_LATENCY_TIME_TO_FIRST_TOKEN = "gen_ai.latency.time_to_first_token"
     GEN_AI_LATENCY_E2E = "gen_ai.latency.e2e"
     GEN_AI_LATENCY_TIME_IN_QUEUE = "gen_ai.latency.time_in_queue"
+    GEN_AI_LATENCY_KV_CACHE_TRANSFER_TIME = "gen_ai.latency.kv_cache_transfer_time"
     GEN_AI_RESPONSE_FINISH_REASONS = "gen_ai.response.finish_reasons"
 
 
+class SpanEvents:
+    KV_CACHE_TRANSFER_START = "kv_cache_transfer_start"
+    KV_CACHE_TRANSFER_END = "kv_cache_transfer_end"
+    CTX_SERVER_SELECTED = "ctx_server.selected"
+    GEN_SERVER_SELECTED = "gen_server.selected"
+
+
 def contains_trace_headers(headers: Mapping[str, str]) -> bool:
     lower_keys = {k.lower() for k in headers.keys()}
     return any(h in lower_keys for h in TRACE_HEADERS)

tensorrt_llm/metrics/enums.py

@@ -6,10 +6,14 @@ class MetricNames(Enum):
     TPOT = "tpot"
     E2E = "e2e"
     REQUEST_QUEUE_TIME = "request_queue_time"
+    ARRIVAL_TIMESTAMP = 'arrival_timestamp'
 
 
 class RequestEventTiming(Enum):
     ARRIVAL_TIME = "arrival_time"
     FIRST_TOKEN_TIME = "first_token_time"  # nosec: B105
     FIRST_SCHEDULED_TIME = "first_scheduled_time"
     LAST_TOKEN_TIME = "last_token_time"  # nosec: B105
+    KV_CACHE_TRANSFER_START = "kv_cache_transfer_start"
+    KV_CACHE_TRANSFER_END = "kv_cache_transfer_end"
+    KV_CACHE_SIZE = "kv_cache_size"

tensorrt_llm/serve/openai_disagg_server.py

@@ -257,6 +257,9 @@ async def _send_disagg_request(self, req: Union[CompletionRequest, ChatCompletio
             if need_ctx:
                 ctx_req = copy.deepcopy(req)
                 ctx_server, _ = await self.ctx_router.get_next_server(ctx_req)
+                #todo: rename event to something more descriptive
+                tracing.add_event(tracing.SpanEvents.CTX_SERVER_SELECTED, attributes={"server": str(ctx_server),})
+
                 # TODO: add ctx_server info into generation request for pre-registration
                 ctx_response = await self._send_context_request(ctx_server, ctx_req, trace_headers)
 
@@ -277,13 +280,11 @@ async def _send_disagg_request(self, req: Union[CompletionRequest, ChatCompletio
             else:
                 ctx_response = None
 
-            #todo: rename event to something more descriptive
-            tracing.add_event('picking generation server')
-
             # Pick a generation server if haven't reserved one, and send request
             if gen_server is None:
                 gen_server, _ = await self.gen_router.get_next_server(req)
             logger.debug("Sending request to gen server: %s", gen_server)
+            tracing.add_event(tracing.SpanEvents.GEN_SERVER_SELECTED,attributes={"server": str(gen_server),})
 
             if not req.stream:
                 try:

tensorrt_llm/serve/openai_server.py

@@ -4,7 +4,6 @@
 import re
 import signal
 import traceback
-from collections.abc import Mapping
 from contextlib import asynccontextmanager
 from datetime import datetime
 from http import HTTPStatus
@@ -15,7 +14,6 @@
 from fastapi import FastAPI, Request
 from fastapi.exceptions import RequestValidationError
 from fastapi.responses import JSONResponse, Response, StreamingResponse
-from starlette.datastructures import Headers
 from starlette.routing import Mount
 from transformers import AutoConfig, AutoProcessor
 
@@ -339,7 +337,7 @@ async def create_chat_response(
                 postproc_args=postproc_args,
             )
 
-            trace_headers = (None if raw_request is None else await self._get_trace_headers(raw_request.headers))
+            trace_headers = (None if raw_request is None else tracing.extract_trace_headers(raw_request.headers))
 
             promise = self.llm.generate_async(
                 inputs=prompt,
@@ -476,7 +474,7 @@ async def generator_wrapper(generator: AsyncIterator[Any]):
                     if request.stream else completion_response_post_processor,
                     postproc_args=postproc_args,
                 )
-                trace_headers = (None if raw_request is None else await self._get_trace_headers(raw_request.headers))
+                trace_headers = (None if raw_request is None else tracing.extract_trace_headers(raw_request.headers))
                 promise = self.llm.generate_async(
                     inputs=prompt,
                     sampling_params=sampling_params,
@@ -521,13 +519,3 @@ async def __call__(self, host, port):
                                 log_level="info",
                                 timeout_keep_alive=TIMEOUT_KEEP_ALIVE)
         await uvicorn.Server(config).serve()
-
-    async def _get_trace_headers(
-        self,
-        headers: Headers,
-    ) -> Optional[Mapping[str, str]]:
-        if tracing.is_tracing_enabled():
-            return tracing.extract_trace_headers(headers)
-        if tracing.contains_trace_headers(headers):
-            tracing.log_tracing_disabled_warning()
-        return None