[PD Disaggregation] 1. fix p2p_initialize => async, 2. fix metrics in PD 3. fix uniexecutor interface (#3893)

Author: JimyMa

lmdeploy/pytorch/engine/engine.py

@@ -1094,7 +1094,7 @@ async def _async_loop_migration(self, resp_que: asyncio.Queue, has_runable_event
                         resp=msg.resp,
                         finish=False,
                         token_ids=np.array(token_ids),
-                        metrics_info=req_metrics,
+                        req_metrics=req_metrics,
                     )
                     outputs[session_id] = out
                     self.update_running_migration([msg], np.array([token_ids]), [False], [None])

lmdeploy/pytorch/engine/executor/uni_executor.py

@@ -112,9 +112,9 @@ def p2p_initialize(self, init_request: DistServeInitRequest):
         """
         return [self.model_agent.cache_engine.p2p_initialize(init_request)]
 
-    def p2p_connect(self, conn_request: List[DistServeKVTransferEndpointInfo]):
+    def p2p_connect(self, remote_engine_id: str, conn_request: List[DistServeKVTransferEndpointInfo]):
         """rdma_connect."""
-        self.model_agent.cache_engine.p2p_connect(conn_request)
+        self.model_agent.cache_engine.p2p_connect(remote_engine_id, conn_request)
 
     async def migrate(self, batch: MigrationExecutionBatch):
         """KV Cache Migration."""

lmdeploy/pytorch/engine/mp_engine/base_worker.py

@@ -83,9 +83,9 @@ def get_schedule_metrics(self):
         """Get schedule metrics."""
         return self.engine.get_schedule_metrics()
 
-    def p2p_initialize(self, conn_request: DistServeInitRequest):
+    async def p2p_initialize(self, conn_request: DistServeInitRequest):
         """Init rdma link."""
-        return self.engine.p2p_initialize(conn_request)
+        return await self.engine.p2p_initialize(conn_request)
 
     def p2p_connect(self, conn_request: DistServeConnectionRequest):
         """rdma_connect."""