fix: share shutdown callback budget with HTTP draining

helm/litellm-helm/README.md

@@ -259,3 +259,9 @@ itself.
 `lensWorker.mode` selects `disabled`, `bundled` or `external`. Bundled mode generates separate service and identity-signing credentials. External mode requires `lensWorker.gateway.secretName` and `lensWorker.serviceTokenSecret.name`; their keys default to `gateway-secret` and `service-token`. Provision distinct values matching the external Lens deployment
 
 Follow the [Lens Helm connection guide](https://github.com/BerriAI/lens/blob/main/helm/lens/README.md#connect-a-gateway) for complete values, verification, GitOps credential requirements and existing-data migration. Source-chart installation requires a built Lens image until the first signed Lens release is published
+
+## Shutdown accounting
+
+`GRACEFUL_SHUTDOWN_TIMEOUT` defaults to 30 seconds. Pending success callbacks share the time remaining since SIGTERM or the first `/health/drain` call. Callback flushing reserves the existing 10-second teardown headroom for cancellation, tracked accounting tasks and persistence. Exhausting this budget can discard callbacks, which is logged
+
+Set `terminationGracePeriodSeconds` to at least `GRACEFUL_SHUTDOWN_TIMEOUT + 10`, allowing additional time for preStop hooks and database writes. This is not a hard shutdown bound or a guarantee of persistence. Database retries retain their own limits, and cancellation-resistant callbacks can remain alive. The existing preStop hook can also be followed by Uvicorn's separate connection drain, as described in `values.yaml`

litellm/litellm_core_utils/logging_worker.py

@@ -473,12 +473,19 @@ class LoggingWorker:
         self.start()
         self.enqueue(async_coroutine)
 
-    async def stop(self, *, discard_pending: bool = False) -> None:
+    @property
+    def pending_count(self) -> int:
+        return (self._queue.qsize() if self._queue is not None else 0) + len(self._dequeued_tasks)
+
+    async def stop(self, *, discard_pending: bool = False, timeout: float | None = None) -> None:
         """Stop the logging worker and clean up resources."""
         if discard_pending and self._queue is not None:
-            for task in self._drain_pending(self._queue):
+            discarded: Final = self._drain_pending(self._queue)
+            for task in discarded:
                 task["coroutine"].close()
                 self._queue.task_done()
+            if discarded:
+                verbose_logger.warning("LoggingWorker shutdown: %d queued callbacks discarded", len(discarded))
         if self._timeout_summary_task is not None:
             self._timeout_summary_task.cancel()
             self._timeout_summary_task = None
@@ -498,16 +505,20 @@ class LoggingWorker:
             task.cancel()
 
         # Wait for cancellation to settle; ignore errors raised during shutdown.
-        await asyncio.gather(*tasks_to_cancel, return_exceptions=True)
+        done, pending = await asyncio.wait(tasks_to_cancel, timeout=timeout)
+        await asyncio.gather(*done, return_exceptions=True)
+        if pending:
+            verbose_logger.warning("LoggingWorker shutdown: %d tasks still alive after cancellation", len(pending))
 
         if discard_pending and self._queue is not None:
             for task in self._unstarted_dequeued_tasks():
                 task["coroutine"].close()
                 self._untrack_dequeued(task)
                 self._queue.task_done()
-        self._worker_task = None
+        if self._worker_task is not None and self._worker_task.done():
+            self._worker_task = None
         # Drop references to completed tasks so we can restart cleanly.
-        self._running_tasks.clear()
+        self._running_tasks.difference_update(done)
 
     async def flush(self) -> None:
         """Flush the logging queue.

litellm/proxy/db/db_spend_update_writer.py

@@ -1900,7 +1900,6 @@ class DBSpendUpdateWriter:
             proxy_logging_obj=proxy_logging_obj,
         )
 
-        # NOTE: Daily tag spend is committed by a separate scheduler job.
 
         ################## Daily End-User Spend Update Transactions ##################
         # Aggregate all in memory daily end-user spend transactions and commit to db

litellm/proxy/proxy_server.py

@@ -79,6 +79,7 @@ from litellm.constants import (
     DEFAULT_SHARED_HEALTH_CHECK_LOCK_TTL,
     DEFAULT_SHARED_HEALTH_CHECK_TTL,
     DEFAULT_SLACK_ALERTING_THRESHOLD,
+    GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS,
     INTERNAL_CALLBACK_NAMES,
     LANGFUSE_SHUTDOWN_FLUSH_TIMEOUT_MILLIS,
     LITELLM_EMBEDDING_PROVIDERS_SUPPORTING_INPUT_ARRAY_OF_TOKENS,
@@ -98,6 +99,7 @@ from litellm.litellm_core_utils.litellm_logging import (
     _init_custom_logger_compatible_class,
     initialize_sentry,
 )
+from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER, LoggingWorker
 from litellm.litellm_core_utils.safe_json_dumps import safe_dumps
 from litellm.litellm_core_utils.safe_json_loads import safe_json_loads
 from litellm.litellm_core_utils.token_counter import UnsupportedContentTypeError, offload_token_count
@@ -1278,21 +1280,29 @@ async def _flush_spend_on_shutdown() -> None:
 
 async def _flush_spend_around_logging_drain_on_shutdown(
     flush_spend: Callable[[], Awaitable[None]] = _flush_spend_on_shutdown,
+    *,
+    worker: LoggingWorker = GLOBAL_LOGGING_WORKER,
+    clock: Callable[[], float] = time.monotonic,
+    drain_tasks: Callable[[float], Awaitable[None]] = drain_logging_tasks,
 ) -> None:
-    from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
-
     await flush_spend()
-    verbose_proxy_logger.info("Draining success callbacks before the final spend flush")
+    remaining: Final = GracefulShutdownManager.remaining(clock=clock)
+    callback_timeout: Final = max(0.0, remaining - GRACEFUL_SHUTDOWN_TEARDOWN_HEADROOM_SECONDS)
+    verbose_proxy_logger.info(
+        "Draining success callbacks before the final spend flush remaining_s=%.3f callback_timeout_s=%.3f",
+        remaining,
+        callback_timeout,
+    )
     try:
-        await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=GracefulShutdownManager.get_timeout())
+        await asyncio.wait_for(worker.flush(), timeout=callback_timeout)
     except asyncio.TimeoutError:
-        verbose_proxy_logger.warning("Timed out draining success callbacks on shutdown; accounting may be incomplete")
+        verbose_proxy_logger.warning(
+            "Success callback shutdown budget exhausted with %d unfinished callbacks; accounting may be incomplete",
+            worker.pending_count,
+        )
     finally:
-        try:
-            await asyncio.wait_for(GLOBAL_LOGGING_WORKER.stop(discard_pending=True), timeout=5.0)
-        except asyncio.TimeoutError:
-            verbose_proxy_logger.warning("Timed out cancelling success callbacks on shutdown")
-    await drain_logging_tasks()
+        await worker.stop(discard_pending=True, timeout=GracefulShutdownManager.remaining(2.0, clock=clock))
+    await drain_tasks(GracefulShutdownManager.remaining(5.0, clock=clock))
     await flush_spend()
 
 
@@ -1455,6 +1465,13 @@ async def _connect_to_count_stored_values() -> SupportsRawQueries:
 
 @asynccontextmanager
 async def proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]:
+    with GracefulShutdownManager.capture_sigterm():
+        async with _proxy_startup_event(app) as state:
+            yield state
+
+
+@asynccontextmanager
+async def _proxy_startup_event(app: FastAPI) -> AsyncGenerator[ProxyLifespanState, None]:
     global \
         prisma_client, \
         master_key, \

litellm/proxy/shutdown/graceful_shutdown_manager.py

@@ -5,7 +5,8 @@ Kubernetes terminates a pod by sending ``SIGTERM`` and, after
 ``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.
+runs this drain. The lifespan registers a SIGTERM wrapper before serving
+requests so callback draining can subtract the preceding connection 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
@@ -20,8 +21,12 @@ granularity of ``InFlightRequestsMiddleware``.
 
 import asyncio
 import os
+import signal
+import threading
 import time
-from collections.abc import Callable
+from collections.abc import Callable, Iterator
+from contextlib import contextmanager
+from types import FrameType
 from typing import Final
 
 from litellm._logging import verbose_proxy_logger
@@ -44,6 +49,7 @@ class GracefulShutdownManager:
     _is_shutting_down: bool = False
     _shutdown_started_at: float | None = None
     _drain_performed: bool = False
+    _sigterm_handler: Callable[[int, FrameType | None], None] | None = None
 
     @classmethod
     def is_shutting_down(cls) -> bool:
@@ -71,7 +77,7 @@ class GracefulShutdownManager:
             return DEFAULT_GRACEFUL_SHUTDOWN_TIMEOUT
 
     @classmethod
-    def start_shutdown(cls) -> None:
+    def start_shutdown(cls, clock: Callable[[], float] = time.monotonic) -> None:
         """
         Mark the worker as draining. Idempotent — repeated calls (e.g. SIGTERM
         followed by a preStop hit on /health/drain) do not reset the clock.
@@ -79,12 +85,42 @@ class GracefulShutdownManager:
         if cls._is_shutting_down:
             return
         cls._is_shutting_down = True
-        cls._shutdown_started_at = time.monotonic()
+        cls._shutdown_started_at = clock()
         verbose_proxy_logger.info(
             "graceful_shutdown_started in_flight_requests=%s",
             get_in_flight_requests(),
         )
 
+    @classmethod
+    def remaining(cls, cap: float | None = None, *, clock: Callable[[], float] = time.monotonic) -> float:
+        elapsed: Final = 0.0 if cls._shutdown_started_at is None else clock() - cls._shutdown_started_at
+        remaining: Final = max(0.0, cls.get_timeout() - elapsed)
+        return remaining if cap is None else min(max(0.0, cap), remaining)
+
+    @classmethod
+    @contextmanager
+    def capture_sigterm(cls, clock: Callable[[], float] = time.monotonic) -> Iterator[None]:
+        if threading.current_thread() is not threading.main_thread():
+            yield
+            return
+        previous: Final = signal.getsignal(signal.SIGTERM)
+        if not callable(previous) or previous is cls._sigterm_handler:
+            yield
+            return
+
+        def handle_sigterm(signum: int, frame: FrameType | None) -> None:
+            cls.start_shutdown(clock)
+            previous(signum, frame)
+
+        cls._sigterm_handler = handle_sigterm
+        signal.signal(signal.SIGTERM, handle_sigterm)
+        try:
+            yield
+        finally:
+            if signal.getsignal(signal.SIGTERM) is handle_sigterm:
+                signal.signal(signal.SIGTERM, previous)
+            cls._sigterm_handler = None
+
     @classmethod
     async def wait_for_drain(
         cls,
@@ -93,13 +129,14 @@ class GracefulShutdownManager:
         count_fn: Callable[[], int] | None = None,
         poll_interval: float = _DRAIN_POLL_INTERVAL,
         log_interval: float = _DRAIN_LOG_INTERVAL,
+        clock: Callable[[], float] = time.monotonic,
     ) -> int:
         """
         Poll the in-flight request counter until it reaches the drain target or
         ``timeout`` seconds elapse.
 
         Args:
-            timeout: Max seconds to wait. Defaults to ``get_timeout()``.
+            timeout: Max seconds to wait. Defaults to the remaining shutdown budget.
             exclude_self: When the caller is itself an in-flight HTTP request
                 (the /health/drain endpoint), set this so the caller's own
                 request is not counted as outstanding work.
@@ -120,7 +157,7 @@ class GracefulShutdownManager:
         cls._drain_performed = True
 
         if timeout is None:
-            timeout = cls.get_timeout()
+            timeout = cls.remaining(clock=clock)
         if count_fn is None:
             count_fn = get_in_flight_requests
 
@@ -128,7 +165,7 @@ class GracefulShutdownManager:
         # and so counts itself; treat <=1 as "drained" in that case.
         target: Final = 1 if exclude_self else 0
 
-        start: Final = time.monotonic()
+        start: Final = clock()
         initial: Final = count_fn()
         last_log = start
 
@@ -142,11 +179,11 @@ class GracefulShutdownManager:
                 verbose_proxy_logger.info(
                     "graceful_shutdown_complete drained_requests=%s elapsed_s=%.2f",
                     drained,
-                    time.monotonic() - start,
+                    clock() - start,
                 )
                 return drained
 
-            elapsed = time.monotonic() - start
+            elapsed = clock() - start
             if elapsed >= timeout:
                 verbose_proxy_logger.warning(
                     "graceful_shutdown_timeout in_flight_requests=%s elapsed_s=%.2f "
@@ -157,7 +194,7 @@ class GracefulShutdownManager:
                 )
                 return max(0, initial - current)
 
-            now = time.monotonic()
+            now = clock()
             if now - last_log >= log_interval:
                 verbose_proxy_logger.info(
                     "drain_waiting in_flight_requests=%s elapsed_s=%.2f",

litellm/responses/streaming_iterator.py

@@ -222,7 +222,7 @@ async def drain_logging_tasks(timeout: float = 5.0) -> None:
         return
     _, pending = await asyncio.wait(tuple(_PENDING_LOGGING_TASKS), timeout=timeout)
     if pending:
-        verbose_logger.warning("Timed out draining %d logging tasks on shutdown, dropping them", len(pending))
+        verbose_logger.warning("Timed out draining logging tasks on shutdown: %d tasks still alive", len(pending))
 
 
 _ERROR_CODE_HTTP_STATUS: Final[Mapping[str, int]] = MappingProxyType(

tests/integration/spend/test_shutdown_flush.py

@@ -1,12 +1,14 @@
 import hashlib
 import json
 import os
+import re
 import threading
+import time
 import uuid
 from collections.abc import Callable, Iterator
 from concurrent.futures import ThreadPoolExecutor
 from contextlib import contextmanager
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from pathlib import Path
 from typing import Final
 
@@ -29,6 +31,7 @@ BURST_REQUESTS: Final = 30
 
 _GATED_LOGGER: Final = """
 import os
+import re
 import httpx
 from litellm.integrations.custom_logger import CustomLogger
 
@@ -139,7 +142,10 @@ def test_callback_backlog_is_accounted_before_shutdown_disconnects(
 ) -> None:
     release: Final = threading.Event()
     tags: Final = (f"backlog-a-{uuid.uuid4().hex}", f"backlog-b-{uuid.uuid4().hex}")
-    with wire_server(_provider) as wire, wire_server(lambda _: Reply(chunks=(b"{}",), gate_before_first=release)) as logger:
+    with (
+        wire_server(_provider) as wire,
+        wire_server(lambda _: Reply(chunks=(b"{}",), gate_before_first=release)) as logger,
+    ):
         config: Final = _tag_shutdown_config(tmp_path, wire.url, redis_buffer, gated=True)
         with owned_proxy_process(
             gateway,
@@ -176,6 +182,70 @@ def test_callback_backlog_is_accounted_before_shutdown_disconnects(
                 release.set()
 
 
+@pytest.mark.parametrize("gunicorn", [False, True])
+def test_sigterm_http_drain_consumes_callback_budget(gateway: Gateway, tmp_path: Path, gunicorn: bool) -> None:
+    release_request: Final = threading.Event()
+    release_callback: Final = threading.Event()
+    tags: Final = (f"deadline-a-{uuid.uuid4().hex}", f"deadline-b-{uuid.uuid4().hex}")
+
+    def held_provider(request: Request) -> Reply:
+        response: Final = _provider(request)
+        return replace(response, chunks=(response.body,), gate_before_first=release_request)
+
+    with (
+        wire_server(held_provider) as wire,
+        wire_server(lambda _: Reply(chunks=(b"{}",), gate_before_first=release_callback)) as logger,
+    ):
+        config: Final = _tag_shutdown_config(tmp_path, wire.url, False, gated=True)
+        with (
+            owned_proxy_process(
+                gateway,
+                tmp_path,
+                {
+                    "DATABASE_URL": os.environ["DATABASE_URL"],
+                    "SHUTDOWN_CALLBACK_URL": logger.url,
+                    "GRACEFUL_SHUTDOWN_TIMEOUT": "30",
+                    "LITELLM_LOG": "DEBUG",
+                },
+                config=config,
+                extra_arguments=("--run_gunicorn",) if gunicorn else (),
+            ) as owned,
+            ThreadPoolExecutor(max_workers=1) as executor,
+        ):
+            key: Final = string_value(
+                owned.gateway.post("/key/generate", {"user_id": f"deadline-{uuid.uuid4().hex}"})["key"]
+            )
+            request: Final = executor.submit(_tagged_chat, owned.gateway, key, tags)
+            try:
+                eventually(wire.received.qsize, lambda count: count == 1, seconds=10)
+                owned.begin_shutdown()
+                eventually(
+                    lambda: "graceful_shutdown_started" in owned.log.read_text(), lambda stamped: stamped, seconds=5
+                )
+                assert not request.done(), "The shutdown clock must start before the HTTP request drains"
+                stamped: Final = time.monotonic()
+                eventually(time.monotonic, lambda now: now - stamped >= 2, seconds=5)
+                held_seconds: Final = time.monotonic() - stamped
+                release_request.set()
+                request.result(timeout=10)
+                budget: Final = eventually(
+                    lambda: re.search(r"remaining_s=([\d.]+) callback_timeout_s=([\d.]+)", owned.log.read_text()),
+                    lambda match: match is not None,
+                    seconds=15,
+                )
+                assert budget is not None
+                remaining: Final = float(budget[1])
+                allowance: Final = float(budget[2])
+                assert remaining <= 30 - held_seconds + 0.05
+                assert allowance == pytest.approx(max(0, remaining - 10), abs=0.002)
+                release_callback.set()
+                owned.process.wait(timeout=45)
+                _assert_shutdown_accounting(key, tags, 1)
+            finally:
+                release_request.set()
+                release_callback.set()
+
+
 def _api_requests(table: str, column: str, identity: str) -> int:
     rows: Final = read_rows(
         f'SELECT coalesce(sum(api_requests), 0)::int AS total FROM "{table}" WHERE {column}=%s', (identity,)

tests/unit/proxy/shutdown/test_graceful_shutdown_manager.py

@@ -179,3 +179,107 @@ async def test_emits_periodic_drain_waiting_log_while_waiting():
         timeout=10, count_fn=count_fn, poll_interval=0, log_interval=0
     )
     assert drained == 3
+
+
+@pytest.mark.parametrize("elapsed, cap, expected", [(4, None, 26), (25, 2, 2), (29, 5, 1), (31, None, 0)])
+def test_remaining_budget_includes_time_before_lifespan(
+    monkeypatch: pytest.MonkeyPatch, elapsed: float, cap: float | None, expected: float
+) -> None:
+    monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30")
+    GracefulShutdownManager.start_shutdown(clock=lambda: 100.0)
+    GracefulShutdownManager.start_shutdown(clock=lambda: 100.0 + elapsed)
+    assert GracefulShutdownManager.remaining(cap, clock=lambda: 100.0 + elapsed) == expected
+
+
+def test_remaining_budget_before_shutdown(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30")
+    assert GracefulShutdownManager.remaining(clock=lambda: 100.0) == 30.0
+    assert GracefulShutdownManager.remaining(-1, clock=lambda: 100.0) == 0.0
+
+
+@pytest.mark.parametrize("prestop", [False, True])
+def test_sigterm_stamps_before_chaining_and_restores_server_handler(
+    monkeypatch: pytest.MonkeyPatch, prestop: bool
+) -> None:
+    import signal
+    from queue import SimpleQueue
+    from types import FrameType
+    from typing import Final
+
+    monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30")
+    calls: Final[SimpleQueue[float]] = SimpleQueue()
+
+    def server_handler(signum: int, frame: FrameType | None) -> None:
+        assert signum == signal.SIGTERM
+        calls.put(GracefulShutdownManager.remaining(clock=lambda: 105.0))
+
+    original: Final = signal.signal(signal.SIGTERM, server_handler)
+    try:
+        if prestop:
+            GracefulShutdownManager.start_shutdown(clock=lambda: 95.0)
+        with GracefulShutdownManager.capture_sigterm(clock=lambda: 100.0):
+            installed: Final = signal.getsignal(signal.SIGTERM)
+            with GracefulShutdownManager.capture_sigterm(clock=lambda: 102.0):
+                assert signal.getsignal(signal.SIGTERM) is installed
+                signal.raise_signal(signal.SIGTERM)
+                signal.raise_signal(signal.SIGTERM)
+            assert calls.get_nowait() == (20.0 if prestop else 25.0)
+            assert calls.get_nowait() == (20.0 if prestop else 25.0)
+            assert calls.empty()
+        assert signal.getsignal(signal.SIGTERM) is server_handler
+    finally:
+        signal.signal(signal.SIGTERM, original)
+
+
+def test_sigterm_capture_preserves_a_replacement_handler() -> None:
+    import signal
+    from types import FrameType
+    from typing import Final
+
+    def server_handler(signum: int, frame: FrameType | None) -> None:
+        pass
+
+    original: Final = signal.signal(signal.SIGTERM, server_handler)
+    try:
+        with GracefulShutdownManager.capture_sigterm():
+            signal.signal(signal.SIGTERM, signal.SIG_IGN)
+        assert signal.getsignal(signal.SIGTERM) is signal.SIG_IGN
+    finally:
+        signal.signal(signal.SIGTERM, original)
+
+
+@pytest.mark.parametrize("handler", ["SIG_DFL", "SIG_IGN"])
+def test_sigterm_capture_preserves_default_and_ignored_dispositions(handler: str) -> None:
+    import signal
+    from typing import Final
+
+    disposition: Final = signal.SIG_DFL if handler == "SIG_DFL" else signal.SIG_IGN
+    original: Final = signal.signal(signal.SIGTERM, disposition)
+    try:
+        with GracefulShutdownManager.capture_sigterm():
+            assert signal.getsignal(signal.SIGTERM) is disposition
+    finally:
+        signal.signal(signal.SIGTERM, original)
+
+
+def test_sigterm_capture_in_worker_thread_leaves_signal_ownership_unchanged() -> None:
+    import signal
+    from concurrent.futures import ThreadPoolExecutor
+    from typing import Final
+
+    original: Final = signal.getsignal(signal.SIGTERM)
+
+    def run_lifespan() -> None:
+        with GracefulShutdownManager.capture_sigterm():
+            assert signal.getsignal(signal.SIGTERM) is original
+
+    with ThreadPoolExecutor(max_workers=1) as executor:
+        executor.submit(run_lifespan).result()
+    assert signal.getsignal(signal.SIGTERM) is original
+
+
+@pytest.mark.asyncio
+async def test_lifespan_request_drain_does_not_restart_an_exhausted_budget(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30")
+    GracefulShutdownManager.start_shutdown(clock=lambda: 100.0)
+    assert await GracefulShutdownManager.wait_for_drain(count_fn=lambda: 2, clock=lambda: 130.0) == 2

tests/unit/proxy/test_proxy_server.py

@@ -263,9 +263,10 @@ token = MASTER_KEY
 
 @pytest.mark.asyncio
 async def test_shutdown_persists_accounting_from_worker_callbacks_before_returning() -> None:
-    from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+    from litellm.litellm_core_utils.logging_worker import LoggingWorker
     from litellm.responses.streaming_iterator import spawn_logging_task
 
+    worker: Final = LoggingWorker()
     started: Final = asyncio.Event()
     release: Final = asyncio.Event()
     first_flush: Final = asyncio.Event()
@@ -285,10 +286,10 @@ async def test_shutdown_persists_accounting_from_worker_callbacks_before_returni
             persisted.put_nowait(pending.get_nowait())
         first_flush.set()
 
-    GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(callback())
+    worker.ensure_initialized_and_enqueue(callback())
     await started.wait()
     shutdown: Final = asyncio.create_task(
-        litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend)
+        litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend, worker=worker)
     )
     try:
         await first_flush.wait()
@@ -300,20 +301,23 @@ async def test_shutdown_persists_accounting_from_worker_callbacks_before_returni
         assert pending.empty()
     finally:
         release.set()
-        await GLOBAL_LOGGING_WORKER.flush()
-        await GLOBAL_LOGGING_WORKER.stop()
+        await worker.flush()
+        await worker.stop()
 
 
 @pytest.mark.asyncio
 async def test_shutdown_reports_worker_timeout_and_flushes_completed_accounting(
     monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
 ) -> None:
-    from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+    from litellm.litellm_core_utils.logging_worker import LoggingWorker
+    from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
 
-    monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "0")
+    monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30")
+    worker: Final = LoggingWorker()
     started: Final = asyncio.Event()
     cancelled: Final = asyncio.Event()
     flushed_after_cancellation: Final = asyncio.Event()
+    GracefulShutdownManager.start_shutdown(clock=lambda: 100.0)
 
     async def callback() -> None:
         started.set()
@@ -326,63 +330,100 @@ async def test_shutdown_reports_worker_timeout_and_flushes_completed_accounting(
         if cancelled.is_set():
             flushed_after_cancellation.set()
 
-    GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(callback())
+    worker.ensure_initialized_and_enqueue(callback())
     await started.wait()
     try:
         with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
-            await litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend)
+            await litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(
+                flush_spend, worker=worker, clock=lambda: 120.0
+            )
         assert cancelled.is_set()
         assert flushed_after_cancellation.is_set()
-        assert "Timed out draining success callbacks" in caplog.text
+        assert "budget exhausted with 1 unfinished callbacks" in caplog.text
     finally:
-        await GLOBAL_LOGGING_WORKER.stop()
+        GracefulShutdownManager.reset()
+        await worker.stop()
 
 
 @pytest.mark.asyncio
 async def test_shutdown_bounds_callback_cancellation_and_flushes_accounting(
     monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
 ) -> None:
-    from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
+    from litellm.litellm_core_utils.logging_worker import LoggingWorker
+    from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
 
-    monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "0")
+    monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30")
+    worker: Final = LoggingWorker()
     started: Final = asyncio.Event()
     cancelling: Final = asyncio.Event()
-    cancelled: Final = asyncio.Event()
+    release: Final = asyncio.Event()
+    finished: Final = asyncio.Event()
     flushed: Final = asyncio.Event()
-    loop: Final = asyncio.get_running_loop()
-    clock: Final = loop.time
+    GracefulShutdownManager.start_shutdown(clock=lambda: 100.0)
 
     async def callback() -> None:
         started.set()
         try:
-            await asyncio.Event().wait()
+            await release.wait()
         except asyncio.CancelledError:
             cancelling.set()
-            try:
-                await asyncio.Event().wait()
-            finally:
-                cancelled.set()
+            await release.wait()
+        finally:
+            finished.set()
+
+    async def drain_tasks(timeout: float) -> None:
+        assert timeout == 0.0
+        await cancelling.wait()
 
     async def flush_spend() -> None:
-        if cancelled.is_set():
+        if cancelling.is_set():
             flushed.set()
 
-    GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(callback())
+    worker.ensure_initialized_and_enqueue(callback())
     await started.wait()
     try:
-        with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"):
-            shutdown: Final = asyncio.create_task(
-                litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend)
+        with caplog.at_level(logging.WARNING, logger="LiteLLM"):
+            await litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(
+                flush_spend, worker=worker, clock=lambda: 130.0, drain_tasks=drain_tasks
             )
-            await cancelling.wait()
-            monkeypatch.setattr(loop, "time", lambda: clock() + 6)
-            await shutdown
-        assert cancelled.is_set()
+        assert not finished.is_set(), "The cancellation-resistant callback must remain observable"
         assert flushed.is_set()
-        assert "Timed out cancelling success callbacks" in caplog.text
+        assert "tasks still alive after cancellation" in caplog.text
     finally:
-        monkeypatch.setattr(loop, "time", clock)
-        await GLOBAL_LOGGING_WORKER.stop()
+        release.set()
+        await worker.flush()
+        await worker.stop()
+        GracefulShutdownManager.reset()
+
+
+@pytest.mark.asyncio
+async def test_shutdown_preserves_persistence_reserve_after_partial_http_drain(
+    monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture
+) -> None:
+    from litellm.litellm_core_utils.logging_worker import LoggingWorker
+    from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
+
+    monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30")
+    worker: Final = LoggingWorker()
+    persisted: Final = asyncio.Event()
+    GracefulShutdownManager.start_shutdown(clock=lambda: 100.0)
+
+    async def flush_spend() -> None:
+        persisted.set()
+
+    async def drain_tasks(timeout: float) -> None:
+        assert timeout == 5.0
+
+    try:
+        with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"):
+            await litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(
+                flush_spend, worker=worker, clock=lambda: 106.0, drain_tasks=drain_tasks
+            )
+        assert "remaining_s=24.000 callback_timeout_s=14.000" in caplog.text
+        assert persisted.is_set()
+    finally:
+        await worker.stop()
+        GracefulShutdownManager.reset()
 
 
 headers = {"Authorization": f"Bearer {token}"}