Fix the pause issue under extremely aggressive scheduling. (#1057)

Author: shihaobai

lightllm/server/router/model_infer/infer_batch.py

@@ -199,16 +199,20 @@ def pause_reqs(self, pause_reqs: List["InferReq"], is_master_in_dp: bool):
             g_infer_state_lock.release()
         return self
 
-    def recover_paused_reqs(self, paused_reqs: List["InferReq"], is_master_in_dp: bool):
+    def recover_paused_reqs(self, paused_reqs: List["InferReq"], is_master_in_dp: bool, can_alloc_token_num: int):
         if paused_reqs:
             g_infer_state_lock.acquire()
 
             for req in paused_reqs:
+                prefill_need_token_num = req.get_cur_total_len()
+                if prefill_need_token_num > can_alloc_token_num:
+                    break
                 req._match_radix_cache()
                 assert req.paused is True
                 req.paused = False
                 if is_master_in_dp:
                     req.shm_req.is_paused = False
+                can_alloc_token_num -= prefill_need_token_num
 
             g_infer_state_lock.release()
         return

lightllm/server/router/model_infer/mode_backend/base_backend.py

@@ -464,7 +464,9 @@ def _get_classed_reqs(
         g_infer_context.pause_reqs(wait_pause_reqs, is_master_in_dp=self.is_master_in_dp)
 
         if recover_paused:
-            g_infer_context.recover_paused_reqs(paused_reqs=paused_reqs, is_master_in_dp=self.is_master_in_dp)
+            g_infer_context.recover_paused_reqs(
+                paused_reqs=paused_reqs, is_master_in_dp=self.is_master_in_dp, can_alloc_token_num=can_alloc_token_num
+            )
 
         return prefill_reqs, decode_reqs