fix

Author: hiworldwzj

lightllm/server/httpserver_for_pd_master/manager.py

@@ -55,11 +55,11 @@ def __init__(
         return
 
     async def register_pd(self, pd_info_json, websocket):
-        await self.pd_manager.register_pd(pd_info_json, websocket)
+        self.pd_manager.register_pd(pd_info_json, websocket)
         return
 
     async def remove_pd(self, pd_info_json):
-        await self.pd_manager.remove_pd(pd_info_json)
+        self.pd_manager.remove_pd(pd_info_json)
         return
 
     async def update_req_status(self, upkv_status: UpKVStatus):
@@ -92,7 +92,7 @@ def tokens(self, prompt, multimodal_params, samping_params: SamplingParams, kwar
     async def select_p_d_node(
         self, prompt: Union[str, List[int]], sampling_params: SamplingParams, multimodal_params: MultimodalParams
     ) -> Tuple[PD_Client_Obj, PD_Client_Obj]:
-        return await self.pd_manager.select_p_d_node(prompt, sampling_params, multimodal_params)
+        return self.pd_manager.select_p_d_node(prompt, sampling_params, multimodal_params)
 
     async def generate(
         self,
@@ -411,7 +411,7 @@ def __init__(self, args):
         )
         return
 
-    async def register_pd(self, pd_info_json, websocket):
+    def register_pd(self, pd_info_json, websocket):
         pd_client = PD_Client_Obj(**pd_info_json)
         pd_client.websocket = websocket
         self.url_to_pd_nodes[pd_client.client_ip_port] = pd_client
@@ -425,19 +425,19 @@ async def register_pd(self, pd_info_json, websocket):
         else:
             assert False, f"mode must in ['prefill', 'decode'], but get {pd_client.mode}"
 
-        await self.selector.update_nodes(self.prefill_nodes, self.decode_nodes)
+        self.selector.update_nodes(self.prefill_nodes, self.decode_nodes)
 
         logger.info(f"mode: {pd_client.mode} url: {pd_client.client_ip_port} registed")
         return
 
-    async def remove_pd(self, pd_info_json):
+    def remove_pd(self, pd_info_json):
         pd_client = PD_Client_Obj(**pd_info_json)
 
         self.url_to_pd_nodes.pop(pd_client.client_ip_port, None)
         self.prefill_nodes = [e for e in self.prefill_nodes if e.client_ip_port != pd_client.client_ip_port]
         self.decode_nodes = [e for e in self.decode_nodes if e.client_ip_port != pd_client.client_ip_port]
 
-        await self.selector.update_nodes(self.prefill_nodes, self.decode_nodes)
+        self.selector.update_nodes(self.prefill_nodes, self.decode_nodes)
 
         logger.info(f"mode: {pd_client.mode} url: {pd_client.client_ip_port} removed")
         return
@@ -461,13 +461,8 @@ def update_node_load_info(self, load_info: Optional[dict]):
             logger.warning(f"udpate node load info failed, load_info: {load_info} error: {str(e)}")
         return
 
-    def get_predict_node_infos(self):
-        """获取所有节点的预测负载信息"""
-        return self.node_info_recorder.get_predict_node_infos()
-
-    async def select_p_d_node(
+    def select_p_d_node(
         self, prompt: Union[str, List[int]], sampling_params: SamplingParams, multimodal_params: MultimodalParams
     ) -> Tuple[PD_Client_Obj, PD_Client_Obj]:
-        p_node, d_node = await self.selector.select_p_d_node(prompt, sampling_params, multimodal_params)
-        self.node_info_recorder.update_predict_node_info(p_node, d_node, prompt, sampling_params, multimodal_params)
+        p_node, d_node = self.selector.select_p_d_node(prompt, sampling_params, multimodal_params)
         return p_node, d_node

lightllm/server/httpserver_for_pd_master/pd_selector/pd_selector.py

@@ -12,18 +12,22 @@ def __init__(self, prefill_nodes: List[PD_Client_Obj], decode_nodes: List[PD_Cli
         self.decode_nodes: List[PD_Client_Obj] = decode_nodes
         self.pd_manager = pd_manager
 
-    async def update_nodes(self, prefill_nodes, decode_nodes):
+    def update_nodes(self, prefill_nodes, decode_nodes):
         self.prefill_nodes = prefill_nodes
         self.decode_nodes = decode_nodes
 
-    async def select_p_d_node(self, prompt: Union[str, List[int]], sampling_params: SamplingParams, multimodal_params: MultimodalParams) -> Tuple[PD_Client_Obj, PD_Client_Obj]:
+    def select_p_d_node(
+        self, prompt: Union[str, List[int]], sampling_params: SamplingParams, multimodal_params: MultimodalParams
+    ) -> Tuple[PD_Client_Obj, PD_Client_Obj]:
         raise NotImplementedError("Subclass must implement this method")
 
 
 class RandomSelector(PDSelector):
     """随机选择器"""
 
-    async def select_p_d_node(self, prompt: Union[str, List[int]], sampling_params: SamplingParams, multimodal_params: MultimodalParams) -> Tuple[PD_Client_Obj, PD_Client_Obj]:
+    def select_p_d_node(
+        self, prompt: Union[str, List[int]], sampling_params: SamplingParams, multimodal_params: MultimodalParams
+    ) -> Tuple[PD_Client_Obj, PD_Client_Obj]:
         p_node = random.choice(self.prefill_nodes)
         d_node = random.choice(self.decode_nodes)
         return p_node, d_node
@@ -37,7 +41,9 @@ def __init__(self, prefill_nodes: List[PD_Client_Obj], decode_nodes: List[PD_Cli
         self.prefill_node_index: int = 0
         self.decode_node_index: int = 0
 
-    async def select_p_d_node(self, prompt: Union[str, List[int]], sampling_params: SamplingParams, multimodal_params: MultimodalParams) -> Tuple[PD_Client_Obj, PD_Client_Obj]:
+    def select_p_d_node(
+        self, prompt: Union[str, List[int]], sampling_params: SamplingParams, multimodal_params: MultimodalParams
+    ) -> Tuple[PD_Client_Obj, PD_Client_Obj]:
         p_node = self.prefill_nodes[self.prefill_node_index]
         d_node = self.decode_nodes[self.decode_node_index]
         self.prefill_node_index = (self.prefill_node_index + 1) % len(self.prefill_nodes)
@@ -48,7 +54,9 @@ async def select_p_d_node(self, prompt: Union[str, List[int]], sampling_params:
 class MemorySelector(PDSelector):
     """基于内存使用情况的选择器"""
 
-    async def select_p_d_node(self, prompt: Union[str, List[int]], sampling_params: SamplingParams, multimodal_params: MultimodalParams) -> Tuple[PD_Client_Obj, PD_Client_Obj]:
+    def select_p_d_node(
+        self, prompt: Union[str, List[int]], sampling_params: SamplingParams, multimodal_params: MultimodalParams
+    ) -> Tuple[PD_Client_Obj, PD_Client_Obj]:
         def _get_min_node(nodes: List[PD_Client_Obj], node_infos: Dict[str, dict], key: str) -> PD_Client_Obj:
             min_node, min_node_value = None, float("inf")
             for node in nodes: