Reject requests to protect server (#1275)

This PR monitors the request queue within the shortfin server and
rejects requests with a 503 error when the queue is filled with more
than max batch size + 2

TODO: Add load testing to integration tests

Author: zeeshanhaque21

shortfin/python/shortfin_apps/llm/components/generate.py

@@ -18,6 +18,8 @@
 
 # TODO: Have a generic "Responder" interface vs just the concrete impl.
 from shortfin.interop.fastapi import FastAPIResponder
+from fastapi.responses import JSONResponse
+from fastapi import status
 
 from .config_struct import DecodeConfig
 from .io_struct import (
@@ -130,6 +132,7 @@ class ClientGenerateBatchProcess(sf.Process):
         "responder",
         "tokenizer",
         "decode_config",
+        "service",
     ]
 
     def __init__(
@@ -140,6 +143,7 @@ def __init__(
         fiber: sf.Fiber | None = None,
     ):
         super().__init__(fiber=service.main_fiber if fiber is None else fiber)
+        self.service = service
         self.gen_req = gen_req
         self.responder = responder
         self.tokenizer = service.tokenizer
@@ -151,12 +155,29 @@ def __init__(
 
     async def run(self):
         logger.debug("Started ClientBatchGenerateProcess: %r", self)
-        streaming = self.gen_req.stream
-        self.responder.start_response()
-        if streaming:
-            self.responder.stream_start()
+
+        # Try to add request to queue
+        # TODO(@zphoenixrises): Add load testing and integration tests for this.
+        if not self.service.add_to_queue():
+            error_response = JSONResponse(
+                status_code=status.HTTP_503_SERVICE_UNAVAILABLE,
+                content={
+                    "error": "Server queue is full. Please try again later.",
+                    "code": "QUEUE_FULL",
+                    "current_size": self.service.current_queue_size,
+                    "max_size": self.service.max_queue_size,
+                },
+            )
+            self.responder.send_response(error_response)
+            self.responder.ensure_response()
+            return
 
         try:
+            streaming = self.gen_req.stream
+            self.responder.start_response()
+            if streaming:
+                self.responder.stream_start()
+
             # Launch all individual generate processes and wait for them to finish.
             gen_processes = []
             input_ids = self.gen_req.input_ids
@@ -166,6 +187,7 @@ async def run(self):
                 input_batch = [input_ids] if self.gen_req.is_single else input_ids
             else:
                 input_batch = self.tokenize()
+
             for index, input_tokens in enumerate(input_batch):
                 decode_config = copy(self.decode_config)
                 decode_config.update_from_sampling_params(
@@ -189,7 +211,10 @@ async def run(self):
 
             await asyncio.gather(*gen_processes)
             self.generate_response(gen_processes, streaming)
+
         finally:
+            # Remove request from queue when done
+            self.service.remove_from_queue()
             self.responder.ensure_response()
 
     def generate_response(

shortfin/python/shortfin_apps/llm/components/service.py

@@ -24,7 +24,6 @@
 
 from ...utils import GenerateService
 
-
 logger = logging.getLogger(__name__)
 
 
@@ -44,18 +43,41 @@ def __init__(
         model_params: ModelParams,
         server_params: "ServerParams",
         program_isolation: str = "per_call",
+        max_queue_size: int = 3,  # Maximum number of requests in queue
     ):
         super().__init__(sysman)
         self.name = name
         self.tokenizer = tokenizer
         self.model_params = model_params
         self.server_params = server_params
+        self.max_queue_size = max_queue_size
+        self.current_queue_size = 0
 
         self.set_isolation(program_isolation)
         self.initialize_worker_and_fiber()
+        self.initialize_queues()
         self.initialize_page_cache()
 
+    def initialize_queues(self):
+        """Initialize request and response queues"""
+        if self.model_params.decode_batch_sizes:
+            self.max_queue_size = max(self.model_params.decode_batch_sizes) + 2
+            print(f"Max queue size: {self.max_queue_size}")
+
+    def add_to_queue(self) -> bool:
+        """Try to add a request to the queue. Returns True if successful, False if queue is full."""
+        if self.current_queue_size >= self.max_queue_size:
+            return False
+        self.current_queue_size += 1
+        return True
+
+    def remove_from_queue(self):
+        """Remove a request from the queue."""
+        if self.current_queue_size > 0:
+            self.current_queue_size -= 1
+
     def initialize_worker_and_fiber(self):
+
         self.main_worker = self.sysman.ls.create_worker(f"{self.name}-inference")
         self.main_fiber = self.sysman.ls.create_fiber(self.main_worker)
         self.prefill_fiber = self.sysman.ls.create_fiber(self.main_worker)

shortfin/python/shortfin_apps/llm/server.py

@@ -68,7 +68,7 @@ def parse_args(argv):
     return parser.parse_args(argv)
 
 
-def main(argv, log_config=uvicorn.config.LOGGING_CONFIG):
+def run_server(argv, log_config=uvicorn.config.LOGGING_CONFIG, port: int | None = None):
     args = parse_args(argv)
     if args.tokenizer_config_json is None:
         # this is only used for the EOS token
@@ -84,7 +84,7 @@ def main(argv, log_config=uvicorn.config.LOGGING_CONFIG):
     uvicorn.run(
         get_app(lifecycle_manager.fastapi_lifespan),
         host=args.host,
-        port=args.port,
+        port=port or args.port,
         log_config=log_config,
         timeout_keep_alive=args.timeout_keep_alive,
     )
@@ -94,7 +94,7 @@ def main(argv, log_config=uvicorn.config.LOGGING_CONFIG):
     from shortfin.support.logging_setup import configure_main_logger
 
     logger = configure_main_logger("server")
-    main(
+    run_server(
         sys.argv[1:],
         # Make logging defer to the default shortfin logging config.
         log_config=UVICORN_LOG_CONFIG,