support more PD node select func. such as random or roundrobin. (#1018)

Author: hiworldwzj

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_strategy",
+        type=str,
+        default="round_robin",
+        choices=["random", "round_robin", "adaptive_load"],
+        help="pd master use this strategy to select p d node, can be round_robin, random or adaptive_load",
+    )
     parser.add_argument(
         "--config_server_host",
         type=str,

lightllm/server/httpserver/pd_loop.py

@@ -180,5 +180,29 @@ async def _pd_process_generate(
 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)))
+            load_info: dict = _get_load_info()
+            await websocket.send(pickle.dumps((ObjType.TOKEN_PACKS, handle_list, load_info)))
+
+
+# 获取节点负载信息
+def _get_load_info() -> dict:
+
+    from lightllm.server.api_http import g_objs
+
+    assert g_objs.shared_token_load is not None, "shared_token_load is not initialized"
+    args = g_objs.args
+    dp_size_in_node = max(1, args.dp // args.nnodes)
+
+    # 获取当前每个 dp 的负载，数值含义为当前的 token 总容量使用率， 上报给 PD_Master 用于做
+    # 调度决策。
+    current_load = [
+        float(g_objs.shared_token_load.get_dynamic_max_load(dp_index)) for dp_index in range(dp_size_in_node)
+    ]
+    mean_node_load = sum(current_load) / len(current_load)
+    load_info = {
+        "total_token_usage_rate": mean_node_load,
+        "client_ip_port": f"{g_objs.httpserver_manager.host_ip}:{g_objs.args.port}",
+    }
+    return load_info

lightllm/server/httpserver_for_pd_master/manager.py

@@ -12,7 +12,7 @@
 import pickle
 
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
-from typing import Union, List, Tuple, Dict
+from typing import Union, List, Tuple, Dict, Optional
 from lightllm.server.core.objs import FinishStatus
 from ..pd_io_struct import PD_Client_Obj, UpKVStatus, ObjType
 from lightllm.server.core.objs import SamplingParams
@@ -25,6 +25,7 @@
 from lightllm.utils.statics_utils import MovingAverage
 from lightllm.server.httpserver.manager import AsyncQueue
 from lightllm.utils.error_utils import ServerBusyError
+from .pd_selector import create_selector
 
 logger = init_logger(__name__)
 
@@ -38,9 +39,8 @@ def __init__(
         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.pd_manager = PDManager(args)
 
         self.req_id_to_out_inf: Dict[int, ReqStatus] = {}
         self.infos_queues = None  # 这个需要延迟初始化，否则使用的loop不对
@@ -52,30 +52,11 @@ def __init__(
         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
-        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
-
-        logger.info(f"mode: {pd_client.mode} url: {pd_client.client_ip_port} registed")
+        self.pd_manager.register_pd(pd_info_json, websocket)
         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.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]
-        logger.info(f"mode: {pd_client.mode} url: {pd_client.client_ip_port} removed")
+        self.pd_manager.remove_pd(pd_info_json)
         return
 
     async def update_req_status(self, upkv_status: UpKVStatus):
@@ -108,11 +89,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 self.pd_manager.select_p_d_node(prompt, sampling_params, multimodal_params)
 
     async def generate(
         self,
@@ -264,7 +241,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 +345,10 @@ 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]:
+                        token_list, node_load_info = obj[1], obj[2]
+                        self.pd_manager.update_node_load_info(node_load_info)
+
+                        for sub_req_id, text, metadata, finish_status in token_list:
                             finish_status: FinishStatus = finish_status
                             group_req_id = convert_sub_id_to_group_id(sub_req_id)
                             try:
@@ -415,3 +395,69 @@ async def pop_all_tokens(self):
             ans = self.out_token_info_list.copy()
             self.out_token_info_list.clear()
         return ans
+
+
+class PDManager:
+    def __init__(self, args):
+        self.args = args
+        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.selector = create_selector(args.select_p_d_node_strategy, self)
+        return
+
+    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
+
+        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, f"mode must in ['prefill', 'decode'], but get {pd_client.mode}"
+
+        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
+
+    def remove_pd(self, pd_info_json):
+        pd_client = PD_Client_Obj(**pd_info_json)
+
+        self.url_to_pd_nodes.pop(pd_client.client_ip_port, None)
+        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]
+
+        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: Optional[dict]):
+        """更新节点负载信息
+        load_info: 节点负载信息字典，内容格式如下，可以为 None
+        {
+        "total_token_usage_rate": xxxx,
+        "client_ip_port": xxxx,
+        }
+        """
+        try:
+            if load_info is None:
+                return
+            client_ip_port = load_info["client_ip_port"]
+            total_token_usage_rate = load_info["total_token_usage_rate"]
+            pd_client = self.url_to_pd_nodes.get(client_ip_port)
+            pd_client.run_status.total_token_usage_rate = total_token_usage_rate
+        except BaseException as e:
+            logger.warning(f"udpate node load info failed, load_info: {load_info} error: {str(e)}")
+        return
+
+    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 = self.selector.select_p_d_node(prompt, sampling_params, multimodal_params)
+        return p_node, d_node

lightllm/server/httpserver_for_pd_master/pd_selector/__init__.py

@@ -0,0 +1,12 @@
+from .pd_selector import PDSelector, RandomSelector, RoundRobinSelector, AdaptiveLoadSelector
+
+
+def create_selector(selector_type: str, pd_manager) -> PDSelector:
+    if selector_type == "random":
+        return RandomSelector(pd_manager)
+    elif selector_type == "round_robin":
+        return RoundRobinSelector(pd_manager)
+    elif selector_type == "adaptive_load":
+        return AdaptiveLoadSelector(pd_manager)
+    else:
+        raise ValueError(f"Invalid selector type: {selector_type}")

lightllm/server/httpserver_for_pd_master/pd_selector/pd_selector.py

@@ -0,0 +1,67 @@
+import random
+from typing import Union, List, Tuple, Dict
+from lightllm.server.pd_io_struct import PD_Client_Obj
+from lightllm.server.core.objs import SamplingParams
+from lightllm.server.multimodal_params import MultimodalParams
+
+
+class PDSelector:
+    def __init__(self, pd_manager):
+        self.prefill_nodes: List[PD_Client_Obj] = []
+        self.decode_nodes: List[PD_Client_Obj] = []
+        self.pd_manager = pd_manager
+
+    def update_nodes(self, prefill_nodes, decode_nodes):
+        self.prefill_nodes = prefill_nodes
+        self.decode_nodes = decode_nodes
+
+    def select_p_d_node(
+        self, prompt: Union[str, List[int]], sampling_params: SamplingParams, multimodal_params: MultimodalParams
+    ) -> Tuple[PD_Client_Obj, PD_Client_Obj]:
+        raise NotImplementedError("Subclass must implement this method")
+
+
+class RandomSelector(PDSelector):
+    """随机选择器"""
+
+    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 = random.choice(self.prefill_nodes)
+        d_node = random.choice(self.decode_nodes)
+        return p_node, d_node
+
+
+class RoundRobinSelector(PDSelector):
+    """轮询选择器"""
+
+    def __init__(self, pd_manager):
+        super().__init__(pd_manager)
+        self.prefill_node_index: int = 0
+        self.decode_node_index: int = 0
+
+    def select_p_d_node(
+        self, prompt: Union[str, List[int]], sampling_params: SamplingParams, multimodal_params: MultimodalParams
+    ) -> Tuple[PD_Client_Obj, PD_Client_Obj]:
+        self.prefill_node_index = self.prefill_node_index % len(self.prefill_nodes)
+        self.decode_node_index = self.decode_node_index % len(self.decode_nodes)
+        p_node = self.prefill_nodes[self.prefill_node_index]
+        d_node = self.decode_nodes[self.decode_node_index]
+        self.prefill_node_index += 1
+        self.decode_node_index += 1
+        return p_node, d_node
+
+
+class AdaptiveLoadSelector(PDSelector):
+    """基于负载使用情况的选择器"""
+
+    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 = self._importance_sampling(self.prefill_nodes)
+        d_node = self._importance_sampling(self.decode_nodes)
+
+        return p_node, d_node
+
+    def _importance_sampling(self, nodes: List[PD_Client_Obj]):
+        return random.choices(nodes, weights=[max(1.0 - e.run_status.total_token_usage_rate, 0.02) for e in nodes])

lightllm/server/pd_io_struct.py

@@ -1,6 +1,6 @@
 import enum
 import time
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Dict, List, Optional, Tuple, Union
 from lightllm.server.req_id_generator import convert_sub_id_to_group_id
 from fastapi import WebSocket
@@ -38,13 +38,19 @@ class ObjType(enum.Enum):
     TOKEN_PACKS = 3
 
 
+@dataclass
+class _PD_Client_RunStatus:
+    total_token_usage_rate: float = 0.0  # pd 节点上的 token 使用率
+
+
 @dataclass
 class PD_Client_Obj:
     node_id: int
     client_ip_port: str
     mode: str  # 只能是 prefill 或者 decode 节点
     start_args: object  # 节点的启动参数信息，用于做匹配性的校验，防止运行过程中出现问题。
     websocket: WebSocket = None  # 用于通信的 websocket 连接对象
+    run_status: _PD_Client_RunStatus = field(default_factory=_PD_Client_RunStatus)
 
     def __post_init__(self):
         if self.mode not in ["prefill", "decode"]: