fix tests

Signed-off-by: Lizhi Zhou <[email protected]>

Author: reasonsolo

tensorrt_llm/serve/disagg_auto_scaling.py

@@ -4,7 +4,7 @@
 import random
 import time
 from dataclasses import asdict, dataclass
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from typing import Any, Awaitable, Callable, Dict, List, Optional, Tuple
 
 from tensorrt_llm.llmapi.disagg_utils import DisaggClusterConfig, ServerRole
 from tensorrt_llm.logger import logger
@@ -97,7 +97,7 @@ async def watch_workers(
         self,
         get_existing_first: bool = True,
         on_event: Optional[Callable[[WorkerInfo, WatchEventType],
-                                    None]] = None):
+                                    Awaitable[Any]]] = None):
         if self._watch_handle:
             logger.error("Watch handle is already initialized")
             return []
@@ -121,25 +121,28 @@ async def watch_workers(
         self._watch_handle = await self._cluster_storage.watch(
             self.worker_key_prefix)
 
-        async def on_event_wrapper():
-            logger.warning(
-                f"Initializing watch task with {len(workers)} existing workers")
+        async def worker_event_loop():
+            logger.info(
+                f"Start watching worker events with {len(workers)} existing workers"
+            )
             for worker_info in workers:
                 await on_event(worker_info, WatchEventType.SET)
-            logger.warning("Start watching worker events")
             while True:
                 try:
                     worker_events = await self._watch_handle.drain()
                     for event in worker_events:
                         worker_info = self._parse_worker_info(event)
                         await on_event(worker_info, event.event_type)
+                except asyncio.CancelledError:
+                    break
                 except Exception as e:
                     logger.error(
                         f"Error updating routers by worker events: {e}")
                     await asyncio.sleep(1)
+            logger.info("Stop watching worker events")
 
         if on_event:
-            self._watch_task = asyncio.create_task(on_event_wrapper())
+            self._watch_task = asyncio.create_task(worker_event_loop())
         return workers
 
     async def unwatch_workers(self) -> None:

tensorrt_llm/serve/openai_client.py

@@ -1,7 +1,22 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.
+#
+# 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.
+
 # yapf: disable
 import asyncio
+import traceback
 from abc import ABC, abstractmethod
-from typing import Any, Dict, List, Optional, Tuple, Type, Union
+from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Type
 
 import aiohttp
 
@@ -16,8 +31,8 @@
 )
 from tensorrt_llm.serve.perf_metrics import ClientMetricsCollector, DisaggPerfMetricsCollector
 from tensorrt_llm.serve.responses_utils import (
-    CompletionResponseGenerator,
     ResponseHooks,
+    UCompletionResponseOrGenerator,
     get_steady_clock_now_in_seconds,
 )
 from tensorrt_llm.serve.router import Router
@@ -28,7 +43,7 @@
 class OpenAIClient(ABC):
     async def send_request(
         self, server: str, request: UCompletionRequest, hooks: Optional[ResponseHooks] = None
-    ) -> Union[UCompletionResponse, CompletionResponseGenerator]:
+    ) -> UCompletionResponseOrGenerator:
         if isinstance(request, CompletionRequest):
             return await self._send_request(
                 server, "v1/completions", request, CompletionResponse, hooks
@@ -48,8 +63,9 @@ async def _send_request(
         request: UCompletionRequest,
         response_type: Type[UCompletionResponse],
         hooks: Optional[ResponseHooks] = None,
-    ) -> Union[UCompletionResponse, CompletionResponseGenerator]:
-        """Send a request to the server and return the response and the body iterator.
+    ) -> UCompletionResponseOrGenerator:
+        """Send a request to the server and return the response and the body generator.
+
         The request is finished (in routers) when the generator is exhausted or there is an error.
         """
         ...
@@ -59,7 +75,7 @@ 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"""
+        """Return the list of ready servers and the list of unready servers."""
         ...
 
     async def shutdown(self) -> None: ...
@@ -97,28 +113,32 @@ async def _send_request(
         request: UCompletionRequest,
         response_type: Type[UCompletionResponse],
         hooks: Optional[ResponseHooks] = None,
-    ) -> Union[UCompletionResponse, CompletionResponseGenerator]:
+    ) -> UCompletionResponseOrGenerator:
         if len(server) == 0:
             server, _ = await self._router.get_next_server(request)
         url = f"http://{server}/{endpoint}"
+        logger.debug(
+            f"Sending {self._client_type} request {request.disaggregated_params.ctx_request_id} to {url}"
+        )
         try:
             self._metrics_collector.inc("total_requests")
             resp_generator = self._post_with_retry(server, url, request, hooks)
             if request.stream:
                 return resp_generator
             else:
                 # consume the generator to get the response and return it directly when it's not streaming
-                resp_json = await anext(resp_generator)
-                response = response_type(**resp_json)
-                if hooks:
-                    if self._client_type == "ctx":
-                        hooks.on_ctx_resp(server, response)
-                    hooks.on_first_token(server, request)
-                    hooks.on_resp_done(server, request, response)
+                response = None
+                async for resp_json in resp_generator:
+                    response = response_type(**resp_json)
+                    if hooks:
+                        if self._client_type == "ctx":
+                            hooks.on_ctx_resp(server, response)
+                        else:
+                            hooks.on_first_token(server, request)
+                            hooks.on_resp_done(server, request, response)
                 return response
         except Exception:
             self._metrics_collector.inc("error_requests")
-            await self._finish_request(request)
             raise
 
     async def _post_with_retry(
@@ -127,7 +147,7 @@ async def _post_with_retry(
         url: str,
         request: UCompletionRequest,
         hooks: Optional[ResponseHooks] = None,
-    ) -> Tuple[aiohttp.ClientResponse, Dict[str, Any]]:
+    ) -> AsyncGenerator[Any, None]:
         json_data = request.model_dump(exclude_unset=True)
         is_stream = request.stream
         for attempt in range(self._max_retries + 1):
@@ -149,22 +169,31 @@ async def _post_with_retry(
                     else:
                         http_response.raise_for_status()
                         response_dict = await http_response.json()
-                        # do yield here until python allows return statements in async generators
+                        # yield here since python forbids return statements in async generators
                         yield response_dict
+                break  # break and skip retries if the whole response is processed without exception
             except (aiohttp.ClientError, OSError) as e:
                 if attempt == self._max_retries:
+                    logger.error(
+                        f"{self._client_type} client error to {url}: {e} - last retry {attempt} of {self._max_retries}"
+                        "failed",
+                        traceback.format_exc(),
+                    )
                     raise
-                import traceback
 
                 logger.error(
-                    f"Client error: {e} - retry {attempt} of {self._max_retries}",
+                    f"{self._client_type} client error to {url}: {e} - retry {attempt} of {self._max_retries}",
                     traceback.format_exc(),
                 )
                 await asyncio.sleep(self._retry_interval)
                 self._metrics_collector.inc("retry_requests")
             except Exception as e:
-                logger.error(f"Error encountered while processing request to {url}: {e}")
+                logger.error(
+                    f"Unexpected error while processing {self._client_type} request to {url}: {e}"
+                )
                 raise
+            finally:
+                await self._finish_request(request)
 
     async def _response_generator(
         self,
@@ -173,7 +202,7 @@ async def _response_generator(
         start_time: float,
         hooks: Optional[ResponseHooks] = None,
         server: str = "",
-    ) -> CompletionResponseGenerator:
+    ) -> AsyncGenerator[Any, None]:
         assert request.stream, "Request is not streaming"
         assert "text/event-stream" in http_response.headers.get("Content-Type", ""), (
             "Response is not streaming"

tensorrt_llm/serve/openai_disagg_server.py

@@ -1,4 +1,19 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.
+#
+# 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.
 #!/usr/bin/env python
+
+# yapf: disable
 import asyncio
 import signal
 import traceback
@@ -10,7 +25,7 @@
 from fastapi import FastAPI, HTTPException, Request
 from fastapi.exceptions import RequestValidationError
 from fastapi.responses import JSONResponse, Response, StreamingResponse
-from prometheus_client import CollectorRegistry, make_asgi_app
+from prometheus_client import make_asgi_app
 
 # yapf: disable
 from tensorrt_llm.executor import CppExecutorError
@@ -49,14 +64,18 @@ def on_req_begin(self, request: UCompletionRequest):
 
     def on_ctx_resp(self, ctx_server: str, response: UCompletionResponse):
         self.ctx_server = ctx_server
+        logger.debug(f"Received context response from {ctx_server} for request {response.choices[0].disaggregated_params.ctx_request_id}")
 
     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()
+        logger.debug(f"Received first token from {gen_server} for request {request.disaggregated_params.ctx_request_id}")
 
     def on_resp_done(self, gen_server: str, request: UCompletionRequest, response: UCompletionResponse = None):
-        ctx_req_id = request.disaggregated_params.ctx_request_id
-        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))
+        if request.disaggregated_params:
+            ctx_req_id = request.disaggregated_params.ctx_request_id
+            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))
+            logger.debug(f"Request {ctx_req_id} completed")
 
 
 class OpenAIDisaggServer:
@@ -81,7 +100,14 @@ def __init__(self,
 
         self._disagg_cluster_storage = create_cluster_storage(config.disagg_cluster_config.cluster_uri, config.disagg_cluster_config.cluster_name) if config.disagg_cluster_config else None
 
-        self._service = OpenAIDisaggregatedService(self._config, self._ctx_router, self._gen_router, self._create_client, self._metadata_server, self._req_timeout_secs, self._server_start_timeout_secs, self._perf_metrics_collector, self._disagg_cluster_storage)
+        self._service = OpenAIDisaggregatedService(
+            self._config, self._ctx_router, self._gen_router, self._create_client,
+            metadata_server=self._metadata_server,
+            metadata_config=self._metadata_server_cfg,
+            req_timeout_secs=self._req_timeout_secs,
+            server_start_timeout_secs=self._server_start_timeout_secs,
+            perf_metrics_collector=self._perf_metrics_collector,
+            disagg_cluster_storage=self._disagg_cluster_storage)
 
         try:
             otlp_cfg = config.otlp_config
@@ -123,9 +149,7 @@ def register_routes(self):
         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"])
-        registry = CollectorRegistry()
-        metrics_app = make_asgi_app(registry=registry)
-        self.app.mount("/prometheus/metrics", metrics_app)
+        self.app.mount("/prometheus/metrics", make_asgi_app())
         if self._disagg_cluster_storage and isinstance(self._disagg_cluster_storage, HttpClusterStorageServer):
             self._disagg_cluster_storage.add_routes(self.app)