fix(proxy): bound SIGTERM connection wait by GRACEFUL_SHUTDOWN_TIMEOUT (#46186)

Resolves LIT-9448

Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

helm/litellm-helm/values.yaml

@@ -440,10 +440,10 @@ db:
 # Prefer the native /health/drain preStop hook over a fixed `sleep`: it marks
 # the pod NotReady and blocks only until in-flight requests actually finish
 # (bounded by GRACEFUL_SHUTDOWN_TIMEOUT, default 30s), instead of always
-# waiting the worst-case duration. The drain runs once (the preStop hook and
-# the SIGTERM handler share it), so set terminationGracePeriodSeconds a few
-# seconds above GRACEFUL_SHUTDOWN_TIMEOUT to leave room for teardown before
-# SIGKILL.
+# waiting the worst-case duration. After SIGTERM the server waits up to
+# GRACEFUL_SHUTDOWN_TIMEOUT again for connections still open and then cancels
+# them, so set terminationGracePeriodSeconds a few seconds above twice
+# GRACEFUL_SHUTDOWN_TIMEOUT to leave room for teardown before SIGKILL.
 #
 # /health/drain is off by default; enable it with
 # general_settings.enable_drain_endpoint: true. The kubelet calls preStop

litellm/constants.py

@@ -104,6 +104,7 @@ S3_PARTITION_GRANULARITY_ENV_VAR: Final = "S3_PARTITION_GRANULARITY"
 MAX_FILE_LIST_LIMIT: Final = 10000
 DEFAULT_SQS_FLUSH_INTERVAL_SECONDS: Final = int(os.getenv("DEFAULT_SQS_FLUSH_INTERVAL_SECONDS", 10))
 DEFAULT_NUM_WORKERS_LITELLM_PROXY: Final = int(os.getenv("DEFAULT_NUM_WORKERS_LITELLM_PROXY", 1))
+GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS: Final = 10
 budget_reservation_disabled_info_emitted = False
 DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE = int(os.getenv("DYNAMIC_RATE_LIMIT_ERROR_THRESHOLD_PER_MINUTE", 1))
 DEFAULT_SQS_BATCH_SIZE: Final = int(os.getenv("DEFAULT_SQS_BATCH_SIZE", 512))

litellm/proxy/proxy_cli.py

@@ -308,12 +308,14 @@ class ProxyInitializationHelpers:
 
         import litellm
         from litellm._logging import _get_uvicorn_json_log_config, resolve_log_level
+        from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
 
         uvicorn_args: Final = {
             "app": "litellm.proxy.proxy_server:app",
             "host": host,
             "port": port,
             "server_header": False,
+            "timeout_graceful_shutdown": GracefulShutdownManager.get_timeout(),
         }
         if log_config is not None:
             print(f"Using log_config: {log_config}")
@@ -571,6 +573,12 @@ class ProxyInitializationHelpers:
         """
         Run litellm with `gunicorn`
         """
+        import math
+
+        from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS
+        from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
+        from litellm.proxy.shutdown.graceful_uvicorn_worker import build_graceful_shutdown_uvicorn_worker
+
         if os.name == "nt":
             pass
         else:
@@ -627,10 +635,12 @@ class ProxyInitializationHelpers:
                 return self.application
 
         print(f"\033[1;32mLiteLLM Proxy: Starting server on {host}:{port} with {num_workers} workers\033[0m\n")
+        graceful_shutdown_timeout: Final = GracefulShutdownManager.get_timeout()
         gunicorn_options: Final = {
             "bind": f"{host}:{port}",
             "workers": num_workers,  # default is 1
-            "worker_class": "uvicorn.workers.UvicornWorker",
+            "worker_class": build_graceful_shutdown_uvicorn_worker(graceful_shutdown_timeout),
+            "graceful_timeout": math.ceil(graceful_shutdown_timeout) + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS,
             "preload": True,  # Add the preload flag,
             "accesslog": "-",  # Log to stdout
             "timeout": 600,  # default to very high number, bedrock/anthropic.claude-v2:1 can take 30+ seconds for the 1st chunk to come in

litellm/proxy/shutdown/graceful_shutdown_manager.py

@@ -2,9 +2,10 @@
 Application-level graceful shutdown coordination for the LiteLLM proxy.
 
 Kubernetes terminates a pod by sending ``SIGTERM`` and, after
-``terminationGracePeriodSeconds``, ``SIGKILL``. By default LiteLLM delegates
-the signal to uvicorn and tears down immediately, dropping any in-flight
-requests (streaming, batch inference, long-lived calls).
+``terminationGracePeriodSeconds``, ``SIGKILL``. On ``SIGTERM`` uvicorn stops
+accepting connections and waits up to ``GRACEFUL_SHUTDOWN_TIMEOUT`` for open
+ones, cancelling whatever is still running, before the lifespan teardown that
+runs this drain.
 
 A fixed ``preStop`` sleep can not solve this: it has to be sized for the
 *worst-case* request, so it either wastes time on every routine shutdown or is

litellm/proxy/shutdown/graceful_uvicorn_worker.py

@@ -0,0 +1,10 @@
+from uvicorn.workers import UvicornWorker
+
+
+def build_graceful_shutdown_uvicorn_worker(timeout_graceful_shutdown: float) -> type[UvicornWorker]:
+    class GracefulShutdownUvicornWorker(UvicornWorker):
+        def __init__(self, *args: object) -> None:
+            self.CONFIG_KWARGS = {**UvicornWorker.CONFIG_KWARGS, "timeout_graceful_shutdown": timeout_graceful_shutdown}
+            super().__init__(*args)
+
+    return GracefulShutdownUvicornWorker

tests/unit/proxy/test_proxy_cli.py

@@ -1,4 +1,5 @@
 import inspect
+import logging
 import os
 from contextlib import nullcontext
 from pathlib import Path
@@ -22,8 +23,10 @@ import yaml
 from uvicorn.config import LOOP_FACTORIES
 from uvicorn.importer import import_from_string
 
+from litellm.constants import GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS
 from litellm.proxy import proxy_cli
 from litellm.proxy.proxy_cli import ProxyInitializationHelpers, run_server
+from litellm.proxy.shutdown.graceful_shutdown_manager import DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT
 
 
 @pytest.fixture(autouse=True)
@@ -153,6 +156,22 @@ class TestProxyInitializationHelpers:
             )
             assert args["timeout_worker_healthcheck"] == 15
 
+    @pytest.mark.parametrize(
+        ("raw_timeout", "expected"),
+        [("7.5", 7.5), ("0", 0.0), (None, DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT)],
+    )
+    def test_uvicorn_graceful_shutdown_is_bounded_by_graceful_shutdown_timeout(
+        self, monkeypatch, raw_timeout, expected
+    ):
+        if raw_timeout is None:
+            monkeypatch.delenv("GRACEFUL_SHUTDOWN_TIMEOUT", raising=False)
+        else:
+            monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", raw_timeout)
+
+        args = ProxyInitializationHelpers._get_default_unvicorn_init_args("localhost", 8000)
+
+        assert uvicorn.Config(**args).timeout_graceful_shutdown == expected
+
     @staticmethod
     def _uvicorn_access_info_enabled(args: dict) -> bool:
         import logging
@@ -2050,6 +2069,49 @@ class TestProxyInitializationHelpers:
         assert captured["options"]["max_requests"] == 1000
         assert captured["options"]["max_requests_jitter"] == 50
 
+    @pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows")
+    def test_gunicorn_workers_bound_graceful_shutdown_by_graceful_shutdown_timeout(self, monkeypatch):
+        pytest.importorskip("gunicorn")
+        import gunicorn.config
+        import gunicorn.glogging
+        from uvicorn.workers import UvicornWorker
+
+        monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "7.5")
+        captured: dict = {}
+
+        def capture_run(self):
+            captured["options"] = dict(self.options)
+
+        with (
+            patch("gunicorn.app.base.BaseApplication.run", capture_run),
+            patch(  # test-quality-ok: option tests must not start a thread or change the pytest worker's child ownership
+                "litellm.proxy.proxy_cli.start_query_engine_reaper"
+            ),
+        ):
+            ProxyInitializationHelpers._run_gunicorn_server(
+                host="127.0.0.1",
+                port=4011,
+                app=MagicMock(),
+                num_workers=2,
+                ssl_certfile_path=None,
+                ssl_keyfile_path=None,
+            )
+
+        for name in ("uvicorn.error", "uvicorn.access"):
+            uvicorn_logger = logging.getLogger(name)
+            for attr in ("handlers", "level", "propagate"):
+                monkeypatch.setattr(uvicorn_logger, attr, getattr(uvicorn_logger, attr))
+        gunicorn_cfg = gunicorn.config.Config()
+        worker = captured["options"]["worker_class"](
+            0, os.getpid(), [], MagicMock(), 30, gunicorn_cfg, gunicorn.glogging.Logger(gunicorn_cfg)
+        )
+
+        assert isinstance(worker, UvicornWorker)
+        assert worker.config.timeout_graceful_shutdown == 7.5
+        assert worker.config.loop == UvicornWorker.CONFIG_KWARGS["loop"]
+        assert worker.config.http == UvicornWorker.CONFIG_KWARGS["http"]
+        assert captured["options"]["graceful_timeout"] == 8 + GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS
+
     @pytest.mark.skipif(os.name == "nt", reason="gunicorn server path skips Windows")
     def test_gunicorn_master_is_reserved_for_forking_before_it_runs(self, fork_reservation):
         """preload forks workers from the master, so native routes are forbidden there first"""