实现 推理引擎的 shutdown方法

Author: shell-nlp

gpt_server/model_backend/base.py

@@ -6,3 +6,6 @@ class ModelBackend(ABC):
     @abstractmethod
     def stream_chat(self, params: Dict[str, Any]):
         pass
+
+    def shutdown(self):
+        pass

gpt_server/model_backend/hf_backend.py

@@ -54,6 +54,9 @@ def __init__(self, tokenizer: PreTrainedTokenizer, model: torch.nn.Module) -> No
                     continue
                 self.model.load_adapter(model_id=lora_path, adapter_name=lora_name)
 
+    def shutdown(self):
+        pass
+
     async def stream_chat(self, params: Dict[str, Any]):
         # params 已不需要传入 prompt
         messages = params["messages"]

gpt_server/model_backend/lmdeploy_backend.py

@@ -102,6 +102,9 @@ def __init__(self, model_path, tokenizer: PreTrainedTokenizer) -> None:
         # 自定义日志
         self.async_engine.request_logger = CustomRequestLogger(max_log_len=None)
 
+    def shutdown(self):
+        pass
+
     async def stream_chat(self, params: Dict[str, Any]) -> AsyncGenerator:
         # params 已不需要传入 prompt
         messages = params["messages"]

gpt_server/model_backend/sglang_backend.py

@@ -8,7 +8,7 @@
 import sglang as sgl
 from transformers import PreTrainedTokenizer
 from sglang.utils import convert_json_schema_to_str
-
+from sglang.srt.entrypoints.engine import Engine
 from qwen_vl_utils import process_vision_info
 from sglang.srt.managers.io_struct import GenerateReqInput
 from gpt_server.settings import get_model_config
@@ -48,7 +48,7 @@ def __init__(self, model_path, tokenizer: PreTrainedTokenizer) -> None:
         logger.info(f"model_config: {model_config}")
         self.lora_requests = []
         # ---
-        self.async_engine = sgl.Engine(
+        self.async_engine: Engine = sgl.Engine(
             model_path=model_path,
             trust_remote_code=True,
             mem_fraction_static=model_config.gpu_memory_utilization,
@@ -60,6 +60,9 @@ def __init__(self, model_path, tokenizer: PreTrainedTokenizer) -> None:
         )
         self.tokenizer = tokenizer
 
+    def shutdown(self):
+        self.async_engine.shutdown()
+
     async def stream_chat(self, params: Dict[str, Any]) -> AsyncGenerator:
         # params 已不需要传入 prompt
         messages = params["messages"]

gpt_server/model_backend/vllm_backend.py

@@ -61,6 +61,9 @@ def __init__(self, model_path, tokenizer: PreTrainedTokenizer) -> None:
         self.tokenizer = tokenizer
         self.reasoning_parser_cache = {}
 
+    def shutdown(self):
+        self.engine.shutdown()
+
     async def stream_chat(self, params: Dict[str, Any]) -> AsyncGenerator:
         # params 已不需要传入 prompt
         messages = params["messages"]

gpt_server/model_worker/base/model_worker_base.py

@@ -375,6 +375,12 @@ async def startup():
                 limit_worker_concurrency=limit_worker_concurrency,
             )
 
+        @app.on_event("shutdown")
+        async def shutdown():
+            global worker
+            # 优雅推出
+            worker.backend.shutdown()
+
         uvicorn.run(app, host=host, port=port)