working on a stream request hang issue

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

Author: reasonsolo

docker/Makefile

@@ -139,7 +139,7 @@ CCACHE_DIR        ?= $(CODE_DIR)/cpp/.ccache
 CONAN_DIR         ?= $(CODE_DIR)/cpp/.conan
 USER_CACHE_DIR    ?= $(shell readlink -f "${HOME_DIR}/.cache")
 RUN_CMD           ?=
-CONTAINER_NAME    ?= tensorrt_llm_bug_wksp
+CONTAINER_NAME    ?= tensorrt_llm
 WORK_DIR          ?= $(CODE_DIR)
 DOCKER_PULL       ?= 0
 
@@ -157,6 +157,7 @@ endif
     		$(GPU_OPTS) \
     		--volume $(SOURCE_DIR):$(CODE_DIR) \
     		$(EXTRA_VOLUMES) \
+    		$(if $(and $(filter 1,$(LOCAL_USER)),$(shell [ -w "$(USER_CACHE_DIR)" ] && echo 1)),--volume $(USER_CACHE_DIR):/home/$(USER_NAME)/.cache:rw) \
     		--env "CCACHE_DIR=$(CCACHE_DIR)" \
     		--env "CCACHE_BASEDIR=$(CODE_DIR)" \
     		--env "CONAN_HOME=$(CONAN_DIR)" \

tensorrt_llm/serve/disagg_auto_scaling.py

@@ -4,7 +4,7 @@
 import random
 import time
 from dataclasses import asdict, dataclass
-from typing import Any, Dict, List, Tuple
+from typing import Any, Callable, Dict, List, Optional, Tuple
 
 from tensorrt_llm.llmapi.disagg_utils import DisaggClusterConfig, ServerRole
 from tensorrt_llm.logger import logger
@@ -44,6 +44,7 @@ def __init__(self, config: DisaggClusterConfig, storage: ClusterStorage):
         self._current_ctx_workers = {}  # worker_id -> WorkerInfo
         self._current_gen_workers = {}  # worker_id -> WorkerInfo
         self._watch_handle = None
+        self._watch_task = None
 
     def __del__(self):
         try:
@@ -92,7 +93,14 @@ def current_gen_worker_num(self) -> int:
     def worker_key_prefix(self) -> str:
         return get_worker_key_prefix(self._config.cluster_name)
 
-    async def watch_workers(self, get_existing_first: bool = True):
+    async def watch_workers(
+        self,
+        get_existing_first: bool = True,
+        on_event: Optional[Callable[[WorkerInfo, WatchEventType],
+                                    None]] = None):
+        if self._watch_handle:
+            logger.error("Watch handle is already initialized")
+            return []
         workers = []
         if get_existing_first:
             # There is a tiny gap between getting existing workers and watching the key,
@@ -106,12 +114,35 @@ async def watch_workers(self, get_existing_first: bool = True):
                 workers.append(self._parse_worker_info(event))
         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")
+            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 Exception as e:
+                    logger.error(
+                        f"Error updating routers by worker events: {e}")
+                    await asyncio.sleep(1)
+
+        if on_event:
+            self._watch_task = asyncio.create_task(on_event_wrapper())
         return workers
 
     async def unwatch_workers(self) -> None:
         if self._watch_handle:
             await self._cluster_storage.unwatch(self.worker_key_prefix)
             self._watch_handle = None
+        if self._watch_task:
+            self._watch_task.cancel()
+            self._watch_task = None
 
     async def get_worker_events(
             self) -> List[Tuple[WorkerInfo, WatchEventType]]:

tensorrt_llm/serve/openai_client.py

@@ -1,7 +1,7 @@
 # yapf: disable
 import asyncio
 from abc import ABC, abstractmethod
-from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Type
+from typing import Any, Dict, List, Optional, Tuple, Type, Union
 
 import aiohttp
 
@@ -12,15 +12,15 @@
                                                 CompletionResponse,
                                                 UCompletionRequest,
                                                 UCompletionResponse)
-from tensorrt_llm.serve.perf_metrics import DisaggPerfMetricsCollector
-from tensorrt_llm.serve.responses_utils import (ResponseHooks,
+from tensorrt_llm.serve.perf_metrics import (ClientMetricsCollector,
+                                             DisaggPerfMetricsCollector)
+from tensorrt_llm.serve.responses_utils import (CompletionResponseGenerator,
+                                                ResponseHooks,
                                                 get_steady_clock_now_in_seconds)
 from tensorrt_llm.serve.router import Router
 
 # yapf: enable
 
-CompletionResponseGenerator = AsyncGenerator[bytes, None]
-
 
 class OpenAIClient(ABC):
 
@@ -29,7 +29,7 @@ async def send_request(
         server: str,
         request: UCompletionRequest,
         hooks: Optional[ResponseHooks] = None
-    ) -> Tuple[UCompletionResponse, AsyncGenerator[bytes, None]]:
+    ) -> Union[UCompletionResponse, CompletionResponseGenerator]:
         if isinstance(request, CompletionRequest):
             return await self._send_request(server, "v1/completions", request,
                                             CompletionResponse, hooks)
@@ -48,7 +48,7 @@ async def _send_request(
         request: UCompletionRequest,
         response_type: Type[UCompletionResponse],
         hooks: Optional[ResponseHooks] = None,
-    ) -> Tuple[UCompletionResponse, AsyncGenerator[bytes, None]]:
+    ) -> Union[UCompletionResponse, CompletionResponseGenerator]:
         """
         Send a request to the server and return the response and the body iterator.
         The request is finished (in routers) when the generator is exhausted or there is an error.
@@ -83,16 +83,19 @@ def __init__(self,
                  router: Router,
                  client_type: str,
                  timeout_secs: int = 180,
+                 max_retries: int = 1,
                  perf_metrics_collector: DisaggPerfMetricsCollector = None):
         assert client_type in ["ctx", "gen"]
         self._router = router
         self._client_type = client_type
-        self._perf_metrics_collector = perf_metrics_collector
+        self._metrics_collector = ClientMetricsCollector(client_type)
         self._session = 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
 
     async def _send_request(
         self,
@@ -101,79 +104,122 @@ async def _send_request(
         request: UCompletionRequest,
         response_type: Type[UCompletionResponse],
         hooks: Optional[ResponseHooks] = None,
-    ) -> Tuple[UCompletionResponse, CompletionResponseGenerator]:
+    ) -> Union[UCompletionResponse, CompletionResponseGenerator]:
         if len(server) == 0:
             server, _ = await self._router.get_next_server(request)
         url = f"http://{server}/{endpoint}"
         try:
-            start_time = get_steady_clock_now_in_seconds()
-            self._perf_metrics_collector.inc(
-                f"{self._client_type}_total_requests")
-            async with self._session.post(
-                    url, json=request.model_dump(
-                        exclude_unset=True)) as http_response:
-                content_type = http_response.headers.get("Content-Type", "")
-                if not request.stream and "text/event-stream" in content_type:
-                    raise ValueError(
-                        "Received an event-stream although request stream was False"
-                    )
-
-                response_dict = await http_response.json()
-                if not http_response.ok:
-                    logger.error(f"Received failed response {response_dict}")
-                    http_response.raise_for_status()
-                response = response_type(**response_dict)
-
-                return response, self._response_generator(
-                    request, http_response, response, start_time, hooks)
+            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)
+                return response
         except Exception:
-            self._perf_metrics_collector.inc(
-                f"{self._client_type}_error_requests")
+            self._metrics_collector.inc("error_requests")
             await self._finish_request(request)
             raise
 
+    async def _post_with_retry(
+        self,
+        server: str,
+        url: str,
+        request: UCompletionRequest,
+        hooks: Optional[ResponseHooks] = None
+    ) -> Tuple[aiohttp.ClientResponse, Dict[str, Any]]:
+        json_data = request.model_dump(exclude_unset=True)
+        is_stream = request.stream
+        for attempt in range(self._max_retries + 1):
+            try:
+                start_time = get_steady_clock_now_in_seconds()
+                async with self._session.post(url,
+                                              json=json_data) as http_response:
+                    content_type = http_response.headers.get("Content-Type", "")
+                    if not is_stream and "text/event-stream" in content_type:
+                        raise ValueError(
+                            "Received an event-stream although request stream was False"
+                        )
+                    if is_stream:
+                        # 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):
+                            yield line
+                    else:
+                        http_response.raise_for_status()
+                        response_dict = await http_response.json()
+                        # do yield here until python allows return statements in async generators
+                        yield response_dict
+            except (aiohttp.ClientError, OSError) as e:
+                if attempt == self._max_retries:
+                    raise
+                import traceback
+                logger.error(
+                    f"Client error: {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}")
+                raise
+
     async def _response_generator(
             self,
             request: UCompletionRequest,
             http_response: aiohttp.ClientResponse,
-            response: UCompletionResponse,
             start_time: float,
-            hooks: Optional[ResponseHooks] = None
-    ) -> CompletionResponseGenerator:
+            hooks: Optional[ResponseHooks] = None,
+            server: str = "") -> CompletionResponseGenerator:
+        """
+        If the request is streaming, yield the response line by line,
+        otherwise, yield nothing because the generator won't be used and the response will be returned directly.
+        """
+        assert request.stream, "Request is not streaming"
+        assert "text/event-stream" in http_response.headers.get(
+            "Content-Type", ""), "Response is not streaming"
         try:
-            if request.stream and "text/event-stream" in http_response.headers.get(
-                    "Content-Type", ""):
-                last_token_time = start_time
-                async for i, line in enumerate(
-                        http_response.content.iter_any()):
-                    now_time = get_steady_clock_now_in_seconds()
-                    if i == 0:
-                        if hooks and hooks.on_first_token:
-                            hooks.on_first_token(request, response)
-                        self._perf_metrics_collector.observe(
-                            f"{self._client_type}_first_token_latency_seconds",
-                            now_time - last_token_time,
-                        )
-                    else:
-                        self._perf_metrics_collector.observe(
-                            f"{self._client_type}_per_token_latency_seconds",
-                            now_time - last_token_time,
-                        )
-                    if line:
-                        yield line
-                        await asyncio.sleep(0)
-                    last_token_time = now_time
-                if hooks and hooks.on_resp_done:
-                    hooks.on_resp_done(request, response)
-                self._perf_metrics_collector.inc(
-                    f"{self._client_type}_completed_requests")
-                self._perf_metrics_collector.observe(
-                    f"{self._client_type}_complete_latency_seconds",
-                    get_steady_clock_now_in_seconds() - start_time,
-                )
+            last_token_time = start_time
+            i = 0
+            async for line in http_response.content.iter_any():
+                now_time = get_steady_clock_now_in_seconds()
+                if i == 0:
+                    if hooks:
+                        hooks.on_first_token(server, request)
+                    self._metrics_collector.observe(
+                        "first_token_latency_seconds",
+                        now_time - last_token_time)
+                else:
+                    self._metrics_collector.observe("per_token_latency_seconds",
+                                                    now_time - last_token_time)
+                i += 1
+                if line:
+                    yield line
+                    await asyncio.sleep(0)
+                last_token_time = now_time
+
+            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,
+            )
+        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")
+            raise
         except Exception:
-            self._perf_metrics_collector.inc(
-                f"{self._client_type}_error_requests")
+            self._metrics_collector.inc("error_requests")
             raise
         finally:
             await self._finish_request(request)