fix failed CI tests

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

Author: reasonsolo

tensorrt_llm/serve/openai_client.py

@@ -29,7 +29,7 @@
     UCompletionRequest,
     UCompletionResponse,
 )
-from tensorrt_llm.serve.perf_metrics import ClientMetricsCollector, DisaggPerfMetricsCollector
+from tensorrt_llm.serve.perf_metrics import ClientMetricsCollector
 from tensorrt_llm.serve.responses_utils import (
     ResponseHooks,
     UCompletionResponseOrGenerator,
@@ -93,13 +93,13 @@ def __init__(
         client_type: str,
         timeout_secs: int = 180,
         max_retries: int = 1,
-        perf_metrics_collector: DisaggPerfMetricsCollector = None,
+        session: Optional[aiohttp.ClientSession] = None,
     ):
         assert client_type in ["ctx", "gen"]
         self._router = router
         self._client_type = client_type
         self._metrics_collector = ClientMetricsCollector(client_type)
-        self._session = aiohttp.ClientSession(
+        self._session = session or aiohttp.ClientSession(
             connector=aiohttp.TCPConnector(limit=0, limit_per_host=0, force_close=False),
             timeout=aiohttp.ClientTimeout(total=timeout_secs),
         )
@@ -263,16 +263,25 @@ async def shutdown(self) -> None:
         await self._session.close()
 
     async def check_ready(self) -> Tuple[List[str], List[str]]:
+        return await OpenAIHttpClient.check_ready_for_servers(self._session, self._router.servers)
+
+    @staticmethod
+    async def check_ready_for_servers(
+        session: aiohttp.ClientSession, servers: List[str]
+    ) -> Tuple[List[str], List[str]]:
         async def check_server_ready(server: str) -> bool:
             try:
-                async with self._session.get(f"http://{server}/health") as response:
+                url = (
+                    f"{server}/health"
+                    if server.startswith("http://")
+                    else f"http://{server}/health"
+                )
+                async with session.get(url) as response:
                     return response.status == 200
             except Exception:
                 return False
 
-        servers_ready = await asyncio.gather(
-            *[check_server_ready(server) for server in self._router.servers]
-        )
-        return [server for server, ready in zip(self._router.servers, servers_ready) if ready], [
-            server for server, ready in zip(self._router.servers, servers_ready) if not ready
+        servers_ready = await asyncio.gather(*[check_server_ready(server) for server in servers])
+        return [server for server, ready in zip(servers, servers_ready) if ready], [
+            server for server, ready in zip(servers, servers_ready) if not ready
         ]

tensorrt_llm/serve/openai_disagg_service.py

@@ -226,7 +226,7 @@ async def setup(self) -> None:
                 await self._gen_router.start_server_monitoring(
                     self._metadata_config.refresh_interval
                 )
-            await self._wait_for_servers_ready()
+            await self._wait_for_all_servers_ready()
 
     async def teardown(self) -> None:
         await self._ctx_client.shutdown()
@@ -239,7 +239,7 @@ async def teardown(self) -> None:
             await self._ctx_router.stop_server_monitoring()
             await self._gen_router.stop_server_monitoring()
 
-    async def _wait_for_servers_ready(self) -> None:
+    async def _wait_for_all_servers_ready(self) -> None:
         async def check_servers_ready():
             elapsed_time = 0
             interval = 3

tests/integration/defs/disaggregated/test_workers.py

@@ -14,7 +14,7 @@
 from transformers import AutoTokenizer
 
 from tensorrt_llm import logger
-from tensorrt_llm.serve.openai_disagg_server import OpenAIDisaggServer
+from tensorrt_llm.serve.openai_client import OpenAIHttpClient
 from tensorrt_llm.serve.openai_protocol import (CompletionRequest,
                                                 DisaggregatedParams)
 from tensorrt_llm.serve.router import (KvCacheAwareRouter,
@@ -66,6 +66,34 @@ def run_disaggregated_workers(
 DEFAULT_TIMEOUT_REQUEST = 180
 
 
+async def wait_until_all_servers_ready(
+    session: aiohttp.ClientSession,
+    servers: List[str],
+    server_start_timeout_secs: int = 180,
+) -> None:
+
+    async def check_all_servers_ready():
+        elapsed_time = 0
+        interval = 3
+        while elapsed_time < server_start_timeout_secs:
+            _, unready_servers = await OpenAIHttpClient.check_ready_for_servers(
+                session, servers)
+            if len(unready_servers) == 0:
+                return
+            await asyncio.sleep(interval)
+            elapsed_time += interval
+            logger.info(
+                f"[{elapsed_time}] Waiting for servers, {unready_servers}...")
+
+    try:
+        await asyncio.wait_for(check_all_servers_ready(),
+                               timeout=server_start_timeout_secs)
+    except asyncio.TimeoutError:
+        raise TimeoutError(
+            f"Timeout waiting for all servers to be ready in {server_start_timeout_secs} seconds"
+        )
+
+
 class BasicWorkerTester:
 
     def __init__(self,
@@ -82,9 +110,9 @@ async def new_session(self):
         session = aiohttp.ClientSession(
             connector=aiohttp.TCPConnector(force_close=True),
             timeout=aiohttp.ClientTimeout(total=self.req_timeout_secs))
-        await OpenAIDisaggServer.wait_for_all_servers_ready(
-            session, self.ctx_servers, self.gen_servers,
-            self.server_start_timeout_secs)
+        await wait_until_all_servers_ready(session,
+                                           self.ctx_servers + self.gen_servers,
+                                           self.server_start_timeout_secs)
         return session
 
     async def send_request(self, session: aiohttp.ClientSession, url: str,