add option --enable-abort-handling

Author: lvhan028

lmdeploy/cli/serve.py

@@ -73,6 +73,7 @@ def add_parser_api_server():
         ArgumentHelper.max_log_len(parser)
         ArgumentHelper.disable_fastapi_docs(parser)
         ArgumentHelper.allow_terminate_by_client(parser)
+        ArgumentHelper.enable_abort_handling(parser)
         # chat template args
         ArgumentHelper.chat_template(parser)
 
@@ -266,6 +267,7 @@ def api_server(args):
                            allow_methods=args.allow_methods,
                            allow_headers=args.allow_headers,
                            allow_terminate_by_client=args.allow_terminate_by_client,
+                           enable_abort_handling=args.enable_abort_handling,
                            log_level=args.log_level.upper(),
                            api_keys=args.api_keys,
                            ssl=args.ssl,
@@ -293,6 +295,7 @@ def api_server(args):
                           allow_methods=args.allow_methods,
                           allow_headers=args.allow_headers,
                           allow_terminate_by_client=args.allow_terminate_by_client,
+                          enable_abort_handling=args.enable_abort_handling,
                           log_level=args.log_level.upper(),
                           api_keys=args.api_keys,
                           ssl=args.ssl,

lmdeploy/cli/utils.py

@@ -454,6 +454,16 @@ def allow_terminate_by_client(parser):
                                    default=False,
                                    help='Enable server to be terminated by request from client')
 
+    @staticmethod
+    def enable_abort_handling(parser):
+        """Add --enable-abort-handling argument to configure server abort
+        request processing."""
+
+        return parser.add_argument('--enable-abort-handling',
+                                   action='store_true',
+                                   default=False,
+                                   help='Enable server to handle client abort requests')
+
     @staticmethod
     def cache_max_entry_count(parser):
         """Add argument cache_max_entry_count to parser."""

lmdeploy/serve/openai/api_server.py

@@ -65,6 +65,7 @@ class VariableInterface:
     # following is for tool parsers
     tool_parser: Optional[ToolParser] = None
     allow_terminate_by_client: bool = False
+    enable_abort_handling: bool = False
 
 
 router = APIRouter()
@@ -1152,6 +1153,11 @@ async def free_cache(cache_free_request: DistServeCacheFreeRequest) -> JSONRespo
 @router.post('/abort_request')
 async def abort_request(request: AbortRequest, raw_request: Request = None):
     """Abort an ongoing request."""
+    if not VariableInterface.enable_abort_handling:
+        return Response(
+            status_code=501,
+            content='This server does not support abort requests. Enable with --enable-abort-handling flag.')
+
     if request.abort_all:
         await VariableInterface.async_engine.stop_all_session()
     else:
@@ -1323,6 +1329,7 @@ def serve(model_path: str,
           reasoning_parser: Optional[str] = None,
           tool_call_parser: Optional[str] = None,
           allow_terminate_by_client: bool = False,
+          enable_abort_handling: bool = False,
           **kwargs):
     """An example to perform model inference through the command line
     interface.
@@ -1381,6 +1388,7 @@ def serve(model_path: str,
     logger.setLevel(log_level)
 
     VariableInterface.allow_terminate_by_client = allow_terminate_by_client
+    VariableInterface.enable_abort_handling = enable_abort_handling
     if api_keys is not None:
         if isinstance(api_keys, str):
             api_keys = api_keys.split(',')