async sleep and sync wakeup

Author: lvhan028

lmdeploy/pytorch/engine/engine.py

@@ -443,9 +443,9 @@ def update_params(self, request: Any):
         """Update params."""
         self.executor.update_params(request)
 
-    def sleep(self, level: int = 1):
+    async def sleep(self, level: int = 1):
         """Sleep."""
-        self.executor.sleep(level)
+        await self.executor.sleep(level)
 
     def wakeup(self, tags: list[str] | None = None):
         """Wakeup."""

lmdeploy/pytorch/engine/executor/mp_executor.py

@@ -373,6 +373,14 @@ def warmup(self):
         """Build cache engine."""
         self.collective_rpc('warmup')
 
+    async def sleep(self, level: int = 1):
+        """Sleep."""
+        await self.collective_rpc_async('sleep', args=(level, ), return_mask=0)
+
+    def wakeup(self, tags: list[str] | None = None):
+        """Wakeup."""
+        self.collective_rpc('wakeup', args=(tags, ), return_mask=0)
+
     async def _prefetch_outputs(self):
         while True:
             out = (await self.collective_rpc_async('get_outputs', receiver_mask=1, return_mask=1))[0]

lmdeploy/pytorch/engine/executor/ray_executor.py

@@ -321,6 +321,18 @@ def collective_rpc(self,
             kwargs = dict()
         return ray.get([getattr(worker, method).remote(*args, **kwargs) for worker in self.workers], timeout=timeout)
 
+    async def collective_rpc_async(self,
+                                   method: str,
+                                   args: tuple[Any] = None,
+                                   kwargs: dict[str, Any] = None):
+        """Collective async rpc."""
+        if args is None:
+            args = list()
+        if kwargs is None:
+            kwargs = dict()
+        tasks = [getattr(worker, method).remote(*args, **kwargs) for worker in self.workers]
+        return await asyncio.gather(*tasks)
+
     def build_model(self):
         """Build model."""
         self.collective_rpc('build_model')
@@ -353,9 +365,9 @@ def warmup(self):
         """Build cache engine."""
         self.collective_rpc('warmup')
 
-    def sleep(self, level: int = 1):
+    async def sleep(self, level: int = 1):
         """Sleep."""
-        self.collective_rpc('sleep', (level, ))
+        await self.collective_rpc_async('sleep', (level, ))
 
     def wakeup(self, tags: list[str] | None = None):
         """Wakeup."""

lmdeploy/pytorch/engine/executor/uni_executor.py

@@ -108,6 +108,14 @@ async def get_output_async(self, dp_rank: int = 0):
         assert dp_rank == 0
         return await self.model_agent.get_output_async()
 
+    async def sleep(self, level: int = 1):
+        """Sleep."""
+        await self.model_agent.sleep(level)
+
+    def wakeup(self, tags: list[str] | None = None):
+        """Wakeup."""
+        self.model_agent.wakeup(tags)
+
     def get_input_processor(self):
         """Get input processor."""
         return self.model_agent.get_input_processor()

lmdeploy/pytorch/engine/mp_engine/base.py

@@ -53,9 +53,9 @@ def end_session(self, session_id: int):
         """End session."""
         return self._collective_rpc('end_session', session_id)
 
-    def sleep(self, level: int):
+    async def sleep(self, level: int):
         """sleep."""
-        return self._collective_rpc('sleep', level)
+        return await self._collective_rpc_async('sleep', level)
 
     def wakeup(self, tags: list[str] | None = None):
         """Wakeup."""

lmdeploy/pytorch/engine/mp_engine/base_worker.py

@@ -100,9 +100,9 @@ def p2p_drop_connect(self, drop_conn_request: DistServeDropConnectionRequest):
         """
         return self.engine.p2p_drop_connect(drop_conn_request)
 
-    def sleep(self, level: int = 1):
+    async def sleep(self, level: int = 1):
         """sleep."""
-        return self.engine.sleep(level)
+        return await self.engine.sleep(level)
 
     def wakeup(self, tags: list[str] | None = None):
         """Wakeup."""

lmdeploy/serve/core/async_engine.py

@@ -209,8 +209,8 @@ def _if_session_stale(self, session: Session,
         epoch = session.epoch
         if epoch is None or epoch == self.epoch:
             return None
-        logger.info(
-            f'[generate] session {session.session_id} dropped (session.epoch={epoch}, epoch={self.epoch})')
+        logger.info(f'[generate] drop stale session {session.session_id} '
+                    f'(session.epoch={epoch}, async_engine.epoch={self.epoch})')
         return GenOut(response='',
                       history_token_len=session.step,
                       input_token_len=input_token_len,
@@ -241,15 +241,15 @@ def prepare_sleep(self):
         self.sleeping_tags = {'weights', 'kv_cache'}
         self.is_sleeping = True
 
-    def sleep(self, level: int = 1):
+    async def sleep(self, level: int = 1):
         """Sleep the model.
 
         Args:
             level (int): The sleep level. Level 1 sleep will offload the model
                 weights and discard the kv cache. Level 2 sleep will
                 discard both the model weights and the kv cache.
         """
-        self.engine.sleep(level)
+        await self.engine.sleep(level)
         self.sleeping_tags = {'weights', 'kv_cache'}
         self.is_sleeping = True
 
@@ -460,7 +460,7 @@ def is_error(status):
         async with session.request_handle() as handle:
             if session.epoch is not None and session.epoch != self.epoch:
                 logger.info(f'[generate] session {session_id} got aborted before starting inference, '
-                               f'session.epoch={session.epoch}, epoch={self.epoch}')
+                               f'session.epoch={session.epoch}, async_engine.epoch={self.epoch}')
                 metrics_processor.increase_failed_requests('abort')
                 yield GenOut(response='',
                              history_token_len=0,

lmdeploy/serve/managers/session_manager.py

@@ -105,7 +105,7 @@ async def request_handle(self):
 
     async def async_abort(self):
         """Abort the session."""
-        logger.info(f'[session] Aborting session {self.session_id}, epoch={self.epoch}')
+        logger.debug(f'[session] Aborting session {self.session_id}, epoch={self.epoch}')
         if self._handle is not None:
             await self._handle.async_cancel(self.session_id)
 
@@ -216,6 +216,7 @@ def get(self, session_id: int | None = None, **kwargs) -> Session:
 
     async def async_abort_all(self):
         """Abort all sessions."""
+        logger.info(f'[SessionManager] aborting all {len(self.sessions)} sessions')
         tasks = []
         for session in list(self.sessions.values()):
             tasks.append(session.async_abort())

lmdeploy/serve/openai/api_server.py

@@ -1192,7 +1192,7 @@ async def sleep(raw_request: Request = None):
     async_engine = VariableInterface.async_engine
     async_engine.prepare_sleep()
     await async_engine.stop_all_session()
-    async_engine.sleep(level)
+    await async_engine.sleep(level)
     return Response(status_code=200)