tested

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
+CONTAINER_NAME    ?= tensorrt_llm_bug_wksp
 WORK_DIR          ?= $(CODE_DIR)
 DOCKER_PULL       ?= 0
 
@@ -157,7 +157,6 @@ 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/llmapi/disagg_utils.py

@@ -80,16 +80,16 @@ class MetadataServerConfig():
     refresh_interval: float = 10.0
 
 
-def get_ctx_gen_server_urls(
+def get_ctx_gen_server_addrs(
         server_configs: list[CtxGenServerConfig]
 ) -> tuple[list[str], list[str]]:
     ctx_server_urls = []
     gen_server_urls = []
     for cfg in server_configs:
         if cfg.type == "ctx":
-            ctx_server_urls.append(f"http://{cfg.hostname}:{cfg.port}")
+            ctx_server_urls.append(f"{cfg.hostname}:{cfg.port}")
         else:
-            gen_server_urls.append(f"http://{cfg.hostname}:{cfg.port}")
+            gen_server_urls.append(f"{cfg.hostname}:{cfg.port}")
 
     return ctx_server_urls, gen_server_urls
 

tensorrt_llm/serve/openai_client.py

@@ -1,8 +1,7 @@
-# yapf disagrees with isort in pre-commit hooks
 # yapf: disable
 import asyncio
 from abc import ABC, abstractmethod
-from typing import Any, AsyncGenerator, Dict, List, Tuple, Type, Union
+from typing import Any, AsyncGenerator, Dict, List, Optional, Tuple, Type
 
 import aiohttp
 
@@ -14,44 +13,41 @@
                                                 UCompletionRequest,
                                                 UCompletionResponse)
 from tensorrt_llm.serve.perf_metrics import DisaggPerfMetricsCollector
-from tensorrt_llm.serve.responses_utils import (CompletionResponseIterator,
+from tensorrt_llm.serve.responses_utils import (ResponseHooks,
                                                 get_steady_clock_now_in_seconds)
 from tensorrt_llm.serve.router import Router
 
 # yapf: enable
 
+CompletionResponseGenerator = AsyncGenerator[bytes, None]
+
 
 class OpenAIClient(ABC):
 
     async def send_request(
-        self, server: str, request: UCompletionRequest
+        self,
+        server: str,
+        request: UCompletionRequest,
+        hooks: Optional[ResponseHooks] = None
     ) -> Tuple[UCompletionResponse, AsyncGenerator[bytes, None]]:
         if isinstance(request, CompletionRequest):
-            return await self.send_completion_request(server, request)
+            return await self._send_request(server, "v1/completions", request,
+                                            CompletionResponse, hooks)
         elif isinstance(request, ChatCompletionRequest):
-            return await self.send_chat_request(server, request)
+            return await self._send_request(server, "v1/chat/completions",
+                                            request, ChatCompletionResponse,
+                                            hooks)
         else:
             raise ValueError(f"Invalid request type: {type(request)}")
 
-    async def send_completion_request(
-        self, server: str, request: CompletionRequest
-    ) -> Tuple[CompletionResponse, AsyncGenerator[bytes, None]]:
-        return await self._send_request(server, "v1/completions", request,
-                                        CompletionResponse)
-
-    async def send_chat_request(
-        self, server: str, request: ChatCompletionRequest
-    ) -> Tuple[ChatCompletionResponse, AsyncGenerator[bytes, None]]:
-        return await self._send_request(server, "v1/chat/completions", request,
-                                        ChatCompletionResponse)
-
     @abstractmethod
     async def _send_request(
         self,
         server: str,
         endpoint: str,
-        request: Union[CompletionRequest, ChatCompletionRequest],
+        request: UCompletionRequest,
         response_type: Type[UCompletionResponse],
+        hooks: Optional[ResponseHooks] = None,
     ) -> Tuple[UCompletionResponse, AsyncGenerator[bytes, None]]:
         """
         Send a request to the server and return the response and the body iterator.
@@ -95,7 +91,7 @@ def __init__(self,
         self._session = aiohttp.ClientSession(
             connector=aiohttp.TCPConnector(limit=0,
                                            limit_per_host=0,
-                                           force_close=True),
+                                           force_close=False),
             timeout=aiohttp.ClientTimeout(total=timeout_secs))
 
     async def _send_request(
@@ -104,7 +100,8 @@ async def _send_request(
         endpoint: str,
         request: UCompletionRequest,
         response_type: Type[UCompletionResponse],
-    ) -> Tuple[UCompletionResponse, CompletionResponseIterator]:
+        hooks: Optional[ResponseHooks] = None,
+    ) -> Tuple[UCompletionResponse, CompletionResponseGenerator]:
         if len(server) == 0:
             server, _ = await self._router.get_next_server(request)
         url = f"http://{server}/{endpoint}"
@@ -113,36 +110,46 @@ async def _send_request(
             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 response:
-                content_type = response.headers.get("Content-Type", "")
+                    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 response.json()
-                if not response.ok:
+                response_dict = await http_response.json()
+                if not http_response.ok:
                     logger.error(f"Received failed response {response_dict}")
-                    response.raise_for_status()
-                return response_type(**response_dict), self._response_generator(
-                    response, start_time)
+                    http_response.raise_for_status()
+                response = response_type(**response_dict)
+
+                return response, self._response_generator(
+                    request, http_response, response, start_time, hooks)
         except Exception:
             self._perf_metrics_collector.inc(
                 f"{self._client_type}_error_requests")
-            self.finish_request(request)
+            await self._finish_request(request)
             raise
 
     async def _response_generator(
-            self, request: UCompletionRequest, response: aiohttp.ClientResponse,
-            start_time: float) -> AsyncGenerator[bytes, None]:
+            self,
+            request: UCompletionRequest,
+            http_response: aiohttp.ClientResponse,
+            response: UCompletionResponse,
+            start_time: float,
+            hooks: Optional[ResponseHooks] = None
+    ) -> CompletionResponseGenerator:
         try:
-            if request.stream and "text/event-stream" in response.headers.get(
+            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(response.content.iter_any()):
+                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,
@@ -152,10 +159,12 @@ async def _response_generator(
                             f"{self._client_type}_per_token_latency_seconds",
                             now_time - last_token_time,
                         )
-                    last_token_time = now_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(
@@ -167,10 +176,10 @@ async def _response_generator(
                 f"{self._client_type}_error_requests")
             raise
         finally:
-            self.finish_request(request)
+            await self._finish_request(request)
 
-    async def finish_request(self, request: UCompletionRequest) -> None:
-        self._router.finish_request(request)
+    async def _finish_request(self, request: UCompletionRequest) -> None:
+        await self._router.finish_request(request)
 
     async def collect_metrics(self) -> Dict[str, Any]:
         metrics = {}

tensorrt_llm/serve/openai_disagg_server.py

@@ -1,11 +1,5 @@
 #!/usr/bin/env python
-import asyncio
-import copy
-import itertools
-import os
-import signal
 import traceback
-from collections import deque
 from contextlib import asynccontextmanager
 from typing import Callable, Optional, Union
 
@@ -18,19 +12,19 @@
 
 # yapf: disable
 from tensorrt_llm.llmapi.disagg_utils import (DisaggServerConfig,
-                                              MetadataServerConfig)
+                                              MetadataServerConfig,
+                                              get_ctx_gen_server_addrs)
 from tensorrt_llm.logger import logger
 from tensorrt_llm.serve.metadata_server import create_metadata_server
 from tensorrt_llm.serve.openai_client import OpenAIClient, OpenAIHttpClient
 from tensorrt_llm.serve.openai_disagg_service import (
     OpenAIDisaggregatedService, ResponseHooks)
 from tensorrt_llm.serve.openai_protocol import (ChatCompletionRequest,
                                                 CompletionRequest)
+from tensorrt_llm.serve.perf_metrics import DisaggPerfMetricsCollector
 from tensorrt_llm.serve.responses_utils import (ServerArrivalTimeMiddleware,
                                                 get_steady_clock_now_in_seconds)
 from tensorrt_llm.serve.router import Router, create_router
-from tensorrt_llm.tensorrt_llm.serve.perf_metrics import \
-    DisaggPerfMetricsCollector
 from tensorrt_llm.version import __version__ as VERSION
 
 # yapf: enale
@@ -44,21 +38,22 @@ def __init__(self,
                  server_start_timeout_secs: int = 180,
                  metadata_server_cfg: Optional[MetadataServerConfig] = None,
                  metrics_interval_secs: int = 0):
-        self.config = config
-        self.req_timeout_secs = req_timeout_secs
-        self.server_start_timeout_secs = server_start_timeout_secs
-        self.metadata_server_cfg = metadata_server_cfg
-        self.metrics_interval_secs = metrics_interval_secs
-
-        self._ctx_router = create_router(config.ctx_router_config, config.ctx_servers, metadata_server_cfg, create_metadata_server(metadata_server_cfg))
-        self._gen_router = create_router(config.gen_router_config, config.gen_servers, metadata_server_cfg, create_metadata_server(metadata_server_cfg))
+        self._config = config
+        self._req_timeout_secs = req_timeout_secs
+        self._server_start_timeout_secs = server_start_timeout_secs
+        self._metadata_server_cfg = metadata_server_cfg
+        self._metrics_interval_secs = metrics_interval_secs
+
+        self._ctx_servers, self._gen_servers = get_ctx_gen_server_addrs(config.server_configs)
+        self._ctx_router = create_router(config.ctx_router_config, self._ctx_servers, metadata_server_cfg, create_metadata_server(metadata_server_cfg))
+        self._gen_router = create_router(config.gen_router_config, self._gen_servers, metadata_server_cfg, create_metadata_server(metadata_server_cfg))
         self._metadata_server = create_metadata_server(metadata_server_cfg)
-        self._perf_metrics_collector = DisaggPerfMetricsCollector(config.perf_metrics_max_requests, [])
+        self._perf_metrics_collector = DisaggPerfMetricsCollector(config.perf_metrics_max_requests)
 
-        self._service = OpenAIDisaggregatedService(config, self._ctx_router, self._gen_router, self._create_client, self.metadata_server, req_timeout_secs, server_start_timeout_secs, self._perf_metrics_collector)
+        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)
 
         @asynccontextmanager
-        async def lifespan() -> None:
+        async def lifespan(app) -> None:
             await self._service.setup()
             yield
             await self._service.teardown()
@@ -72,10 +67,9 @@ async def validation_exception_handler(_, exc):
             return JSONResponse(status_code=400, content={"error": str(exc)})
 
         self.register_routes()
-        self.mount_metrics()
 
     def _create_client(self, router: Router, client_type: str, perf_metrics_collector: DisaggPerfMetricsCollector) -> OpenAIClient:
-        return OpenAIHttpClient(router, client_type, self.req_timeout_secs, perf_metrics_collector)
+        return OpenAIHttpClient(router, client_type, self._req_timeout_secs, perf_metrics_collector)
 
 
     def register_routes(self):
@@ -90,19 +84,27 @@ def register_routes(self):
         self.app.mount("/prometheus/metrics", metrics_app)
 
     def _wrap_entry_point(self, entry_point: Callable) -> Callable:
-        async def wrapper(req: Union[CompletionRequest, ChatCompletionRequest], raw_request: Request) -> Response:
+        async def wrapper(req: Union[CompletionRequest, ChatCompletionRequest], raw_req: Request) -> Response:
+            def update_arrival_time(req: Union[CompletionRequest, ChatCompletionRequest]):
+                raw_req.state.server_arrival_time = get_steady_clock_now_in_seconds()
+            def update_first_token_time(req: Union[CompletionRequest, ChatCompletionRequest], response: Response):
+                raw_req.state.server_first_token_time = get_steady_clock_now_in_seconds()
             try:
-                hooks = ResponseHooks();
-                hooks.on_req_begin = lambda req: raw_request.state.server_arrival_time = get_steady_clock_now_in_seconds()
-                hooks.on_first_token = lambda req, response: raw_request.state.server_first_token_time = get_steady_clock_now_in_seconds()
-                response, iterator = await entry_point(req, hooks)
+                hooks = ResponseHooks()
+                hooks.on_req_begin = update_arrival_time
+                hooks.on_first_token = update_first_token_time
+                response, generator = await entry_point(req, hooks)
                 if req.stream:
-                    return StreamingResponse(content=iterator, media_type="text/event-stream")
+                    return StreamingResponse(content=generator, media_type="text/event-stream")
                 else:
                     return response
             except Exception as e:
                 logger.error(f"Error in entry point: {e}")
-                return Response(status_code=500, content=f"Internal server error: {e}")
+                print(traceback.format_exc())
+                import sys
+                sys.exit(1)
+                #return Response(status_code=500, content=f"Internal server error: {e}")
+                raise e
         return wrapper
 
 
@@ -117,11 +119,11 @@ async def cluster_info(self) -> JSONResponse:
     async def version(self) -> JSONResponse:
         return JSONResponse(content={"version": VERSION})
 
-    async def __call__(self, host, port):
+    async def __call__(self, host: str, port: int):
         config = uvicorn.Config(self.app,
                                 host=host,
                                 port=port,
-                                log_level="info",
+                                log_level=logger.level,
                                 timeout_keep_alive=TIMEOUT_KEEP_ALIVE)
         await uvicorn.Server(config).serve()
 
@@ -150,7 +152,7 @@ async def set_steady_clock_offset(server_url: str, offset: float) -> None:
             async with session.post(server_url + STEADY_CLOCK_OFFSET_ENDPOINT, json=payload) as response:
                 if response.status != 200:
                     logger.warning(f"Cannot set disagg server steady clock offset for server {server_url}, the perf metrics timestamps could be mis-aligned")
-        for server_url in self.ctx_servers + self.gen_servers:
+        for server_url in self._ctx_servers + self._gen_servers:
             delay, offset = await query_steady_clock_offset(server_url)
             if delay is None or offset is None:
                 logger.warning(f"Unable to measure steady clock offset for {server_url}; skipping adjustment")