feat: Dynamically remove servers in PD (NVIDIA#5270)

Signed-off-by: Shunkang <[email protected]>

Author: Shunkangz

tensorrt_llm/llmapi/disagg_utils.py

@@ -58,6 +58,7 @@ class MetadataServerConfig():
     hostname: str = "localhost"
     port: int = 2379
     health_check_timeout: float = 5.0
+    refresh_interval: float = 10.0
 
 
 def parse_disagg_config_file(yaml_config_file: str):

tensorrt_llm/serve/openai_disagg_server.py

@@ -35,20 +35,20 @@
 class OpenAIDisaggServer:
 
     def __init__(self,
-                 ctx_servers: List[str] = None,
-                 gen_servers: List[str] = None,
+                 ctx_servers: List[str],
+                 gen_servers: List[str],
                  req_timeout_secs: int = 180,
                  server_start_timeout_secs: int = 180,
                  ctx_router_config: Optional[RouterConfig] = None,
                  gen_router_config: Optional[RouterConfig] = None,
                  conditional_disagg_config: Optional[ConditionalDisaggConfig] = None,
-                 metadata_server_cfg: MetadataServerConfig = None):
+                 metadata_server_cfg: Optional[MetadataServerConfig] = None):
 
         self.ctx_servers = ctx_servers
         self.gen_servers = gen_servers
         self.metadata_server = create_metadata_server(metadata_server_cfg)
-        self.ctx_router = create_router(ctx_router_config, ctx_servers, self.metadata_server)
-        self.gen_router = create_router(gen_router_config, gen_servers, self.metadata_server)
+        self.ctx_router = create_router(ctx_router_config, ctx_servers, metadata_server_cfg, self.metadata_server)
+        self.gen_router = create_router(gen_router_config, gen_servers, metadata_server_cfg, self.metadata_server)
         self.conditional_disagg_config = conditional_disagg_config
 
 
@@ -76,8 +76,8 @@ async def lifespan(app: FastAPI):
 
             if self.metadata_server:
                 logger.info("Starting server monitoring via metadata service")
-                await self.ctx_router.start_server_monitoring()
-                await self.gen_router.start_server_monitoring()
+                await self.ctx_router.start_server_monitoring(metadata_server_cfg.refresh_interval)
+                await self.gen_router.start_server_monitoring(metadata_server_cfg.refresh_interval)
 
             yield