[TRTLLM-7843][feat] implement disagg cluster auto-scaling (NVIDIA#8215)

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

Author: reasonsolo

requirements.txt

@@ -67,7 +67,7 @@ nvtx
 matplotlib # FIXME: this is added to make nvtx happy
 meson
 ninja
-etcd3
+etcd3 @ git+https://github.com/kragniz/python-etcd3.git@e58a899579ba416449c4e225b61f039457c8072a
 blake3
 soundfile
 triton==3.3.1; platform_machine == "x86_64"

tensorrt_llm/commands/serve.py

@@ -21,7 +21,9 @@
 from tensorrt_llm.llmapi import (BuildConfig, CapacitySchedulerPolicy,
                                  DynamicBatchConfig, KvCacheConfig,
                                  SchedulerConfig)
-from tensorrt_llm.llmapi.disagg_utils import (MetadataServerConfig, ServerRole,
+from tensorrt_llm.llmapi.disagg_utils import (DisaggClusterConfig,
+                                              MetadataServerConfig, ServerRole,
+                                              extract_disagg_cluster_config,
                                               parse_disagg_config_file,
                                               parse_metadata_server_config_file)
 from tensorrt_llm.llmapi.llm_utils import update_llm_args_with_extra_dict
@@ -140,7 +142,8 @@ def launch_server(host: str,
                   port: int,
                   llm_args: dict,
                   metadata_server_cfg: Optional[MetadataServerConfig] = None,
-                  server_role: Optional[ServerRole] = None):
+                  server_role: Optional[ServerRole] = None,
+                  disagg_cluster_config: Optional[DisaggClusterConfig] = None):
 
     backend = llm_args["backend"]
     model = llm_args["model"]
@@ -161,7 +164,8 @@ def launch_server(host: str,
     server = OpenAIServer(llm=llm,
                           model=model,
                           server_role=server_role,
-                          metadata_server_cfg=metadata_server_cfg)
+                          metadata_server_cfg=metadata_server_cfg,
+                          disagg_cluster_config=disagg_cluster_config)
 
     # Optionally disable GC (default: not disabled)
     if os.getenv("TRTLLM_SERVER_DISABLE_GC", "0") == "1":
@@ -313,6 +317,10 @@ def convert(self, value: Any, param: Optional["click.Parameter"],
     help=
     "Exit with runtime error when attention window is too large to fit even a single sequence in the KV cache."
 )
+@click.option("--disagg_cluster_uri",
+              type=str,
+              default=None,
+              help="URI of the disaggregated cluster.")
 @click.option("--enable_chunked_prefill",
               is_flag=True,
               default=False,
@@ -327,7 +335,7 @@ def serve(
         extra_llm_api_options: Optional[str], reasoning_parser: Optional[str],
         metadata_server_config_file: Optional[str], server_role: Optional[str],
         fail_fast_on_attention_window_too_large: bool,
-        enable_chunked_prefill: bool):
+        enable_chunked_prefill: bool, disagg_cluster_uri: Optional[str]):
     """Running an OpenAI API compatible server
 
     MODEL: model name | HF checkpoint path | TensorRT engine path
@@ -364,14 +372,27 @@ def serve(
     metadata_server_cfg = parse_metadata_server_config_file(
         metadata_server_config_file)
 
-    if metadata_server_cfg is not None:
-        assert server_role is not None, "server_role is required when metadata_server_cfg is provided"
+    # Specify disagg_cluster_config in config file or through command line "--disagg_cluster_uri",
+    # but disagg_cluster_uri takes precedence over cluster uri in config file
+    disagg_cluster_config = llm_args.pop("disagg_cluster", None)
+    if disagg_cluster_config:
+        disagg_cluster_config = extract_disagg_cluster_config(
+            disagg_cluster_config, disagg_cluster_uri)
+    elif disagg_cluster_uri:
+        disagg_cluster_config = DisaggClusterConfig(
+            cluster_uri=disagg_cluster_uri)
+
+    if metadata_server_cfg is not None or disagg_cluster_config is not None:
+        assert (
+            server_role is not None
+        ), "server_role is required when metadata_server_cfg or disagg_cluster_config is provided"
         try:
             server_role = ServerRole[server_role.upper()]
         except ValueError:
             raise ValueError(f"Invalid server role: {server_role}. " \
                              f"Must be one of: {', '.join([role.name for role in ServerRole])}")
-    launch_server(host, port, llm_args, metadata_server_cfg, server_role)
+    launch_server(host, port, llm_args, metadata_server_cfg, server_role,
+                  disagg_cluster_config)
 
 
 @click.command("mm_embedding_serve")

tensorrt_llm/llmapi/disagg_utils.py

@@ -1,7 +1,7 @@
 import logging
 from dataclasses import dataclass, field
 from enum import IntEnum
-from typing import Any, List, Literal, Optional, Tuple
+from typing import Any, Dict, List, Literal, Optional, Tuple
 
 import yaml
 from mpi4py.MPI import COMM_WORLD, Comm
@@ -68,6 +68,7 @@ class DisaggServerConfig():
     conditional_disagg_config: Optional[ConditionalDisaggConfig] = None
     max_retries: int = 1
     perf_metrics_max_requests: int = 0
+    disagg_cluster_config: Optional[DisaggClusterConfig] = None
 
 
 @dataclass
@@ -111,6 +112,7 @@ def extract_disagg_cfg(hostname: str = 'localhost',
                        context_servers: Optional[dict] = None,
                        generation_servers: Optional[dict] = None,
                        conditional_disagg_config: Optional[dict] = None,
+                       disagg_cluster: Optional[dict] = None,
                        **kwargs: Any) -> DisaggServerConfig:
     context_servers = context_servers or {}
     generation_servers = generation_servers or {}
@@ -131,23 +133,27 @@ def extract_disagg_cfg(hostname: str = 'localhost',
                 # Inherit the value from the top-level
                 servers[key] = value
 
+    server_configs = []
+    disagg_cluster_config = None
     ctx_router_config = extract_router_config(context_servers)
     gen_router_config = extract_router_config(generation_servers)
-
-    server_configs = extract_ctx_gen_cfgs(
-        type="ctx", **context_servers) + extract_ctx_gen_cfgs(
-            type="gen", **generation_servers)
-
     ctx_router_config.server_role = ServerRole.CONTEXT
     gen_router_config.server_role = ServerRole.GENERATION
+    if disagg_cluster:
+        disagg_cluster_config = extract_disagg_cluster_config(disagg_cluster)
+    else:
+        server_configs = extract_ctx_gen_cfgs(
+            type="ctx", **context_servers) + extract_ctx_gen_cfgs(
+                type="gen", **generation_servers)
 
     conditional_disagg_config = ConditionalDisaggConfig(
         **conditional_disagg_config) if conditional_disagg_config else None
 
     config = DisaggServerConfig(server_configs, hostname, port,
                                 ctx_router_config, gen_router_config,
                                 conditional_disagg_config, max_retries,
-                                perf_metrics_max_requests)
+                                perf_metrics_max_requests,
+                                disagg_cluster_config)
 
     return config
 
@@ -235,6 +241,33 @@ def get_server_configs_dict(
     return num_workers, server_dict
 
 
+def extract_disagg_cluster_config(
+        cluster_config_dict: Dict[str, Any],
+        cluster_uri: Optional[str] = None) -> DisaggClusterConfig:
+    """
+    Build the DisaggClusterConfig from the cluster_config_dict.
+    Use the default value of DisaggClusterConfig and MinimalInstances if the corresponding fields are not provided.
+    If cluster_uri is provided, it will override the cluster_uri in the cluster_config_dict.
+    """
+
+    def update_dataclass(obj, data_dict: Dict[str, Any]):
+        for key, value in data_dict.items():
+            if key not in obj.__dataclass_fields__:
+                raise KeyError(
+                    f"Key {key} not found in {obj.__class__.__name__}")
+            if value is not None:
+                setattr(obj, key, value)
+        return obj
+
+    cluster_config_dict["minimal_instances"] = update_dataclass(
+        MinimalInstances(), cluster_config_dict.get("minimal_instances", {}))
+    cluster_config = update_dataclass(
+        DisaggClusterConfig(cluster_uri or cluster_config_dict["cluster_uri"]),
+        cluster_config_dict,
+    )
+    return cluster_config
+
+
 def split_world_comm(
         server_configs: List[CtxGenServerConfig]) -> Tuple[bool, int, Comm]:
 

tensorrt_llm/serve/cluster_storage.py

@@ -281,9 +281,10 @@ async def _check_expired(self):
                         self._storage.pop(k)
                 for k, v in kv_to_delete.items():
                     await self._notify_watch_event(k, v, WatchEventType.DELETE)
-                logger.debug(
-                    f"Checked expired, {before_len} -> {len(self._storage)}, keys to delete: {kv_to_delete.keys()}"
-                )
+                if len(kv_to_delete) > 0:
+                    logger.debug(
+                        f"Checked expired, {before_len} -> {len(self._storage)}, keys to delete: {kv_to_delete.keys()}"
+                    )
             except Exception as e:
                 logger.error(f"Error checking expired: {e}")
 
@@ -298,9 +299,12 @@ def __init__(self, cluster_uri, cluster_name):
         self._cluster_name = cluster_name
 
     def __del__(self):
-        if asyncio.get_event_loop():
-            asyncio.run_coroutine_threadsafe(self._session.close(),
-                                             asyncio.get_event_loop())
+        try:
+            if asyncio.get_event_loop():
+                asyncio.run_coroutine_threadsafe(self._session.close(),
+                                                 asyncio.get_event_loop())
+        except RuntimeError:
+            pass
 
     def _url_for(self, endpoint: str) -> str:
         return f"{self._cluster_uri}/{endpoint}"

tensorrt_llm/serve/disagg_auto_scaling.py

@@ -46,9 +46,13 @@ def __init__(self, config: DisaggClusterConfig, storage: ClusterStorage):
         self._watch_handle = None
 
     def __del__(self):
-        if asyncio.get_event_loop():
-            asyncio.run_coroutine_threadsafe(self.stop(),
-                                             asyncio.get_event_loop())
+        try:
+            if asyncio.get_event_loop():
+                asyncio.run_coroutine_threadsafe(self.stop(),
+                                                 asyncio.get_event_loop())
+        except RuntimeError:
+            # the event loop may not be running when the cluster manager is destroyed
+            pass
 
     async def start(self) -> None:
         await self._cluster_storage.start()
@@ -208,9 +212,13 @@ def __init__(self, role: ServerRole, host: str, port: int,
         self._worker_id = f"{role.name}-{host}:{port}-{int(time.time()*1000)}-{os.getpid()}-{random.randint(0, 1000):03}"
 
     def __del__(self):
-        if asyncio.get_event_loop():
-            asyncio.run_coroutine_threadsafe(self.deregister_worker(),
-                                             asyncio.get_event_loop())
+        try:
+            if asyncio.get_event_loop():
+                asyncio.run_coroutine_threadsafe(self.deregister_worker(),
+                                                 asyncio.get_event_loop())
+        except RuntimeError:
+            # the event loop may not be running when the worker is destroyed
+            pass
 
     @property
     def worker_id(self) -> str: