diff --git a/docs/source/features/multi-modality.md b/docs/source/features/multi-modality.md index dc272451c7a8..e4ad9c025808 100644 --- a/docs/source/features/multi-modality.md +++ b/docs/source/features/multi-modality.md @@ -74,6 +74,52 @@ You can then send OpenAI-compatible requests, such as via curl or API clients, t `trtllm-serve` renders the chat template itself and marks the resulting `TextPrompt` with `chat_template_applied`, as do the multimodal prompt builders in `tensorrt_llm.inputs.utils`. An input processor that would otherwise apply the chat template to a raw `prompt` should use such a prompt verbatim; see `TextPrompt` in `tensorrt_llm/inputs/data.py`. +### Input processing and encoder metrics + +Enable the existing metrics options in the serving YAML: + +```yaml +return_perf_metrics: true +enable_iter_perf_stats: true +``` + +`/prometheus/metrics` exposes these histograms; `/metrics` serves iteration statistics as JSON. + +| Metric | Measurement | +| --- | --- | +| `trtllm_input_processing_seconds{stage="total"}` | Wall time from entry into a chat, completion, encoder, Harmony or Responses handler through synchronous input preparation and request submission, before waiting for generation. Includes message parsing, rendering, media resolution, preprocessing and their waits. | +| `trtllm_input_processing_seconds{stage="media_load"}` | Time awaiting the common media resolver, including concurrent downloads, decodes and worker waits. | +| `trtllm_input_processing_seconds{stage="chat_template"}` | Time awaiting chat template rendering, when rendering is needed. | +| `trtllm_input_processing_seconds{stage="processor"}` | Input processor wall time inside its worker, including tokenization, model preprocessing and hashing. | +| `trtllm_input_processing_queue_seconds` | Wait for an input processor worker, separate from processor execution. | +| `trtllm_media_processing_seconds{stage="download", modality="image"}` | Per-item HTTP(S) download wall time, excluding shared HTTP session initialization. Also supports `audio` and `video`. Local and inline data inputs do not create download observations. | +| `trtllm_media_processing_seconds{stage="decode", modality="image"}` | Per-item decode execution in the common media worker, separate from queue time. Local file decoding includes file IO; inline decoding includes base64 conversion. | +| `trtllm_media_processing_seconds{stage="decode_queue", modality="image"}` | Per-item wait for a media decode worker. | +| `trtllm_multimodal_encoder_gpu_time_seconds` | Rank-0 CUDA-event elapsed time per encoder call through `MultimodalModelMixin`, including item scheduling, inline encoding and side-stream prefetch. | + +Input timings are captured in a request context independently of Prometheus. When a route emits a completed executor metrics record, it also carries them in `perf_metrics.input_processing_metrics` in the existing JSONL output (`perf_metrics_output_dir`). Values are seconds: request-level stages use keys such as `total`, `media_load`, `chat_template`, `processor` and `processor_queue`. Per-item times are summed by stage and modality within each request, for example `media_download_image` and `media_decode_image`; these sums can exceed elapsed request time when items overlap. Batched completion inputs share the HTTP request's frontend timings. Anthropic Message Batch items have separate request contexts and JSONL records. + +`/health_generate` probes do not contribute to the input-processing histograms; their request-local timings and existing executor metrics remain available. For Anthropic Messages, `total` starts in the delegated chat handler and excludes Anthropic request conversion and validation. Shared HTTP session initialization remains included in `media_load` and `total`. + +Encoder-only servers expose the same Prometheus endpoint and completed-request input timings when metrics are enabled. Applications that configure a `ProcessPoolExecutor` for media loading retain the direct worker submission path; per-worker `decode` and `decode_queue` timings are available only for thread executors. + +For selected requests, send `X-TRTLLM-return-metrics: 1` to receive the same values in `Server-Timing`, for example `server_input_total;dur=150.000000` (milliseconds). Streaming responses use the existing `trtllm.perf_metrics` SSE event. These fields survive context/generation worker handoff and remain separate from executor timing and batch-level encoder GPU measurements. Stage histograms include failed attempts; request JSONL records follow the existing completed-request reporting path. + +Processor metrics cover the chat and text completion paths that offload preprocessing to the input processor pool. Encoder metrics exclude warmup, cache-only lookups and H2D transfers performed before the encoder call. Calls using precomputed embeddings or legacy paths outside `MultimodalModelMixin` do not produce encoder observations. The count is encoder calls, not requests or media items. The timings describe rank 0; they are not a sum or maximum across GPUs. Completed encoder observations may appear in a later iteration's JSON field, `multimodalEncoderStats.gpuTimesMS`, because collection does not wait for the GPU. + +Media loading and chat rendering can overlap. Encoder work can overlap LLM execution. Do not add these histogram averages to reconstruct TTFT or divide a shared encoder call across requests. Use client TTFT and throughput alongside the stage distributions to identify bottlenecks. + +[AIPerf server metrics](https://docs.nvidia.com/aiperf/server-metrics/server-metrics-collection) can scrape these metrics without a TRT-LLM-specific adapter. Add the explicit endpoint to your profile command: + +```bash +aiperf profile --model MODEL --url http://localhost:8000 \ + --endpoint-type chat --endpoint /v1/chat/completions \ + --server-metrics http://localhost:8000/prometheus/metrics \ + --server-metrics-formats json csv jsonl +``` + +The histogram sum and count provide total measured time and mean; percentile values are bucket estimates. Grafana can query the same Prometheus series. + ### Run with [`trtllm-bench`](../../source/commands/trtllm-bench.rst) Evaluate offline inference performance with multimodal inputs using the `trtllm-bench` tool. For detailed instructions, see the [benchmarking guide](../developer-guide/perf-benchmarking.md). diff --git a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py index 242c6cedda90..661f677daef2 100644 --- a/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py +++ b/tensorrt_llm/_torch/models/modeling_multimodal_mixin.py @@ -433,6 +433,7 @@ class MultimodalModelMixin: model_config: ModelConfig _multimodal_encoder_cache: Optional[TensorLRUCache] = None + _encoder_forward_context: Callable[[], contextlib.AbstractContextManager[None]] | None = None @classmethod def _cast_multimodal_encoder_dtype( @@ -482,6 +483,16 @@ def validate_multimodal_request_data(self, mm_data: dict) -> None: """ return None + def _run_multimodal_encoder( + self, + multimodal_params: Sequence[MultimodalParams], + ) -> torch.Tensor: + """Run encoder work with optional runtime instrumentation after H2D.""" + if self._encoder_forward_context is None: + return self.encode_multimodal_inputs(multimodal_params) + with self._encoder_forward_context(): + return self.encode_multimodal_inputs(multimodal_params) + def prepare_multimodal_encoder_inputs( self, selected_items: Sequence[tuple[MultimodalParams, int]], @@ -599,7 +610,7 @@ def forward_multimodal_encoder_items( def flush_group() -> None: if not group_params: return - embeddings = self.encode_multimodal_inputs(group_params) + embeddings = self._run_multimodal_encoder(group_params) expected_length = sum(group_lengths) if embeddings.shape[0] != expected_length: raise MultimodalEncoderContractError( @@ -1133,7 +1144,7 @@ def _get_or_encode_multimodal_embeddings( self._encode_with_partial_cache(partial_hits, encoder_cache) embeddings = get_multimodal_embeddings( - encoder_forward_fn=self.encode_multimodal_inputs, + encoder_forward_fn=self._run_multimodal_encoder, multimodal_params=list(multimodal_params), ) if encoder_cache is not None: @@ -1435,7 +1446,7 @@ def _encode_with_partial_cache( ) # Keep one call across requests and modalities so EncoderGroup can # batch modalities that share an encoder. - batched_output = self.encode_multimodal_inputs([param for param, _, _ in encoder_inputs]) + batched_output = self._run_multimodal_encoder([param for param, _, _ in encoder_inputs]) output_lengths = [length for _, lengths, _ in encoder_inputs for length in lengths] encoded_items = iter(torch.split(batched_output, output_lengths, dim=0)) @@ -1734,7 +1745,7 @@ def _dispatch_cross_iter_prefetch( if partial_hits and encoder_cache is not None: model._encode_with_partial_cache(partial_hits, encoder_cache) if cache_misses: - encoder_output = model.encode_multimodal_inputs(cache_misses).clone() + encoder_output = model._run_multimodal_encoder(cache_misses).clone() _store_chunked_prefill_embeddings(cache_misses, [encoder_output]) if encoder_cache is not None: for param in cache_misses: diff --git a/tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py b/tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py index d129095b0045..16eb177b20e0 100644 --- a/tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py +++ b/tensorrt_llm/_torch/pyexecutor/perf_metrics_manager.py @@ -8,6 +8,7 @@ readability and separation of concerns. """ +from collections.abc import Iterator from contextlib import contextmanager from types import SimpleNamespace from typing import Optional @@ -33,6 +34,8 @@ def __init__(self, enabled: bool): self._perf_events = None self._perf_event_idx = 0 self._forward_event_pool = [] + self._encoder_events: list[tuple[torch.cuda.Event, torch.cuda.Event]] = [] + self._encoder_gpu_times_ms: list[float] = [] # ------------------------------------------------------------------ # GPU event helpers @@ -74,6 +77,44 @@ def release_forward_timing_events(self, start_event, end_event) -> None: if start_event is not None and end_event is not None: self._forward_event_pool.append((start_event, end_event)) + @contextmanager + def record_encoder_forward(self) -> Iterator[None]: + """Time one encoder call on the current stream, including graph replay.""" + self._collect_encoder_gpu_times() + start_event, end_event = self.borrow_forward_timing_events() + start_event.record() + try: + yield + finally: + end_event.record() + self._encoder_events.append((start_event, end_event)) + + def take_encoder_gpu_times_ms(self) -> list[float]: + """Drain completed encoder calls without waiting for a side stream. + + Pending events stay owned here until they complete; they must not be + reused by a later encoder or LLM forward before their timing is read. + """ + self._collect_encoder_gpu_times() + completed, self._encoder_gpu_times_ms = self._encoder_gpu_times_ms, [] + return completed + + def _collect_encoder_gpu_times(self) -> None: + # Reclaim completed events even when iteration stats are downsampled. + pending = [] + for start_event, end_event in self._encoder_events: + try: + if not end_event.query(): + pending.append((start_event, end_event)) + continue + elapsed_ms = start_event.elapsed_time(end_event) + except RuntimeError as error: + logger.warning("Failed to compute encoder GPU elapsed time: %s", error) + else: + self._encoder_gpu_times_ms.append(float(elapsed_ms)) + self.release_forward_timing_events(start_event, end_event) + self._encoder_events = pending + @contextmanager def record_perf_events( self, start_event: Optional[torch.cuda.Event], end_event: Optional[torch.cuda.Event] diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py index cd5703847df2..260d0db613a8 100644 --- a/tensorrt_llm/_torch/pyexecutor/py_executor.py +++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py @@ -73,8 +73,8 @@ from ..disaggregation.orchestration.transfer_manager import AsyncTransferManager from ..distributed import Distributed from ..distributed.communicator import ReduceOp -from ..models.modeling_multimodal_mixin import \ - maybe_prefetch_mm_encoder_for_next_iter +from ..models.modeling_multimodal_mixin import ( + MultimodalModelMixin, maybe_prefetch_mm_encoder_for_next_iter) from ..models.modeling_utils import DecoderModelForCausalLM from ..modules.decoder_layer import DecoderLayer from ..modules.dwdp.dense.manager import DenseDwdpManager @@ -1372,6 +1372,9 @@ def _initialize_execution(self, resources: PyExecutorResources) -> None: from tensorrt_llm._torch.late_compile import mark_server_ready mark_server_ready() + # Install only after warmup/capture so dummy encoder work is excluded. + self._setup_encoder_perf_metrics() + # Snapshot some cumulative KV cache counters so that stats reported to # users exclude blocks reused and missed during warmup dummy requests. if hasattr(self.kv_cache_manager, 'snapshot_warmup_baseline'): @@ -1642,6 +1645,15 @@ def rebind_resources(self, resources: PyExecutorResources) -> None: self.shutdown() raise + def _setup_encoder_perf_metrics(self) -> None: + """Install optional encoder timing on the module that owns encoder methods.""" + # torch.compile delegates encoder methods with self bound to _orig_mod. + model = getattr(self.model_engine.model, "_orig_mod", + self.model_engine.model) + if (self.enable_iter_perf_stats and self.dist.rank == 0 + and isinstance(model, MultimodalModelMixin)): + model._encoder_forward_context = self.perf_manager.record_encoder_forward + @property def is_generation_model(self) -> bool: """Whether requests generate tokens, rather than finish after one forward.""" @@ -3580,6 +3592,16 @@ def _append_iter_stats(self, or self.dist.pp_size > 1) scheduler_mode = "overlap" if is_overlap_like else "non_overlap" + # These are completed rank-0 encoder calls, including side-stream + # prefetch. They can finish after the iteration that launched them. + # Attach them once, rather than copying them to every ADP fanout row. + encoder_gpu_times_ms = (self.perf_manager.take_encoder_gpu_times_ms() + if attention_dp_rank in (None, 0) else []) + + mm_encoder_stats = { + "gpuTimesMS": encoder_gpu_times_ms + } if encoder_gpu_times_ms else None + tp_size = getattr(self.dist, "tp_size", 1) gather_all_ranks = os.environ.get("TLLM_METRICS_ALL_RANKS", "0") == "1" if (gather_all_ranks and self.enable_iter_perf_stats and tp_size > 1 @@ -3597,6 +3619,8 @@ def _append_iter_stats(self, local_dict["prevDeviceStepTimeMS"] = prev_device_step_time_ms if gpu_forward_time_ms is not None: local_dict["gpuForwardTimeMS"] = gpu_forward_time_ms + if mm_encoder_stats is not None: + local_dict["multimodalEncoderStats"] = mm_encoder_stats local_dict["schedulerMode"] = scheduler_mode local_dict["rank"] = self.dist.tp_rank # Buffer for _flush_iter_stats_synced; in-place tp_allgather would @@ -3615,6 +3639,7 @@ def _append_iter_stats(self, # [5] prev_device_step_time_ms: Optional[float] # [6] scheduler_mode: "overlap" | "non_overlap" # [7] gpu_forward_time_ms: Optional[float] + # [8] mm_encoder_stats: optional completed rank-0 encoder observations with self.stats_lock: if (not _stats_buffer_is_unbounded(self.max_stats_len) and len(self.stats) > self.max_stats_len): @@ -3622,7 +3647,7 @@ def _append_iter_stats(self, self.stats.append( (stats, req_stats, kv_iter_stats, attention_dp_rank, host_step_time_ms, prev_device_step_time_ms, scheduler_mode, - gpu_forward_time_ms)) + gpu_forward_time_ms, mm_encoder_stats)) def _process_iter_stats( self, diff --git a/tensorrt_llm/executor/base_worker.py b/tensorrt_llm/executor/base_worker.py index 2fcf106677c3..4f1c75079e5c 100644 --- a/tensorrt_llm/executor/base_worker.py +++ b/tensorrt_llm/executor/base_worker.py @@ -1301,6 +1301,7 @@ def _stats_serializer(stats) -> str: prev_device_step_time_ms = stats[5] if len(stats) > 5 else None scheduler_mode = stats[6] if len(stats) > 6 else None gpu_forward_time_ms = stats[7] if len(stats) > 7 else None + mm_encoder_stats = stats[8] if len(stats) > 8 else None stats_dict = json.loads(iteration_stats.to_json_str()) # Always tag the row so Dynamo's adapter can read @@ -1337,6 +1338,8 @@ def _stats_serializer(stats) -> str: # ForwardPassMetrics wall_time. if gpu_forward_time_ms is not None: stats_dict["gpuForwardTimeMS"] = gpu_forward_time_ms + if mm_encoder_stats is not None: + stats_dict["multimodalEncoderStats"] = mm_encoder_stats # Scheduler mode for this record. "overlap" means iterLatencyMS # spans ~2 loops (use hostStepTimeMS for clean per-loop cost); # "non_overlap" means iterLatencyMS is itself the clean per-loop diff --git a/tensorrt_llm/inputs/media_io.py b/tensorrt_llm/inputs/media_io.py index 9a94952ffd0d..49cf4bc9fc89 100644 --- a/tensorrt_llm/inputs/media_io.py +++ b/tensorrt_llm/inputs/media_io.py @@ -14,9 +14,9 @@ import threading from abc import ABC, abstractmethod from collections.abc import Callable, Iterator -from concurrent.futures import Executor +from concurrent.futures import Executor, ProcessPoolExecutor from contextlib import contextmanager -from contextvars import ContextVar +from contextvars import ContextVar, copy_context from io import BytesIO from pathlib import Path from stat import S_ISREG @@ -48,6 +48,11 @@ from tensorrt_llm.inputs.multimodal_data import AudioData, MediaModality, VideoData from tensorrt_llm.logger import logger +from tensorrt_llm.metrics.perf_utils import ( + _record_input_processing, + _start_input_processing, + _time_input_processing, +) class MediaFetchError(RuntimeError): @@ -348,11 +353,14 @@ async def _read_aiohttp_content(response: aiohttp.ClientResponse) -> bytes: content.write(chunk) -async def _safe_aiohttp_get(url: str, timeout_sec: int = 30) -> bytes: +async def _safe_aiohttp_get( + url: str, timeout_sec: int = 30, *, session: Optional[aiohttp.ClientSession] = None +) -> bytes: """Aiohttp GET wrapper that validates every redirect hop before following.""" await asyncio.to_thread(_validate_url, url) timeout = aiohttp.ClientTimeout(total=timeout_sec) - session = await _get_aiohttp_session() + if session is None: + session = await _get_aiohttp_session() current_url = url for _ in range(_MAX_REDIRECTS + 1): try: @@ -1294,6 +1302,7 @@ class BaseMediaIO(ABC, Generic[_MediaT]): _work_tracker: ClassVar[ContextVar[Optional[Callable[[asyncio.Future], None]]]] = ContextVar( "media_io_work_tracker", default=None ) + _modality: ClassVar[str] = "other" # Required MIME type prefix for `data:` URIs handled by this modality # (e.g. `"image/"`). A data URI whose declared media type falls outside @@ -1381,7 +1390,9 @@ async def async_load(self, url: str) -> _MediaT: """ parsed = urlparse(url) if parsed.scheme in ("http", "https"): - data = await _safe_aiohttp_get(url) + session = await _get_aiohttp_session() + with _time_input_processing("download", self._modality): + data = await _safe_aiohttp_get(url, session=session) return await self._run_in_executor(self.load_bytes, data) elif parsed.scheme == "data": data_spec, b64_data = parsed.path.split(",", 1) @@ -1432,14 +1443,28 @@ def _load_allowed_local_file(self, url: str) -> _MediaT: finally: os.close(fd) - @staticmethod - async def _run_in_executor(fn, *args, **kwargs): + @classmethod + async def _run_in_executor( + cls, fn: Callable[..., _MediaT], *args: Any, **kwargs: Any + ) -> _MediaT: """Run a blocking decode callable on the configured executor.""" if kwargs: fn = functools.partial(fn, **kwargs) loop = asyncio.get_event_loop() executor = BaseMediaIO._request_executor.get() or BaseMediaIO._executor - future = loop.run_in_executor(executor, fn, *args) + queued = _start_input_processing() + if queued is None or isinstance(executor, ProcessPoolExecutor): + future = loop.run_in_executor(executor, fn, *args) + else: + modality = cls._modality + + def decode() -> _MediaT: + _record_input_processing("decode_queue", queued, modality) + with _time_input_processing("decode", modality): + return fn(*args) + + # run_in_executor does not propagate request ContextVars to workers. + future = loop.run_in_executor(executor, copy_context().run, decode) tracker = BaseMediaIO._work_tracker.get() if tracker is not None: tracker(future) @@ -1452,6 +1477,7 @@ class ImageMediaIO(BaseMediaIO[Union[Image.Image, torch.Tensor, np.ndarray]]): """I/O for the image modality.""" _data_uri_mime_prefix: ClassVar[Optional[str]] = "image/" + _modality = "image" def __init__(self, format: str = "pt", device: str = "cpu") -> None: if format not in _SUPPORTED_IMAGE_FORMATS: @@ -1504,6 +1530,7 @@ class AudioMediaIO(BaseMediaIO[Tuple[np.ndarray, int]]): """I/O for the audio modality.""" _data_uri_mime_prefix: ClassVar[Optional[str]] = "audio/" + _modality = "audio" def __init__( self, decoded_memory_reserver: Optional[_DecodedMemoryReserver] = None, **kwargs @@ -1564,6 +1591,7 @@ class VideoMediaIO(BaseMediaIO[VideoData]): """Video I/O with a 200 MiB encoded-input limit and coupled sampling kwargs.""" _data_uri_mime_prefix: ClassVar[Optional[str]] = "video/" + _modality = "video" def __init__( self, diff --git a/tensorrt_llm/inputs/utils.py b/tensorrt_llm/inputs/utils.py index ab77c168b863..d21001b53039 100644 --- a/tensorrt_llm/inputs/utils.py +++ b/tensorrt_llm/inputs/utils.py @@ -37,6 +37,7 @@ from tensorrt_llm.inputs.registry import (MULTIMODAL_PLACEHOLDER_REGISTRY, MultimodalPlaceholderPlacement) from tensorrt_llm.llmapi.llm_utils import ModelLoader +from tensorrt_llm.metrics.perf_utils import _time_input_processing from tensorrt_llm.tokenizer import TokenizerBase, TransformersTokenizer from tensorrt_llm.tokenizer.deepseek_v4 import DeepseekV4Tokenizer from tensorrt_llm.tokenizer.deepseek_v32 import DeepseekV32Tokenizer @@ -410,9 +411,12 @@ async def _retrieve( out[modality].append(result) return dict(out) + if not self._data and not self._embeddings: + return None, None # _data and _embeddings also gathered concurrently - data_result, embed_result = await asyncio.gather( - _retrieve(self._data), _retrieve(self._embeddings)) + with _time_input_processing("media_load"): + data_result, embed_result = await asyncio.gather( + _retrieve(self._data), _retrieve(self._embeddings)) return data_result, embed_result async def discard_all_async( @@ -918,21 +922,22 @@ async def async_apply_chat_template( injected_chat_template_kwargs: Optional[Collection[str]] = None, ) -> (str | List[str]): """Apply chat template without blocking the event loop.""" - return await asyncio.to_thread( - apply_chat_template, - model_type=model_type, - tokenizer=tokenizer, - processor=processor, - conversation=conversation, - add_generation_prompt=add_generation_prompt, - mm_placeholder_counts=mm_placeholder_counts, - tools=tools, - documents=documents, - chat_template=chat_template, - chat_template_kwargs=chat_template_kwargs, - enable_tokenize=enable_tokenize, - injected_chat_template_kwargs=injected_chat_template_kwargs, - ) + with _time_input_processing("chat_template"): + return await asyncio.to_thread( + apply_chat_template, + model_type=model_type, + tokenizer=tokenizer, + processor=processor, + conversation=conversation, + add_generation_prompt=add_generation_prompt, + mm_placeholder_counts=mm_placeholder_counts, + tools=tools, + documents=documents, + chat_template=chat_template, + chat_template_kwargs=chat_template_kwargs, + enable_tokenize=enable_tokenize, + injected_chat_template_kwargs=injected_chat_template_kwargs, + ) def apply_mm_placeholders( diff --git a/tensorrt_llm/metrics/collector.py b/tensorrt_llm/metrics/collector.py index d96309c4ebcd..a24943d51778 100644 --- a/tensorrt_llm/metrics/collector.py +++ b/tensorrt_llm/metrics/collector.py @@ -36,6 +36,26 @@ # makes the metrics-side limit its own explicit concern. MAX_SPEC_DECODE_POSITION_LABELS = 16 +_PROCESSING_TIME_BUCKETS = ( + 0.0001, + 0.00025, + 0.0005, + 0.001, + 0.0025, + 0.005, + 0.01, + 0.025, + 0.05, + 0.1, + 0.25, + 0.5, + 1.0, + 2.5, + 5.0, + 10.0, + 30.0, +) + # Adapted from https://github.com/vllm-project/vllm/blob/v0.10.0rc1/vllm/engine/metrics.py#L30 class MetricsCollector: @@ -68,6 +88,9 @@ class MetricsCollector: trtllm_prefill_perplexity trtllm_generation_perplexity trtllm_request_error_total + trtllm_input_processing_seconds + trtllm_input_processing_queue_seconds + trtllm_media_processing_seconds Iteration-level metrics: trtllm_kv_cache_hit_rate @@ -113,6 +136,7 @@ class MetricsCollector: trtllm_spec_decode_draft_overhead trtllm_prefill_batch_occupancy trtllm_prefill_batch_tokens + trtllm_multimodal_encoder_gpu_time_seconds Config info metrics (logged once at startup via log_config_info): trtllm_model_config_info @@ -570,6 +594,38 @@ def __init__( buckets=[64, 128, 256, 512, 1024, 2048, 4096, 8192, 16384, 32768], labelnames=iteration_labelnames) + self.histogram_input_processing = Histogram( + name=self.metric_prefix + "input_processing_seconds", + documentation=( + "Server input processing wall time per stage. Media loading " + "and chat template stages can overlap and include async waits; " + "processor time excludes its thread-pool queue wait. Includes " + "failed attempts. Stages must not be added to obtain TTFT."), + buckets=_PROCESSING_TIME_BUCKETS, + labelnames=[*self.labels, "stage"]) + self.histogram_input_processing_queue = Histogram( + name=self.metric_prefix + "input_processing_queue_seconds", + documentation="Time waiting for a server input-processor thread.", + buckets=_PROCESSING_TIME_BUCKETS, + labelnames=self.labels.keys()) + self.histogram_media_processing = Histogram( + name=self.metric_prefix + "media_processing_seconds", + documentation= + ("Per-media-item download, decode execution and decode " + "worker queue times. Local file decode includes file IO; " + "data-URI decode includes base64 conversion. Includes failed attempts." + ), + buckets=_PROCESSING_TIME_BUCKETS, + labelnames=[*self.labels, "stage", "modality"]) + self.histogram_multimodal_encoder_gpu_time = Histogram( + name=self.metric_prefix + "multimodal_encoder_gpu_time_seconds", + documentation= + ("Rank-0 GPU elapsed time per completed multimodal encoder " + "call through MultimodalModelMixin. Excludes pre-encoder H2D, cache-only " + "lookups and warmup. Batched work is not per-request latency."), + buckets=_PROCESSING_TIME_BUCKETS, + labelnames=self.labels.keys()) + # HTTP error counter self.labelname_http_code = "http_code" self.labels_with_http_code = { @@ -940,6 +996,25 @@ def log_request_metrics_dict(self, metrics_dict: dict) -> None: self.last_log_time = time.time() + def log_input_processing( + self, + stage: str, + duration_seconds: float, + modality: str | None = None, + ) -> None: + """Observe the same input-stage sample stored in the request record.""" + if modality is not None: + self.histogram_media_processing.labels(**self.labels, + stage=stage, + modality=modality).observe( + duration_seconds) + elif stage == "processor_queue": + self._log_histogram(self.histogram_input_processing_queue, + duration_seconds) + else: + self.histogram_input_processing.labels( + **self.labels, stage=stage).observe(duration_seconds) + def log_iteration_stats(self, iteration_stats: dict) -> None: """Log iteration-level statistics from TRTLLM engine. @@ -981,6 +1056,11 @@ def _log_rank0_iteration_stats(self, iteration_stats: dict) -> None: Args: iteration_stats: Serialized iteration measurements from rank 0. """ + for elapsed_ms in (iteration_stats.get("multimodalEncoderStats") + or {}).get("gpuTimesMS", ()): + self._log_histogram(self.histogram_multimodal_encoder_gpu_time, + elapsed_ms / 1000.0) + # Top-level queue & load metrics if "numActiveRequests" in iteration_stats: self._log_gauge(self.num_requests_running, diff --git a/tensorrt_llm/metrics/perf_utils.py b/tensorrt_llm/metrics/perf_utils.py index 3bbf7f7a8130..e76bded1e8dc 100644 --- a/tensorrt_llm/metrics/perf_utils.py +++ b/tensorrt_llm/metrics/perf_utils.py @@ -18,10 +18,41 @@ dependencies and can be unit-tested without a full TensorRT-LLM install. """ +import time +from collections.abc import Callable, Iterator +from contextlib import contextmanager +from contextvars import ContextVar from typing import Optional from .enums import MetricNames, RequestEventTiming +_input_processing_observer: ContextVar[Callable[[str, float, str | None], None] | None] = ( + ContextVar("input_processing_observer", default=None) +) + + +def _start_input_processing() -> float | None: + return time.perf_counter() if _input_processing_observer.get() is not None else None + + +def _record_input_processing( + stage: str, started: float | None, modality: str | None = None +) -> None: + if started is not None: + observer = _input_processing_observer.get() + if observer is not None: + observer(stage, time.perf_counter() - started, modality) + + +@contextmanager +def _time_input_processing(stage: str, modality: str | None = None) -> Iterator[None]: + """Measure an input stage only within an opted-in request context.""" + started = _start_input_processing() + try: + yield + finally: + _record_input_processing(stage, started, modality) + def process_req_perf_metrics( req_perf_metrics_dict: Optional[dict], output_length: int diff --git a/tensorrt_llm/serve/_perf_metrics_schema.py b/tensorrt_llm/serve/_perf_metrics_schema.py index 67370c57ba26..563e4bb98e8d 100644 --- a/tensorrt_llm/serve/_perf_metrics_schema.py +++ b/tensorrt_llm/serve/_perf_metrics_schema.py @@ -69,6 +69,7 @@ class SpeculativeDecodingMetrics(TypedDict): class PerfMetrics(TypedDict): timing_metrics: TimingMetrics + input_processing_metrics: NotRequired[dict[str, float]] first_iter: NotRequired[int] last_iter: NotRequired[int] kv_cache_metrics: NotRequired[KvCacheMetrics] diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 233512d9788b..2b21ae957bcb 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -17,13 +17,15 @@ import traceback import uuid from collections import deque +from collections.abc import Callable from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager +from contextvars import copy_context from datetime import datetime from http import HTTPStatus from pathlib import Path from typing import (TYPE_CHECKING, Annotated, Any, AsyncGenerator, - AsyncIterator, Dict, List, Optional, Tuple, Union) + AsyncIterator, Dict, List, Optional, Tuple, TypeVar, Union) import msgspec import uvicorn @@ -65,6 +67,9 @@ from tensorrt_llm.media.encoding import image_to_bytes from tensorrt_llm.media.tensor_payload import is_tensor_format from tensorrt_llm.metrics.collector import MetricsCollector +from tensorrt_llm.metrics.perf_utils import (_record_input_processing, + _start_input_processing, + _time_input_processing) from tensorrt_llm.runtime.kv_cache_hash import \ get_effective_kv_cache_event_hash_algo from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams @@ -116,6 +121,7 @@ from tensorrt_llm.serve.openai_video_routes import _VideoRoutesMixin from tensorrt_llm.serve.perf_metrics import (PerfMetricsJsonlWriter, PerfMetricsMiddleware, + _input_processing_context, build_request_metrics_record) from tensorrt_llm.serve.postprocess_handlers import ( ChatCompletionPostprocArgs, ChatPostprocArgs, CompletionPostprocArgs, @@ -574,6 +580,9 @@ def resolve_spec_decode_num_spec_tokens(args: Any) -> Optional[int]: return getattr(spec_config, "max_draft_len", None) +_InputProcessingResult = TypeVar("_InputProcessingResult") + + class OpenAIServer(_VideoRoutesMixin): @staticmethod @@ -938,10 +947,14 @@ async def validation_exception_handler(request, exc): self.register_routes() if self._collect_perf_metrics: - self.app.add_middleware(PerfMetricsMiddleware, - expose_headers=self._expose_perf_metrics, - writer=self._perf_metrics_writer, - adjusted_clock=self._adjusted_steady_clock) + self.app.add_middleware( + PerfMetricsMiddleware, + expose_headers=self._expose_perf_metrics, + writer=self._perf_metrics_writer, + adjusted_clock=self._adjusted_steady_clock, + input_processing_observer=( + self.metrics_collector.log_input_processing + if self.metrics_collector else None)) self.app.add_middleware(ServerArrivalTimeMiddleware, adjusted_clock=self._adjusted_steady_clock) @@ -1395,6 +1408,8 @@ def register_mm_encoder_routes(self): self.app.add_api_route("/v1/chat/completions", self.openai_mm_encoder, methods=["POST"]) + if self.generator.args.return_perf_metrics: + self.mount_metrics() async def _require_rl_control_auth(self, raw_request: Request) -> None: body = await raw_request.body() @@ -2065,6 +2080,25 @@ def _record_iteration_stat(self, llm_stat: dict) -> None: if self._iteration_stats_buffer is not None: self._iteration_stats_buffer.append(llm_stat) + async def _run_input_processor( + self, + process: Callable[[], _InputProcessingResult], + ) -> _InputProcessingResult: + """Run preprocessing with separate thread queue and execution times.""" + loop = asyncio.get_running_loop() + queued = _start_input_processing() + if queued is None: + return await loop.run_in_executor(self._input_proc_executor, + process) + + def run() -> _InputProcessingResult: + _record_input_processing("processor_queue", queued) + with _time_input_processing("processor"): + return process() + + return await loop.run_in_executor(self._input_proc_executor, + copy_context().run, run) + async def openai_chat(self, request: ChatCompletionRequest, raw_request: Request) -> Response: @@ -2107,6 +2141,7 @@ async def chat_stream_generator( logger.error(traceback.format_exc()) raise + input_started = _start_input_processing() try: # Extension preprocessing, the required-tool check and tool dumping # are the first stage of the shared prompt-preparation pipeline. @@ -2481,9 +2516,7 @@ async def chat_stream_generator( generate_inputs = prompt preprocess_fn = getattr(self.generator, "preprocess", None) if preprocess_fn is not None: - loop = asyncio.get_event_loop() - generate_inputs = await loop.run_in_executor( - self._input_proc_executor, + generate_inputs = await self._run_input_processor( functools.partial(preprocess_fn, prompt, sampling_params, disaggregated_params)) @@ -2503,6 +2536,9 @@ async def chat_stream_generator( priority=request.priority if request.priority is not None else DEFAULT_REQUEST_PRIORITY, ) + _record_input_processing("total", input_started) + input_started = None + asyncio.create_task(self.await_disconnected(raw_request, promise)) if raw_request is not None: # The executor client id, which the perf metrics, the @@ -2546,6 +2582,8 @@ async def chat_stream_generator( except Exception as e: logger.error(traceback.format_exc()) return self.create_error_response(str(e)) + finally: + _record_input_processing("total", input_started) async def anthropic_messages(self, request: AnthropicMessagesRequest, raw_request: Request) -> Response: @@ -2727,7 +2765,16 @@ async def _run_batched_message( raw_request = self._synthetic_request( chat_request.model_dump(exclude_none=True)) - response = await self.openai_chat(chat_request, raw_request) + if self._collect_perf_metrics: + observer = (self.metrics_collector.log_input_processing + if self.metrics_collector else None) + with _input_processing_context(raw_request.scope["state"], + observer): + response = await self.openai_chat(chat_request, raw_request) + for record in raw_request.state.perf_metrics_records: + self._perf_metrics_writer.submit(record) + else: + response = await self.openai_chat(chat_request, raw_request) if response is None: return "errored", { "type": "api_error", @@ -2893,6 +2940,7 @@ async def create_mm_embedding_response(promise: RequestOutput): ), ) + input_started = _start_input_processing() try: ensure_request_chat_template_allowed( request, self.allow_request_chat_template) @@ -2953,9 +3001,13 @@ async def create_mm_embedding_response(promise: RequestOutput): mm_item_order) promise = self.generator.generate_async(inputs=prompt, ) + _record_input_processing("total", input_started) + input_started = None + asyncio.create_task(self.await_disconnected(raw_request, promise)) response = await create_mm_embedding_response(promise) + await self._extract_metrics(promise, raw_request) return JSONResponse(content=response.model_dump()) except CppExecutorError: @@ -2966,6 +3018,8 @@ async def create_mm_embedding_response(promise: RequestOutput): except Exception as e: logger.error(traceback.format_exc()) return self.create_error_response(str(e)) + finally: + _record_input_processing("total", input_started) async def openai_completion(self, request: CompletionRequest, raw_request: Request) -> Response: @@ -3075,6 +3129,7 @@ async def generator_wrapper(generator: AsyncIterator[Any]): yield output yield "data: [DONE]\n\n" + input_started = _start_input_processing() try: if isinstance(request.prompt, list) and not request.prompt: return self.create_error_response("'prompt' must not be empty.") @@ -3133,9 +3188,7 @@ async def generator_wrapper(generator: AsyncIterator[Any]): prompt = prompt_inputs(prompt) if prompt.get("prompt") is not None: - loop = asyncio.get_event_loop() - prompt_token_ids, extra_processed_inputs = await loop.run_in_executor( - self._input_proc_executor, + prompt_token_ids, extra_processed_inputs = await self._run_input_processor( functools.partial(self.generator.input_processor, prompt, sampling_params)) tokens_prompt = TokensPrompt( @@ -3165,6 +3218,9 @@ async def generator_wrapper(generator: AsyncIterator[Any]): postproc_params_collection.append( None if self.postproc_worker_enabled else postproc_params) + _record_input_processing("total", input_started) + input_started = None + if request.stream: generators = [ completion_generator(promise, params) for promise, params in @@ -3191,6 +3247,8 @@ async def generator_wrapper(generator: AsyncIterator[Any]): except Exception as e: logger.error(traceback.format_exc()) return self.create_error_response(str(e)) + finally: + _record_input_processing("total", input_started) async def chat_harmony(self, request: ChatCompletionRequest, raw_request: Request) -> Response: @@ -3227,6 +3285,7 @@ async def create_streaming_generator(promise: RequestOutput, logger.error(traceback.format_exc()) raise + input_started = _start_input_processing() try: ensure_request_chat_template_allowed( request, self.allow_request_chat_template) @@ -3327,6 +3386,9 @@ async def create_streaming_generator(promise: RequestOutput, priority=request.priority if request.priority is not None else DEFAULT_REQUEST_PRIORITY, ) + _record_input_processing("total", input_started) + input_started = None + if not self.postproc_worker_enabled: postproc_args.num_prompt_tokens = len(promise.prompt_token_ids) @@ -3350,6 +3412,8 @@ async def create_streaming_generator(promise: RequestOutput, logger.debug("Error details: %s", traceback.format_exc()) return self.create_error_response(message=str(e), err_type="internal_error") + finally: + _record_input_processing("total", input_started) def _responses_prompt_tokens_excluded_from_usage( self, request: ResponsesRequest) -> int: @@ -3432,6 +3496,7 @@ async def create_streaming_generator(promise: RequestOutput, yield pp_res await self._extract_metrics(res, raw_request) + input_started = _start_input_processing() try: if request.background: logger.warning( @@ -3573,6 +3638,9 @@ async def create_streaming_generator(promise: RequestOutput, _postproc_params=postproc_params if self.postproc_worker_enabled else None, ) + _record_input_processing("total", input_started) + input_started = None + if not self.postproc_worker_enabled: # Set by the executor only for postprocessing it runs itself. postproc_args.num_prompt_tokens = len(promise.prompt_token_ids) @@ -3616,6 +3684,9 @@ async def create_streaming_generator(promise: RequestOutput, logger.error(traceback.format_exc()) return self.create_error_response(str(e)) + finally: + _record_input_processing("total", input_started) + return JSONResponse(content={"detail": "None"}) async def openai_responses_get_response(self, diff --git a/tensorrt_llm/serve/perf_metrics.py b/tensorrt_llm/serve/perf_metrics.py index 5c6c7b2715b3..f388c194864a 100644 --- a/tensorrt_llm/serve/perf_metrics.py +++ b/tensorrt_llm/serve/perf_metrics.py @@ -52,7 +52,10 @@ import math import os import socket +import threading import traceback +from collections.abc import Callable, Iterator +from contextlib import contextmanager from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path @@ -61,6 +64,7 @@ from tensorrt_llm._utils import AdjustedSteadyClock from tensorrt_llm.llmapi.disagg_utils import ServerRole from tensorrt_llm.logger import logger +from tensorrt_llm.metrics.perf_utils import _input_processing_observer from tensorrt_llm.serve._perf_metrics_schema import ( SERVER_TIMING_HEADER, DisaggPerfMetricsRecord, @@ -278,6 +282,10 @@ def build_request_metrics_record( "timing_metrics": timing_metrics, "kv_cache_metrics": {name: getattr(metrics.kv_cache_metrics, name) for name in _KV_FIELDS}, } + if raw_request is not None: + input_metrics = getattr(raw_request.state, "input_processing_metrics", None) + if input_metrics: + phase_record["input_processing_metrics"] = dict(input_metrics) speculative = metrics.speculative_decoding if speculative.total_draft_tokens > 0: phase_record["speculative_decoding"] = { @@ -359,6 +367,11 @@ def build_metrics_headers(records: List[Dict[str, Any]]) -> Dict[str, str]: if duration is not None: values[SERVER_TIMING_HEADER].append(f"{phase}_{name};dur={duration:.6f}") + for name, seconds in phase_record.get("input_processing_metrics", {}).items(): + values[SERVER_TIMING_HEADER].append( + f"{phase}_input_{name};dur={seconds * 1000:.6f}" + ) + breakdown = phase_record.get("time_breakdown_metrics") or {} for header, key, label in ( (STEP_METRICS_HEADER, "step_metrics", "step"), @@ -445,6 +458,13 @@ def build_metrics_record_from_headers( timing_metrics.setdefault(field, arrival_time + durations[name]) phase_record = {"metrics_headers": metrics_headers} + input_metrics = { + name[len("input_") :]: seconds + for name, seconds in durations.items() + if name.startswith("input_") + } + if input_metrics: + phase_record["input_processing_metrics"] = input_metrics if timing_metrics: phase_record["timing_metrics"] = timing_metrics return { @@ -561,6 +581,8 @@ def _jsonl_perf_metrics(phase_record: Dict[str, Any]) -> PerfMetrics: perf_metrics["kv_cache_metrics"] = phase_record["kv_cache_metrics"] if "speculative_decoding" in phase_record: perf_metrics["speculative_decoding"] = phase_record["speculative_decoding"] + if "input_processing_metrics" in phase_record: + perf_metrics["input_processing_metrics"] = phase_record["input_processing_metrics"] timing_metrics = dict(perf_metrics.get("timing_metrics", {})) if not timing_metrics.get("kv_cache_size"): @@ -788,6 +810,31 @@ def build_metrics_sse_event(headers: Dict[str, str]) -> bytes: return (f"event: {SSE_METRICS_EVENT}\ndata: {payload}\n\n").encode() +@contextmanager +def _input_processing_context( + state: dict[str, Any], + observer: Callable[[str, float, str | None], None] | None = None, +) -> Iterator[None]: + """Keep input timings local to one HTTP request or background batch item.""" + input_metrics: dict[str, float] = {} + state["input_processing_metrics"] = input_metrics + input_lock = threading.Lock() + + def observe_input(stage: str, seconds: float, modality: str | None) -> None: + name = f"media_{stage}_{modality}" if modality is not None else stage + # Concurrent media decode workers share this bounded request record. + with input_lock: + input_metrics[name] = input_metrics.get(name, 0.0) + seconds + if observer is not None: + observer(stage, seconds, modality) + + token = _input_processing_observer.set(observe_input) + try: + yield + finally: + _input_processing_observer.reset(token) + + class PerfMetricsMiddleware: """Expose request metrics and optionally persist completed records.""" @@ -797,11 +844,13 @@ def __init__( expose_headers: bool, writer: Optional[PerfMetricsJsonlWriter] = None, adjusted_clock: Optional[AdjustedSteadyClock] = None, + input_processing_observer: Callable[[str, float, str | None], None] | None = None, ): self._app = app self._expose_headers = expose_headers self._writer = writer self._adjusted_clock = adjusted_clock + self._input_processing_observer = input_processing_observer async def __call__(self, scope: Dict[str, Any], receive: Any, send: Any) -> None: if scope["type"] != "http": @@ -857,4 +906,9 @@ async def send_metrics(message: Dict[str, Any]) -> None: return await send(message) - await self._app(scope, receive, send_metrics) + # Keep liveness probes out of the input-processing histograms. + input_observer = ( + self._input_processing_observer if scope.get("path") != "/health_generate" else None + ) + with _input_processing_context(scope["state"], input_observer): + await self._app(scope, receive, send_metrics) diff --git a/tests/integration/test_lists/test-db/l0_cpu.yml b/tests/integration/test_lists/test-db/l0_cpu.yml index ff302874d5d6..67dff01d7c30 100644 --- a/tests/integration/test_lists/test-db/l0_cpu.yml +++ b/tests/integration/test_lists/test-db/l0_cpu.yml @@ -333,12 +333,16 @@ l0_cpu: - unittest/llmapi/apps/test_harmony_channel_validation.py - unittest/llmapi/apps/test_harmony_parsing.py - unittest/llmapi/apps/test_kimi_serve_extensions.py + - unittest/llmapi/apps/test_media_io.py::test_media_metrics_separate_download_decode_and_queue + - unittest/llmapi/apps/test_media_io.py::test_media_decode_failure_is_recorded + - unittest/llmapi/apps/test_media_io.py::test_media_process_pool_keeps_decode_picklable # test_musespark_chat_template.py reads the Muse Spark Open tokenizer files from LLM_MODELS_ROOT. - unittest/llmapi/apps/test_musespark_chat_template.py - unittest/llmapi/apps/test_musespark_serve_extensions.py - unittest/llmapi/apps/test_musespark_tool_parser.py - unittest/llmapi/apps/test_openai_chat_disagg_multimodal.py - unittest/llmapi/apps/test_openai_chat_multimodal_prompt.py + - unittest/llmapi/apps/test_openai_input_metrics.py - unittest/llmapi/apps/test_openai_protocol_chat_agent_fields.py - unittest/llmapi/apps/test_openai_protocol_chat_store.py - unittest/llmapi/apps/test_openai_server_graceful_shutdown.py @@ -354,6 +358,8 @@ l0_cpu: - unittest/llmapi/apps/test_request_metrics.py::test_header_derived_record_keeps_kv_transfer_without_kv_cache_size - unittest/llmapi/apps/test_request_metrics.py::test_unpopulated_kv_transfer_timestamps_stay_absent - unittest/llmapi/apps/test_request_metrics.py::TestJsonlWriterSurvival + - unittest/llmapi/apps/test_request_metrics.py::test_input_metrics_survive_headers_disagg_and_jsonl + - unittest/llmapi/apps/test_request_metrics.py::test_input_metrics_are_request_local_and_exported - unittest/llmapi/apps/test_responses_custom_tools.py - unittest/llmapi/apps/test_responses_deepseek_v41_namespaces.py - unittest/llmapi/apps/test_responses_input_preprocess.py diff --git a/tests/unittest/_torch/multimodal/test_multimodal_mixin.py b/tests/unittest/_torch/multimodal/test_multimodal_mixin.py index 80668cb785e7..7d00f8daa299 100644 --- a/tests/unittest/_torch/multimodal/test_multimodal_mixin.py +++ b/tests/unittest/_torch/multimodal/test_multimodal_mixin.py @@ -15,7 +15,7 @@ # SPDX-License-Identifier: Apache-2.0 from types import SimpleNamespace -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest import torch @@ -415,6 +415,7 @@ def test_encoder_cache_creation_logs_byte_capacity_without_embedding_metadata(): ) +@pytest.mark.cpu_only def test_encoder_cache_repeated_item_across_requests_skips_encoder(): model = CountingEncoderMultimodalModel( make_embedding(hidden_size=4), @@ -423,11 +424,15 @@ def test_encoder_cache_repeated_item_across_requests_skips_encoder(): ) first = make_keyed_multimodal_param() second = make_keyed_multimodal_param() + model._encoder_forward_context = MagicMock() first_embeddings = model._get_or_encode_multimodal_embeddings([first]) second_embeddings = model._get_or_encode_multimodal_embeddings([second]) assert model.encode_calls == 1 + model._encoder_forward_context.assert_called_once_with() + model._encoder_forward_context.return_value.__enter__.assert_called_once_with() + model._encoder_forward_context.return_value.__exit__.assert_called_once() torch.testing.assert_close(second_embeddings, first_embeddings) diff --git a/tests/unittest/_torch/pyexecutor/test_perf_metrics_manager.py b/tests/unittest/_torch/pyexecutor/test_perf_metrics_manager.py index fa7d4787c7e3..bf95d815198e 100644 --- a/tests/unittest/_torch/pyexecutor/test_perf_metrics_manager.py +++ b/tests/unittest/_torch/pyexecutor/test_perf_metrics_manager.py @@ -1,17 +1,49 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 +from threading import Lock from types import SimpleNamespace from unittest.mock import MagicMock, call import pytest +from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin from tensorrt_llm._torch.pyexecutor.perf_metrics_manager import PerfMetricsManager from tensorrt_llm._torch.pyexecutor.py_executor import PyExecutor pytestmark = pytest.mark.cpu_only +@pytest.mark.parametrize( + "enabled,rank,compiled,multimodal", + [ + (True, 0, False, True), + (False, 0, False, True), + (True, 1, False, True), + (True, 0, True, True), + (True, 0, False, False), + ], +) +def test_encoder_timing_callback_installation(enabled, rank, compiled, multimodal): + model = MultimodalModelMixin() if multimodal else SimpleNamespace(_encoder_forward_context=None) + model._encoder_forward_context = None + executor = SimpleNamespace( + model_engine=SimpleNamespace(model=SimpleNamespace(_orig_mod=model) if compiled else model), + enable_iter_perf_stats=enabled, + dist=SimpleNamespace(rank=rank), + perf_manager=MagicMock(), + ) + + PyExecutor._setup_encoder_perf_metrics(executor) + + expected = ( + executor.perf_manager.record_encoder_forward + if enabled and rank == 0 and multimodal + else None + ) + assert model._encoder_forward_context == expected + + class _FakeEvent: def __init__(self, elapsed_time_ms=0.0): self.elapsed_time_ms = elapsed_time_ms @@ -47,6 +79,78 @@ def _make_request(events, forward_start_time): return SimpleNamespace(py_perf_timing=timing), metric +def test_encoder_timing_retains_unready_events_and_reports_each_call_once(): + manager = PerfMetricsManager(enabled=False) + start, end = MagicMock(), MagicMock() + start.elapsed_time.return_value = 2.5 + end.query.side_effect = [False, True] + manager.borrow_forward_timing_events = MagicMock(return_value=(start, end)) + + with manager.record_encoder_forward(): + start.record.assert_called_once_with() + end.record.assert_not_called() + end.record.assert_called_once_with() + + assert manager.take_encoder_gpu_times_ms() == [] + assert manager._forward_event_pool == [] + start.elapsed_time.assert_not_called() + + assert manager.take_encoder_gpu_times_ms() == [2.5] + assert manager._forward_event_pool == [(start, end)] + assert manager.take_encoder_gpu_times_ms() == [] + start.synchronize.assert_not_called() + end.synchronize.assert_not_called() + + +def test_encoder_timing_discards_failed_events_without_reusing_them(): + manager = PerfMetricsManager(enabled=False) + start, end = MagicMock(), MagicMock() + end.query.return_value = True + start.elapsed_time.side_effect = RuntimeError("invalid event") + manager._encoder_events = [(start, end)] + + assert manager.take_encoder_gpu_times_ms() == [] + assert manager._encoder_events == [] + assert manager._forward_event_pool == [] + + +def test_encoder_events_are_reused_between_downsampled_stats_records(): + manager = PerfMetricsManager(enabled=False) + start, end = MagicMock(), MagicMock() + start.elapsed_time.return_value = 2.5 + end.query.return_value = True + manager._forward_event_pool = [(start, end)] + + for _ in range(3): + with manager.record_encoder_forward(): + pass + + assert manager.take_encoder_gpu_times_ms() == [2.5, 2.5, 2.5] + assert manager._forward_event_pool == [(start, end)] + + +@pytest.mark.parametrize("adp_rank", [None, 0, 1]) +def test_encoder_samples_are_not_duplicated_by_adp_fanout(adp_rank, monkeypatch): + monkeypatch.delenv("TLLM_METRICS_ALL_RANKS", raising=False) + manager = MagicMock() + manager.take_encoder_gpu_times_ms.return_value = [2.5] + executor = SimpleNamespace( + enable_attention_dp=True, + disable_overlap_scheduler=False, + dist=SimpleNamespace(tp_size=2, pp_size=1), + perf_manager=manager, + stats_lock=Lock(), + stats=[], + max_stats_len=100, + ) + + PyExecutor._append_iter_stats(executor, MagicMock(), attention_dp_rank=adp_rank) + + expected = None if adp_rank == 1 else {"gpuTimesMS": [2.5]} + assert executor.stats[0][8] == expected + assert manager.take_encoder_gpu_times_ms.call_count == (adp_rank != 1) + + def _make_emitting_request(request_id): request = MagicMock() request.py_request_id = request_id diff --git a/tests/unittest/executor/test_stats_serializer.py b/tests/unittest/executor/test_stats_serializer.py index 71a0970f6411..4ba7d4fdd93e 100644 --- a/tests/unittest/executor/test_stats_serializer.py +++ b/tests/unittest/executor/test_stats_serializer.py @@ -122,6 +122,15 @@ def _make_mock_kv_iter_stats( class TestStatsSerializer: + @pytest.mark.parametrize("times", [[], [0.0, 2.5]]) + def test_serializer_encoder_gpu_samples(self, times): + iter_stats = _make_mock_iteration_stats() + stats = (iter_stats, None, None, 0, None, None, "overlap", None, {"gpuTimesMS": times}) + result = json.loads(BaseWorker._stats_serializer(stats)) + assert result["multimodalEncoderStats"]["gpuTimesMS"] == times + legacy = json.loads(BaseWorker._stats_serializer(stats[:8])) + assert "multimodalEncoderStats" not in legacy + def test_serializer_without_kv_iter_stats(self): """Legacy 2-tuple and 3-tuple with None should produce same output.""" iter_stats = _make_mock_iteration_stats() diff --git a/tests/unittest/llmapi/apps/test_media_io.py b/tests/unittest/llmapi/apps/test_media_io.py index d139ea2cd4a6..421679b73eff 100644 --- a/tests/unittest/llmapi/apps/test_media_io.py +++ b/tests/unittest/llmapi/apps/test_media_io.py @@ -1,16 +1,106 @@ # SPDX-License-Identifier: Apache-2.0 # Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. -from unittest.mock import AsyncMock, patch +import base64 +from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor +from io import BytesIO +from multiprocessing import get_context +from unittest.mock import AsyncMock, MagicMock, patch import pytest +from PIL import Image from tensorrt_llm.inputs import MultimodalDataTracker from tensorrt_llm.inputs.media_io import AudioMediaIO, BaseMediaIO, ImageMediaIO, VideoMediaIO +from tensorrt_llm.metrics.perf_utils import _input_processing_observer from tensorrt_llm.serve.chat_utils import parse_chat_message_content_part pytestmark = pytest.mark.cpu_only +@pytest.mark.asyncio +@pytest.mark.parametrize( + "media_cls, mime", + [(ImageMediaIO, "image/png"), (AudioMediaIO, "audio/wav"), (VideoMediaIO, "video/mp4")], +) +@pytest.mark.parametrize("scheme", ["http", "data", "file"]) +async def test_media_metrics_separate_download_decode_and_queue(media_cls, mime, scheme): + media = media_cls() + observer = MagicMock() + token = _input_processing_observer.set(observer) + now = [1.0] + + async def open_session(): + now[0] += 5.0 + + async def download(*args, **kwargs): + now[0] += 2.0 + return b"data" + + url, method = { + "http": ("http://example.test/media", "load_bytes"), + "data": (f"data:{mime};base64,AA==", "load_base64"), + "file": ("file:///unused", "load_file"), + }[scheme] + try: + with ThreadPoolExecutor(max_workers=1) as pool: + with ( + patch.object(BaseMediaIO, "_executor", pool), + patch.object(media_cls, method, return_value="decoded"), + patch( + "tensorrt_llm.metrics.perf_utils.time.perf_counter", side_effect=lambda: now[0] + ), + patch( + "tensorrt_llm.inputs.media_io._get_aiohttp_session", side_effect=open_session + ), + patch( + "tensorrt_llm.inputs.media_io._safe_aiohttp_get", + side_effect=download, + ), + ): + assert await media.async_load(url) == "decoded" + finally: + _input_processing_observer.reset(token) + samples = [args.args for args in observer.call_args_list] + assert [stage for stage, _, _ in samples] == (["download"] if scheme == "http" else []) + [ + "decode_queue", + "decode", + ] + assert all(seconds >= 0 and modality == mime.split("/")[0] for _, seconds, modality in samples) + if scheme == "http": + assert samples[0][1] == 2.0 + + +@pytest.mark.asyncio +async def test_media_decode_failure_is_recorded(): + observer = MagicMock() + token = _input_processing_observer.set(observer) + try: + with patch.object(ImageMediaIO, "load_base64", side_effect=ValueError("bad media")): + with pytest.raises(ValueError, match="bad media"): + await ImageMediaIO().async_load("data:image/png;base64,AA==") + finally: + _input_processing_observer.reset(token) + assert [args.args[0] for args in observer.call_args_list] == ["decode_queue", "decode"] + + +@pytest.mark.asyncio +async def test_media_process_pool_keeps_decode_picklable(): + data = BytesIO() + Image.new("RGB", (2, 3), color="red").save(data, format="PNG") + url = "data:image/png;base64," + base64.b64encode(data.getvalue()).decode() + observer = MagicMock() + token = _input_processing_observer.set(observer) + try: + with ProcessPoolExecutor(max_workers=1, mp_context=get_context("spawn")) as pool: + with patch.object(BaseMediaIO, "_executor", pool): + image = await ImageMediaIO(format="pil").async_load(url) + assert image.size == (2, 3) + assert image.getpixel((0, 0)) == (255, 0, 0) + observer.assert_not_called() + finally: + _input_processing_observer.reset(token) + + class CustomError(Exception): pass diff --git a/tests/unittest/llmapi/apps/test_openai_input_metrics.py b/tests/unittest/llmapi/apps/test_openai_input_metrics.py new file mode 100644 index 000000000000..af2d6ec2dfa8 --- /dev/null +++ b/tests/unittest/llmapi/apps/test_openai_input_metrics.py @@ -0,0 +1,315 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 + +import asyncio +import json +from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace +from unittest.mock import AsyncMock, MagicMock, call, patch + +import pytest +from fastapi.responses import JSONResponse +from test_common.serve.app_factory import AggConfig, isolated_serve_env + +from tensorrt_llm.bindings.executor import RequestPerfMetrics +from tensorrt_llm.llmapi.disagg_utils import ServerRole +from tensorrt_llm.metrics.perf_utils import _input_processing_observer, _time_input_processing +from tensorrt_llm.serve.anthropic_protocol import AnthropicCreateBatchRequest +from tensorrt_llm.serve.openai_protocol import ChatCompletionRequest, ResponsesRequest +from tensorrt_llm.serve.perf_metrics import ( + RETURN_METRICS_HEADER, + SERVER_TIMING_HEADER, + PerfMetricsMiddleware, +) + +pytestmark = pytest.mark.cpu_only + + +@pytest.fixture +def observer(): + observer = MagicMock() + token = _input_processing_observer.set(observer) + yield observer + _input_processing_observer.reset(token) + + +@pytest.fixture +def serve_env(tmp_path): + with isolated_serve_env(tmp_path) as env: + yield env + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fails", [False, True]) +async def test_processor_metrics_separate_queue_and_execution(fails, observer, serve_env): + server = serve_env.build_openai_server(AggConfig()) + + def process(): + if fails: + raise ValueError("bad input") + return "processed" + + with ThreadPoolExecutor(max_workers=1) as pool: + server._input_proc_executor = pool + with patch( + "tensorrt_llm.metrics.perf_utils.time.perf_counter", side_effect=[1.0, 1.5, 1.5, 1.75] + ): + if fails: + with pytest.raises(ValueError, match="bad input"): + await server._run_input_processor(process) + else: + assert await server._run_input_processor(process) == "processed" + + assert observer.call_args_list == [ + call("processor_queue", 0.5, None), + call("processor", 0.25, None), + ] + + +@pytest.mark.asyncio +async def test_input_stage_preserves_failure_and_records_elapsed_time(observer): + with patch("tensorrt_llm.metrics.perf_utils.time.perf_counter", side_effect=[1.0, 2.0]): + with pytest.raises(ValueError, match="bad media"): + with _time_input_processing("media_load"): + raise ValueError("bad media") + + observer.assert_called_once_with("media_load", 1.0, None) + + +@pytest.mark.asyncio +async def test_disabled_input_metrics_do_not_read_clock(serve_env): + server = serve_env.build_openai_server(AggConfig()) + + async def load(): + return "media" + + with ThreadPoolExecutor(max_workers=1) as pool: + server._input_proc_executor = pool + with patch("tensorrt_llm.metrics.perf_utils.time.perf_counter") as clock: + assert await server._run_input_processor(lambda: "processed") == "processed" + with _time_input_processing("media_load"): + assert await load() == "media" + clock.assert_not_called() + + +@pytest.mark.parametrize("enabled", [False, True]) +def test_mm_encoder_prometheus_route_is_optional(enabled, serve_env): + server = serve_env.build_openai_server( + AggConfig(role=ServerRole.MM_ENCODER, return_perf_metrics=enabled) + ) + + response = serve_env.client(server.app).get("/prometheus/metrics") + assert response.status_code == (200 if enabled else 404) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", ["openai_mm_encoder", "openai_responses"]) +@pytest.mark.parametrize("fails", [False, True]) +async def test_total_includes_inline_preprocessing_but_not_generation( + route, fails, observer, serve_env +): + server = serve_env.build_openai_server( + AggConfig(role=ServerRole.MM_ENCODER if route == "openai_mm_encoder" else None) + ) + server.model = "test-model" + server.allow_request_chat_template = False + server.model_config = None + server.multimodal_server_config = None + server.enable_store = False + server.conversation_store = None + server.use_harmony = False + server.tokenizer = None + server.processor = None + server.chat_template = None + server.tool_parser = None + server.await_disconnected = AsyncMock() + if route != "openai_mm_encoder": + server._extract_metrics = AsyncMock() + now = [1.0] + server._collect_perf_metrics = True + server._adjusted_steady_clock = SimpleNamespace( + now=lambda: now[0], to_reference_time=lambda value: value + ) + server.metrics_collector = None + + async def complete(): + now[0] += 10.0 + + promise = SimpleNamespace( + request_id="1", + finished=True, + outputs=[SimpleNamespace(request_perf_metrics=RequestPerfMetrics())], + candidate_metrics=[], + metrics_dict={}, + prompt_token_ids=[7, 8], + aresult=AsyncMock(side_effect=complete), + disaggregated_params=SimpleNamespace( + multimodal_embedding_handles=[{"tensor_size": [4, 8]}] + ), + ) + + def generate_async(**kwargs): + now[0] += 0.05 + if fails: + raise ValueError("bad preprocessing") + return promise + + server.generator = SimpleNamespace( + args=SimpleNamespace(reasoning_parser=None, num_postprocess_workers=0), + generate_async=generate_async, + ) + if route == "openai_mm_encoder": + request = ChatCompletionRequest( + model=server.model, + messages=[{"role": "user", "content": "hi"}], + prompt_token_ids=[7, 8], + ) + else: + request = ResponsesRequest(model=server.model, input="hi") + raw_request = server._synthetic_request(request.model_dump(exclude_none=True)) + raw_request.scope["headers"].append((RETURN_METRICS_HEADER.encode(), b"1")) + writer = MagicMock() + messages = [] + response = None + + async def app(scope, receive, send): + nonlocal response + response = await getattr(server, route)(request, raw_request) + await response(scope, receive, send) + + async def send(message): + messages.append(message) + + middleware = PerfMetricsMiddleware( + app, expose_headers=True, writer=writer, input_processing_observer=observer + ) + + with ( + patch("tensorrt_llm.metrics.perf_utils.time.perf_counter", side_effect=lambda: now[0]), + patch( + "tensorrt_llm.serve.openai_server.parse_chat_messages_coroutines", + side_effect=lambda *args, **kwargs: ([], asyncio.sleep(0, result=(None, None)), {}, []), + ), + patch( + "tensorrt_llm.serve.openai_server.responses_api_request_preprocess", + new=AsyncMock(return_value=({"prompt": "hi"}, None)), + ), + patch( + "tensorrt_llm.serve.openai_server.responses_api_create_response", + new=AsyncMock(return_value=SimpleNamespace(model_dump=lambda **kwargs: {})), + ), + ): + await middleware(raw_request.scope, raw_request.receive, send) + + assert response.status_code == (400 if fails else 200) + totals = [sample.args[1] for sample in observer.call_args_list if sample.args[0] == "total"] + assert totals == pytest.approx([0.05]) + if route == "openai_mm_encoder" and not fails: + headers = {name.decode(): value.decode() for name, value in messages[0]["headers"]} + assert "server_input_total;dur=50.000000" in headers[SERVER_TIMING_HEADER] + writer.submit.assert_called_once() + record = writer.submit.call_args.args[0] + assert record["request_id"] == "1" + assert record["phases"]["server"]["input_processing_metrics"]["total"] == pytest.approx( + 0.05 + ) + else: + writer.submit.assert_not_called() + + +@pytest.mark.asyncio +@pytest.mark.parametrize("collect_metrics", [False, True]) +async def test_background_batch_items_have_separate_input_metrics_and_jsonl( + collect_metrics, serve_env +): + server = serve_env.build_openai_server(AggConfig()) + server.model = "test-model" + server._collect_perf_metrics = collect_metrics + server.metrics_collector = MagicMock() + server._perf_metrics_writer = MagicMock() + gate = asyncio.Event() + item_states = {} + + async def chat(request, raw_request): + await gate.wait() + key = str(request.max_completion_tokens) + item_states[key] = raw_request.scope["state"] + observer = _input_processing_observer.get() + if collect_metrics: + await asyncio.to_thread(observer, "processor", request.max_completion_tokens / 10, None) + raw_request.state.perf_metrics_records.append( + {"request_id": key, "input": dict(raw_request.state.input_processing_metrics)} + ) + else: + assert observer is None + return JSONResponse( + { + "id": key, + "object": "chat.completion", + "created": 1, + "model": server.model, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 2, "completion_tokens": 1, "total_tokens": 3}, + } + ) + + server.openai_chat = chat + request = AnthropicCreateBatchRequest( + requests=[ + { + "custom_id": str(index), + "params": { + "model": server.model, + "max_tokens": index, + "messages": [{"role": "user", "content": "hi"}], + }, + } + for index in (1, 2) + ] + ) + batch_ids = [] + + async def app(scope, receive, send): + response = await server.anthropic_create_batch(request) + batch_ids.append(json.loads(response.body)["id"]) + await response(scope, receive, send) + + async def send(message): + pass + + scope = {"type": "http", "headers": []} + if collect_metrics: + await PerfMetricsMiddleware(app, expose_headers=True)(scope, None, send) + else: + await app(scope, None, send) + gate.set() + store = server._batch_store() + await asyncio.wait_for(store._batches[batch_ids[0]].task, timeout=5) + + assert store.get(batch_ids[0]).request_counts.succeeded == 2 + assert _input_processing_observer.get() is None + if collect_metrics: + assert scope["state"]["input_processing_metrics"] == {} + assert item_states["1"]["input_processing_metrics"] == {"processor": 0.1} + assert item_states["2"]["input_processing_metrics"] == {"processor": 0.2} + server._perf_metrics_writer.submit.assert_has_calls( + [ + call({"request_id": key, "input": {"processor": int(key) / 10}}) + for key in ("1", "2") + ], + any_order=True, + ) + assert server._perf_metrics_writer.submit.call_count == 2 + server.metrics_collector.log_input_processing.assert_has_calls( + [call("processor", 0.1, None), call("processor", 0.2, None)], any_order=True + ) + else: + assert all("input_processing_metrics" not in state for state in item_states.values()) + server._perf_metrics_writer.submit.assert_not_called() + server.metrics_collector.log_input_processing.assert_not_called() diff --git a/tests/unittest/llmapi/apps/test_request_metrics.py b/tests/unittest/llmapi/apps/test_request_metrics.py index cffed81b75e1..2579de2395ec 100644 --- a/tests/unittest/llmapi/apps/test_request_metrics.py +++ b/tests/unittest/llmapi/apps/test_request_metrics.py @@ -16,10 +16,12 @@ import json import logging import threading +from unittest.mock import MagicMock import pytest from tensorrt_llm.logger import logger +from tensorrt_llm.metrics.perf_utils import _input_processing_observer from tensorrt_llm.serve import perf_metrics as perf_metrics_module from tensorrt_llm.serve.perf_metrics import ( CTX_CHUNK_METRICS_HEADER, @@ -40,6 +42,83 @@ pytestmark = pytest.mark.cpu_only +def test_input_metrics_survive_headers_disagg_and_jsonl(): + record = _record() + inputs = { + "total": 0.15, + "media_load": 0.1, + "processor": 0.04, + "media_download_image": 0.08, + "media_decode_image": 0.02, + } + record["phases"]["server"]["input_processing_metrics"] = inputs + headers = build_metrics_headers([record]) + assert "server_input_total;dur=150.000000" in headers[SERVER_TIMING_HEADER] + ctx = build_metrics_record_from_headers(headers, "ctx", request_id="42") + assert ctx["phases"]["ctx"]["input_processing_metrics"] == pytest.approx(inputs) + assert _jsonl_record(record)["perf_metrics"]["input_processing_metrics"] == inputs + combined = combine_disagg_metrics("1", {"timing_metrics": {}}, ctx, None) + assert combined["phases"]["ctx"]["input_processing_metrics"] == pytest.approx(inputs) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("path", ["/v1/chat/completions", "/health_generate"]) +async def test_input_metrics_are_request_local_and_exported(stream, path): + async def app(scope, receive, send): + await asyncio.sleep(0) + observer = _input_processing_observer.get() + await asyncio.to_thread(observer, "decode", scope["seconds"], "image") + record = _record() + record["phases"]["server"]["input_processing_metrics"] = dict( + scope["state"]["input_processing_metrics"] + ) + scope["state"]["perf_metrics_records"].append(record) + await send( + { + "type": "http.response.start", + "status": 200, + "headers": [ + (b"content-type", b"text/event-stream" if stream else b"application/json") + ], + } + ) + await send({"type": "http.response.body", "body": b"data: [DONE]\n\n" if stream else b"{}"}) + + collector = MagicMock() + middleware = PerfMetricsMiddleware( + app, expose_headers=True, input_processing_observer=collector + ) + + async def request(seconds): + messages = [] + + async def capture(message): + messages.append(message) + + scope = { + "type": "http", + "path": path, + "headers": [(RETURN_METRICS_HEADER.encode(), b"1")], + "seconds": seconds, + } + await middleware(scope, None, capture) + assert scope["state"]["input_processing_metrics"] == {"media_decode_image": seconds} + if stream: + data = messages[-1]["body"].split(b"data: ")[-1].strip() + headers = json.loads(data) + else: + headers = {key.decode(): value.decode() for key, value in messages[0]["headers"]} + assert ( + f"server_input_media_decode_image;dur={seconds * 1000:.6f}" + in headers[SERVER_TIMING_HEADER] + ) + + await asyncio.gather(request(0.1), request(0.2)) + assert collector.call_count == (0 if path == "/health_generate" else 2) + assert _input_processing_observer.get() is None + + def _record(status="complete"): return { "schema_version": 1, diff --git a/tests/unittest/metrics/test_collector.py b/tests/unittest/metrics/test_collector.py index 4f722d80df4d..ed872eff4c39 100644 --- a/tests/unittest/metrics/test_collector.py +++ b/tests/unittest/metrics/test_collector.py @@ -125,6 +125,40 @@ def test_iteration_latency_ms_to_seconds(self, collector): collector.log_iteration_stats(SAMPLE_ITERATION_STATS) assert _get_gauge_value(collector, "iteration_latency_seconds") == pytest.approx(0.0155) + def test_encoder_samples_preserve_calls_and_convert_units(self, collector): + collector.log_iteration_stats({"multimodalEncoderStats": {"gpuTimesMS": [2.0, 4.0]}}) + collector.log_iteration_stats({}) + labels = collector.labels + assert ( + REGISTRY.get_sample_value("trtllm_multimodal_encoder_gpu_time_seconds_count", labels) + == 2 + ) + assert REGISTRY.get_sample_value( + "trtllm_multimodal_encoder_gpu_time_seconds_sum", labels + ) == pytest.approx(0.006) + + def test_input_processing_separates_stages_and_worker_queue(self, collector): + collector.log_input_processing("media_load", 0.25) + collector.log_input_processing("processor", 0.04) + collector.log_input_processing("processor_queue", 0.1) + for stage, elapsed in (("media_load", 0.25), ("processor", 0.04)): + assert REGISTRY.get_sample_value( + "trtllm_input_processing_seconds_sum", {**collector.labels, "stage": stage} + ) == pytest.approx(elapsed) + assert REGISTRY.get_sample_value( + "trtllm_input_processing_queue_seconds_sum", collector.labels + ) == pytest.approx(0.1) + + def test_media_processing_is_labeled_per_item(self, collector): + collector.log_input_processing("download", 0.1, "image") + collector.log_input_processing("decode", 0.02, "image") + collector.log_input_processing("decode", 0.03, "image") + labels = {**collector.labels, "stage": "decode", "modality": "image"} + assert REGISTRY.get_sample_value("trtllm_media_processing_seconds_count", labels) == 2 + assert REGISTRY.get_sample_value( + "trtllm_media_processing_seconds_sum", labels + ) == pytest.approx(0.05) + def test_memory_usage(self, collector): collector.log_iteration_stats(SAMPLE_ITERATION_STATS) diff --git a/tests/unittest/serve/test_middleware_order.py b/tests/unittest/serve/test_middleware_order.py index aec17dd4f79d..4221d09d422d 100644 --- a/tests/unittest/serve/test_middleware_order.py +++ b/tests/unittest/serve/test_middleware_order.py @@ -51,21 +51,27 @@ def _perf(kwargs, expose_headers): AggConfig(return_perf_metrics=True), [ _ARRIVAL_WITH_CLOCK, - _perf(["adjusted_clock", "expose_headers", "writer"], True), + _perf( + ["adjusted_clock", "expose_headers", "input_processing_observer", "writer"], True + ), ], ), "llm_perf_metrics_output_dir_only": ( AggConfig(perf_metrics_output_dir=True), [ _ARRIVAL_WITH_CLOCK, - _perf(["adjusted_clock", "expose_headers", "writer"], False), + _perf( + ["adjusted_clock", "expose_headers", "input_processing_observer", "writer"], False + ), ], ), "llm_perf_metrics_enable_metrics_env": ( AggConfig(return_perf_metrics=True, env=(("ENABLE_METRICS", "true"),)), [ _ARRIVAL_WITH_CLOCK, - _perf(["adjusted_clock", "expose_headers", "writer"], True), + _perf( + ["adjusted_clock", "expose_headers", "input_processing_observer", "writer"], True + ), {"class": "PrometheusInstrumentatorMiddleware"}, ], ),