fix

Author: hiworldwzj

lightllm/server/router/batch.py

@@ -54,10 +54,10 @@ def filter_out_finished_req(self, shm_req_manager: ShmReqManager):
         self.id_to_reqs = {req.request_id: req for req in self.reqs}
         return
 
-    def pop_req(self, req_id):
+    def pop_req(self, req_id) -> Req:
         self.reqs = [req for req in self.reqs if req.request_id != req_id]
-        self.id_to_reqs.pop(req_id)
-        return
+        req = self.id_to_reqs.pop(req_id)
+        return req
 
     def is_clear(self):
         return len(self.reqs) == 0

lightllm/server/router/manager.py

@@ -45,6 +45,8 @@ def __init__(self, args: StartArgs, router_port, detokenization_port, metric_por
         # 兼容多机纯tp的运行模式，这时候 1 // 2 == 0, 需要兼容
         self.dp_size_in_node = max(1, args.dp // self.nnodes)
         self.is_multinode_tp = args.nnodes > 1 and args.dp == 1
+        self.is_multinode_tp_master = self.is_multinode_tp and args.node_rank == 0
+        self.is_multinode_tp_slave = self.is_multinode_tp and args.node_rank != 0
         self.is_multinode_and_multidp = args.nnodes > 1 and args.dp > 1
         # 判断是否是保守调度，保守调度不会发生暂停 req 的情况，但是有些场景可能影响吞吐
         self.is_safe_schedule = args.router_token_ratio == 0.0
@@ -359,21 +361,73 @@ def _add_req(self, group_req_indexes: GroupReqIndexes):
         return
 
     def _generate_new_batch(self):
-        limit_router_queue_length = None
-        if self.is_multinode_tp:
-            # 使用 all_reduce 获取最小值
-            limit_router_queue_length = len(self.req_queue.waiting_req_list)
-            limit_router_queue_length_tensor = torch.tensor(limit_router_queue_length, dtype=torch.int32, device="cpu")
-            dist.all_reduce(limit_router_queue_length_tensor, op=dist.ReduceOp.MIN, group=self.mulitnode_group)
-            limit_router_queue_length = limit_router_queue_length_tensor.item()
-
         # 调度的时候需要考虑当前运行的batch，和调度了但是暂时还没有推理的部分请求。
         new_batch = self.req_queue.generate_new_batch(
-            Batch.merge_two_batch(self.running_batch, self.schedule_new_batch), limit_router_queue_length
+            Batch.merge_two_batch(self.running_batch, self.schedule_new_batch)
         )
         self.schedule_new_batch = Batch.merge_two_batch(self.schedule_new_batch, new_batch)
         return
 
+    def _multinode_tp_generate_new_batch(self):
+        dist.barrier(group=self.mulitnode_group)
+
+        # 调度的时候需要考虑当前运行的batch，和调度了但是暂时还没有推理的部分请求。
+        if self.is_multinode_tp_master:
+            new_batch = self.req_queue.generate_new_batch(
+                Batch.merge_two_batch(self.running_batch, self.schedule_new_batch)
+            )
+            if new_batch is not None:
+                req_ids = [req.request_id for req in new_batch.reqs]
+            else:
+                req_ids = []
+            dist.broadcast_object_list([len(req_ids)], src=0, group=self.mulitnode_group)
+            dist.broadcast_object_list(req_ids, src=0, group=self.mulitnode_group)
+            req_id_select_mark = [1 for _ in range(len(req_ids))]
+            req_id_select_mark = torch.tensor(req_id_select_mark, dtype=torch.int32, device="cpu")
+            dist.all_reduce(req_id_select_mark, op=dist.ReduceOp.MIN, group=self.mulitnode_group)
+            back_req_list = []
+            for req_id, select in zip(req_ids, req_id_select_mark.numpy()):
+                if select == 0:
+                    req = new_batch.pop_req(req_id)
+                    back_req_list.append(req)
+            self.req_queue.waiting_req_list = back_req_list + self.req_queue.waiting_req_list
+            if new_batch.is_clear():
+                new_batch = None
+        else:
+            req_nums = [None]
+            dist.broadcast_object_list(req_nums, src=0, group=self.mulitnode_group)
+            req_num = req_nums[0]
+            req_ids = [None for _ in range(req_num)]
+            dist.broadcast_object_list(req_ids, src=0, group=self.mulitnode_group)
+            all_req_id_set = set([req.request_id for req in self.req_queue.waiting_req_list])
+            req_id_select_mark = []
+            for req_id in req_ids:
+                req_id_select_mark.append(1 if req_id in all_req_id_set else 0)
+            req_id_select_mark = torch.tensor(req_id_select_mark, dtype=torch.int32, device="cpu")
+            dist.all_reduce(req_id_select_mark, op=dist.ReduceOp.MIN, group=self.mulitnode_group)
+            select_req_ids = []
+            for req_id, select in zip(req_ids, req_id_select_mark.numpy()):
+                if select == 1:
+                    select_req_ids.append(req_id)
+
+            select_reqs = []
+            for req_id in select_req_ids:
+                for req in self.req_queue.waiting_req_list:
+                    if req.request_id == req_id:
+                        select_reqs.append(req)
+
+            for req in select_reqs:
+                self.req_queue.waiting_req_list.remove(req)
+            if select_reqs:
+                new_batch = Batch(-1, reqs=select_reqs, dp_size_in_node=self.dp_size_in_node)
+            else:
+                new_batch = None
+
+        self.schedule_new_batch = Batch.merge_two_batch(self.schedule_new_batch, new_batch)
+
+        dist.barrier(group=self.mulitnode_group)
+        return
+
     async def _recv_new_reqs_and_schedule(self):
         if not hasattr(self, "recv_max_count"):
             self.recv_max_count = 64
@@ -394,9 +448,11 @@ async def _recv_new_reqs_and_schedule(self):
             # 当队列已经开始清空的时候，将一次接受的数量下调
             self.recv_max_count = 64
 
-        # 只有当推理侧没有发生暂停的时候，才执行新的调度
-        if self._get_paused_req_num() == 0:
-            self._generate_new_batch()
+        if self.is_multinode_tp:
+            self._multinode_tp_generate_new_batch()
+        else:
+            if self._get_paused_req_num() == 0:
+                self._generate_new_batch()
         return
 
     def clean_up(self):

lightllm/server/router/req_queue/base_queue.py

@@ -58,12 +58,10 @@ def get_batch_dp_req_size(self, current_batch: Batch):
 
         return len([req for req in current_batch.reqs if req.sample_params.suggested_dp_index == self.dp_index])
 
-    def generate_new_batch(self, current_batch: Batch, limit_router_queue_length: int = None):
+    def generate_new_batch(self, current_batch: Batch):
         """
         args:
             current_batch: current batch
-            limit_router_queue_length: the least length of waiting list across all nodes.
-            It only works when nnodes > 1 and dp_size == 1.
         return:
             new batch
         """

lightllm/server/router/req_queue/chunked_prefill/beam_impl.py

@@ -69,12 +69,12 @@ def _can_add_new_group_reqs(self, cur_handle_group_reqs: List[Req], is_busy, new
             return False, new_batch_first_router_need_tokens
 
     # @calculate_time(show=True, min_cost_ms=10)
-    def generate_new_batch(self, current_batch: Batch, limit_router_queue_length: int = None):
+    def generate_new_batch(self, current_batch: Batch):
         if len(self.waiting_req_list) == 0:
             return None
 
         # 如果当前已经被调度的请求数量超过了上限，直接不调度新的请求了。
-        exist_req_num = self.get_batch_dp_req_size(current_batch) + len(self.pause_req_dict)
+        exist_req_num = self.get_batch_dp_req_size(current_batch)
         req_is_full = exist_req_num >= self.running_max_req_size
         if req_is_full:
             return None

lightllm/server/router/req_queue/chunked_prefill/impl.py

@@ -54,7 +54,7 @@ def _can_add_new_req(self, req: Req, is_busy, new_batch_first_router_need_tokens
                 return False, new_batch_first_router_need_tokens
 
     # @calculate_time(show=True, min_cost_ms=10)
-    def generate_new_batch(self, current_batch: Batch, limit_router_queue_length: int = None):
+    def generate_new_batch(self, current_batch: Batch):
         if len(self.waiting_req_list) == 0:
             return None
 
@@ -75,10 +75,7 @@ def generate_new_batch(self, current_batch: Batch, limit_router_queue_length: in
         abort_req_list = []
         aborted_count = 0
 
-        if limit_router_queue_length is None:
-            waiting_queue = self.waiting_req_list
-        else:
-            waiting_queue = self.waiting_req_list[:limit_router_queue_length]
+        waiting_queue = self.waiting_req_list
 
         for req in waiting_queue:
             if req.is_aborted:

lightllm/server/router/req_queue/chunked_prefill/impl_for_pd_decode.py

@@ -24,7 +24,7 @@ def _init_cache_list(self, current_batch: Batch, is_busy):
         return
 
     # @calculate_time(show=True, min_cost_ms=10)
-    def generate_new_batch(self, current_batch: Batch, limit_router_queue_length: int = None):
+    def generate_new_batch(self, current_batch: Batch):
         if len(self.waiting_req_list) == 0:
             return None
 

lightllm/server/router/req_queue/dp_base_queue.py

@@ -30,7 +30,7 @@ def get_wait_req_num(self):
         return sum(queue.get_wait_req_num() for queue in self.inner_queues)
 
     # @calculate_time(show=True, min_cost_ms=10)
-    def generate_new_batch(self, current_batch: Batch, limit_router_queue_length: int = None):
+    def generate_new_batch(self, current_batch: Batch):
         batches = [
             self.inner_queues[dp_index].generate_new_batch(current_batch) for dp_index in range(self.dp_size_in_node)
         ]