adding shutdown endpoint part 2

Author: kumarrah2002

src/litserve/server.py

@@ -34,7 +34,7 @@
 
 import uvicorn
 import uvicorn.server
-from fastapi import Depends, FastAPI, HTTPException, Request, Response
+from fastapi import Depends, FastAPI, HTTPException, Request, Response, status
 from fastapi.responses import JSONResponse, StreamingResponse
 from fastapi.security import APIKeyHeader
 from starlette.formparsers import MultiPartParser
@@ -337,7 +337,7 @@ async def handle_request(self, request, request_type) -> StreamingResponse:
                 self.server.data_streamer(response_queue, data_available=event),
                 self.server._callback_runner.trigger_event,
                 EventTypes.ON_RESPONSE.value,
-                litserver=self.server,
+                litserve=self.server,
             )
 
             return StreamingResponse(response_generator)
@@ -357,6 +357,7 @@ def __init__(
         timeout: Union[float, bool] = 30,
         healthcheck_path: str = "/health",
         info_path: str = "/info",
+        shutdown_path: str = "/shutdown",
         model_metadata: Optional[dict] = None,
         spec: Optional[LitSpec] = None,
         max_payload_size=None,
@@ -370,6 +371,7 @@ def __init__(
         stream: bool = False,
         api_path: Optional[str] = None,
         loop: Optional[Union[str, LitLoop]] = None,
+        uvicorn_graceful_timeout: int = 30,
     ):
         """Initialize a LitServer instance for high-performance model inference.
 
@@ -394,6 +396,9 @@ def __init__(
 
             info_path (str, optional):
                 Server info endpoint path showing metadata and configuration. Defaults to "/info".
+                
+            shutdown_path (str, optional):
+                Server shutdown endpoint path that terminates and cleans up all worker and server processes. Defaults to "/shutdown".
 
             model_metadata (dict, optional):
                 Model metadata displayed at info endpoint (e.g., {"version": "1.0"}). Defaults to None.
@@ -419,6 +424,10 @@ def __init__(
             max_batch_size, batch_timeout, stream, spec, api_path, loop:
                 **Deprecated**: Configure these in your LitAPI implementation instead.
 
+            uvicorn_graceful_timeout (int, optional):
+                Timeout in seconds for Uvicorn to gracefully shut down its workers.
+                Defaults to 30.
+
         Example:
             >>> # Basic
             >>> server = LitServer(MyLitAPI())
@@ -441,7 +450,7 @@ def __init__(
                 "Old usage:\n"
                 "    server = LitServer(api, max_batch_size=N, batch_timeout=T, ...)\n\n"
                 "New usage:\n"
-                "    api = LitAPI(max_batch_size=N, batch_timeout=T, ...)\n"
+                f"    api = LitAPI(max_batch_size=N, batch_timeout=T, ...)\n"
                 "    server = LitServer(api, ...)",
                 DeprecationWarning,
                 stacklevel=2,
@@ -507,6 +516,7 @@ def __init__(
 
         self.healthcheck_path = healthcheck_path
         self.info_path = info_path
+        self._shutdown_path = shutdown_path
         self.track_requests = track_requests
         self.timeout = timeout
         self.litapi_connector.set_request_timeout(timeout)
@@ -531,6 +541,12 @@ def __init__(
         self.use_zmq = fast_queue
         self.transport_config = None
         self.litapi_request_queues = {}
+        self.inference_workers: List[mp.Process] = []
+        self.uvicorn_workers: List[Union[mp.Process, threading.Thread]] = []
+        self.manager: Optional[mp.Manager] = None
+        self._shutdown_event: Optional[mp.Event] = None
+        self._server_port: Optional[int] = None
+        self.uvicorn_graceful_timeout = uvicorn_graceful_timeout
 
         accelerator = self._connector.accelerator
         devices = self._connector.devices
@@ -542,8 +558,8 @@ def __init__(
                 device_list = range(devices)
             self.devices = [self.device_identifiers(accelerator, device) for device in device_list]
 
-        self.inference_workers = self.devices * self.workers_per_device
-        self.transport_config = TransportConfig(transport_config="zmq" if self.use_zmq else "mp")
+        self.inference_workers_config = self.devices * self.workers_per_device
+        self.transport_config = TransportConfig(transport_type="zmq" if self.use_zmq else "mp")
         self.register_endpoints()
 
     def launch_inference_worker(self, lit_api: LitAPI):
@@ -557,7 +573,7 @@ def launch_inference_worker(self, lit_api: LitAPI):
 
         process_list = []
         endpoint = lit_api.api_path.split("/")[-1]
-        for worker_id, device in enumerate(self.inference_workers):
+        for worker_id, device in enumerate(self.inference_workers_config):
             if len(device) == 1:
                 device = device[0]
 
@@ -578,6 +594,7 @@ def launch_inference_worker(self, lit_api: LitAPI):
             )
             process.start()
             process_list.append(process)
+        self.inference_workers.extend(process_list)
         return process_list
 
     @asynccontextmanager
@@ -698,6 +715,15 @@ def register_endpoints(self):
                 response_type = Response
             self._register_api_endpoints(lit_api, request_type, response_type)
 
+        # Register the new shutdown endpoint
+        @self.app.post(self._shutdown_path, status_code=status.HTTP_200_OK, dependencies=[Depends(self.setup_auth())])
+        async def shutdown_endpoint():
+            logger.info("Received shutdown request via /shutdown endpoint. Signaling main process.")
+            if self._shutdown_event:
+                self._shutdown_event.set()
+            return Response(content="Server is initiating graceful shutdown.", status_code=status.HTTP_200_OK)
+
+
     def _get_request_queue(self, api_path: str):
         return self.litapi_request_queues[api_path]
 
@@ -770,19 +796,105 @@ def verify_worker_status(self):
         logger.debug("One or more workers are ready to serve requests")
 
     def _init_manager(self, num_api_servers: int):
-        manager = self.transport_config.manager = mp.Manager()
+        self.manager = mp.Manager()
+        self.transport_config.manager = self.manager
         self.transport_config.num_consumers = num_api_servers
-        self.workers_setup_status = manager.dict()
+        self.workers_setup_status = self.manager.dict()
+        self._shutdown_event = self.manager.Event()
 
         # create request queues for each unique lit_api api_path
         for lit_api in self.litapi_connector:
-            self.litapi_request_queues[lit_api.api_path] = manager.Queue()
+            self.litapi_request_queues[lit_api.api_path] = self.manager.Queue()
 
         if self._logger_connector._loggers:
-            self.logger_queue = manager.Queue()
+            self.logger_queue = self.manager.Queue()
         self._logger_connector.run(self)
         self._transport = create_transport_from_config(self.transport_config)
-        return manager
+        return self.manager
+    
+    def _perform_graceful_shutdown(self):
+        """Encapsulates the graceful shutdown logic for LitServe."""
+        logger.info("Starting graceful shutdown of LitServe components.")
+
+        # 1. Close the message transport first to stop new messages
+        if self._transport:
+            logger.info("Closing message transport...")
+            self._transport.close()
+
+        # 2. Terminate and join inference worker processes
+        if self.inference_workers:
+            logger.info(f"Terminating {len(self.inference_workers)} inference workers...")
+            for i, iw in enumerate(self.inference_workers):
+                worker_pid = iw.pid
+                logger.info(f"Worker {i} (PID: {worker_pid}): Checking status BEFORE termination attempt. is_alive(): {iw.is_alive()}")
+
+                if iw.is_alive():
+                    try:
+                        iw.terminate()
+                        logger.info(f"Worker {i} (PID: {worker_pid}): Sent SIGTERM. Joining...")
+                        iw.join(timeout=5)
+
+                        logger.info(f"Worker {i} (PID: {worker_pid}): Status AFTER SIGTERM & JOIN. is_alive(): {iw.is_alive()}")
+                        if iw.is_alive():
+                            logger.warning(f"Worker {i} (PID: {worker_pid}): Did not terminate gracefully. Forcibly killing (SIGKILL).")
+                            iw.kill()
+                            logger.info(f"Worker {i} (PID: {worker_pid}): Status AFTER SIGKILL. is_alive(): {iw.is_alive()}")
+                        else:
+                            logger.info(f"Worker {i} (PID: {worker_pid}): Terminated gracefully.")
+                    except Exception as e:
+                        logger.error(f"Error during termination of worker {i} (PID: {worker_pid}): {e}")
+                else:
+                    logger.info(f"Worker {i} (PID: {worker_pid}): Was already not alive before termination attempt.")
+            logger.info("Inference workers termination loop completed.")
+
+        # 3. Terminate Uvicorn API server workers tracked by LitServe (the master processes/threads)
+        if self.uvicorn_workers:
+            logger.info(f"Terminating {len(self.uvicorn_workers)} Uvicorn API server master processes/threads...")
+            for i, uw in enumerate(self.uvicorn_workers):
+                logger.info(f"PID: {uw.pid}; {type(uw)}")
+                uvicorn_runner_pid = uw.pid
+                if uvicorn_runner_pid:
+                    log_prefix = f"Uvicorn Master {'Process' if uvicorn_runner_pid else 'Thread'} {i}"
+                    log_prefix += f" (PID: {uvicorn_runner_pid})"
+                    
+                if isinstance(uw, threading.Thread):
+                    logger.info(f"{log_prefix}: will terminate with the main process.")
+                else:
+                    logger.info(f"{log_prefix}: Checking status BEFORE termination. is_alive(): {uw.is_alive()}")
+                    if uw.is_alive():
+                        try:
+                            uw.terminate()
+                            logger.info(f"{log_prefix}: Sent SIGTERM. Joining (timeout={self.uvicorn_graceful_timeout}s)...")
+                            uw.join(timeout=self.uvicorn_graceful_timeout)
+                            logger.info(f"{log_prefix}: Status AFTER SIGTERM & JOIN. is_alive(): {uw.is_alive()}")
+                            if uw.is_alive():
+                                logger.warning(f"{log_prefix}: Did not terminate gracefully. Forcibly killing.")
+                                uw.kill()
+                                logger.info(f"{log_prefix}: Status AFTER SIGKILL. is_alive(): {uw.is_alive()}")
+                            else:
+                                logger.info(f"{log_prefix}: Terminated gracefully.")
+                        except Exception as e:
+                            logger.error(f"Error during termination of {log_prefix}: {e}")
+                    else:
+                        logger.info(f"{log_prefix}: Already not alive.")
+            logger.info("Uvicorn API server master processes/threads termination loop completed.")
+
+
+        # 5. Shut down the multiprocessing manager
+        if self.manager:
+            logger.info("Shutting down multiprocessing manager...")
+            self.manager.shutdown()
+            logger.info("Multiprocessing manager shut down.")
+
+        logger.info("All LitServe components gracefully shut down.")
+
+        # 6. Exit the main process
+        def exit_process():
+            time.sleep(0.5)
+            logger.info("Exiting main LitServe process.")
+            os._exit(0)
+
+        threading.Thread(target=exit_process).start()
 
     def run(
         self,
@@ -824,21 +936,6 @@ def run(
             **kwargs:
                 Additional uvicorn server options (ssl_keyfile, ssl_certfile, etc.). See uvicorn docs.
 
-        Example:
-            >>> server.run()  # Basic
-
-            >>> server.run(  # Production
-            ...     port=8080,
-            ...     num_api_servers=4,
-            ...     log_level="warning"
-            ... )
-
-            >>> server.run(  # Development
-            ...     log_level="debug",
-            ...     pretty_logs=True,
-            ...     generate_client_file=True
-            ... )
-
         """
         if generate_client_file:
             LitServer.generate_client_file(port=port)
@@ -859,9 +956,11 @@ def run(
         configure_logging(log_level, use_rich=pretty_logs)
         config = uvicorn.Config(app=self.app, host=host, port=port, log_level=log_level, **kwargs)
         sockets = [config.bind_socket()]
+        
+        self._server_port = port
 
         if num_api_servers is None:
-            num_api_servers = len(self.inference_workers)
+            num_api_servers = len(self.inference_workers_config)
 
         if num_api_servers < 1:
             raise ValueError("num_api_servers must be greater than 0")
@@ -874,34 +973,26 @@ def run(
         elif api_server_worker_type is None:
             api_server_worker_type = "process"
 
-        manager = self._init_manager(num_api_servers)
+        self._init_manager(num_api_servers)
         self._logger_connector.run(self)
-        inference_workers = []
         for lit_api in self.litapi_connector:
-            _inference_workers = self.launch_inference_worker(lit_api)
-            inference_workers.extend(_inference_workers)
+            self.launch_inference_worker(lit_api)
 
         self.verify_worker_status()
         try:
-            uvicorn_workers = self._start_server(
+            self.uvicorn_workers = self._start_server(
                 port, num_api_servers, log_level, sockets, api_server_worker_type, **kwargs
             )
             print(f"Swagger UI is available at http://0.0.0.0:{port}/docs")
-            # On Linux, kill signal will be captured by uvicorn.
-            # => They will join and raise a KeyboardInterrupt, allowing to Shutdown server.
-            for i, uw in enumerate(uvicorn_workers):
-                uw: Union[Process, Thread]
-                if isinstance(uw, Process):
-                    print(f"Uvicorn worker {i} : [{uw.pid}]")
-                uw.join()
+            
+            while not self._shutdown_event.is_set():
+                time.sleep(0.1)
+
+        except KeyboardInterrupt:
+            logger.info("KeyboardInterrupt received. Initiating graceful shutdown.")
         finally:
-            print("Shutting down LitServe")
-            self._transport.close()
-            for iw in inference_workers:
-                iw: Process
-                iw.terminate()
-                iw.join()
-            manager.shutdown()
+            self._perform_graceful_shutdown()
+
 
     def _prepare_app_run(self, app: FastAPI):
         # Add middleware to count active requests
@@ -920,17 +1011,17 @@ def _start_server(self, port, num_uvicorn_servers, log_level, sockets, uvicorn_w
             app: FastAPI = copy.copy(self.app)
 
             self._prepare_app_run(app)
-            config = uvicorn.Config(app=app, host="0.0.0.0", port=port, log_level=log_level, **kwargs)
-            if sys.platform == "win32" and num_uvicorn_servers > 1:
-                logger.debug("Enable Windows explicit socket sharing...")
-                # We make sure sockets is listening...
-                # It prevents further [WinError 10022]
-                for sock in sockets:
-                    sock.listen(config.backlog)
-                # We add worker to say unicorn to use a shared socket (win32)
-                # https://github.com/encode/uvicorn/pull/802
-                config.workers = num_uvicorn_servers
-            server = uvicorn.Server(config=config)
+            uvicorn_config = uvicorn.Config(
+                app=app,
+                host="0.0.0.0",
+                port=port,
+                log_level=log_level,
+                workers=num_uvicorn_servers if uvicorn_worker_type == "process" else 1,
+                timeout_graceful_shutdown=self.uvicorn_graceful_timeout,
+                **kwargs
+            )
+
+            server = uvicorn.Server(config=uvicorn_config)
             if uvicorn_worker_type == "process":
                 ctx = mp.get_context("fork")
                 w = ctx.Process(target=server.run, args=(sockets,))
@@ -947,4 +1038,4 @@ def setup_auth(self):
             return self.lit_api.authorize
         if LIT_SERVER_API_KEY:
             return api_key_auth
-        return no_auth
+        return no_auth