[TRTLLM-9181][feat] improve disagg-server prometheus metrics; synchronize workers' clocks when workers are dynamic (#9726)

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

Author: reasonsolo

tensorrt_llm/_utils.py

@@ -473,10 +473,20 @@ def dim_resolve_negative(dim, ndim):
     return tuple(pos)
 
 
-def get_free_port():
-    with socket.socket() as sock:
-        sock.bind(("", 0))
-        return sock.getsockname()[1]
+def get_free_port() -> int:
+    return get_free_ports(1)[0]
+
+
+def get_free_ports(num=1) -> List[int]:
+    sockets = [
+        socket.socket(socket.AF_INET, socket.SOCK_STREAM) for _ in range(num)
+    ]
+    for s in sockets:
+        s.bind(('', 0))
+    ports = [s.getsockname()[1] for s in sockets]
+    for s in sockets:
+        s.close()
+    return ports
 
 
 # mpi4py only exports MPI_COMM_TYPE_SHARED, so we define OMPI_COMM_TYPE_HOST here

tensorrt_llm/serve/disagg_auto_scaling.py

@@ -2,6 +2,7 @@
 import json
 import os
 import random
+import socket
 import time
 from dataclasses import asdict, dataclass
 from typing import Any, Awaitable, Callable, Dict, List, Optional, Tuple
@@ -29,6 +30,18 @@ def get_worker_key(name: str, role: ServerRole, worker_id: str = "") -> str:
     return f"{get_worker_key_prefix(name)}/{worker_id}"
 
 
+def get_host_from_uri(uri: str) -> str:
+    return uri.split("://")[1].split(":")[0]
+
+
+# Get the local ip address from a remote host,
+# if remote host is not provided, use Google's public DNS server "8.8.8.8"
+def get_local_ip(remote_host: str = "8.8.8.8") -> str:
+    with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as s:
+        s.connect((remote_host, 80))
+        return s.getsockname()[0]
+
+
 class DisaggClusterManager:
     """
     The cluster manager is responsible for managing the workers in the cluster.
@@ -238,18 +251,25 @@ class DisaggClusterWorker:
     It will send heartbeat to the cluster storage every heartbeat_interval_sec seconds.
     If the worker heartbeat fails, it will re-register itself.
     """
+    LOCALHOST_IPS = ["localhost", "127.0.0.1", "0.0.0.0", "::1",
+                     "::"]  # nosec B104
 
     def __init__(self, role: ServerRole, host: str, port: int,
                  config: DisaggClusterConfig, storage: ClusterStorage):
         self._role = role
-        self._host = host
         self._port = port
         self._config = config
         self._cluster_storage = storage
         self._stop = False
         self._heartbeat_task = None
         self._last_heartbeat = 0
-        self._worker_id = f"{role.name}-{host}:{port}-{int(time.time()*1000)}-{os.getpid()}-{random.randint(0, 1000):03}"
+        register_host = host
+        # if the host is localhost and the cluster uri is not localhost, use the hostname to register the worker
+        disagg_host = get_host_from_uri(self._config.cluster_uri)
+        if host in self.LOCALHOST_IPS and disagg_host not in self.LOCALHOST_IPS:
+            register_host = get_local_ip(disagg_host)
+        self._host = register_host
+        self._worker_id = f"{role.name}-{register_host}:{port}-{int(time.time()*1000)}-{os.getpid()}-{random.randint(0, 1000):03}"
 
     def __del__(self):
         try:

tensorrt_llm/serve/openai_client.py

@@ -183,6 +183,9 @@ async def _post_with_retry(
                         yield response_dict
                         # finish the request after the successful response
                         await self._finish_request(request)
+                        self._metrics_collector.complete_latency_seconds.observe(
+                            get_steady_clock_now_in_seconds() - start_time
+                        )
                 break  # break and skip retries if the whole response is processed without exception
             except (aiohttp.ClientError, OSError) as e:
                 if lines_yielded > 0:
@@ -227,25 +230,24 @@ async def _response_generator(
             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.first_token_latency_seconds.observe(
-                        now_time - last_token_time
-                    )
-                else:
-                    self._metrics_collector.per_token_latency_seconds.observe(
-                        now_time - last_token_time
-                    )
-                i += 1
                 if line:
+                    if i == 0:
+                        if hooks:
+                            hooks.on_first_token(server, request)
+                        self._metrics_collector.first_token_latency_seconds.observe(
+                            now_time - last_token_time
+                        )
+                    else:
+                        self._metrics_collector.per_token_latency_seconds.observe(
+                            now_time - last_token_time
+                        )
+                    i += 1
                     yield line
                     await asyncio.sleep(0)
                 last_token_time = now_time
 
             if hooks:
                 hooks.on_resp_done(server, request, None)
-            self._metrics_collector.completed_requests.inc()
             self._metrics_collector.complete_latency_seconds.observe(
                 get_steady_clock_now_in_seconds() - start_time
             )
@@ -262,6 +264,7 @@ async def _response_generator(
             await self._finish_request(request)
 
     async def _finish_request(self, request: UCompletionRequest) -> None:
+        self._metrics_collector.completed_requests.inc()
         await self._router.finish_request(request)
 
     async def collect_metrics(self) -> Dict[str, Any]:

tensorrt_llm/serve/openai_disagg_server.py

@@ -57,11 +57,12 @@ def __init__(self, raw_req: Request, perf_metrics_collector: DisaggPerfMetricsCo
         self.raw_req = raw_req
         self.ctx_server = ""
         self.gen_server = ""
+        self.request_arrival_time = raw_req.state.server_arrival_time
         self.server_first_token_time = 0
         self.perf_metrics_collector = perf_metrics_collector
 
     def on_req_begin(self, request: UCompletionRequest):
-        ...
+        self.perf_metrics_collector.queue_latency_seconds.observe(get_steady_clock_now_in_seconds() - self.request_arrival_time)
 
     def on_ctx_resp(self, ctx_server: str, response: UCompletionResponse):
         self.ctx_server = ctx_server
@@ -93,8 +94,8 @@ def __init__(self,
         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._ctx_router = create_router(config.ctx_router_config, self._ctx_servers, metadata_server_cfg, create_metadata_server(metadata_server_cfg), self._sync_server_clock)
+        self._gen_router = create_router(config.gen_router_config, self._gen_servers, metadata_server_cfg, create_metadata_server(metadata_server_cfg), self._sync_server_clock)
         self._metadata_server = create_metadata_server(metadata_server_cfg)
         self._perf_metrics_collector = DisaggPerfMetricsCollector(config.perf_metrics_max_requests)
 
@@ -122,8 +123,10 @@ def __init__(self,
 
         @asynccontextmanager
         async def lifespan(app) -> None:
+            # Prepare servers (sync server clock) when static ctx/gen server list is used
+            await self._ctx_router.prepare_servers()
+            await self._gen_router.prepare_servers()
             await self._service.setup()
-            await self._set_steady_clock_offsets()
             yield
             await self._service.teardown()
 
@@ -133,6 +136,7 @@ async def lifespan(app) -> None:
 
         @self.app.exception_handler(RequestValidationError)
         async def validation_exception_handler(_, exc):
+            self._perf_metrics_collector.validation_exceptions.inc()
             return JSONResponse(status_code=400, content={"error": str(exc)})
 
         self.register_routes()
@@ -158,8 +162,14 @@ def register_routes(self):
     def _wrap_entry_point(self, entry_point: Callable) -> Callable:
         async def wrapper(req: UCompletionRequest, raw_req: Request) -> Response:
             try:
+                self._perf_metrics_collector.total_requests.inc()
+                if req.stream:
+                    self._perf_metrics_collector.stream_requests.inc()
+                else:
+                    self._perf_metrics_collector.nonstream_requests.inc()
                 hooks = RawRequestResponseHooks(raw_req, self._perf_metrics_collector)
                 response_or_generator = await entry_point(req, hooks)
+                self._perf_metrics_collector.total_responses.inc()
                 if req.stream:
                     return StreamingResponse(content=response_or_generator, media_type="text/event-stream")
                 else:
@@ -173,9 +183,11 @@ def _handle_exception(self, exception):
             logger.error("CppExecutorError: ", traceback.format_exc())
             signal.raise_signal(signal.SIGINT)
         elif isinstance(exception, HTTPException):
+            self._perf_metrics_collector.http_exceptions.inc()
             logger.error(f"HTTPException {exception.status_code} {exception.detail}: ", traceback.format_exc())
             raise exception
         else:
+            self._perf_metrics_collector.internal_errors.inc()
             logger.error("Internal server error: ", traceback.format_exc())
             raise HTTPException(status_code=500, detail=f"Internal server error {str(exception)}")
 
@@ -199,13 +211,12 @@ async def __call__(self, host: str, port: int, sockets: list[socket.socket] | No
                                 timeout_keep_alive=TIMEOUT_KEEP_ALIVE)
         await uvicorn.Server(config).serve(sockets=sockets)
 
-    # TODO: rework this for service discovery, now it's only for static server list
-    async def _set_steady_clock_offsets(self):
-        STEADY_CLOCK_OFFSET_ENDPOINT = "/steady_clock_offset"
+    async def _sync_server_clock(self, server: str):
+        """ Sync the ctx/gen server's steady clock with the disagg-server's steady clock (in case NTP service is not running). """
         async def query_steady_clock_offset(session: aiohttp.ClientSession, server_url: str) -> tuple[Optional[float], Optional[float]]:
             try:
                 originate_ts = get_steady_clock_now_in_seconds()
-                async with session.get(server_url + STEADY_CLOCK_OFFSET_ENDPOINT) as response:
+                async with session.get(server_url) as response:
                     destination_ts = get_steady_clock_now_in_seconds()
                     if response.status == 200:
                         response_content = await response.json()
@@ -222,12 +233,11 @@ async def query_steady_clock_offset(session: aiohttp.ClientSession, server_url:
 
         async def set_steady_clock_offset(session: aiohttp.ClientSession, server_url: str, offset: float) -> None:
             payload = {"offset": offset}
-            async with session.post(server_url + STEADY_CLOCK_OFFSET_ENDPOINT, json=payload) as response:
+            async with session.post(server_url, 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")
 
         async def align_steady_clock_offset(session: aiohttp.ClientSession, server_url: str) -> None:
-            server_url = f"http://{server_url}" if not server_url.startswith("http://") else server_url
             delay, offset = await query_steady_clock_offset(session, server_url)
             if delay is None or offset is None:
                 logger.warning(f"Unable to measure steady clock offset for {server_url}; skipping adjustment")
@@ -236,7 +246,13 @@ async def align_steady_clock_offset(session: aiohttp.ClientSession, server_url:
             # Negate the offset so that worker servers can adjust their steady clock by adding the new offset
             await set_steady_clock_offset(session, server_url, -offset)
 
-        async with aiohttp.ClientSession(
-            connector=aiohttp.TCPConnector(limit=0, limit_per_host=0, force_close=True),
-            timeout=aiohttp.ClientTimeout(total=self._req_timeout_secs)) as session:
-            await asyncio.gather(*[align_steady_clock_offset(session, server_url) for server_url in self._ctx_servers + self._gen_servers])
+        server_scheme = "http://" if not server.startswith("http://") else ""
+        server_url = f"{server_scheme}{server}/steady_clock_offset"
+
+        try:
+            async with aiohttp.ClientSession(
+                connector=aiohttp.TCPConnector(limit=0, limit_per_host=0, force_close=True),
+                timeout=aiohttp.ClientTimeout(total=self._req_timeout_secs)) as session:
+                await align_steady_clock_offset(session, server_url)
+        except (aiohttp.ClientError, OSError) as e:
+            logger.warning(f"Unable to align steady clock offset for {server_url}: {e}; skipping adjustment")

tensorrt_llm/serve/perf_metrics.py

@@ -15,7 +15,7 @@
 import asyncio
 from collections import defaultdict, deque
 from dataclasses import dataclass
-from typing import Any, Dict, List, Literal, Optional, Union
+from typing import Any, Dict, List, Literal, Optional
 
 from tensorrt_llm.llmapi.disagg_utils import ServerRole
 
@@ -64,7 +64,7 @@ class MetricsDefinition:
     buckets: Optional[List[float]] = None
 
 
-METRICS_DEFINITIONS = [
+CLIENT_METRICS_DEFINITIONS = [
     MetricsDefinition("total_requests", "Total number of requests", "counter"),
     MetricsDefinition("error_requests", "Total number of error requests", "counter"),
     MetricsDefinition("retry_requests", "Total number of retry requests", "counter"),
@@ -96,23 +96,29 @@ class MetricsDefinition:
 }
 
 
+def instance_metric(definition: MetricsDefinition, role: Optional[ServerRole] = None):
+    # import lazily to avoid breaking `set_prometheus_multiproc_dir`
+    from prometheus_client import Counter, Histogram
+
+    name = (
+        f"{ROLE_TO_CLIENT_TYPE[role]}_{definition.name}"
+        if role in ROLE_TO_CLIENT_TYPE
+        else definition.name
+    )
+    if definition.type == "counter":
+        return Counter(name, definition.description)
+    elif definition.type == "histogram":
+        return Histogram(name, definition.description, buckets=definition.buckets)
+    else:
+        raise ValueError(f"Invalid metric type: {definition.type}")
+
+
 class ClientMetricsCollector:
     def __init__(self, role: ServerRole):
         self._role = role
-        # import lazily to avoid breaking `set_prometheus_multiproc_dir`
-        from prometheus_client import Counter, Histogram
-
-        def instance_metric(definition: MetricsDefinition) -> Union[Counter | Histogram]:
-            name = f"{ROLE_TO_CLIENT_TYPE[role]}_{definition.name}"
-            if definition.type == "counter":
-                return Counter(name, definition.description)
-            elif definition.type == "histogram":
-                return Histogram(name, definition.description, buckets=definition.buckets)
-            else:
-                raise ValueError(f"Invalid metric type: {definition.type}")
-
         self._metrics = {
-            definition.name: instance_metric(definition) for definition in METRICS_DEFINITIONS
+            definition.name: instance_metric(definition, role)
+            for definition in CLIENT_METRICS_DEFINITIONS
         }
 
     def __getattr__(
@@ -121,17 +127,41 @@ def __getattr__(
         return self._metrics[key]
 
 
+SERVER_METRICS_DEFINITIONS = [
+    MetricsDefinition("total_requests", "Total number of requests", "counter"),
+    MetricsDefinition("stream_requests", "Total number of stream requests", "counter"),
+    MetricsDefinition("nonstream_requests", "Total number of non-stream requests", "counter"),
+    MetricsDefinition("validation_exceptions", "Total number of validation exceptions", "counter"),
+    MetricsDefinition("http_exceptions", "Total number of HTTP exceptions", "counter"),
+    MetricsDefinition("internal_errors", "Total number of internal errors", "counter"),
+    MetricsDefinition("total_responses", "Total number of responses", "counter"),
+    MetricsDefinition(
+        "queue_latency_seconds",
+        "Histogram of latency from request arrival to being processed in seconds",
+        "histogram",
+        SHORT_TIME_BUCKETS,
+    ),
+]
+
+
 class DisaggPerfMetricsCollector:
     def __init__(self, max_requests: int):
         self._max_requests = max_requests
         self._request_meteics = deque(maxlen=max_requests)
         self._server_metrics = defaultdict(dict)
         self._lock = asyncio.Lock()
         self._clients = []
+        self._metrics = {
+            definition.name: instance_metric(definition)
+            for definition in SERVER_METRICS_DEFINITIONS
+        }
 
     def add_client(self, client):
         self._clients.append(client)
 
+    def __getattr__(self, key: str):
+        return self._metrics[key]
+
     async def add_per_request_metrics(
         self,
         ctx_server: str,