dp balancer abstract

Author: shihaobai

lightllm/server/api_cli.py

@@ -112,7 +112,7 @@ def make_argument_parser() -> argparse.ArgumentParser:
         help="tool call parser type",
     )
     parser.add_argument(
-        "--running_max_req_size", type=int, default=1000, help="the max size for forward requests in the same time"
+        "--running_max_req_size", type=int, default=2048, help="the max size for forward requests in the same time"
     )
     parser.add_argument("--nnodes", type=int, default=1, help="the number of nodes")
     parser.add_argument("--node_rank", type=int, default=0, help="the rank of the current node")
@@ -137,6 +137,12 @@ def make_argument_parser() -> argparse.ArgumentParser:
                         using the deepseekv2 model, set dp to be equal to the tp parameter. In other cases, please
                         do not set it and keep the default value as 1.""",
     )
+    parser.add_argument(
+        "--dp_balancer",
+        type=str,
+        default="round_robin",
+        help="the dp balancer type, default is round_robin",
+    )
     parser.add_argument(
         "--max_req_total_len", type=int, default=16384, help="the max value for req_input_len + req_output_len"
     )

lightllm/server/router/req_queue/dp_base_queue.py

@@ -2,6 +2,7 @@
 from typing import List
 from ..batch import Batch, Req
 from lightllm.server.router.req_queue.base_queue import BaseQueue
+from lightllm.server.router.req_queue.dp_balancer import get_dp_balancer
 from lightllm.common.basemodel.infer_lock import g_router_lock
 from lightllm.utils.log_utils import init_logger
 
@@ -12,14 +13,13 @@ class DpQueue:
     def __init__(self, args, router, base_queue_class, dp_size_in_node) -> None:
         self.dp_size_in_node = dp_size_in_node
         self.base_queue_class = base_queue_class
-        self.pre_select_dp_index = self.dp_size_in_node - 1
         from lightllm.server.router.manager import RouterManager
 
         self.router: RouterManager = router
         self.inner_queues: List[BaseQueue] = [
             base_queue_class(args, router, dp_index, dp_size_in_node) for dp_index in range(self.dp_size_in_node)
         ]
-
+        self.dp_balancer = get_dp_balancer(args, dp_size_in_node, self.inner_queues)
         return
 
     def get_dp_queue(self, dp_index: int):
@@ -49,8 +49,7 @@ def append(self, req: Req):
         suggested_dp_index = req.sample_params.suggested_dp_index
         if suggested_dp_index >= self.dp_size_in_node or suggested_dp_index < 0:
             logger.warning(f"input req {req.request_id} dp index {suggested_dp_index} is invalid")
-            suggested_dp_index = self._get_suggest_dp_index()
-            self.pre_select_dp_index = suggested_dp_index
+            suggested_dp_index = self.dp_balancer.get_suggest_dp_index()
             req.sample_params.suggested_dp_index = suggested_dp_index
             self.inner_queues[suggested_dp_index].append(req)
         else:
@@ -59,12 +58,11 @@ def append(self, req: Req):
 
     def extend(self, req_group: List[Req]):
         # 同一个组的，要分配在同一个 dp 上，效率最高
-        index = self._get_suggest_dp_index()
+        index = self.dp_balancer.get_suggest_dp_index()
         for req in req_group:
             suggested_dp_index = req.sample_params.suggested_dp_index
             if suggested_dp_index >= self.dp_size_in_node or suggested_dp_index < 0:
                 logger.warning(f"input req {req.request_id} dp index {suggested_dp_index} is invalid")
-                self.pre_select_dp_index = index
                 req.sample_params.suggested_dp_index = index
                 self.inner_queues[index].append(req)
             else:
@@ -87,21 +85,3 @@ def update_token_load(self, current_batch: Batch, force_update=False):
                     self.router.shared_token_load.set_estimated_peak_token_count(estimated_peak_token_count, dp_index)
                     self.router.shared_token_load.set_dynamic_max_load(dynamic_max_load, dp_index)
         return
-
-    def _get_suggest_dp_index(self):
-        min_length = min(len(queue.waiting_req_list) for queue in self.inner_queues)
-        select_dp_indexes = [
-            i for i, queue in enumerate(self.inner_queues) if len(queue.waiting_req_list) == min_length
-        ]
-
-        # multi thread safe keep
-        if not select_dp_indexes:
-            return random.randint(0, self.dp_size_in_node - 1)
-
-        # round_robin select.
-        for i in range(self.dp_size_in_node):
-            next_dp_index = (self.pre_select_dp_index + i + 1) % self.dp_size_in_node
-            if next_dp_index in select_dp_indexes:
-                return next_dp_index
-
-        return random.choice(select_dp_indexes)