diff --git a/examples/disaggregated/slurm/benchmark/run_benchmark_aiperf.sh b/examples/disaggregated/slurm/benchmark/run_benchmark_aiperf.sh index 1f2d19b18417..385936d67b8b 100755 --- a/examples/disaggregated/slurm/benchmark/run_benchmark_aiperf.sh +++ b/examples/disaggregated/slurm/benchmark/run_benchmark_aiperf.sh @@ -92,11 +92,5 @@ for concurrency in ${concurrency_list}; do echo "Benchmark with concurrency ${concurrency} done" done -# Fetch perf metrics from disagg server -echo "Fetching perf metrics from http://${hostname}:${port}/perf_metrics ..." -curl -s "http://${hostname}:${port}/perf_metrics" > ${log_path}/perf_metrics.json 2>&1 || true -if [ -s "${log_path}/perf_metrics.json" ]; then - echo "Perf metrics saved to ${log_path}/perf_metrics.json" -else - echo "Warning: perf_metrics response was empty or endpoint not available" -fi +# Configure perf_metrics_output_dir on each server to persist per-request JSONL. +echo "Per-request metrics are available in the configured server-side JSONL output." diff --git a/tensorrt_llm/llmapi/disagg_utils.py b/tensorrt_llm/llmapi/disagg_utils.py index a8272b494645..901bf7ad1c19 100644 --- a/tensorrt_llm/llmapi/disagg_utils.py +++ b/tensorrt_llm/llmapi/disagg_utils.py @@ -92,6 +92,8 @@ class DisaggServerConfig(): otlp_config: Optional[OtlpConfig] = None max_retries: int = 1 perf_metrics_max_requests: int = 0 + return_perf_metrics: bool = False + perf_metrics_output_dir: Optional[str] = None disagg_cluster_config: Optional[DisaggClusterConfig] = None node_id: int = uuid.getnode( ) % 256 # Assuming only one disagg-server is running on a machine, modulo 256. @@ -189,6 +191,8 @@ def extract_disagg_cfg(hostname: str = 'localhost', port: int = 8000, max_retries: int = 1, perf_metrics_max_requests: int = 0, + return_perf_metrics: bool = False, + perf_metrics_output_dir: Optional[str] = None, context_servers: Optional[dict] = None, generation_servers: Optional[dict] = None, conditional_disagg_config: Optional[dict] = None, @@ -207,10 +211,12 @@ def extract_disagg_cfg(hostname: str = 'localhost', context_servers = context_servers or {} generation_servers = generation_servers or {} + inherited_args = dict(kwargs) + # If parameters are specified outside the context_severs and generation_servers sections, # make sure they match # Also inherit the values from the top-level - for key, value in kwargs.items(): + for key, value in inherited_args.items(): for server_type, servers in [("context_servers", context_servers), ("generation_servers", generation_servers) ]: @@ -241,11 +247,19 @@ def extract_disagg_cfg(hostname: str = 'localhost', otlp_config = OtlpConfig(**otlp_config) if otlp_config else None - config = DisaggServerConfig(server_configs, hostname, port, - ctx_router_config, gen_router_config, - conditional_disagg_config, otlp_config, - max_retries, perf_metrics_max_requests, - disagg_cluster_config) + config = DisaggServerConfig( + server_configs=server_configs, + hostname=hostname, + port=port, + ctx_router_config=ctx_router_config, + gen_router_config=gen_router_config, + conditional_disagg_config=conditional_disagg_config, + otlp_config=otlp_config, + max_retries=max_retries, + perf_metrics_max_requests=perf_metrics_max_requests, + return_perf_metrics=return_perf_metrics, + perf_metrics_output_dir=perf_metrics_output_dir, + disagg_cluster_config=disagg_cluster_config) if node_id is not None: node_id_space = 1 << DISAGG_NODE_ID_BITS if not 0 <= node_id < node_id_space: diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 9c679c5d8bd0..5549e1583df6 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/tensorrt_llm/llmapi/llm_args.py @@ -4390,15 +4390,27 @@ class BaseLlmArgs(StrictBaseModel): status="deprecated", telemetry=TelemetryField.categorical('pytorch', '_autodeploy')) - return_perf_metrics: bool = Field(default=False, - description="Return perf metrics.", - status="prototype") + return_perf_metrics: bool = Field( + default=False, + description= + "Allow serving responses to include per-request performance metrics when " + "the request sets X-TRTLLM-return-metrics: 1.", + status="prototype") + + perf_metrics_output_dir: Optional[str] = Field( + default=None, + description="Directory for per-process performance metrics JSONL " + "files. Setting this enables collection even when " + "return_perf_metrics is false.", + status="prototype", + telemetry=False) perf_metrics_max_requests: NonNegativeInt = Field( default=0, description= - "The maximum number of requests for perf metrics. Must also set return_perf_metrics to true to get perf metrics.", - status="prototype") + "Deprecated compatibility field. Completed per-request metrics are no " + "longer retained in memory.", + status="deprecated") prometheus_metrics_config: Optional[PrometheusMetricsConfig] = Field( default=None, diff --git a/tensorrt_llm/scaffolding/task.py b/tensorrt_llm/scaffolding/task.py index 82c92032cfcc..c77767eada81 100644 --- a/tensorrt_llm/scaffolding/task.py +++ b/tensorrt_llm/scaffolding/task.py @@ -1,3 +1,17 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + from __future__ import annotations import json @@ -116,8 +130,8 @@ class GenerationTask(Task): # Server-side request id captured from the OpenAI client's streaming # chunks (every chunk's ``chunk.id`` field). Trace-replay clients use - # this to look up the matching record in trtllm-serve's - # ``/perf_metrics`` drain and attach per-request KV-cache statistics + # this to correlate response-carried per-request metrics and attach + # KV-cache statistics # (``num_reused_blocks`` / ``num_missed_blocks`` / ``free_num_blocks``) # to the per-LLM-call row in the step JSON. ``None`` until the worker # observes the first chunk (or when not running against an diff --git a/tensorrt_llm/scaffolding/trace_replay/replay.py b/tensorrt_llm/scaffolding/trace_replay/replay.py index 9c4278760692..8561acd56685 100644 --- a/tensorrt_llm/scaffolding/trace_replay/replay.py +++ b/tensorrt_llm/scaffolding/trace_replay/replay.py @@ -1,3 +1,17 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import asyncio import hashlib import random @@ -256,7 +270,7 @@ class DropPathStats: POSTs to the server, one immediately after — for every conv id this branch owns. Each probe's ``request_id`` (server-assigned, captured by the worker from the OpenAI streaming chunk) is recorded here and - later joined with trtllm-serve's ``/perf_metrics`` drain to recover + later joined with response-carried metrics to recover ``num_reused_blocks`` (the per-request KV-cache hit count) and ``free_num_blocks`` (the post-call snapshot of the free-block pool). @@ -399,7 +413,7 @@ async def _run(self): if event is None: # sentinel # Fire end-of-branch retention probes BEFORE returning, # so the parent ``wait_all_done`` covers them and the - # client's downstream /perf_metrics drain sees the + # client receives the response-carried metrics for the # probe records. For a child branch this is at # parallel_end (immediately after any drop_kv_cache); # for the root branch it is at end-of-session. @@ -503,7 +517,7 @@ async def _handle_drop_kv_cache(self, event: TraceEvent): # ``max_tokens=1, ignore_eos=True`` probe per (conv_id, phase) pair # immediately before and after the truncate. Each probe's # request_id is recorded here; the per-request KV-cache hit - # accounting is joined in later via /perf_metrics. The probes + # accounting is joined later from response-carried metrics. The probes # serialize the truncate against a real before/after measurement # so downstream verification can mechanically prove the truncate # actually freed the blocks the engine claimed it did (P0.6.a). diff --git a/tensorrt_llm/scaffolding/worker.py b/tensorrt_llm/scaffolding/worker.py index ad1432934201..4c95167164e0 100644 --- a/tensorrt_llm/scaffolding/worker.py +++ b/tensorrt_llm/scaffolding/worker.py @@ -1,3 +1,17 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + import asyncio import copy import json @@ -341,8 +355,8 @@ async def generation_handler(self, task: GenerationTask) -> TaskStatus: now = time.perf_counter() # Every chunk carries the server-assigned request id. Capture # it once (the first non-None value) so callers can later - # correlate per-request perf metrics drained from - # ``/perf_metrics`` with the GenerationTask that issued them. + # correlate response-carried per-request metrics with the + # GenerationTask that issued them. if request_id is None: cid = getattr(chunk, "id", None) if cid is not None: diff --git a/tensorrt_llm/serve/_perf_metrics_schema.py b/tensorrt_llm/serve/_perf_metrics_schema.py new file mode 100644 index 000000000000..5a9d606d85ac --- /dev/null +++ b/tensorrt_llm/serve/_perf_metrics_schema.py @@ -0,0 +1,97 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from typing import List, Optional, Union + +from typing_extensions import NotRequired, TypedDict + + +class StepMetrics(TypedDict): + forward_start_time: float + forward_end_time: float + sample_start_time: float + sample_end_time: float + gpu_forward_time: float + gpu_sample_time: float + token_time: NotRequired[float] + scheduled_time: NotRequired[float] + prev_batch_token_time: NotRequired[float] + iter: NotRequired[int] + + +class TimeBreakdownMetrics(TypedDict): + step_metrics: NotRequired[List[StepMetrics]] + ctx_chunk_metrics: NotRequired[List[StepMetrics]] + ctx_gpu_forward_time: NotRequired[float] + ctx_gpu_sample_time: NotRequired[float] + + +class TimingMetrics(TypedDict): + arrival_time: Optional[float] + first_scheduled_time: NotRequired[Optional[float]] + first_token_time: NotRequired[Optional[float]] + last_token_time: Optional[float] + server_arrival_time: NotRequired[Optional[float]] + server_first_token_time: NotRequired[Optional[float]] + kv_cache_size: NotRequired[int] + kv_cache_transfer_start: NotRequired[Optional[float]] + kv_cache_transfer_end: NotRequired[Optional[float]] + + +class KvCacheMetrics(TypedDict): + num_total_allocated_blocks: int + num_new_allocated_blocks: int + num_reused_blocks: int + num_missed_blocks: int + + +class SpeculativeDecodingMetrics(TypedDict): + acceptance_rate: float + total_accepted_draft_tokens: int + total_draft_tokens: int + + +class PerfMetrics(TypedDict): + timing_metrics: TimingMetrics + first_iter: NotRequired[int] + last_iter: NotRequired[int] + kv_cache_metrics: NotRequired[KvCacheMetrics] + speculative_decoding: NotRequired[SpeculativeDecodingMetrics] + + +class WorkerPerfMetrics(TypedDict): + request_id: Union[int, str] + perf_metrics: PerfMetrics + ctx_request_id: NotRequired[int] + time_breakdown_metrics: NotRequired[TimeBreakdownMetrics] + + +class WorkerPerfMetricsRecord(WorkerPerfMetrics): + status: str + disagg_request_id: NotRequired[int] + + +class DisaggPerfMetricsRecord(TypedDict): + ctx_server: str + gen_server: str + disagg_server_arrival_time: float + disagg_ctx_dispatch_time: Optional[float] + disagg_server_first_token_time: Optional[float] + status: str + disagg_request_id: NotRequired[int] + ctx_perf_metrics: NotRequired[WorkerPerfMetrics] + gen_perf_metrics: NotRequired[WorkerPerfMetrics] + + +PerfMetricsRecord = Union[WorkerPerfMetricsRecord, DisaggPerfMetricsRecord] diff --git a/tensorrt_llm/serve/openai_client.py b/tensorrt_llm/serve/openai_client.py index 26646a54344c..89bb75c0dc1f 100644 --- a/tensorrt_llm/serve/openai_client.py +++ b/tensorrt_llm/serve/openai_client.py @@ -14,10 +14,11 @@ # yapf: disable import asyncio +import json import os import traceback from abc import ABC, abstractmethod -from typing import Any, AsyncGenerator, Awaitable, Callable, Dict, List, Optional, Tuple, Type +from typing import Any, AsyncGenerator, Awaitable, Callable, List, Optional, Tuple, Type import aiohttp @@ -31,7 +32,13 @@ UCompletionRequest, UCompletionResponse, ) -from tensorrt_llm.serve.perf_metrics import ClientMetricsCollector +from tensorrt_llm.serve.perf_metrics import ( + _PERF_METRICS_HEADER_BUDGET_BYTES, + RETURN_METRICS_HEADER, + SSE_METRICS_EVENT, + ClientMetricsCollector, + build_metrics_record_from_headers, +) from tensorrt_llm.serve.responses_utils import ( ResponseHooks, UCompletionResponseOrGenerator, @@ -58,6 +65,10 @@ _msgpack_encoder = msgspec.msgpack.Encoder() +def _metrics_phase(role: ServerRole) -> str: + return "ctx" if role is ServerRole.CONTEXT else "gen" + + class OpenAIClient(ABC): async def send_request( self, @@ -98,9 +109,6 @@ async def _send_request( """ ... - @abstractmethod - async def collect_metrics(self) -> Dict[str, Any]: ... - @abstractmethod async def check_ready(self) -> Tuple[List[str], List[str]]: """Return the list of ready servers and the list of unready servers.""" @@ -133,6 +141,7 @@ def __init__( retry_interval_sec: int = 1, session: Optional[aiohttp.ClientSession] = None, disagg_id_generator: Optional[Callable[[], Awaitable[int]]] = None, + request_perf_metrics: bool = False, ): self._router = router self._role = role @@ -146,10 +155,12 @@ def __init__( keepalive_timeout=1, ), timeout=aiohttp.ClientTimeout(total=timeout_secs), + max_field_size=_PERF_METRICS_HEADER_BUDGET_BYTES, ) self._max_retries = max_retries self._retry_interval_sec = retry_interval_sec self._disagg_id_generator = disagg_id_generator + self._request_perf_metrics = request_perf_metrics async def _send_request( self, @@ -215,6 +226,8 @@ async def _post_with_retry( dp = getattr(request, "disaggregated_params", None) if dp is not None and getattr(dp, "disagg_request_id", None) is not None: dp.disagg_request_id = await self._disagg_id_generator() + if hooks: + hooks.on_disagg_request_id(dp.disagg_request_id) # Serialize once on the orchestrator's single event-loop thread. if _MSGSPEC_ENABLED: # msgspec msgpack: encode the request dict to msgpack bytes. Keep @@ -227,6 +240,8 @@ async def _post_with_retry( else: body = request.model_dump_json(exclude_unset=True) req_headers = {"Content-Type": "application/json"} + if self._request_perf_metrics: + req_headers[RETURN_METRICS_HEADER] = "1" try: lines_yielded = 0 start_time = get_steady_clock_now_in_seconds() @@ -236,6 +251,27 @@ async def _post_with_retry( headers=req_headers, ) as http_response: content_type = http_response.headers.get("Content-Type", "") + if self._request_perf_metrics: + role = _metrics_phase(self._role) + disagg_params = getattr(request, "disaggregated_params", None) + request_id = "" + if disagg_params is not None: + request_id = str( + disagg_params.disagg_request_id + or disagg_params.ctx_request_id + or "" + ) + response_metrics = build_metrics_record_from_headers( + http_response.headers, + role, + request_id=request_id, + ) + if hooks and response_metrics: + hooks.on_perf_metrics( + server, + role, + response_metrics, + ) if not is_stream and "text/event-stream" in content_type: raise ValueError( "Received an event-stream although request stream was False" @@ -322,24 +358,64 @@ async def _response_generator( success = True try: last_token_time = start_time - i = 0 - async for line in http_response.content.iter_any(): + chunk_count = 0 + marker = f"event: {SSE_METRICS_EVENT}\n".encode() + pending = b"" + metrics_event = b"" + async for chunk in http_response.content.iter_any(): now_time = get_steady_clock_now_in_seconds() - if line: - if i == 0: - if hooks: - hooks.on_first_token(server, request) - self._metrics_collector.first_token_latency_seconds.observe( - now_time - last_token_time - ) - else: - self._metrics_collector.per_token_latency_seconds.observe( - now_time - last_token_time - ) - i += 1 - yield line + if chunk: + if chunk_count == 0 and hooks: + hooks.on_first_token(server, request) + latency = now_time - last_token_time + metric = ( + self._metrics_collector.first_token_latency_seconds + if chunk_count == 0 + else self._metrics_collector.per_token_latency_seconds + ) + metric.observe(latency) + chunk_count += 1 + last_token_time = now_time + + if not self._request_perf_metrics: + yield chunk await asyncio.sleep(0) - last_token_time = now_time + continue + + if metrics_event: + metrics_event += chunk + continue + pending += chunk + marker_index = pending.find(marker) + if marker_index >= 0: + if marker_index: + yield pending[:marker_index] + metrics_event = pending[marker_index:] + pending = b"" + continue + emit_size = len(pending) - len(marker) + 1 + if emit_size > 0: + yield pending[:emit_size] + pending = pending[emit_size:] + await asyncio.sleep(0) + + if pending: + yield pending + if hooks and metrics_event: + data_prefix = b"data: " + data_index = metrics_event.find(data_prefix) + if data_index >= 0: + data = metrics_event[data_index + len(data_prefix) :].split(b"\n", 1)[0] + try: + headers = json.loads(data) + metrics = build_metrics_record_from_headers( + headers, _metrics_phase(self._role) + ) + except (TypeError, ValueError) as error: + logger.warning("Ignoring malformed perf metrics event: %s", error) + metrics = None + if metrics: + hooks.on_perf_metrics(server, _metrics_phase(self._role), metrics) if hooks: hooks.on_resp_done(server, request, None) @@ -374,17 +450,6 @@ async def _finish_request( request, self._session, success=success, req_id=req_id ) - async def collect_metrics(self) -> Dict[str, Any]: - metrics = {} - for server in self._router.servers: - try: - async with self._session.get(f"http://{server}/perf_metrics") as response: - metrics[server] = await response.json() - except Exception: - logger.error(f"Failed to collect metrics from {server}") - continue - return metrics - async def shutdown(self) -> None: await self._session.close() diff --git a/tensorrt_llm/serve/openai_disagg_server.py b/tensorrt_llm/serve/openai_disagg_server.py index 6df7651fb1dd..78de077034f1 100644 --- a/tensorrt_llm/serve/openai_disagg_server.py +++ b/tensorrt_llm/serve/openai_disagg_server.py @@ -46,7 +46,10 @@ from tensorrt_llm.serve.openai_protocol import ( ChatCompletionRequest, CompletionRequest, UCompletionRequest, UCompletionResponse, ensure_request_chat_template_allowed) -from tensorrt_llm.serve.perf_metrics import DisaggPerfMetricsCollector +from tensorrt_llm.serve.perf_metrics import (DisaggPerfMetricsCollector, + PerfMetricsJsonlWriter, + PerfMetricsMiddleware, + combine_disagg_metrics) from tensorrt_llm.serve.responses_utils import (ServerArrivalTimeMiddleware, get_steady_clock_now_in_seconds) from tensorrt_llm.serve.router import Router @@ -60,44 +63,79 @@ } class RawRequestResponseHooks(ResponseHooks): - def __init__(self, raw_req: Request, perf_metrics_collector: DisaggPerfMetricsCollector): + def __init__(self, raw_req: Request, queue_latency_metric, + collect_perf_metrics: bool): self.raw_req = raw_req + self.queue_latency_metric = queue_latency_metric + self.collect_perf_metrics = collect_perf_metrics self.ctx_server = "" self.gen_server = "" + self.request_id = "" + self.disagg_request_id = None self.request_arrival_time = raw_req.state.server_arrival_time self.server_first_token_time = 0 self.ctx_dispatch_time = 0 - self.perf_metrics_collector = perf_metrics_collector + self.ctx_metrics = None + self.gen_metrics = None def on_req_begin(self, request: UCompletionRequest): - self.perf_metrics_collector.queue_latency_seconds.observe(get_steady_clock_now_in_seconds() - self.request_arrival_time) + params = request.disaggregated_params + if params is not None: + self.disagg_request_id = params.disagg_request_id + request_id = params.disagg_request_id or params.ctx_request_id + self.request_id = str(request_id or "") + self.queue_latency_metric.observe( + get_steady_clock_now_in_seconds() - self.request_arrival_time) + + def on_disagg_request_id(self, disagg_request_id: int): + self.disagg_request_id = disagg_request_id + self.request_id = str(disagg_request_id) def on_ctx_dispatch(self, request: UCompletionRequest): self.ctx_dispatch_time = get_steady_clock_now_in_seconds() + def on_perf_metrics(self, server: str, role: str, metrics: dict): + if role == "ctx": + self.ctx_server = server + self.ctx_metrics = metrics + elif role == "gen": + self.gen_server = server + self.gen_metrics = metrics + def on_ctx_resp(self, ctx_server: str, response: UCompletionResponse): self.ctx_server = ctx_server - def on_first_token(self, gen_server: str, request: UCompletionRequest, response: UCompletionResponse = None): + def on_first_token( + self, gen_server: str, request: UCompletionRequest, + response: UCompletionResponse = None): self.gen_server = gen_server self.server_first_token_time = get_steady_clock_now_in_seconds() - def on_resp_done(self, gen_server: str, request: UCompletionRequest, response: UCompletionResponse = None): - if request.disaggregated_params: - ctx_req_id = request.disaggregated_params.ctx_request_id - task = asyncio.create_task( - self.perf_metrics_collector.add_per_request_metrics( - self.ctx_server, - gen_server, - ctx_req_id, - self.raw_req.state.server_arrival_time, - self.server_first_token_time, - self.ctx_dispatch_time, - ) - ) - background_tasks = self.perf_metrics_collector._background_tasks - background_tasks.add(task) - task.add_done_callback(background_tasks.discard) + def on_resp_done( + self, gen_server: str, request: UCompletionRequest, + response: UCompletionResponse = None): + self.gen_server = gen_server + if not self.collect_perf_metrics: + return + disagg_phase = { + "ctx_server": self.ctx_server, + "gen_server": self.gen_server, + "timing_metrics": { + "arrival_time": self.request_arrival_time, + "last_token_time": get_steady_clock_now_in_seconds(), + "server_arrival_time": self.request_arrival_time, + "ctx_dispatch_time": self.ctx_dispatch_time or None, + "server_first_token_time": self.server_first_token_time or None, + }, + } + self.raw_req.state.perf_metrics_records.append( + combine_disagg_metrics( + self.request_id, + disagg_phase, + self.ctx_metrics, + self.gen_metrics, + disagg_request_id=self.disagg_request_id, + )) class OpenAIDisaggServer: @@ -120,7 +158,14 @@ def __init__(self, # process owns the routers + cluster state (DisaggCoordinatorService). self._coordinator_url = coordinator_url - self._perf_metrics_collector = DisaggPerfMetricsCollector(config.perf_metrics_max_requests) + self._perf_metrics_collector = DisaggPerfMetricsCollector( + config.perf_metrics_max_requests) + self._expose_perf_metrics = config.return_perf_metrics + self._collect_perf_metrics = ( + config.return_perf_metrics + or config.perf_metrics_output_dir is not None) + self._perf_metrics_writer = PerfMetricsJsonlWriter( + config.perf_metrics_output_dir, "disagg") self._disagg_cluster_storage = None if config.disagg_cluster_config: @@ -148,8 +193,7 @@ def __init__(self, self._service = OpenAIDisaggregatedService( self._config, self._coordinator, self._create_client, - req_timeout_secs=self._req_timeout_secs, - perf_metrics_collector=self._perf_metrics_collector) + req_timeout_secs=self._req_timeout_secs) try: otlp_cfg = config.otlp_config @@ -165,17 +209,19 @@ def __init__(self, @asynccontextmanager async def lifespan(app) -> None: # The cluster manager (via setup) owns server preparation + monitoring. + await self._perf_metrics_writer.start() await self._service.setup() yield await self._service.teardown() - if self._perf_metrics_collector._background_tasks: - await asyncio.gather( - *self._perf_metrics_collector._background_tasks, - return_exceptions=True, - ) + await self._perf_metrics_writer.close() self.app = FastAPI(lifespan=lifespan) + if self._collect_perf_metrics: + self.app.add_middleware( + PerfMetricsMiddleware, + expose_headers=self._expose_perf_metrics, + writer=self._perf_metrics_writer) self.app.add_middleware(ServerArrivalTimeMiddleware) # Log request-body validation failures so a client/server schema mismatch @@ -206,11 +252,10 @@ async def validation_exception_handler(request: Request, exc): def _create_client(self, router: Router, role: ServerRole, max_retries: int = 1) -> OpenAIClient: async def disagg_id_generator(): return await self._coordinator.get_disagg_request_id() - client = OpenAIHttpClient( + return OpenAIHttpClient( router, role, self._req_timeout_secs, max_retries, - disagg_id_generator=disagg_id_generator) - self._perf_metrics_collector.add_client(client) - return client + disagg_id_generator=disagg_id_generator, + request_perf_metrics=self._collect_perf_metrics) def register_routes(self): # The disagg service owns only the request-serving endpoints (/v1/*) and @@ -221,7 +266,6 @@ def register_routes(self): self.app.add_api_route("/health", self.health, methods=["GET"]) self.app.add_api_route("/cluster_info", self.cluster_info, methods=["GET"]) self.app.add_api_route("/version", self.version, methods=["GET"]) - self.app.add_api_route("/perf_metrics", self._perf_metrics_collector.get_perf_metrics, methods=["GET"]) # import prometheus_client lazily to break the `set_prometheus_multiproc_dir` from prometheus_client import make_asgi_app self.app.mount("/prometheus/metrics", make_asgi_app()) @@ -272,6 +316,7 @@ def _wrap_entry_point(self, entry_point: Callable, request_type: type = UComplet # The bare Union UCompletionRequest (no discriminator) makes Pydantic try # CompletionRequest first and 400 every chat body, so override the wrapper's # annotation with request_type (as openai_server.py does). + @tracing.trace_span("disaggregated_request") async def wrapper(req: request_type, raw_req: Request) -> Response: try: self._perf_metrics_collector.total_requests.inc() @@ -285,13 +330,16 @@ async def wrapper(req: request_type, raw_req: Request) -> Response: except ValueError as e: raise HTTPException(status_code=400, detail=str(e)) from e self._extract_conversation_id(req, raw_req) - hooks = RawRequestResponseHooks(raw_req, self._perf_metrics_collector) + hooks = RawRequestResponseHooks( + raw_req, self._perf_metrics_collector.queue_latency_seconds, + self._collect_perf_metrics) response_or_generator = await entry_point(req, hooks) self._perf_metrics_collector.total_responses.inc() if req.stream: - return StreamingResponse(content=response_or_generator, media_type="text/event-stream") - else: - return JSONResponse(content=response_or_generator.model_dump()) + return StreamingResponse( + content=response_or_generator, + media_type="text/event-stream") + return JSONResponse(content=response_or_generator.model_dump()) except Exception as e: self._handle_exception(e) return wrapper diff --git a/tensorrt_llm/serve/openai_disagg_service.py b/tensorrt_llm/serve/openai_disagg_service.py index e2ee00257e98..8450a2c0f1a6 100644 --- a/tensorrt_llm/serve/openai_disagg_service.py +++ b/tensorrt_llm/serve/openai_disagg_service.py @@ -29,7 +29,6 @@ UCompletionResponse, ) from tensorrt_llm.serve.openai_service import OpenAIService -from tensorrt_llm.serve.perf_metrics import DisaggPerfMetricsCollector from tensorrt_llm.serve.responses_utils import ( ResponseHooks, UCompletionResponseOrGenerator, @@ -49,7 +48,6 @@ def __init__( coordinator: "DisaggCoordinator", client_factory: Callable[[Router, ServerRole], OpenAIClient], req_timeout_secs: int = 180, - perf_metrics_collector: Optional[DisaggPerfMetricsCollector] = None, ): self._config = config # The service drives the coordinator's ctx/gen routers uniformly, so serving @@ -60,7 +58,6 @@ def __init__( self._gen_router = coordinator.gen_router self._client_factory = client_factory self._req_timeout_secs = req_timeout_secs - self._perf_metrics_collector = perf_metrics_collector # Opt-in body-shrink for generation_only requests; see _get_gen_request. self._strip_gen_message_history = config.gen_strip_message_history # Opt-in: ask context workers to return prompt_token_ids as base64 int32. @@ -123,6 +120,8 @@ async def _send_disagg_request_ctx_first( # empty server means client decides which server to use ctx_server = None disagg_request_id = await self._coordinator.get_disagg_request_id() + if hooks: + hooks.on_disagg_request_id(disagg_request_id) # reserve a gen_server if conditional disagg is needed gen_server, need_ctx = await self._check_conditional_disagg(request, disagg_request_id) # Context retries may replace disagg_request_id for the KV-transfer @@ -150,6 +149,8 @@ async def _send_disagg_request_ctx_first( ctx_response_disagg_params = ctx_response.choices[0].disaggregated_params if ctx_response_disagg_params.disagg_request_id is not None: disagg_request_id = ctx_response_disagg_params.disagg_request_id + if hooks: + hooks.on_disagg_request_id(disagg_request_id) gen_req = self._get_gen_request(request, ctx_response, disagg_request_id) except Exception: if gen_server: @@ -180,6 +181,8 @@ async def _send_disagg_request_ctx_first( else: if gen_server: await self._gen_router.finish_request(request, req_id=gen_reservation_id) + if hooks: + hooks.on_resp_done("", request, ctx_response) if request.stream: # ctx client will never return a generator when streaming is requested # make up for this by returning a done generator @@ -392,6 +395,8 @@ async def _send_disagg_request_gen_first( # Single-issuer disagg id (see _send_disagg_request_ctx_first): fetch from # the coordinator so fleet workers never mint colliding ids. disagg_request_id = await self._coordinator.get_disagg_request_id() + if hooks: + hooks.on_disagg_request_id(disagg_request_id) if need_ctx: # arrival->here = pre-ctx wait in the orchestrator/fleet. if hooks: diff --git a/tensorrt_llm/serve/openai_server.py b/tensorrt_llm/serve/openai_server.py index 153b35b5f6cb..cc69721700d0 100644 --- a/tensorrt_llm/serve/openai_server.py +++ b/tensorrt_llm/serve/openai_server.py @@ -77,6 +77,9 @@ ensure_request_chat_template_allowed, to_llm_conversation_params, to_llm_disaggregated_params) from tensorrt_llm.serve.openai_video_routes import _VideoRoutesMixin +from tensorrt_llm.serve.perf_metrics import (PerfMetricsJsonlWriter, + PerfMetricsMiddleware, + build_request_metrics_record) from tensorrt_llm.serve.postprocess_handlers import ( ChatCompletionPostprocArgs, ChatPostprocArgs, CompletionPostprocArgs, ResponsesAPIPostprocArgs, chat_harmony_post_processor, @@ -315,8 +318,21 @@ def __init__( else: self.model = model self.metrics_collector = None - self.perf_metrics = None - self.perf_metrics_lock = None + args = getattr(self.generator, "args", None) + self._expose_perf_metrics = bool( + args and getattr(args, "return_perf_metrics", False)) + perf_metrics_output_dir = (getattr(args, "perf_metrics_output_dir", + None) if args else None) + self._collect_perf_metrics = (self._expose_perf_metrics + or perf_metrics_output_dir is not None) + # AsyncLLM uses this flag to request engine-level snapshots. Preserve the + # original value separately because only it controls public headers. + if self._collect_perf_metrics and args is not None: + args.return_perf_metrics = True + server_kind = server_role.name.lower( + ) if server_role is not None else "server" + self._perf_metrics_writer = PerfMetricsJsonlWriter( + perf_metrics_output_dir, server_kind) self._iteration_stats_collector_task = None self._iteration_stats_wakeup_event = asyncio.Event() # Bounded snapshot of iteration stats for the GET /metrics handler. @@ -346,6 +362,7 @@ def __init__( @asynccontextmanager async def lifespan(app: FastAPI): + await self._perf_metrics_writer.start() if self.metadata_server is not None: metadata = { "model": self.model, @@ -416,6 +433,7 @@ async def lifespan(app: FastAPI): yield + await self._perf_metrics_writer.close() if self.embedding_batcher is not None: await self.embedding_batcher.shutdown() logger.info("Stopped encode dynamic batcher") @@ -473,6 +491,10 @@ async def validation_exception_handler(_, exc): else: self.register_routes() + if self._collect_perf_metrics: + self.app.add_middleware(PerfMetricsMiddleware, + expose_headers=self._expose_perf_metrics, + writer=self._perf_metrics_writer) self.app.add_middleware(ServerArrivalTimeMiddleware) def _init_visual_gen(self): @@ -573,10 +595,6 @@ def _init_llm(self, chat_template: Optional[str] = None): pmc.request_inference_time_buckets if pmc else None), ) self._log_config_info_metrics() - max_perf_metrics = self.generator.args.perf_metrics_max_requests - if max_perf_metrics > 0: - self.perf_metrics = deque(maxlen=max_perf_metrics) - self.perf_metrics_lock = asyncio.Lock() @staticmethod def _ensure_post_processor_hook_supported( @@ -799,9 +817,6 @@ def register_routes(self): self.app.add_api_route("/metrics", self.get_iteration_stats, methods=["GET"]) - self.app.add_api_route("/perf_metrics", - self.get_perf_metrics, - methods=["GET"]) self.app.add_api_route("/energy_metrics", self.get_energy_metrics, methods=["GET"]) @@ -1246,79 +1261,6 @@ async def get_steady_clock_offset(self) -> JSONResponse: "transmit_ts": transmit_ts }) - async def get_perf_metrics(self) -> JSONResponse: - if self.perf_metrics is None: - return JSONResponse(content=[]) - async with self.perf_metrics_lock: - perf_metrics = self.perf_metrics - self.perf_metrics = deque( - maxlen=self.generator.args.perf_metrics_max_requests) - for metrics_dict in perf_metrics: - metrics = metrics_dict["perf_metrics"] - timing_metrics = metrics.timing_metrics - kv_cache_metrics = metrics.kv_cache_metrics - speculative_decoding = metrics.speculative_decoding - metrics_json = { - "first_iter": metrics.first_iter, - "last_iter": metrics.last_iter, - # exclude metrics.iter since it is only meaningful when the request is not finished - } - server_arrival_time = metrics_dict.pop("server_arrival_time", None) - if server_arrival_time is not None: - server_arrival_time += self.disagg_server_steady_clock_offset - server_first_token_time = metrics_dict.pop( - "server_first_token_time", None) - if server_first_token_time is not None: - server_first_token_time += self.disagg_server_steady_clock_offset - metrics_json["timing_metrics"] = { - "server_arrival_time": - server_arrival_time, - "arrival_time": - timing_metrics.arrival_time.total_seconds() + - self.disagg_server_steady_clock_offset, - "first_scheduled_time": - timing_metrics.first_scheduled_time.total_seconds() + - self.disagg_server_steady_clock_offset, - "first_token_time": - timing_metrics.first_token_time.total_seconds() + - self.disagg_server_steady_clock_offset, - "server_first_token_time": - server_first_token_time, - "last_token_time": - timing_metrics.last_token_time.total_seconds() + - self.disagg_server_steady_clock_offset, - } - metrics_json["kv_cache_metrics"] = { - "num_total_allocated_blocks": - kv_cache_metrics.num_total_allocated_blocks, - "num_new_allocated_blocks": - kv_cache_metrics.num_new_allocated_blocks, - "num_reused_blocks": kv_cache_metrics.num_reused_blocks, - "num_missed_blocks": kv_cache_metrics.num_missed_blocks, - } - if timing_metrics.kv_cache_size > 0: - metrics_json["timing_metrics"].update({ - # TODO: move to kv_cache_metrics - "kv_cache_size": - timing_metrics.kv_cache_size, - "kv_cache_transfer_start": - timing_metrics.kv_cache_transfer_start.total_seconds() + - self.disagg_server_steady_clock_offset, - "kv_cache_transfer_end": - timing_metrics.kv_cache_transfer_end.total_seconds() + - self.disagg_server_steady_clock_offset, - }) - if speculative_decoding.total_draft_tokens > 0: - metrics_json["speculative_decoding"] = { - "acceptance_rate": speculative_decoding.acceptance_rate, - "total_accepted_draft_tokens": - speculative_decoding.total_accepted_draft_tokens, - "total_draft_tokens": - speculative_decoding.total_draft_tokens, - } - metrics_dict["perf_metrics"] = metrics_json - return JSONResponse(content=list(perf_metrics)) - async def get_kv_cache_events(self) -> JSONResponse: events = [] try: @@ -1332,6 +1274,17 @@ async def get_kv_cache_events(self) -> JSONResponse: async def _extract_metrics(self, res: RequestOutput, raw_request: Request): if not res.finished: return + if self._collect_perf_metrics: + if raw_request and not getattr(raw_request.state, + "server_first_token_time", None): + raw_request.state.server_first_token_time = ( + get_steady_clock_now_in_seconds()) + record = build_request_metrics_record( + res, + raw_request, + steady_clock_offset=self.disagg_server_steady_clock_offset) + if record is not None and raw_request is not None: + raw_request.state.perf_metrics_records.append(record) if self.metrics_collector: if res.candidate_metrics: for candidate_m in res.candidate_metrics: @@ -1345,31 +1298,6 @@ async def _extract_metrics(self, res: RequestOutput, raw_request: Request): # Wake up the stats collector to drain iteration stats if getattr(self.generator.args, "enable_iter_perf_stats", True): self._iteration_stats_wakeup_event.set() - if self.generator.args.return_perf_metrics: - output = res.outputs[0] - item = { - "request_id": res.request_id, - "perf_metrics": res.outputs[0].request_perf_metrics - } - if raw_request: - item["server_arrival_time"] = getattr(raw_request.state, - "server_arrival_time", - None) - if not getattr(raw_request.state, "server_first_token_time", - None): - raw_request.state.server_first_token_time = get_steady_clock_now_in_seconds( - ) - item[ - "server_first_token_time"] = raw_request.state.server_first_token_time - if output.disaggregated_params: - item[ - "ctx_request_id"] = output.disaggregated_params.ctx_request_id - # Request-level time breakdown (on GenerationResult/RequestOutput, not CompletionOutput) - if getattr(res, 'time_breakdown_metrics', None) is not None: - item["time_breakdown_metrics"] = res.time_breakdown_metrics - if self.perf_metrics is not None: - async with self.perf_metrics_lock: - self.perf_metrics.append(item) async def _create_chat_response( self, @@ -2004,9 +1932,6 @@ async def create_streaming_generator(promise: RequestOutput, try: if not self.postproc_worker_enabled: post_processor, args = postproc_params.post_processor, postproc_params.postproc_args - # Stamp first-token time on the first response, then append a - # /perf_metrics entry after [DONE]. The deque is only - # populated inside _extract_metrics. first_response = await anext(promise) raw_request.state.server_first_token_time = ( get_steady_clock_now_in_seconds()) @@ -2081,8 +2006,8 @@ async def create_streaming_generator(promise: RequestOutput, "thinking_token_budget is not supported by the Harmony " "GPT-OSS serving path; use reasoning_effort instead") sampling_params.detokenize = False # Harmony adapter handles detokenization - # Enable per-request perf metrics when the env var is set. - # Otherwise the /perf_metrics deque stays empty on this path. + # The server-level effective flag already enables engine metrics + # for header or JSONL output. if len(os.getenv("TRTLLM_KVCACHE_TIME_OUTPUT_PATH", "")) > 0: sampling_params.return_perf_metrics = True disaggregated_params = to_llm_disaggregated_params( @@ -2165,6 +2090,7 @@ async def create_response( tool_parser=args.tool_parser, ) + await self._extract_metrics(promise, raw_request) return response async def create_streaming_generator(promise: RequestOutput, @@ -2181,6 +2107,7 @@ async def create_streaming_generator(promise: RequestOutput, res, args) for pp_res in pp_results: yield pp_res + await self._extract_metrics(res, raw_request) try: if request.background: diff --git a/tensorrt_llm/serve/perf_metrics.py b/tensorrt_llm/serve/perf_metrics.py index 4eaed580ba8f..275ad005c3fd 100644 --- a/tensorrt_llm/serve/perf_metrics.py +++ b/tensorrt_llm/serve/perf_metrics.py @@ -12,12 +12,52 @@ # See the License for the specific language governing permissions and # limitations under the License. +"""Prometheus and per-request serving metrics. + +Non-streaming request opt-in and response example:: + + POST /v1/completions HTTP/1.1 + X-TRTLLM-return-metrics: 1 + + HTTP/1.1 200 OK + Content-Type: application/json + Server-Timing: server_queue;dur=1.250000, server_ttft;dur=8.500000, server_e2e;dur=24.000000 + X-TRTLLM-Start-End-Time: server-start;ts=12345.123456, server-end;ts=12345.147456 + X-TRTLLM-Step-Metrics: server-step-0-forward;dur=2.100000, server-step-0-sample;dur=0.400000 + X-TRTLLM-Ctx-Chunk-Metrics: server-ctx-chunk-0-forward;dur=4.200000 + +Streaming responses carry the same fields in a named SSE event after ``[DONE]``:: + + data: [DONE] + + event: trtllm.perf_metrics + data: { + "Server-Timing": "server_queue;dur=1.250000, server_ttft;dur=8.500000", + "X-TRTLLM-Start-End-Time": "server-start;ts=12345.123456, server-end;ts=12345.147456", + "X-TRTLLM-Step-Metrics": "server-step-0-forward;dur=2.100000", + "X-TRTLLM-Ctx-Chunk-Metrics": "server-ctx-chunk-0-forward;dur=4.200000", + } +""" + import asyncio -from collections import defaultdict, deque +import json +import math +import os +import socket from dataclasses import dataclass +from datetime import datetime, timezone +from pathlib import Path from typing import Any, Dict, List, Literal, Optional from tensorrt_llm.llmapi.disagg_utils import ServerRole +from tensorrt_llm.logger import logger +from tensorrt_llm.serve._perf_metrics_schema import ( + DisaggPerfMetricsRecord, + PerfMetrics, + PerfMetricsRecord, + WorkerPerfMetrics, + WorkerPerfMetricsRecord, +) COUNTER_METRICS = [ ("total_requests", "Total number of requests"), @@ -65,28 +105,10 @@ class MetricsDefinition: CLIENT_METRICS_DEFINITIONS = [ - MetricsDefinition("total_requests", "Total number of requests", "counter"), - MetricsDefinition("error_requests", "Total number of error requests", "counter"), - MetricsDefinition("retry_requests", "Total number of retry requests", "counter"), - MetricsDefinition("completed_requests", "Total number of completed requests", "counter"), - MetricsDefinition( - "first_token_latency_seconds", - "Histogram of latency from first token to completion in seconds", - "histogram", - SHORT_TIME_BUCKETS, - ), - MetricsDefinition( - "complete_latency_seconds", - "Histogram of latency from request arrival to last token in seconds", - "histogram", - LONG_TIME_BUCKETS, - ), - MetricsDefinition( - "per_token_latency_seconds", - "Histogram of latency from request arrival to completion in seconds", - "histogram", - SHORT_TIME_BUCKETS, - ), + MetricsDefinition(name, description, "counter") for name, description in COUNTER_METRICS +] + [ + MetricsDefinition(name, description, "histogram", buckets) + for name, description, buckets in HISTOGRAM_METRICS ] ROLE_TO_CLIENT_TYPE = { @@ -127,125 +149,576 @@ def __getattr__( return self._metrics[key] +SERVER_COUNTER_METRICS = ( + ("total_requests", "Total number of requests"), + ("stream_requests", "Total number of stream requests"), + ("nonstream_requests", "Total number of non-stream requests"), + ("validation_exceptions", "Total number of validation exceptions"), + ("http_exceptions", "Total number of HTTP exceptions"), + ("internal_errors", "Total number of internal errors"), + ("total_responses", "Total number of responses"), +) SERVER_METRICS_DEFINITIONS = [ - MetricsDefinition("total_requests", "Total number of requests", "counter"), - MetricsDefinition("stream_requests", "Total number of stream requests", "counter"), - MetricsDefinition("nonstream_requests", "Total number of non-stream requests", "counter"), - MetricsDefinition("validation_exceptions", "Total number of validation exceptions", "counter"), - MetricsDefinition("http_exceptions", "Total number of HTTP exceptions", "counter"), - MetricsDefinition("internal_errors", "Total number of internal errors", "counter"), - MetricsDefinition("total_responses", "Total number of responses", "counter"), + MetricsDefinition(name, description, "counter") for name, description in SERVER_COUNTER_METRICS +] + [ MetricsDefinition( "queue_latency_seconds", "Histogram of latency from request arrival to being processed in seconds", "histogram", SHORT_TIME_BUCKETS, - ), + ) ] class DisaggPerfMetricsCollector: - def __init__(self, max_requests: int): - self._max_requests = max_requests - self._request_meteics = deque(maxlen=max_requests) - self._server_metrics = defaultdict(dict) - self._lock = asyncio.Lock() - self._collect_lock = asyncio.Lock() - self._clients = [] - self._background_tasks: set[asyncio.Task] = set() + """Prometheus metrics owned by one disaggregated HTTP server process.""" + + def __init__(self, max_requests: int = 0): + # Kept for compatibility; per-request retention now belongs to JSONL. + del max_requests self._metrics = { definition.name: instance_metric(definition) for definition in SERVER_METRICS_DEFINITIONS } - def add_client(self, client): - self._clients.append(client) - def __getattr__(self, key: str): return self._metrics[key] - async def add_per_request_metrics( - self, - ctx_server: str, - gen_server: str, - ctx_request_id: int, - server_arrival_time: float, - server_first_token_time: float, - ctx_dispatch_time: float = 0, - ): - async with self._lock: - self._request_meteics.append( + +SERVER_TIMING_HEADER = "Server-Timing" +START_END_TIME_HEADER = "X-TRTLLM-Start-End-Time" +STEP_METRICS_HEADER = "X-TRTLLM-Step-Metrics" +CTX_CHUNK_METRICS_HEADER = "X-TRTLLM-Ctx-Chunk-Metrics" +SSE_METRICS_EVENT = "trtllm.perf_metrics" +RETURN_METRICS_HEADER = "X-TRTLLM-return-metrics" +_RETURN_METRICS_HEADER_BYTES = RETURN_METRICS_HEADER.lower().encode() + +_SCHEMA_VERSION = 1 +_PERF_METRICS_HEADER_BUDGET_BYTES = 80 * 1024 +_WRITER_QUEUE_SIZE = 1024 +_WRITER_BATCH_SIZE = 64 +_WRITER_SHUTDOWN_TIMEOUT_SECONDS = 5 + + +_TIMING_FIELDS = ( + "arrival_time", + "first_scheduled_time", + "first_token_time", + "last_token_time", + "kv_cache_transfer_start", + "kv_cache_transfer_end", +) +_KV_FIELDS = ( + "num_total_allocated_blocks", + "num_new_allocated_blocks", + "num_reused_blocks", + "num_missed_blocks", + "kv_cache_hit_rate", +) +_SPEC_FIELDS = ("acceptance_rate", "total_accepted_draft_tokens", "total_draft_tokens") + + +def _as_seconds(value: Any, offset: float = 0) -> Optional[float]: + try: + seconds = float(value.total_seconds()) + except (AttributeError, TypeError, ValueError): + return None + return seconds + offset if seconds > 0 else None + + +def build_request_metrics_record( + result: Any, + raw_request: Any = None, + phase: str = "server", + steady_clock_offset: float = 0, +) -> Optional[Dict[str, Any]]: + """Convert a completed RequestOutput metrics snapshot to JSON-safe data.""" + if not result or not getattr(result, "outputs", None): + return None + output = result.outputs[0] + metrics = getattr(output, "request_perf_metrics", None) + if metrics is None: + return None + + timing = metrics.timing_metrics + timing_metrics = { + name: _as_seconds(getattr(timing, name), steady_clock_offset) for name in _TIMING_FIELDS + } + timing_metrics["kv_cache_size"] = timing.kv_cache_size + if raw_request is not None: + for name in ("server_arrival_time", "server_first_token_time"): + value = getattr(raw_request.state, name, None) + timing_metrics[name] = value + steady_clock_offset if value is not None else None + + phase_record: Dict[str, Any] = { + "first_iter": metrics.first_iter, + "last_iter": metrics.last_iter, + "timing_metrics": timing_metrics, + "kv_cache_metrics": {name: getattr(metrics.kv_cache_metrics, name) for name in _KV_FIELDS}, + } + speculative = metrics.speculative_decoding + if speculative.total_draft_tokens > 0: + phase_record["speculative_decoding"] = { + name: getattr(speculative, name) for name in _SPEC_FIELDS + } + if getattr(result, "time_breakdown_metrics", None) is not None: + phase_record["time_breakdown_metrics"] = result.time_breakdown_metrics + + record: Dict[str, Any] = { + "schema_version": _SCHEMA_VERSION, + "request_id": str(getattr(result, "request_id", "")), + "status": "complete", + "phases": {phase: phase_record}, + } + disagg = getattr(output, "disaggregated_params", None) + if disagg: + for name in ("ctx_request_id", "disagg_request_id"): + value = getattr(disagg, name, None) + if value is not None: + record[name] = value + return record + + +def _elapsed_ms(values: Dict[str, Any], start: str, end: Optional[str] = None) -> Optional[float]: + try: + value = float(values[start]) + if end is not None: + value = (float(values[end]) - value) * 1000 + except (KeyError, TypeError, ValueError): + return None + return value if value >= 0 and math.isfinite(value) else None + + +def build_metrics_headers(records: List[Dict[str, Any]]) -> Dict[str, str]: + """Format completed records as Server-Timing-style metric lists.""" + values = { + SERVER_TIMING_HEADER: [], + START_END_TIME_HEADER: [], + STEP_METRICS_HEADER: [], + CTX_CHUNK_METRICS_HEADER: [], + } + for record in records: + for phase, phase_record in record.get("phases", {}).items(): + timing = phase_record.get("timing_metrics", {}) + for name, field in ( + ("start", "arrival_time"), + ("end", "last_token_time"), + ): + timestamp = timing.get(field) + if timestamp is not None: + values[START_END_TIME_HEADER].append( + f"{phase}-{name};ts={float(timestamp):.9f}" + ) + timing_ranges = ( ( - ctx_server, - gen_server, - ctx_request_id, - server_arrival_time, - server_first_token_time, - ctx_dispatch_time, + ("queue", "server_arrival_time", "ctx_dispatch_time"), + ("ttft", "server_arrival_time", "server_first_token_time"), + ) + if phase == "disagg" + else ( + ("queue", "arrival_time", "first_scheduled_time"), + ("ttft", "arrival_time", "first_token_time"), + ("e2e", "arrival_time", "last_token_time"), + ("kv_transfer", "kv_cache_transfer_start", "kv_cache_transfer_end"), ) ) + for name, start, end in timing_ranges: + duration = _elapsed_ms(timing, start, end) + if duration is not None: + values[SERVER_TIMING_HEADER].append(f"{phase}_{name};dur={duration:.6f}") - async def get_perf_metrics(self) -> List[Dict[str, Any]]: - async with self._collect_lock: - perf_metrics = {} - for client in self._clients: - metrics_dict = await client.collect_metrics() - perf_metrics.update(metrics_dict) - - return_metrics = [] - async with self._lock: - for server, metrics_data in perf_metrics.items(): - server_metrics = self._server_metrics[server] - # avoid metrics map inflation by limiting the number of requests to add - available_req_num = min( - max(0, self._max_requests - len(server_metrics)), - len(metrics_data), + breakdown = phase_record.get("time_breakdown_metrics") or {} + for header, key, label in ( + (STEP_METRICS_HEADER, "step_metrics", "step"), + (CTX_CHUNK_METRICS_HEADER, "ctx_chunk_metrics", "ctx-chunk"), + ): + for index, metrics in enumerate(breakdown.get(key, [])): + item = metrics.get("iter", index) if key == "step_metrics" else index + durations = ( + ("forward", _elapsed_ms(metrics, "forward_start_time", "forward_end_time")), + ("sample", _elapsed_ms(metrics, "sample_start_time", "sample_end_time")), + ("gpu-forward", _elapsed_ms(metrics, "gpu_forward_time")), + ("gpu-sample", _elapsed_ms(metrics, "gpu_sample_time")), ) - req_metrics_map = { - req_metrics["ctx_request_id"]: req_metrics - for req_metrics in metrics_data[:available_req_num] - if "ctx_request_id" in req_metrics - } - server_metrics.update(req_metrics_map) - - remain_keys = [] - for ( - ctx_server, - gen_server, - ctx_request_id, - server_arrival_time, - server_first_token_time, - ctx_dispatch_time, - ) in self._request_meteics: - gen_perf_metrics = self._server_metrics[gen_server].pop(ctx_request_id, None) - if gen_perf_metrics is None: - # generation not finished - remain_keys.append( - ( - ctx_server, - gen_server, - ctx_request_id, - server_arrival_time, - server_first_token_time, - ctx_dispatch_time, - ) - ) - continue - ctx_perf_metrics = self._server_metrics[ctx_server].pop(ctx_request_id, None) - # TODO: strip the keys for less repeating and use table style response - return_metrics.append( - { - "ctx_server": ctx_server, - "gen_server": gen_server, - "disagg_server_arrival_time": server_arrival_time, - # arrival->ctx_dispatch = pre-ctx wait in the - # orchestrator/fleet (accept queue + event loop + - # pipeline), the dominant TTFT term under fleet load. - "disagg_ctx_dispatch_time": ctx_dispatch_time, - "disagg_server_first_token_time": server_first_token_time, - "ctx_perf_metrics": ctx_perf_metrics, - "gen_perf_metrics": gen_perf_metrics, - } + values[header].extend( + f"{phase}-{label}-{item}-{name};dur={duration:.6f}" + for name, duration in durations + if duration is not None ) - self._request_meteics = deque(remain_keys, maxlen=self._max_requests) - return return_metrics + + inherited = phase_record.get("metrics_headers") or {} + for header in (STEP_METRICS_HEADER, CTX_CHUNK_METRICS_HEADER): + if inherited.get(header): + values[header].append(inherited[header]) + + return {header: ", ".join(items) for header, items in values.items() if items} + + +def build_metrics_record_from_headers( + headers: Any, + phase: str, + request_id: str = "", +) -> Optional[Dict[str, Any]]: + """Build a request-local phase from standard metrics fields.""" + metrics_headers = {} + for header_name in ( + SERVER_TIMING_HEADER, + START_END_TIME_HEADER, + STEP_METRICS_HEADER, + CTX_CHUNK_METRICS_HEADER, + ): + value = headers.get(header_name) + if value: + metrics_headers[header_name] = value.replace("server_", f"{phase}_").replace( + "server-", f"{phase}-" + ) + if not metrics_headers: + return None + timing_metrics = {} + fields = { + f"{phase}-start": "arrival_time", + f"{phase}-end": "last_token_time", + } + for item in metrics_headers.get(START_END_TIME_HEADER, "").split(","): + name, separator, timestamp = item.strip().partition(";ts=") + if separator and name in fields: + try: + timing_metrics[fields[name]] = float(timestamp) + except ValueError: + logger.warning("Ignoring invalid %s timestamp: %s", name, timestamp) + + durations = {} + prefix = f"{phase}_" + for item in metrics_headers.get(SERVER_TIMING_HEADER, "").split(","): + name, separator, duration = item.strip().partition(";dur=") + if separator and name.startswith(prefix): + try: + durations[name[len(prefix) :]] = float(duration) / 1000 + except ValueError: + logger.warning("Ignoring invalid %s duration: %s", name, duration) + + arrival_time = timing_metrics.get("arrival_time") + if arrival_time is not None: + for name, field in ( + ("queue", "first_scheduled_time"), + ("ttft", "first_token_time"), + ("e2e", "last_token_time"), + ): + if name in durations: + timing_metrics.setdefault(field, arrival_time + durations[name]) + + phase_record = {"metrics_headers": metrics_headers} + if timing_metrics: + phase_record["timing_metrics"] = timing_metrics + return { + "schema_version": _SCHEMA_VERSION, + "request_id": request_id, + "status": "complete", + "metrics_headers": metrics_headers, + "phases": {phase: phase_record}, + } + + +def _limit_metrics_headers(headers: Dict[str, str]) -> Dict[str, str]: + def size(values: Dict[str, str]) -> int: + return sum(len(name.encode()) + len(value.encode()) + 4 for name, value in values.items()) + + if size(headers) <= _PERF_METRICS_HEADER_BUDGET_BYTES: + return headers + logger.warning( + "Performance metrics payload exceeds %d bytes; omitting step and context-chunk metrics", + _PERF_METRICS_HEADER_BUDGET_BYTES, + ) + return { + name: value + for name, value in headers.items() + if name in (SERVER_TIMING_HEADER, START_END_TIME_HEADER) + } + + +def combine_disagg_metrics( + request_id: str, + disagg_phase: Dict[str, Any], + ctx_record: Optional[Dict[str, Any]], + gen_record: Optional[Dict[str, Any]], + disagg_request_id: Optional[int] = None, +) -> Dict[str, Any]: + phases: Dict[str, Any] = {"disagg": disagg_phase} + + def add_worker_phase(name: str, record: Dict[str, Any]) -> None: + phase_record = dict(next(iter(record.get("phases", {}).values()), {})) + phase_record["request_id"] = record.get("request_id") + if record.get("ctx_request_id") is not None: + phase_record["ctx_request_id"] = record["ctx_request_id"] + if record.get("metrics_headers"): + phase_record["metrics_headers"] = record["metrics_headers"] + phases[name] = phase_record + + if ctx_record: + add_worker_phase("ctx", ctx_record) + if gen_record: + add_worker_phase("gen", gen_record) + combined = { + "schema_version": _SCHEMA_VERSION, + "request_id": request_id, + "status": "complete", + "phases": phases, + } + if disagg_request_id is not None: + combined["disagg_request_id"] = disagg_request_id + return combined + + +def _jsonl_perf_metrics(phase_record: Dict[str, Any]) -> PerfMetrics: + perf_metrics: PerfMetrics = { + "timing_metrics": dict(phase_record.get("timing_metrics", {})), + } + if "first_iter" in phase_record: + perf_metrics["first_iter"] = phase_record["first_iter"] + if "last_iter" in phase_record: + perf_metrics["last_iter"] = phase_record["last_iter"] + if "kv_cache_metrics" in phase_record: + perf_metrics["kv_cache_metrics"] = phase_record["kv_cache_metrics"] + if "speculative_decoding" in phase_record: + perf_metrics["speculative_decoding"] = phase_record["speculative_decoding"] + + timing_metrics = dict(perf_metrics.get("timing_metrics", {})) + if not timing_metrics.get("kv_cache_size"): + for name in ("kv_cache_size", "kv_cache_transfer_start", "kv_cache_transfer_end"): + timing_metrics.pop(name, None) + perf_metrics["timing_metrics"] = timing_metrics + + kv_cache_metrics = dict(perf_metrics.get("kv_cache_metrics", {})) + kv_cache_metrics.pop("kv_cache_hit_rate", None) + if kv_cache_metrics: + perf_metrics["kv_cache_metrics"] = kv_cache_metrics + return perf_metrics + + +def _jsonl_worker_metrics( + record: Dict[str, Any], phase_record: Dict[str, Any] +) -> WorkerPerfMetrics: + request_id = record["request_id"] + try: + request_id = int(request_id) + except (TypeError, ValueError): + pass + worker_metrics: WorkerPerfMetrics = { + "request_id": request_id, + "perf_metrics": _jsonl_perf_metrics(phase_record), + } + if record.get("ctx_request_id") is not None: + worker_metrics["ctx_request_id"] = record["ctx_request_id"] + if phase_record.get("time_breakdown_metrics") is not None: + worker_metrics["time_breakdown_metrics"] = phase_record["time_breakdown_metrics"] + return worker_metrics + + +def _jsonl_record(record: Dict[str, Any]) -> PerfMetricsRecord: + phases = record.get("phases", {}) + if "disagg" not in phases: + worker_metrics = _jsonl_worker_metrics(record, phases["server"]) + jsonl_record: WorkerPerfMetricsRecord = { + **worker_metrics, + "status": record.get("status", "complete"), + } + if record.get("disagg_request_id") is not None: + jsonl_record["disagg_request_id"] = record["disagg_request_id"] + return jsonl_record + + disagg_phase = phases["disagg"] + disagg_timing = disagg_phase["timing_metrics"] + disagg_record: DisaggPerfMetricsRecord = { + "ctx_server": disagg_phase["ctx_server"], + "gen_server": disagg_phase["gen_server"], + "disagg_server_arrival_time": disagg_timing["server_arrival_time"], + "disagg_ctx_dispatch_time": disagg_timing["ctx_dispatch_time"], + "disagg_server_first_token_time": disagg_timing["server_first_token_time"], + "status": record.get("status", "complete"), + } + for phase, field in ( + ("ctx", "ctx_perf_metrics"), + ("gen", "gen_perf_metrics"), + ): + phase_record = phases.get(phase) + if phase_record: + worker_record = { + "request_id": phase_record.get("request_id", record["request_id"]), + "ctx_request_id": phase_record.get("ctx_request_id"), + } + worker_metrics = _jsonl_worker_metrics(worker_record, phase_record) + if field == "ctx_perf_metrics": + disagg_record["ctx_perf_metrics"] = worker_metrics + else: + disagg_record["gen_perf_metrics"] = worker_metrics + if record.get("disagg_request_id") is not None: + disagg_record["disagg_request_id"] = record["disagg_request_id"] + return disagg_record + + +class PerfMetricsJsonlWriter: + """Best-effort bounded JSONL writer shared by both serving apps.""" + + def __init__(self, output_dir: Optional[str], server_kind: str): + self._output_dir = Path(output_dir) if output_dir else None + self._server_kind = server_kind + self._queue: asyncio.Queue = asyncio.Queue(maxsize=_WRITER_QUEUE_SIZE) + self._task: Optional[asyncio.Task] = None + self._path: Optional[Path] = None + self.dropped_records = 0 + self._write_error_count = 0 + + async def start(self) -> None: + if self._output_dir is None or self._task is not None: + return + try: + self._output_dir.mkdir(parents=True, exist_ok=True) + timestamp = datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ") + filename = ( + f"perf_metrics-{self._server_kind}-{socket.gethostname()}-" + f"{os.getpid()}-{timestamp}.jsonl" + ) + self._path = self._output_dir / filename + self._task = asyncio.create_task(self._run()) + except OSError as error: + logger.error("Disabling performance metrics JSONL output: %s", error) + self._output_dir = None + + def submit(self, record: Dict[str, Any]) -> None: + if self._task is None: + return + try: + self._queue.put_nowait(record) + except asyncio.QueueFull: + self.dropped_records += 1 + if self.dropped_records == 1 or self.dropped_records % 1000 == 0: + logger.warning("Dropped %d performance metrics records", self.dropped_records) + + async def close(self) -> None: + if self._task is None: + return + task = self._task + try: + await asyncio.wait_for( + self._queue.put(None), + timeout=_WRITER_SHUTDOWN_TIMEOUT_SECONDS, + ) + await asyncio.wait_for(task, timeout=_WRITER_SHUTDOWN_TIMEOUT_SECONDS) + except asyncio.TimeoutError: + logger.warning( + "Timed out flushing performance metrics JSONL; dropping remaining records" + ) + task.cancel() + await asyncio.gather(task, return_exceptions=True) + finally: + self._task = None + + async def _run(self) -> None: + stop = False + while not stop: + item = await self._queue.get() + if item is None: + return + records = [item] + while len(records) < _WRITER_BATCH_SIZE: + try: + item = self._queue.get_nowait() + except asyncio.QueueEmpty: + break + if item is None: + stop = True + break + records.append(item) + serialized = [] + for record in records: + try: + item = _jsonl_record(record) + serialized.append( + json.dumps(item, separators=(",", ":"), allow_nan=False) + "\n" + ) + except (KeyError, TypeError, ValueError) as error: + self.dropped_records += 1 + if self.dropped_records == 1 or self.dropped_records % 1000 == 0: + logger.warning("Dropped malformed performance metrics record: %s", error) + if not serialized: + continue + try: + data = "".join(serialized) + await asyncio.to_thread(self._write, data) + except OSError as error: + self.dropped_records += len(serialized) + self._write_error_count += 1 + if self._write_error_count == 1: + logger.warning("Failed to write performance metrics JSONL: %s", error) + + def _write(self, data: str) -> None: + if self._path is not None: + with self._path.open("a", encoding="utf-8") as output: + output.write(data) + + +def build_metrics_sse_event(headers: Dict[str, str]) -> bytes: + headers = _limit_metrics_headers(headers) + if not headers: + return b"" + payload = json.dumps(headers, separators=(",", ":")) + return (f"event: {SSE_METRICS_EVENT}\ndata: {payload}\n\n").encode() + + +class PerfMetricsMiddleware: + """Expose request metrics and optionally persist completed records.""" + + def __init__( + self, app: Any, expose_headers: bool, writer: Optional[PerfMetricsJsonlWriter] = None + ): + self._app = app + self._expose_headers = expose_headers + self._writer = writer + + async def __call__(self, scope: Dict[str, Any], receive: Any, send: Any) -> None: + if scope["type"] != "http": + await self._app(scope, receive, send) + return + is_stream = False + metrics_headers = None + records = scope.setdefault("state", {}).setdefault("perf_metrics_records", []) + return_metrics = self._expose_headers and any( + name.lower() == _RETURN_METRICS_HEADER_BYTES and value.strip() == b"1" + for name, value in scope.get("headers", []) + ) + + async def send_metrics(message: Dict[str, Any]) -> None: + nonlocal is_stream, metrics_headers + if message["type"] == "http.response.start": + headers = list(message.get("headers", [])) + is_stream = any( + key.lower() == b"content-type" and b"text/event-stream" in value.lower() + for key, value in headers + ) + if records and return_metrics and not is_stream: + metrics_headers = build_metrics_headers(records) + public_headers = _limit_metrics_headers(metrics_headers) + headers.extend( + (name.encode(), value.encode()) for name, value in public_headers.items() + ) + message["headers"] = headers + + elif message["type"] == "http.response.body" and not message.get("more_body", False): + if records: + if metrics_headers is None: + metrics_headers = build_metrics_headers(records) + if self._writer is not None: + for record in records: + self._writer.submit(record) + if return_metrics and is_stream: + message["body"] = message.get("body", b"") + build_metrics_sse_event( + metrics_headers + ) + try: + await send(message) + except OSError: + pass + return + await send(message) + + await self._app(scope, receive, send_metrics) diff --git a/tensorrt_llm/serve/responses_utils.py b/tensorrt_llm/serve/responses_utils.py index bac9c7615230..edd36f8624f9 100644 --- a/tensorrt_llm/serve/responses_utils.py +++ b/tensorrt_llm/serve/responses_utils.py @@ -2036,6 +2036,9 @@ class ResponseHooks(ABC): def on_req_begin(self, request: UCompletionRequest): pass + def on_disagg_request_id(self, disagg_request_id: int): + """Receive the request ID immediately after the service allocates it.""" + def on_ctx_dispatch(self, request: UCompletionRequest): """Record when the disaggregated service starts context placement. @@ -2043,6 +2046,9 @@ def on_ctx_dispatch(self, request: UCompletionRequest): or fleet. The default is a no-op for non-instrumented implementations. """ + def on_perf_metrics(self, server: str, role: str, metrics: dict): + """Receive request-local metrics carried by an upstream response.""" + @abstractmethod def on_ctx_resp(self, ctx_server: str, response: UCompletionResponse): pass diff --git a/tensorrt_llm/serve/scripts/backend_request_func.py b/tensorrt_llm/serve/scripts/backend_request_func.py index e2496a6754d1..9f01abe94917 100644 --- a/tensorrt_llm/serve/scripts/backend_request_func.py +++ b/tensorrt_llm/serve/scripts/backend_request_func.py @@ -35,7 +35,7 @@ async def _iter_sse_data(response_content): continue payload = line.removeprefix("data:").lstrip() if payload == "[DONE]": - continue + return yield payload diff --git a/tensorrt_llm/serve/scripts/benchmark_serving.py b/tensorrt_llm/serve/scripts/benchmark_serving.py index 0374f96aab20..f3ddffd21f6b 100644 --- a/tensorrt_llm/serve/scripts/benchmark_serving.py +++ b/tensorrt_llm/serve/scripts/benchmark_serving.py @@ -28,6 +28,7 @@ from collections.abc import AsyncGenerator, Iterable from dataclasses import dataclass from datetime import datetime +from pathlib import Path from typing import Any, Optional import aiohttp @@ -736,32 +737,59 @@ async def fetch_energy_metrics(base_url: str) -> Optional[dict]: return None -async def fetch_perf_metrics(base_url: str) -> dict: - """ - Fetch performance metrics from the /perf_metrics endpoint. - - Args: - base_url: The base URL of the server +def _snapshot_perf_metrics(output_dir: str) -> dict[Path, int]: + directory = Path(output_dir) + if not directory.exists(): + return {} + if not directory.is_dir(): + raise ValueError( + f"Performance metrics output path is not a directory: {output_dir}") + return { + path: path.stat().st_size + for path in directory.glob("perf_metrics-*.jsonl") + } - Returns: - Dictionary containing the performance metrics - """ - perf_url = f"{base_url}/perf_metrics" - async with aiohttp.ClientSession(trust_env=True, - timeout=AIOHTTP_TIMEOUT) as session: - try: - async with session.get(perf_url) as response: - if response.status == 200: - return await response.json() - else: - print( - f"Failed to fetch performance metrics. Status: {response.status}" - ) - return {} - except Exception as e: - print(f"Error fetching performance metrics: {e}") - return {} +def _perf_metrics_files(output_dir: str, offsets: dict[Path, + int]) -> list[Path]: + paths = sorted(Path(output_dir).glob("perf_metrics-*.jsonl")) + by_kind = {} + for path in paths: + if path.stat().st_size <= offsets.get(path, 0): + continue + kind = path.name.removeprefix("perf_metrics-").split("-", 1)[0] + by_kind.setdefault(kind, []).append(path) + if "disagg" in by_kind: + return by_kind["disagg"] + if "server" in by_kind: + return by_kind["server"] + return [] + + +def _read_new_perf_metrics( + output_dir: str, + offsets: dict[Path, int], + expected_count: int, + timeout: float = 10, +) -> list[dict]: + deadline = time.monotonic() + timeout + records = [] + while time.monotonic() < deadline: + records = [] + for path in _perf_metrics_files(output_dir, offsets): + with path.open("r", encoding="utf-8") as metrics_file: + metrics_file.seek(offsets.get(path, 0)) + for line in metrics_file: + if not line.strip(): + continue + try: + records.append(json.loads(line)) + except json.JSONDecodeError: + continue + if len(records) >= expected_count: + return records + time.sleep(0.1) + return records def main(args: argparse.Namespace): @@ -965,6 +993,10 @@ def create_dataset_and_sample(dataset_name: str): # Avoid GC - reduce pause times. gc.disable() + perf_metrics_output_dir = getattr(args, 'save_request_time_breakdown', None) + perf_metrics_offsets = (_snapshot_perf_metrics(perf_metrics_output_dir) + if perf_metrics_output_dir else {}) + benchmark_result = asyncio.run( benchmark( backend=backend, @@ -1045,54 +1077,42 @@ def create_dataset_and_sample(dataset_name: str): json.dump(result_json, outfile) save_to_pytorch_benchmark_format(args, result_json, file_name) - # Save per-request breakdown if requested - if args.save_request_time_breakdown: - print("Fetching request performance metrics...") - perf_metrics = asyncio.run(fetch_perf_metrics(base_url)) - - if perf_metrics: - # Generate filename for perf metrics - current_dt = datetime.now().strftime("%Y%m%d-%H%M%S") - base_model_id = model_id.split("/")[-1] - max_concurrency_str = (f"-concurrency{args.max_concurrency}" - if args.max_concurrency is not None else "") - perf_filename = f"{backend}-{args.request_rate}qps{max_concurrency_str}-{base_model_id}-{current_dt}-perf_metrics.json" - - if args.result_dir: - perf_filename = os.path.join(args.result_dir, perf_filename) - - # Save perf metrics to JSON file - with open(perf_filename, "w", encoding='utf-8') as outfile: - try: - json.dump(perf_metrics, outfile, indent=2) - except Exception as e: - print(f"Failed to save perf metrics: {e}") - - print(f"Request performance metrics saved to: {perf_filename}") - - # Create timing diagram from the saved JSON file - try: - analyzer = RequestTimeBreakdown() - - print("Creating time diagram from request time breakdown...") - timing_data = analyzer.parse_json_file(perf_filename) - - if timing_data: - # Generate HTML filename for the timing diagram - diagram_filename = f"{os.path.splitext(perf_filename)[0]}-time_diagram.html" - analyzer.create_timing_diagram(timing_data, - diagram_filename) - - print(f"Time diagram saved to: {diagram_filename}") - else: - print( - "No time data found in request time breakdown - skipping diagram creation." - ) - except Exception as e: - print(f"Failed to create time diagram: {e}") - print("Performance metrics were still saved successfully.") + if perf_metrics_output_dir: + expected_count = benchmark_result["completed"] + int( + not args.no_test_input) + perf_metrics = _read_new_perf_metrics(perf_metrics_output_dir, + perf_metrics_offsets, + expected_count) + if not perf_metrics: + print("No new public-server performance metrics found; " + "skipping time breakdown report.") + return + if len(perf_metrics) < expected_count: + print(f"Warning: found {len(perf_metrics)} of " + f"{expected_count} expected performance metrics records.") + + current_dt = datetime.now().strftime("%Y%m%d-%H%M%S") + base_model_id = model_id.split("/")[-1] + max_concurrency_str = (f"-concurrency{args.max_concurrency}" + if args.max_concurrency is not None else "") + output_stem = (f"{backend}-{args.request_rate}qps{max_concurrency_str}-" + f"{base_model_id}-{current_dt}-perf_metrics") + if args.result_dir: + output_stem = os.path.join(args.result_dir, output_stem) + perf_filename = f"{output_stem}.jsonl" + with open(perf_filename, "w", encoding="utf-8") as outfile: + for record in perf_metrics: + outfile.write(json.dumps(record, separators=(",", ":")) + "\n") + print(f"Request performance metrics saved to: {perf_filename}") + + analyzer = RequestTimeBreakdown() + timing_data = analyzer.parse_json_file(perf_filename) + if timing_data: + diagram_filename = f"{output_stem}-time_diagram.html" + analyzer.create_timing_diagram(timing_data, diagram_filename) + print(f"Time diagram saved to: {diagram_filename}") else: - print("Failed to fetch per-request performance metrics.") + print("No time data found; skipping time breakdown diagram.") if __name__ == "__main__": @@ -1487,9 +1507,14 @@ def create_dataset_and_sample(dataset_name: str): parser.add_argument( "--save-request-time-breakdown", - action="store_true", - help= - "After benchmarking, call the /perf_metric endpoint, save the result as JSON, and create an interactive time breakdown diagram.", + nargs="?", + const=".", + default=None, + metavar="PERF_METRICS_OUTPUT_DIR", + help=("Read JSONL records dumped by the server's " + "perf_metrics_output_dir, save the benchmark records, and " + "create an interactive time breakdown diagram. If no directory " + "is provided, use the current directory."), ) args = parser.parse_args() diff --git a/tensorrt_llm/serve/scripts/time_breakdown/README.md b/tensorrt_llm/serve/scripts/time_breakdown/README.md index cdee49212d50..d6b6ecc60eb4 100644 --- a/tensorrt_llm/serve/scripts/time_breakdown/README.md +++ b/tensorrt_llm/serve/scripts/time_breakdown/README.md @@ -376,49 +376,33 @@ The tool expects a JSON file containing an array of request performance metrics ## Usage -### Integration with Benchmark Serving +### Server-side JSONL workflow + +Configure the server-side writer: -Step 1: Set in `extra-llm-api-config.yaml`: ```yaml -return_perf_metrics: True -perf_metrics_max_requests: +perf_metrics_output_dir: ``` -If running disaggregated serving, add configs for all servers (disagg, context and generation server). -Step 2: Add `--save-request-time-breakdown` when running `benchmark_serving.py`: -```bash -python -m tensorrt_llm.serve.scripts.benchmark_serving \ - --model ${model_name} \ - --dataset-name random \ - --ignore-eos \ - --num-prompts 1000 \ - --random-input-len 1024 \ - --random-output-len 2048 \ - --random-ids \ - --max-concurrency 64 \ - --save-result \ - --result-dir \ - --percentile-metrics "ttft,tpot,itl,e2e" \ - --save-request-time-breakdown -``` +Run the workload, then copy the generated `perf_metrics-*.jsonl` file from the server. For disaggregated serving, use the disagg server file; it contains the combined disagg, context, and generation phases. ### As a CLI Tool ```bash # Basic usage -python time_breakdown.py perf_metrics.json +python time_breakdown.py perf_metrics-disagg.jsonl # Specify output file -python time_breakdown.py perf_metrics.json -o my_time_diagram.html +python time_breakdown.py perf_metrics-disagg.jsonl -o my_time_diagram.html # Limit max requests and sort by E2E latency -python time_breakdown.py perf_metrics.json --max-requests 100 --sort-by e2e +python time_breakdown.py perf_metrics-disagg.jsonl --max-requests 100 --sort-by e2e # Show statistics only -python time_breakdown.py perf_metrics.json --stats-only +python time_breakdown.py perf_metrics-disagg.jsonl --stats-only # Create diagram and show statistics -python time_breakdown.py perf_metrics.json --show-stats +python time_breakdown.py perf_metrics-disagg.jsonl --show-stats ``` ### CLI Options diff --git a/tensorrt_llm/serve/scripts/time_breakdown/time_breakdown.py b/tensorrt_llm/serve/scripts/time_breakdown/time_breakdown.py index 19ad6767d6fd..19c8fe3fd978 100644 --- a/tensorrt_llm/serve/scripts/time_breakdown/time_breakdown.py +++ b/tensorrt_llm/serve/scripts/time_breakdown/time_breakdown.py @@ -14,12 +14,12 @@ - Hover to show individual segment details Usage as CLI: - python time_breakdown.py [options] + python time_breakdown.py [options] Usage as library: from time_breakdown import RequestTimeBreakdown analyzer = RequestTimeBreakdown() - timing_data = analyzer.parse_json_file("perf_metrics.json") + timing_data = analyzer.parse_json_file("perf_metrics.jsonl") analyzer.create_timing_diagram(timing_data, "output.html") """ @@ -33,6 +33,8 @@ import numpy as np import plotly.graph_objects as go +from tensorrt_llm.serve._perf_metrics_schema import PerfMetricsRecord + @dataclass class TimingMetric: @@ -181,7 +183,7 @@ def get_metric_by_name(self, name: str) -> Optional[TimingMetric]: class RequestDataParser: """Parser for disaggregated format with ctx_perf_metrics and gen_perf_metrics.""" - def parse_request(self, request_data: Dict, + def parse_request(self, request_data: PerfMetricsRecord, request_index: int) -> Dict[str, Any]: # Check if both ctx_perf_metrics and gen_perf_metrics exist and are not None ctx_perf = request_data.get('ctx_perf_metrics') @@ -207,15 +209,18 @@ def parse_request(self, request_data: Dict, float('nan')) ctx_first_token_time = ctx_metrics.get('first_token_time', float('nan')) ctx_server_arrival_time = ctx_metrics.get('server_arrival_time', - float('nan')) - ctx_server_first_token_time = ctx_metrics.get('server_first_token_time', - float('nan')) + ctx_arrival_time) + ctx_server_first_token_time = ctx_metrics.get( + 'server_first_token_time', + ctx_metrics.get('last_token_time', float('nan'))) # Generation timing - gen_server_first_token_time = gen_metrics.get('server_first_token_time', - float('nan')) - gen_server_arrival_time = gen_metrics.get('server_arrival_time', - float('nan')) + gen_server_first_token_time = gen_metrics.get( + 'server_first_token_time', + gen_metrics.get('last_token_time', float('nan'))) + gen_server_arrival_time = gen_metrics.get( + 'server_arrival_time', gen_metrics.get('arrival_time', + float('nan'))) gen_arrival_time = gen_metrics.get('arrival_time', float('nan')) gen_first_token_time = gen_metrics.get('first_token_time', float('nan')) gen_first_scheduled_time = gen_metrics.get('first_scheduled_time', @@ -240,15 +245,13 @@ def parse_request(self, request_data: Dict, else: request_id = request_data.get('request_id', request_index) - # Time breakdown metrics - check new unified structure first, then fall back to legacy + # Time breakdown metrics step_metrics = None ctx_gpu_forward_time = None ctx_gpu_sample_time = None ctx_chunk_metrics = None - # Try new unified time_breakdown_metrics structure if is_disaggregated: - # time_breakdown_metrics is at gen_perf_metrics top level, not inside perf_metrics time_breakdown = (gen_perf or {}).get('time_breakdown_metrics') if time_breakdown: step_metrics = time_breakdown.get('step_metrics') @@ -260,7 +263,6 @@ def parse_request(self, request_data: Dict, # Legacy: step_metrics inside perf_metrics gen_perf_data = (gen_perf or {}).get('perf_metrics') or {} step_metrics = gen_perf_data.get('step_metrics') - # ctx GPU timing / chunk metrics from ctx_perf if ctx_gpu_forward_time is None: ctx_time_breakdown = (ctx_perf or {}).get('time_breakdown_metrics') @@ -278,7 +280,6 @@ def parse_request(self, request_data: Dict, ctx_gpu_sample_time = (ctx_perf or {}).get('ctx_gpu_sample_time') else: - # Try time_breakdown_metrics at top level first (new structure) time_breakdown = request_data.get('time_breakdown_metrics') if time_breakdown: step_metrics = time_breakdown.get('step_metrics') @@ -334,26 +335,55 @@ def __init__(self, config: Optional[TimingMetricsConfig] = None): def parse_json_file(self, json_file_path: str) -> List[Dict]: """Parse JSON performance metrics file and extract timing information.""" - try: - with open(json_file_path, 'r') as f: - data = json.load(f) - except FileNotFoundError: - print(f"Error: File '{json_file_path}' not found.") - sys.exit(1) - except json.JSONDecodeError as e: - print(f"Error parsing JSON file '{json_file_path}': {e}") - sys.exit(1) + + def iter_records(json_file): + if json_file_path.endswith('.jsonl'): + for line_number, line in enumerate(json_file, start=1): + if not line.strip(): + continue + try: + record = json.loads(line) + except json.JSONDecodeError as error: + raise ValueError( + f"Error parsing JSONL file '{json_file_path}' at " + f"line {line_number}: {error}") from error + if not isinstance(record, dict): + raise ValueError( + f"Expected a JSON object at line {line_number}: " + f"{json_file_path}") + yield record + return + + try: + data = json.load(json_file) + except json.JSONDecodeError as error: + json_file.seek(0) + if not any(line.strip() for line in json_file): + return + raise ValueError( + f"Error parsing JSON file '{json_file_path}': {error}" + ) from error + + if isinstance(data, dict): + yield data + elif isinstance(data, list): + yield from data + else: + raise ValueError( + "Expected a JSON array, JSON object, or JSONL file: " + f"{json_file_path}") timing_data = [] - for i, request in enumerate(data): - parsed_data = self.parser.parse_request(request, i) + with open(json_file_path, 'r') as json_file: + for i, request in enumerate(iter_records(json_file)): + parsed_data = self.parser.parse_request(request, i) - # Calculate durations for each metric - for metric in self.config.metrics: - duration = metric.calculate_duration(parsed_data) - parsed_data[f'{metric.name}_time'] = duration + # Calculate durations for each metric + for metric in self.config.metrics: + duration = metric.calculate_duration(parsed_data) + parsed_data[f'{metric.name}_time'] = duration - timing_data.append(parsed_data) + timing_data.append(parsed_data) if timing_data: has_gen_metrics = any(not math.isnan( @@ -2141,16 +2171,17 @@ def main(): formatter_class=argparse.RawDescriptionHelpFormatter, epilog=""" Examples: - python time_breakdown.py perf_metrics.json - python time_breakdown.py perf_metrics.json -o my_timing.html - python time_breakdown.py perf_metrics.json --stats-only - python time_breakdown.py perf_metrics.json --max-requests 50 --sort-by e2e - python time_breakdown.py perf_metrics.json --max-requests 100 --sort-by arrival + python time_breakdown.py perf_metrics.jsonl + python time_breakdown.py perf_metrics.jsonl -o my_timing.html + python time_breakdown.py perf_metrics.jsonl --stats-only + python time_breakdown.py perf_metrics.jsonl --max-requests 50 --sort-by e2e + python time_breakdown.py perf_metrics.jsonl --max-requests 100 --sort-by arrival """) - parser.add_argument('json_file', - type=str, - help='Path to JSON performance metrics file') + parser.add_argument( + 'json_file', + type=str, + help='Path to a JSON or server-produced JSONL performance metrics file') parser.add_argument('-o', '--output', type=str, @@ -2180,11 +2211,18 @@ def main(): analyzer = RequestTimeBreakdown() print(f"Parsing: {args.json_file}") - timing_data = analyzer.parse_json_file(args.json_file) + try: + timing_data = analyzer.parse_json_file(args.json_file) + except FileNotFoundError: + print(f"Error: File '{args.json_file}' not found.") + return 1 + except ValueError as error: + print(error) + return 1 if not timing_data: print("No timing data found.") - sys.exit(1) + return 1 if args.stats_only or args.show_stats: analyzer.show_statistics(timing_data) @@ -2194,7 +2232,8 @@ def main(): args.output, max_requests=args.max_requests, sort_by=args.sort_by) + return 0 if __name__ == '__main__': - main() + sys.exit(main()) diff --git a/tests/integration/defs/accuracy/test_disaggregated_serving.py b/tests/integration/defs/accuracy/test_disaggregated_serving.py index 59f4bcaa87f9..345361c51536 100644 --- a/tests/integration/defs/accuracy/test_disaggregated_serving.py +++ b/tests/integration/defs/accuracy/test_disaggregated_serving.py @@ -166,7 +166,7 @@ def launch_disaggregated_llm( print( f"Using unified tp parameter for testing is not recommended. Please use server configs instead." ) - perf_max_requests = 50 + perf_metrics_output_dir = os.path.join(temp_dir.name, "perf_metrics") def _apply_perf_flags(cfg: Optional[Dict[str, Any]]): if not isinstance(cfg, dict): @@ -175,7 +175,7 @@ def _apply_perf_flags(cfg: Optional[Dict[str, Any]]): # Only set these if the switch is enabled. # Use `setdefault` so explicit per-test overrides are preserved. cfg.setdefault("return_perf_metrics", True) - cfg.setdefault("perf_metrics_max_requests", perf_max_requests) + cfg.setdefault("perf_metrics_output_dir", perf_metrics_output_dir) _apply_perf_flags(disaggregated_server_config) _apply_perf_flags(ctx_server_config) @@ -479,27 +479,6 @@ def generate_async(prompt: str, thread_pool.futures.append(future) return future - def _get_perf_metrics(): - path = "/perf_metrics" - perf_url = f"http://localhost:{serve_port}{path}" - try: - print(f"Fetching perf metrics from {perf_url}") - resp = requests.get(perf_url, timeout=10) - if resp.status_code == 200: - try: - metrics = resp.json() - print("perf_metrics JSON:") - print(json.dumps(metrics, indent=2, ensure_ascii=False)) - except ValueError: - print("perf_metrics returned non-JSON response:", - resp.text) - else: - print( - f"perf_metrics returned status {resp.status_code}: {resp.text}" - ) - except requests.exceptions.RequestException as e: - print(f"Error fetching {perf_url}: {e}") - def _show_kvcache_time(kv_cache_perf_dir, max_lines=100): print(f"kv_cache_perf_dir: {kv_cache_perf_dir}") for file in os.listdir(kv_cache_perf_dir): @@ -515,7 +494,6 @@ def _show_kvcache_time(kv_cache_perf_dir, max_lines=100): finally: if enable_perf: _show_kvcache_time(kv_cache_perf_dir) - _get_perf_metrics() # Gracefully shut down all server processes all_processes = list( diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional_deepseek_v3_v2.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional_deepseek_v3_v2.yaml index d33e1541406e..5b21530c62d3 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional_deepseek_v3_v2.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_conditional_deepseek_v3_v2.yaml @@ -7,9 +7,9 @@ conditional_disagg_config: max_local_prefill_length: 100 disable_overlap_scheduler: true enable_autotuner: false -# Enable per-request perf metrics so the disagg /perf_metrics endpoint (used by -# the test's _check_bypass_signal) records routed requests; bypassed ones absent. +# Enable response-carried metrics so the test can validate routed phases. perf_metrics_max_requests: 1000 +return_perf_metrics: true context_servers: num_instances: 1 tensor_parallel_size: 1 diff --git a/tests/integration/defs/disaggregated/test_configs/disagg_config_metrics.yaml b/tests/integration/defs/disaggregated/test_configs/disagg_config_metrics.yaml index 48fc0a072f02..4bb52cc134f9 100644 --- a/tests/integration/defs/disaggregated/test_configs/disagg_config_metrics.yaml +++ b/tests/integration/defs/disaggregated/test_configs/disagg_config_metrics.yaml @@ -5,6 +5,7 @@ backend: "pytorch" cuda_graph_config: null disable_overlap_scheduler: True perf_metrics_max_requests: 1000 +return_perf_metrics: true context_servers: num_instances: 1 tensor_parallel_size: 1 diff --git a/tests/integration/defs/disaggregated/test_disaggregated.py b/tests/integration/defs/disaggregated/test_disaggregated.py index 05c905db31bb..4e1cc89d4713 100644 --- a/tests/integration/defs/disaggregated/test_disaggregated.py +++ b/tests/integration/defs/disaggregated/test_disaggregated.py @@ -40,7 +40,8 @@ run_disagg_server, run_gen_worker, terminate, wait_for_disagg_server_ready) from test_common.perf_metrics_utils import (get_timing_metrics, - validate_timing_metrics) + validate_timing_metrics, + wait_for_perf_metrics_jsonl) from tensorrt_llm._utils import mpi_disabled from tensorrt_llm.logger import logger @@ -654,6 +655,7 @@ def setup_disagg_cluster( save_log: bool = False, startup_callback=None, startup_tick: int = 30, + perf_metrics_output_dir: str | None = None, ) -> tuple[dict[str, Any], list[ProcessWrapper], list[ProcessWrapper], ProcessWrapper, int, str]: """Load config, launch workers + disagg server, wait for ready. @@ -671,6 +673,13 @@ def setup_disagg_cluster( with open(config_file, 'r') as f: config = yaml.safe_load(f) + if perf_metrics_output_dir is not None: + config["perf_metrics_output_dir"] = perf_metrics_output_dir + for servers in ("context_servers", "generation_servers"): + config.setdefault( + servers, + {})["perf_metrics_output_dir"] = perf_metrics_output_dir + speculative_config = config.get("speculative_config") if isinstance(speculative_config, dict): speculative_model = speculative_config.get("speculative_model") @@ -792,6 +801,10 @@ def setup_disagg_cluster( config.get("conditional_disagg_config", None), "perf_metrics_max_requests": config.get("perf_metrics_max_requests", 0), + "perf_metrics_output_dir": + config.get("perf_metrics_output_dir", None), + "return_perf_metrics": + config.get("return_perf_metrics", False), } if schedule_style: server_config["schedule_style"] = schedule_style @@ -915,7 +928,8 @@ def run_disaggregated_test(example_dir, cwd=None, disagg_schedule_style=None, post_client_test=None, - assert_gen_log_contains=None): + assert_gen_log_contains=None, + perf_metrics_output_dir=None): """Run disaggregated test using service discovery instead of MPI. If assert_gen_log_contains is set, the generation-worker logs are captured and, after the @@ -935,7 +949,8 @@ def run_disaggregated_test(example_dir, config, ctx_workers, gen_workers, disagg_server, server_port, work_dir = \ setup_disagg_cluster(config_file, model_name=model_path, env=run_env, cwd=cwd, schedule_style=disagg_schedule_style, - save_log=assert_gen_log_contains is not None) + save_log=assert_gen_log_contains is not None, + perf_metrics_output_dir=perf_metrics_output_dir) server_host = config.get("hostname", "localhost") @@ -1062,28 +1077,13 @@ def test_disaggregated_benchmark_gen_only(disaggregated_test_root, indirect=True) def test_disaggregated_router(disaggregated_test_root, disaggregated_example_root, llm_venv, - llama_model_root, router_type, tmp_path): + llama_model_root, router_type): setup_model_symlink(llm_venv, llama_model_root, "TinyLlama/TinyLlama-1.1B-Chat-v1.0") - metrics_file = tmp_path / f"perf_metrics_{router_type}.json" - - def fetch_perf_metrics(server_url: str): - import json - - import requests as http_requests - resp = http_requests.get(f"{server_url}/perf_metrics", timeout=10) - assert resp.status_code == 200, \ - f"Failed to fetch perf_metrics: {resp.status_code}" - metrics = resp.json() - metrics_file.write_text(json.dumps(metrics, indent=2)) - logger.info(f"Router={router_type}: saved {len(metrics)} perf metrics " - f"to {metrics_file}") - run_disaggregated_test(disaggregated_example_root, router_type, env=llm_venv._new_env, - extra_endpoints_test=fetch_perf_metrics, model_path=llama_model_root, cwd=llm_venv.get_working_directory()) @@ -1503,12 +1503,14 @@ def post_client_test(server_url: str): indirect=True) def test_disaggregated_perf_metrics(disaggregated_test_root, llm_venv, disaggregated_example_root, - llama_model_root): + llama_model_root, tmp_path): setup_model_symlink(llm_venv, llama_model_root, "TinyLlama/TinyLlama-1.1B-Chat-v1.0") - def extra_endpoints_test(server_url: str): - item = get_timing_metrics(server_url) + perf_metrics_output_dir = str(tmp_path / "perf_metrics") + + def extra_endpoints_test(_server_url: str): + item = get_timing_metrics(perf_metrics_output_dir) # Use helper function to validate all timing metrics comprehensively validate_timing_metrics(item, "perf_metrics test") @@ -1517,7 +1519,8 @@ def extra_endpoints_test(server_url: str): env=llm_venv._new_env, extra_endpoints_test=extra_endpoints_test, model_path=llama_model_root, - cwd=llm_venv.get_working_directory()) + cwd=llm_venv.get_working_directory(), + perf_metrics_output_dir=perf_metrics_output_dir) @pytest.mark.parametrize("llama_model_root", ['TinyLlama-1.1B-Chat-v1.0'], @@ -2103,45 +2106,29 @@ def test_disaggregated_deepseek_v3_lite_bf16_conditional( indirect=True) def test_disaggregated_deepseek_v3_lite_bf16_conditional_v2( disaggregated_test_root, disaggregated_example_root, llm_venv, - deepseek_v3_model_root): + deepseek_v3_model_root, tmp_path): setup_model_symlink(llm_venv, deepseek_v3_model_root, "DeepSeek-V3-Lite/bf16") - # Conditional disagg handles short-prefill requests locally on the gen - # server (bypassing the ctx handoff + add_per_request_metrics), while routed - # requests are recorded in the disagg /perf_metrics. Verify ONCE after all - # client iterations via post_client_test (not per-iteration): routed-request - # metrics are recorded asynchronously (add_per_request_metrics via - # create_task on response completion) and surface only after the client - # traffic settles, so a per-iteration read races that lag; /perf_metrics is - # also consume-on-read, so query it exactly once at the end. - def _check_routed_recorded(server_url: str): - import requests as http_requests - metrics = [] - deadline = time.time() + 60 - while True: - resp = http_requests.get(f"{server_url}/perf_metrics", timeout=10) - assert resp.status_code == 200, \ - f"perf_metrics fetch failed: {resp.status_code}" - metrics = resp.json() - if metrics or time.time() >= deadline: - break - time.sleep(2) - logger.info(f"conditional_v2 perf_metrics len={len(metrics)} " - f"(routed requests recorded; bypassed ones absent)") - # With short prompts every prompt's first occurrence routes through the - # context server (match=0 -> need_ctx), so at least one routed request - # must be recorded; an empty result means conditional routing never - # engaged. - assert metrics, \ - "no per-request metrics recorded after client runs; conditional routing may be misconfigured" + perf_metrics_output_dir = str(tmp_path / "perf_metrics") + + def _check_routed_recorded(_server_url: str): + records = wait_for_perf_metrics_jsonl(perf_metrics_output_dir, + expected_count=3, + timeout=60) + assert any( + record.get("worker", {}).get("server_kind") == "disagg" + and {"ctx", "gen"} <= record.get("phases", {}).keys() + for record in records + ), "new prompt did not produce routed context and generation metrics" run_disaggregated_test(disaggregated_example_root, "deepseek_v3_lite_bf16_conditional_v2", env=llm_venv._new_env, post_client_test=_check_routed_recorded, model_path=deepseek_v3_model_root, - cwd=llm_venv.get_working_directory()) + cwd=llm_venv.get_working_directory(), + perf_metrics_output_dir=perf_metrics_output_dir) @skip_no_hopper diff --git a/tests/test_common/perf_metrics_utils.py b/tests/test_common/perf_metrics_utils.py index c63faa6d81f8..7ed5269d5c35 100644 --- a/tests/test_common/perf_metrics_utils.py +++ b/tests/test_common/perf_metrics_utils.py @@ -1,19 +1,81 @@ +import json +import time +from pathlib import Path + import requests -def get_timing_metrics(server_url: str): - response = requests.get(f"{server_url}/perf_metrics", timeout=10) - assert response.status_code == 200 - perf_metrics = response.json() - assert len(perf_metrics) > 0 - return perf_metrics[0] +def _read_perf_metrics_jsonl(output_dir): + records = [] + for path in Path(output_dir).glob("perf_metrics-*.jsonl"): + for line in path.read_text().splitlines(): + if not line: + continue + try: + records.append((path, json.loads(line))) + except json.JSONDecodeError: + continue + return records + + +def read_perf_metrics_jsonl(output_dir): + return [record for _, record in _read_perf_metrics_jsonl(output_dir)] + + +def wait_for_perf_metrics_jsonl(output_dir, expected_count: int, timeout: float = 10): + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + records = read_perf_metrics_jsonl(output_dir) + if len(records) >= expected_count: + return records + time.sleep(0.1) + raise AssertionError( + f"Timed out waiting for {expected_count} performance metrics JSONL records" + ) + + +def get_timing_metrics(output_dir, timeout: float = 10): + """Read and join one completed disaggregated request from JSONL files.""" + deadline = time.monotonic() + timeout + while time.monotonic() < deadline: + records = _read_perf_metrics_jsonl(output_dir) + by_request = {} + for path, record in records: + disagg_request_id = record.get("disagg_request_id") + server_kind = path.name.removeprefix("perf_metrics-").split("-", 1)[0] + if disagg_request_id is not None and server_kind: + by_request.setdefault(disagg_request_id, {})[server_kind] = record + + for records_by_kind in by_request.values(): + if not {"context", "generation", "disagg"} <= records_by_kind.keys(): + continue + ctx_record = records_by_kind["context"] + gen_record = records_by_kind["generation"] + disagg_record = records_by_kind["disagg"] + ctx_metrics = { + key: value + for key, value in ctx_record.items() + if key not in ("disagg_request_id", "status") + } + gen_metrics = { + key: value + for key, value in gen_record.items() + if key not in ("disagg_request_id", "status") + } + return { + **disagg_record, + "ctx_perf_metrics": ctx_metrics, + "gen_perf_metrics": gen_metrics, + } + time.sleep(0.1) + raise AssertionError("Timed out waiting for joined performance metrics JSONL records") def validate_timing_metrics(perf_metrics_item, request_context="", time_tolerance_seconds=0.005): """Helper function to validate timing metrics relationships. Args: - perf_metrics_item: A single performance metrics item from the /perf_metrics endpoint + perf_metrics_item: Joined context, generation, and disagg JSONL metrics request_context: String context for error messages (e.g., "request 1", "streaming") """ # Validate basic structure diff --git a/tests/unittest/api_stability/references/llm.yaml b/tests/unittest/api_stability/references/llm.yaml index 959f249183e8..4c06bbdb664e 100644 --- a/tests/unittest/api_stability/references/llm.yaml +++ b/tests/unittest/api_stability/references/llm.yaml @@ -35,6 +35,10 @@ methods: annotation: bool default: False status: prototype + perf_metrics_output_dir: + annotation: Optional[str] + default: null + status: prototype # Bindings and mirrored configs peft_cache_config: annotation: Optional[tensorrt_llm.llmapi.llm_args.PeftCacheConfig] @@ -202,7 +206,7 @@ methods: perf_metrics_max_requests: annotation: int default: 0 - status: prototype + status: deprecated prometheus_metrics_config: annotation: Optional[tensorrt_llm.llmapi.llm_args.PrometheusMetricsConfig] default: null diff --git a/tests/unittest/disaggregated/test_disagg_openai_client.py b/tests/unittest/disaggregated/test_disagg_openai_client.py index cead306ae3ff..9b8144fd9a2e 100644 --- a/tests/unittest/disaggregated/test_disagg_openai_client.py +++ b/tests/unittest/disaggregated/test_disagg_openai_client.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. +# Copyright (c) 2025-2026, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -26,6 +26,8 @@ DisaggregatedParams, UsageInfo, ) +from tensorrt_llm.serve.perf_metrics import _PERF_METRICS_HEADER_BUDGET_BYTES, SSE_METRICS_EVENT +from tensorrt_llm.serve.responses_utils import ResponseHooks from tensorrt_llm.serve.router import Router @@ -116,6 +118,16 @@ def test_initialization(self, mock_router, mock_session): assert client._session == mock_session assert client._max_retries == 5 + @pytest.mark.asyncio + async def test_internal_client_accepts_perf_metrics_header_size(self, mock_router): + with ( + patch("tensorrt_llm.serve.openai_client.ClientMetricsCollector"), + patch("tensorrt_llm.serve.openai_client.aiohttp.ClientSession") as session, + ): + OpenAIHttpClient(router=mock_router, role=ServerRole.GENERATION) + + assert session.call_args.kwargs["max_field_size"] == _PERF_METRICS_HEADER_BUDGET_BYTES + @pytest.mark.asyncio async def test_non_streaming_completion_request( self, openai_client, completion_request, mock_session, mock_router @@ -189,6 +201,45 @@ async def mock_iter_any(): streaming_completion_request, mock_session, success=True ) + @pytest.mark.asyncio + async def test_malformed_streaming_metrics_do_not_fail_request( + self, openai_client, streaming_completion_request, mock_session, mock_router + ): + openai_client._request_perf_metrics = True + mock_http_response = AsyncMock() + mock_http_response.status = 200 + mock_http_response.headers = {"Content-Type": "text/event-stream"} + + response_data = b'data: "Hello"\n\ndata: [DONE]\n\n' + metrics_data = f"event: {SSE_METRICS_EVENT}\ndata: not-json\n\n".encode() + + async def mock_iter_any(): + yield b"" + yield response_data + yield metrics_data + + mock_http_response.content = AsyncMock() + mock_http_response.content.iter_any = mock_iter_any + mock_http_response.__aenter__ = AsyncMock(return_value=mock_http_response) + mock_http_response.__aexit__ = AsyncMock() + mock_session.post.return_value = mock_http_response + hooks = MagicMock(spec=ResponseHooks) + + response_generator = await openai_client.send_request( + streaming_completion_request, hooks=hooks + ) + chunks = [chunk async for chunk in response_generator] + + assert b"".join(chunks) == response_data + hooks.on_first_token.assert_called_once_with("localhost:8000", streaming_completion_request) + hooks.on_perf_metrics.assert_not_called() + hooks.on_resp_done.assert_called_once_with( + "localhost:8000", streaming_completion_request, None + ) + mock_router.finish_request.assert_called_once_with( + streaming_completion_request, mock_session, success=True + ) + @pytest.mark.asyncio async def test_request_with_custom_server( self, openai_client, completion_request, mock_session, mock_router diff --git a/tests/unittest/disaggregated/test_disagg_utils.py b/tests/unittest/disaggregated/test_disagg_utils.py index f3556bae1fe5..e1fbc4a8645a 100644 --- a/tests/unittest/disaggregated/test_disagg_utils.py +++ b/tests/unittest/disaggregated/test_disagg_utils.py @@ -127,6 +127,22 @@ def test_extract_disagg_cfg(sample_yaml_config): assert config.disagg_coordinator_url == "http://coordinator:7999" +def test_extract_disagg_metrics_controls(): + yaml_config = get_yaml_config() + yaml_config["context_servers"]["return_perf_metrics"] = False + yaml_config["generation_servers"]["return_perf_metrics"] = False + config = extract_disagg_cfg( + **yaml_config, + return_perf_metrics=True, + perf_metrics_output_dir="/tmp/perf", + ) + + assert config.return_perf_metrics is True + assert config.perf_metrics_output_dir == "/tmp/perf" + assert all("perf_metrics_output_dir" not in server.other_args + for server in config.server_configs) + + @pytest.mark.parametrize("node_id", [-1, 256]) def test_extract_disagg_cfg_rejects_out_of_range_node_id(node_id): with pytest.raises(ValueError, match="node_id must be in range"): diff --git a/tests/unittest/disaggregated/test_openai_disagg_service.py b/tests/unittest/disaggregated/test_openai_disagg_service.py index 04defdd1cd37..f19d36739dca 100644 --- a/tests/unittest/disaggregated/test_openai_disagg_service.py +++ b/tests/unittest/disaggregated/test_openai_disagg_service.py @@ -382,6 +382,31 @@ async def _delayed_gen_response(*_args, **_kwargs): ) +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +async def test_context_only_response_finishes_hooks(stream): + service = _make_service("context_first") + service._ctx_client = AsyncMock() + service._gen_client = AsyncMock() + service._coordinator.get_disagg_request_id = AsyncMock(return_value=42) + service._check_conditional_disagg = AsyncMock(return_value=(None, True)) + service._check_gen_only_disagg = AsyncMock(return_value=False) + service._ctx_router.get_next_server = AsyncMock(return_value=("ctx:9000", {"server_info": {}})) + ctx_response = _make_completion_response("done", finish_reason="stop", context_only=True) + service._ctx_client.send_request = AsyncMock(return_value=ctx_response) + hooks = mock.Mock() + request = CompletionRequest(model="test-model", prompt="hello", stream=stream) + + result = await service._send_disagg_request(request, hooks) + if stream: + assert [chunk async for chunk in result] == [b"data: [DONE]\n\n"] + else: + assert result is ctx_response + + hooks.on_resp_done.assert_called_once_with("", request, ctx_response) + service._gen_client.send_request.assert_not_awaited() + + @pytest.mark.asyncio @pytest.mark.parametrize("schedule_style", ["context_first", "generation_first"]) async def test_send_disagg_request_leaves_streaming_usage_to_gen_server(schedule_style): diff --git a/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py b/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py index 20f7ef29eea6..02b2b4421b1b 100644 --- a/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py +++ b/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes.py @@ -7,8 +7,6 @@ import pytest from test_common.http_utils import (wait_for_endpoint_down, wait_for_endpoint_ready) -from test_common.perf_metrics_utils import (get_timing_metrics, - validate_timing_metrics) from tensorrt_llm._utils import get_sm_version @@ -121,8 +119,6 @@ def worker(model_name: str, ctx_tp_pp_size: tuple, gen_tp_pp_size: tuple): "enable_block_reuse": False, }, "disable_overlap_scheduler": True, - "perf_metrics_max_requests": 1000, - "return_perf_metrics": True, } if is_ctx_node(): print(f"starting ctx_server for rank {RANK} node rank {NODE_RANK}") @@ -203,11 +199,6 @@ def test_completion(client: openai.OpenAI, message = completion.choices[0].text assert message.startswith('2.') - perf_metrics = get_timing_metrics(disagg_server.url_root) - # allow 5ms leniency when comparing the time points from disagg and ctx/gen servers - validate_timing_metrics(perf_metrics, - "multinode test_completion", - time_tolerance_seconds=0.005) # sleep 10 seconds to ensure a successful wait_for_endpoint_ready on rank1 time.sleep(10) disagg_server.terminate() diff --git a/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes_service_discovery.py b/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes_service_discovery.py index 0e6c997a41a4..780844f2a77d 100644 --- a/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes_service_discovery.py +++ b/tests/unittest/llmapi/apps/_test_disagg_serving_multi_nodes_service_discovery.py @@ -6,7 +6,6 @@ import openai import pytest -from test_common.perf_metrics_utils import get_timing_metrics, validate_timing_metrics from tensorrt_llm._utils import get_free_port from tensorrt_llm.llmapi.disagg_utils import ServerRole @@ -112,8 +111,6 @@ def worker(model_name: str, disagg_cluster_config: dict): "enable_block_reuse": False, }, "disable_overlap_scheduler": True, - "return_perf_metrics": True, - "perf_metrics_max_requests": 1000, } # start workers on 0.0.0.0:, then the workers should be able to # report their correct hostname:port to the disagg server @@ -165,7 +162,6 @@ def disagg_server(disagg_cluster_config: dict): "disagg_cluster": disagg_cluster_config, "port": DISAGG_SERVER_PORT, "hostname": "0.0.0.0", - "perf_metrics_max_requests": 1000, } print(f"starting disagg_server for rank {RANK} node rank {NODE_RANK}") # ctx/gen servers are unnecessary for service discovery test @@ -217,9 +213,6 @@ def test_completion( message = completion.choices[0].text assert message.startswith("2.") - perf_metrics = get_timing_metrics(disagg_server.url_root) - validate_timing_metrics(perf_metrics, "multinode test_completion") - disagg_server.terminate() elif is_gen_node(): diff --git a/tests/unittest/llmapi/apps/_test_openai_chat_harmony_perf_metrics.py b/tests/unittest/llmapi/apps/_test_openai_chat_harmony_perf_metrics.py index 85cd2a073b1a..7bd3cfa17019 100644 --- a/tests/unittest/llmapi/apps/_test_openai_chat_harmony_perf_metrics.py +++ b/tests/unittest/llmapi/apps/_test_openai_chat_harmony_perf_metrics.py @@ -1,25 +1,14 @@ # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 -"""Tests for /perf_metrics population on the Harmony (GPT-OSS) chat path. - -Asserts two behaviors of the Harmony chat endpoint: - * The streaming generator stamps `server_first_token_time` and calls - `_extract_metrics(...)` after `data: [DONE]`, so the `/perf_metrics` - deque is populated for streaming chat requests. - * When `TRTLLM_KVCACHE_TIME_OUTPUT_PATH` is set, the per-request - `sampling_params.return_perf_metrics` flag is enabled so the engine - emits metrics for each request. Without it, the deque stays empty - even with `return_perf_metrics: True` at the LLM-args level. -""" - -import json +"""Tests JSONL performance metrics for the Harmony chat path.""" + import os import tempfile -from urllib.request import urlopen import openai import pytest import yaml +from test_common.perf_metrics_utils import read_perf_metrics_jsonl, wait_for_perf_metrics_jsonl from utils.llm_data import llm_datasets_root from ..test_llm import get_model_path @@ -46,15 +35,14 @@ def kv_cache_time_output_dir(tmp_path_factory): @pytest.fixture(scope="module") -def extra_llm_api_options_file(): +def extra_llm_api_options_file(kv_cache_time_output_dir: str): fd, path = tempfile.mkstemp(suffix=".yaml", prefix="extra_llm_api_options_") os.close(fd) try: with open(path, "w") as f: yaml.dump( { - "return_perf_metrics": True, - "perf_metrics_max_requests": 16, + "perf_metrics_output_dir": kv_cache_time_output_dir, }, f, ) @@ -92,17 +80,10 @@ def async_client(server: RemoteOpenAIServer): return server.get_async_client() -def _drain_perf_metrics(server: RemoteOpenAIServer): - response = urlopen(f"{server.url_root}/perf_metrics") - assert response.status == 200 - return json.loads(response.read()) - - def _assert_perf_metrics_entry_well_formed(entry: dict): assert "request_id" in entry - assert "perf_metrics" in entry + assert entry["status"] == "complete" pm = entry["perf_metrics"] - assert "first_iter" in pm and "last_iter" in pm assert pm["first_iter"] <= pm["last_iter"] tm = pm["timing_metrics"] @@ -115,10 +96,12 @@ def _assert_perf_metrics_entry_well_formed(entry: dict): @pytest.mark.asyncio(loop_scope="module") async def test_non_streaming_perf_metrics( - async_client: openai.AsyncOpenAI, server: RemoteOpenAIServer, model: str + async_client: openai.AsyncOpenAI, + server: RemoteOpenAIServer, + model: str, + kv_cache_time_output_dir: str, ): - # Drain anything from prior tests in this module session. - _drain_perf_metrics(server) + previous_count = len(read_perf_metrics_jsonl(kv_cache_time_output_dir)) response = await async_client.chat.completions.create( model=model, messages=[{"role": "user", "content": "Reply with exactly the single word: PONG."}], @@ -126,20 +109,19 @@ async def test_non_streaming_perf_metrics( ) assert response.choices[0].message.content is not None - entries = _drain_perf_metrics(server) - assert len(entries) == 1, ( - "Expected exactly one /perf_metrics entry after a single non-streaming " - f"harmony chat completion, got {len(entries)}: {entries}" - ) - _assert_perf_metrics_entry_well_formed(entries[0]) + records = wait_for_perf_metrics_jsonl(kv_cache_time_output_dir, previous_count + 1) + entry = records[-1] + _assert_perf_metrics_entry_well_formed(entry) @pytest.mark.asyncio(loop_scope="module") async def test_streaming_perf_metrics( - async_client: openai.AsyncOpenAI, server: RemoteOpenAIServer, model: str + async_client: openai.AsyncOpenAI, + server: RemoteOpenAIServer, + model: str, + kv_cache_time_output_dir: str, ): - # Drain anything from prior tests in this module session. - _drain_perf_metrics(server) + previous_count = len(read_perf_metrics_jsonl(kv_cache_time_output_dir)) stream = await async_client.chat.completions.create( model=model, messages=[{"role": "user", "content": "Explain transformers in one sentence."}], @@ -153,10 +135,6 @@ async def test_streaming_perf_metrics( saw_done = True assert saw_done, "Streaming chat returned no chunks" - entries = _drain_perf_metrics(server) - assert len(entries) == 1, ( - "Expected exactly one /perf_metrics entry after a single streaming " - f"harmony chat completion, got {len(entries)}: {entries}. " - "This usually means _extract_metrics did not run after [DONE]." - ) - _assert_perf_metrics_entry_well_formed(entries[0]) + records = wait_for_perf_metrics_jsonl(kv_cache_time_output_dir, previous_count + 1) + entry = records[-1] + _assert_perf_metrics_entry_well_formed(entry) diff --git a/tests/unittest/llmapi/apps/_test_openai_perf_metrics.py b/tests/unittest/llmapi/apps/_test_openai_perf_metrics.py index c1049939953b..b285f2190e3f 100644 --- a/tests/unittest/llmapi/apps/_test_openai_perf_metrics.py +++ b/tests/unittest/llmapi/apps/_test_openai_perf_metrics.py @@ -1,11 +1,13 @@ -import json import logging import os import tempfile -from urllib.request import urlopen import pytest +import requests import yaml +from test_common.perf_metrics_utils import wait_for_perf_metrics_jsonl + +from tensorrt_llm.serve import perf_metrics from ..test_llm import get_model_path from .openai_server import RemoteOpenAIServer @@ -21,13 +23,18 @@ def model_name(): @pytest.fixture(scope="module") -def temp_extra_llm_api_options_file(request): +def perf_metrics_output_dir(tmp_path_factory): + return tmp_path_factory.mktemp("perf_metrics") + + +@pytest.fixture(scope="module") +def temp_extra_llm_api_options_file(perf_metrics_output_dir): temp_dir = tempfile.gettempdir() temp_file_path = os.path.join(temp_dir, "extra_llm_api_options.yaml") try: extra_llm_api_options_dict = { "return_perf_metrics": True, - "perf_metrics_max_requests": 10 + "perf_metrics_output_dir": str(perf_metrics_output_dir), } with open(temp_file_path, 'w') as f: @@ -51,34 +58,43 @@ def server(model_name: str, logger.info("Tests completed, shutting down server") -def test_metrics_endpoint(server: RemoteOpenAIServer): - - client = server.get_client() - client.completions.create( - model="Server", - prompt="Hello, my name is", - max_tokens=25, - stream=False, +def test_return_perf_metrics_and_jsonl_dump(server: RemoteOpenAIServer, + perf_metrics_output_dir): + response = requests.post( + f"{server.url_root}/v1/completions", + headers={perf_metrics.RETURN_METRICS_HEADER: "1"}, + json={ + "model": "Server", + "prompt": "Hello, my name is", + "max_tokens": 2, + }, + timeout=120, ) - - response = urlopen(f'{server.url_root}/perf_metrics') - assert response.status is 200 - - data_list = json.loads(response.read()) - assert len(data_list) == 1 - assert "perf_metrics" in data_list[0] - assert "request_id" in data_list[0] - - data = data_list[0]["perf_metrics"] - assert "first_iter" in data - assert "last_iter" in data - assert data["first_iter"] <= data["last_iter"] - - timing_metrics = data["timing_metrics"] - assert "arrival_time" in timing_metrics - assert "first_scheduled_time" in timing_metrics - assert "first_token_time" in timing_metrics - assert "last_token_time" in timing_metrics + assert response.status_code == 200 + + for header in ( + perf_metrics.SERVER_TIMING_HEADER, + perf_metrics.START_END_TIME_HEADER, + perf_metrics.STEP_METRICS_HEADER, + perf_metrics.CTX_CHUNK_METRICS_HEADER, + ): + assert response.headers.get(header) + + records = wait_for_perf_metrics_jsonl(perf_metrics_output_dir, + expected_count=1) + data = records[-1] + assert data["status"] == "complete" + assert set(data) == { + "request_id", + "perf_metrics", + "time_breakdown_metrics", + "status", + } + + request_metrics = data["perf_metrics"] + assert request_metrics["first_iter"] <= request_metrics["last_iter"] + + timing_metrics = request_metrics["timing_metrics"] assert timing_metrics["arrival_time"] < timing_metrics[ "first_scheduled_time"] assert timing_metrics["first_scheduled_time"] < timing_metrics[ @@ -86,16 +102,35 @@ def test_metrics_endpoint(server: RemoteOpenAIServer): assert timing_metrics["first_token_time"] <= timing_metrics[ "last_token_time"] - kv_cache_metrics = data["kv_cache_metrics"] - assert "num_total_allocated_blocks" in kv_cache_metrics - assert "num_new_allocated_blocks" in kv_cache_metrics - assert "num_reused_blocks" in kv_cache_metrics - assert "num_missed_blocks" in kv_cache_metrics + kv_cache_metrics = request_metrics["kv_cache_metrics"] assert kv_cache_metrics["num_new_allocated_blocks"] <= kv_cache_metrics[ "num_total_allocated_blocks"] - # exclude disagg specific metrics - assert "ctx_request_id" not in data_list[0] - assert "kv_cache_size" not in timing_metrics + assert "ctx_request_id" not in data assert "kv_cache_transfer_start" not in timing_metrics assert "kv_cache_transfer_end" not in timing_metrics + + +def test_streaming_metrics_require_request_opt_in(server: RemoteOpenAIServer): + payload = { + "model": "Server", + "prompt": "Hello, my name is", + "max_tokens": 2, + "stream": True, + } + response = requests.post(f"{server.url_root}/v1/completions", + json=payload, + timeout=120) + assert response.status_code == 200 + assert "data: [DONE]" in response.text + assert f"event: {perf_metrics.SSE_METRICS_EVENT}" not in response.text + + response = requests.post( + f"{server.url_root}/v1/completions", + headers={perf_metrics.RETURN_METRICS_HEADER: "1"}, + json=payload, + timeout=120, + ) + assert response.status_code == 200 + assert "data: [DONE]" in response.text + assert f"event: {perf_metrics.SSE_METRICS_EVENT}" in response.text diff --git a/tests/unittest/llmapi/apps/test_disagg_perf_metrics_collector.py b/tests/unittest/llmapi/apps/test_disagg_perf_metrics_collector.py deleted file mode 100644 index 70ab2432d98f..000000000000 --- a/tests/unittest/llmapi/apps/test_disagg_perf_metrics_collector.py +++ /dev/null @@ -1,68 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -import asyncio - -import pytest - -from tensorrt_llm.serve import perf_metrics - - -class _DummyMetric: - def inc(self): - pass - - def observe(self, _value): - pass - - -class _BlockingClient: - def __init__(self): - self.entered = asyncio.Event() - self.release = asyncio.Event() - self.active_collectors = 0 - self.max_active_collectors = 0 - self.calls = 0 - - async def collect_metrics(self): - self.calls += 1 - self.active_collectors += 1 - self.max_active_collectors = max(self.max_active_collectors, self.active_collectors) - self.entered.set() - await self.release.wait() - self.active_collectors -= 1 - return {} - - -@pytest.mark.asyncio -async def test_disagg_perf_metrics_collection_is_serialized(monkeypatch): - monkeypatch.setattr(perf_metrics, "instance_metric", lambda _definition: _DummyMetric()) - collector = perf_metrics.DisaggPerfMetricsCollector(max_requests=8) - client = _BlockingClient() - collector.add_client(client) - - first_task = asyncio.create_task(collector.get_perf_metrics()) - await client.entered.wait() - - second_task = asyncio.create_task(collector.get_perf_metrics()) - await asyncio.sleep(0) - - assert client.calls == 1 - assert client.max_active_collectors == 1 - - client.release.set() - assert await first_task == [] - assert await second_task == [] - assert client.calls == 2 - assert client.max_active_collectors == 1 diff --git a/tests/unittest/llmapi/apps/test_disagg_serving_perf_metrics.py b/tests/unittest/llmapi/apps/test_disagg_serving_perf_metrics.py index 9b18e3025d73..fc48349d3dfa 100644 --- a/tests/unittest/llmapi/apps/test_disagg_serving_perf_metrics.py +++ b/tests/unittest/llmapi/apps/test_disagg_serving_perf_metrics.py @@ -3,15 +3,25 @@ import openai import pytest +import requests from test_common.http_utils import wait_for_endpoint_ready from test_common.perf_metrics_utils import ( get_prometheus_metrics, get_timing_metrics, validate_timing_metrics, + wait_for_perf_metrics_jsonl, ) +from utils.llm_data import llm_models_root from tensorrt_llm._utils import get_free_ports -from tests.unittest.utils.llm_data import llm_models_root +from tensorrt_llm.serve.perf_metrics import ( + CTX_CHUNK_METRICS_HEADER, + RETURN_METRICS_HEADER, + SERVER_TIMING_HEADER, + SSE_METRICS_EVENT, + START_END_TIME_HEADER, + STEP_METRICS_HEADER, +) from ..test_llm import get_model_path from .openai_server import RemoteDisaggOpenAIServer, RemoteOpenAIServer @@ -44,6 +54,11 @@ def model_name(): return model_path +@pytest.fixture +def perf_metrics_output_dir(tmp_path): + return tmp_path / "perf_metrics" + + @pytest.fixture def disagg_cluster_config(disagg_port: int): return { @@ -52,7 +67,7 @@ def disagg_cluster_config(disagg_port: int): } -def worker_config(model_name: str, disagg_cluster_config: dict): +def worker_config(model_name: str, disagg_cluster_config: dict, perf_metrics_output_dir): return { "model": model_name, "disagg_cluster": disagg_cluster_config, @@ -66,14 +81,21 @@ def worker_config(model_name: str, disagg_cluster_config: dict): "disable_overlap_scheduler": True, "cuda_graph_config": None, "return_perf_metrics": True, + "perf_metrics_output_dir": str(perf_metrics_output_dir), "perf_metrics_max_requests": 1000, } @pytest.fixture -def workers(model_name: str, disagg_cluster_config: dict, ctx_port: int, gen_port: int): +def workers( + model_name: str, + disagg_cluster_config: dict, + ctx_port: int, + gen_port: int, + perf_metrics_output_dir, +): model_path = get_model_path(model_name) - extra_config = worker_config(model_name, disagg_cluster_config) + extra_config = worker_config(model_name, disagg_cluster_config, perf_metrics_output_dir) def worker(server_role: str, port: int): return RemoteOpenAIServer( @@ -92,12 +114,14 @@ def worker(server_role: str, port: int): @pytest.fixture -def disagg_server(disagg_cluster_config: dict, workers, disagg_port: int): +def disagg_server(disagg_cluster_config: dict, workers, disagg_port: int, perf_metrics_output_dir): disagg_config = { "hostname": "localhost", "port": disagg_port, "disagg_cluster": disagg_cluster_config, "perf_metrics_max_requests": 1000, + "return_perf_metrics": True, + "perf_metrics_output_dir": str(perf_metrics_output_dir), } with RemoteDisaggOpenAIServer( ctx_servers=[], @@ -149,6 +173,62 @@ def check_historgram(metrics_dict: dict, count: int, range: tuple[float, float]) assert mean > range[0] and mean < range[1] +@pytest.mark.timeout(300) +def test_return_perf_metrics_and_jsonl_dump( + workers: Tuple[RemoteOpenAIServer, RemoteOpenAIServer], + disagg_server: RemoteDisaggOpenAIServer, + model_name: str, + perf_metrics_output_dir, +): + assert len(workers) == 2 + for worker in workers: + worker.wait_for_server(timeout=120) + wait_for_endpoint_ready(disagg_server.url_root + "/health") + + payload = { + "model": model_name, + "prompt": "Reply with one token.", + "max_tokens": 1, + "temperature": 0.0, + } + response = requests.post( + f"{disagg_server.url_root}/v1/completions", + headers={RETURN_METRICS_HEADER: "1"}, + json=payload, + timeout=120, + ) + assert response.status_code == 200 + assert response.json()["id"] is not None + for header in ( + SERVER_TIMING_HEADER, + START_END_TIME_HEADER, + STEP_METRICS_HEADER, + CTX_CHUNK_METRICS_HEADER, + ): + assert response.headers.get(header) + + timing_metrics = get_timing_metrics(perf_metrics_output_dir) + validate_timing_metrics(timing_metrics, "test_return_perf_metrics_and_jsonl_dump") + records = wait_for_perf_metrics_jsonl(perf_metrics_output_dir, expected_count=3) + disagg_record = next(record for record in records if "ctx_server" in record) + disagg_request_id = disagg_record["disagg_request_id"] + for record in records: + assert record["status"] == "complete" + assert record["disagg_request_id"] == disagg_request_id + + assert "ctx_perf_metrics" in disagg_record + assert "gen_perf_metrics" in disagg_record + + response = requests.post( + f"{disagg_server.url_root}/v1/completions", + json={**payload, "stream": True}, + timeout=120, + ) + assert response.status_code == 200 + assert "data: [DONE]" in response.text + assert f"event: {SSE_METRICS_EVENT}" not in response.text + + @pytest.mark.asyncio @pytest.mark.timeout(300) async def test_completion_metrics( @@ -171,9 +251,6 @@ async def test_completion_metrics( max_token=max_token, model_name=model_name, ) - timing_metrics = get_timing_metrics(disagg_server.url_root) - validate_timing_metrics(timing_metrics, "test_completion_metrics") - metrics = get_prometheus_metrics(disagg_server.url_root) print(metrics) @@ -183,7 +260,11 @@ async def test_completion_metrics( assert metrics[f"{role}_error_requests"] == 0 assert f"{role}_retry_requests" in metrics - check_historgram(metrics["gen_first_token_latency_seconds"], total_requests, (0.0, 0.3)) + check_historgram( + metrics["gen_first_token_latency_seconds"], + total_requests, + (0.0, 0.3), + ) check_historgram(metrics["gen_complete_latency_seconds"], total_requests, (0.0, 0.6)) assert metrics["total_requests"] == total_requests @@ -219,5 +300,9 @@ async def test_completion_metrics( assert metrics["http_exceptions"] == 0 assert metrics["internal_errors"] == 0 - check_historgram(metrics["gen_complete_latency_seconds"], total_requests * 2, (0.0, 0.6)) + check_historgram( + metrics["gen_complete_latency_seconds"], + total_requests * 2, + (0.0, 0.6), + ) check_historgram(metrics["queue_latency_seconds"], total_requests * 2, (0.0, 0.03)) diff --git a/tests/unittest/llmapi/apps/test_request_metrics.py b/tests/unittest/llmapi/apps/test_request_metrics.py new file mode 100644 index 000000000000..7429e50aaf47 --- /dev/null +++ b/tests/unittest/llmapi/apps/test_request_metrics.py @@ -0,0 +1,400 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json + +import pytest + +from tensorrt_llm.serve.perf_metrics import ( + CTX_CHUNK_METRICS_HEADER, + RETURN_METRICS_HEADER, + SERVER_TIMING_HEADER, + SSE_METRICS_EVENT, + START_END_TIME_HEADER, + STEP_METRICS_HEADER, + PerfMetricsJsonlWriter, + PerfMetricsMiddleware, + _jsonl_record, + build_metrics_headers, + build_metrics_record_from_headers, + combine_disagg_metrics, +) +from tensorrt_llm.serve.scripts.time_breakdown import RequestDataParser + + +def _record(status="complete"): + return { + "schema_version": 1, + "request_id": "42", + "status": status, + "phases": { + "server": { + "timing_metrics": { + "arrival_time": 1.0, + "first_scheduled_time": 1.01, + "first_token_time": 1.02, + "last_token_time": 1.05, + "kv_cache_transfer_start": None, + "kv_cache_transfer_end": None, + }, + "time_breakdown_metrics": { + "step_metrics": [ + { + "iter": 3, + "forward_start_time": 2.0, + "forward_end_time": 2.002, + "sample_start_time": 2.002, + "sample_end_time": 2.003, + "gpu_forward_time": 1.5, + "gpu_sample_time": 0.5, + } + ], + "ctx_chunk_metrics": [ + { + "forward_start_time": 1.0, + "forward_end_time": 1.004, + "sample_start_time": 1.004, + "sample_end_time": 1.005, + "gpu_forward_time": 3.0, + "gpu_sample_time": 0.25, + } + ], + }, + } + }, + } + + +def test_metrics_headers_use_metric_list_syntax(): + headers = build_metrics_headers([_record()]) + + assert "server_queue;dur=10.000000" in headers[SERVER_TIMING_HEADER] + assert "server_ttft;dur=20.000000" in headers[SERVER_TIMING_HEADER] + assert "server_e2e;dur=50.000000" in headers[SERVER_TIMING_HEADER] + assert "server-start;ts=1.000000000" in headers[START_END_TIME_HEADER] + assert "server-end;ts=1.050000000" in headers[START_END_TIME_HEADER] + assert "server-step-3-forward;dur=2.000000" in headers[STEP_METRICS_HEADER] + assert "server-step-3-gpu-sample;dur=0.500000" in headers[STEP_METRICS_HEADER] + assert "server-ctx-chunk-0-forward;dur=4.000000" in headers[CTX_CHUNK_METRICS_HEADER] + + +def test_combine_disagg_metrics_is_request_local(): + ctx = { + "request_id": "ctx-7", + "ctx_request_id": 7, + "metrics_headers": {CTX_CHUNK_METRICS_HEADER: "ctx-ctx-chunk-0-forward;dur=4.000000"}, + "phases": {"server": {"timing_metrics": {"arrival_time": 1.0}}}, + } + gen = { + "request_id": "gen-7", + "ctx_request_id": 7, + "metrics_headers": {STEP_METRICS_HEADER: "gen-step-3-forward;dur=2.000000"}, + "phases": {"server": {"timing_metrics": {"arrival_time": 2.0}}}, + } + + record = combine_disagg_metrics( + "7", + {"timing_metrics": {}}, + ctx, + gen, + disagg_request_id=7, + ) + + assert set(record["phases"]) == {"disagg", "ctx", "gen"} + assert record["disagg_request_id"] == 7 + assert record["phases"]["ctx"]["request_id"] == "ctx-7" + assert record["phases"]["gen"]["ctx_request_id"] == 7 + headers = build_metrics_headers([record]) + assert "ctx-ctx-chunk-0-forward;dur=4.000000" in headers[CTX_CHUNK_METRICS_HEADER] + assert "gen-step-3-forward;dur=2.000000" in headers[STEP_METRICS_HEADER] + + +def test_time_breakdown_parser_accepts_header_derived_disagg_record(): + headers = build_metrics_headers([_record()]) + ctx = build_metrics_record_from_headers(headers, "ctx", request_id="42") + gen = build_metrics_record_from_headers(headers, "gen", request_id="42") + record = combine_disagg_metrics( + "42", + { + "ctx_server": "ctx:8000", + "gen_server": "gen:8000", + "timing_metrics": { + "server_arrival_time": 0.99, + "ctx_dispatch_time": 1.0, + "server_first_token_time": 1.03, + }, + }, + ctx, + gen, + disagg_request_id=42, + ) + + parsed = RequestDataParser().parse_request(_jsonl_record(record), 0) + combined_headers = build_metrics_headers([record]) + + assert parsed["ctx_arrival_time"] == pytest.approx(1.0) + assert parsed["ctx_first_scheduled_time"] == pytest.approx(1.01) + assert parsed["ctx_first_token_time"] == pytest.approx(1.02) + assert parsed["ctx_server_arrival_time"] == pytest.approx(1.0) + assert parsed["ctx_server_first_token_time"] == pytest.approx(1.05) + assert parsed["gen_arrival_time"] == pytest.approx(1.0) + assert parsed["gen_first_scheduled_time"] == pytest.approx(1.01) + assert parsed["gen_first_token_time"] == pytest.approx(1.02) + assert parsed["gen_server_arrival_time"] == pytest.approx(1.0) + assert parsed["gen_server_first_token_time"] == pytest.approx(1.05) + assert parsed["disagg_server_arrival_time"] == pytest.approx(0.99) + assert combined_headers[START_END_TIME_HEADER].count("ctx-start;") == 1 + assert combined_headers[SERVER_TIMING_HEADER].count("ctx_queue;") == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("expose_headers", "request_metrics", "expected"), + [ + (False, True, False), + (True, False, False), + (True, True, True), + ], +) +async def test_middleware_controls_public_headers(expose_headers, request_metrics, expected): + sent = [] + + async def app(scope, receive, send): + scope["state"]["perf_metrics_records"].extend([_record(), _record()]) + await send( + { + "type": "http.response.start", + "status": 200, + "headers": [(b"content-type", b"application/json")], + } + ) + await send( + { + "type": "http.response.body", + "body": b"{}", + "more_body": False, + } + ) + + middleware = PerfMetricsMiddleware(app, expose_headers=expose_headers) + headers = [(RETURN_METRICS_HEADER.encode(), b"1")] if request_metrics else [] + scope = {"type": "http", "headers": headers, "state": {}} + + async def capture(message): + sent.append(message) + + await middleware(scope, None, capture) + + header_names = {key.lower() for key, _ in sent[0]["headers"]} + assert (SERVER_TIMING_HEADER.lower().encode() in header_names) is expected + assert (STEP_METRICS_HEADER.lower().encode() in header_names) is expected + assert (CTX_CHUNK_METRICS_HEADER.lower().encode() in header_names) is expected + if expected: + headers = dict(sent[0]["headers"]) + assert headers[SERVER_TIMING_HEADER.encode()].count(b"server_queue;") == 2 + + +@pytest.mark.asyncio +async def test_middleware_limits_non_streaming_metrics_headers(): + sent = [] + + async def app(scope, receive, send): + record = _record() + breakdown = record["phases"]["server"]["time_breakdown_metrics"] + breakdown["step_metrics"] *= 2000 + breakdown["ctx_chunk_metrics"] *= 2000 + scope["state"]["perf_metrics_records"].append(record) + await send( + { + "type": "http.response.start", + "status": 200, + "headers": [(b"content-type", b"application/json")], + } + ) + await send( + { + "type": "http.response.body", + "body": b"{}", + "more_body": False, + } + ) + + async def capture(message): + sent.append(message) + + middleware = PerfMetricsMiddleware(app, expose_headers=True) + await middleware( + { + "type": "http", + "headers": [(RETURN_METRICS_HEADER.encode(), b"1")], + "state": {}, + }, + None, + capture, + ) + + header_names = {key.lower() for key, _ in sent[0]["headers"]} + assert SERVER_TIMING_HEADER.lower().encode() in header_names + assert START_END_TIME_HEADER.lower().encode() in header_names + assert STEP_METRICS_HEADER.lower().encode() not in header_names + assert CTX_CHUNK_METRICS_HEADER.lower().encode() not in header_names + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("expose_headers", "request_metrics", "expected"), + [ + (False, True, False), + (True, False, False), + (True, True, True), + ], +) +async def test_stream_metrics_follow_done(expose_headers, request_metrics, expected): + sent = [] + + async def app(scope, receive, send): + await send( + { + "type": "http.response.start", + "status": 200, + "headers": [(b"content-type", b"text/event-stream")], + } + ) + await send( + { + "type": "http.response.body", + "body": b"data: [DONE]\n\n", + "more_body": True, + } + ) + scope["state"]["perf_metrics_records"].append(_record()) + await send( + { + "type": "http.response.body", + "body": b"", + "more_body": False, + } + ) + + async def capture(message): + sent.append(message) + + middleware = PerfMetricsMiddleware(app, expose_headers=expose_headers) + headers = [(RETURN_METRICS_HEADER.encode(), b"1")] if request_metrics else [] + await middleware({"type": "http", "headers": headers, "state": {}}, None, capture) + + assert sent[-2]["body"] == b"data: [DONE]\n\n" + has_metrics_event = f"event: {SSE_METRICS_EVENT}".encode() in sent[-1]["body"] + assert has_metrics_event is expected + + +@pytest.mark.asyncio +async def test_disconnect_after_done_is_ignored(): + async def app(scope, receive, send): + await send( + { + "type": "http.response.start", + "status": 200, + "headers": [(b"content-type", b"text/event-stream")], + } + ) + await send( + { + "type": "http.response.body", + "body": b"data: [DONE]\n\n", + "more_body": True, + } + ) + scope["state"]["perf_metrics_records"].append(_record()) + await send( + { + "type": "http.response.body", + "body": b"", + "more_body": False, + } + ) + + async def disconnect(message): + if message["type"] == "http.response.body" and not message.get("more_body", False): + raise OSError("client disconnected") + + middleware = PerfMetricsMiddleware(app, expose_headers=True) + await middleware( + { + "type": "http", + "headers": [(RETURN_METRICS_HEADER.encode(), b"1")], + "state": {}, + }, + None, + disconnect, + ) + + +@pytest.mark.asyncio +async def test_file_middleware_intercepts_detail_headers(tmp_path): + writer = PerfMetricsJsonlWriter(str(tmp_path), "test") + await writer.start() + + async def app(scope, receive, send): + records = [_record(), _record()] + records[0]["disagg_request_id"] = 17 + records[1]["disagg_request_id"] = 18 + scope["state"]["perf_metrics_records"].extend(records) + await send( + { + "type": "http.response.start", + "status": 200, + "headers": [(b"content-type", b"application/json")], + } + ) + await send( + { + "type": "http.response.body", + "body": b"{}", + "more_body": False, + } + ) + + sent = [] + + async def capture(message): + sent.append(message) + + middleware = PerfMetricsMiddleware(app, expose_headers=False, writer=writer) + await middleware({"type": "http", "headers": [], "state": {}}, None, capture) + await writer.close() + + header_names = {key.lower() for key, _ in sent[0]["headers"]} + assert STEP_METRICS_HEADER.lower().encode() not in header_names + assert CTX_CHUNK_METRICS_HEADER.lower().encode() not in header_names + + output_file = next(tmp_path.glob("perf_metrics-test-*.jsonl")) + saved = [json.loads(line) for line in output_file.read_text().splitlines()] + assert [record["disagg_request_id"] for record in saved] == [17, 18] + assert saved[0]["time_breakdown_metrics"]["step_metrics"] + assert saved[0]["time_breakdown_metrics"]["ctx_chunk_metrics"] + + +@pytest.mark.asyncio +async def test_jsonl_writer_drops_only_malformed_record(tmp_path): + writer = PerfMetricsJsonlWriter(str(tmp_path), "test") + await writer.start() + writer.submit({"phases": {}}) + writer.submit(_record()) + await writer.close() + + output_file = next(tmp_path.glob("perf_metrics-test-*.jsonl")) + records = [json.loads(line) for line in output_file.read_text().splitlines()] + assert writer.dropped_records == 1 + assert records[0]["request_id"] == 42 diff --git a/tests/unittest/others/test_time_breakdown.py b/tests/unittest/others/test_time_breakdown.py index ae45744cade5..412b95734160 100644 --- a/tests/unittest/others/test_time_breakdown.py +++ b/tests/unittest/others/test_time_breakdown.py @@ -16,6 +16,8 @@ import unittest from unittest.mock import patch +from tensorrt_llm.serve.scripts.benchmark_serving import ( + _read_new_perf_metrics, _snapshot_perf_metrics) from tensorrt_llm.serve.scripts.time_breakdown import (RequestDataParser, RequestTimeBreakdown, TimingMetric, @@ -393,9 +395,50 @@ def test_parse_json_file(self): finally: os.unlink(temp_file) + def test_parse_jsonl_file(self): + """Test parsing server-produced JSONL records.""" + with tempfile.NamedTemporaryFile(mode='w', + suffix='.jsonl', + delete=False) as f: + for record in self.test_data: + f.write(json.dumps(record) + '\n') + temp_file = f.name + + try: + timing_data = self.analyzer.parse_json_file(temp_file) + self.assertEqual(len(timing_data), 2) + self.assertEqual(timing_data[0]['ctx_server_arrival_time'], 1.0) + self.assertEqual(timing_data[1]['ctx_server_arrival_time'], 2.0) + finally: + os.unlink(temp_file) + + def test_read_new_disagg_metrics_for_benchmark(self): + """Test reading only new combined records from a metrics directory.""" + with tempfile.TemporaryDirectory() as output_dir: + disagg_path = os.path.join( + output_dir, "perf_metrics-disagg-host-1-start.jsonl") + with open(disagg_path, "w", encoding="utf-8") as output: + output.write('{"disagg_request_id":0}\n') + offsets = _snapshot_perf_metrics(output_dir) + + with open(disagg_path, "a", encoding="utf-8") as output: + output.write('{"disagg_request_id":1}\n') + for kind in ("context", "generation"): + worker_path = os.path.join( + output_dir, f"perf_metrics-{kind}-host-2-start.jsonl") + with open(worker_path, "w", encoding="utf-8") as output: + output.write('{"request_id":1}\n') + + records = _read_new_perf_metrics(output_dir, + offsets, + expected_count=1, + timeout=0.1) + + self.assertEqual(records, [{"disagg_request_id": 1}]) + def test_parse_json_file_not_found(self): """Test parsing a non-existent file.""" - with self.assertRaises(SystemExit): + with self.assertRaises(FileNotFoundError): self.analyzer.parse_json_file('non_existent_file.json') def test_parse_json_file_invalid_json(self): @@ -406,7 +449,7 @@ def test_parse_json_file_invalid_json(self): temp_file = f.name try: - with self.assertRaises(SystemExit): + with self.assertRaises(ValueError): self.analyzer.parse_json_file(temp_file) finally: os.unlink(temp_file)