fix

Author: hiworldwzj

lightllm/server/router/manager.py

@@ -200,7 +200,7 @@ async def wait_to_model_ready(self):
 
         return
 
-    async def add_req(self, group_req_indexes: GroupReqIndexes):
+    def add_req(self, group_req_indexes: GroupReqIndexes):
         req_group = []
         for req_index in group_req_indexes.shm_req_indexes:
             req = self.shm_req_manager.get_req_obj_by_index(req_index)
@@ -211,6 +211,7 @@ async def add_req(self, group_req_indexes: GroupReqIndexes):
             logger.info(f"router recive req id {req.request_id} cost time {time.time() - req.start_time} s")
         self.req_queue.extend(req_group)
         self.send_to_detokenization.send_pyobj(group_req_indexes, protocol=pickle.HIGHEST_PROTOCOL)
+
         return
 
     async def loop_for_fwd(
@@ -262,18 +263,18 @@ async def get_schedule_result(self, running_batch: Batch):
         if self.schedule_task is None:
 
             def get_new_batch():
-                current_waiting_num = None
+                limit_router_queue_length = None
                 if self.nnodes > 1 and self.args.dp == 1:
                     # 使用 all_reduce 获取最小值
-                    current_waiting_num = len(self.req_queue.waiting_req_list)
-                    current_waiting_num_tensor = torch.tensor(current_waiting_num, dtype=torch.int32, device="cpu")
-                    dist.all_reduce(current_waiting_num_tensor, op=dist.ReduceOp.MIN, group=self.mulitnode_group)
-                    current_waiting_num = current_waiting_num_tensor.item()
+                    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()
 
                 self.overlap_event.wait(timeout=0.020)
                 self.overlap_event.clear()
                 time.sleep(0.003)
-                new_batch = self.req_queue.generate_new_batch(running_batch, current_waiting_num)
+                new_batch = self.req_queue.generate_new_batch(running_batch, limit_router_queue_length)
                 return new_batch
 
             self.schedule_task = asyncio.get_running_loop().run_in_executor(self.overlap_thread_pool, get_new_batch)
@@ -399,7 +400,7 @@ async def loop_for_netio_req(self):
         while True:
             recv_req: GroupReqIndexes = await self.recv_from_httpserver.recv_pyobj()
             if isinstance(recv_req, GroupReqIndexes):
-                await self.add_req(recv_req)
+                self.add_req(recv_req)
             else:
                 assert False, f"Error Req Inf {recv_req}"
 
@@ -408,7 +409,6 @@ def clean_up(self):
 
 
 def start_router_process(args, router_port, detokenization_port, model_rpc_ports, metric_port, pipe_writer):
-
     # 注册 graceful 退出的处理
     graceful_registry(inspect.currentframe().f_code.co_name)
     start_parent_check_thread()

lightllm/server/router/req_queue/base_queue.py

@@ -69,11 +69,11 @@ 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, current_waiting_num: int = None):
+    def generate_new_batch(self, current_batch: Batch, limit_router_queue_length: int = None):
         """
         args:
             current_batch: current batch
-            current_waiting_num: the least length of waiting list across all nodes.
+            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/impl.py

@@ -56,7 +56,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, current_waiting_num: int = None):
+    def generate_new_batch(self, current_batch: Batch, limit_router_queue_length: int = None):
 
         # 如果当前已经被调度的请求数量超过了上限，直接不调度新的请求了。
         exist_req_num = self.get_batch_dp_req_size(current_batch) + len(self.pause_req_dict)
@@ -74,7 +74,13 @@ def generate_new_batch(self, current_batch: Batch, current_waiting_num: int = No
         can_run_list = []
         abort_req_list = []
         aborted_count = 0
-        for req in self.waiting_req_list[:current_waiting_num]:
+        
+        if limit_router_queue_length is None:
+            waiting_queue = self.waiting_req_list
+        else:
+            waiting_queue = self.waiting_req_list[:limit_router_queue_length]
+        
+        for req in waiting_queue:
             if req.is_aborted and not req.is_paused:
                 # 由于管理的复杂性，只有没有被调度运行过的请求可以因为abort直接在队列中忽略掉.
                 # 暂停的请求需要恢复后，由 router manager 部分来过滤。暂时保持这种处理方法, 否则会导致管理token的泄漏

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

@@ -76,7 +76,7 @@ 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, current_waiting_num: int = None):
+    def generate_new_batch(self, current_batch: Batch, limit_router_queue_length: int = None):
         # 如果当前已经被调度的请求数量超过了上限，直接不调度新的请求了。
         exist_req_num = self.get_batch_dp_req_size(current_batch) + len(self.pause_req_dict)
         req_is_full = exist_req_num >= self.running_max_req_size

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

@@ -61,7 +61,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, current_waiting_num: int = None):
+    def generate_new_batch(self, current_batch: Batch, limit_router_queue_length: int = None):
         # 如果当前已经被调度的请求数量超过了上限，直接不调度新的请求了。
         exist_req_num = self.get_batch_dp_req_size(current_batch) + len(self.pause_req_dict)
         req_is_full = exist_req_num >= self.running_max_req_size
@@ -76,7 +76,13 @@ def generate_new_batch(self, current_batch: Batch, current_waiting_num: int = No
         abort_req_list = []
         new_batch_first_router_need_tokens = 0  # 主要是对 prefill 大块计算时候的token数量限制
         aborted_count = 0
-        for req in self.waiting_req_list[:current_waiting_num]:
+
+        if limit_router_queue_length is None:
+            waiting_queue = self.waiting_req_list
+        else:
+            waiting_queue = self.waiting_req_list[:limit_router_queue_length]
+
+        for req in waiting_queue:
             if req.is_aborted and not req.is_paused:
                 # 由于管理的复杂性，只有没有被调度运行过的请求可以因为abort直接在队列中忽略掉.
                 # 暂停的请求需要恢复后，由 router manager 部分来过滤。暂时保持这种处理方法, 否则会导致管理token和管理req对象的泄漏

lightllm/server/router/req_queue/continues_batch/pd_decode_impl.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, current_waiting_num: int = None):
+    def generate_new_batch(self, current_batch: Batch, limit_router_queue_length: int = None):
         # 如果当前已经被调度的请求数量超过了上限，直接不调度新的请求了。
         exist_req_num = self.get_batch_dp_req_size(current_batch) + len(self.pause_req_dict)
         req_is_full = exist_req_num >= self.running_max_req_size

lightllm/server/router/req_queue/dp_base_queue.py

@@ -27,7 +27,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, current_waiting_num: int = None):
+    def generate_new_batch(self, current_batch: Batch, limit_router_queue_length: int = None):
         batches = [self.inner_queues[dp_index].generate_new_batch(current_batch) for dp_index in range(self.dp_size)]
         return self._merge_batch(batches)