fix lint

Author: Weichao Luo

lightllm/server/pd_io_struct.py

@@ -127,6 +127,7 @@ class RemotePrefillServerInfo:
     prefill_server_ip: str
     prefill_server_port: int
 
+
 @dataclass
 class DistInfo:
     world_size: int
@@ -136,6 +137,7 @@ class DistInfo:
     dp_size_in_node: int
     node_world_size: int
 
+
 @dataclass
 class PDTransLeaveInfo:
     decode_id: int

lightllm/server/router/manager.py

@@ -210,12 +210,18 @@ async def wait_to_model_ready(self):
                 start_pd_remote_prefill_server_process,
             )
 
-            dist_info = DistInfo(self.world_size, self.nnodes, self.dp_size,
-                                 self.dp_world_size, self.dp_size_in_node, self.node_world_size)
+            dist_info = DistInfo(
+                self.world_size,
+                self.nnodes,
+                self.dp_size,
+                self.dp_world_size,
+                self.dp_size_in_node,
+                self.node_world_size,
+            )
 
             start_pd_remote_prefill_server_process(
                 self.args.pd_node_id,
-                dist_info = dist_info,
+                dist_info=dist_info,
                 http_server_port=self.args.pd_remote_prefill_http_port,
                 server_port=self.args.pd_remote_prefill_port,
                 from_backend_queue=self.info_queue,
@@ -235,8 +241,15 @@ async def wait_to_model_ready(self):
             from lightllm.server.router.model_infer.mode_backend.pd_nixl.pd_remote_prefill import (
                 start_pd_remote_prefill_client_process,
             )
-            dist_info = DistInfo(self.world_size, self.nnodes, self.dp_size,
-                                 self.dp_world_size, self.dp_size_in_node, self.node_world_size)
+
+            dist_info = DistInfo(
+                self.world_size,
+                self.nnodes,
+                self.dp_size,
+                self.dp_world_size,
+                self.dp_size_in_node,
+                self.node_world_size,
+            )
 
             start_pd_remote_prefill_client_process(
                 self.args.pd_node_id,

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

@@ -160,7 +160,7 @@ def _transfer_kv_to_remote(self, req: InferReq):
 
         kv_transfer_req = KVMoveRequest(
             group_req_id=group_req_id,
-            token_ids=token_index[ : req.cur_kv_len].tolist(),
+            token_ids=token_index[: req.cur_kv_len].tolist(),
             prev_kv_len=transfer_state.current_kv_len,
             cur_kv_len=req.cur_kv_len,
         )
@@ -176,7 +176,6 @@ def _transfer_kv_to_remote(self, req: InferReq):
         transfer_state.current_kv_len = req.cur_kv_len
         transfer_state.current_chunk_id += 1
 
-
     def _decode_filter_reqs(
         self, prefill_reqs: List[InferReq], aborted_reqs: List[InferReq], decode_reqs: List[InferReq]
     ):

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

@@ -25,6 +25,7 @@ def init_custom(self):
 
         self.reduce_tensor = torch.tensor([0], dtype=torch.int32, device="cuda", requires_grad=False)
         from lightllm.server.router.model_infer.mode_backend.dp_backend.pre_process import padded_prepare_prefill_inputs
+
         kwargs, run_reqs, padded_req_num = padded_prepare_prefill_inputs([], 1, is_multimodal=self.is_multimodal)
         self.model.forward(**kwargs)
         assert len(run_reqs) == 0 and padded_req_num == 1

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

@@ -73,7 +73,11 @@ 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=lambda req, _1, _2: self._transfer_kv_to_remote(req),
             )
 
@@ -104,6 +108,10 @@ 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=lambda req, _1, _2: self._transfer_kv_to_remote(req),
             )

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

@@ -7,9 +7,13 @@
 from lightllm.utils.log_utils import init_logger
 
 from .pd_remote_prefill_obj import (
-    RemoteAgent, KVMoveRequest, PrefillRequest,
-    RemotePrefillStatus, ThreadSafeDict, KVMoveRequestState
-    )
+    RemoteAgent,
+    KVMoveRequest,
+    PrefillRequest,
+    RemotePrefillStatus,
+    ThreadSafeDict,
+    KVMoveRequestState,
+)
 
 
 logger = init_logger(__name__)
@@ -120,10 +124,10 @@ def write_blocks(self, request: KVMoveRequest, prefill_request: PrefillRequest,
             return
 
         kv_move_start = max(skip_kv_move_len, request.prev_kv_len)
-        kv_move_end   = request.cur_kv_len
+        kv_move_end = request.cur_kv_len
 
         src_token_ids = request.token_ids[kv_move_start:]
-        dst_token_ids = prefill_request.data.token_ids[kv_move_start - skip_kv_move_len: kv_move_end]
+        dst_token_ids = prefill_request.data.token_ids[kv_move_start - skip_kv_move_len : kv_move_end]
 
         remote_agent: RemoteAgent = self.remote_agents[prefill_request.decode_id][
             self.tp_idx
@@ -140,7 +144,8 @@ def write_blocks(self, request: KVMoveRequest, prefill_request: PrefillRequest,
                 group_req_id=group_reqeust_id,
                 status=1,
                 chunk_id=prefill_request.transfer_state.current_chunk_id,
-                is_last=is_finished)
+                is_last=is_finished,
+            )
 
             handle = self.nixl_agent.make_prepped_xfer(
                 "WRITE", src_handle, src_token_descs, dst_handle, dst_token_descs, notify_status.serialize()
@@ -151,10 +156,7 @@ def write_blocks(self, request: KVMoveRequest, prefill_request: PrefillRequest,
 
             if group_reqeust_id not in self.inflight_transfers:
                 self.inflight_transfers[group_reqeust_id] = KVMoveRequestState(
-                    handles=[],
-                    done_handles=[],
-                    remote_agent=remote_agent,
-                    abort=False
+                    handles=[], done_handles=[], remote_agent=remote_agent, abort=False
                 )
             self.inflight_transfers[group_reqeust_id].handles.append(handle)
 
@@ -199,7 +201,9 @@ def get_done_tranfers(self):
                     logger.warning(f"{req_id} Transfer failed with state {xfer_state}")
                     failed = True
                     kv_move_state.done_handles.append(handle)
-                    notify_failed_status = RemotePrefillStatus(group_req_id=req_id, status=-1, chunk_id=-1, is_last=True)
+                    notify_failed_status = RemotePrefillStatus(
+                        group_req_id=req_id, status=-1, chunk_id=-1, is_last=True
+                    )
                     self.nixl_agent.send_notif(remote_agent.name, notify_failed_status.serialize())
 
             kv_move_state.handles = left_handles

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

@@ -112,13 +112,17 @@ def main_loop(self):
                     if self.dist_info.dp_size_in_node > 1:
                         group_req_id = request.data.sampling_params.group_request_id
                         suggested_dp_index = request.data.sampling_params.suggested_dp_index
-                        if suggested_dp_index < 0: # not likely to happen
+                        if suggested_dp_index < 0:  # not likely to happen
                             suggested_dp_index = random.randint(0, self.dist_info.dp_size_in_node)
                             request.data.sampling_params.suggested_dp_index = suggested_dp_index
-                            logger.warning(f"Suggested dp index is negative for {group_req_id}, set to {suggested_dp_index}")
+                            logger.warning(
+                                f"Suggested dp index is negative for {group_req_id}, set to {suggested_dp_index}"
+                            )
 
-                        for local_rank in range(suggested_dp_index * self.dist_info.dp_world_size,
-                                                (suggested_dp_index + 1) * self.dist_info.dp_world_size):
+                        for local_rank in range(
+                            suggested_dp_index * self.dist_info.dp_world_size,
+                            (suggested_dp_index + 1) * self.dist_info.dp_world_size,
+                        ):
                             self.to_backend_queues[local_rank].put(request)
                     else:
                         for queue in self.to_backend_queues:
@@ -217,7 +221,11 @@ def main_loop(self):
     def remote_prefill(self, server_id: int, prefill_request: RemotePrefillRequest):
         socket, _ = self.remote_prefill_servers[server_id]
         prefill_request.sampling_params.max_new_tokens = 1
-        socket.send_pyobj(PrefillRequest(type=RemoteRequstType.REMOTE_PREFILL, decode_id=self.id, data=prefill_request, transfer_state=None))
+        socket.send_pyobj(
+            PrefillRequest(
+                type=RemoteRequstType.REMOTE_PREFILL, decode_id=self.id, data=prefill_request, transfer_state=None
+            )
+        )
 
 
 def remote_prefill_server_loop(
@@ -256,7 +264,11 @@ def start_pd_remote_prefill_server_process(
 
 
 def remote_prefill_client_loop(
-    id: int, dist_info: DistInfo, from_backend_queue: mp.Queue, to_backend_queues: List[mp.Queue], agent_meta_queues: List[mp.Queue]
+    id: int,
+    dist_info: DistInfo,
+    from_backend_queue: mp.Queue,
+    to_backend_queues: List[mp.Queue],
+    agent_meta_queues: List[mp.Queue],
 ):
     graceful_registry(inspect.currentframe().f_code.co_name)
 
@@ -271,11 +283,16 @@ def remote_prefill_client_loop(
 
 
 def start_pd_remote_prefill_client_process(
-    id: int, dist_info: DistInfo, from_backend_queue: mp.Queue, to_backend_queues: List[mp.Queue], agent_meta_queues: List[mp.Queue]
+    id: int,
+    dist_info: DistInfo,
+    from_backend_queue: mp.Queue,
+    to_backend_queues: List[mp.Queue],
+    agent_meta_queues: List[mp.Queue],
 ):
 
     proc = mp.Process(
-        target=remote_prefill_client_loop, args=(id, dist_info, from_backend_queue, to_backend_queues, agent_meta_queues)
+        target=remote_prefill_client_loop,
+        args=(id, dist_info, from_backend_queue, to_backend_queues, agent_meta_queues),
     )
     proc.start()
     assert proc.is_alive()

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

@@ -53,12 +53,14 @@ class ConnectRequest(RemoteRequest):
     agent_metadatas: List[bytes]
     agent_mem_descs: List[bytes]
 
+
 @dataclass
 class TransferState:
     start_time: float
     current_kv_len: int
     current_chunk_id: int
 
+
 @dataclass
 class PrefillRequest(RemoteRequest):
     decode_id: int
@@ -82,6 +84,7 @@ class RemoteAgent:
     kv_mem_desc: nixlBind.nixlRegDList
     kv_xfer_handles: nixl_prepped_dlist_handle
 
+
 @dataclass
 class KVMoveRequestState:
     handles: List[nixl_xfer_handle]
@@ -184,4 +187,3 @@ def bind(self, addr: str):
 
     def connect(self, addr: str):
         return self.sock.connect(addr)
-