Merge branch 'overlap' of https://github.com/ModelTC/lightllm into overlap

Author: shihaobai

.pre-commit-config.yaml

@@ -11,4 +11,4 @@ repos:
     hooks:
       - id: flake8
         additional_dependencies: [flake8-typing-imports==1.9.0]
-        args: ['--config=.flake8', '--max-line-length=120', '--ignore=TYP001, E722, C901, E203, E266, E402, E302, E241, E902, E731, F403, E701, F405, F401, W292, W293, W503, W606']
+        args: ['--config=.flake8', '--max-line-length=120', '--ignore=TYP001, E722, C901, E203, E266, E402, E302, E241, E902, E731, F403, E701, F405, F401, W292, W293, W503, W606, E231']

lightllm/server/router/batch.py

@@ -67,12 +67,6 @@ def is_clear(self):
         return len(self.reqs) == 0
 
     def merge(self, mini_batch: "Batch"):
-        for _req in mini_batch.reqs:
-            self.reqs.append(_req)
-        self.id_to_reqs = {req.request_id: req for req in self.reqs}
-        return
-
-    def dp_merge(self, mini_batch: "Batch"):
         if mini_batch is None:
             return
 
@@ -81,6 +75,18 @@ def dp_merge(self, mini_batch: "Batch"):
         self.id_to_reqs = {req.request_id: req for req in self.reqs}
         return
 
+    @staticmethod
+    def merge_two_batch(batch1: "Batch", batch2: "Batch") -> "Batch":
+        if batch1 is None and batch2 is None:
+            return None
+
+        not_none_batch = batch1 if batch1 is not None else batch2
+
+        merge_batch = Batch(-1, [], not_none_batch.dp_size_in_node)
+        merge_batch.merge(batch1)
+        merge_batch.merge(batch2)
+        return merge_batch
+
     def __repr__(self):
         return f"batch_id={self.batch_id}, " f"reqs={self.reqs}, "
 

lightllm/server/router/manager.py

@@ -1,26 +1,20 @@
-import copy
 import time
-import uuid
 import uvloop
 import asyncio
 import torch
-import rpyc
 import pickle
-import threading
 import inspect
 
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
-import concurrent.futures
 import zmq
 import zmq.asyncio
 import torch.multiprocessing as mp
 import torch.distributed as dist
 import multiprocessing
 from typing import Dict, List, Optional
-from .batch import Batch
+from .batch import Batch, Req
 from .model_infer.model_rpc import start_model_process, ModelRpcClient
 from .req_queue import build_req_queue
-from lightllm.utils.infer_utils import calculate_time
 from lightllm.server.core.objs.io_objs import GroupReqIndexes
 from lightllm.server.core.objs import ShmReqManager, StartArgs
 from .dynamic_prompt.radix_cache import RadixCacheReadOnlyClient
@@ -108,8 +102,8 @@ def __init__(self, args: StartArgs, router_port, detokenization_port, metric_por
         g_router_lock.obj = self.router_lock
 
         # 调度和推理进行折叠使用的线程池
-        # self.overlap_thread_pool = concurrent.futures.ThreadPoolExecutor(max_workers=1)
-        # self.schedule_task = None
+        self.schedule_new_batch: Batch = None
+        self.schedule_event = asyncio.Event()
         return
 
     async def wait_to_model_ready(self):
@@ -140,8 +134,6 @@ async def wait_to_model_ready(self):
             self.model_rpc_servers.append(rpc_model)
 
         self.model_rpc_client = ModelRpcClient(
-            model_infer_servers=self.model_rpc_servers,
-            world_size=self.world_size,
             rpc_event=self.rpc_event,
             rpc_finished_event=self.rpc_finished_event,
         )
@@ -223,7 +215,6 @@ 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(
@@ -285,7 +276,7 @@ async def loop_for_fwd(
             if self.running_batch is None:
                 await asyncio.sleep(0.01)  # 10ms
 
-    def get_new_batch(self):
+    def generate_new_batch(self):
         limit_router_queue_length = None
         if self.is_multinode_tp:
             # 使用 all_reduce 获取最小值
@@ -294,25 +285,30 @@ def get_new_batch(self):
             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()
 
-        new_batch = self.req_queue.generate_new_batch(self.running_batch, limit_router_queue_length)
-        return new_batch
+        # 调度的时候需要考虑当前运行的batch，和调度了但是暂时还没有推理的部分请求。
+        new_batch = self.req_queue.generate_new_batch(
+            Batch.merge_two_batch(self.running_batch, self.schedule_new_batch), limit_router_queue_length
+        )
+        self.schedule_new_batch = Batch.merge_two_batch(self.schedule_new_batch, new_batch)
+        return
 
     async def _step(self):
         """
         事件处理循环
         """
-        # 删除所有已经 finished 的 req
-        # 当前无运行请求时
-        new_batch = None
-        if not self.batch_queue.empty():
-            new_batch = self.batch_queue.get_nowait()
-        if new_batch is not None:
+        # 判断是否有新请求加入推理
+        # 激进调度满足，有新的推理batch就需要进行加入。
+        # 或者延迟step的步数满足了当前条件，也需要进行新的推理batch的加入。
+        if (self.schedule_new_batch is not None) and (
+            (not self.args.disable_aggressive_schedule) or (self.has_wait_tokens >= self.max_wait_tokens)
+        ):
+            new_batch = self.schedule_new_batch
+            self.schedule_new_batch = None
+            self._add_new_batch_to_running_batch(new_batch=new_batch)
             await self._prefill_batch(new_batch)
-            self._filter_runing_batch()
-            if self.running_batch is None:
-                self.running_batch = new_batch
-            else:
-                self.running_batch.merge(new_batch)
+            self.stats_tool.count_prompt_tokens(new_batch)
+            self._filter_reqs_from_running_batch()
+            self.has_wait_tokens = 0
 
         # Check if need pause some requests for decode.
         for dp_index in range(self.dp_size_in_node):
@@ -327,41 +323,43 @@ async def _step(self):
 
         # Decode
         self.stats_tool.count_output_tokens(self.running_batch)
-        await self._decode_batch(self.running_batch)
-        if self.world_size // self.nnodes == 1:
-            # node_world_size == 1 时，协程不会让出来，导致无法调度新请求，所以sleep 1ms，可以修改一下
-            await asyncio.sleep(0.001)
-        self._filter_runing_batch()
+        await self._decode_batch()
+        self._filter_reqs_from_running_batch()
         self.has_wait_tokens += 1
         return
 
     async def _prefill_batch(self, batch: Batch):
+        # 添加新请求
         reqs = [r.to_router_rpc_obj() for r in batch.reqs]
         await self.model_rpc_client.prefill(reqs)
-        batch.filter_out_finished_req(self.shm_req_manager)
         self._send_detokenization_pack()
-
         logger.debug(f"Prefill Batch: {batch.simple_log()} \n")
         return
 
-    async def _decode_batch(self, batch: Batch):
+    async def _decode_batch(self):
+        self.schedule_event.set()
         await self.model_rpc_client.decode()
-        # 在 self.is_multinode_and_multidp 为 True 时，传入的 batch 对象可能为 None。
-        if batch is not None:
-            batch.filter_out_finished_req(self.shm_req_manager)
-
         self._send_detokenization_pack()
         return
 
-    async def _pause_reqs(self, pasue_reqs):
+    async def _pause_reqs(self, pasue_reqs: List[Req]):
         pasue_req_ids = [r.request_id for r in pasue_reqs]
         await self.model_rpc_client.pause_reqs(pasue_req_ids)
         return
 
-    def _filter_runing_batch(self):
-        if self.running_batch is not None and self.running_batch.is_clear():
-            self.running_batch = None
-            return
+    def _add_new_batch_to_running_batch(self, new_batch: Batch):
+        if self.running_batch is None:
+            self.running_batch = new_batch
+        else:
+            self.running_batch.merge(new_batch)
+        return
+
+    def _filter_reqs_from_running_batch(self):
+        if self.running_batch is not None:
+            self.running_batch.filter_out_finished_req(self.shm_req_manager)
+            if self.running_batch.is_clear():
+                self.running_batch = None
+        return
 
     def _can_decode(self, batch: Batch, dp_index: int):
         if self.is_pd_run_mode or self.is_safe_schedule or batch is None:
@@ -389,19 +387,35 @@ def get_used_tokens(self, dp_index):
             return self.max_total_token_num - self.read_only_statics_mem_manager.get_unrefed_token_num(dp_index)
 
     async def loop_for_netio_req(self):
+        recv_max_count = 66
+
         while True:
             try:
-                recv_req: GroupReqIndexes = self.recv_from_httpserver.recv_pyobj(zmq.NOBLOCK)
-                if isinstance(recv_req, GroupReqIndexes):
-                    self.add_req(recv_req)
-                else:
-                    assert False, f"Error Req Inf {recv_req}"
+                # 一次最多从 zmq 中取 recv_max_count 个请求，防止 zmq 队列中请求数量过多导致阻塞了主循环。
+                for _ in range(recv_max_count):
+                    recv_req: GroupReqIndexes = self.recv_from_httpserver.recv_pyobj(zmq.NOBLOCK)
+                    if isinstance(recv_req, GroupReqIndexes):
+                        self.add_req(recv_req)
+                    else:
+                        assert False, f"Error Req Inf {recv_req}"
+
+                # 当队列中存在较多的请求时，将一次接受的数量上调
+                recv_max_count = min(int(recv_max_count * 1.3), 300)
+
             except zmq.ZMQError:
-                new_batch = self.get_new_batch()
-                if new_batch is not None:
-                    self.batch_queue.put_nowait(new_batch)
-                await asyncio.sleep(0.005)
-                continue
+                # 当队列已经开始清空的时候，将一次接受的数量下调
+                recv_max_count = 66
+
+            try:
+                await asyncio.wait_for(self.schedule_event.wait(), timeout=0.02)
+            except asyncio.TimeoutError:
+                pass
+
+            if self.schedule_event.is_set():
+                self.generate_new_batch()
+                self.schedule_event.clear()
+
+        return
 
     def clean_up(self):
         return
@@ -440,7 +454,6 @@ def handle_exception(loop, context):
     loop = asyncio.new_event_loop()
     loop.set_exception_handler(handle_exception)
     asyncio.set_event_loop(loop)
-    router.batch_queue = asyncio.Queue()
 
     loop.create_task(router.loop_for_fwd())
     loop.run_until_complete(router.loop_for_netio_req())