tensorrt_llm/serve/openai_server.py (2)

85-87: Bound the in-memory perf metrics buffer to avoid unbounded growth

If clients don’t poll /perf_metrics frequently (or ever), self.perf_metrics can grow without bound. Consider a bounded deque to cap memory usage.

Apply within the shown range:

-        self.perf_metrics = []
+        self.perf_metrics = deque(maxlen=1000)
         self.perf_metrics_lock = asyncio.Lock()

Add the import at the top of the file (outside the shown range):

from collections import deque

---

245-281: Harden serialization against optional/absent fields

speculative_decoding and some timing/KV-cache fields may be absent or None depending on backend/features. Accessing attributes unguarded can raise. Recommend using getattr to guard optional fields.

Apply within the shown range:

-        for metrics_dict in perf_metrics:
+        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
             # TODO: other metrics
             metrics_json = {
                 "first_iter": metrics.first_iter,
                 "last_iter": metrics.last_iter,
-                "arrival_time": timing_metrics.arrival_time.total_seconds(),
-                "first_scheduled_time": timing_metrics.first_scheduled_time.total_seconds(),
-                "first_token_time": timing_metrics.first_token_time.total_seconds(),
-                "last_token_time": timing_metrics.last_token_time.total_seconds(),
+                "arrival_time": getattr(timing_metrics, "arrival_time").total_seconds(),
+                "first_scheduled_time": getattr(timing_metrics, "first_scheduled_time").total_seconds(),
+                "first_token_time": getattr(timing_metrics, "first_token_time").total_seconds(),
+                "last_token_time": getattr(timing_metrics, "last_token_time").total_seconds(),
                 "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:
+            kv_cache_size = getattr(timing_metrics, "kv_cache_size", 0)
+            if kv_cache_size and kv_cache_size > 0:
                 metrics_json.update({
-                    "kv_cache_size": timing_metrics.kv_cache_size,
-                    "kv_cache_transfer_start": timing_metrics.kv_cache_transfer_start.total_seconds(),
-                    "kv_cache_transfer_end": timing_metrics.kv_cache_transfer_end.total_seconds(),
+                    "kv_cache_size": kv_cache_size,
+                    **(
+                        {"kv_cache_transfer_start": getattr(timing_metrics, "kv_cache_transfer_start").total_seconds()}
+                        if getattr(timing_metrics, "kv_cache_transfer_start", None) is not None else {}
+                    ),
+                    **(
+                        {"kv_cache_transfer_end": getattr(timing_metrics, "kv_cache_transfer_end").total_seconds()}
+                        if getattr(timing_metrics, "kv_cache_transfer_end", None) is not None else {}
+                    ),
                 })
-            if speculative_decoding.total_draft_tokens > 0:
+            if getattr(speculative_decoding, "total_draft_tokens", 0) > 0:
                 metrics_json.update({
-                    "acceptance_rate": speculative_decoding.acceptance_rate,
-                    "total_accepted_draft_tokens": speculative_decoding.total_accepted_draft_tokens,
-                    "total_draft_tokens": speculative_decoding.total_draft_tokens,
+                    "acceptance_rate": getattr(speculative_decoding, "acceptance_rate"),
+                    "total_accepted_draft_tokens": getattr(speculative_decoding, "total_accepted_draft_tokens"),
+                    "total_draft_tokens": getattr(speculative_decoding, "total_draft_tokens"),
                 })
             metrics_dict["perf_metrics"] = metrics_json
         return JSONResponse(content=perf_metrics)

tests/unittest/llmapi/apps/_test_openai_perf_metrics.py (1)

61-66: Use equality (==) instead of identity (is) and add a short poll to avoid flakes

response.status is 200 relies on CPython small-int interning and is not the intended check; use ==. Also, adding a brief retry loop reduces flakiness if metrics are exposed slightly after the completion returns.

Apply within the shown range:

-    response = urlopen(f'{server.url_root}/perf_metrics')
-    assert response.status is 200
-
-    data_list = json.loads(response.read())
-    assert len(data_list) == 1
+    # Poll a few times in case metric extraction lags the response
+    data_list = []
+    for _ in range(10):
+        with urlopen(f'{server.url_root}/perf_metrics') as response:
+            assert response.status == 200
+            data_list = json.loads(response.read())
+        if len(data_list) >= 1:
+            break
+        time.sleep(0.2)
+    assert len(data_list) == 1

Add this import at the top of the file (outside the shown range):

import time

tensorrt_llm/serve/openai_disagg_server.py (1)

137-146: Consider parallelizing fan-out GETs and handling non-200s gracefully

The current loop fetches each server’s /perf_metrics sequentially and assumes success. For lower latency and resilience under partial failures, gather requests concurrently and skip/record servers that return non-200 or raise.

If helpful, I can provide a concurrent asyncio.gather version with per-server error handling.

tests/integration/defs/disaggregated/test_disaggregated.py (5)

20-20: Import Optional for accurate type annotation usage

You assign None to a Callable-typed parameter later. Bring in Optional to annotate it correctly.

Apply this diff:

-from typing import Callable
+from typing import Callable, Optional

---

39-45: Add return type, docstring, and explicit UTF-8 decoding

Small polish to improve clarity and robustness when reading YAML.

Apply this diff:

-def get_disagg_server_url_from_cfg(config_file: str):
-    with open(config_file, 'r') as file:
+def get_disagg_server_url_from_cfg(config_file: str) -> str:
+    """Derive the disaggregated server base URL from a YAML config."""
+    with open(config_file, 'r', encoding='utf-8') as file:
         config = yaml.safe_load(file)
     server_host = config.get('hostname', 'localhost')
     server_port = config.get('port', 8000)
     return f"http://{server_host}:{server_port}"

---

165-166: Make extra_endpoints_test annotation Optional to match default None

Type checkers will flag Callable[[str], None] with a None default. Use Optional[Callable[…]].

Apply this diff:

-                           prompt_file="prompts.json",
-                           extra_endpoints_test: Callable[[str], None] = None):
+                           prompt_file="prompts.json",
+                           extra_endpoints_test: Optional[Callable[[str], None]] = None):

---

246-249: Avoid re-reading the YAML config on every iteration

Minor inefficiency: the server URL is constant for the test. Compute it once outside the loop.

Apply this diff within the shown lines:

-                if extra_endpoints_test is not None:
-                    extra_endpoints_test(
-                        get_disagg_server_url_from_cfg(config_file))
+                if extra_endpoints_test is not None:
+                    extra_endpoints_test(server_url)

And add this one-liner just before entering the workers/server context (or right after computing server_cmd), outside the per-iteration loop:

server_url = get_disagg_server_url_from_cfg(config_file)

---

561-571: Relax strict temporal assertions to reduce flakiness across schedulers/clocks

Timing can be equal at boundaries due to clock resolution and ordering. Using < may intermittently fail; prefer <= for adjacent events. Keep strict ordering where causality must hold.

Apply this diff:

-        assert ctx_metrics["first_token_time"] == ctx_metrics["last_token_time"]
-        assert ctx_metrics["last_token_time"] < gen_metrics["arrival_time"]
+        assert ctx_metrics["first_token_time"] == ctx_metrics["last_token_time"]
+        assert ctx_metrics["last_token_time"] <= gen_metrics["arrival_time"]
@@
-        assert gen_metrics["arrival_time"] < gen_metrics[
+        assert gen_metrics["arrival_time"] <= gen_metrics[
             "kv_cache_transfer_start"]
-        assert gen_metrics["kv_cache_transfer_start"] < gen_metrics[
+        assert gen_metrics["kv_cache_transfer_start"] <= gen_metrics[
             "kv_cache_transfer_end"]
-        assert gen_metrics["kv_cache_transfer_end"] < gen_metrics[
+        assert gen_metrics["kv_cache_transfer_end"] <= gen_metrics[
             "first_scheduled_time"]

If ctx/gen servers ever run on different hosts with unsynchronized clocks, consider validating relative durations (non-negative deltas) instead of absolute cross-server timestamps.

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

tensorrt_llm/serve/openai_disagg_server.py (1)

418-445: Python 3.8 compatibility: replace list[str] with List[str]

Same PEP 585 generics issue here; update annotations for method and inner helper.

Apply this diff:

-    async def wait_for_all_servers_ready(cls, session: aiohttp.ClientSession,
-                                         ctx_servers: list[str],
-                                         gen_servers: list[str],
+    async def wait_for_all_servers_ready(cls, session: aiohttp.ClientSession,
+                                         ctx_servers: List[str],
+                                         gen_servers: List[str],
                                          server_start_timeout_secs: int = 180):
-        async def get_unready_servers(servers: list[str]) -> list[str]:
+        async def get_unready_servers(servers: List[str]) -> List[str]:

Ensure typing imports include List (already present).

tensorrt_llm/serve/openai_disagg_server.py (2)

57-63: Python 3.8 compatibility: replace PEP 585 generics with typing.List/Dict/Tuple (duplicate of prior review)

Using list[...]/dict[...] requires Python 3.9+. The project targets 3.8+. Also tighten the nested mapping type for ctx_request_id keys.

Apply within the shown range:

-        # record corresponding keys of context and generation servers for perf metrics
-        self.perf_metrics_keys: list[tuple[str, str, int]] = [] # (ctx_server_key, gen_server_key, ctx_request_id)
+        # record corresponding keys of context and generation servers for perf metrics
+        self.perf_metrics_keys: List[Tuple[str, str, int]] = []  # (ctx_server_key, gen_server_key, ctx_request_id)
         self.perf_metrics_keys_lock = asyncio.Lock()
-        self.server_perf_metrics: dict[str, dict] = {} # server_key -> {ctx_request_id: perf_metrics}
+        self.server_perf_metrics: Dict[str, Dict[int, dict]] = {}  # server_key -> {ctx_request_id: perf_metrics}
         for server in self.ctx_servers + self.gen_servers:
             self.server_perf_metrics[server] = {}

And update typing imports at the top of the file:

-from typing import List, Optional, Type, Union
+from typing import Dict, List, Optional, Tuple, Type, Union

---

145-166: Fix data loss and partial-pair emission when correlating ctx/gen metrics (duplicate of prior review)

Current logic pops gen first, then appends even if ctx is missing (ctx_perf_metrics can be None). This both drops ctx data prematurely and can emit incomplete pairs. Peek (get) until both present, then pop both.

Apply this diff:

-            remain_keys = []
-            for ctx_server, gen_server, ctx_request_id in self.perf_metrics_keys:
-                gen_perf_metrics = self.server_perf_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))
-                    continue
-                ctx_perf_metrics = self.server_perf_metrics[ctx_server].pop(ctx_request_id, None)
-                return_metrics.append({
-                    "ctx_server": ctx_server,
-                    "gen_server": gen_server,
-                    "ctx_perf_metrics": ctx_perf_metrics,
-                    "gen_perf_metrics": gen_perf_metrics})
-            self.perf_metrics_keys = remain_keys
+            remain_keys = []
+            for ctx_server, gen_server, ctx_request_id in self.perf_metrics_keys:
+                ctx_perf_metrics = self.server_perf_metrics.get(ctx_server, {}).get(ctx_request_id)
+                gen_perf_metrics = self.server_perf_metrics.get(gen_server, {}).get(ctx_request_id)
+                if not (ctx_perf_metrics and gen_perf_metrics):
+                    # one side not finished yet
+                    remain_keys.append((ctx_server, gen_server, ctx_request_id))
+                    continue
+                # both present; remove from stores to avoid duplication
+                self.server_perf_metrics[ctx_server].pop(ctx_request_id, None)
+                self.server_perf_metrics[gen_server].pop(ctx_request_id, None)
+                return_metrics.append({
+                    "ctx_server": ctx_server,
+                    "gen_server": gen_server,
+                    "ctx_perf_metrics": ctx_perf_metrics,
+                    "gen_perf_metrics": gen_perf_metrics,
+                })
+            self.perf_metrics_keys = remain_keys

tests/integration/defs/disaggregated/test_disaggregated.py (1)

529-577: Complete prior fix: prefer json.load(resp) and verify gen metrics schema

You already fixed the urllib import and added a timeout. Finish by parsing JSON directly from the response. Also, confirm whether kv_cache_size resides at gen_perf_metrics top-level or under perf_metrics.

Apply this diff for JSON parsing:

-        with urllib.request.urlopen(f"{server_url}/perf_metrics",
-                                    timeout=10) as resp:
+        with urllib.request.urlopen(f"{server_url}/perf_metrics", timeout=10) as resp:
             assert resp.status == 200
-            perf_metrics = json.loads(resp.read())
+            perf_metrics = json.load(resp)

If kv_cache_size is a gen-level field (as described in the PR), adjust assertions accordingly:

• Use gen_item = item["gen_perf_metrics"]
• Keep gen_metrics = gen_item["perf_metrics"] for timing fields
• Assert "kv_cache_size" in gen_item and related ordering on gen_item fields.

To verify the schema location from the codebase, run:

#!/bin/bash
# Inspect where kv_cache_size is attached in server responses
rg -n -C3 --type=py -P '\bkv_cache_size\b' 

This should show whether kv_cache_size is included inside "perf_metrics" or at the top level of the gen-server metrics object so the test can assert the right path.

tensorrt_llm/serve/openai_disagg_server.py (2)

306-309: Avoid duplicate correlation keys

Repeated appends of the same (ctx_server, gen_server, ctx_request_id) can cause unnecessary work. Guard against duplicates.

Apply this diff:

-            if need_ctx:
-                async with self.perf_metrics_keys_lock:
-                    self.perf_metrics_keys.append((ctx_server, gen_server, req.disaggregated_params.ctx_request_id))
+            if need_ctx:
+                async with self.perf_metrics_keys_lock:
+                    key = (ctx_server, gen_server, req.disaggregated_params.ctx_request_id)
+                    if key not in self.perf_metrics_keys:
+                        self.perf_metrics_keys.append(key)

---

1-1: Missing NVIDIA copyright header

Per repository guidelines, prepend the NVIDIA SPDX header to this source file.

Example:

# SPDX-FileCopyrightText: Copyright (c) 2022-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0

tests/integration/defs/disaggregated/test_disaggregated.py (1)

39-45: Make YAML read robust and add return type annotation

Add explicit return type, set encoding, and guard against empty files.

Apply this diff:

-def get_disagg_server_url_from_cfg(config_file: str):
-    with open(config_file, 'r') as file:
-        config = yaml.safe_load(file)
+def get_disagg_server_url_from_cfg(config_file: str) -> str:
+    with open(config_file, 'r', encoding='utf-8') as file:
+        config = yaml.safe_load(file) or {}
     server_host = config.get('hostname', 'localhost')
     server_port = config.get('port', 8000)
     return f"http://{server_host}:{server_port}"

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

tests/integration/defs/disaggregated/test_disaggregated.py (1)

529-577: Tighten JSON parsing and set a slightly more forgiving timeout

Use json.load on the response file-like object and keep the explicit urllib.request import. Your structure assertions are good.

Apply this diff:

-        import json
-        import urllib.request
+        import json
+        import urllib.request
@@
-        with urllib.request.urlopen(f"{server_url}/perf_metrics",
-                                    timeout=10) as resp:
+        with urllib.request.urlopen(f"{server_url}/perf_metrics",
+                                    timeout=30) as resp:
             assert resp.status == 200
-            perf_metrics = json.loads(resp.read())
+            perf_metrics = json.load(resp)

tensorrt_llm/serve/openai_disagg_server.py (1)

11-11: Use typing generics for Python 3.8 compatibility

The codebase targets Python 3.8+, where PEP 585 built-in generics (list[str], dict[str, …], tuple[...]) are not supported. Replace them with typing.List/Dict/Tuple and import the needed symbols.

Apply this diff to the imports and changed annotations in this file:

-from typing import Optional, Type, Union
+from typing import Dict, List, Optional, Tuple, Type, Union

And update usages within the shown ranges:

-            self.server_perf_metrics: dict[str, dict[int, dict]] = {}
+            self.server_perf_metrics: Dict[str, Dict[int, dict]] = {}

For other occurrences in this file (e.g., method annotations using list[str]), similarly convert to List[str]/Tuple[...].

</blockquote></details>

</blockquote></details>

<details>
<summary>🧹 Nitpick comments (9)</summary><blockquote>

<details>
<summary>tensorrt_llm/serve/openai_server.py (4)</summary><blockquote>

`88-100`: **Bounded queue for metrics is good; guard against None config**

Deque with maxlen addresses unbounded growth nicely. To avoid surprises if the config remains Optional in some environments, coerce None to 0 before comparison.

Apply this diff:

```diff
-            max_perf_metrics = self.llm.args.perf_metrics_max_requests
+            max_perf_metrics = int(self.llm.args.perf_metrics_max_requests or 0)

---

169-169: Endpoint naming is fine; consider consistency with config knob

Route name /perf_metrics matches functionality. If you want stronger alignment with the controlling flag (return_perf_metrics), leaving as-is is ok. A rename to /request_perf_metrics could be considered later, but not required now.

---

266-309: Safe swap-and-drain pattern; minor structural suggestion

Swapping out the deque under lock and then formatting outside the lock is correct. A minor ergonomics improvement: instead of mutating elements in-place (metrics_dict["perf_metrics"] = metrics_json), build the final list to avoid mixing raw and formatted entries if future code reads the deque during formatting.

---

320-336: Ensure request_perf_metrics is present before accessing

Accessing res.outputs[0].request_perf_metrics assumes metrics were collected. If return_perf_metrics is true but upstream didn’t populate request_perf_metrics, this will raise. Add a guard to skip appending when it’s None.

Apply this diff:

-            item = {
-                "request_id": res.request_id,
-                "perf_metrics": res.outputs[0].request_perf_metrics
-            }
+            req_perf = res.outputs[0].request_perf_metrics
+            if req_perf is None:
+                return
+            item = {
+                "request_id": res.request_id,
+                "perf_metrics": req_perf,
+            }

tensorrt_llm/serve/openai_disagg_server.py (3)

146-157: Harden remote fetches: check HTTP status and skip failures

A single failing upstream /perf_metrics currently bubbles up as an exception after buffering. Consider skipping servers that return non-200 or raise client errors, logging a warning, so the endpoint remains robust.

Apply this diff:

-        try:
-            for server in self.ctx_servers + self.gen_servers:
-                async with self.session.get(f"{server}/perf_metrics") as response:
-                    server_perf_metrics = await response.json()
-                    perf_metrics[server] = server_perf_metrics
-        except Exception as e:
-            # Keep the exception to raise it after saving perf metrics
-            exc = e
+        try:
+            for server in self.ctx_servers + self.gen_servers:
+                try:
+                    async with self.session.get(f"{server}/perf_metrics") as response:
+                        if response.status != 200:
+                            logger.warning(f"Skipping {server}/perf_metrics: HTTP {response.status}")
+                            continue
+                        server_perf_metrics = await response.json()
+                        perf_metrics[server] = server_perf_metrics
+                except aiohttp.ClientError as e:
+                    logger.error(f"Failed to fetch {server}/perf_metrics: {e}")
+                    continue
+        except Exception as e:
+            # Keep the exception to raise it after saving perf metrics
+            exc = e

---

177-191: Avoid data loss when correlating: pop only after both sides present

You already pop the gen side only after checking its presence; consider mirroring the get-then-pop pattern for both sides to avoid races. This is fine currently but can be made more explicit and symmetric.

Apply this diff:

-                gen_perf_metrics = self.server_perf_metrics[gen_server].pop(ctx_request_id, None)
-                if gen_perf_metrics is None:
+                gen_perf_metrics = self.server_perf_metrics[gen_server].get(ctx_request_id)
+                if gen_perf_metrics is None:
                     # generation not finished
                     remain_keys.append((ctx_server, gen_server, ctx_request_id))
                     continue
-                ctx_perf_metrics = self.server_perf_metrics[ctx_server].pop(ctx_request_id, None)
+                ctx_perf_metrics = self.server_perf_metrics[ctx_server].get(ctx_request_id)
+                # both present; remove to avoid duplication next time
+                self.server_perf_metrics[gen_server].pop(ctx_request_id, None)
+                self.server_perf_metrics[ctx_server].pop(ctx_request_id, None)

---

118-123: create_error_response type is misleading and always raises

Signature says -> ErrorResponse, but it unconditionally raises HTTPException. Either return a JSONResponse with an ErrorResponse body or change the signature/name to reflect raising.

Apply this diff to return a proper response:

-    def create_error_response(
-            message: str,
-            status_code: HTTPStatus = HTTPStatus.BAD_REQUEST) -> ErrorResponse:
-        raise HTTPException(status_code=500, detail=f"Internal server error {message}")
+    def create_error_response(
+            message: str,
+            status_code: HTTPStatus = HTTPStatus.BAD_REQUEST) -> JSONResponse:
+        payload = {"error": message, "code": status_code.value}
+        return JSONResponse(content=payload, status_code=status_code.value)

And ensure call sites use this response (or keep raising where appropriate).

tensorrt_llm/commands/serve.py (2)

1-1: Missing NVIDIA copyright header.

Per repository guidelines, prepend the current year NVIDIA header to all source files.

Add at the very top:

# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.

---

206-210: Help text contradicts the default backend.

Help says “Default is cpp path” while default is "pytorch".

Consider:

-              help="Set to 'pytorch' for pytorch path. Default is cpp path.")
+              help="Set to 'pytorch' for PyTorch path. Default is 'pytorch'.")

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

tensorrt_llm/serve/openai_disagg_server.py (3)

150-160: Harden cross-server fetch: check HTTP status and skip failures (don’t raise for a single bad server)

A failing server currently bubbles an exception after caching metrics, breaking the whole endpoint. Skip failed/non-200 servers and continue aggregating.

Apply this diff:

-        perf_metrics = {}
-        exc = None
-        try:
-            for server in self.ctx_servers + self.gen_servers:
-                async with self.session.get(f"{server}/perf_metrics") as response:
-                    server_perf_metrics = await response.json()
-                    perf_metrics[server] = server_perf_metrics
-        except Exception as e:
-            # Keep the exception to raise it after saving perf metrics
-            exc = e
+        perf_metrics = {}
+        for server in self.ctx_servers + self.gen_servers:
+            try:
+                async with self.session.get(f"{server}/perf_metrics") as response:
+                    if response.status != 200:
+                        logger.warning(f"Skipping {server}/perf_metrics: HTTP {response.status}")
+                        continue
+                    server_perf_metrics = await response.json()
+                    perf_metrics[server] = server_perf_metrics
+            except aiohttp.ClientError as e:
+                logger.error(f"Failed to fetch {server}/perf_metrics: {e}")
+                continue

And remove the exc handling below since we no longer re-raise.

---

11-12: Use typing.Dict/List/Tuple for Python 3.8 compatibility

PEP 585 built-in generics in type hints aren’t supported on 3.8. Import and use typing equivalents.

Apply this diff:

-from typing import Optional, Type, Union
+from typing import Dict, List, Optional, Tuple, Type, Union

Additionally, update annotations below as suggested in related diffs.

---

178-191: Avoid data loss when correlating ctx/gen metrics (don’t pop until both present)

If gen is ready but ctx isn’t, current code pops gen and drops the pair. Peek for both; only pop when both are present.

Apply this diff:

-            remain_keys = []
-            for ctx_server, gen_server, ctx_request_id in self.perf_metrics_keys:
-                gen_perf_metrics = self.server_perf_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))
-                    continue
-                ctx_perf_metrics = self.server_perf_metrics[ctx_server].pop(ctx_request_id, None)
-                return_metrics.append({
-                    "ctx_server": ctx_server,
-                    "gen_server": gen_server,
-                    "ctx_perf_metrics": ctx_perf_metrics,
-                    "gen_perf_metrics": gen_perf_metrics})
+            remain_keys = []
+            for ctx_server, gen_server, ctx_request_id in self.perf_metrics_keys:
+                ctx_perf_metrics = self.server_perf_metrics.get(ctx_server, {}).get(ctx_request_id)
+                gen_perf_metrics = self.server_perf_metrics.get(gen_server, {}).get(ctx_request_id)
+                if not (ctx_perf_metrics and gen_perf_metrics):
+                    # one side not finished yet
+                    remain_keys.append((ctx_server, gen_server, ctx_request_id))
+                    continue
+                # both present; remove and return
+                self.server_perf_metrics[ctx_server].pop(ctx_request_id, None)
+                self.server_perf_metrics[gen_server].pop(ctx_request_id, None)
+                return_metrics.append({
+                    "ctx_server": ctx_server,
+                    "gen_server": gen_server,
+                    "ctx_perf_metrics": ctx_perf_metrics,
+                    "gen_perf_metrics": gen_perf_metrics})

tensorrt_llm/llmapi/disagg_utils.py (1)

71-76: Validate server configs to avoid invalid URLs (http://None:None)

If hostname/port are unset, this builds invalid URLs. Fail fast with a clear error.

Apply this diff:

     ctx_server_urls = []
     gen_server_urls = []
     for cfg in server_configs:
-        if cfg.type == "ctx":
-            ctx_server_urls.append(f"http://{cfg.hostname}:{cfg.port}")
-        else:
-            gen_server_urls.append(f"http://{cfg.hostname}:{cfg.port}")
+        if cfg.hostname is None or cfg.port is None:
+            raise ValueError(f"Missing hostname/port for server config: {cfg}")
+        if cfg.type == "ctx":
+            ctx_server_urls.append(f"http://{cfg.hostname}:{cfg.port}")
+        elif cfg.type == "gen":
+            gen_server_urls.append(f"http://{cfg.hostname}:{cfg.port}")
+        else:
+            raise ValueError(f"Unknown server type: {cfg.type}")

tensorrt_llm/serve/openai_disagg_server.py (2)

36-36: Fix comment typo

Nit: “enale” -> “enable”.

Apply this diff:

-# yapf: enale
+# yapf: enable

---

401-401: Use typing.Callable for create_generator type hint

Minor typing cleanup; also import Callable from typing.

Apply this diff:

-                           response_type: Type[Union[CompletionResponse, ChatCompletionResponse]],
-                           create_generator: callable) -> Union[CompletionResponse, ChatCompletionResponse, StreamingResponse]:
+                           response_type: Type[Union[CompletionResponse, ChatCompletionResponse]],
+                           create_generator: Callable) -> Union[CompletionResponse, ChatCompletionResponse, StreamingResponse]:

Outside this range (imports):

from typing import Callable

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

tensorrt_llm/serve/openai_disagg_server.py (7)

446-453: Python 3.8 typing: replace list[str] with List[str]

Adhere to the 3.8+ target and repo policy.

-    async def wait_for_all_servers_ready(cls, session: aiohttp.ClientSession,
-                                         ctx_servers: list[str],
-                                         gen_servers: list[str],
+    async def wait_for_all_servers_ready(cls, session: aiohttp.ClientSession,
+                                         ctx_servers: List[str],
+                                         gen_servers: List[str],
                                          server_start_timeout_secs: int = 180):
-        async def get_unready_servers(servers: list[str]) -> list[str]:
+        async def get_unready_servers(servers: List[str]) -> List[str]:

---

11-11: Typing imports incomplete; add typing generics for 3.8 compatibility

You use Dict/List/Tuple/Deque/Callable/AsyncIterator in suggested fixes and to avoid PEP 585 built-ins. Extend imports.

-from typing import Optional, Type, Union
+from typing import AsyncIterator, Callable, Deque, Dict, List, Optional, Tuple, Type, Union

---

59-63: Python 3.8 compatibility: avoid PEP 585 built-in generics and add precise annotations

Replace built-in generics and annotate the deque precisely. Also keep server_perf_metrics’ value type stable.

-            # (ctx_server, gen_server, ctx_request_id)
-            self.perf_metrics_keys = deque(maxlen=self.perf_metrics_max_requests)
+            # (ctx_server, gen_server, ctx_request_id)
+            self.perf_metrics_keys: Deque[Tuple[str, str, int]]
+            self.perf_metrics_keys = deque(maxlen=self.perf_metrics_max_requests)
             self.perf_metrics_keys_lock = asyncio.Lock()
             # server_key -> {ctx_request_id: perf_metrics}
-            self.server_perf_metrics: dict[str, dict[int, dict]] = {}
+            self.server_perf_metrics: Dict[str, Dict[int, dict]] = {}

---

180-193: Fix data loss when correlating ctx/gen metrics: only pop after both sides are present

Current flow pops gen first; if ctx isn’t ready yet, you emit a record with ctx=None and drop gen, losing correlation forever. Peek with get(), only pop when both exist, and retain the key otherwise.

-            remain_keys = []
-            for ctx_server, gen_server, ctx_request_id in self.perf_metrics_keys:
-                gen_perf_metrics = self.server_perf_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))
-                    continue
-                ctx_perf_metrics = self.server_perf_metrics[ctx_server].pop(ctx_request_id, None)
-                return_metrics.append({
-                    "ctx_server": ctx_server,
-                    "gen_server": gen_server,
-                    "ctx_perf_metrics": ctx_perf_metrics,
-                    "gen_perf_metrics": gen_perf_metrics})
+            remain_keys = []
+            for ctx_server, gen_server, ctx_request_id in self.perf_metrics_keys:
+                ctx_store = self.server_perf_metrics.get(ctx_server, {})
+                gen_store = self.server_perf_metrics.get(gen_server, {})
+                ctx_perf_metrics = ctx_store.get(ctx_request_id)
+                gen_perf_metrics = gen_store.get(ctx_request_id)
+                if not (ctx_perf_metrics and gen_perf_metrics):
+                    # one side not finished yet; keep the key
+                    remain_keys.append((ctx_server, gen_server, ctx_request_id))
+                    continue
+                # both present; now remove from stores to avoid duplication
+                ctx_store.pop(ctx_request_id, None)
+                gen_store.pop(ctx_request_id, None)
+                return_metrics.append({
+                    "ctx_server": ctx_server,
+                    "gen_server": gen_server,
+                    "ctx_perf_metrics": ctx_perf_metrics,
+                    "gen_perf_metrics": gen_perf_metrics,
+                })

---

1-1: Add NVIDIA SPDX header per repo guidelines

Source files must include the NVIDIA copyright/SPDX header.

+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
 #!/usr/bin/env python

---

59-63: Audit and update PEP 585 built-in generics for Python 3.8 compatibility

The repository currently uses built-in generics (e.g. list[str], dict[str, …], tuple[int, …], deque[int]) which are only supported on Python 3.9+. Since our target is Python 3.8, these must be replaced with the corresponding types from the typing module.

• In tensorrt_llm/serve/openai_disagg_server.py
– Lines 59–63:
```diff

• self.server_perf_metrics: dict[str, dict[int, dict]] = {}

• from typing import Dict
• self.server_perf_metrics: Dict[str, Dict[int, Dict]] = {}

– Lines 446–453 (method signature):
```diff

• async def get_unready_servers(servers: list[str]) -> list[str]:

• from typing import List
• async def get_unready_servers(servers: List[str]) -> List[str]:

• Apply the same conversion across the codebase wherever you see:

• list[...] → typing.List[...]
• dict[...] → typing.Dict[...]
• set[...] → typing.Set[...]
• tuple[...] → typing.Tuple[...]
• deque[...] → typing.Deque[...]

• Ensure you add the necessary imports, e.g.

from typing import List, Dict, Tuple, Set, Deque
from collections import deque  # if using deque instances

This refactor is mandatory to maintain compatibility with Python 3.8.

---

146-160: Harden /perf_metrics aggregation: check HTTP status, handle errors per-server, don’t raise global exception

As-is, any exception in the loop defers raising and then 500s the endpoint after partially updating caches. Prefer skipping failed servers and proceeding. Also check response.status before parsing JSON.

-        perf_metrics = {}
-        exc = None
-        try:
-            for server in self.ctx_servers + self.gen_servers:
-                async with self.session.get(f"{server}/perf_metrics") as response:
-                    server_perf_metrics = await response.json()
-                    perf_metrics[server] = server_perf_metrics
-        except Exception as e:
-            # Keep the exception to raise it after saving perf metrics
-            exc = e
+        perf_metrics: Dict[str, list] = {}
+        for server in self.ctx_servers + self.gen_servers:
+            try:
+                async with self.session.get(f"{server}/perf_metrics") as response:
+                    if response.status != 200:
+                        logger.warning(f"Skipping {server}/perf_metrics: HTTP {response.status}")
+                        continue
+                    try:
+                        server_perf_metrics = await response.json()
+                    except Exception as je:
+                        logger.warning(f"Invalid JSON from {server}/perf_metrics: {je}")
+                        continue
+                    perf_metrics[server] = server_perf_metrics
+            except aiohttp.ClientError as ce:
+                logger.error(f"Failed to fetch {server}/perf_metrics: {ce}")
+                continue

Optional (latency): fetch all servers concurrently with asyncio.gather and per-task error handling. Want a suggested patch?

tensorrt_llm/serve/openai_disagg_server.py (5)

399-405: Tighten type of create_generator parameter

Use typing.Callable with an explicit signature that matches your generator helpers.

-    async def send_request(self, url: str,
+    async def send_request(self, url: str,
                            request: Union[CompletionRequest, ChatCompletionRequest],
                            endpoint: str,
                            response_type: Type[Union[CompletionResponse, ChatCompletionResponse]],
-                           create_generator: callable) -> Union[CompletionResponse, ChatCompletionResponse, StreamingResponse]:
+                           create_generator: Callable[[str, Union[CompletionRequest, ChatCompletionRequest]], AsyncIterator[bytes]]
+                           ) -> Union[CompletionResponse, ChatCompletionResponse, StreamingResponse]:

---

334-337: Consider awaiting key registration to reduce race window

Using asyncio.create_task makes correlation key insertion eventual. If /perf_metrics is called immediately, that request may be missed until the next poll. If acceptable, ignore; otherwise await self._add_perf_metrics_keys(...) here. Minor.

If you want to quantify any impact, I can script a stress test that triggers /perf_metrics concurrently with request dispatch to measure missed correlations.

---

36-36: Nit: fix typo in yapf directive

Spelling: “enable”.

-# yapf: enale
+# yapf: enable

---

118-123: create_error_response ignores provided status and always raises 500

Either return a proper JSONResponse with the given status or raise HTTPException(status_code=status_code, ...). Otherwise callers can’t control status codes.

-    def create_error_response(
-            message: str,
-            status_code: HTTPStatus = HTTPStatus.BAD_REQUEST) -> ErrorResponse:
-        raise HTTPException(status_code=500, detail=f"Internal server error {message}")
+    def create_error_response(
+            message: str,
+            status_code: HTTPStatus = HTTPStatus.BAD_REQUEST) -> ErrorResponse:
+        raise HTTPException(status_code=int(status_code), detail=str(message))

---

146-196: Add tests for /perf_metrics error-paths and pruning

• One server returns non-200; verify endpoint still returns aggregated results.
• JSON parse error from a server; verify skip.
• Correlation path when ctx arrives before gen and vice versa; verify no data loss.
• Pruning: exceed perf_metrics_max_requests; verify oldest entries are dropped deterministically.

I can draft pytest cases mirroring tests/unittest/llmapi/apps/_test_openai_perf_metrics.py style.

Configuration used: .coderabbit.yaml
Review profile: CHILL
Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

Learnt from: galagam
PR: NVIDIA/TensorRT-LLM#6487
File: tests/unittest/_torch/auto_deploy/unit/singlegpu/test_ad_trtllm_bench.py:1-12
Timestamp: 2025-08-06T13:58:07.506Z
Learning: In TensorRT-LLM, test files (files under tests/ directories) do not require NVIDIA copyright headers, unlike production source code files. Test files typically start directly with imports, docstrings, or code.

tests/unittest/api_stability/references/llm.yaml (1)

134-137: Consider grouping with the Stats section for readability

perf_metrics_max_requests is conceptually tied to return_perf_metrics and other stats controls. For maintainers scanning this reference, placing it under the “# Stats” block (near return_perf_metrics) would improve discoverability. No functional change—purely organizational.

Proposed edit (move block near the Stats entries):

@@
-      print_iter_log:
-        annotation: bool
-        default: False
-        status: beta
-      perf_metrics_max_requests:
-        annotation: int
-        default: 0
-        status: prototype
+      print_iter_log:
+        annotation: bool
+        default: False
+        status: beta
@@
       return_perf_metrics:
         annotation: bool
         default: False
         status: prototype
+      perf_metrics_max_requests:
+        annotation: int
+        default: 0
+        status: prototype

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

tests/integration/defs/disaggregated/test_disaggregated.py (1)

543-551: Past feedback addressed: correct urllib import, timeout, and JSON parsing

Thanks for switching to urllib.request, adding a timeout, and using json.load on the response object. This resolves the previously raised runtime and robustness issues.

tests/integration/defs/disaggregated/test_disaggregated.py (3)

39-45: Add return type annotation and a brief docstring; open YAML with explicit encoding

Minor polish to improve readability and consistency with the rest of the file where files are opened with encoding.

-def get_disagg_server_url_from_cfg(config_file: str):
-    with open(config_file, 'r') as file:
+def get_disagg_server_url_from_cfg(config_file: str) -> str:
+    """Return the disaggregated server base URL (http://host:port) from a YAML config."""
+    with open(config_file, 'r', encoding='utf-8') as file:
         config = yaml.safe_load(file)
     server_host = config.get('hostname', 'localhost')
     server_port = config.get('port', 8000)
     return f"http://{server_host}:{server_port}"

---

246-249: Avoid re-parsing YAML each loop iteration; compute server_url once

We open and parse the config file on every iteration just to derive the URL. Cache it once for slight speedup and less IO.

-                if extra_endpoints_test is not None:
-                    extra_endpoints_test(
-                        get_disagg_server_url_from_cfg(config_file))
+                if extra_endpoints_test is not None:
+                    extra_endpoints_test(server_url)

Add this initialization near where config_file is determined (outside the loop):

# After obtaining config_file and before entering the for-loop
server_url = get_disagg_server_url_from_cfg(config_file)

---

547-552: Add a retry loop when polling /perf_metrics to reduce flakes

I’ve verified that all metric producers emit both timing_metrics and kv_cache_metrics (see tensorrt_llm/serve/openai_server.py around lines 273–298), so it’s safe to wait briefly for the aggregator to combine them.

Please consider updating the test at tests/integration/defs/disaggregated/test_disaggregated.py (around lines 547–552) as follows:

-        with urllib.request.urlopen(f"{server_url}/perf_metrics",
-                                    timeout=10) as resp:
-            assert resp.status == 200
-            perf_metrics = json.load(resp)
-        assert len(perf_metrics) > 0
+        import time
+        perf_metrics = []
+        for _ in range(10):  # up to ~10s total
+            with urllib.request.urlopen(f"{server_url}/perf_metrics", timeout=10) as resp:
+                assert resp.status == 200
+                perf_metrics = json.load(resp)
+            if perf_metrics:
+                break
+            time.sleep(1)
+        assert perf_metrics, "No perf metrics returned after retries"

• Location to update:
  • tests/integration/defs/disaggregated/test_disaggregated.py lines 547–552
• Verified producers include kv_cache_metrics:
  • tensorrt_llm/serve/openai_server.py lines 273–298 [run_scripts output]

Optional refactor to improve test stability without major changes to production code.

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

tests/integration/defs/disaggregated/test_disaggregated.py (3)

20-20: Type annotations import looks good

Adding Callable is appropriate for the new callback parameter.

---

165-166: Signature extension is clear and typed

Adding prompt_file default and the typed extra_endpoints_test callback reads well.

---

69-69: perf_metrics config reference verified – LGTM

Verified that disagg_config_metrics.yaml exists under tests/integration/defs/disaggregated/test_configs/ and that the mapping in tests/integration/defs/disaggregated/test_disaggregated.py (line 69) correctly references it. No further changes needed.

tests/integration/defs/disaggregated/test_disaggregated.py (2)

544-551: Fix urllib import and add timeout for robustness

The code uses urllib.request.urlopen but only imports urllib, which will cause an AttributeError at runtime. Also, add a timeout to prevent indefinite blocking.

Apply this diff:

     def extra_endpoints_test(server_url: str):
         import json
-        import urllib.request
+        import urllib.request

         with urllib.request.urlopen(f"{server_url}/perf_metrics",
-                                    timeout=10) as resp:
+                                    timeout=30) as resp:
             assert resp.status == 200
             perf_metrics = json.load(resp)

---

559-571: Fix: kv_cache_size is not in timing_metrics and timestamps need robust checks

The test incorrectly tries to access kv_cache_size from gen_metrics (which are timing_metrics). The field is actually in kv_cache_metrics. Also, the strict timing comparisons could be flaky.

Apply this diff:

         ctx_metrics = item["ctx_perf_metrics"]["perf_metrics"]["timing_metrics"]
         gen_metrics = item["gen_perf_metrics"]["perf_metrics"]["timing_metrics"]
+        gen_kv_metrics = item["gen_perf_metrics"]["perf_metrics"].get("kv_cache_metrics", {})
         # only one token is generated in ctx
         assert ctx_metrics["last_token_time"] - ctx_metrics[
             "first_token_time"] < 1e-3
         assert ctx_metrics["last_token_time"] < gen_metrics["arrival_time"]
-        assert gen_metrics["kv_cache_size"] > 0
+        # Check kv_cache_metrics if available
+        if gen_kv_metrics:
+            # Note: kv_cache_size might be in timing_metrics for legacy reasons
+            if "kv_cache_size" in gen_metrics:
+                assert gen_metrics["kv_cache_size"] > 0
         assert gen_metrics["arrival_time"] < gen_metrics[
             "kv_cache_transfer_start"]
         assert gen_metrics["kv_cache_transfer_start"] < gen_metrics[
             "kv_cache_transfer_end"]
         assert gen_metrics["kv_cache_transfer_end"] < gen_metrics[
             "first_scheduled_time"]

tensorrt_llm/serve/openai_disagg_server.py (2)

55-66: Fix type annotations for Python 3.8 compatibility

The code uses PEP 585 built-in generic dict[str, dict[int, dict]] which is not supported in Python 3.8. Replace with typing module equivalents.

Apply this diff:

+from typing import Dict
 
             # server_key -> {ctx_request_id: perf_metrics}
-            self.server_perf_metrics: dict[str, dict[int, dict]] = {}
+            self.server_perf_metrics: Dict[str, Dict[int, dict]] = {}

---

150-159: Add error handling for individual server failures

If any server's /perf_metrics endpoint fails or returns non-200, the entire aggregation fails. Consider skipping failed servers to keep the endpoint robust.

Apply this diff:

         perf_metrics = {}
         exc = None
         try:
             for server in self.ctx_servers + self.gen_servers:
-                async with self.session.get(f"{server}/perf_metrics") as response:
-                    server_perf_metrics = await response.json()
-                    perf_metrics[server] = server_perf_metrics
+                try:
+                    async with self.session.get(f"{server}/perf_metrics") as response:
+                        if response.status != 200:
+                            logger.warning(f"Skipping {server}/perf_metrics: HTTP {response.status}")
+                            continue
+                        server_perf_metrics = await response.json()
+                        perf_metrics[server] = server_perf_metrics
+                except Exception as e:
+                    logger.error(f"Failed to fetch {server}/perf_metrics: {e}")
+                    # Don't set exc here - allow partial aggregation
         except Exception as e:
             # Keep the exception to raise it after saving perf metrics
             exc = e

tensorrt_llm/serve/openai_server.py (1)

266-308: Consider adding error handling for missing perf_metrics fields

The function accesses nested attributes without checking if they exist. Consider adding defensive checks to handle cases where some metrics might be None or missing.

Add defensive checks:

     for metrics_dict in perf_metrics:
         metrics = metrics_dict["perf_metrics"]
+        if metrics is None:
+            continue
         timing_metrics = metrics.timing_metrics
         kv_cache_metrics = metrics.kv_cache_metrics
         speculative_decoding = metrics.speculative_decoding

tensorrt_llm/serve/openai_disagg_server.py (1)

177-178: Consider not re-raising the exception to allow partial results

Re-raising the exception here causes the endpoint to return 500 even when some metrics were successfully collected. Consider returning partial results instead.

-            if exc is not None:
-                raise exc
+            # Log the exception but continue with partial results
+            if exc is not None:
+                logger.error(f"Error during perf metrics collection: {exc}")

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

💡 Knowledge Base configuration:

• MCP integration is disabled by default for public repositories
• Jira integration is disabled by default for public repositories
• Linear integration is disabled by default for public repositories

You can enable these sources in your CodeRabbit configuration.

Learnt from: galagam
PR: NVIDIA/TensorRT-LLM#6487
File: tests/unittest/_torch/auto_deploy/unit/singlegpu/test_ad_trtllm_bench.py:1-12
Timestamp: 2025-08-06T13:58:07.506Z
Learning: In TensorRT-LLM, test files (files under tests/ directories) do not require NVIDIA copyright headers, unlike production source code files. Test files typically start directly with imports, docstrings, or code.