diff --git a/litellm/litellm_core_utils/attempt_settlement.py b/litellm/litellm_core_utils/attempt_settlement.py new file mode 100644 index 00000000000..02a6497cff8 --- /dev/null +++ b/litellm/litellm_core_utils/attempt_settlement.py @@ -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, + ) diff --git a/litellm/litellm_core_utils/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py index 138a01c6524..de17bcbdfb9 100644 --- a/litellm/litellm_core_utils/internal_call_metadata.py +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -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. diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index eef454862f9..b7c9a96c9ff 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -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 { diff --git a/litellm/router_strategy/least_busy.py b/litellm/router_strategy/least_busy.py index dfcb4cccf54..9d3767268ca 100644 --- a/litellm/router_strategy/least_busy.py +++ b/litellm/router_strategy/least_busy.py @@ -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) diff --git a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py index 846cc9f3d61..4b618aba476 100644 --- a/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/io_token_rate_limit_check.py @@ -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( diff --git a/litellm/utils.py b/litellm/utils.py index b9bf872e853..ce5e4129ede 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -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, diff --git a/tests/unit/router_strategy/test_least_busy.py b/tests/unit/router_strategy/test_least_busy.py index f156a5ee405..98dadb8c40e 100644 --- a/tests/unit/router_strategy/test_least_busy.py +++ b/tests/unit/router_strategy/test_least_busy.py @@ -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 diff --git a/tests/unit/test_router/test_io_token_rate_limits.py b/tests/unit/test_router/test_io_token_rate_limits.py index f09ad5973e2..bb4813a273b 100644 --- a/tests/unit/test_router/test_io_token_rate_limits.py +++ b/tests/unit/test_router/test_io_token_rate_limits.py @@ -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) diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 44bdd187da4..6374c8bef59 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -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,