refomat

Author: shihaobai

lightllm/server/core/objs/req.py

@@ -95,12 +95,14 @@ class Req(ctypes.Structure):
         # 请求回复累计概率和
         ("cumlogprob", ctypes.c_float),
     ]
-    
+
     def get_str(self):
-        return (f"request_id:{self.request_id}, input_len:{self.input_len},"
-               f"shm_cur_kv_len:{self.shm_cur_kv_len},"
-               f"shm_cur_output_len:{self.shm_cur_output_len},"
-               f"finish_status:{self.finish_status.is_finished()}")
+        return (
+            f"request_id:{self.request_id}, input_len:{self.input_len},"
+            f"shm_cur_kv_len:{self.shm_cur_kv_len},"
+            f"shm_cur_output_len:{self.shm_cur_output_len},"
+            f"finish_status:{self.finish_status.is_finished()}"
+        )
 
     def init(
         self,

lightllm/server/router/model_infer/model_rpc.py

@@ -32,9 +32,9 @@ class ModelRpcServer:
     def __init__(
         self,
         args,
-        tp_rank: int, 
-        local_tp_rank: int, 
-        world_size: int, 
+        tp_rank: int,
+        local_tp_rank: int,
+        world_size: int,
         local_world_size: int,
         rpc_event: multiprocessing.Event,
         rpc_finished_event: multiprocessing.Event,
@@ -286,7 +286,9 @@ def _init_env(
 
     g_router_lock.obj = router_lock
 
-    model_rpc_server = ModelRpcServer(args, tp_rank, local_tp_rank, world_size, local_world_size, rpc_event, rpc_finished_event, info_queue, mem_queue)
+    model_rpc_server = ModelRpcServer(
+        args, tp_rank, local_tp_rank, world_size, local_world_size, rpc_event, rpc_finished_event, info_queue, mem_queue
+    )
     success_event.set()
 
     model_rpc_server.loop_thread.join()
@@ -309,12 +311,34 @@ async def start_model_process(
 
     # 单卡时不使用 rpc
     if world_size == 1:
-        return ModelRpcServer(args, tp_rank, local_tp_rank, world_size, local_world_size, rpc_event, rpc_finished_event, info_queue, mem_queue)
+        return ModelRpcServer(
+            args,
+            tp_rank,
+            local_tp_rank,
+            world_size,
+            local_world_size,
+            rpc_event,
+            rpc_finished_event,
+            info_queue,
+            mem_queue,
+        )
 
     success_event = mp.Event()
     proc = mp.Process(
         target=_init_env,
-        args=(args, tp_rank, local_tp_rank, world_size, local_world_size, info_queue, mem_queue, router_lock, rpc_event, rpc_finished_event, success_event),
+        args=(
+            args,
+            tp_rank,
+            local_tp_rank,
+            world_size,
+            local_world_size,
+            info_queue,
+            mem_queue,
+            router_lock,
+            rpc_event,
+            rpc_finished_event,
+            success_event,
+        ),
     )
     proc.start()
     success_event.wait(timeout=40)

lightllm/server/router/req_queue/base_queue.py

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

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):
+    def generate_new_batch(self, current_batch: Batch, current_waiting_num: int = -1):
         # 如果当前已经被调度的请求数量超过了上限，直接不调度新的请求了。
         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):
+    def generate_new_batch(self, current_batch: Batch, current_waiting_num: int = -1):
         # 如果当前已经被调度的请求数量超过了上限，直接不调度新的请求了。
         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/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):
+    def generate_new_batch(self, current_batch: Batch, current_waiting_num: int = -1):
         # 如果当前已经被调度的请求数量超过了上限，直接不调度新的请求了。
         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):
+    def generate_new_batch(self, current_batch: Batch, current_waiting_num: int = -1):
         batches = [self.inner_queues[dp_index].generate_new_batch(current_batch) for dp_index in range(self.dp_size)]
         return self._merge_batch(batches)