mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-30 06:43:27 +00:00
fix(proxy): persist tag spend and drain success callbacks within the shutdown budget (#48988)
* fix: drain shutdown callbacks and persist pending tag spend * fix: catch shutdown timeouts on Python 3.10 * fix: share shutdown callback budget with HTTP draining * fix: type shutdown cancellation tasks and document its clock * test: retain callbacks after bounded worker cancellation * test(proxy): inspect the lifespan body after the SIGTERM wrapper split Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(utils): keep the success callback enqueue on its own task The synchronous enqueue spawned the logging worker loop inside the awaited call and ran callbacks that per-call loops used to drop, which hung the SDK callback tests on the autouse worker flush. The shutdown worker flush covers the queued callback on its own, so the one-tick task is restored Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(logging_worker): keep the queue's unfinished count exact when a full queue rejects a callback The aggressive clear ran the rejected callback and marked it done against the queue it was never put in, so flush() could return while a real callback was still running and the shutdown drain cancelled it early. The rejected callback is now re-enqueued once the clear made room, and only runs uncounted when the queue is still full Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(logging_worker): build the aggressive-clear run set immutably Final tuple instead of an appended list, with the re-enqueue in a _requeue helper, and Final annotations on the overflow regression test Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(logging): track overflow callbacks through shutdown * chore(logging): explain shared callback ownership * test(logging): isolate deferred callback completion --------- Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com> Co-authored-by: joshua <joshua@berri.ai> Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot]
GitHub
parent
c98641fd5a
commit
c45b264024
14 changed files with 1011 additions and 143 deletions
|
|
@ -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`
|
||||
|
|
|
|||
|
|
@ -6,7 +6,7 @@ import atexit
|
|||
import contextvars
|
||||
import inspect
|
||||
import logging
|
||||
from collections.abc import Coroutine, Iterator
|
||||
from collections.abc import Awaitable, Callable, Coroutine, Iterator
|
||||
from typing import Final
|
||||
|
||||
from typing_extensions import TypedDict
|
||||
|
|
@ -66,14 +66,19 @@ class LoggingWorker:
|
|||
max_queue_size: int = LOGGING_WORKER_MAX_QUEUE_SIZE,
|
||||
concurrency: int = LOGGING_WORKER_CONCURRENCY,
|
||||
timeout_summary_window: float = LOGGING_WORKER_TIMEOUT_SUMMARY_WINDOW_SECONDS,
|
||||
*,
|
||||
retry_sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
|
||||
):
|
||||
self.timeout = timeout
|
||||
self.max_queue_size = max_queue_size
|
||||
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._overflow_tasks: set[asyncio.Task[None]] = set() # mutable-ok: done callbacks discard refs
|
||||
self._pending_enqueues: dict[int, LoggingTask] = {} # mutable-ok: retries and shutdown share callback ownership
|
||||
self._retry_sleep = retry_sleep
|
||||
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
|
||||
|
|
@ -168,6 +173,7 @@ class LoggingWorker:
|
|||
self._sem = None
|
||||
self._worker_task = None
|
||||
self._running_tasks.clear()
|
||||
self._overflow_tasks.clear()
|
||||
self._timeout_summary_task = None
|
||||
self._timeout_burst_count = 0
|
||||
self._queue = new_queue
|
||||
|
|
@ -185,6 +191,10 @@ class LoggingWorker:
|
|||
self._sem = asyncio.Semaphore(self.concurrency)
|
||||
if self._worker_task is None or self._worker_task.done():
|
||||
self._worker_task = asyncio.create_task(self._worker_loop())
|
||||
if not self._overflow_tasks:
|
||||
self._aggressive_clear_in_progress = False
|
||||
for pending in tuple(self._pending_enqueues.values()):
|
||||
self._schedule_delayed_enqueue_retry(pending)
|
||||
|
||||
async def _process_log_task(
|
||||
self, task: LoggingTask, sem: asyncio.Semaphore, queue: "asyncio.Queue[LoggingTask]"
|
||||
|
|
@ -278,12 +288,24 @@ class LoggingWorker:
|
|||
# Capture the current context when enqueueing
|
||||
task: Final = LoggingTask(coroutine=coroutine, context=contextvars.copy_context())
|
||||
|
||||
self._enqueue_task(task)
|
||||
|
||||
def _enqueue_task(self, task: LoggingTask) -> None:
|
||||
if self._queue is None:
|
||||
return
|
||||
try:
|
||||
self._queue.put_nowait(task)
|
||||
except asyncio.QueueFull:
|
||||
# Queue is full - handle it appropriately
|
||||
verbose_logger.exception("LoggingWorker queue is full")
|
||||
self._pending_enqueues[id(task)] = task
|
||||
self._handle_queue_full(task)
|
||||
else:
|
||||
self._pending_enqueues.pop(id(task), None)
|
||||
|
||||
def _start_overflow_task(self, coroutine: Coroutine[object, object, None]) -> None:
|
||||
task: Final = asyncio.create_task(coroutine)
|
||||
self._overflow_tasks.add(task)
|
||||
task.add_done_callback(self._overflow_tasks.discard)
|
||||
|
||||
def _should_start_aggressive_clear(self) -> bool:
|
||||
"""
|
||||
|
|
@ -326,7 +348,7 @@ class LoggingWorker:
|
|||
if self._should_start_aggressive_clear():
|
||||
self._mark_aggressive_clear_started()
|
||||
# Schedule clearing as async task so enqueue returns immediately (non-blocking)
|
||||
asyncio.create_task(self._aggressively_clear_queue_async(task))
|
||||
self._start_overflow_task(self._aggressively_clear_queue_async(task))
|
||||
else:
|
||||
# Cooldown active or clear in progress, schedule a delayed retry
|
||||
self._schedule_delayed_enqueue_retry(task)
|
||||
|
|
@ -363,7 +385,7 @@ class LoggingWorker:
|
|||
delay: Final = self._calculate_retry_delay()
|
||||
|
||||
# Schedule the retry as a background task
|
||||
asyncio.create_task(self._retry_enqueue_task(task, delay))
|
||||
self._start_overflow_task(self._retry_enqueue_task(task, delay))
|
||||
except RuntimeError:
|
||||
# No event loop, drop the task as we can't schedule a retry
|
||||
pass
|
||||
|
|
@ -373,17 +395,8 @@ class LoggingWorker:
|
|||
Retry enqueueing the task after delay, preserving original context.
|
||||
This is called as a background task from _schedule_delayed_enqueue_retry.
|
||||
"""
|
||||
await asyncio.sleep(delay)
|
||||
|
||||
# Try to enqueue the task directly, preserving its original context
|
||||
if self._queue is None:
|
||||
return
|
||||
|
||||
try:
|
||||
self._queue.put_nowait(task)
|
||||
except asyncio.QueueFull:
|
||||
# Still full - handle it appropriately (clear or retry again)
|
||||
self._handle_queue_full(task)
|
||||
await self._retry_sleep(delay)
|
||||
self._enqueue_task(task)
|
||||
|
||||
def _extract_tasks_from_queue(self) -> list[LoggingTask]:
|
||||
"""
|
||||
|
|
@ -423,15 +436,10 @@ class LoggingWorker:
|
|||
return
|
||||
|
||||
extracted_tasks: Final = self._extract_tasks_from_queue()
|
||||
|
||||
# Add new task to extracted tasks to process directly
|
||||
if new_task is not None:
|
||||
self._track_dequeued(new_task)
|
||||
extracted_tasks.append(new_task)
|
||||
|
||||
# Process extracted tasks directly
|
||||
self._enqueue_task(new_task)
|
||||
if extracted_tasks:
|
||||
await self._process_extracted_tasks(extracted_tasks)
|
||||
await asyncio.gather(*(self._process_single_task(task) for task in extracted_tasks))
|
||||
except Exception as e:
|
||||
verbose_logger.exception("LoggingWorker error during aggressive clear: %s", e)
|
||||
finally:
|
||||
|
|
@ -439,7 +447,7 @@ class LoggingWorker:
|
|||
self._aggressive_clear_in_progress = False
|
||||
|
||||
async def _process_single_task(self, task: LoggingTask) -> None:
|
||||
"""Process a single task and mark it done."""
|
||||
"""Process a task extracted from the queue and mark it done."""
|
||||
if self._queue is None:
|
||||
return
|
||||
|
||||
|
|
@ -455,17 +463,6 @@ class LoggingWorker:
|
|||
self._untrack_dequeued(task)
|
||||
self._queue.task_done()
|
||||
|
||||
async def _process_extracted_tasks(self, tasks: list[LoggingTask]) -> None:
|
||||
"""
|
||||
Process tasks that were extracted from the queue to make room.
|
||||
Processes them concurrently without semaphore limits for maximum speed.
|
||||
"""
|
||||
if not tasks or self._queue is None:
|
||||
return
|
||||
|
||||
# Process all tasks concurrently for maximum speed
|
||||
await asyncio.gather(*[self._process_single_task(task) for task in tasks])
|
||||
|
||||
def ensure_initialized_and_enqueue(self, async_coroutine: Coroutine[object, None, None]) -> None:
|
||||
"""
|
||||
Ensure the logging worker is initialized and enqueue the coroutine.
|
||||
|
|
@ -473,32 +470,63 @@ class LoggingWorker:
|
|||
self.start()
|
||||
self.enqueue(async_coroutine)
|
||||
|
||||
async def stop(self) -> None:
|
||||
@property
|
||||
def pending_count(self) -> int:
|
||||
return (
|
||||
(self._queue.qsize() if self._queue is not None else 0)
|
||||
+ len(self._dequeued_tasks)
|
||||
+ len(self._pending_enqueues)
|
||||
)
|
||||
|
||||
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:
|
||||
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
|
||||
self._emit_timeout_summary()
|
||||
|
||||
if self._worker_task is None and not self._running_tasks:
|
||||
# No worker launched and no in-flight tasks to drain.
|
||||
return
|
||||
if discard_pending:
|
||||
rejected: Final = tuple(self._pending_enqueues.values())
|
||||
self._pending_enqueues.clear()
|
||||
for task in rejected:
|
||||
task["coroutine"].close()
|
||||
if rejected:
|
||||
verbose_logger.warning("LoggingWorker shutdown: %d overflow callbacks discarded", len(rejected))
|
||||
|
||||
tasks_to_cancel: Final[list[asyncio.Task]] = 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)
|
||||
tasks_to_cancel: Final = (
|
||||
*self._running_tasks,
|
||||
*self._overflow_tasks,
|
||||
*((self._worker_task,) if self._worker_task is not None else ()),
|
||||
)
|
||||
if not tasks_to_cancel:
|
||||
return
|
||||
|
||||
for task in tasks_to_cancel:
|
||||
# Propagate cancellation to every pending task.
|
||||
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))
|
||||
|
||||
self._worker_task = None
|
||||
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()
|
||||
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.
|
||||
|
|
@ -516,6 +544,9 @@ class LoggingWorker:
|
|||
return
|
||||
self.start()
|
||||
await self._queue.join()
|
||||
while self._overflow_tasks:
|
||||
await asyncio.wait(tuple(self._overflow_tasks))
|
||||
await self._queue.join()
|
||||
|
||||
async def clear_queue(self):
|
||||
"""
|
||||
|
|
|
|||
|
|
@ -1900,8 +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
|
||||
await self._flush_daily_spend_queue(
|
||||
|
|
|
|||
|
|
@ -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 (
|
||||
|
|
@ -1274,15 +1276,44 @@ async def _flush_spend_logs_queue_on_shutdown() -> None:
|
|||
|
||||
|
||||
async def _flush_spend_on_shutdown() -> None:
|
||||
from litellm.proxy.utils import update_daily_tag_spend
|
||||
|
||||
await flush_spend_counters_on_shutdown()
|
||||
if prisma_client is not None:
|
||||
await update_daily_tag_spend(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj)
|
||||
await _flush_spend_logs_queue_on_shutdown()
|
||||
|
||||
|
||||
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:
|
||||
await flush_spend()
|
||||
await drain_logging_tasks()
|
||||
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,
|
||||
)
|
||||
|
||||
async def drain_callbacks() -> None:
|
||||
await drain_tasks(callback_timeout)
|
||||
await worker.flush()
|
||||
|
||||
try:
|
||||
await asyncio.wait_for(drain_callbacks(), timeout=callback_timeout)
|
||||
except asyncio.TimeoutError:
|
||||
verbose_proxy_logger.warning(
|
||||
"Success callback shutdown budget exhausted with %d unfinished callbacks; accounting may be incomplete",
|
||||
worker.pending_count,
|
||||
)
|
||||
finally:
|
||||
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()
|
||||
|
||||
|
||||
|
|
@ -1445,6 +1476,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, \
|
||||
|
|
|
|||
|
|
@ -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, Generator
|
||||
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) -> Generator[None, None, 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.
|
||||
|
|
@ -107,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).
|
||||
|
|
@ -120,7 +158,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 +166,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 +180,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 +195,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",
|
||||
|
|
|
|||
|
|
@ -227,7 +227,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(
|
||||
|
|
|
|||
|
|
@ -1632,15 +1632,19 @@ def _schedule_async_success_logging(
|
|||
return
|
||||
|
||||
def _enqueue_async_logging() -> None:
|
||||
from litellm.responses.streaming_iterator import spawn_logging_task
|
||||
|
||||
logging_obj.mark_async_success_log_enqueued()
|
||||
asyncio.create_task(
|
||||
spawn_logging_task(
|
||||
asyncio.get_running_loop(),
|
||||
_client_async_logging_helper(
|
||||
logging_obj=logging_obj,
|
||||
result=result,
|
||||
start_time=start_time,
|
||||
end_time=end_time,
|
||||
is_completion_with_fallbacks=is_completion_with_fallbacks,
|
||||
)
|
||||
),
|
||||
task_name="success callback enqueue",
|
||||
)
|
||||
|
||||
if not logging_obj.defer_async_logging:
|
||||
|
|
|
|||
|
|
@ -1,11 +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
|
||||
|
||||
|
|
@ -26,6 +29,225 @@ MODEL_INSERT_ARRIVED_LOG_LINE: Final = "path=/model/new"
|
|||
COMMIT_DELAY_SECONDS: Final = 15
|
||||
BURST_REQUESTS: Final = 30
|
||||
|
||||
_GATED_LOGGER: Final = """
|
||||
import os
|
||||
import re
|
||||
import httpx
|
||||
from litellm.integrations.custom_logger import CustomLogger
|
||||
|
||||
class ShutdownLogger(CustomLogger):
|
||||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||||
async with httpx.AsyncClient(timeout=25, trust_env=False) as client:
|
||||
await client.post(os.environ["SHUTDOWN_CALLBACK_URL"])
|
||||
|
||||
logger = ShutdownLogger()
|
||||
"""
|
||||
|
||||
|
||||
def _tag_shutdown_config(directory: Path, upstream: str, redis_buffer: bool, *, gated: bool = False) -> Path:
|
||||
path: Final = directory / "tag-shutdown.yaml"
|
||||
if gated:
|
||||
(directory / "shutdown_callbacks.py").write_text(_GATED_LOGGER)
|
||||
path.write_text(
|
||||
yaml.safe_dump(
|
||||
{
|
||||
"model_list": [
|
||||
{
|
||||
"model_name": "shutdown-model",
|
||||
"litellm_params": {
|
||||
"model": "openai/gpt-4o-mini",
|
||||
"api_base": upstream + "/v1",
|
||||
"api_key": "synthetic-key",
|
||||
"input_cost_per_token": 0.01,
|
||||
"output_cost_per_token": 0.02,
|
||||
},
|
||||
}
|
||||
],
|
||||
"general_settings": {
|
||||
"master_key": "os.environ/LITELLM_MASTER_KEY",
|
||||
"database_url": "os.environ/DATABASE_URL",
|
||||
"proxy_batch_write_at": 3600,
|
||||
"use_redis_transaction_buffer": redis_buffer,
|
||||
"disable_model_info_refresh": True,
|
||||
},
|
||||
"litellm_settings": {
|
||||
"cache": True,
|
||||
"cache_params": {"type": "redis", "host": "os.environ/REDIS_HOST", "port": "os.environ/REDIS_PORT"},
|
||||
"callbacks": ["shutdown_callbacks.logger"] if gated else [],
|
||||
},
|
||||
}
|
||||
)
|
||||
)
|
||||
return path
|
||||
|
||||
|
||||
def _tagged_chat(gateway: Gateway, key: str, tags: tuple[str, str]) -> None:
|
||||
response: Final = gateway.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{
|
||||
"model": "shutdown-model",
|
||||
"messages": [{"role": "user", "content": uuid.uuid4().hex}],
|
||||
"metadata": {"tags": list(tags)},
|
||||
},
|
||||
key=key,
|
||||
)
|
||||
assert response.status_code == 200, response.text
|
||||
|
||||
|
||||
def _assert_shutdown_accounting(key: str, tags: tuple[str, str], requests: int) -> None:
|
||||
hashed: Final = hashlib.sha256(key.encode()).hexdigest()
|
||||
for table, count in (("LiteLLM_SpendLogs", "count(*)"), ("LiteLLM_DailyUserSpend", "sum(api_requests)")):
|
||||
rows: Final = read_rows(
|
||||
f'SELECT ({count})::int AS requests, sum(spend) AS spend FROM "{table}" WHERE api_key=%s', (hashed,)
|
||||
)
|
||||
assert rows[0]["requests"] == requests, (table, rows)
|
||||
assert rows[0]["spend"] == pytest.approx(requests * 0.05, abs=1e-9), (table, rows)
|
||||
key_rows: Final = read_rows('SELECT spend FROM "LiteLLM_VerificationToken" WHERE token=%s', (hashed,))
|
||||
assert key_rows[0]["spend"] == pytest.approx(requests * 0.05, abs=1e-9)
|
||||
for tag in tags:
|
||||
tag_rows: Final = read_rows(
|
||||
'SELECT sum(api_requests)::int AS requests, sum(spend) AS spend FROM "LiteLLM_DailyTagSpend" WHERE api_key=%s AND tag=%s',
|
||||
(hashed, tag),
|
||||
)
|
||||
assert tag_rows[0]["requests"] == requests, (tag, tag_rows)
|
||||
assert tag_rows[0]["spend"] == pytest.approx(requests * 0.05, abs=1e-9), (tag, tag_rows)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("redis_buffer", [False, True])
|
||||
def test_tagged_request_survives_shutdown_before_the_tag_timer(
|
||||
gateway: Gateway, tmp_path: Path, redis_buffer: bool
|
||||
) -> None:
|
||||
tags: Final = (f"shutdown-a-{uuid.uuid4().hex}", f"shutdown-b-{uuid.uuid4().hex}")
|
||||
overrides: Final = {"DATABASE_URL": os.environ["DATABASE_URL"], "LITELLM_LOG": "DEBUG"}
|
||||
with wire_server(_provider) as wire:
|
||||
config: Final = _tag_shutdown_config(tmp_path, wire.url, redis_buffer)
|
||||
with owned_proxy_process(gateway, tmp_path, overrides, config=config) as owned:
|
||||
key: Final = string_value(
|
||||
owned.gateway.post("/key/generate", {"user_id": f"shutdown-{uuid.uuid4().hex}"})["key"]
|
||||
)
|
||||
_tagged_chat(owned.gateway, key, tags)
|
||||
owned.begin_shutdown()
|
||||
owned.process.wait(timeout=90)
|
||||
with owned_proxy_process(gateway, tmp_path, overrides, config=config) as restarted:
|
||||
_assert_shutdown_accounting(key, tags, 1)
|
||||
restarted.begin_shutdown()
|
||||
restarted.process.wait(timeout=90)
|
||||
_assert_shutdown_accounting(key, tags, 1)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("redis_buffer", [False, True])
|
||||
@pytest.mark.parametrize("overflow", [False, True])
|
||||
def test_callback_backlog_is_accounted_before_shutdown_disconnects(
|
||||
gateway: Gateway, tmp_path: Path, redis_buffer: bool, overflow: bool
|
||||
) -> None:
|
||||
release: Final = threading.Event()
|
||||
requests: Final = 6 if overflow else 3
|
||||
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,
|
||||
):
|
||||
config: Final = _tag_shutdown_config(tmp_path, wire.url, redis_buffer, gated=True)
|
||||
with owned_proxy_process(
|
||||
gateway,
|
||||
tmp_path,
|
||||
{
|
||||
"DATABASE_URL": os.environ["DATABASE_URL"],
|
||||
"LOGGING_WORKER_CONCURRENCY": "2",
|
||||
"LOGGING_WORKER_MAX_QUEUE_SIZE": "2" if overflow else "50000",
|
||||
"SHUTDOWN_CALLBACK_URL": logger.url,
|
||||
"GRACEFUL_SHUTDOWN_TIMEOUT": "15",
|
||||
"LITELLM_LOG": "DEBUG",
|
||||
},
|
||||
config=config,
|
||||
) as owned:
|
||||
key: Final = string_value(
|
||||
owned.gateway.post("/key/generate", {"user_id": f"backlog-{uuid.uuid4().hex}"})["key"]
|
||||
)
|
||||
try:
|
||||
for _ in range(requests):
|
||||
_tagged_chat(owned.gateway, key, tags)
|
||||
eventually(logger.received.qsize, lambda count: count >= 2, seconds=10)
|
||||
owned.begin_shutdown()
|
||||
eventually(
|
||||
lambda: (
|
||||
owned.process.poll() is not None
|
||||
or "Draining success callbacks before the final spend flush" in owned.log.read_text()
|
||||
),
|
||||
lambda ready: ready,
|
||||
seconds=30,
|
||||
)
|
||||
release.set()
|
||||
owned.process.wait(timeout=90)
|
||||
_assert_shutdown_accounting(key, tags, requests)
|
||||
finally:
|
||||
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(
|
||||
|
|
|
|||
|
|
@ -82,6 +82,252 @@ class _RecordCollector(logging.Handler):
|
|||
self.records.append(record)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("running", [False, True])
|
||||
async def test_stop_discards_waiting_callbacks_without_executing_them(running: bool) -> None:
|
||||
worker: Final = LoggingWorker(concurrency=1)
|
||||
started: Final = asyncio.Event()
|
||||
cancelled: Final = asyncio.Event()
|
||||
executed: Final = asyncio.Event()
|
||||
|
||||
async def blocked() -> None:
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
finally:
|
||||
cancelled.set()
|
||||
|
||||
async def waiting() -> None:
|
||||
executed.set()
|
||||
|
||||
if running:
|
||||
worker.ensure_initialized_and_enqueue(blocked())
|
||||
await started.wait()
|
||||
else:
|
||||
worker._ensure_queue()
|
||||
worker.enqueue(waiting())
|
||||
await worker.stop(discard_pending=True)
|
||||
assert not executed.is_set()
|
||||
assert cancelled.is_set() == running
|
||||
assert worker._queue is not None
|
||||
assert worker._queue.empty()
|
||||
await asyncio.wait_for(worker._queue.join(), timeout=1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_does_not_replay_a_callback_dequeued_before_its_processor_started() -> None:
|
||||
worker: Final = LoggingWorker(concurrency=1)
|
||||
executed: Final = asyncio.Event()
|
||||
|
||||
async def callback() -> None:
|
||||
executed.set()
|
||||
|
||||
worker.ensure_initialized_and_enqueue(callback())
|
||||
await asyncio.sleep(0)
|
||||
await worker.stop(discard_pending=True)
|
||||
assert not executed.is_set()
|
||||
worker._flush_on_exit()
|
||||
assert not executed.is_set()
|
||||
assert worker._queue is not None
|
||||
await asyncio.wait_for(worker._queue.join(), timeout=1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stop_timeout_keeps_callbacks_observable_until_queue_cleanup_finishes() -> None:
|
||||
worker: Final = LoggingWorker(concurrency=1)
|
||||
started: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
completed: Final = asyncio.Event()
|
||||
|
||||
async def running() -> None:
|
||||
started.set()
|
||||
await asyncio.Event().wait()
|
||||
|
||||
async def queued() -> None:
|
||||
await release.wait()
|
||||
completed.set()
|
||||
|
||||
worker.ensure_initialized_and_enqueue(running())
|
||||
await started.wait()
|
||||
worker.enqueue(queued())
|
||||
try:
|
||||
await worker.stop(timeout=0)
|
||||
assert not completed.is_set()
|
||||
release.set()
|
||||
await worker.flush()
|
||||
assert completed.is_set(), "Queue cleanup must survive a bounded cancellation wait"
|
||||
finally:
|
||||
release.set()
|
||||
await worker.stop()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_overflow_callback_does_not_undercount_the_unfinished_queue() -> None:
|
||||
"""A callback rejected by a full queue is run by the aggressive clear. It must not consume a
|
||||
queue.task_done() it never contributed, otherwise flush() returns while a real callback is still
|
||||
running and the shutdown drain cancels it with budget to spare.
|
||||
"""
|
||||
worker: Final = LoggingWorker(timeout=5.0, max_queue_size=2, concurrency=1)
|
||||
gates: Final = {name: asyncio.Event() for name in ("first", "second", "third")}
|
||||
completed: Final[list[str]] = []
|
||||
|
||||
async def gated(name: str) -> None:
|
||||
await gates[name].wait()
|
||||
completed.append(name)
|
||||
|
||||
async def instant() -> None:
|
||||
completed.append("overflow")
|
||||
|
||||
worker.start()
|
||||
worker.enqueue(gated("first"))
|
||||
await asyncio.sleep(0)
|
||||
worker.enqueue(gated("second"))
|
||||
worker.enqueue(gated("third"))
|
||||
assert worker._queue is not None and worker._queue.full(), "precondition: the queue is full"
|
||||
worker.enqueue(instant())
|
||||
for _ in range(5):
|
||||
await asyncio.sleep(0)
|
||||
|
||||
gates["first"].set()
|
||||
gates["second"].set()
|
||||
asyncio.get_running_loop().call_later(0.2, gates["third"].set)
|
||||
await asyncio.wait_for(worker.flush(), timeout=2.0)
|
||||
assert "third" in completed, "flush returned while a queued callback was still unfinished"
|
||||
assert sorted(completed) == ["first", "overflow", "second", "third"]
|
||||
await worker.stop()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("queue_size", (1, 2))
|
||||
async def test_flush_waits_for_delayed_overflow_and_preserves_callback_context(queue_size: int) -> None:
|
||||
retry_started: Final = asyncio.Event()
|
||||
release_retry: Final = asyncio.Event()
|
||||
completed: Final[asyncio.Queue[tuple[int, str]]] = asyncio.Queue()
|
||||
identity: Final = contextvars.ContextVar("overflow_identity", default="outside")
|
||||
|
||||
async def retry_sleep(delay: float) -> None:
|
||||
retry_started.set()
|
||||
await release_retry.wait()
|
||||
|
||||
worker: Final = LoggingWorker(max_queue_size=queue_size, concurrency=1, retry_sleep=retry_sleep)
|
||||
|
||||
async def callback(index: int) -> None:
|
||||
completed.put_nowait((index, identity.get()))
|
||||
|
||||
worker.start()
|
||||
token: Final = identity.set("request")
|
||||
for index in range(4):
|
||||
worker.enqueue(callback(index))
|
||||
identity.reset(token)
|
||||
flushing: Final = asyncio.create_task(worker.flush())
|
||||
try:
|
||||
await asyncio.wait_for(retry_started.wait(), timeout=2)
|
||||
assert worker._queue is not None
|
||||
await asyncio.wait_for(worker._queue.join(), timeout=2)
|
||||
await asyncio.sleep(0)
|
||||
assert not flushing.done(), "Queue completion must not hide delayed overflow callbacks"
|
||||
assert worker.pending_count > 0
|
||||
release_retry.set()
|
||||
await asyncio.wait_for(flushing, timeout=2)
|
||||
assert sorted(completed.get_nowait() for _ in range(completed.qsize())) == [
|
||||
(index, "request") for index in range(4)
|
||||
]
|
||||
assert worker.pending_count == 0
|
||||
finally:
|
||||
release_retry.set()
|
||||
await worker.stop(discard_pending=True)
|
||||
await asyncio.gather(flushing, return_exceptions=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("start_retry", (False, True))
|
||||
async def test_shutdown_discards_overflow_without_late_enqueue_or_replay(start_retry: bool) -> None:
|
||||
import inspect
|
||||
|
||||
retry_started: Final = asyncio.Event()
|
||||
release_retry: Final = asyncio.Event()
|
||||
executed: Final[asyncio.Queue[int]] = asyncio.Queue()
|
||||
|
||||
async def retry_sleep(delay: float) -> None:
|
||||
retry_started.set()
|
||||
await release_retry.wait()
|
||||
|
||||
worker: Final = LoggingWorker(max_queue_size=2, concurrency=1, retry_sleep=retry_sleep)
|
||||
|
||||
async def callback(index: int) -> None:
|
||||
executed.put_nowait(index)
|
||||
|
||||
callbacks: Final = tuple(callback(index) for index in range(4))
|
||||
worker._ensure_queue()
|
||||
for callback_coro in callbacks:
|
||||
worker.enqueue(callback_coro)
|
||||
if start_retry:
|
||||
await asyncio.wait_for(retry_started.wait(), timeout=2)
|
||||
await worker.stop(discard_pending=True, timeout=1)
|
||||
completed_at_stop: Final = executed.qsize()
|
||||
release_retry.set()
|
||||
await worker.flush()
|
||||
try:
|
||||
assert executed.qsize() == completed_at_stop, "Discarded overflow must not run when the worker restarts"
|
||||
assert worker.pending_count == 0
|
||||
assert all(inspect.getcoroutinestate(coro) == inspect.CORO_CLOSED for coro in callbacks)
|
||||
finally:
|
||||
await worker.stop(discard_pending=True)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shutdown_cancels_aggressive_clear_callbacks_before_returning() -> None:
|
||||
started: Final = asyncio.Event()
|
||||
cancelled: Final = asyncio.Event()
|
||||
worker: Final = LoggingWorker(max_queue_size=2, concurrency=1)
|
||||
|
||||
async def blocked() -> None:
|
||||
started.set()
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
finally:
|
||||
cancelled.set()
|
||||
|
||||
async def noop() -> None:
|
||||
pass
|
||||
|
||||
worker._ensure_queue()
|
||||
worker.enqueue(blocked())
|
||||
worker.enqueue(noop())
|
||||
worker.enqueue(noop())
|
||||
await asyncio.wait_for(started.wait(), timeout=2)
|
||||
await worker.stop(discard_pending=True, timeout=1)
|
||||
assert cancelled.is_set(), "Overflow execution must participate in shutdown cancellation"
|
||||
assert worker.pending_count == 0
|
||||
assert worker._queue is not None
|
||||
await asyncio.wait_for(worker._queue.join(), timeout=1)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_stopped_overflow_can_be_flushed_when_not_discarded() -> None:
|
||||
release_retry: Final = asyncio.Event()
|
||||
completed: Final[asyncio.Queue[int]] = asyncio.Queue()
|
||||
|
||||
async def retry_sleep(delay: float) -> None:
|
||||
await release_retry.wait()
|
||||
|
||||
worker: Final = LoggingWorker(max_queue_size=2, concurrency=1, retry_sleep=retry_sleep)
|
||||
|
||||
async def callback(index: int) -> None:
|
||||
completed.put_nowait(index)
|
||||
|
||||
worker._ensure_queue()
|
||||
for index in range(4):
|
||||
worker.enqueue(callback(index))
|
||||
await worker.stop()
|
||||
release_retry.set()
|
||||
try:
|
||||
await asyncio.wait_for(worker.flush(), timeout=2)
|
||||
assert sorted(completed.get_nowait() for _ in range(completed.qsize())) == list(range(4))
|
||||
finally:
|
||||
await worker.stop(discard_pending=True)
|
||||
|
||||
|
||||
class TestLoggingWorker:
|
||||
"""Test cases for LoggingWorker functionality."""
|
||||
|
||||
|
|
|
|||
|
|
@ -75,11 +75,13 @@ class _RecordingLogger(CustomLogger):
|
|||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.standard_logging_object: StandardLoggingPayload | None = None
|
||||
self.logged: asyncio.Event = asyncio.Event()
|
||||
|
||||
async def async_log_success_event(
|
||||
self, kwargs: Mapping[str, object], response_obj: object, start_time: datetime, end_time: datetime
|
||||
) -> None:
|
||||
self.standard_logging_object = cast(StandardLoggingPayload, kwargs["standard_logging_object"])
|
||||
self.logged.set()
|
||||
|
||||
|
||||
class PostCallGuardrail(CustomGuardrail):
|
||||
|
|
@ -239,52 +241,33 @@ class TestHasPostCallGuardrailsForPassthrough:
|
|||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_deferred_flag_stores_and_executes_closure():
|
||||
"""
|
||||
When defer_async_logging is True on logging_obj:
|
||||
1. wrapper_async stores a callable closure instead of calling create_task
|
||||
2. Calling the closure fires create_task
|
||||
3. Sync callbacks fire immediately (not deferred)
|
||||
"""
|
||||
mock_logging_obj = MagicMock()
|
||||
mock_logging_obj.defer_async_logging = True
|
||||
mock_logging_obj.enqueue_deferred_logging = None
|
||||
|
||||
await litellm.acompletion(
|
||||
model="gpt-3.5-turbo",
|
||||
@pytest.mark.parametrize("deferred", (False, True))
|
||||
async def test_success_logging_waits_for_deferred_closure_when_requested(deferred: bool) -> None:
|
||||
recorder: Final = _RecordingLogger()
|
||||
logging_obj: Final = Logging(
|
||||
model="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
stream=False,
|
||||
call_type="acompletion",
|
||||
start_time=datetime(2026, 1, 1),
|
||||
litellm_call_id="deferred-success",
|
||||
function_id="deferred-success",
|
||||
dynamic_async_success_callbacks=[recorder],
|
||||
)
|
||||
logging_obj.defer_async_logging = deferred
|
||||
response: Final = await litellm.acompletion(
|
||||
model="gpt-4o-mini",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
mock_response="Hello!",
|
||||
litellm_logging_obj=mock_logging_obj,
|
||||
litellm_logging_obj=logging_obj,
|
||||
)
|
||||
|
||||
# Closure was stored
|
||||
enqueue_fn = mock_logging_obj.enqueue_deferred_logging
|
||||
assert callable(enqueue_fn), "Closure should be stored on logging_obj"
|
||||
|
||||
# Sync callbacks fired immediately
|
||||
mock_logging_obj.handle_sync_success_callbacks_for_async_calls.assert_called_once()
|
||||
|
||||
# Calling the closure fires create_task
|
||||
created_tasks = []
|
||||
real_create_task = asyncio.create_task
|
||||
|
||||
def tracking_create_task(coro):
|
||||
task = real_create_task(coro)
|
||||
created_tasks.append(task)
|
||||
return task
|
||||
|
||||
with patch("asyncio.create_task", side_effect=tracking_create_task):
|
||||
enqueue_fn()
|
||||
|
||||
assert len(created_tasks) >= 1, "Closure should fire asyncio.create_task"
|
||||
|
||||
for task in created_tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except (asyncio.CancelledError, Exception):
|
||||
pass
|
||||
if deferred:
|
||||
assert recorder.standard_logging_object is None
|
||||
assert logging_obj.enqueue_deferred_logging is not None
|
||||
logging_obj.enqueue_deferred_logging()
|
||||
await asyncio.wait_for(recorder.logged.wait(), timeout=5)
|
||||
assert recorder.standard_logging_object is not None
|
||||
assert recorder.standard_logging_object["id"] == response.id
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
|
|
@ -401,40 +384,6 @@ async def test_deferred_anthropic_messages_bridged_to_the_responses_api_logs_the
|
|||
assert recorder.standard_logging_object["response_cost"] == pytest.approx(3 * 2e-7 + 7333 * 2e-8 + 23 * 1.25e-6)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. Non-streaming regression: without flag, create_task fires normally
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_no_flag_fires_create_task_normally():
|
||||
"""Without defer_async_logging, wrapper_async calls create_task as before."""
|
||||
created_tasks = []
|
||||
real_create_task = asyncio.create_task
|
||||
|
||||
def tracking_create_task(coro):
|
||||
task = real_create_task(coro)
|
||||
created_tasks.append(task)
|
||||
return task
|
||||
|
||||
with patch("asyncio.create_task", side_effect=tracking_create_task):
|
||||
await litellm.acompletion(
|
||||
model="gpt-3.5-turbo",
|
||||
messages=[{"role": "user", "content": "hi"}],
|
||||
mock_response="Hello!",
|
||||
)
|
||||
|
||||
assert len(created_tasks) >= 1
|
||||
|
||||
for task in created_tasks:
|
||||
if not task.done():
|
||||
task.cancel()
|
||||
try:
|
||||
await task
|
||||
except (asyncio.CancelledError, Exception):
|
||||
pass
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4. Non-streaming: deferred success log is suppressed when guardrail raises
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -35,6 +35,7 @@ from typing_extensions import TypedDict
|
|||
import litellm.proxy.proxy_server as ps
|
||||
from litellm.proxy.management_helpers.audit_logs import track_audit_task
|
||||
from litellm.proxy.proxy_server import (
|
||||
_proxy_startup_event,
|
||||
ProxyStartupEvent,
|
||||
_resolve_pydantic_type,
|
||||
_resolve_typed_dict_type,
|
||||
|
|
@ -1059,7 +1060,7 @@ def test_otel_global_provider_published_after_callback_init():
|
|||
span was orphaned. The publish (``publish_global_otel_v2_provider``) must
|
||||
therefore appear after ``_initialize_startup_logging`` in the lifespan source.
|
||||
"""
|
||||
wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event)
|
||||
wrapped = getattr(_proxy_startup_event, "__wrapped__", _proxy_startup_event)
|
||||
source = inspect.getsource(wrapped)
|
||||
init_pos = source.find("ProxyStartupEvent.initialize_startup_logging(")
|
||||
publish_pos = source.find("publish_global_otel_v2_provider(")
|
||||
|
|
@ -1145,7 +1146,7 @@ def test_proxy_startup_event_warns_for_global_budget_without_database():
|
|||
The call must follow Prisma setup so DB-backed deployments do not false-positive.
|
||||
Direct ``_warn_budget_without_db`` tests cover the warning behavior itself.
|
||||
"""
|
||||
wrapped = getattr(proxy_startup_event, "__wrapped__", proxy_startup_event)
|
||||
wrapped = getattr(_proxy_startup_event, "__wrapped__", _proxy_startup_event)
|
||||
source = inspect.getsource(wrapped)
|
||||
budget_check_pos = source.find("if prisma_client is not None and litellm.max_budget > 0:")
|
||||
warn_pos = source.find("warn_budget_without_db(")
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -263,6 +263,174 @@ async def test_shutdown_commits_spend_a_logging_task_queued_while_the_first_comm
|
|||
|
||||
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 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()
|
||||
pending: Final[asyncio.Queue[str]] = asyncio.Queue()
|
||||
persisted: Final[asyncio.Queue[str]] = asyncio.Queue()
|
||||
|
||||
async def enqueue_spend() -> None:
|
||||
pending.put_nowait("accounted request")
|
||||
|
||||
async def callback() -> None:
|
||||
started.set()
|
||||
await release.wait()
|
||||
spawn_logging_task(asyncio.get_running_loop(), enqueue_spend(), task_name="accounting")
|
||||
|
||||
async def flush_spend() -> None:
|
||||
while not pending.empty():
|
||||
persisted.put_nowait(pending.get_nowait())
|
||||
first_flush.set()
|
||||
|
||||
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, worker=worker)
|
||||
)
|
||||
try:
|
||||
await first_flush.wait()
|
||||
waited_for_callback: Final = not shutdown.done()
|
||||
release.set()
|
||||
await shutdown
|
||||
assert waited_for_callback, "Shutdown returned while accounting callbacks were still running"
|
||||
assert persisted.get_nowait() == "accounted request"
|
||||
assert pending.empty()
|
||||
finally:
|
||||
release.set()
|
||||
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 LoggingWorker
|
||||
from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
|
||||
|
||||
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()
|
||||
try:
|
||||
await asyncio.Event().wait()
|
||||
finally:
|
||||
cancelled.set()
|
||||
|
||||
async def flush_spend() -> None:
|
||||
if cancelled.is_set():
|
||||
flushed_after_cancellation.set()
|
||||
|
||||
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, worker=worker, clock=lambda: 120.0
|
||||
)
|
||||
assert cancelled.is_set()
|
||||
assert flushed_after_cancellation.is_set()
|
||||
assert "budget exhausted with 1 unfinished callbacks" in caplog.text
|
||||
finally:
|
||||
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 LoggingWorker
|
||||
from litellm.proxy.shutdown.graceful_shutdown_manager import GracefulShutdownManager
|
||||
|
||||
monkeypatch.setenv("GRACEFUL_SHUTDOWN_TIMEOUT", "30")
|
||||
worker: Final = LoggingWorker()
|
||||
started: Final = asyncio.Event()
|
||||
cancelling: Final = asyncio.Event()
|
||||
release: Final = asyncio.Event()
|
||||
finished: Final = asyncio.Event()
|
||||
flushed: Final = asyncio.Event()
|
||||
GracefulShutdownManager.start_shutdown(clock=lambda: 100.0)
|
||||
|
||||
async def callback() -> None:
|
||||
started.set()
|
||||
try:
|
||||
await release.wait()
|
||||
except asyncio.CancelledError:
|
||||
cancelling.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 cancelling.is_set():
|
||||
flushed.set()
|
||||
|
||||
worker.ensure_initialized_and_enqueue(callback())
|
||||
await started.wait()
|
||||
try:
|
||||
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
|
||||
)
|
||||
assert not finished.is_set(), "The cancellation-resistant callback must remain observable"
|
||||
assert flushed.is_set()
|
||||
assert "tasks still alive after cancellation" in caplog.text
|
||||
finally:
|
||||
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()
|
||||
drain_budgets: Final[asyncio.Queue[float]] = asyncio.Queue()
|
||||
GracefulShutdownManager.start_shutdown(clock=lambda: 100.0)
|
||||
|
||||
async def flush_spend() -> None:
|
||||
persisted.set()
|
||||
|
||||
async def drain_tasks(timeout: float) -> None:
|
||||
drain_budgets.put_nowait(timeout)
|
||||
|
||||
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()
|
||||
assert tuple(drain_budgets.get_nowait() for _ in range(drain_budgets.qsize())) == (14.0, 5.0)
|
||||
finally:
|
||||
await worker.stop()
|
||||
GracefulShutdownManager.reset()
|
||||
|
||||
|
||||
headers = {"Authorization": f"Bearer {token}"}
|
||||
|
||||
example_completion_result = {
|
||||
|
|
@ -4059,3 +4227,34 @@ def test_detailed_debug_logs_request_body_only_as_received_and_as_sent(
|
|||
assert len(body_dumps) == 2, [message[:120] for message in body_dumps]
|
||||
assert body_dumps[0].startswith("Request received by LiteLLM:")
|
||||
assert "POST Request Sent from LiteLLM" in body_dumps[1]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_shutdown_waits_for_enqueue_producers_before_stopping_worker() -> None:
|
||||
from litellm.litellm_core_utils.logging_worker import LoggingWorker
|
||||
from litellm.responses.streaming_iterator import spawn_logging_task
|
||||
|
||||
worker: Final = LoggingWorker()
|
||||
persisted: Final[asyncio.Queue[str]] = asyncio.Queue()
|
||||
accounting: Final[asyncio.Queue[str]] = asyncio.Queue()
|
||||
|
||||
async def callback() -> None:
|
||||
accounting.put_nowait("completed request")
|
||||
|
||||
async def producer() -> None:
|
||||
await asyncio.sleep(0)
|
||||
await asyncio.sleep(0)
|
||||
worker.ensure_initialized_and_enqueue(callback())
|
||||
|
||||
async def flush_spend() -> None:
|
||||
while not accounting.empty():
|
||||
persisted.put_nowait(accounting.get_nowait())
|
||||
|
||||
spawn_logging_task(asyncio.get_running_loop(), producer(), task_name="success enqueue")
|
||||
try:
|
||||
await litellm.proxy.proxy_server._flush_spend_around_logging_drain_on_shutdown(flush_spend, worker=worker)
|
||||
assert persisted.qsize() == 1, "An enqueue producer must finish before callback draining stops"
|
||||
assert persisted.get_nowait() == "completed request"
|
||||
assert worker.pending_count == 0
|
||||
finally:
|
||||
await worker.stop(discard_pending=True)
|
||||
|
|
|
|||
|
|
@ -11294,6 +11294,7 @@ def test_get_optional_params_forwards_a_tool_config_cache_point_to_a_downstream_
|
|||
assert optional_params["extra_body"] == {"kept": True, "cache_control_injection_points": [point]}
|
||||
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("provider", "expected_warnings"),
|
||||
(("streamlake", 1), ("veniceai", 1), ("my-own-llm", 0)),
|
||||
|
|
@ -11312,3 +11313,34 @@ def test_custom_provider_named_like_a_built_in_warns_once(
|
|||
warnings: Final = [r for r in caplog.records if f"provider '{provider}'" in r.getMessage()]
|
||||
assert len(warnings) == expected_warnings
|
||||
assert provider in litellm._custom_providers
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("deferred", (False, True))
|
||||
async def test_success_enqueue_is_tracked_without_starting_callbacks_synchronously(deferred: bool) -> None:
|
||||
from litellm.litellm_core_utils.litellm_logging import Logging
|
||||
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
|
||||
from litellm.responses.streaming_iterator import drain_logging_tasks
|
||||
from litellm.utils import _schedule_async_success_logging
|
||||
|
||||
now: Final = datetime(2026, 1, 1, tzinfo=timezone.utc)
|
||||
logging_obj: Final = Logging(
|
||||
model="gpt-4o-mini",
|
||||
messages=[],
|
||||
stream=False,
|
||||
call_type="acompletion",
|
||||
start_time=now,
|
||||
litellm_call_id="shutdown-enqueue",
|
||||
function_id="shutdown-enqueue",
|
||||
)
|
||||
logging_obj.defer_async_logging = deferred
|
||||
_schedule_async_success_logging(logging_obj, ModelResponse(), now, now, False)
|
||||
if deferred:
|
||||
assert not logging_obj.async_success_log_enqueued
|
||||
assert logging_obj.enqueue_deferred_logging is not None
|
||||
logging_obj.enqueue_deferred_logging()
|
||||
assert "has_logged_async_success" not in logging_obj.model_call_details
|
||||
await drain_logging_tasks()
|
||||
assert GLOBAL_LOGGING_WORKER.pending_count > 0, "Tracked enqueue must finish before worker draining starts"
|
||||
await GLOBAL_LOGGING_WORKER.flush()
|
||||
assert logging_obj.model_call_details["has_logged_async_success"] is True
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue