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,5 +1,20 @@
+# 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
 
@@ -18,6 +33,7 @@
 from tensorrt_llm.serve.responses_utils import (
     CompletionResponseGenerator,
     ResponseHooks,
+    UCompletionResponseOrGenerator,
     get_steady_clock_now_in_seconds,
 )
 from tensorrt_llm.serve.router import Router
@@ -48,8 +64,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 +76,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: ...
@@ -101,24 +118,28 @@ async def _send_request(
         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(
@@ -149,22 +170,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,

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,16 @@ def on_req_begin(self, request: UCompletionRequest):
 
     def on_ctx_resp(self, ctx_server: str, response: UCompletionResponse):
         self.ctx_server = ctx_server
+        logger.info(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()
 
     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))
 
 
 class OpenAIDisaggServer:
@@ -81,7 +98,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 +147,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)
 

tensorrt_llm/serve/openai_disagg_service.py

@@ -1,3 +1,17 @@
+# 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.
+
 import asyncio
 import copy
 import os
@@ -7,6 +21,7 @@
     ConditionalDisaggConfig,
     DisaggClusterConfig,
     DisaggServerConfig,
+    MetadataServerConfig,
     ServerRole,
 )
 from tensorrt_llm.logger import logger
@@ -37,8 +52,9 @@ def __init__(
         config: DisaggServerConfig,
         ctx_router: Router,
         gen_router: Router,
-        client_factory: Callable[[Router, str], OpenAIClient] = None,
+        client_factory: Callable[[Router, str], OpenAIClient],
         metadata_server: Optional[JsonDictionary] = None,
+        metadata_config: Optional[MetadataServerConfig] = None,
         req_timeout_secs: int = 180,
         server_start_timeout_secs: int = 180,
         perf_metrics_collector: Optional[DisaggPerfMetricsCollector] = None,
@@ -49,6 +65,7 @@ def __init__(
         self._gen_router = gen_router
         self._client_factory = client_factory
         self._metadata_server = metadata_server
+        self._metadata_config = metadata_config
         self._req_timeout_secs = req_timeout_secs
         self._server_start_timeout_secs = server_start_timeout_secs
         self._perf_metrics_collector = perf_metrics_collector
@@ -86,11 +103,6 @@ async def openai_chat_completion(
     async def _send_disagg_request(
         self, request: UCompletionRequest, hooks: Optional[ResponseHooks] = None
     ) -> Union[UCompletionResponse, CompletionResponseGenerator]:
-        """This is the main disaggregated serving logic:
-        1. send context request to the context server if ctx is needed, return the context response if gen is not needed
-        2. build a generation request based on the context response and send it to the generation server if gen is needed,
-         return the generation response
-        """
         if hooks:
             hooks.on_req_begin(request)
         # empty server means client decides which server to use
@@ -104,7 +116,7 @@ async def _send_disagg_request(
         if need_ctx:
             ctx_req = self._get_ctx_request(request)
             # ctx generator is empty
-            ctx_response = await self._ctx_client.send_request(ctx_server, ctx_req)
+            ctx_response = await self._ctx_client.send_request(ctx_server, ctx_req, hooks)
             await self._verify_ctx_response(ctx_response)
             gen_req = self._get_gen_request(request, ctx_response)
         if ctx_response is None or self._need_gen(ctx_response):
@@ -206,13 +218,13 @@ async def setup(self) -> None:
             await self._disagg_cluster_manager.watch_workers(on_event=self._on_worker_event)
             logger.info("Disagg cluster manager started")
         else:
-            if self._metadata_server:
+            if self._metadata_server and self._metadata_config:
                 logger.info("Starting server monitoring via metadata service")
                 await self._ctx_router.start_server_monitoring(
-                    self.metadata_server.refresh_interval
+                    self._metadata_config.refresh_interval
                 )
                 await self._gen_router.start_server_monitoring(
-                    self.metadata_server.refresh_interval
+                    self._metadata_config.refresh_interval
                 )
             await self._wait_for_servers_ready()
 
@@ -274,7 +286,11 @@ async def _verify_ctx_response(self, ctx_response: UCompletionResponse) -> None:
 # FIXME: This is a demo to show the basic idea of disagg-service with pre-allocating generation
 class OpenAIDisaggregatedPreAllocService(OpenAIDisaggregatedService):
     def _need_gen(self, request: UCompletionRequest) -> bool:
-        return request.max_tokens > 1
+        if isinstance(request, CompletionRequest) and request.max_tokens is not None:
+            return request.max_tokens > 1
+        if isinstance(request, ChatCompletionRequest) and request.max_completion_tokens is not None:
+            return request.max_completion_tokens > 1
+        return False
 
     async def _send_disagg_request(
         self, request: UCompletionRequest, hooks: Optional[ResponseHooks] = None
@@ -290,23 +306,24 @@ async def _send_disagg_request(
         need_gen = self._need_gen(request)
         # send ctx and gen requests in parallel
         assert need_gen or need_ctx, "Neither generation nor context is required"
-        with asyncio.TaskGroup() as tg:
-            if need_ctx:
-
-                async def _run_ctx_task():
-                    # send ctx request and gen request in parallel
-                    ctx_req = self._get_ctx_request(request)
-                    ctx_response = await self._ctx_client.send_completion_request(
-                        ctx_server, ctx_req
-                    )
-                    return ctx_response
-
-                ctx_task = tg.create_task(_run_ctx_task())
-            if need_gen:
-                gen_task = tg.create_task(
-                    self._gen_client.send_completion_request(gen_server, request, hooks)
-                )
+        gen_task = None
+        ctx_task = None
+        tasks = []
+
+        async def _run_ctx_task():
+            # send ctx request and gen request in parallel
+            ctx_req = self._get_ctx_request(request)
+            ctx_response = await self._ctx_client.send_request(ctx_server, ctx_req, hooks)
+            return ctx_response
+
+        if need_ctx:
+            ctx_task = asyncio.create_task(_run_ctx_task())
+        if need_gen:
+            gen_task = asyncio.create_task(
+                self._gen_client.send_request(gen_server, request, hooks)
+            )
+            tasks.append(gen_task)
+        await asyncio.gather(*tasks)
         if need_gen:
             return gen_task.result()
-        else:
-            return ctx_task.result()
+        return ctx_task.result()