allow api server terminated through requests from clients (#3533)

* allow terminate

* fix engine role for tubomind

Author: RunningLeon

lmdeploy/cli/serve.py

@@ -133,7 +133,7 @@ def add_parser_api_server():
         ArgumentHelper.model_name(parser)
         ArgumentHelper.max_log_len(parser)
         ArgumentHelper.disable_fastapi_docs(parser)
-
+        ArgumentHelper.allow_terminate_by_client(parser)
         # chat template args
         ArgumentHelper.chat_template(parser)
 
@@ -361,6 +361,7 @@ def api_server(args):
                        allow_credentials=args.allow_credentials,
                        allow_methods=args.allow_methods,
                        allow_headers=args.allow_headers,
+                       allow_terminate_by_client=args.allow_terminate_by_client,
                        log_level=args.log_level.upper(),
                        api_keys=args.api_keys,
                        ssl=args.ssl,

lmdeploy/cli/utils.py

@@ -406,6 +406,15 @@ def tool_call_parser(parser):
             default=None,
             help=f'The registered tool parser name {ToolParserManager.module_dict.keys()}. Default to None.')
 
+    @staticmethod
+    def allow_terminate_by_client(parser):
+        """Add argument allow_terminate_by_client to parser."""
+
+        return parser.add_argument('--allow-terminate-by-client',
+                                   action='store_true',
+                                   default=False,
+                                   help='Enable server to be terminated by request from client')
+
     @staticmethod
     def cache_max_entry_count(parser):
         """Add argument cache_max_entry_count to parser."""

lmdeploy/serve/openai/api_server.py

@@ -54,6 +54,7 @@ class VariableInterface:
     reasoning_parser: Optional[ReasoningParser] = None
     # following is for tool parsers
     tool_parser: Optional[ToolParser] = None
+    allow_terminate_by_client: bool = False
 
 
 router = APIRouter()
@@ -234,6 +235,19 @@ async def health() -> Response:
     return Response(status_code=200)
 
 
+@router.get('/terminate')
+async def terminate():
+    """terminate server."""
+    import signal
+
+    if not VariableInterface.allow_terminate_by_client:
+        return create_error_response(
+            HTTPStatus.BAD_REQUEST,
+            'The server can not be terminated. Please add --allow-terminate-by-client when start the server.')
+    os.kill(os.getpid(), signal.SIGTERM)
+    return Response(status_code=200)
+
+
 # modified from https://github.com/vllm-project/vllm/blob/v0.5.4/vllm/entrypoints/openai/logits_processors.py#L51  # noqa
 def logit_bias_logits_processor(logit_bias: Union[Dict[int, float], Dict[str, float]], tokenizer) -> LogitsProcessor:
     try:
@@ -1056,21 +1070,16 @@ async def startup_event():
     try:
         import requests
         engine_config = VariableInterface.async_engine.engine.engine_config
+        engine_role = engine_config.role.value if hasattr(engine_config, 'role') else 1
         url = f'{VariableInterface.proxy_url}/nodes/add'
-        data = {
-            'url': VariableInterface.api_server_url,
-            'status': {
-                'models': get_model_list(),
-                'role': engine_config.role.value
-            }
-        }
+        data = {'url': VariableInterface.api_server_url, 'status': {'models': get_model_list(), 'role': engine_role}}
         headers = {'accept': 'application/json', 'Content-Type': 'application/json'}
         response = requests.post(url, headers=headers, json=data)
 
         if response.status_code != 200:
             raise HTTPException(status_code=400, detail='Service registration failed')
     except Exception as e:
-        print(f'Service registration failed: {e}')
+        logger.error(f'Service registration failed: {e}')
 
 
 class ConcurrencyLimitMiddleware(BaseHTTPMiddleware):
@@ -1127,6 +1136,7 @@ def serve(model_path: str,
           max_concurrent_requests: Optional[int] = None,
           reasoning_parser: Optional[str] = None,
           tool_call_parser: Optional[str] = None,
+          allow_terminate_by_client: bool = False,
           **kwargs):
     """An example to perform model inference through the command line
     interface.
@@ -1178,6 +1188,7 @@ def serve(model_path: str,
             clients concurrently during that time. Default to None.
         reasoning_parser (str): The reasoning parser name.
         tool_call_parser (str): The tool call parser name.
+        allow_terminate_by_client (bool): Allow request from client to terminate server.
     """
     if os.getenv('TM_LOG_LEVEL') is None:
         os.environ['TM_LOG_LEVEL'] = log_level
@@ -1207,6 +1218,7 @@ def serve(model_path: str,
     if max_concurrent_requests is not None:
         app.add_middleware(ConcurrencyLimitMiddleware, max_concurrent_requests=max_concurrent_requests)
 
+    VariableInterface.allow_terminate_by_client = allow_terminate_by_client
     if api_keys is not None:
         if isinstance(api_keys, str):
             api_keys = api_keys.split(',')

lmdeploy/serve/proxy/proxy.py

@@ -179,6 +179,35 @@ def remove(self, node_url: str):
             for conn in dropped_conn:
                 self.pd_connection_pool.drop(*conn)
 
+    def terminate_node(self, node_url: str):
+        """terminate a node."""
+        success = True
+        if node_url in self.nodes:
+            self.nodes.pop(node_url)
+            headers = {'accept': 'application/json'}
+            try:
+                response = requests.get(f'{node_url}/terminate', headers=headers)
+                if response.status_code != 200:
+                    success = False
+                    logger.error(f'Failed to terminate node {node_url}, '
+                                 f'error_code={response.status_code}, '
+                                 f'error_msg={response.text}')
+            except:  # noqa
+                success = False
+        else:
+            success = False
+        self.update_config_file()
+        return success
+
+    def terminate_all_nodes(self):
+        """terminate all nodes."""
+        node_url_li = list(self.nodes.keys())
+        all_success = True
+        for node_url in node_url_li:
+            if not self.terminate_node(node_url):
+                all_success = False
+        return all_success
+
     def remove_stale_nodes_by_expiration(self):
         """remove stale nodes."""
         to_be_deleted = []
@@ -433,9 +462,10 @@ def add_node(node: Node, raw_request: Request = None):
 
 
 @app.post('/nodes/remove', dependencies=[Depends(check_api_key)])
-def remove_node(node_url: str):
+def remove_node(node: Node):
     """Show available models."""
     try:
+        node_url = node.url
         node_manager.remove(node_url)
         logger.info(f'delete node {node_url} successfully')
         return 'Deleted successfully'
@@ -444,6 +474,33 @@ def remove_node(node_url: str):
         return 'Failed to delete, please check the input url.'
 
 
+@app.post('/nodes/terminate', dependencies=[Depends(check_api_key)])
+def terminate_node(node: Node):
+    """terminate nodes."""
+    try:
+        node_url = node.url
+        success = node_manager.terminate_node(node_url)
+        if not success:
+            return f'Failed to terminate node {node_url}'
+        return 'Terminated successfully'
+    except:  # noqa
+        logger.error(f'Terminate node {node_url} failed.')
+        return 'Failed to terminate node {node_url}, please check the input url.'
+
+
+@app.get('/nodes/terminate_all', dependencies=[Depends(check_api_key)])
+def terminate_node_all():
+    """terminate nodes."""
+    try:
+        success = node_manager.terminate_all_nodes()
+        if not success:
+            return 'Failed to terminate all nodes'
+        return 'All nodes terminated successfully'
+    except:  # noqa
+        logger.error('Failed to terminate all nodes')
+        return 'Failed to terminate all nodes.'
+
+
 @app.post('/distserve/connection_warmup')
 async def connection_warmup():
     await asyncio.gather(*[