fix lint.

Author: Weichao Luo

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

@@ -261,7 +261,7 @@ def _post_handle(
         is_chuncked_mode: bool,
         do_filter_finished_reqs: bool,
         extra_post_req_handle_func: Optional[Callable[[InferReq, int, float], None]] = None,
-        call_post_handle_for_chunk: bool = False ,
+        call_post_handle_for_chunk: bool = False,
     ) -> List[int]:
         """
         extra_post_req_handle_func 用于提供在一个请求确定输出的时候，给出额外的后处理操作，主要是用于

lightllm/server/router/model_infer/mode_backend/continues_batch/impl.py

@@ -46,7 +46,7 @@ def normal_prefill_reqs(
         ok_finished_reqs: List[InferReq],
         mask_func: Optional[Callable[[List[InferReq], torch.Tensor], None]] = None,
         extra_post_req_handle_func: Optional[Callable[[InferReq, int, float], None]] = None,
-        call_post_handle_for_chunk: bool = False
+        call_post_handle_for_chunk: bool = False,
     ):
         model_input, run_reqs = prepare_prefill_inputs(
             prefill_reqs, is_chuncked_mode=not self.disable_chunked_prefill, is_multimodal=self.is_multimodal

lightllm/server/router/model_infer/mode_backend/dp_backend/impl.py

@@ -52,9 +52,15 @@ def decode(self):
         self._overlap_req_init_and_filter(uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True)
         return
 
-    def normal_prefill_reqs(self, prefill_reqs: List[InferReq], max_prefill_num: int, uninit_reqs, ok_finished_reqs,
-                            extra_post_req_handle_func: Optional[Callable[[InferReq, int, float], None]] = None,
-                            call_post_handle_for_chunk: bool = False):
+    def normal_prefill_reqs(
+        self,
+        prefill_reqs: List[InferReq],
+        max_prefill_num: int,
+        uninit_reqs,
+        ok_finished_reqs,
+        extra_post_req_handle_func: Optional[Callable[[InferReq, int, float], None]] = None,
+        call_post_handle_for_chunk: bool = False,
+    ):
         model_input, run_reqs, padded_req_num = padded_prepare_prefill_inputs(
             prefill_reqs, is_multimodal=self.is_multimodal
         )
@@ -67,9 +73,13 @@ def normal_prefill_reqs(self, prefill_reqs: List[InferReq], max_prefill_num: int
             next_token_ids = next_token_ids.detach().cpu().numpy()
             next_token_logprobs = torch.log(next_token_probs).detach().cpu().numpy()
             self._post_handle(
-                run_reqs, next_token_ids, next_token_logprobs, is_chuncked_mode=True, do_filter_finished_reqs=False,
+                run_reqs,
+                next_token_ids,
+                next_token_logprobs,
+                is_chuncked_mode=True,
+                do_filter_finished_reqs=False,
                 extra_post_req_handle_func=extra_post_req_handle_func,
-                call_post_handle_for_chunk=call_post_handle_for_chunk
+                call_post_handle_for_chunk=call_post_handle_for_chunk,
             )
         return
 
@@ -121,9 +131,15 @@ def overlap_decode(self, decode_reqs: List[InferReq], max_decode_num: int, unini
             )
         return
 
-    def overlap_prefill_reqs(self, prefill_reqs: List[InferReq], max_prefill_num: int, uninit_reqs, ok_finished_reqs,
-                            extra_post_req_handle_func: Optional[Callable[[InferReq, int, float], None]] = None,
-                            call_post_handle_for_chunk: bool = False):
+    def overlap_prefill_reqs(
+        self,
+        prefill_reqs: List[InferReq],
+        max_prefill_num: int,
+        uninit_reqs,
+        ok_finished_reqs,
+        extra_post_req_handle_func: Optional[Callable[[InferReq, int, float], None]] = None,
+        call_post_handle_for_chunk: bool = False,
+    ):
         (
             micro_input,
             run_reqs,
@@ -148,8 +164,12 @@ def overlap_prefill_reqs(self, prefill_reqs: List[InferReq], max_prefill_num: in
             next_token_ids = next_token_ids.detach().cpu().numpy()
             next_token_logprobs = torch.log(next_token_probs).detach().cpu().numpy()
             self._post_handle(
-                all_run_reqs, next_token_ids, next_token_logprobs, is_chuncked_mode=True, do_filter_finished_reqs=False,
+                all_run_reqs,
+                next_token_ids,
+                next_token_logprobs,
+                is_chuncked_mode=True,
+                do_filter_finished_reqs=False,
                 extra_post_req_handle_func=extra_post_req_handle_func,
-                call_post_handle_for_chunk=call_post_handle_for_chunk
+                call_post_handle_for_chunk=call_post_handle_for_chunk,
             )
         return

lightllm/server/router/model_infer/mode_backend/pd_nixl/impl_for_pd_base.py

@@ -62,16 +62,20 @@ def init_custom(self):
         self.page_scheduer = SafePageIndexScheduler(self.nixl_agent.num_pages)
 
         self.nixl_meta_queue.put(
-            (self.nixl_agent.agent_metadata, self.nixl_agent.num_tokens, self.nixl_agent.num_pages,
-             self.nixl_agent.local_mem_desc, self.nixl_agent.local_page_mem_desc)
+            (
+                self.nixl_agent.agent_metadata,
+                self.nixl_agent.num_tokens,
+                self.nixl_agent.num_pages,
+                self.nixl_agent.local_mem_desc,
+                self.nixl_agent.local_page_mem_desc,
+            )
         )
 
     def _start_async_loop(self, async_loop_func):
         loop = asyncio.new_event_loop()
         asyncio.set_event_loop(loop)
         loop.run_until_complete(async_loop_func())
 
-
     async def _handle_remote_prefill(self, req_status: RemotePrefillStatus):
         group_req_id = req_status.group_req_id
         status = req_status.status
@@ -80,29 +84,36 @@ async def _handle_remote_prefill(self, req_status: RemotePrefillStatus):
 
         ret = None
         if run_req := self.remote_prefilled_reqs.get(group_req_id, None):
-            if req_status.transfer_type == RemoteTransferType.PAGE_TRANSFER and status == RemoteTransferStatusType.SUCCESS:
+            if (
+                req_status.transfer_type == RemoteTransferType.PAGE_TRANSFER
+                and status == RemoteTransferStatusType.SUCCESS
+            ):
                 kv_start, kv_len = req_status.kv_start, req_status.kv_len
-                token_ids = g_infer_context.req_manager.req_to_token_indexs[run_req.req_idx][kv_start: kv_start + kv_len] # gpu tensor
-                self.model.mem_manager.kv_buffer[:, token_ids, :, :] = self.model.mem_manager.kv_move_buffer[req_status.page_id][:kv_len].transpose(0, 1)
+                token_ids = g_infer_context.req_manager.req_to_token_indexs[run_req.req_idx][
+                    kv_start : kv_start + kv_len
+                ]  # gpu tensor
+                self.model.mem_manager.kv_buffer[:, token_ids, :, :] = self.model.mem_manager.kv_move_buffer[
+                    req_status.page_id
+                ][:kv_len].transpose(0, 1)
                 ret = PageTransferAck(group_req_id=group_req_id, page_id=req_status.page_id)
 
             if req_status.is_last or status != RemoteTransferStatusType.SUCCESS:
-                    shm_req: PDNIXLChunkedPrefillReq = run_req.shm_req
-                    shm_req.set_pd_req_rank_state(self.rank_in_dp, status.value)
-                    self.remote_prefilled_reqs.pop(group_req_id)
-                    self.request_to_first_token[group_req_id] = (req_status.next_token_id, req_status.next_token_logprob)
+                shm_req: PDNIXLChunkedPrefillReq = run_req.shm_req
+                shm_req.set_pd_req_rank_state(self.rank_in_dp, status.value)
+                self.remote_prefilled_reqs.pop(group_req_id)
+                self.request_to_first_token[group_req_id] = (req_status.next_token_id, req_status.next_token_logprob)
 
-                    if self.is_master_in_dp:
-                        # return page ids
-                        if group_req_id in self.request_to_page_ids:
-                            self.page_scheduer.return_(self.request_to_page_ids[group_req_id])
-                            del self.request_to_page_ids[group_req_id]
-
-                        logger.info(
-                            f"remote prefill reqeust: {group_req_id} done with status: {status} "
-                            f"took: {time.time() - run_req.remote_prefill_start} seconds"
-                        )
-                    ret = None
+                if self.is_master_in_dp:
+                    # return page ids
+                    if group_req_id in self.request_to_page_ids:
+                        self.page_scheduer.return_(self.request_to_page_ids[group_req_id])
+                        del self.request_to_page_ids[group_req_id]
+
+                    logger.info(
+                        f"remote prefill reqeust: {group_req_id} done with status: {status} "
+                        f"took: {time.time() - run_req.remote_prefill_start} seconds"
+                    )
+                ret = None
 
         else:
             if self.is_master_in_dp:
@@ -112,7 +123,7 @@ async def _handle_remote_prefill(self, req_status: RemotePrefillStatus):
 
     async def _prefill_wait_loop_async(self):
         while True:
-             # from local
+            # from local
             try:
                 req_status = self.from_remote_queue.get_nowait()
                 await self._handle_remote_prefill(req_status)
@@ -141,7 +152,7 @@ async def _prefill_wait_loop_async(self):
 
             await asyncio.sleep(PDNIXLBackendBase._THREAD_WAIT_INTERVAL)
 
-    def _handle_chunked_transfer(self, req: InferReq, next_token_id: int=None, next_token_logprob: float=None):
+    def _handle_chunked_transfer(self, req: InferReq, next_token_id: int = None, next_token_logprob: float = None):
         if next_token_id:
             next_token_id = int(next_token_id)
             next_token_logprob = float(next_token_logprob)
@@ -164,7 +175,7 @@ def _handle_chunked_transfer(self, req: InferReq, next_token_id: int=None, next_
                 free_page_ids=remote_request.data.page_ids.copy(),
                 next_token_id=next_token_id,
                 next_token_logprob=next_token_logprob,
-                lock=threading.Lock()
+                lock=threading.Lock(),
             )
             shm_req.set_pd_req_rank_state(self.rank_in_dp, RemoteTransferStatusType.IN_PROGRESS.value)
             req.in_prefill_or_transfer = True
@@ -179,7 +190,6 @@ def _handle_chunked_transfer(self, req: InferReq, next_token_id: int=None, next_
                     transfer_state.next_token_id = next_token_id
                     transfer_state.next_token_logprob = next_token_logprob
 
-
     async def _transfer_kv_to_remote_paged_batch(self, transfer_reqs: List[KVMoveRequest]):
         start = time.time()
         requests_by_agents = dict()
@@ -198,26 +208,30 @@ async def _transfer_kv_to_remote_paged_batch(self, transfer_reqs: List[KVMoveReq
 
                 start_kv_len = transfer_state.transfered_kv_len
                 trans_kv_len = min(trans_req.cur_kv_len - trans_req.prev_kv_len, self.nixl_agent.page_size)
-                trans_kv_index = transfer_state.token_index[start_kv_len: start_kv_len + trans_kv_len]
-                self.model.mem_manager.kv_move_buffer[page_index][:trans_kv_len] = self.model.mem_manager.kv_buffer[:,trans_kv_index, :, : ].transpose(0, 1)
+                trans_kv_index = transfer_state.token_index[start_kv_len : start_kv_len + trans_kv_len]
+                self.model.mem_manager.kv_move_buffer[page_index][:trans_kv_len] = self.model.mem_manager.kv_buffer[
+                    :, trans_kv_index, :, :
+                ].transpose(0, 1)
 
                 receive_page = transfer_state.free_page_ids.pop(0)
                 requests_by_agents[decode_id][0].append(page_index)
                 requests_by_agents[decode_id][1].append(receive_page)
-                is_last = (transfer_state.is_finished and start_kv_len + trans_kv_len == transfer_state.current_kv_len)
-
-                requests_by_agents[decode_id][2].append(RemotePrefillStatus(
-                    transfer_type=RemoteTransferType.PAGE_TRANSFER,
-                    group_req_id=group_req_id,
-                    status=RemoteTransferStatusType.SUCCESS,
-                    chunk_id=transfer_state.current_chunk_id,
-                    is_last=is_last,
-                    page_id=receive_page,
-                    kv_start=start_kv_len,
-                    kv_len=trans_kv_len,
-                    next_token_id=transfer_state.next_token_id,
-                    next_token_logprob=transfer_state.next_token_logprob
-                ))
+                is_last = transfer_state.is_finished and start_kv_len + trans_kv_len == transfer_state.current_kv_len
+
+                requests_by_agents[decode_id][2].append(
+                    RemotePrefillStatus(
+                        transfer_type=RemoteTransferType.PAGE_TRANSFER,
+                        group_req_id=group_req_id,
+                        status=RemoteTransferStatusType.SUCCESS,
+                        chunk_id=transfer_state.current_chunk_id,
+                        is_last=is_last,
+                        page_id=receive_page,
+                        kv_start=start_kv_len,
+                        kv_len=trans_kv_len,
+                        next_token_id=transfer_state.next_token_id,
+                        next_token_logprob=transfer_state.next_token_logprob,
+                    )
+                )
                 transfer_state.transfered_kv_len += trans_kv_len
 
         # wait copy done
@@ -227,11 +241,7 @@ async def _transfer_kv_to_remote_paged_batch(self, transfer_reqs: List[KVMoveReq
             # transfer
             self.nixl_agent.write_blocks_paged(decode_id, transfer_pages, receive_pages, notifications)
 
-
-        logger.info(
-            f"transfer kv to remote paged batch: {len(transfer_reqs)} "
-            f"took: {time.time() - start} seconds"
-        )
+        logger.info(f"transfer kv to remote paged batch: {len(transfer_reqs)} " f"took: {time.time() - start} seconds")
 
     async def _handle_transfer_loop(self):
         while True:
@@ -312,7 +322,6 @@ async def _wait_page_transfer_loop(self):
 
             await asyncio.sleep(PDNIXLBackendBase._THREAD_WAIT_INTERVAL)
 
-
     async def _wait_transfer_loop(self):
         while True:
             done_req_ids = self.nixl_agent.get_done_tranfers()
@@ -375,7 +384,7 @@ def _transfer_kv_to_remote(self, req: InferReq, group_req_id: int, cur_kv_len: i
 
         kv_transfer_req = KVMoveRequest(
             group_req_id=group_req_id,
-            token_ids=token_index[: cur_kv_len].tolist(),
+            token_ids=token_index[:cur_kv_len].tolist(),
             prev_kv_len=transfer_state.current_kv_len,
             cur_kv_len=cur_kv_len,
         )
@@ -403,11 +412,11 @@ def _post_remote_prefill(self, req: InferReq, success: bool = True):
         if self.is_master_in_dp:
             req.shm_req.shm_cur_kv_len = req.cur_kv_len
 
+        group_req_id = req.shm_req.group_req_id
         if not success:
             self.request_to_first_token.pop(group_req_id, None)
             return
 
-        group_req_id = req.shm_req.group_req_id
         assert group_req_id in self.request_to_first_token
         token_id, token_logprob = self.request_to_first_token.pop(group_req_id)
 
@@ -520,14 +529,13 @@ def _prepare_remote_prefill_inputs(self, req_objs: List[InferReq]):
             g_infer_context.radix_cache.free_radix_cache_to_get_enough_token(input_ids.shape[0])
         mem_indexes = g_infer_context.req_manager.mem_manager.alloc(input_ids.shape[0])
 
-
         req_to_token_indexs = g_infer_context.req_manager.req_to_token_indexs
         for idx, req_idx in enumerate(nopad_b_req_idx):
             cur_kv_len = req_objs[idx].cur_kv_len
             seq_len = nopad_b_seq_len[idx]
             mem_start = nopad_b_start_loc[idx]
-            mem_end = nopad_b_start_loc[idx+1]
-            req_to_token_indexs[req_idx, cur_kv_len:nopad_b_seq_len[idx]] = mem_indexes[mem_start:mem_end]
+            mem_end = nopad_b_start_loc[idx + 1]
+            req_to_token_indexs[req_idx, cur_kv_len : nopad_b_seq_len[idx]] = mem_indexes[mem_start:mem_end]
 
         kwargs = {
             "batch_size": len(run_reqs),
@@ -547,4 +555,4 @@ def _prefill_abort_remote(self, req_objs: List[InferReq]):
                 self.nixl_agent.send_abort_notify(self.remote_prefill_requests[group_req_id].decode_id, group_req_id)
                 del self.remote_prefill_requests[group_req_id]
                 if group_req_id in self.inflght_transfer_requests:
-                    del self.inflght_transfer_requests[group_req_id]
+                    del self.inflght_transfer_requests[group_req_id]

lightllm/server/router/model_infer/mode_backend/pd_nixl/impl_for_pd_decode.py

@@ -9,7 +9,12 @@
 from lightllm.utils.log_utils import init_logger
 from lightllm.server.multimodal_params import MultimodalParams
 
-from .pd_remote_prefill_obj import RemotePrefillTask, RemotePrefillServerInfo, RemotePrefillRequest, RemoteTransferStatusType
+from .pd_remote_prefill_obj import (
+    RemotePrefillTask,
+    RemotePrefillServerInfo,
+    RemotePrefillRequest,
+    RemoteTransferStatusType,
+)
 
 from .impl_for_pd_base import PDNIXLBackendBase
 
@@ -22,9 +27,9 @@ def __init__(self, prefill_task_queue: mp.Queue, prefill_done_queue: mp.Queue, n
 
     def init_custom(self):
         super().init_custom()
-        self.wait_prefill_thread = threading.Thread(target=self._start_async_loop,
-                                                    args=(self._prefill_wait_loop_async,),
-                                                    daemon=True)
+        self.wait_prefill_thread = threading.Thread(
+            target=self._start_async_loop, args=(self._prefill_wait_loop_async,), daemon=True
+        )
         self.wait_move_page_pool = ThreadPoolExecutor(max_workers=4)
         self.wait_prefill_thread.start()
         return
@@ -45,7 +50,7 @@ def _build_remote_prefill_task(self, index: int, kwargs: Dict, req: InferReq):
             multimodal_params=MultimodalParams.from_dict(req.multimodal_params),
             local_cached_len=req.cur_kv_len,
             token_ids=mem_indexes[b_start_loc[index] : b_start_loc[index + 1]],
-            page_ids=self.page_scheduer.borrow() # get page ids for this request, blocking when not enough pages
+            page_ids=self.page_scheduer.borrow(),  # get page ids for this request, blocking when not enough pages
         )
         return RemotePrefillTask(server_info=prefill_node_info, prefill_request=prefill_request)
 
@@ -82,16 +87,20 @@ def decode(self):
                 if self.is_master_in_dp:
                     run_req.remote_prefill_start = time.time()
                     # since this function may blocking the calling thread, so we do it in a thread pool
-                    self.wait_move_page_pool.submit(self._trigger_remote_prefill,
-                                                    shm_req.group_req_id, idx, kwargs, run_req)
+                    self.wait_move_page_pool.submit(
+                        self._trigger_remote_prefill, shm_req.group_req_id, idx, kwargs, run_req
+                    )
 
-                shm_req.set_pd_req_rank_state(self.rank_in_dp, RemoteTransferStatusType.IN_PROGRESS.value)  # set in progress state
+                shm_req.set_pd_req_rank_state(
+                    self.rank_in_dp, RemoteTransferStatusType.IN_PROGRESS.value
+                )  # set in progress state
                 run_req.in_prefill_or_transfer = True
                 self.remote_prefilled_reqs[shm_req.group_req_id] = run_req
 
         if decode_reqs:
             ContinuesBatchBackend.normal_decode(
-                self, decode_reqs=decode_reqs, uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs)
+                self, decode_reqs=decode_reqs, uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs
+            )
 
         self._overlap_req_init_and_filter(uninit_reqs=uninit_reqs, ok_finished_reqs=ok_finished_reqs, clear_list=True)
         return