scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-12 08:42:18 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 207

fix(router): a cancelled or retried attempt releases its io-token reservation and least-busy count (#50486)

Browse source 
* fix(router): a cancelled or retried attempt releases its io-token reservation and least-busy count

A routed attempt cancelled mid-call (client disconnect with cancel_on_disconnect) kept its ITPM/OTPM
reservation and its least-busy in-flight count, and retries sharing one logging object leaked the
count of every attempt after the first. Each async attempt now runs under its own attempt id, so the
attempt and its own callbacks release what it holds exactly once, in tasks a cancellation can't stop.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(router): settle an io-token reservation once even through a logged copy of its metadata

/v1/messages and /v1/responses stash the reservation in litellm_metadata, which the logging object
copies. A cancelled attempt refunded it, then the proxy's failure logging refunded the copy again.
Each stash now carries an id, and a reservation already settled under that id is not settled again.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(router): keep spent tokens counted, refund a previous attempt's stash, and fit every call in flight

A request cancelled after its provider answered now keeps the tokens it spent counted instead of
refunding them. A retry refunds the reservation its previous attempt left behind even though it runs
under a new attempt. The least-busy reservation store holds twice the keys, since each call in flight
now keeps one more.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(router): never evict a settled reservation id before its counters' window has passed

Settled ids lived in a capacity-bounded cache, so past 10,000 settlements in two minutes a late copy
of an evicted reservation was refunded a second time. They now leave only when their window expires.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(router): start releases eagerly and keep attempt settlement out of router_utils

A release that never waits (an in-memory counter) now finishes before its caller goes on, as it did
when the caller awaited it inline, so a websocket turn's usage is counted by the time its logging is
flushed. Each reservation carries its own id, and the attempt settlement module moves to
litellm_core_utils, since litellm/utils.py may not import from router_utils.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

* fix(router): count only finished answers as spent, and declare the least-busy redis family

A stream's wrapper returning is not an answer: a stream cancelled before its first chunk now hands
back its whole reservation. Only a non-streaming result marks its attempt as answered. The module
level least-busy shared decrement carries its redis family, as every direct redis caller must.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

---------

Co-authored-by: Joshua Valluru <326636767+joshua-berri@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
joshua-berri 2026-11-07 16:54:26 -08:00 • committed by GitHub
parent c024a4d2e3
commit 3c3f266322
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
9 changed files with 995 additions and 189 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

90
litellm/litellm_core_utils/attempt_settlement.py Normal file
View file

@ -0,0 +1,90 @@
"""
What a routed attempt holds (io-token reservations, least-busy in-flight counts) is released exactly
once, by whichever of the attempt itself or its logging callbacks gets there first, and the release
runs to completion even when the request is cancelled mid-release.
"""
import asyncio
import contextlib
import sys
import uuid
from collections.abc import Coroutine, Generator, Sequence
from contextvars import ContextVar
from typing import Final
from litellm._logging import verbose_router_logger
ATTEMPT_SETTLEMENT_GRACE_SECONDS: Final = 5.0
class RoutedAttempt:
__slots__ = ("answered", "attempt_id")
def __init__(self) -> None:
self.attempt_id: Final = uuid.uuid4().hex
self.answered: bool = False
_ATTEMPT: Final[ContextVar[RoutedAttempt | None]] = ContextVar("routed_attempt", default=None)
_RELEASING: Final[set["asyncio.Task[None]"]] = set() # mutable-ok: a release leaves it when it finishes
def current_attempt_id() -> str | None:
attempt: Final = _ATTEMPT.get()
return None if attempt is None else attempt.attempt_id
def mark_attempt_answered() -> None:
"""The provider answered the current routed attempt: what it reserved was spent, not just held."""
attempt: Final = _ATTEMPT.get()
if attempt is not None:
attempt.answered = True
@contextlib.contextmanager
def routed_attempt() -> Generator[RoutedAttempt, None, None]:
attempt: Final = RoutedAttempt()
token: Final = _ATTEMPT.set(attempt)
try:
yield attempt
finally:
_ATTEMPT.reset(token)
def _finished(task: "asyncio.Task[None]") -> None:
_RELEASING.discard(task)
if task.cancelled():
verbose_router_logger.warning("an attempt's release was cancelled before it finished")
return
error: Final = task.exception()
if error is not None:
verbose_router_logger.warning("an attempt's release failed: %s", error)
def _start(release: Coroutine[object, object, None]) -> "asyncio.Task[None]":
"""Run the release up to its first real wait right away, so one that never waits (an in-memory
counter) is done before its caller goes on, as it was when the caller awaited it inline."""
loop: Final = asyncio.get_running_loop()
if sys.version_info >= (3, 12):
return asyncio.Task(release, loop=loop, eager_start=True)
return loop.create_task(release)
def release_in_background(release: Coroutine[object, object, None]) -> "asyncio.Task[None]":
task: Final = _start(release)
_RELEASING.add(task)
task.add_done_callback(_finished)
return task
async def released(releases: Sequence["asyncio.Task[None]"], grace: float | None = None) -> None:
"""Wait for ``releases`` without ever cancelling them: a cancelled waiter leaves them running."""
if all(release.done() for release in releases):
return
_, pending = await asyncio.wait(releases, timeout=grace)
if pending:
verbose_router_logger.warning(
"%s attempt release(s) still running after %ss; they keep running until they finish",
len(pending),
grace,
)

13
litellm/litellm_core_utils/internal_call_metadata.py
View file

@ -46,7 +46,18 @@ ITPM_RESERVED_KEY: Final = "_litellm_itpm_reserved"
OTPM_RESERVED_KEY: Final = "_litellm_otpm_reserved"
ITPM_CACHE_KEY: Final = "_litellm_itpm_cache_key"
OTPM_CACHE_KEY: Final = "_litellm_otpm_cache_key"
IO_TOKEN_RESERVATION_KEYS: Final = frozenset({ITPM_RESERVED_KEY, OTPM_RESERVED_KEY, ITPM_CACHE_KEY, OTPM_CACHE_KEY})
IO_TOKEN_ATTEMPT_KEY: Final = "_litellm_io_token_attempt"
IO_TOKEN_RESERVATION_ID_KEY: Final = "_litellm_io_token_reservation_id"
IO_TOKEN_RESERVATION_KEYS: Final = frozenset(
{
ITPM_RESERVED_KEY,
OTPM_RESERVED_KEY,
ITPM_CACHE_KEY,
OTPM_CACHE_KEY,
IO_TOKEN_ATTEMPT_KEY,
IO_TOKEN_RESERVATION_ID_KEY,
}
)
MODEL_ACCESS_GROUP_METADATA_KEY: Final = "user_api_key_matched_model_access_groups"
"""Where auth records the model access groups that authorized the request, for the spend writer.

55
litellm/router_rust/core.py
View file

@ -47,6 +47,12 @@ from litellm.litellm_core_utils.add_retry_fallback_headers import (
prepare_response_for_header_attachment,
)
from litellm.litellm_core_utils.agentic_followup_failure import is_agentic_followup_failure
from litellm.litellm_core_utils.attempt_settlement import (
ATTEMPT_SETTLEMENT_GRACE_SECONDS,
RoutedAttempt,
released,
routed_attempt,
)
from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type
from litellm.litellm_core_utils.core_helpers import (
as_str_mapping,
@ -209,7 +215,14 @@ from litellm.router_rust.surface.settings import (
request_timeout,
settings_view,
)
from litellm.router_strategy.least_busy import CALL_IDENTITY, call_id_of, owning_picks, pick_reserves, picks_reserve
from litellm.router_strategy.least_busy import (
CALL_IDENTITY,
call_id_of,
owning_picks,
pick_reserves,
picks_reserve,
release_started_attempt,
)
from litellm.router_utils.cooldown_handlers import get_cooldown_deployment_info, mark_deployment_selection_failure
from litellm.router_utils.endpoint_eligibility import endpoint_for_call, selected_endpoint
from litellm.router_utils.kubernetes_pod_discovery import (
@ -218,7 +231,7 @@ from litellm.router_utils.kubernetes_pod_discovery import (
resolve_pods_after_bound,
)
from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import EncryptedContentAffinityCheck
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_io_token_refund_failure
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import release_io_token_reservation
from litellm.router_utils.router_callbacks.track_deployment_metrics import find_deployment_metadata
from litellm.router_utils.routing_read_batch import RoutingReadBatch
from litellm.types.router import (
@ -5150,6 +5163,28 @@ class _Driver:
async def attempt(self, step: Mapping[str, object]) -> Mapping[str, object]:
if self._passthrough is not None and not step.get("deployment_id"):
return await self._passthrough_attempt(step)
with routed_attempt() as attempt:
return await self._routed_attempt(step, attempt)
def _release(
self, call: AttemptCall, attempt: RoutedAttempt, *, refund_io_tokens: bool
) -> "tuple[asyncio.Task[None], ...]":
"""What a failed or cancelled attempt still holds, whatever its logging object (shared across
attempts, it logs only the first failure). Tokens a provider already spent are kept counted."""
core: Final = self._core
releases: Final = (
release_started_attempt(core._selectors.reservations, core._cache, attempt.attempt_id), # pyright: ignore[reportPrivateUsage] # the driver serves its own core
release_io_token_reservation(
core._cache, # pyright: ignore[reportPrivateUsage] # see above
call.kwargs, # pyright: ignore[reportArgumentType] # the call's kwargs dict
parent_otel_span=get_parent_otel_span_from_kwargs(call.kwargs),
)
if refund_io_tokens
else None,
)
return tuple(release for release in releases if release is not None)
async def _routed_attempt(self, step: Mapping[str, object], attempt: RoutedAttempt) -> Mapping[str, object]:
step = await self._resolve_pods(step) # rebind-ok: the picked deployment, resolved to a pod
deployment: Mapping[str, object] = as_str_mapping(step.get("deployment")) or {} # rebind-ok: the picked one
call: AttemptCall | None = None # rebind-ok: built inside the try
@ -5215,16 +5250,16 @@ class _Driver:
# `self.success_calls[model_name] += 1`, inside the try: a counter replaced with a
# plain dict fails the attempt.
self._core._endpoints.counters.succeeded(self._plan, call) # pyright: ignore[reportPrivateUsage] # see above
except asyncio.CancelledError:
if call is not None:
await released(
self._release(call, attempt, refund_io_tokens=not attempt.answered),
ATTEMPT_SETTLEMENT_GRACE_SECONDS,
)
raise
except Exception as error: # noqa: BLE001 # every attempt failure is classified and handed to the Rust call loop
if call is not None:
# `make_call`: a failed call hands back the io tokens it reserved, whatever its
# logging object (shared across attempts, it logs only the first failure).
with contextlib.suppress(Exception):
await async_io_token_refund_failure(
self._core._cache, # pyright: ignore[reportPrivateUsage] # the driver serves its own core
call.kwargs, # pyright: ignore[reportArgumentType] # the call's kwargs dict
parent_otel_span=get_parent_otel_span_from_kwargs(call.kwargs),
)
await released(self._release(call, attempt, refund_io_tokens=True))
failed: Final = self._failed(step, deployment, call, error)
if reading_stream is None or not logs_failure:
return {

144
litellm/router_strategy/least_busy.py
View file

@ -1,3 +1,4 @@
import asyncio
import contextlib
import logging
import threading
@ -15,14 +16,16 @@ from litellm._logging import verbose_router_logger
from litellm.caching.caching import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import log_redis_failure
from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.attempt_settlement import current_attempt_id, release_in_background, released
from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type
from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET
from litellm.types.utils import CallTypes
IN_FLIGHT_COUNT_TTL_SECONDS: Final = 60 * 60
SELECTION_RESERVATION_TTL_SECONDS: Final = 30.0
SELECTION_RESERVATION_MAX_KEYS: Final = 10_000
SELECTION_RESERVATION_MAX_KEYS: Final = 20_000
_RESERVATIONS_LOCK: Final = threading.Lock()
_PICK_RESERVES: Final[ContextVar[bool]] = ContextVar("least_busy_pick_reserves")
_PICK_OWNER: Final[ContextVar[str | None]] = ContextVar("least_busy_pick_owner", default=None)
@ -76,7 +79,9 @@ _CALL_KWARGS: Final = TypeAdapter(_CallKwargs)
CALL_IDENTITY: Final = TypeAdapter(CallIdentity)
_DEPLOYMENTS: Final = TypeAdapter(list[_Deployment])
_MEMORY_COUNTS: Final[TypeAdapter[tuple[float | None, ...] | None]] = TypeAdapter(tuple[float | None, ...] | None)
_ATTEMPT_TIMES: Final[TypeAdapter[tuple[float, ...] | None]] = TypeAdapter(tuple[float, ...] | None)
_Start: TypeAlias = tuple[float, str | None]
_STARTS: Final[TypeAdapter[tuple[_Start, ...] | None]] = TypeAdapter(tuple[_Start, ...] | None)
_STARTED_ATTEMPT: Final[TypeAdapter[tuple[str, str, str] | None]] = TypeAdapter(tuple[str, str, str] | None)
_Pick: TypeAlias = tuple[float, str | None]
_PICKS: Final[TypeAdapter[tuple[_Pick, ...] | None]] = TypeAdapter(tuple[_Pick, ...] | None)
_PICKED_DEPLOYMENTS: Final[TypeAdapter[tuple[str, ...] | None]] = TypeAdapter(tuple[str, ...] | None)
@ -100,7 +105,11 @@ def _routed_deployment_id(metadata: _Metadata | None) -> str | int | None:
def _deployment_ref(kwargs: Mapping[str, object]) -> tuple[str, str] | None:
call_type: Final = kwargs.get("call_type")
if is_batch_retrieve_call_type(call_type) or call_type == CallTypes.arealtime.value:
if (
is_batch_retrieve_call_type(call_type)
or call_type == CallTypes.arealtime.value
or kwargs.get(LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL)
):
return None
try:
call: Final = _CALL_KWARGS.validate_python(kwargs)
@ -133,6 +142,10 @@ def _started_attempts_key(call_id: str, ref: tuple[str, str]) -> str:
return f"least_busy_started:{call_id}:{_request_count_key(*ref)}"
def _started_attempt_key(attempt_id: str) -> str:
return f"least_busy_attempt:{attempt_id}"
def _picked_deployments_key(call_id: str) -> str:
return f"least_busy_picked:{call_id}"
@ -211,8 +224,8 @@ def _picks(store: _ReservationStore, deployment_id: str) -> tuple[_Pick, ...]:
return _PICKS.validate_python(store.get_cache(_reservation_key(deployment_id))) or ()
def _attempt_times(store: _ReservationStore, key: str) -> tuple[float, ...]:
return _ATTEMPT_TIMES.validate_python(store.get_cache(key)) or ()
def _starts(store: _ReservationStore, key: str) -> tuple[_Start, ...]:
return _STARTS.validate_python(store.get_cache(key)) or ()
def _picked_deployments(store: _ReservationStore, call_id: str) -> tuple[str, ...]:
@ -231,6 +244,60 @@ def _release_call_on(store: _ReservationStore, deployment_id: str, call_id: str)
_put(store, _reservation_key(deployment_id), kept, SELECTION_RESERVATION_TTL_SECONDS)
def _end_started_attempt(store: _ReservationStore, attempt_id: str, *, every: bool) -> tuple[str, str, int] | None:
"""Under ``_RESERVATIONS_LOCK``: end one of the attempt's own starts (a call retried inside one
attempt starts more than once), or ``every`` one still open. Returns the deployment and how many ended."""
started: Final = _STARTED_ATTEMPT.validate_python(store.get_cache(_started_attempt_key(attempt_id)))
if started is None:
return None
call_id, model_group, deployment_id = started
key: Final = _started_attempts_key(call_id, (model_group, deployment_id))
starts: Final = _starts(store, key)
own: Final = tuple(index for index, start in enumerate(starts) if start[1] == attempt_id)
ending: Final = frozenset(own if every else own[:1])
_put(
store,
key,
tuple(start for index, start in enumerate(starts) if index not in ending),
IN_FLIGHT_COUNT_TTL_SECONDS,
)
if len(ending) == len(own):
store.delete_cache(_started_attempt_key(attempt_id))
return (model_group, deployment_id, len(ending)) if ending else None
def _count_locally(cache: DualCache, key: str, delta: int) -> None:
try:
local: Final = cache.increment_cache(key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS)
if local < 0:
cache.set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS)
except TypeError as e:
_warn_unwritable(key, e)
@with_service_target(ROUTER_USAGE_TARGET)
async def _async_count_shared(cache: DualCache, key: str, delta: int) -> None:
redis_cache: Final = cache.redis_cache
if redis_cache is None:
return
try:
await redis_cache.async_increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS)
except Exception as e:
_warn_unwritable(key, e)
def release_started_attempt(store: _ReservationStore, cache: DualCache, attempt_id: str) -> "asyncio.Task[None] | None":
"""End a routed attempt that never reached its success or failure callback (refused, failed
without logging, or cancelled): its in-flight count drops now, the shared one in the background."""
with _RESERVATIONS_LOCK:
ended: Final = _end_started_attempt(store, attempt_id, every=True)
if ended is None:
return None
model_group, deployment_id, count = ended
_count_locally(cache, _request_count_key(model_group, deployment_id), -count)
return release_in_background(_async_count_shared(cache, _request_count_key(model_group, deployment_id), -count))
def release_unstarted_picks(store: _ReservationStore, call_id: str | None) -> None:
if call_id is None:
return
@ -312,29 +379,47 @@ class LeastBusyLoggingHandler(CustomLogger):
return healthy_deployments[index]
def _start_attempt(self, ref: tuple[str, str], call_id: str | None) -> None:
attempt_id: Final = current_attempt_id()
with _RESERVATIONS_LOCK:
now: Final = self._clock()
self._take_reservation(ref[1], call_id, now)
self._count_locally(_request_count_key(*ref), 1)
_count_locally(self.router_cache, _request_count_key(*ref), 1)
if call_id is not None:
key: Final = _started_attempts_key(call_id, ref)
started: Final = (*_attempt_times(self._reservations, key), now)
started: Final = (*_starts(self._reservations, key), (now, attempt_id))
_put(self._reservations, key, started, IN_FLIGHT_COUNT_TTL_SECONDS)
if attempt_id is not None:
_put(
self._reservations,
_started_attempt_key(attempt_id),
(call_id, *ref),
IN_FLIGHT_COUNT_TTL_SECONDS,
)
def _end_attempt(self, ref: tuple[str, str], call_id: str | None) -> bool:
def _end_attempt(self, ref: tuple[str, str], call_id: str | None) -> tuple[str, str] | None:
if call_id is None:
with _RESERVATIONS_LOCK:
self._count_locally(_request_count_key(*ref), -1)
return True
_count_locally(self.router_cache, _request_count_key(*ref), -1)
return ref
attempt_id: Final = current_attempt_id()
if attempt_id is not None:
with _RESERVATIONS_LOCK:
ended: Final = _end_started_attempt(self._reservations, attempt_id, every=False)
if ended is not None:
_count_locally(self.router_cache, _request_count_key(ended[0], ended[1]), -1)
return None if ended is None else (ended[0], ended[1])
key: Final = _started_attempts_key(call_id, ref)
with _RESERVATIONS_LOCK:
started: Final = _attempt_times(self._reservations, key)
started: Final = _starts(self._reservations, key)
if started:
oldest_attempt: Final = started[0][1]
if oldest_attempt is not None:
self._reservations.delete_cache(_started_attempt_key(oldest_attempt))
_put(self._reservations, key, started[1:], IN_FLIGHT_COUNT_TTL_SECONDS)
self._count_locally(_request_count_key(*ref), -1)
_count_locally(self.router_cache, _request_count_key(*ref), -1)
else:
self._take_reservation(ref[1], call_id, self._clock())
return bool(started)
return ref if started else None
@with_service_target(ROUTER_USAGE_TARGET)
def log_pre_api_call(self, model: str, messages: object, kwargs: Mapping[str, object]) -> None:
@ -418,15 +503,17 @@ class LeastBusyLoggingHandler(CustomLogger):
def _settle(self, kwargs: Mapping[str, object]) -> None:
ref: Final = _deployment_ref(kwargs)
if ref is None or not self._end_attempt(ref, call_id_of(kwargs)):
return
self._count_shared(_request_count_key(*ref), -1)
ended: Final = None if ref is None else self._end_attempt(ref, call_id_of(kwargs))
if ended is not None:
self._count_shared(_request_count_key(*ended), -1)
async def _async_settle(self, kwargs: Mapping[str, object]) -> None:
ref: Final = _deployment_ref(kwargs)
if ref is None or not self._end_attempt(ref, call_id_of(kwargs)):
return
await self._async_count_shared(_request_count_key(*ref), -1)
ended: Final = None if ref is None else self._end_attempt(ref, call_id_of(kwargs))
if ended is not None:
await released(
(release_in_background(_async_count_shared(self.router_cache, _request_count_key(*ended), -1)),)
)
def _this_workers_counts(self, keys: tuple[str, ...]) -> tuple[int, ...]:
return _local_counts(self.router_cache.batch_get_cache(list(keys), local_only=True), keys)
@ -434,16 +521,6 @@ class LeastBusyLoggingHandler(CustomLogger):
def _at_least_this_workers(self, keys: tuple[str, ...], shared: tuple[int, ...]) -> tuple[int, ...]:
return tuple(max(counts) for counts in zip(shared, self._this_workers_counts(keys), strict=True))
def _count_locally(self, key: str, delta: int) -> None:
try:
local: Final = self.router_cache.increment_cache(
key, delta, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS
)
if local < 0:
self.router_cache.set_cache(key, 0, local_only=True, ttl=IN_FLIGHT_COUNT_TTL_SECONDS)
except TypeError as e:
_warn_unwritable(key, e)
def _count_shared(self, key: str, delta: int) -> None:
redis_cache: Final = self.router_cache.redis_cache
if redis_cache is None:
@ -452,12 +529,3 @@ class LeastBusyLoggingHandler(CustomLogger):
redis_cache.increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS)
except Exception as e:
_warn_unwritable(key, e)
async def _async_count_shared(self, key: str, delta: int) -> None:
redis_cache: Final = self.router_cache.redis_cache
if redis_cache is None:
return
try:
await redis_cache.async_increment_with_floor(key, delta, IN_FLIGHT_COUNT_TTL_SECONDS)
except Exception as e:
_warn_unwritable(key, e)

269
litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py
View file

@ -11,9 +11,13 @@ Used by ModelRateLimitingCheck when a deployment sets itpm/otpm.
from __future__ import annotations
import asyncio
import contextlib
import contextvars
from collections.abc import Mapping, MutableMapping
import threading
import time
from collections import OrderedDict
from collections.abc import Callable, Mapping, MutableMapping
from dataclasses import dataclass
from datetime import datetime
from types import MappingProxyType
@ -26,11 +30,15 @@ import litellm
from litellm import token_counter
from litellm._internal_context import with_service_target
from litellm._logging import verbose_router_logger
from litellm._uuid import uuid
from litellm.caching.dual_cache import DualCache
from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError
from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION
from litellm.litellm_core_utils.attempt_settlement import current_attempt_id, release_in_background, released
from litellm.litellm_core_utils.core_helpers import coerce_token_limit
from litellm.litellm_core_utils.internal_call_metadata import (
IO_TOKEN_ATTEMPT_KEY,
IO_TOKEN_RESERVATION_ID_KEY,
IO_TOKEN_RESERVATION_KEYS,
ITPM_CACHE_KEY,
ITPM_RESERVED_KEY,
@ -55,6 +63,31 @@ else:
RoutingArgsTTL: Final = 60
class _SettledReservations:
"""The ids of reservations settled within the counters' window, so a copy of a stash that outlives
the dict it was claimed from (the logging object copies litellm_metadata) is never settled twice.
Ids leave only once their window has passed, never to make room."""
def __init__(self, window: float, clock: Callable[[], float] = time.monotonic) -> None:
self._window: Final = window
self._clock: Final = clock
self._expiries: Final[OrderedDict[str, float]] = OrderedDict() # mutable-ok: pruned oldest first as ids expire
self._lock: Final = threading.Lock()
def claim(self, reservation_id: str) -> bool:
with self._lock:
now: Final = self._clock()
while self._expiries and next(iter(self._expiries.values())) <= now:
_ = self._expiries.popitem(last=False)
if reservation_id in self._expiries:
return False
self._expiries[reservation_id] = now + self._window
return True
_SETTLED_RESERVATIONS: Final = _SettledReservations(window=RoutingArgsTTL * 2)
@dataclass(frozen=True, slots=True)
class _IoTokenRateLimitRequest:
kwargs: dict[str, object] | None
@ -355,9 +388,10 @@ def _resolve_reconcile_usage_tokens(
def _reservation(
itpm_reserved: int, otpm_reserved: int, itpm_cache_key: str | None, otpm_cache_key: str | None
) -> Mapping[str, object]:
) -> Mapping[str, int | str | None]:
return MappingProxyType(
{
IO_TOKEN_RESERVATION_ID_KEY: uuid.uuid4().hex,
ITPM_RESERVED_KEY: itpm_reserved,
OTPM_RESERVED_KEY: otpm_reserved,
ITPM_CACHE_KEY: itpm_cache_key,
@ -370,6 +404,7 @@ def _stash_reservation_in_metadata(request: _IoTokenRateLimitRequest | None, res
metadata: Final = None if request is None or not request.kwargs else request.kwargs.get(request.metadata_channel)
if isinstance(metadata, dict):
metadata.update(reservation)
metadata[IO_TOKEN_ATTEMPT_KEY] = current_attempt_id()
def _key_for_this_minute(reserved_key: str) -> str:
@ -387,12 +422,16 @@ def _reconcile_increments(
)
def _extract_reservation(reservation: Mapping[str, int | str | None]) -> tuple[int, int, str | None, str | None]:
def _reserved_tokens(value: object) -> int:
return value if isinstance(value, int) and not isinstance(value, bool) else 0
def _extract_reservation(reservation: Mapping[str, object]) -> tuple[int, int, str | None, str | None]:
itpm_cache_key: Final = reservation.get(ITPM_CACHE_KEY)
otpm_cache_key: Final = reservation.get(OTPM_CACHE_KEY)
return (
int(reservation.get(ITPM_RESERVED_KEY, 0) or 0),
int(reservation.get(OTPM_RESERVED_KEY, 0) or 0),
_reserved_tokens(reservation.get(ITPM_RESERVED_KEY)),
_reserved_tokens(reservation.get(OTPM_RESERVED_KEY)),
itpm_cache_key if isinstance(itpm_cache_key, str) else None,
otpm_cache_key if isinstance(otpm_cache_key, str) else None,
)
@ -421,11 +460,31 @@ def _reservation_channels(kwargs: Mapping[str, object] | None) -> tuple[object,
return top_level + from_params + from_logging_object
def _read_reservation_from_kwargs(kwargs: Mapping[str, object] | None) -> tuple[int, int, str | None, str | None]:
for channel_dict in _reservation_channels(kwargs):
if isinstance(channel_dict, dict) and ITPM_RESERVED_KEY in channel_dict:
return _extract_reservation(channel_dict)
return 0, 0, None, None
def _stashed_reservation(kwargs: Mapping[str, object] | None) -> Mapping[str, object] | None:
channels: Final = (_as_mutable_mapping(channel) for channel in _reservation_channels(kwargs))
return next((channel for channel in channels if channel is not None and ITPM_RESERVED_KEY in channel), None)
def _claim_reservation_from_kwargs(
kwargs: Mapping[str, object] | None, *, any_attempt: bool = False
) -> tuple[int, int, str | None, str | None]:
"""
Take the stashed reservation so nothing else settles it. Under a routed attempt only that attempt's
own reservation is taken, unless ``any_attempt``: a late callback of an earlier attempt sharing these
kwargs leaves the next attempt's reservation alone.
"""
stashed: Final = _stashed_reservation(kwargs)
if stashed is None:
return 0, 0, None, None
attempt: Final = current_attempt_id()
if not any_attempt and attempt is not None and stashed.get(IO_TOKEN_ATTEMPT_KEY) not in (None, attempt):
return 0, 0, None, None
reservation: Final = _extract_reservation(stashed)
reservation_id: Final = stashed.get(IO_TOKEN_RESERVATION_ID_KEY)
_clear_reservation_from_kwargs(kwargs)
if isinstance(reservation_id, str) and not _SETTLED_RESERVATIONS.claim(reservation_id):
return 0, 0, None, None
return reservation
def _clear_reservation_from_kwargs(kwargs: Mapping[str, object] | None) -> None:
@ -626,11 +685,50 @@ async def async_io_token_pre_call_check(
return deployment
itpm_key, otpm_key = cache_keys
itpm_reserved = 0
otpm_reserved = 0
# The reservation runs in a task of its own: a request cancelled mid-reservation never knows how much
# it took, so the task finishes and then hands back whatever it reserved.
reserving: Final = asyncio.get_running_loop().create_task(
_async_reserve(
dual_cache,
itpm_key,
otpm_key,
_reservation_value(estimated_input, itpm_limit) if itpm_limit is not None else 0,
otpm_to_reserve,
itpm_limit,
otpm_limit,
parent_otel_span,
)
)
try:
reservation: Final = await asyncio.shield(reserving)
except asyncio.CancelledError:
_ = release_in_background(_refund_once_reserved(dual_cache, reserving, parent_otel_span))
raise
_stash_reservation_in_metadata(request, reservation)
return deployment
async def _refund_once_reserved(
dual_cache: DualCache, reserving: asyncio.Task[Mapping[str, int | str | None]], parent_otel_span: Span | None
) -> None:
try:
reservation: Final = await reserving
except Exception: # noqa: BLE001 # a refused or failed reservation holds nothing
return
await _async_increment_all(dual_cache, _refunds(_extract_reservation(reservation)), parent_otel_span)
async def _async_reserve(
dual_cache: DualCache,
itpm_key: str,
otpm_key: str,
itpm_reserved: int,
otpm_reserved: int,
itpm_limit: int | None,
otpm_limit: int | None,
parent_otel_span: Span | None,
) -> Mapping[str, int | str | None]:
if itpm_limit is not None:
itpm_reserved = _reservation_value(estimated_input, itpm_limit)
await _increment_with_rollback(
dual_cache,
itpm_key,
@ -641,7 +739,6 @@ async def async_io_token_pre_call_check(
)
if otpm_limit is not None:
otpm_reserved = otpm_to_reserve
try:
await _increment_with_rollback(
dual_cache,
@ -664,16 +761,12 @@ async def async_io_token_pre_call_check(
)
raise
_stash_reservation_in_metadata(
request,
_reservation(
itpm_reserved,
otpm_reserved,
itpm_key if itpm_limit is not None else None,
otpm_key if otpm_limit is not None else None,
),
return _reservation(
itpm_reserved,
otpm_reserved if otpm_limit is not None else 0,
itpm_key if itpm_limit is not None else None,
otpm_key if otpm_limit is not None else None,
)
return deployment
async def _async_redis_usage(redis_cache: RedisCache | None, key: str, parent_otel_span: Span | None) -> int | None:
@ -732,25 +825,22 @@ def io_token_reconcile_success(
) -> None:
request_kwargs: Final[Mapping[str, object] | None] = kwargs
response: Final[object] = response_obj
reservation: Final = _read_reservation_from_kwargs(request_kwargs)
reservation: Final = _claim_reservation_from_kwargs(request_kwargs)
itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation
if itpm_key is None and otpm_key is None:
return
billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(request_kwargs, response)
try:
if usage_resolved:
for key, delta in _reconcile_increments(reservation, billable_input, completion_tokens):
dual_cache.increment_cache(key=key, value=delta, ttl=RoutingArgsTTL)
else:
verbose_router_logger.debug(
"[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)",
itpm_reserved,
otpm_reserved,
)
finally:
_clear_reservation_from_kwargs(request_kwargs)
if usage_resolved:
for key, delta in _reconcile_increments(reservation, billable_input, completion_tokens):
dual_cache.increment_cache(key=key, value=delta, ttl=RoutingArgsTTL)
else:
verbose_router_logger.debug(
"[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)",
itpm_reserved,
otpm_reserved,
)
verbose_router_logger.debug(
"[IO TOKEN LIMIT] reconciled (usage_resolved=%s, itpm_reserved=%s, billable_input=%s, otpm_reserved=%s, output=%s)",
@ -772,7 +862,7 @@ async def async_io_token_reconcile_success(
) -> None:
request_kwargs: Final[Mapping[str, object] | None] = kwargs
response: Final[object] = response_obj
reservation: Final = _read_reservation_from_kwargs(request_kwargs)
reservation: Final = _claim_reservation_from_kwargs(request_kwargs)
itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation
if itpm_key is None and otpm_key is None:
return
@ -781,22 +871,25 @@ async def async_io_token_reconcile_success(
# A refund goes back to the minute that held the reservation, while usage above the reservation is charged
# to the minute the response finished in, so a call that runs past its minute still counts against a live
# window. Always clear the stash afterwards (even if an increment throws) so a retry or a duplicate success
# event can't re-process it.
try:
if usage_resolved:
for key, delta in _reconcile_increments(reservation, billable_input, completion_tokens):
await dual_cache.async_increment_cache(
key=key, value=delta, ttl=RoutingArgsTTL, parent_otel_span=parent_otel_span
)
else:
verbose_router_logger.debug(
"[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)",
itpm_reserved,
otpm_reserved,
# window.
if usage_resolved:
await released(
(
release_in_background(
_async_increment_all(
dual_cache,
_reconcile_increments(reservation, billable_input, completion_tokens),
parent_otel_span,
)
),
)
finally:
_clear_reservation_from_kwargs(request_kwargs)
)
else:
verbose_router_logger.debug(
"[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)",
itpm_reserved,
otpm_reserved,
)
verbose_router_logger.debug(
"[IO TOKEN LIMIT] reconciled (usage_resolved=%s, itpm_reserved=%s, billable_input=%s, otpm_reserved=%s, output=%s)",
@ -812,24 +905,14 @@ async def async_io_token_reconcile_success(
def io_token_refund_failure(
dual_cache: DualCache,
kwargs: Mapping[str, object] | None,
*,
any_attempt: bool = False,
) -> None:
request_kwargs: Final[Mapping[str, object] | None] = kwargs
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs)
if itpm_key is None and otpm_key is None:
return
if itpm_key is not None and itpm_reserved > 0:
dual_cache.increment_cache(
key=itpm_key,
value=-itpm_reserved,
ttl=RoutingArgsTTL,
)
if otpm_key is not None and otpm_reserved > 0:
dual_cache.increment_cache(
key=otpm_key,
value=-otpm_reserved,
ttl=RoutingArgsTTL,
)
_clear_reservation_from_kwargs(request_kwargs)
reservation: Final = _claim_reservation_from_kwargs(request_kwargs, any_attempt=any_attempt)
itpm_reserved, otpm_reserved, _, _ = reservation
for key, delta in _refunds(reservation):
dual_cache.increment_cache(key=key, value=delta, ttl=RoutingArgsTTL)
verbose_router_logger.debug("[IO TOKEN LIMIT] refunded ITPM=%s OTPM=%s", itpm_reserved, otpm_reserved)
@ -857,7 +940,7 @@ def refund_stale_reservation_before_retry(dual_cache: DualCache, kwargs: Mapping
"""
if not kwargs:
return
io_token_refund_failure(dual_cache, kwargs)
io_token_refund_failure(dual_cache, kwargs, any_attempt=True)
@with_service_target(ROUTER_USAGE_TARGET)
@ -867,26 +950,38 @@ async def async_io_token_refund_failure(
*,
parent_otel_span: Span | None = None,
) -> None:
request_kwargs: Final[Mapping[str, object] | None] = kwargs
itpm_reserved, otpm_reserved, itpm_key, otpm_key = _read_reservation_from_kwargs(request_kwargs)
if itpm_key is None and otpm_key is None:
return
if itpm_key is not None and itpm_reserved > 0:
refund: Final = release_io_token_reservation(dual_cache, kwargs, parent_otel_span=parent_otel_span)
if refund is not None:
await released((refund,))
def release_io_token_reservation(
dual_cache: DualCache,
kwargs: Mapping[str, object] | None,
*,
parent_otel_span: Span | None = None,
) -> asyncio.Task[None] | None:
reservation: Final = _claim_reservation_from_kwargs(kwargs)
refunds: Final = _refunds(reservation)
if not refunds:
return None
verbose_router_logger.debug("[IO TOKEN LIMIT] refunding ITPM=%s OTPM=%s", reservation[0], reservation[1])
return release_in_background(_async_increment_all(dual_cache, refunds, parent_otel_span))
def _refunds(reservation: tuple[int, int, str | None, str | None]) -> tuple[tuple[str, int], ...]:
itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation
held: Final = ((itpm_key, itpm_reserved), (otpm_key, otpm_reserved))
return tuple((key, -reserved) for key, reserved in held if key is not None and reserved > 0)
async def _async_increment_all(
dual_cache: DualCache, increments: tuple[tuple[str, int], ...], parent_otel_span: Span | None
) -> None:
for key, delta in increments:
await dual_cache.async_increment_cache(
key=itpm_key,
value=-itpm_reserved,
ttl=RoutingArgsTTL,
parent_otel_span=parent_otel_span,
key=key, value=delta, ttl=RoutingArgsTTL, parent_otel_span=parent_otel_span
)
if otpm_key is not None and otpm_reserved > 0:
await dual_cache.async_increment_cache(
key=otpm_key,
value=-otpm_reserved,
ttl=RoutingArgsTTL,
parent_otel_span=parent_otel_span,
)
_clear_reservation_from_kwargs(request_kwargs)
verbose_router_logger.debug("[IO TOKEN LIMIT] refunded ITPM=%s OTPM=%s", itpm_reserved, otpm_reserved)
def build_io_token_rate_limit_headers(

6
litellm/utils.py
View file

@ -88,6 +88,7 @@ from litellm.litellm_core_utils.allowed_tools_tool_choice import (
CHAT_ALLOWED_TOOLS_LAYOUT,
parse_chat_allowed_tools_tool_choice,
)
from litellm.litellm_core_utils.attempt_settlement import mark_attempt_answered
from litellm.litellm_core_utils.core_helpers import (
bind_budget_reservation_to_callbacks,
has_attributes,
@ -2515,7 +2516,9 @@ def client(original_function):
exception=deployment_error,
call_type=call_type,
)
except BaseException: # noqa: S110, BLE001 # hook dispatch - including cancellation mid-await - must never replace the real deployment failure, so there is nothing to do with what it raises
except asyncio.CancelledError:
raise
except BaseException: # noqa: S110, BLE001 # a failing hook must never replace the real deployment failure, so there is nothing to do with what it raises
pass
raise
end_time = datetime.datetime.now()
@ -2557,6 +2560,7 @@ def client(original_function):
)
elif call_type in (CallTypes.arealtime.value, CallTypes.aresponses_websocket.value):
return result
mark_attempt_answered()
### POST-CALL RULES ###
post_call_processing(
original_response=result,

203
tests/unit/router_strategy/test_least_busy.py
View file

@ -16,11 +16,11 @@ import respx
import litellm
from litellm import Router
from litellm.caching.caching import DualCache
from litellm.caching.in_memory_cache import DEFAULT_MAX_SIZE_IN_MEMORY, InMemoryCache
from litellm.caching.redis_cache import RedisCircuitBreakerOpenError
from litellm.constants import LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL
from litellm.litellm_core_utils.attempt_settlement import routed_attempt
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.router_strategy.least_busy import (
IN_FLIGHT_COUNT_TTL_SECONDS,
@ -28,6 +28,7 @@ from litellm.router_strategy.least_busy import (
LeastBusyLoggingHandler,
owning_picks,
picks_reserve,
release_started_attempt,
release_unstarted_picks,
)
from litellm.types.utils import CallTypes, ModelResponse
@ -287,7 +288,9 @@ class OpenBreakerRedis(SharedRedisCounters):
@pytest.mark.asyncio
async def test_an_open_circuit_breaker_falls_back_without_a_warning_per_request(caplog: pytest.LogCaptureFixture) -> None:
async def test_an_open_circuit_breaker_falls_back_without_a_warning_per_request(
caplog: pytest.LogCaptureFixture,
) -> None:
worker: Final = _worker(OpenBreakerRedis())
with caplog.at_level(logging.DEBUG, logger="LiteLLM Router"):
@ -463,9 +466,7 @@ async def test_router_picks_least_busy_deployment_and_completion_restores_counts
if use_async
else router.get_available_deployment(model=ROUTER_GROUP, messages=None)
)
response: Final = await router.acompletion(
model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}]
)
response: Final = await router.acompletion(model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}])
for _ in range(10):
await asyncio.sleep(0)
await GLOBAL_LOGGING_WORKER.flush()
@ -604,7 +605,9 @@ def test_two_threads_picking_at_once_reserve_different_deployments() -> None:
worker: Final = LeastBusyLoggingHandler(router_cache=DualCache(), reservation_store=store)
with store.first_reads_meet(), ThreadPoolExecutor(max_workers=2) as pool:
picks: Final = tuple(pool.map(lambda _: worker.get_available_deployments(GROUP, HEALTHY, reserve=True), range(2)))
picks: Final = tuple(
pool.map(lambda _: worker.get_available_deployments(GROUP, HEALTHY, reserve=True), range(2))
)
assert sorted(pick["model_info"]["id"] for pick in picks if pick is not None) == ["dep-a", "dep-b"]
@ -626,7 +629,6 @@ def test_a_reservation_and_a_release_on_one_deployment_at_once_both_count() -> N
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B
class _SlowLocalRead(InMemoryCache):
"""Answers a count read with the counts from when the read started, once another thread has begun its request."""
@ -706,7 +708,9 @@ async def test_looking_up_a_pass_through_deployment_does_not_hold_it_for_later_p
for _ in range(2)
]
sync_picks: Final = [
router.get_available_deployment_for_pass_through(ROUTER_GROUP, request_kwargs={"metadata": {}})["model_info"]["id"]
router.get_available_deployment_for_pass_through(ROUTER_GROUP, request_kwargs={"metadata": {}})["model_info"][
"id"
]
for _ in range(2)
]
@ -733,7 +737,9 @@ def test_a_cache_hit_leaves_requests_in_flight_on_its_deployment_counted() -> No
assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_B
@pytest.mark.parametrize("outcome", [LeastBusyLoggingHandler.log_success_event, LeastBusyLoggingHandler.log_failure_event])
@pytest.mark.parametrize(
"outcome", [LeastBusyLoggingHandler.log_success_event, LeastBusyLoggingHandler.log_failure_event]
)
def test_a_call_that_ends_before_reaching_the_provider_frees_its_pick(
outcome: Callable[[LeastBusyLoggingHandler, dict[str, object], None, None, None], None],
) -> None:
@ -755,6 +761,16 @@ def test_a_call_the_proxy_rejected_after_logging_its_start_leaves_no_count_behin
assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A
def test_a_proxy_error_logged_on_a_request_whose_failure_was_already_logged_leaves_no_count_behind() -> None:
worker: Final = _worker(None)
worker.log_pre_api_call(
model="m", messages=[], kwargs={**_call_kwargs("dep-a"), LITELLM_LOGGING_NO_UPSTREAM_LLM_CALL: True}
)
assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A
def test_a_call_logged_under_its_callers_group_name_frees_the_pick_made_for_another_group() -> None:
worker: Final = _worker(None)
assert worker.get_available_deployments("tier-group", HEALTHY, reserve=True) is DEPLOYMENT_A
@ -916,7 +932,9 @@ _REJECTED_BATCH: Final = {"error": {"message": "Invalid input file", "type": "in
@pytest.mark.asyncio
@pytest.mark.parametrize("provider_response", [httpx.Response(200, json=_BATCH), httpx.Response(400, json=_REJECTED_BATCH)])
@pytest.mark.parametrize(
"provider_response", [httpx.Response(200, json=_BATCH), httpx.Response(400, json=_REJECTED_BATCH)]
)
async def test_a_batch_create_does_not_keep_its_deployment_busy(
provider_response: httpx.Response, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
@ -926,7 +944,10 @@ async def test_a_batch_create_does_not_keep_its_deployment_busy(
with contextlib.suppress(openai.BadRequestError):
await router.acreate_batch(
model=ROUTER_GROUP, completion_window="24h", endpoint="/v1/chat/completions", input_file_id="file-least-busy"
model=ROUTER_GROUP,
completion_window="24h",
endpoint="/v1/chat/completions",
input_file_id="file-least-busy",
)
await _drain_logging()
@ -1013,7 +1034,6 @@ def test_a_sync_text_completion_does_not_keep_its_deployment_busy(
assert _router_pick(router, TEXT_GROUP) == _called_deployment(route)
@pytest.mark.parametrize("call_identity", _CALL_IDENTITIES)
def test_a_sync_responses_call_that_fails_before_reaching_its_provider_frees_its_pick(
call_identity: dict[str, str],
@ -1025,6 +1045,7 @@ def test_a_sync_responses_call_that_fails_before_reaching_its_provider_frees_its
assert _router_pick(router) == ROUTER_DEPLOYMENT_IDS[0]
@pytest.mark.asyncio
async def test_a_batch_cancel_does_not_keep_its_deployment_busy(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
@ -1047,3 +1068,161 @@ async def test_picking_a_deployment_without_calling_it_does_not_keep_it_busy() -
picked: Final = await router.select(ROUTER_GROUP, {"metadata": {}})
assert _router_pick(router) == picked["model_info"]["id"]
def _started_in_its_own_attempt(worker: LeastBusyLoggingHandler, kwargs: dict[str, object]) -> contextvars.Context:
def start() -> contextvars.Context:
with routed_attempt():
worker.log_pre_api_call(model="m", messages=[], kwargs=kwargs)
return contextvars.copy_context()
return contextvars.copy_context().run(start)
def test_a_late_duplicate_failure_of_one_attempt_leaves_the_next_attempt_counted() -> None:
worker: Final = _worker(None)
attempt: Final = _call_kwargs("dep-a", "retried-call")
first: Final = _started_in_its_own_attempt(worker, attempt)
first.run(worker.log_failure_event, attempt, None, None, None)
_ = _started_in_its_own_attempt(worker, attempt)
first.run(worker.log_failure_event, attempt, None, None, None)
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_B
class _HeldSharedDecrement(SharedRedisCounters):
def __init__(self) -> None:
super().__init__()
self.decrementing: Final = asyncio.Event()
self.release: Final = asyncio.Event()
async def async_increment_with_floor(self, key: str, value: int, ttl: int) -> int:
if value < 0:
self.decrementing.set()
await self.release.wait()
return self.increment_with_floor(key, value, ttl)
@pytest.mark.asyncio
async def test_a_failure_callback_cancelled_mid_decrement_still_frees_the_shared_count() -> None:
shared: Final = _HeldSharedDecrement()
worker: Final = _worker(shared)
attempt: Final = _call_kwargs("dep-a")
worker.log_pre_api_call(model="m", messages=[], kwargs=attempt)
settling: Final = asyncio.create_task(worker.async_log_failure_event(attempt, None, None, None))
await shared.decrementing.wait()
settling.cancel()
settling.cancel()
with pytest.raises(asyncio.CancelledError):
await settling
shared.release.set()
await asyncio.sleep(0)
assert shared.count(f"{GROUP}_request_count:dep-a") == 0
_CHAT_REPLY: Final = {
"id": "chatcmpl-least-busy",
"object": "chat.completion",
"created": 1,
"model": "gpt-4.1-mini",
"choices": [{"index": 0, "message": {"role": "assistant", "content": "done"}, "finish_reason": "stop"}],
"usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2},
}
@pytest.mark.parametrize(
"attempts",
[("hang",), ("503", "hang"), ("503", "503")],
ids=["cancelled_while_waiting_on_the_provider", "cancelled_on_its_retry", "every_attempt_failed"],
)
@pytest.mark.asyncio
async def test_a_request_sharing_one_logging_object_that_ends_without_an_answer_leaves_its_deployment_idle(
attempts: tuple[str, ...], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
waiting_on_provider: Final = asyncio.Event()
script: Final = iter((*attempts, "ok"))
async def upstream(_: httpx.Request) -> httpx.Response:
match next(script):
case "503":
return httpx.Response(503, json={"error": {"message": "overloaded"}})
case "hang":
waiting_on_provider.set()
await asyncio.Event().wait()
return httpx.Response(200, json=_CHAT_REPLY)
respx_mock.post("https://api.openai.com/v1/chat/completions").mock(side_effect=upstream)
router: Final = Router(
model_list=[
{
"model_name": ROUTER_GROUP,
"litellm_params": {"model": "openai/gpt-4.1-mini", "api_key": "key-1"},
"model_info": {"id": "1"},
}
],
routing_strategy="least-busy",
num_retries=1,
retry_after=0,
)
logging_obj, request = litellm.utils.function_setup(
original_function="acompletion",
rules_obj=litellm.utils.Rules(),
start_time=datetime.datetime(2026, 1, 1),
model=ROUTER_GROUP,
messages=[{"role": "user", "content": "hi"}],
litellm_call_id="unanswered-call",
)
call: Final["asyncio.Task[object]"] = asyncio.create_task(
router.acompletion(**request, litellm_logging_obj=logging_obj)
)
if "hang" in attempts:
await waiting_on_provider.wait()
call.cancel()
with pytest.raises((asyncio.CancelledError, litellm.ServiceUnavailableError)):
await call
await _drain_logging()
assert router.cache.get_cache(key=f"{ROUTER_GROUP}_request_count:1") == 0
assert isinstance(
await router.acompletion(model=ROUTER_GROUP, messages=[{"role": "user", "content": "hi"}]), ModelResponse
)
@pytest.mark.asyncio
async def test_an_attempt_that_retried_inside_itself_frees_every_start_once_whatever_order_they_end_in() -> None:
shared: Final = SharedRedisCounters()
store: Final = InMemoryCache()
cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=shared) # pyright: ignore[reportArgumentType] # duck-typed Redis double
worker: Final = LeastBusyLoggingHandler(router_cache=cache, reservation_store=store)
attempt: Final = _call_kwargs("dep-a")
async def retried_then_cancelled() -> None:
with routed_attempt() as routed:
worker.log_pre_api_call(model="m", messages=[], kwargs=attempt)
worker.log_pre_api_call(model="m", messages=[], kwargs=attempt)
await worker.async_log_failure_event(attempt, None, None, None)
release: Final = release_started_attempt(store, cache, routed.attempt_id)
if release is not None:
await release
await retried_then_cancelled()
assert shared.count(f"{GROUP}_request_count:dep-a") == 0
assert worker.get_available_deployments(GROUP, HEALTHY, reserve=True) is DEPLOYMENT_A
def test_every_call_in_flight_at_once_frees_its_count_with_the_default_store_size() -> None:
worker: Final = _worker(None)
in_flight: Final = 7_000
calls: Final = tuple(_call_kwargs("dep-a", f"call-{n}") for n in range(in_flight))
attempts: Final = tuple(_started_in_its_own_attempt(worker, call) for call in calls)
for attempt, call in zip(attempts, calls, strict=True):
attempt.run(worker.log_failure_event, call, None, None, None)
assert worker.get_available_deployments(GROUP, HEALTHY) is DEPLOYMENT_A

341
tests/unit/test_router/test_io_token_rate_limits.py
View file

@ -3,12 +3,15 @@ Tests for separate ITPM/OTPM deployment rate limits (enforce_model_rate_limits).
"""
import asyncio
from collections.abc import Awaitable, Callable, Mapping
import contextvars
from collections.abc import AsyncIterator, Awaitable, Callable, Mapping
from datetime import datetime, timezone
from typing import Final
from unittest.mock import AsyncMock, MagicMock
import httpx
import pytest
import respx
from fakeredis import FakeRedis, FakeServer
from fakeredis.aioredis import FakeRedis as FakeAsyncRedis
@ -18,6 +21,8 @@ from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache, RedisCircuitBreakerOpenError
from litellm.constants import NO_MAX_TOKENS_TPM_RESERVATION_FRACTION
from litellm.litellm_core_utils.attempt_settlement import routed_attempt
from litellm.litellm_core_utils.internal_call_metadata import IO_TOKEN_RESERVATION_ID_KEY
from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject
from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
@ -950,11 +955,8 @@ class TestModelRateLimitingCheckIOTokens:
usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3),
)
with pytest.raises(RuntimeError):
await async_io_token_reconcile_success(dual_cache, kwargs, response)
await async_io_token_reconcile_success(dual_cache, kwargs, response)
# The stash is cleared even though reconciliation raised, so a duplicate
# success event can't re-process it.
assert ITPM_RESERVED_KEY not in metadata
assert ITPM_CACHE_KEY not in metadata
@ -1853,3 +1855,332 @@ async def test_a_later_turn_is_admitted_on_this_instance_s_count_when_the_shared
admitted: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT)
assert admitted[ITPM_CACHE_KEY if kind == "itpm" else OTPM_CACHE_KEY] == _turn_key(kind)
_HELD_DEPLOYMENT: Final = {
"litellm_params": {"model": ROUTER_IO_MODEL, "itpm": 50000, "otpm": 50000},
"model_info": {"id": "held-id"},
"model_name": "limited",
}
def _held_keys() -> tuple[str, str]:
minute: Final = _current_minute()
return (
f"global_router:held-id:{ROUTER_IO_MODEL}:itpm:{minute}",
f"global_router:held-id:{ROUTER_IO_MODEL}:otpm:{minute}",
)
class _HeldIncrement(DualCache):
"""Holds the first increment matching ``holds`` until ``release`` is set."""
def __init__(self, holds: Callable[[str, int], bool]) -> None:
super().__init__()
self._holds: Final = holds
self.holding: Final = asyncio.Event()
self.release: Final = asyncio.Event()
async def async_increment_cache(self, key: str, value: float, **kwargs: object) -> float | None:
if not self.holding.is_set() and self._holds(key, int(value)):
self.holding.set()
await self.release.wait()
return await super().async_increment_cache(key=key, value=value, **kwargs)
async def _counts(cache: DualCache) -> tuple[int, int]:
for _ in range(10):
await asyncio.sleep(0)
itpm_key, otpm_key = _held_keys()
itpm: Final[object] = await cache.async_get_cache(key=itpm_key)
otpm: Final[object] = await cache.async_get_cache(key=otpm_key)
return (itpm if isinstance(itpm, int) else 0), (otpm if isinstance(otpm, int) else 0)
async def _cancel_twice_while_held(cache: _HeldIncrement, held: "asyncio.Task[object]") -> None:
await cache.holding.wait()
held.cancel()
held.cancel()
with pytest.raises(asyncio.CancelledError):
await held
cache.release.set()
@pytest.mark.asyncio
async def test_a_request_cancelled_mid_reservation_hands_back_what_it_reserved() -> None:
cache: Final = _HeldIncrement(lambda key, value: ":otpm:" in key and value > 0)
request: Final[dict[str, object]] = {"messages": PROXY_MESSAGES, "max_tokens": 100, "metadata": {}}
async def reserve() -> object:
set_io_token_rate_limit_request_kwargs(request, metadata_channel="metadata")
return await async_io_token_pre_call_check(cache, dict(_HELD_DEPLOYMENT))
await _cancel_twice_while_held(cache, asyncio.create_task(reserve()))
assert await _counts(cache) == (0, 0)
assert ITPM_RESERVED_KEY not in request["metadata"]
def _stashed(itpm_reserved: int, otpm_reserved: int) -> dict[str, object]:
itpm_key, otpm_key = _held_keys()
return {
"metadata": {
ITPM_RESERVED_KEY: itpm_reserved,
OTPM_RESERVED_KEY: otpm_reserved,
ITPM_CACHE_KEY: itpm_key,
OTPM_CACHE_KEY: otpm_key,
}
}
async def _reserved(cache: DualCache, itpm_reserved: int, otpm_reserved: int) -> dict[str, object]:
itpm_key, otpm_key = _held_keys()
await cache.async_increment_cache(key=itpm_key, value=itpm_reserved)
await cache.async_increment_cache(key=otpm_key, value=otpm_reserved)
return _stashed(itpm_reserved, otpm_reserved)
@pytest.mark.asyncio
async def test_a_refund_cancelled_between_its_decrements_still_hands_back_both() -> None:
cache: Final = _HeldIncrement(lambda key, value: ":itpm:" in key and value < 0)
kwargs: Final = await _reserved(cache, 9, 100)
await _cancel_twice_while_held(
cache, asyncio.create_task(ModelRateLimitingCheck(cache).async_log_failure_event(kwargs, None, None, None))
)
assert await _counts(cache) == (0, 0)
class _YieldingIncrement(DualCache):
async def async_increment_cache(self, key: str, value: float, **kwargs: object) -> float | None:
await asyncio.sleep(0)
return await super().async_increment_cache(key=key, value=value, **kwargs)
@pytest.mark.asyncio
async def test_competing_failure_callbacks_hand_a_reservation_back_once() -> None:
cache: Final = _YieldingIncrement()
kwargs: Final = await _reserved(cache, 9, 100)
check: Final = ModelRateLimitingCheck(cache)
await asyncio.gather(*(check.async_log_failure_event(kwargs, None, None, None) for _ in range(2)))
assert await _counts(cache) == (0, 0)
@pytest.mark.asyncio
async def test_a_late_failure_callback_of_one_attempt_leaves_the_next_attempts_reservation_held() -> None:
cache: Final = DualCache()
check: Final = ModelRateLimitingCheck(cache)
request: Final[dict[str, object]] = {"messages": PROXY_MESSAGES, "max_tokens": 100, "metadata": {}}
async def attempt() -> contextvars.Context:
with routed_attempt():
set_io_token_rate_limit_request_kwargs(request, metadata_channel="metadata")
await async_io_token_pre_call_check(cache, dict(_HELD_DEPLOYMENT))
return contextvars.copy_context()
def failure_logged_in(context: contextvars.Context) -> "asyncio.Task[None]":
return context.run(asyncio.create_task, check.async_log_failure_event(request, None, None, None))
first: Final = await attempt()
await failure_logged_in(first)
_ = await attempt()
held_by_second: Final = await _counts(cache)
await failure_logged_in(first)
assert held_by_second[1] == 100
assert await _counts(cache) == held_by_second
def _anthropic_reply() -> dict[str, object]:
return {
"id": "msg_held",
"type": "message",
"role": "assistant",
"model": "claude-haiku-5-5",
"content": [{"type": "text", "text": "done"}],
"stop_reason": "end_turn",
"usage": {"input_tokens": 3, "output_tokens": 1},
}
@pytest.mark.parametrize(
"attempts",
[("hang",), ("503", "hang")],
ids=["cancelled_while_waiting_on_the_provider", "cancelled_on_its_retry"],
)
@pytest.mark.asyncio
async def test_a_cancelled_request_sharing_one_logging_object_hands_back_every_reservation(
attempts: tuple[str, ...], respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
waiting_on_provider: Final = asyncio.Event()
script: Final = iter((*attempts, "ok"))
async def upstream(_: httpx.Request) -> httpx.Response:
match next(script):
case "503":
return httpx.Response(
503, json={"type": "error", "error": {"type": "overloaded_error", "message": "x"}}
)
case "hang":
waiting_on_provider.set()
await asyncio.Event().wait()
return httpx.Response(200, json=_anthropic_reply())
respx_mock.post("https://api.anthropic.com/v1/messages").mock(side_effect=upstream)
router: Final = _router_with_io_limits()
first_minute: Final = _current_minute()
call: Final["asyncio.Task[object]"] = asyncio.create_task(
router.acompletion(
model="limited",
messages=PROXY_MESSAGES,
max_tokens=100,
metadata={"user_api_key_hash": "hashed-key"},
litellm_call_id="proxy-call",
litellm_logging_obj=_logging_obj_shared_across_attempts("acompletion"),
)
)
await waiting_on_provider.wait()
call.cancel()
with pytest.raises(asyncio.CancelledError):
await call
assert await _io_counters_since(router, first_minute) == (0, 0)
assert await _proxy_chat(router, "hi") == await _io_counters_since(router, first_minute)
@pytest.mark.asyncio
async def test_a_reservation_refunded_through_a_logged_copy_of_its_metadata_is_refunded_once() -> None:
cache: Final = DualCache()
check: Final = ModelRateLimitingCheck(cache)
request: Final[dict[str, object]] = {"messages": PROXY_MESSAGES, "max_tokens": 100, "litellm_metadata": {}}
async def cancelled_attempt() -> None:
with routed_attempt():
set_io_token_rate_limit_request_kwargs(request, metadata_channel="litellm_metadata")
await async_io_token_pre_call_check(cache, dict(_HELD_DEPLOYMENT))
logged.update({"litellm_params": {"litellm_metadata": dict(request["litellm_metadata"])}}) # pyright: ignore[reportArgumentType] # the stashed metadata dict
await check.async_log_failure_event(request, None, None, None)
logged: Final[dict[str, object]] = {}
await asyncio.create_task(cancelled_attempt())
await check.async_log_failure_event(logged, None, None, None)
assert await _counts(cache) == (0, 0)
@pytest.mark.asyncio
async def test_a_retry_refunds_the_reservation_its_previous_attempt_left_behind() -> None:
cache: Final = DualCache()
request: Final[dict[str, object]] = {"messages": PROXY_MESSAGES, "max_tokens": 100, "metadata": {}}
async def first_attempt() -> None:
with routed_attempt():
set_io_token_rate_limit_request_kwargs(request, metadata_channel="metadata")
await async_io_token_pre_call_check(cache, dict(_HELD_DEPLOYMENT))
def retry() -> None:
with routed_attempt():
refund_stale_reservation_before_retry(cache, request)
await asyncio.create_task(first_attempt())
contextvars.copy_context().run(retry)
assert await _counts(cache) == (0, 0)
@pytest.mark.asyncio
async def test_a_request_cancelled_after_its_provider_answered_keeps_the_tokens_it_spent_counted(
monkeypatch: pytest.MonkeyPatch,
) -> None:
from litellm.integrations.custom_logger import CustomLogger
answered: Final = asyncio.Event()
class SlowSuccessHook(CustomLogger):
async def async_post_call_success_deployment_hook(self, request_data, response, call_type):
answered.set()
await asyncio.Event().wait()
monkeypatch.setattr(litellm, "callbacks", [*litellm.callbacks, SlowSuccessHook()])
router: Final = _router_with_io_limits()
first_minute: Final = _current_minute()
call: Final["asyncio.Task[tuple[int, int]]"] = asyncio.create_task(_proxy_chat(router, "hi"))
await answered.wait()
call.cancel()
with pytest.raises(asyncio.CancelledError):
await call
_, otpm = await _io_counters_since(router, first_minute)
assert otpm == 100
def _stashed_as(reservation_id: str, itpm_reserved: int, otpm_reserved: int) -> dict[str, object]:
stashed: Final = _stashed(itpm_reserved, otpm_reserved)
return {"metadata": {**stashed["metadata"], IO_TOKEN_RESERVATION_ID_KEY: reservation_id}} # pyright: ignore[reportGeneralTypeIssues] # the stash is a dict
@pytest.mark.asyncio
async def test_a_settled_reservation_stays_settled_however_many_others_settle_after_it() -> None:
cache: Final = DualCache()
check: Final = ModelRateLimitingCheck(cache)
itpm_key, otpm_key = _held_keys()
await cache.async_increment_cache(key=itpm_key, value=9)
await cache.async_increment_cache(key=otpm_key, value=100)
first: Final = _stashed_as("first-reservation", 9, 100)
logged_copy: Final = {"litellm_params": {"metadata": dict(first["metadata"])}} # pyright: ignore[reportArgumentType] # the stash is a dict
await check.async_log_failure_event(first, None, None, None)
for n in range(10_001):
io_token_refund_failure(cache, _stashed_as(f"later-reservation-{n}", 0, 0))
await check.async_log_failure_event(logged_copy, None, None, None)
assert await _counts(cache) == (0, 0)
class _FirstChunkNeverArrives(httpx.AsyncByteStream):
def __init__(self, reading: asyncio.Event) -> None:
self._reading: Final = reading
async def __aiter__(self) -> AsyncIterator[bytes]:
self._reading.set()
await asyncio.Event().wait()
yield b""
@pytest.mark.asyncio
async def test_a_stream_cancelled_before_its_first_chunk_hands_back_its_reservation(
respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch
) -> None:
monkeypatch.setattr(litellm, "disable_aiohttp_transport", True)
reading_first_chunk: Final = asyncio.Event()
respx_mock.post("https://api.anthropic.com/v1/messages").mock(
return_value=httpx.Response(
200, headers={"content-type": "text/event-stream"}, stream=_FirstChunkNeverArrives(reading_first_chunk)
)
)
router: Final = _router_with_io_limits()
first_minute: Final = _current_minute()
call: Final["asyncio.Task[object]"] = asyncio.create_task(
router.acompletion(
model="limited",
messages=PROXY_MESSAGES,
max_tokens=100,
stream=True,
metadata={"user_api_key_hash": "hashed-key"},
)
)
await reading_first_chunk.wait()
call.cancel()
with pytest.raises(asyncio.CancelledError):
await call
assert await _io_counters_since(router, first_minute) == (0, 0)

63
tests/unit/test_utils.py
View file

@ -6511,6 +6511,34 @@ async def test_wrapper_async_raises_original_exception_even_if_hook_callback_err
)
@pytest.mark.asyncio
async def test_wrapper_async_lets_a_cancellation_during_the_failure_hook_cancel_the_call(
monkeypatch: pytest.MonkeyPatch,
) -> None:
hook_entered: Final = asyncio.Event()
class HangingLogger(CustomLogger):
async def async_post_call_failure_deployment_hook(
self, request_data, exception, call_type, fallback_depth=None
):
hook_entered.set()
await asyncio.Event().wait()
monkeypatch.setattr(litellm, "callbacks", [HangingLogger()])
call: Final = asyncio.create_task(
litellm.acompletion(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hi"}],
mock_response=litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o-mini"),
)
)
await hook_entered.wait()
call.cancel()
with pytest.raises(asyncio.CancelledError):
await call
@pytest.mark.asyncio
async def test_router_fallback_chain_reports_increasing_fallback_depth(monkeypatch: pytest.MonkeyPatch) -> None:
"""Regression: a real Router fallback chain must report fallback_depth=None on the
@ -6898,41 +6926,6 @@ async def test_router_fallback_not_skipped_when_failure_hook_callback_touches_at
assert response.choices[0].message.content == "fallback worked"
@pytest.mark.asyncio
async def test_wrapper_async_preserves_original_exception_when_hook_await_is_cancelled(
monkeypatch: pytest.MonkeyPatch,
) -> None:
"""Regression: if the caller's own timeout (e.g. asyncio.wait_for) fires while the
failure hook is still being awaited, the real deployment exception must still reach
the caller - not get replaced by CancelledError/TimeoutError from the hook's own
await getting cancelled."""
hook_entered: Final = asyncio.Event()
hook_released: Final = asyncio.Event()
class SlowLogger(CustomLogger):
async def async_post_call_failure_deployment_hook(
self, request_data, exception, call_type, fallback_depth=None
):
hook_entered.set()
await hook_released.wait()
monkeypatch.setattr(litellm, "callbacks", [SlowLogger()])
call: Final = asyncio.create_task(
litellm.acompletion(
model="gpt-4o-mini",
messages=[{"role": "user", "content": "hi"}],
mock_response=litellm.AuthenticationError(message="bad key", llm_provider="openai", model="gpt-4o-mini"),
)
)
assert await asyncio.wait_for(hook_entered.wait(), timeout=5)
call.cancel()
with pytest.raises(litellm.AuthenticationError):
await call
@pytest.mark.asyncio
async def test_wrapper_async_failure_hook_latency_does_not_inflate_reported_duration(
monkeypatch: pytest.MonkeyPatch,