Repository: Lightning-AI/LitServe

Issue:
shutdown LitServe programmatically
<!--

⚠️ BEFORE SUBMITTING, READ:

We're excited for your request! However, here are things we are not interested in:

- Decorators.
- Doing the same thing in multiple ways.    
- Adding more layers of abstraction... tree-depth should be 1 at most.
- Features that over-engineer or complicate the code internals.    
- Linters, and crud that complicates projects.    

-->

----

## 🚀 Feature

Multiple users have requested for a way to stop a running LitServe programatically rather terminating the process manually via keyboard interrupt such as https://github.com/Lightning-AI/LitServe/issues/491. 

One way is to have an endpoint `/shutdown` which gracefully terminates and cleans up everything. 

---

Set `enable_shutdown_api=True` in `LitServer` to enable the `/shutdown` api endpoint which can be used to stop LitServe programmatically. 

```
import litserve as ls

if __name__ == "__main__":
    api = ls.test_examples.SimpleLitAPI()
    server = ls.LitServer(api, workers_per_device=10, enable_shutdown_api=True)
    server.run(port=8000, pretty_logs=True)
```



Discussion:
Hey @aniketmaurya,
I've done some research and implemented a shutdown endpoint that successfully terminates the processes. However, there are still some exceptions thrown during shutdown due to Uvicorn's inherited run command.

Specifically, when using os.kill(os.getpid(), signal.SIGINT) to stop the server, exceptions are expected, but they can't be suppressed because we don't have control over Uvicorn’s internal run behavior.

Is this acceptable as is, or would you prefer that we find a way to suppress those exception messages?

<img width="644" alt="Image" src="https://github.com/user-attachments/assets/64a60199-3636-432e-b6ae-d4d7e92afcde" />
Hi @kumarrah2002 , 
About the exceptions during shutdown — I think it might be acceptable as long as we ensure that:
No child processes or threads are left running and port is fully released for reuse.

btw, @aniketmaurya 
One question out of curiosity: Incase if this endpoint is accessible publicly, wouldn’t it pose a security risk by allowing anyone to shut down the server?

great point @bhimrazy! let's add additional authorization for `/shutdown`. 
@aniketmaurya An idea I had for security was to generate a unique API key that is only for the shutdown endpoint (e.g. `SHUTDOWN_API_KEY`) that is saved to the user's environment or logged in the terminal. When the shutdown endpoint is called, this key must be present or there will be a 401 Unauthorized error. I discussed this with @bhimrazy but wanted to get your perspective. 
CC: @Borda 
sounds good @kumarrah2002! additionally, shutdown endpoint should be optional and only enable that if the user provides an argument such as `LitServer(enable_shutdown_api=False)`. 



Patch:
diff --git a/src/litserve/loops/base.py b/src/litserve/loops/base.py
index 470cc3d8..557c4460 100644
--- a/src/litserve/loops/base.py
+++ b/src/litserve/loops/base.py
@@ -225,7 +225,7 @@ def __init__(self):
 
     def kill(self):
         try:
-            logger.info(f"Stop Server Requested - Kill parent pid [{self._server_pid}] from [{os.getpid()}]")
+            logger.debug(f"Stop Server Requested - Kill parent pid [{self._server_pid}] from [{os.getpid()}]")
             if sys.platform == "win32":
                 os.kill(self._server_pid, signal.SIGTERM)
         except PermissionError:
diff --git a/src/litserve/loops/loops.py b/src/litserve/loops/loops.py
index 8920bb7e..a9249c87 100644
--- a/src/litserve/loops/loops.py
+++ b/src/litserve/loops/loops.py
@@ -67,7 +67,6 @@ def inference_worker(
     workers_setup_status: Dict[int, str],
     callback_runner: CallbackRunner,
 ):
-    print("workers_setup_status", workers_setup_status)
     lit_spec = lit_api.spec
     loop: LitLoop = lit_api.loop
     stream = lit_api.stream
@@ -84,8 +83,6 @@ def inference_worker(
     lit_api.device = device
     callback_runner.trigger_event(EventTypes.AFTER_SETUP.value, lit_api=lit_api)
 
-    print(f"Setup complete for worker {f'{endpoint}_{worker_id}'}.")
-
     if workers_setup_status:
         workers_setup_status[f"{endpoint}_{worker_id}"] = WorkerSetupStatus.READY
 
diff --git a/src/litserve/server.py b/src/litserve/server.py
index 94e19d5d..7d59320c 100644
--- a/src/litserve/server.py
+++ b/src/litserve/server.py
@@ -19,6 +19,7 @@
 import logging
 import multiprocessing as mp
 import os
+import secrets
 import sys
 import threading
 import time
@@ -27,16 +28,14 @@
 from abc import ABC, abstractmethod
 from collections import deque
 from contextlib import asynccontextmanager
-from multiprocessing.context import Process
 from queue import Queue
-from threading import Thread
 from typing import Callable, Dict, Iterable, List, Literal, Optional, Sequence, Tuple, Union
 
 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 fastapi.security import APIKeyHeader, OAuth2PasswordBearer
 from starlette.formparsers import MultiPartParser
 from starlette.middleware.gzip import GZipMiddleware
 
@@ -58,6 +57,8 @@
 
 # if defined, it will require clients to auth with X-API-Key in the header
 LIT_SERVER_API_KEY = os.environ.get("LIT_SERVER_API_KEY")
+SHUTDOWN_API_KEY = os.environ.get("LIT_SHUTDOWN_API_KEY")
+oauth2_scheme = OAuth2PasswordBearer(tokenUrl="token")
 
 # FastAPI writes form files to disk over 1MB by default, which prevents serialization by multiprocessing
 MultiPartParser.max_file_size = sys.maxsize
@@ -357,6 +358,8 @@ def __init__(
         timeout: Union[float, bool] = 30,
         healthcheck_path: str = "/health",
         info_path: str = "/info",
+        shutdown_path: str = "/shutdown",
+        enable_shutdown_api: bool = False,
         model_metadata: Optional[dict] = None,
         spec: Optional[LitSpec] = None,
         max_payload_size=None,
@@ -395,6 +398,14 @@ 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".
+
+            enable_shutdown_api (bool, optional):
+                Enable the shutdown endpoint. If True, the server will listen for shutdown requests
+                at the specified path. Defaults to False.
+
             model_metadata (dict, optional):
                 Model metadata displayed at info endpoint (e.g., {"version": "1.0"}). Defaults to None.
 
@@ -496,6 +507,24 @@ def __init__(
                 "info_path must start with '/'. Please provide a valid api path like '/info', '/details', or '/v1/info'"
             )
 
+        if enable_shutdown_api and not shutdown_path.startswith("/"):
+            raise ValueError("shutdown_path must start with '/'. Please provide a valid api path like '/shutdown'")
+
+        global SHUTDOWN_API_KEY
+        if enable_shutdown_api and not SHUTDOWN_API_KEY:
+            SHUTDOWN_API_KEY = secrets.token_urlsafe(32)
+            logger.warning(
+                "LitServe's Shutdown API is enabled, but the `LIT_SHUTDOWN_API_KEY` environment variable is missing."
+                f"Generated shutdown API key: {SHUTDOWN_API_KEY}"
+            )
+        if enable_shutdown_api:
+            curl_command = (
+                "curl -X 'POST' 'http://localhost:8000/shutdown' "
+                "-H 'accept: application/json' "
+                f"-H 'Authorization: Bearer {SHUTDOWN_API_KEY}' "
+                "-d ''"
+            )
+            logger.info(f"To shutdown the server, run command: \n{curl_command}\n")
         try:
             json.dumps(model_metadata)
         except (TypeError, ValueError):
@@ -507,6 +536,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)
@@ -523,6 +553,7 @@ def __init__(
         self._logger_connector = _LoggerConnector(self, loggers)
         self.logger_queue = None
         self.lit_api = lit_api
+        self.enable_shutdown_api = enable_shutdown_api
         self.workers_per_device = workers_per_device
         self.max_payload_size = max_payload_size
         self.model_metadata = model_metadata
@@ -531,6 +562,8 @@ def __init__(
         self.use_zmq = fast_queue
         self.transport_config = None
         self.litapi_request_queues = {}
+        self._shutdown_event: Optional[mp.Event] = None
+        self.uvicorn_graceful_timeout = 30
 
         accelerator = self._connector.accelerator
         devices = self._connector.devices
@@ -542,7 +575,7 @@ 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.inference_workers_config = self.devices * self.workers_per_device
         self.transport_config = TransportConfig(transport_config="zmq" if self.use_zmq else "mp")
         self.register_endpoints()
 
@@ -557,7 +590,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]
 
@@ -601,7 +634,7 @@ async def lifespan(self, app: FastAPI):
         )
         task = loop.create_task(future, name=f"response_queue_to_buffer-{app.response_queue_id}")
         task.add_done_callback(
-            lambda _: logger.info(f"Response queue to buffer task terminated for consumer_id {app.response_queue_id}")
+            lambda _: logger.debug(f"Response queue to buffer task terminated for consumer_id {app.response_queue_id}")
         )
 
         try:
@@ -684,6 +717,17 @@ async def info(request: Request) -> Response:
                 }
             )
 
+        if self.enable_shutdown_api:
+
+            @self.app.post(self._shutdown_path, dependencies=[Depends(self.shutdown_api_key_auth)])
+            async def shutdown_endpoint():
+                if not self._shutdown_event:
+                    raise HTTPException(
+                        status_code=status.HTTP_503_SERVICE_UNAVAILABLE, detail="Server is still starting up"
+                    )
+                self._shutdown_event.set()
+                return Response(content="Server is initiating graceful shutdown.", status_code=status.HTTP_200_OK)
+
     def register_endpoints(self):
         self._register_internal_endpoints()
         for lit_api in self.litapi_connector:
@@ -771,9 +815,11 @@ 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()
+        manager = mp.Manager()
+        self.transport_config.manager = manager
         self.transport_config.num_consumers = num_api_servers
         self.workers_setup_status = manager.dict()
+        self._shutdown_event = manager.Event()
 
         # create request queues for each unique lit_api api_path
         for lit_api in self.litapi_connector:
@@ -785,6 +831,63 @@ def _init_manager(self, num_api_servers: int):
         self._transport = create_transport_from_config(self.transport_config)
         return manager
 
+    def _perform_graceful_shutdown(
+        self,
+        manager: mp.Manager,
+        uvicorn_workers: List[Union[mp.Process, threading.Thread]],
+        inference_workers: List[mp.Process],
+        shutdown_reason: str = "normal",
+    ):
+        """Encapsulates the graceful shutdown logic."""
+        logger.info("Shutting down LitServe...")
+
+        # Handle transport closure based on shutdown reason
+        if shutdown_reason == "keyboard_interrupt":
+            logger.debug("KeyboardInterrupt detected - skipping transport cleanup to avoid hanging")
+            self._transport.close(send_sentinel=False)
+        else:
+            self._transport.close(send_sentinel=True)
+
+        # terminate Uvicorn server workers tracked by LitServe (the master processes/threads)
+        if uvicorn_workers:
+            for i, uw in enumerate(uvicorn_workers):
+                uvicorn_pid = uw.ident if isinstance(uw, threading.Thread) else uw.pid
+                uvicorn_name = uw.name
+
+                log_prefix = f"{uvicorn_name} (PID: {uvicorn_pid})"
+
+                if not uw.is_alive():
+                    logger.warning(f"{log_prefix}: Already not alive.")
+                    continue
+                try:
+                    uw.terminate()
+                    uw.join(timeout=self.uvicorn_graceful_timeout)
+                    if uw.is_alive():
+                        logger.warning(f"{log_prefix}: Did not terminate gracefully. Forcibly killing.")
+                        uw.kill()
+                except Exception as e:
+                    logger.error(f"Error during termination of {log_prefix}: {e}")
+
+        # terminate and join inference worker processes
+        logger.debug(f"Terminating {len(inference_workers)} inference workers...")
+        for i, iw in enumerate(inference_workers):
+            worker_pid = iw.pid
+            worker_name = iw.name
+            try:
+                iw.terminate()
+                iw.join(timeout=5)
+                if iw.is_alive():
+                    logger.warning(
+                        f"Worker {worker_name} (PID: {worker_pid}): Did not terminate gracefully. Killing (SIGKILL)."
+                    )
+                    iw.kill()
+                else:
+                    logger.debug(f"Worker {worker_name} (PID: {worker_pid}): Terminated gracefully.")
+            except Exception as e:
+                logger.error(f"Error while terminating worker {worker_name} (PID: {worker_pid}): {e}")
+
+        manager.shutdown()
+
     def run(
         self,
         host: str = "0.0.0.0",
@@ -826,19 +929,19 @@ def run(
                 Additional uvicorn server options (ssl_keyfile, ssl_certfile, etc.). See uvicorn docs.
 
         Example:
-            >>> server.run()  # Basic
+        >>> server.run()  # Basic
 
-            >>> server.run(  # Production
-            ...     port=8080,
-            ...     num_api_servers=4,
-            ...     log_level="warning"
-            ... )
+        >>> 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
-            ... )
+        >>> server.run(  # Development
+        ...     log_level="debug",
+        ...     pretty_logs=True,
+        ...     generate_client_file=True
+        ... )
 
         """
         if generate_client_file:
@@ -862,7 +965,7 @@ def run(
         sockets = [config.bind_socket()]
 
         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")
@@ -883,26 +986,22 @@ def run(
             inference_workers.extend(_inference_workers)
 
         self.verify_worker_status()
+
+        shutdown_reason = "normal"
+        uvicorn_workers = []
         try:
             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()
+
+            self._shutdown_event.wait()
+
+        except KeyboardInterrupt:
+            logger.info("KeyboardInterrupt received. Initiating graceful shutdown.")
+            shutdown_reason = "keyboard_interrupt"
         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(manager, uvicorn_workers, inference_workers, shutdown_reason)
 
     def _prepare_app_run(self, app: FastAPI):
         # Add middleware to count active requests
@@ -921,17 +1020,25 @@ 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)
+            uvicorn_config = uvicorn.Config(
+                app=app,
+                host="0.0.0.0",
+                port=port,
+                log_level=log_level,
+                timeout_graceful_shutdown=self.uvicorn_graceful_timeout,
+                **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)
+                    sock.listen(uvicorn_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.workers = num_uvicorn_servers
+
+            server = uvicorn.Server(config=uvicorn_config)
             if uvicorn_worker_type == "process":
                 ctx = mp.get_context("fork")
                 w = ctx.Process(target=server.run, args=(sockets,), name=f"lit-uvicorn-{response_queue_id}")
@@ -949,3 +1056,11 @@ def setup_auth(self):
         if LIT_SERVER_API_KEY:
             return api_key_auth
         return no_auth
+
+    def shutdown_api_key_auth(self, shutdown_api_key: str = Depends(oauth2_scheme)):
+        if not SHUTDOWN_API_KEY or shutdown_api_key != SHUTDOWN_API_KEY:
+            raise HTTPException(
+                status_code=401,
+                detail="Invalid Bearer token for Shutdown API."
+                " Check that you are passing a correct 'Authorization: Bearer <SHUTDOWN_API_KEY>' in your header.",
+            )
diff --git a/src/litserve/transport/base.py b/src/litserve/transport/base.py
index 673a3eb9..578a7171 100644
--- a/src/litserve/transport/base.py
+++ b/src/litserve/transport/base.py
@@ -13,6 +13,6 @@ async def areceive(self, timeout: Optional[int] = None, consumer_id: Optional[in
         """Receive a message from model workers or any publisher."""
         pass
 
-    def close(self) -> None:
+    def close(self, **kwargs) -> None:
         """Clean up resources if needed (e.g., sockets, processes)."""
         pass
diff --git a/src/litserve/transport/process_transport.py b/src/litserve/transport/process_transport.py
index 170e9e3b..ef75f556 100644
--- a/src/litserve/transport/process_transport.py
+++ b/src/litserve/transport/process_transport.py
@@ -31,14 +31,16 @@ async def areceive(self, consumer_id: int, timeout: Optional[float] = None, bloc
                 raise
             return None
 
-    def close(self) -> None:
+    def close(self, send_sentinel: bool = True) -> None:
         # Mark the transport as closed
         self._closed = True
 
-        # Put sentinel values in the queues as a backup mechanism
-        for queue in self._queues:
-            with suppress(Exception):
-                queue.put(None)
+        # Only send sentinel values if requested and safe to do so
+        if send_sentinel:
+            # Put sentinel values in the queues as a backup mechanism
+            for queue in self._queues:
+                with suppress(Exception):
+                    queue.put(None, block=False)
 
     def __reduce__(self):
         return (MPQueueTransport, (None, self._queues))
diff --git a/src/litserve/transport/zmq_transport.py b/src/litserve/transport/zmq_transport.py
index 860e075e..1d76567d 100644
--- a/src/litserve/transport/zmq_transport.py
+++ b/src/litserve/transport/zmq_transport.py
@@ -32,7 +32,7 @@ async def areceive(self, consumer_id: Optional[int] = None, timeout=None) -> dic
             self.setup(zmq.SUB, consumer_id)
         return await self._zmq.get(timeout=timeout)
 
-    def close(self) -> None:
+    def close(self, **kwargs) -> None:
         if self._zmq:
             self._zmq.close()
         else:
diff --git a/src/litserve/utils.py b/src/litserve/utils.py
index 1e1fc397..e72b5632 100644
--- a/src/litserve/utils.py
+++ b/src/litserve/utils.py
@@ -84,6 +84,7 @@ def wrap_litserve_start(server: "LitServer"):
     try:
         yield server
     finally:
+        server._shutdown_event.set()
         # First close the transport to signal to the response_queue_to_buffer task that it should stop
         server._transport.close()
         for p in processes:


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index ccb4a3c9..7ea8d9dc 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -338,3 +338,33 @@ def openai_embedding_request_data():
 @pytest.fixture
 def openai_embedding_request_data_array():
     return {"input": ["A beautiful sunset over the beach."] * 4, "model": "lit", "encoding_format": "float"}
+
+
+class MockEvent:
+    def set(self):
+        pass
+
+    def wait(self):
+        pass
+
+
+class MockQueue:
+    def put(self, item):
+        pass
+
+    def get(self, block=True, timeout=None):
+        pass
+
+
+@pytest.fixture
+def mock_manager():
+    class MockManager:
+        def __init__(self):
+            self.Queue = MockQueue
+            self.Event = MockEvent
+            self.dict = dict
+
+        def shutdown(self):
+            pass
+
+    return MockManager()
diff --git a/tests/test_lit_server.py b/tests/test_lit_server.py
index e8143e4d..669cf156 100644
--- a/tests/test_lit_server.py
+++ b/tests/test_lit_server.py
@@ -199,7 +199,7 @@ def test_mocked_accelerator():
 
 
 @patch("litserve.server.uvicorn")
-def test_server_run(mock_uvicorn):
+def test_server_run(mock_uvicorn, mock_manager):
     server = LitServer(SimpleLitAPI())
     server.verify_worker_status = MagicMock()
     with pytest.raises(ValueError, match="port must be a value from 1024 to 65535 but got"):
@@ -211,12 +211,20 @@ def test_server_run(mock_uvicorn):
     with pytest.raises(ValueError, match="host must be '0.0.0.0', '127.0.0.1', or '::' but got"):
         server.run(host="127.0.0.2")
 
-    server.run(port=8000)
+    # test port 8000
+    with patch("litserve.server.mp.Manager", return_value=mock_manager):
+        server.run(port=8000)
     mock_uvicorn.Config.assert_called()
     mock_uvicorn.reset_mock()
-    server.run(port="8001")
+
+    # test port 8001
+    with patch("litserve.server.mp.Manager", return_value=mock_manager):
+        server.run(port="8001")
     mock_uvicorn.Config.assert_called()
-    server.run(host="::", port="8000")
+
+    # test host "::" and port 8000
+    with patch("litserve.server.mp.Manager", return_value=mock_manager):
+        server.run(host="::", port="8000")
     mock_uvicorn.Config.assert_called()
 
 
@@ -242,10 +250,11 @@ def server_for_api_worker_test(simple_litapi):
 
 @pytest.mark.skipif(sys.platform == "win32", reason="Test is only for Unix")
 @patch("litserve.server.uvicorn")
-def test_server_run_with_api_server_worker_type(mock_uvicorn, server_for_api_worker_test):
+def test_server_run_with_api_server_worker_type(mock_uvicorn, server_for_api_worker_test, mock_manager):
     server = server_for_api_worker_test
 
-    server.run(api_server_worker_type="process", num_api_servers=10)
+    with patch("litserve.server.mp.Manager", return_value=mock_manager):
+        server.run(api_server_worker_type="process", num_api_servers=10)
     server.launch_inference_worker.assert_called_with(server.lit_api)
 
 
@@ -253,11 +262,12 @@ def test_server_run_with_api_server_worker_type(mock_uvicorn, server_for_api_wor
 @pytest.mark.parametrize(("api_server_worker_type", "num_api_workers"), [(None, 1), ("process", 1)])
 @patch("litserve.server.uvicorn")
 def test_server_run_with_process_api_worker(
-    mock_uvicorn, api_server_worker_type, num_api_workers, server_for_api_worker_test
+    mock_uvicorn, api_server_worker_type, num_api_workers, server_for_api_worker_test, mock_manager
 ):
     server = server_for_api_worker_test
 
-    server.run(api_server_worker_type=api_server_worker_type, num_api_workers=num_api_workers)
+    with patch("litserve.server.mp.Manager", return_value=mock_manager):
+        server.run(api_server_worker_type=api_server_worker_type, num_api_workers=num_api_workers)
     server.launch_inference_worker.assert_called_with(server.lit_api)
     actual = server._start_server.call_args
     assert actual[0][4] == "process", "Server should run in process mode"
@@ -266,9 +276,10 @@ def test_server_run_with_process_api_worker(
 
 @pytest.mark.skipif(sys.platform == "win32", reason="Test is only for Unix")
 @patch("litserve.server.uvicorn")
-def test_server_run_with_thread_api_worker(mock_uvicorn, server_for_api_worker_test):
+def test_server_run_with_thread_api_worker(mock_uvicorn, server_for_api_worker_test, mock_manager):
     server = server_for_api_worker_test
-    server.run(api_server_worker_type="thread")
+    with patch("litserve.server.mp.Manager", return_value=mock_manager):
+        server.run(api_server_worker_type="thread")
     server.launch_inference_worker.assert_called_with(server.lit_api)
     assert server._start_server.call_args[0][4] == "thread", "Server should run in thread mode"
     mock_uvicorn.Config.assert_called()
@@ -287,7 +298,7 @@ def test_server_run_with_invalid_api_worker(simple_litapi):
 
 @pytest.mark.skipif(sys.platform != "win32", reason="Test is only for Windows")
 @patch("litserve.server.uvicorn")
-def test_server_run_windows(mock_uvicorn):
+def test_server_run_windows(mock_uvicorn, mock_manager):
     api = ls.test_examples.SimpleLitAPI()
     server = ls.LitServer(api)
     server.verify_worker_status = MagicMock()
@@ -295,7 +306,8 @@ def test_server_run_windows(mock_uvicorn):
     server._transport = MagicMock()
     server._start_server = MagicMock()
 
-    server.run(api_server_worker_type=None)
+    with patch("litserve.server.mp.Manager", return_value=mock_manager):
+        server.run(api_server_worker_type=None)
     actual = server._start_server.call_args
     assert actual[0][4] == "thread", "Windows only supports thread mode"
 
diff --git a/tests/test_simple.py b/tests/test_simple.py
index 6bbb6d5a..169fc8bb 100644
--- a/tests/test_simple.py
+++ b/tests/test_simple.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 import asyncio
+import random
 import time
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import ExitStack
@@ -19,10 +20,11 @@
 import numpy as np
 import pytest
 from asgi_lifespan import LifespanManager
-from fastapi import Request, Response
+from fastapi import Request, Response, status
 from fastapi.testclient import TestClient
 from httpx import ASGITransport, AsyncClient
 
+import litserve.server
 from litserve import LitAPI, LitServer
 from litserve.utils import wrap_litserve_start
 
@@ -170,6 +172,51 @@ def test_load(lit_server):
             assert el == {"output": i**2}
 
 
+@pytest.fixture
+def shutdown_api_key():
+    api_key = f"test-key-{random.randint(100, 1000000)}"
+    litserve.server.SHUTDOWN_API_KEY = api_key
+    yield api_key
+    litserve.server.SHUTDOWN_API_KEY = None
+
+
+def test_shutdown_endpoint_single_worker(shutdown_api_key):
+    server = LitServer(
+        SimpleLitAPI(),
+        accelerator="cpu",
+        devices=1,
+        workers_per_device=1,
+        enable_shutdown_api=True,
+    )
+
+    with wrap_litserve_start(server) as srv, TestClient(srv.app) as client:
+        response_no_header = client.post("/shutdown")
+        assert response_no_header.status_code == 401
+
+        response_correct_key = client.post("/shutdown", headers={"Authorization": f"Bearer {shutdown_api_key}"})
+        assert response_correct_key.status_code == status.HTTP_200_OK
+
+
+def test_shutdown_endpoint_multiple_workers(shutdown_api_key):
+    server = LitServer(
+        SimpleLitAPI(),
+        accelerator="cpu",
+        devices=1,
+        workers_per_device=3,
+        enable_shutdown_api=True,
+    )
+
+    with wrap_litserve_start(server) as srv, TestClient(srv.app) as client:
+        response_no_header = client.post("/shutdown")
+        assert response_no_header.status_code == 401
+
+        response_wrong_key = client.post("/shutdown", headers={"Authorization": "Bearer wrong_key"})
+        assert response_wrong_key.status_code == 401
+
+        response_correct_key = client.post("/shutdown", headers={"Authorization": f"Bearer {shutdown_api_key}"})
+        assert response_correct_key.status_code == status.HTTP_200_OK
+
+
 class SlowLitAPI(LitAPI):
     def setup(self, device):
         self.model = lambda x: x**2