address review comments

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

Author: reasonsolo

tensorrt_llm/serve/openai_client.py

@@ -20,6 +20,7 @@
 
 import aiohttp
 
+from tensorrt_llm.llmapi.disagg_utils import ServerRole
 from tensorrt_llm.logger import logger
 from tensorrt_llm.serve.openai_protocol import (
     ChatCompletionRequest,
@@ -42,26 +43,29 @@
 
 class OpenAIClient(ABC):
     async def send_request(
-        self, server: str, request: UCompletionRequest, hooks: Optional[ResponseHooks] = None
+        self,
+        request: UCompletionRequest,
+        server: Optional[str] = None,
+        hooks: Optional[ResponseHooks] = None,
     ) -> UCompletionResponseOrGenerator:
         if isinstance(request, CompletionRequest):
             return await self._send_request(
-                server, "v1/completions", request, CompletionResponse, hooks
+                "v1/completions", request, CompletionResponse, server, hooks
             )
         elif isinstance(request, ChatCompletionRequest):
             return await self._send_request(
-                server, "v1/chat/completions", request, ChatCompletionResponse, hooks
+                "v1/chat/completions", request, ChatCompletionResponse, server, hooks
             )
         else:
             raise ValueError(f"Invalid request type: {type(request)}")
 
     @abstractmethod
     async def _send_request(
         self,
-        server: str,
         endpoint: str,
         request: UCompletionRequest,
         response_type: Type[UCompletionResponse],
+        server: Optional[str] = None,
         hooks: Optional[ResponseHooks] = None,
     ) -> UCompletionResponseOrGenerator:
         """Send a request to the server and return the response and the body generator.
@@ -90,55 +94,58 @@ class OpenAIHttpClient(OpenAIClient):
     def __init__(
         self,
         router: Router,
-        client_type: str,
+        role: ServerRole,
         timeout_secs: int = 180,
         max_retries: int = 1,
+        retry_interval_sec: int = 1,
         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._role = role
+        self._metrics_collector = ClientMetricsCollector(role)
         self._session = session or aiohttp.ClientSession(
             connector=aiohttp.TCPConnector(limit=0, limit_per_host=0, force_close=False),
             timeout=aiohttp.ClientTimeout(total=timeout_secs),
         )
         self._max_retries = max_retries
-        self._retry_interval = 1
+        self._retry_interval_sec = retry_interval_sec
 
     async def _send_request(
         self,
-        server: str,
         endpoint: str,
         request: UCompletionRequest,
         response_type: Type[UCompletionResponse],
+        server: Optional[str] = None,
         hooks: Optional[ResponseHooks] = None,
     ) -> UCompletionResponseOrGenerator:
-        if len(server) == 0:
+        if server is None:
             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}"
+            f"Sending {self._role} request {request.disaggregated_params.ctx_request_id} to {url}"
         )
         try:
-            self._metrics_collector.inc("total_requests")
+            self._metrics_collector.total_requests.inc()
             resp_generator = self._post_with_retry(server, url, request, hooks)
             if request.stream:
+                # return the response generator, the request is not done yet
                 return resp_generator
             else:
                 # consume the generator to get the response and return it directly when it's not streaming
                 response = None
                 async for resp_json in resp_generator:
                     response = response_type(**resp_json)
                     if hooks:
-                        if self._client_type == "ctx":
+                        if self._role == ServerRole.CONTEXT:
                             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")
+            self._metrics_collector.error_requests.inc()
+            # finish the request upon error
+            await self._finish_request(request)
             raise
 
     async def _post_with_retry(
@@ -163,45 +170,45 @@ async def _post_with_retry(
                         # do NOT return generator directly here or the response will go
                         # out of scope and get destroyed
                         async for line in self._response_generator(
-                            request, http_response, start_time, hooks, server
+                            request, http_response, start_time, server, hooks
                         ):
                             yield line
+                        # don't finish the request here since the response generator is not done yet
                     else:
                         http_response.raise_for_status()
                         response_dict = await http_response.json()
                         # yield here since python forbids return statements in async generators
                         yield response_dict
+                        # finish the request after the successful response
+                        await self._finish_request(request)
                 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}"
+                        f"Client error to {url}: {e} - last retry {attempt} of {self._max_retries}"
                         "failed",
                         traceback.format_exc(),
                     )
                     raise
-
                 logger.error(
-                    f"{self._client_type} client error to {url}: {e} - retry {attempt} of {self._max_retries}",
+                    f"{self._role} 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")
+                await asyncio.sleep(self._retry_interval_sec)
+                self._metrics_collector.retry_requests.inc()
             except Exception as e:
                 logger.error(
-                    f"Unexpected error while processing {self._client_type} request to {url}: {e}"
+                    f"Unexpected error while processing {self._role} request to {url}: {e}"
                 )
                 raise
-            finally:
-                await self._finish_request(request)
 
     async def _response_generator(
         self,
         request: UCompletionRequest,
         http_response: aiohttp.ClientResponse,
         start_time: float,
+        server: str,
         hooks: Optional[ResponseHooks] = None,
-        server: str = "",
     ) -> AsyncGenerator[Any, None]:
         assert request.stream, "Request is not streaming"
         assert "text/event-stream" in http_response.headers.get("Content-Type", ""), (
@@ -215,12 +222,12 @@ async def _response_generator(
                 if i == 0:
                     if hooks:
                         hooks.on_first_token(server, request)
-                    self._metrics_collector.observe(
-                        "first_token_latency_seconds", now_time - last_token_time
+                    self._metrics_collector.first_token_latency_seconds.observe(
+                        now_time - last_token_time
                     )
                 else:
-                    self._metrics_collector.observe(
-                        "per_token_latency_seconds", now_time - last_token_time
+                    self._metrics_collector.per_token_latency_seconds.observe(
+                        now_time - last_token_time
                     )
                 i += 1
                 if line:
@@ -230,20 +237,20 @@ async def _response_generator(
 
             if hooks:
                 hooks.on_resp_done(server, request, None)
-            self._metrics_collector.inc("completed_requests")
-            self._metrics_collector.observe(
-                "complete_latency_seconds",
-                get_steady_clock_now_in_seconds() - start_time,
+            self._metrics_collector.completed_requests.inc()
+            self._metrics_collector.complete_latency_seconds.observe(
+                get_steady_clock_now_in_seconds() - start_time
             )
         except aiohttp.ClientError as e:
             # a client error is expected when the response stream is done if the connector has close=True
-            logger.error(f"{self._client_type} Client error: {e}")
-            self._metrics_collector.inc("error_requests")
+            logger.error(f"{self._role} client {server} error: {e}")
+            self._metrics_collector.error_requests.inc()
             raise
         except Exception:
-            self._metrics_collector.inc("error_requests")
+            self._metrics_collector.error_requests.inc()
             raise
         finally:
+            # finish the request after streaming response is done or error is raised
             await self._finish_request(request)
 
     async def _finish_request(self, request: UCompletionRequest) -> None:

tensorrt_llm/serve/openai_disagg_server.py

@@ -31,7 +31,7 @@
 from tensorrt_llm.executor.executor import CppExecutorError
 from tensorrt_llm.llmapi import tracing
 from tensorrt_llm.llmapi.disagg_utils import (DisaggServerConfig,
-                                              MetadataServerConfig,
+                                              MetadataServerConfig, ServerRole,
                                               get_ctx_gen_server_addrs)
 from tensorrt_llm.logger import logger
 from tensorrt_llm.serve.cluster_storage import (HttpClusterStorageServer,
@@ -136,8 +136,8 @@ async def validation_exception_handler(_, exc):
 
         self.register_routes()
 
-    def _create_client(self, router: Router, client_type: str, max_retries) -> OpenAIClient:
-        client = OpenAIHttpClient(router, client_type, self._req_timeout_secs, max_retries)
+    def _create_client(self, router: Router, role: ServerRole, max_retries: int = 1) -> OpenAIClient:
+        client = OpenAIHttpClient(router, role, self._req_timeout_secs, max_retries)
         self._perf_metrics_collector.add_client(client)
         return client
 

tensorrt_llm/serve/openai_disagg_service.py

@@ -52,7 +52,7 @@ def __init__(
         config: DisaggServerConfig,
         ctx_router: Router,
         gen_router: Router,
-        client_factory: Callable[[Router, str], OpenAIClient],
+        client_factory: Callable[[Router, ServerRole], OpenAIClient],
         metadata_server: Optional[JsonDictionary] = None,
         metadata_config: Optional[MetadataServerConfig] = None,
         req_timeout_secs: int = 180,
@@ -106,21 +106,25 @@ async def _send_disagg_request(
         if hooks:
             hooks.on_req_begin(request)
         # empty server means client decides which server to use
-        gen_server = ""
-        ctx_server = ""
+        reserved_gen_server = None
+        reserved_ctx_server = None
         # reserve a gen_server if conditional disagg is needed
-        gen_server, need_ctx = await self._check_conditional_disagg(request)
+        reserved_gen_server, need_ctx = await self._check_conditional_disagg(request)
         need_ctx = need_ctx and not await self._check_gen_only_disagg(request)
         ctx_response = None
         gen_req = 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, hooks)
+            ctx_response = await self._ctx_client.send_request(
+                ctx_req, server=reserved_ctx_server, hooks=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):
-            return await self._gen_client.send_request(gen_server, gen_req, hooks)
+            return await self._gen_client.send_request(
+                gen_req, server=reserved_gen_server, hooks=hooks
+            )
         else:
             if request.stream:
                 # ctx client will never return a generator when streaming is requested
@@ -170,7 +174,7 @@ async def _check_conditional_disagg(self, request: UCompletionRequest) -> bool:
             ):
                 return gen_server, True
             return gen_server, False
-        return "", True
+        return None, True
 
     async def _check_gen_only_disagg(self, request: UCompletionRequest) -> bool:
         if os.getenv("TRTLLM_DISAGG_BENCHMARK_GEN_ONLY") == "1":
@@ -206,8 +210,12 @@ def conditional_disagg_config(self) -> Optional[ConditionalDisaggConfig]:
         return self._config.conditional_disagg_config
 
     async def setup(self) -> None:
-        self._ctx_client = self._client_factory(self._ctx_router, "ctx", self._config.max_retries)
-        self._gen_client = self._client_factory(self._gen_router, "gen", self._config.max_retries)
+        self._ctx_client = self._client_factory(
+            self._ctx_router, ServerRole.CONTEXT, self._config.max_retries
+        )
+        self._gen_client = self._client_factory(
+            self._gen_router, ServerRole.GENERATION, self._config.max_retries
+        )
 
         if self.disagg_cluster_config and self._cluster_storage:
             logger.info("Starting disagg cluster manager")
@@ -263,12 +271,17 @@ async def check_servers_ready():
     async def _on_worker_event(self, worker_info: WorkerInfo, event_type: WatchEventType):
         router_map = {ServerRole.CONTEXT: self._ctx_router, ServerRole.GENERATION: self._gen_router}
         worker_addr = f"{worker_info.host}:{worker_info.port}"
-        router = router_map[worker_info.role]
-        if event_type == WatchEventType.SET:
-            await router.add_server(worker_addr)
-        elif event_type == WatchEventType.DELETE:
-            await router.remove_server(worker_addr)
-        logger.info(f"Worker {event_type.name} event: {worker_info.worker_id}, {worker_addr}")
+        try:
+            router = router_map[worker_info.role]
+            if event_type == WatchEventType.SET:
+                await router.add_server(worker_addr)
+            elif event_type == WatchEventType.DELETE:
+                await router.remove_server(worker_addr)
+            logger.info(f"Worker {event_type.name} event: {worker_info.worker_id}, {worker_addr}")
+        except KeyError:
+            logger.error(
+                f"Unknown worker role: {worker_info.role}, Worker {worker_info.worker_id} event: {event_type.name}"
+            )
 
     async def _verify_ctx_response(self, ctx_response: UCompletionResponse) -> None:
         if ctx_response:
@@ -281,49 +294,3 @@ async def _verify_ctx_response(self, ctx_response: UCompletionResponse) -> None:
             if ctx_response.choices[0].disaggregated_params.ctx_request_id is None:
                 raise ValueError("Invalid disaggregated params in context phase response.")
             return ctx_response
-
-
-# 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:
-        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
-    ) -> UCompletionResponseOrGenerator:
-        if hooks:
-            hooks.on_req_begin(request)
-        # empty server means client decides which server to use
-        gen_server = ""
-        ctx_server = ""
-        # reserve a gen_server if conditional disagg is needed
-        gen_server, need_ctx = await self._check_conditional_disagg(request)
-        need_ctx = need_ctx and not await self._check_gen_only_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"
-        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()
-        return ctx_task.result()