feat: support more PD node select func

lightllm/server/api_cli.py

@@ -42,6 +42,13 @@ def make_argument_parser() -> argparse.ArgumentParser:
         default=42000,
         help="p d mode, decode node used for kv move manager rpyc server port",
     )
+    parser.add_argument(
+        "--select_p_d_node_func",
+        type=str,
+        default="round_robin",
+        choices=["random", "round_robin", "memory"],
+        help="select p d node func, can be round_robin, random or memory",
+    )
     parser.add_argument(
         "--config_server_host",
         type=str,

lightllm/server/httpserver/pd_loop.py

@@ -176,9 +176,29 @@ async def _pd_process_generate(
         logger.error(str(e))
 
 
+# 获取节点负载信息
+def _get_load_info(have_finished_req: bool) -> dict:
+    if not have_finished_req:
+        return None
+
+    from lightllm.server.api_http import g_objs
+    assert g_objs.shared_token_load is not None, "shared_token_load is not initialized"
+    current_load = [
+        float(g_objs.shared_token_load.get_dynamic_max_load(dp_index)) for dp_index in range(g_objs.args.dp)
+    ]
+    load_info = {
+        "mem_len": min(current_load),
+        "client_ip_port": f"{g_objs.httpserver_manager.host_ip}:{g_objs.args.port}"
+    }
+    return load_info
+
+
 # 转发token的task
 async def _up_tokens_to_pd_master(forwarding_queue: AsyncQueue, websocket):
     while True:
         handle_list = await forwarding_queue.wait_to_get_all_data()
+
         if handle_list:
-            await websocket.send(pickle.dumps((ObjType.TOKEN_PACKS, handle_list)))
+            have_finished_req = any(finish_status.is_finished() for _, _, _, finish_status in handle_list)
+            load_info: dict = _get_load_info(have_finished_req)
+            await websocket.send(pickle.dumps((ObjType.TOKEN_PACKS, handle_list, load_info)))

lightllm/server/httpserver_for_pd_master/manager.py

@@ -25,59 +25,99 @@
 from lightllm.utils.statics_utils import MovingAverage
 from lightllm.server.httpserver.manager import AsyncQueue
 from lightllm.utils.error_utils import ServerBusyError
+from .node_info_recorder import PredictNodeInfoRecorder
+from .pd_selector import (
+    create_selector,
+    PDSelector,
+)
 
 logger = init_logger(__name__)
 
 
-class HttpServerManagerForPDMaster:
-    def __init__(
-        self,
-        args,
-        metric_port,
-    ):
+class PDManager:
+    def __init__(self, args):
         self.args = args
-        self.metric_client = MetricClient(metric_port)
-        self.id_gen = ReqIDGenerator()
         self.prefill_nodes: List[PD_Client_Obj] = []
         self.decode_nodes: List[PD_Client_Obj] = []
-        self.url_to_pd_nodes: Dict[str, PD_Client_Obj] = {}
-
-        self.req_id_to_out_inf: Dict[int, ReqStatus] = {}
-        self.infos_queues = None  # 这个需要延迟初始化，否则使用的loop不对
-
-        self.tokenizer = get_tokenizer(args.model_dir, args.tokenizer_mode, trust_remote_code=args.trust_remote_code)
-
-        self.first_time_costs = MovingAverage()
-        self.per_token_costs = MovingAverage()
+        self.node_info_recorder: PredictNodeInfoRecorder = PredictNodeInfoRecorder()
+        self.selector: PDSelector = create_selector(args.select_p_d_node_func, self.prefill_nodes, self.decode_nodes, self)
         return
 
     async def register_pd(self, pd_info_json, websocket):
         pd_client = PD_Client_Obj(**pd_info_json)
         pd_client.websocket = websocket
-        self.url_to_pd_nodes[pd_client.client_ip_port] = pd_client
+        self.node_info_recorder.register_node(pd_client)
+
         if pd_client.mode == "prefill":
             self.prefill_nodes = [e for e in self.prefill_nodes if e.client_ip_port != pd_client.client_ip_port]
             self.prefill_nodes.append(pd_client)
         elif pd_client.mode == "decode":
             self.decode_nodes = [e for e in self.decode_nodes if e.client_ip_port != pd_client.client_ip_port]
             self.decode_nodes.append(pd_client)
         else:
-            assert False
+            assert False, f"mode must in ['prefill', 'decode'], but get {pd_client.mode}"
+
+        await self.selector.update_nodes(self.prefill_nodes, self.decode_nodes)
 
         logger.info(f"mode: {pd_client.mode} url: {pd_client.client_ip_port} registed")
         return
 
     async def remove_pd(self, pd_info_json):
         pd_client = PD_Client_Obj(**pd_info_json)
-        try:
-            del self.url_to_pd_nodes[pd_client.client_ip_port]
-        except:
-            pass
+        self.node_info_recorder.remove_node(pd_client)
+
         self.prefill_nodes = [e for e in self.prefill_nodes if e.client_ip_port != pd_client.client_ip_port]
         self.decode_nodes = [e for e in self.decode_nodes if e.client_ip_port != pd_client.client_ip_port]
+
+        await self.selector.update_nodes(self.prefill_nodes, self.decode_nodes)
+
         logger.info(f"mode: {pd_client.mode} url: {pd_client.client_ip_port} removed")
         return
 
+    def update_node_load_info(self, load_info: dict):
+        """更新节点负载信息"""
+        if load_info is None:
+            return
+        self.node_info_recorder.update_node_load_info(load_info)
+
+    def get_predict_node_infos(self):
+        """获取所有节点的预测负载信息"""
+        return self.node_info_recorder.get_predict_node_infos()
+
+    async def select_p_d_node(self, prompt: Union[str, List[int]], sampling_params: SamplingParams, multimodal_params: MultimodalParams) -> Tuple[PD_Client_Obj, PD_Client_Obj]:
+        p_node, d_node = await self.selector.select_p_d_node(prompt, sampling_params, multimodal_params)
+        self.node_info_recorder.update_predict_node_info(p_node, d_node, prompt, sampling_params, multimodal_params)
+        return p_node, d_node
+
+class HttpServerManagerForPDMaster:
+    def __init__(
+        self,
+        args,
+        metric_port,
+    ):
+        self.args = args
+        self.metric_client = MetricClient(metric_port)
+        self.id_gen = ReqIDGenerator()
+
+        self.pd_manager = PDManager(args)
+
+        self.req_id_to_out_inf: Dict[int, ReqStatus] = {}
+        self.infos_queues = None  # 这个需要延迟初始化，否则使用的loop不对
+
+        self.tokenizer = get_tokenizer(args.model_dir, args.tokenizer_mode, trust_remote_code=args.trust_remote_code)
+
+        self.first_time_costs = MovingAverage()
+        self.per_token_costs = MovingAverage()
+        return
+
+    async def register_pd(self, pd_info_json, websocket):
+        await self.pd_manager.register_pd(pd_info_json, websocket)
+        return
+
+    async def remove_pd(self, pd_info_json):
+        await self.pd_manager.remove_pd(pd_info_json)
+        return
+
     async def update_req_status(self, upkv_status: UpKVStatus):
         try:
             group_request_id = convert_sub_id_to_group_id(upkv_status.group_request_id)
@@ -108,11 +148,7 @@ def tokens(self, prompt, multimodal_params, samping_params: SamplingParams, kwar
     async def select_p_d_node(
         self, prompt: Union[str, List[int]], sampling_params: SamplingParams, multimodal_params: MultimodalParams
     ) -> Tuple[PD_Client_Obj, PD_Client_Obj]:
-        import random
-
-        p_node = random.choice(self.prefill_nodes)
-        d_node = random.choice(self.decode_nodes)
-        return p_node, d_node
+        return await self.pd_manager.select_p_d_node(prompt, sampling_params, multimodal_params)
 
     async def generate(
         self,
@@ -264,7 +300,7 @@ async def _wait_to_token_package(
         request: Request,
     ):
         out_token_counter = 0
-        first_token_cost_ms = sys.float_info.max
+        first_token_cost_ms = float('inf')
         group_request_id = sampling_params.group_request_id
         unfinished_count = sampling_params.best_of
         is_first_token = True
@@ -368,7 +404,16 @@ async def handle_loop(self):
             try:
                 for obj in objs:
                     if obj[0] == ObjType.TOKEN_PACKS:
-                        for sub_req_id, text, metadata, finish_status in obj[1]:
+                        # 检查是否包含节点信息
+                        if len(obj) >= 3:
+                            handle_list, load_info = obj[1], obj[2]
+                            # 更新节点负载信息
+                            self.pd_manager.update_node_load_info(load_info)
+                        else:
+                            # 兼容旧格式
+                            handle_list = obj[1]
+
+                        for sub_req_id, text, metadata, finish_status in handle_list:
                             finish_status: FinishStatus = finish_status
                             group_req_id = convert_sub_id_to_group_id(sub_req_id)
                             try:

lightllm/server/httpserver_for_pd_master/node_info_recorder.py

@@ -0,0 +1,91 @@
+import copy
+
+from ..pd_io_struct import PD_Client_Obj
+from lightllm.server.httpserver.manager import SamplingParams, MultimodalParams
+from typing import Union, List, Dict
+from lightllm.utils.log_utils import init_logger
+
+logger = init_logger(__name__)
+
+
+class NodeInfoRecorder:
+    def __init__(self):
+        self.prefill_node_info: dict = {}
+        self.decode_node_info: dict = {}
+
+    def register_node(self, pd_client: PD_Client_Obj):
+        node_info = {
+            "node_id": pd_client.node_id,
+            "client_ip_port": pd_client.client_ip_port,
+            "mode": pd_client.mode,
+            "node": pd_client,
+            "mem_len": 0,
+            # "batch_size": 0,
+        }
+        if pd_client.mode == "prefill":
+            self.prefill_node_info[pd_client.client_ip_port] = node_info
+        elif pd_client.mode == "decode":
+            self.decode_node_info[pd_client.client_ip_port] = node_info
+        else:
+            assert False, f"mode must in ['prefill', 'decode'], but get {pd_client.mode}"
+
+    def remove_node(self, pd_client: PD_Client_Obj):
+        if pd_client.mode == "prefill":
+            del self.prefill_node_info[pd_client.client_ip_port]
+        elif pd_client.mode == "decode":
+            del self.decode_node_info[pd_client.client_ip_port]
+        else:
+            assert False, f"mode must in ['prefill', 'decode'], but get {pd_client.mode}"
+
+    def update_node_load_info(self, load_info: dict):
+        if "client_ip_port" in load_info:
+            ip_port = load_info["client_ip_port"]
+            if ip_port in self.prefill_node_info:
+                self.prefill_node_info[ip_port]["mem_len"] = load_info["mem_len"]
+            elif ip_port in self.decode_node_info:
+                self.decode_node_info[ip_port]["mem_len"] = load_info["mem_len"]
+            else:
+                logger.warning(f"Received load info for unknown node: {ip_port}")
+        else:
+            logger.warning("Received load info without client_ip_port")
+
+
+class PredictNodeInfoRecorder(NodeInfoRecorder):
+    def __init__(self):
+        super().__init__()
+        self.prefill_predict_node_info: dict = {}
+        self.decode_predict_node_info: dict = {}
+
+    def register_node(self, pd_client: PD_Client_Obj):
+        super().register_node(pd_client)
+        if pd_client.mode == "prefill":
+            self.prefill_predict_node_info[pd_client.client_ip_port] = copy.copy(self.prefill_node_info[pd_client.client_ip_port])
+        elif pd_client.mode == "decode":
+            self.decode_predict_node_info[pd_client.client_ip_port] = copy.copy(self.decode_node_info[pd_client.client_ip_port])
+
+    def remove_node(self, pd_client: PD_Client_Obj):
+        super().remove_node(pd_client)
+        if pd_client.mode == "prefill":
+            del self.prefill_predict_node_info[pd_client.client_ip_port]
+        elif pd_client.mode == "decode":
+            del self.decode_predict_node_info[pd_client.client_ip_port]
+
+    def update_node_load_info(self, load_info: dict):
+        super().update_node_load_info(load_info)
+        ip_port = load_info["client_ip_port"]
+        if ip_port in self.prefill_node_info:
+            self.prefill_predict_node_info[ip_port] = copy.copy(self.prefill_node_info[ip_port])
+        elif ip_port in self.decode_node_info:
+            self.decode_predict_node_info[ip_port] = copy.copy(self.decode_node_info[ip_port])
+        else:
+            logger.warning(f"Received load info for unknown node: {ip_port}")
+
+    def update_predict_node_info(self, p_node: PD_Client_Obj, d_node: PD_Client_Obj, prompt: Union[str, List[int]], sampling_params: SamplingParams, multimodal_params: MultimodalParams):
+        self.prefill_predict_node_info[p_node.client_ip_port]["mem_len"] += len(prompt)
+        self.decode_predict_node_info[d_node.client_ip_port]["mem_len"] += sampling_params.max_new_tokens
+
+    def get_predict_node_infos(self) -> Dict[str, dict]:
+        return {
+            "prefill": self.prefill_predict_node_info,
+            "decode": self.decode_predict_node_info,
+        }

lightllm/server/httpserver_for_pd_master/pd_selector/__init__.py

@@ -0,0 +1,25 @@
+from typing import List
+from lightllm.server.httpserver_for_pd_master.manager import PD_Client_Obj
+from .pd_selector import (
+    PDSelector,
+    RandomSelector,
+    RoundRobinSelector,
+    MemorySelector
+)
+
+__all__ = [
+    "PDSelector",
+    "RandomSelector", 
+    "RoundRobinSelector",
+    "MemorySelector"
+]
+
+def create_selector(selector_type: str, prefill_nodes: List[PD_Client_Obj], decode_nodes: List[PD_Client_Obj], pd_manager) -> PDSelector:
+    if selector_type == "random":
+        return RandomSelector(prefill_nodes, decode_nodes, pd_manager)
+    elif selector_type == "round_robin":
+        return RoundRobinSelector(prefill_nodes, decode_nodes, pd_manager)
+    elif selector_type == "memory":
+        return MemorySelector(prefill_nodes, decode_nodes, pd_manager)
+    else:
+        raise ValueError(f"Invalid selector type: {selector_type}")