fix abort

Author: shihaobai

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

@@ -116,15 +116,14 @@ def generate_new_batch(self, current_batch: Batch):
             if ok_insert:
                 can_run_list.extend(cur_group_reqs)
 
+        new_batch = None
         if len(can_run_list) != 0:
             new_batch = Batch(uuid.uuid4().int, can_run_list, dp_size_in_node=self.dp_size_in_node)
-            for req in abort_req_list:
-                self.router.shm_req_manager.put_back_req_obj(req)
 
-            self.waiting_req_list = self.waiting_req_list[len(can_run_list) + aborted_count :]
-            return new_batch
-        else:
-            return None
+        for req in abort_req_list:
+            self.router.shm_req_manager.put_back_req_obj(req)
+        self.waiting_req_list = self.waiting_req_list[len(can_run_list) + aborted_count :]
+        return new_batch
 
     def _add_to_group(self, cur_group_reqs, req: Req):
         if len(cur_group_reqs) == 0:

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

@@ -91,15 +91,13 @@ def generate_new_batch(self, current_batch: Batch):
                 can_run_list.append(req)
             else:
                 break
-
+        new_batch = None
         if len(can_run_list) != 0:
             new_batch = Batch(uuid.uuid4().int, can_run_list, dp_size_in_node=self.dp_size_in_node)
-            for req in abort_req_list:
-                self.router.shm_req_manager.put_back_req_obj(req)
-            self.waiting_req_list = self.waiting_req_list[len(can_run_list) + aborted_count :]
-            return new_batch
-        else:
-            return None
+        for req in abort_req_list:
+            self.router.shm_req_manager.put_back_req_obj(req)
+        self.waiting_req_list = self.waiting_req_list[len(can_run_list) + aborted_count :]
+        return new_batch
 
     def _calcu_batch_token_load_batch_not_none(self, current_batch: Batch):
         is_busy = self.is_busy()

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

@@ -48,15 +48,13 @@ def generate_new_batch(self, current_batch: Batch):
                 can_run_list.append(req)
             else:
                 break
-
+        new_batch = None
         if len(can_run_list) != 0:
             new_batch = Batch(uuid.uuid4().int, can_run_list, dp_size_in_node=self.dp_size_in_node)
-            for req in abort_req_list:
-                self.router.shm_req_manager.put_back_req_obj(req)
-            self.waiting_req_list = self.waiting_req_list[len(can_run_list) + aborted_count :]
-            return new_batch
-        else:
-            return None
+        for req in abort_req_list:
+            self.router.shm_req_manager.put_back_req_obj(req)
+        self.waiting_req_list = self.waiting_req_list[len(can_run_list) + aborted_count :]
+        return new_batch
 
     def _calcu_batch_token_load_batch_not_none(self, current_batch: Batch):
         is_busy = self.is_busy()