reformat.

Author: hiworldwzj

lightllm/server/api_start.py

@@ -283,10 +283,10 @@ def pd_master_start(args):
     set_unique_server_name(args)
     if args.run_mode != "pd_master":
         return
-    
+
     # when use config_server to support multi pd_master node, we
     # need generate unique node id for each pd_master node.
-    # otherwise, we use the 0 for single pd_master node. 
+    # otherwise, we use the 0 for single pd_master node.
     if args.config_server_host and args.config_server_port:
         args.pd_node_id = uuid.uuid4().int
     else:
@@ -344,7 +344,7 @@ def config_server_start(args):
     set_unique_server_name(args)
     if args.run_mode != "config_server":
         return
-    
+
     logger.info(f"all start args:{args}")
 
     set_env_start_args(args)

lightllm/server/config_server/__init__.py

@@ -9,4 +9,4 @@
 - Global configuration query: Provides mechanisms for querying shared 
     configuration data across the system.
 - Designed for distributed systems operating in PD separation mode.
-"""
+"""

lightllm/server/config_server/api_http.py

@@ -12,7 +12,7 @@
 logger = init_logger(__name__)
 app = FastAPI()
 
-registered_pd_master_objs:Dict[str, PD_Master_Obj] = {}
+registered_pd_master_objs: Dict[str, PD_Master_Obj] = {}
 registered_pd_master_obj_lock = Lock()
 
 global_req_id = 0
@@ -61,13 +61,15 @@ async def websocket_endpoint(websocket: WebSocket):
             registered_pd_master_objs.pop(registered_pd_master_obj.node_id, None)
     return
 
+
 @app.get("/registered_objects")
 async def get_registered_objects():
     with registered_pd_master_obj_lock:
         serialized_data = pickle.dumps(registered_pd_master_objs)
-        base64_encoded = base64.b64encode(serialized_data).decode('utf-8')
+        base64_encoded = base64.b64encode(serialized_data).decode("utf-8")
         return {"data": base64_encoded}
 
+
 @app.get("/allocate_global_unique_id_range")
 async def allocate_global_id_range():
     """
@@ -85,8 +87,8 @@ async def allocate_global_id_range():
     global global_req_id
     range_size = 800000
     with global_req_id_lock:
-        if global_req_id + range_size > 2**63 - 1:
-            global_req_id = 0    
+        if global_req_id + range_size > 2 ** 63 - 1:
+            global_req_id = 0
         start_id = global_req_id
         global_req_id += range_size
         end_id = global_req_id

lightllm/server/httpserver/pd_loop.py

@@ -15,6 +15,7 @@
 
 logger = init_logger(__name__)
 
+
 async def timer_log(manager: HttpServerManager):
     while True:
         await asyncio.sleep(30)
@@ -32,13 +33,13 @@ async def pd_handle_loop(manager: HttpServerManager):
 
     asyncio.create_task(timer_log(manager))
 
-    id_to_handle_task:Dict[int, asyncio.Task] = {}
+    id_to_handle_task: Dict[int, asyncio.Task] = {}
 
     while True:
         try:
             id_to_pd_master_obj = await _get_pd_master_objs(manager.args)
             logger.info(f"get pd_master_objs {id_to_pd_master_obj}")
-            
+
             if id_to_pd_master_obj is not None:
                 for node_id, pd_master_obj in id_to_handle_task.items():
                     if node_id not in id_to_pd_master_obj:
@@ -51,7 +52,7 @@ async def pd_handle_loop(manager: HttpServerManager):
                         id_to_handle_task[node_id] = asyncio.create_task(_pd_handle_task(manager, pd_master_obj))
 
             await asyncio.sleep(30)
-        
+
         except Exception as e:
             logger.exception(str(e))
             await asyncio.sleep(10)
@@ -70,7 +71,7 @@ async def _pd_handle_task(manager: HttpServerManager, pd_master_obj: PD_Master_O
         try:
             uri = f"ws://{pd_master_obj.host_ip_port}/pd_register"
             async with websockets.connect(uri, max_queue=(2048 * 1024, 2048 * 1023)) as websocket:
-                
+
                 sock = websocket.transport.get_extra_info("socket")
                 sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
 
@@ -83,35 +84,35 @@ async def _pd_handle_task(manager: HttpServerManager, pd_master_obj: PD_Master_O
                     "mode": manager.pd_mode.value,
                     "start_args": args_dict,
                 }
-                 
+
                 await websocket.send(json.dumps(regist_json))
                 logger.info(f"Sent registration JSON: {regist_json}")
-                
+
                 # 转发任务
-                forwarding_tokens_task = asyncio.create_task(
-                    _up_tokens_to_pd_master(forwarding_queue, websocket)
-                )
-                
+                forwarding_tokens_task = asyncio.create_task(_up_tokens_to_pd_master(forwarding_queue, websocket))
+
                 # 接收 pd master 发来的请求，并推理后，将生成的token转发回pd master。
                 while True:
                     recv_bytes = await websocket.recv()
                     obj = pickle.loads(recv_bytes)
                     if obj[0] == ObjType.REQ:
                         prompt, sampling_params, multimodal_params = obj[1]
-                        asyncio.create_task(_pd_process_generate(manager, prompt, sampling_params, multimodal_params, forwarding_queue))
+                        asyncio.create_task(
+                            _pd_process_generate(manager, prompt, sampling_params, multimodal_params, forwarding_queue)
+                        )
                     elif obj[0] == ObjType.ABORT:
                         group_req_id = obj[1]
                         await manager.abort(group_req_id)
                     else:
                         logger.error(f"recevie error obj {str(obj)}")
-                
+
         except asyncio.CancelledError:
             # 如果任务被取消，则退出循环
             logger.warning(f"forwarding_tokens_task {pd_master_obj} cancelled")
             if forwarding_tokens_task is not None:
                 forwarding_tokens_task.cancel()
             return
-        
+
         except Exception as e:
             logger.error("connetion to pd_master has error")
             logger.exception(str(e))
@@ -122,7 +123,7 @@ async def _pd_handle_task(manager: HttpServerManager, pd_master_obj: PD_Master_O
             logger.info("reconnection to pd_master")
 
 
-async def _get_pd_master_objs(args)->Optional[Dict[int, PD_Master_Obj]]:
+async def _get_pd_master_objs(args) -> Optional[Dict[int, PD_Master_Obj]]:
     """
     get_pd_master_objs 主要负责从 pd master 获取所有的pd master对象。
     """
@@ -135,15 +136,15 @@ async def _get_pd_master_objs(args)->Optional[Dict[int, PD_Master_Obj]]:
         ans = dict()
         ans[0] = PD_Master_Obj(node_id=0, host_ip_port=f"{args.pd_master_ip}:{args.pd_master_port}")
         return ans
-    
+
     # 使用 config_server 服务来发现所有的 pd_master 节点。
     uri = f"ws://{args.config_server_host}:{args.config_server_port}/registered_objects"
 
     try:
         async with httpx.AsyncClient() as client:
             response = await client.get(uri)
             if response.status_code == 200:
-                base64data  = response.json()["data"]
+                base64data = response.json()["data"]
                 id_to_pd_master_obj = pickle.loads(base64.b64decode(base64data))
                 return id_to_pd_master_obj
             else:
@@ -154,8 +155,11 @@ async def _get_pd_master_objs(args)->Optional[Dict[int, PD_Master_Obj]]:
         await asyncio.sleep(10)
         return None
 
+
 # 触发推理的task
-async def _pd_process_generate(manager: HttpServerManager, prompt, sampling_params, multimodal_params, forwarding_queue:AsyncQueue):
+async def _pd_process_generate(
+    manager: HttpServerManager, prompt, sampling_params, multimodal_params, forwarding_queue: AsyncQueue
+):
     try:
         async for sub_req_id, request_output, metadata, finish_status in manager.generate(
             prompt, sampling_params, multimodal_params, None
@@ -175,4 +179,3 @@ async def _up_tokens_to_pd_master(forwarding_queue: AsyncQueue, websocket):
         handle_list = await forwarding_queue.wait_to_get_all_data()
         if handle_list:
             await websocket.send(pickle.dumps((ObjType.TOKEN_PACKS, handle_list)))
-

lightllm/server/httpserver_for_pd_master/register_loop.py

@@ -10,6 +10,7 @@
 
 logger = init_logger(__name__)
 
+
 async def register_loop(manager: HttpServerManagerForPDMaster):
     assert manager.args.host not in ["127.0.0.1", "localhost"], "pd mode must specify host ip"
 
@@ -23,15 +24,17 @@ async def register_loop(manager: HttpServerManagerForPDMaster):
         try:
             uri = f"ws://{manager.args.config_server_host}:{manager.args.config_server_port}/pd_master_register"
             async with websockets.connect(uri, max_queue=(2048 * 1024, 2048 * 1023)) as websocket:
-                
+
                 sock = websocket.transport.get_extra_info("socket")
                 sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
 
-                pd_master_obj = PD_Master_Obj(node_id=manager.args.pd_node_id, host_ip_port=f"{manager.host_ip}:{manager.args.port}")
+                pd_master_obj = PD_Master_Obj(
+                    node_id=manager.args.pd_node_id, host_ip_port=f"{manager.host_ip}:{manager.args.port}"
+                )
 
                 await websocket.send(pickle.dumps(pd_master_obj))
                 logger.info(f"Sent registration pd_master obj: {pd_master_obj}")
-                
+
                 while True:
                     await websocket.send("heartbeat")
                     await asyncio.sleep(60)
@@ -41,4 +44,3 @@ async def register_loop(manager: HttpServerManagerForPDMaster):
             logger.exception(str(e))
             await asyncio.sleep(10)
             logger.info("reconnection to config_server")
-

lightllm/server/pd_io_struct.py

@@ -46,15 +46,17 @@ def __post_init__(self):
 
     def to_llm_url(self):
         return f"http://{self.client_ip_port}/pd_generate_stream"
-    
+
+
 @dataclass
 class PD_Master_Obj:
     node_id: int
     host_ip_port: str
-    
+
     def to_log_str(self):
         return f"PD_MASTER host_ip_port: {self.host_ip_port} node_id: {self.node_id}"
 
+
 @dataclass
 class UpKVStatus:
     type: str = "kv_move_status"
@@ -73,7 +75,7 @@ def __post_init__(self):
             error_info = "group_request_id only can be int"
             logger.error(error_info)
             raise ValueError(error_info)
-        
+
         if not isinstance(self.pd_master_node_id, int):
             error_info = "pd_master_node_id only can be int"
             logger.error(error_info)
@@ -163,4 +165,4 @@ def get_cost_time(self):
 @dataclass
 class KVMoveTaskGroup:
     tasks: List[KVMoveTask]
-    connect_id: str
+    connect_id: str

lightllm/server/req_id_generator.py

@@ -17,9 +17,11 @@ def __init__(self):
         from lightllm.server.core.objs.atomic_lock import AtomicShmLock
         from lightllm.server.core.objs.shm_array import ShmArray
         from lightllm.utils.envs_utils import get_unique_server_name, get_env_start_args
-        
+
         self.args = get_env_start_args()
-        self.use_config_server = self.args.config_server_host and self.args.config_server_port and self.args.run_mode == "pd_master"
+        self.use_config_server = (
+            self.args.config_server_host and self.args.config_server_port and self.args.run_mode == "pd_master"
+        )
         self.current_id = ShmArray(f"{get_unique_server_name()}_req_id_gen", (2,), dtype=np.int64)
         self.current_id.create_shm()
         self.current_id.arr[0] = 0
@@ -43,7 +45,9 @@ def _check_and_set_new_id_range(self):
                             # 保证id满足倍乘关系
                             self.current_id.arr[0] = (id_range["start_id"] // MAX_BEST_OF + 1) * MAX_BEST_OF
                             self.current_id.arr[1] = id_range["end_id"]
-                            assert self.current_id.arr[0] + MAX_BEST_OF < self.current_id.arr[1], f"get id range error {self.current_id.arr[0]} {self.current_id.arr[1]}"
+                            assert (
+                                self.current_id.arr[0] + MAX_BEST_OF < self.current_id.arr[1]
+                            ), f"get id range error {self.current_id.arr[0]} {self.current_id.arr[1]}"
                             return
                         else:
                             raise RuntimeError(f"Failed to fetch ID range from config server: {response.status_code}")

lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/decode_node_impl/decode_trans_obj.py

@@ -160,9 +160,11 @@ def put_to_radix_loop(self):
                         f"{func_name} put kv to radix cache ok, req_id: {task.id()} cost_time {task.get_cost_time()} s"
                     )
                     self.manager.up_status_in_queue.put(
-                        UpKVStatus(group_request_id=task.group_request_id, 
-                                   dp_index=task.decode_dp_index,
-                                   pd_master_node_id=task.decode_node.pd_master_node_id)
+                        UpKVStatus(
+                            group_request_id=task.group_request_id,
+                            dp_index=task.decode_dp_index,
+                            pd_master_node_id=task.decode_node.pd_master_node_id,
+                        )
                     )
                     logger.info(f"{func_name} up kv status req_id: {task.id()} finished")
                 move_tasks.clear()

lightllm/server/router/model_infer/mode_backend/continues_batch/pd_mode/decode_node_impl/up_status.py

@@ -26,21 +26,21 @@ def __init__(self, args, task_in_queue: mp.Queue, task_out_queue: mp.Queue):
 
     def thread_loop(self):
         asyncio.run(self.task_loop())
-    
+
     async def task_loop(self):
 
-        self.id_to_handle_task:Dict[int, asyncio.Task] = {}
-        self.id_to_handle_queue:Dict[int, asyncio.Queue] = {}
+        self.id_to_handle_task: Dict[int, asyncio.Task] = {}
+        self.id_to_handle_queue: Dict[int, asyncio.Queue] = {}
 
         asyncio.create_task(self.dispatch_task_loop())
 
         while True:
             try:
                 from lightllm.server.httpserver.pd_loop import _get_pd_master_objs
-    
+
                 id_to_pd_master_obj = await _get_pd_master_objs(self.args)
                 logger.info(f"get pd_master_objs {id_to_pd_master_obj}")
-                
+
                 if id_to_pd_master_obj is not None:
                     for node_id, pd_master_obj in self.id_to_handle_task.items():
                         if node_id not in id_to_pd_master_obj:
@@ -55,23 +55,23 @@ async def task_loop(self):
                             self.id_to_handle_task[node_id] = asyncio.create_task(self.up_kv_status_task(pd_master_obj))
 
                 await asyncio.sleep(30)
-            
+
             except Exception as e:
                 logger.exception(str(e))
                 await asyncio.sleep(10)
 
     async def dispatch_task_loop(self):
-            while True:
-                try:
-                    loop = asyncio.get_event_loop()
-                    upkv_status: UpKVStatus = await loop.run_in_executor(None, self.task_queue.get)
-                    if upkv_status.pd_master_node_id in self.id_to_handle_queue:
-                        await self.id_to_handle_queue[upkv_status.pd_master_node_id].put(upkv_status)
-                    else:
-                        logger.warning(f"upstatus {upkv_status} no connection to pd_master, drop it")
-                except BaseException as e:
-                    logger.exception(str(e))
-                    await asyncio.sleep(10)
+        while True:
+            try:
+                loop = asyncio.get_event_loop()
+                upkv_status: UpKVStatus = await loop.run_in_executor(None, self.task_queue.get)
+                if upkv_status.pd_master_node_id in self.id_to_handle_queue:
+                    await self.id_to_handle_queue[upkv_status.pd_master_node_id].put(upkv_status)
+                else:
+                    logger.warning(f"upstatus {upkv_status} no connection to pd_master, drop it")
+            except BaseException as e:
+                logger.exception(str(e))
+                await asyncio.sleep(10)
 
     async def up_kv_status_task(self, pd_master_obj: PD_Master_Obj):
         while True:
@@ -98,15 +98,14 @@ async def up_kv_status_task(self, pd_master_obj: PD_Master_Obj):
             except asyncio.CancelledError:
                 logger.info(f"up_kv_status_task {pd_master_obj} cancelled")
                 return
-             
+
             except Exception as e:
                 logger.error(f"connetion to pd_master {pd_master_obj} has error: {str(e)}")
                 logger.exception(str(e))
                 await asyncio.sleep(10)
                 logger.info("reconnection to pd_master")
 
 
-
 def _init_env(args, task_in_queue: mp.Queue, task_out_queue: mp.Queue):
     graceful_registry(inspect.currentframe().f_code.co_name)
     up_kv_manager = UpStatusManager(args, task_in_queue, task_out_queue)