add req pause for dp (ModelTC#822)

Co-authored-by: baishihao <baishihao@sensetime.com>
Co-authored-by: hiworldwzj <30762946+hiworldwzj@users.noreply.github.com>

Author: shihaobai

lightllm/server/router/batch.py

@@ -30,6 +30,16 @@ def get_batch_decode_need_tokens(self):
 
         return new_batch_decode_need_tokens
 
+    def get_req_list_for_dp(self, dp_index: int):
+        if self.dp_size_in_node == 1:
+            return self.reqs
+
+        req_list = []
+        for req in self.reqs:
+            if req.sample_params.suggested_dp_index == dp_index:
+                req_list.append(req)
+        return req_list
+
     def filter_out_finished_req(self, shm_req_manager: ShmReqManager):
         unfinished_req_ids = []
         for req in self.reqs:

lightllm/server/router/manager.py

@@ -244,17 +244,19 @@ async def loop_for_fwd(
                         estimated_peak_token_count = self.shared_token_load.get_estimated_peak_token_count(d_i)
                         logger.debug(
                             f"dp_i {d_i} current batch size: {len(self.running_batch.reqs)} \n"
-                            f"dp_i {d_i} paused req num: {self.req_queue.get_paused_req_num()} \n"
+                            f"dp_i {d_i} paused req num: {self.req_queue.get_paused_req_num(d_i)} \n"
                             f"dp_i {d_i} frozen token num: {frozen_token_num} \n"
                             f"dp_i {d_i} estimated_peak_token_count: {estimated_peak_token_count} \n"
                             f"dp_i {d_i} token used ratio: {token_ratio1} not contain prompt cache tree unrefed token\n"
                             f"dp_i {d_i} token used ratio: {token_ratio2} contain prompt cache tree unrefed token"
                         )
+                        self.metric_client.gauge_set(
+                            "lightllm_batch_pause_size", self.req_queue.get_paused_req_num(d_i)
+                        )
                 # pd decode mode need to update token_load more frequently
                 self.req_queue.update_token_load(self.running_batch, force_update=self.is_pd_decode_mode)
                 self.stats_tool.print_stats()
                 self.metric_client.gauge_set("lightllm_batch_current_size", len(self.running_batch.reqs))
-                self.metric_client.gauge_set("lightllm_batch_pause_size", self.req_queue.get_paused_req_num())
                 self.metric_client.gauge_set("lightllm_queue_size", self.req_queue.get_wait_req_num())
                 self.metric_client.gauge_set(
                     "lightllm_batch_current_max_tokens",
@@ -356,23 +358,22 @@ async def _step(self):
                     self.running_batch.merge(new_mini_batch)
                 return
 
-        # 正常 decode 阶段， 如果可以直接decode就直接decode，否则通过暂停策略暂停一些请求
-        # 释放一些管理的 token
-        if self._can_decode(self.running_batch):
-            self.stats_tool.count_output_tokens(self.running_batch)
-            await self._decode_batch(self.running_batch)
-            self._filter_runing_batch()
-            self.has_wait_tokens += 1
-            return
-        else:
-            # pause strategy
-            paused_reqs = select_paused_reqs(
-                self.running_batch, self.pause_strategy, self.req_queue, self.max_total_token_num
-            )
-            await self._pause_reqs(paused_reqs)
-            logger.debug(f"pasued req num: {self.req_queue.get_paused_req_num()}")
-            self.has_wait_tokens = 0
-            return
+        # Check if need pause some requests for decode.
+        for dp_index in range(self.dp_size_in_node):
+            while not self._can_decode(self.running_batch, dp_index=dp_index):
+                # pause strategy
+                paused_reqs = select_paused_reqs(
+                    self.running_batch, self.pause_strategy, self.req_queue, self.max_total_token_num, dp_index=dp_index
+                )
+                await self._pause_reqs(paused_reqs)
+                logger.debug(f"DP index {dp_index} pasues req num: {self.req_queue.get_paused_req_num(dp_index)}")
+                self.has_wait_tokens = 0
+
+        # Decode
+        self.stats_tool.count_output_tokens(self.running_batch)
+        await self._decode_batch(self.running_batch)
+        self._filter_runing_batch()
+        self.has_wait_tokens += 1
         return
 
     async def _prefill_batch(self, batch: Batch):
@@ -416,16 +417,12 @@ def _filter_runing_batch(self):
             self.running_batch = None
             return
 
-    def _can_decode(self, batch: Batch):
-        # p d 分离模式下，目前只能使用保守调度，保证请求放入进行decode的时候
-        # 显存token肯定是够用的。
-        # deepseekv2 dp 模式下,采用保守调度，也肯定够用
-        if self.is_pd_run_mode or self.dp_size_in_node > 1 or self.is_safe_schedule:
+    def _can_decode(self, batch: Batch, dp_index: int):
+        if self.is_pd_run_mode or self.is_safe_schedule:
             return True
-
-        # 下面的判定条件，只在 dp 为 1 的情况下启用
-        assert self.dp_size_in_node == 1
-        return batch.get_batch_decode_need_tokens()[0] + self.get_used_tokens(0) <= self.max_total_token_num
+        return (
+            batch.get_batch_decode_need_tokens()[dp_index] + self.get_used_tokens(dp_index) <= self.max_total_token_num
+        )
 
     def get_used_tokens(self, dp_index):
         if self.args.use_dynamic_prompt_cache:

lightllm/server/router/pause_strategy.py

@@ -3,6 +3,7 @@
 from typing import List, Tuple
 from .batch import Batch, Req
 from lightllm.server.router.req_queue.base_queue import BaseQueue
+from lightllm.server.router.req_queue.dp_base_queue import DpQueue
 
 
 class Strategy:
@@ -14,13 +15,16 @@ class Fcfs(Strategy):
     def __init__(self) -> None:
         super().__init__()
 
-    def ordering_reqs(self, batch: Batch):
-        reqs = [req for req in batch.reqs]
-        return sorted(reqs, key=lambda req: req.request_id, reverse=True)
+    def ordering_reqs(self, reqs: List):
+        return reqs[::-1]
 
 
-def select_paused_reqs(batch: Batch, strategy: Strategy, req_queue: BaseQueue, max_total_token_num):
-    reqs: List[Req] = strategy.ordering_reqs(batch)
+def select_paused_reqs(
+    batch: Batch, strategy: Strategy, req_queue: BaseQueue, max_total_token_num: int, dp_index: int
+) -> List[Req]:
+    if isinstance(req_queue, DpQueue):
+        req_queue = req_queue.get_dp_queue(dp_index)
+    reqs: List[Req] = strategy.ordering_reqs(batch.get_req_list_for_dp(dp_index))
 
     if len(reqs) == 0:
         return []

lightllm/server/router/req_queue/base_queue.py

@@ -38,7 +38,8 @@ def extend(self, req_group: List[Req]):
         self.waiting_req_list.extend(req_group)
         return
 
-    def get_paused_req_num(self):
+    def get_paused_req_num(self, fake_dp_index: int = 0):
+        assert fake_dp_index == 0
         return len(self.pause_req_dict)
 
     def get_wait_req_num(self):

lightllm/server/router/req_queue/dp_base_queue.py

@@ -19,10 +19,15 @@ def __init__(self, args, router, base_queue_class, dp_size_in_node) -> None:
         self.inner_queues: List[BaseQueue] = [
             base_queue_class(args, router, dp_index, dp_size_in_node) for dp_index in range(self.dp_size_in_node)
         ]
+
         return
 
-    def get_paused_req_num(self):
-        return sum(queue.get_paused_req_num() for queue in self.inner_queues)
+    def get_dp_queue(self, dp_index: int):
+        assert dp_index < self.dp_size_in_node, "dp index out of range"
+        return self.inner_queues[dp_index]
+
+    def get_paused_req_num(self, dp_index: int = 0):
+        return self.inner_queues[dp_index].get_paused_req_num()
 
     def get_wait_req_num(self):
         return sum(queue.get_wait_req_num() for queue in self.inner_queues)