fix: type shutdown cancellation tasks and document its clock

litellm/litellm_core_utils/logging_worker.py

@@ -72,8 +72,8 @@ class LoggingWorker:
         self.concurrency = concurrency
         self.timeout_summary_window = timeout_summary_window
         self._queue: asyncio.Queue[LoggingTask] | None = None
-        self._worker_task: asyncio.Task | None = None
-        self._running_tasks: set[asyncio.Task] = set()
+        self._worker_task: asyncio.Task[None] | None = None
+        self._running_tasks: set[asyncio.Task[None]] = set()
         self._dequeued_tasks: dict[int, LoggingTask] = {}  # mutable-ok: refs so flush can rescue never-started tasks
         self._sem: asyncio.Semaphore | None = None
         self._bound_loop: asyncio.AbstractEventLoop | None = None
@@ -495,7 +495,7 @@ class LoggingWorker:
             # No worker launched and no in-flight tasks to drain.
             return
 
-        tasks_to_cancel: Final[list[asyncio.Task]] = list(self._running_tasks)
+        tasks_to_cancel: Final[list[asyncio.Task[None]]] = list(self._running_tasks)
         if self._worker_task:
             # Include the main worker loop so it stops fetching work.
             tasks_to_cancel.append(self._worker_task)

litellm/proxy/db/db_spend_update_writer.py

@@ -1900,7 +1900,6 @@ class DBSpendUpdateWriter:
             proxy_logging_obj=proxy_logging_obj,
         )
 
-
         ################## Daily End-User Spend Update Transactions ##################
         # Aggregate all in memory daily end-user spend transactions and commit to db
         await self._flush_daily_spend_queue(

litellm/proxy/shutdown/graceful_shutdown_manager.py

@@ -24,7 +24,7 @@ import os
 import signal
 import threading
 import time
-from collections.abc import Callable, Iterator
+from collections.abc import Callable, Generator
 from contextlib import contextmanager
 from types import FrameType
 from typing import Final
@@ -99,7 +99,7 @@ class GracefulShutdownManager:
 
     @classmethod
     @contextmanager
-    def capture_sigterm(cls, clock: Callable[[], float] = time.monotonic) -> Iterator[None]:
+    def capture_sigterm(cls, clock: Callable[[], float] = time.monotonic) -> Generator[None, None, None]:
         if threading.current_thread() is not threading.main_thread():
             yield
             return
@@ -144,6 +144,7 @@ class GracefulShutdownManager:
                 live ``InFlightRequestsMiddleware`` counter; injectable for tests.
             poll_interval: Seconds between counter polls.
             log_interval: Minimum seconds between ``drain_waiting`` log lines.
+            clock: Monotonic clock shared with the shutdown timestamp.
 
         Returns:
             Number of requests that drained while waiting (>= 0).