fix race window

Author: lvhan028

lmdeploy/serve/core/async_engine.py

@@ -236,6 +236,11 @@ async def stop_all_session(self):
         self.epoch += 1
         await self.session_mgr.async_abort_all()
 
+    def prepare_sleep(self):
+        """Reject new inference requests before backend sleep starts."""
+        self.sleeping_tags = {'weights', 'kv_cache'}
+        self.is_sleeping = True
+
     def sleep(self, level: int = 1):
         """Sleep the model.
 

lmdeploy/serve/openai/api_server.py

@@ -1190,6 +1190,7 @@ async def sleep(raw_request: Request = None):
     if level not in (1, 2):
         return create_error_response(HTTPStatus.BAD_REQUEST, 'The "level" query parameter must be 1 or 2.')
     async_engine = VariableInterface.async_engine
+    async_engine.prepare_sleep()
     await async_engine.stop_all_session()
     async_engine.sleep(level)
     return Response(status_code=200)