diff --git a/litellm/_internal_context.py b/litellm/_internal_context.py index e97d79a95bc..39fe789160b 100644 --- a/litellm/_internal_context.py +++ b/litellm/_internal_context.py @@ -197,6 +197,20 @@ def current_request_parent_otel_span() -> Span | None: return _request_parent_otel_span.get() +@contextmanager +def request_parent_otel_span(span: Span | None) -> Generator[None]: + """Anchor the datastore service spans made inside this block, and by the tasks it starts, under the span an SDK + caller passed. A request that already anchored its own, as the proxy does at auth, keeps it.""" + if span is None or _request_parent_otel_span.get() is not None: + yield + return + token: Final = _request_parent_otel_span.set(span) + try: + yield + finally: + _restore(_request_parent_otel_span, token) + + @contextmanager def outside_mirrored_call() -> Generator[None]: """Logging a mirrored call: the calls its callbacks start are not mirrors themselves.""" diff --git a/litellm/caching/redis_batch.py b/litellm/caching/redis_batch.py index d621d09b439..b3e6d2eb291 100644 --- a/litellm/caching/redis_batch.py +++ b/litellm/caching/redis_batch.py @@ -16,7 +16,7 @@ import logging import time import weakref from collections.abc import Awaitable, Callable, Generator, Mapping, Sequence -from contextvars import ContextVar, Token, copy_context +from contextvars import ContextVar, Token from dataclasses import dataclass, field from datetime import timedelta from enum import Enum, auto @@ -528,7 +528,6 @@ class RequestRedisBatches: "_deadline", "_deadline_flush", "_native", - "_native_context", "_native_factory", "_native_flush_hooks", "_native_misses", @@ -552,7 +551,6 @@ class RequestRedisBatches: self._native_factory: Final = native_factory self._native_send_factory: Final = native_send_factory self._native: NativeRequestRedis | None = None - self._native_context: Final = copy_context() self._native_flush_hooks: Final[list[Callable[[], None]]] = [] # mutable-ok: append-only registry self._native_misses: frozenset[str] = frozenset() self._deadline: asyncio.TimerHandle | None = None @@ -582,12 +580,9 @@ class RequestRedisBatches: return batch def native(self) -> NativeRequestRedis: - """The request's native batches, built on first use so requests that touch no native store pay nothing. - - They keep a copy of the context they are built in, and the cycle collector cannot see through native - objects, so they are built in the context this scope was opened from, which cannot reach the scope.""" + """The request's native batches, built on first use so requests that touch no native store pay nothing.""" if self._native is None: - self._native = self._native_context.run(self._native_factory, self.post_call_deadline) + self._native = self._native_factory(self.post_call_deadline) return self._native def native_post_call(self) -> NativeRequestRedis: diff --git a/litellm/litellm_core_utils/internal_call_metadata.py b/litellm/litellm_core_utils/internal_call_metadata.py index de17bcbdfb9..98c009552d1 100644 --- a/litellm/litellm_core_utils/internal_call_metadata.py +++ b/litellm/litellm_core_utils/internal_call_metadata.py @@ -44,16 +44,16 @@ BACKGROUND_RESPONSE_COST_REQUEST_ID_SUFFIX: Final = "_background_cost" 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" +ITPM_COUNTER_KEY: Final = "_litellm_itpm_counter" +OTPM_COUNTER_KEY: Final = "_litellm_otpm_counter" 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, + ITPM_COUNTER_KEY, + OTPM_COUNTER_KEY, IO_TOKEN_ATTEMPT_KEY, IO_TOKEN_RESERVATION_ID_KEY, } diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 0ca09140f40..1f2e6511be7 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -14,6 +14,7 @@ import sys import time import traceback from collections.abc import Awaitable, Callable, Coroutine, Iterable, Iterator, Mapping, Sequence +from contextlib import AbstractContextManager from dataclasses import dataclass from datetime import datetime as dt_object from enum import Enum @@ -28,7 +29,12 @@ from pydantic import BaseModel, JsonValue, TypeAdapter from typing_extensions import TypedDict import litellm -from litellm._internal_context import in_mirrored_call, outside_mirrored_call, post_response_phase +from litellm._internal_context import ( + in_mirrored_call, + outside_mirrored_call, + post_response_phase, + request_parent_otel_span, +) from litellm._lazy_imports import get_token_counter_new from litellm._logging import ( is_debugging_on, @@ -97,6 +103,7 @@ from litellm.litellm_core_utils.core_helpers import ( attribute_or_none, decoded_json_object, finite_cost, + get_parent_otel_span_from_kwargs, get_provider_response_headers_from_hidden_params, has_attributes, is_expected_client_error, @@ -3738,6 +3745,11 @@ class Logging(LiteLLMLoggingBaseClass): if trace_id_var.get() == self._own_trace_id and session_id_var.get() == self._own_session_id: self._restore_correlation_context() + def _parent_otel_span_anchored(self) -> AbstractContextManager[None]: + """Every callback runs in one of the four handlers, streamed or not, so anchoring here gives each + callback's datastore spans the caller's span without the callback knowing.""" + return request_parent_otel_span(get_parent_otel_span_from_kwargs(self.model_call_details)) + def success_handler( self, result: object = None, # heterogeneous response object; varies by call type (ANN401 ignored, see ruff-strict.toml) @@ -3751,7 +3763,7 @@ class Logging(LiteLLMLoggingBaseClass): """Restores trace_id/session_id contextvars once this attempt's own success logging (including any nested calls its callbacks trigger) is fully done.""" try: - with post_response_phase(), outside_mirrored_call(): + with post_response_phase(), outside_mirrored_call(), self._parent_otel_span_anchored(): return self._success_handler_body( result=result, start_time=start_time, @@ -4267,7 +4279,7 @@ class Logging(LiteLLMLoggingBaseClass): """Restores trace_id/session_id contextvars once this attempt's own success logging (including any nested calls its callbacks trigger) is fully done.""" try: - with post_response_phase(), outside_mirrored_call(): + with post_response_phase(), outside_mirrored_call(), self._parent_otel_span_anchored(): return await self._async_success_handler_body( result=result, start_time=start_time, end_time=end_time, cache_hit=cache_hit, **kwargs ) @@ -4841,7 +4853,7 @@ class Logging(LiteLLMLoggingBaseClass): ``for_async_request`` marks the sync run of a failure whose async run also logs it, so a CustomLogger that gets ``async_log_failure_event`` from that async run skips its sync hook.""" try: - with outside_mirrored_call(): + with outside_mirrored_call(), self._parent_otel_span_anchored(): return self._failure_handler_body( exception=exception, traceback_exception=traceback_exception, @@ -5042,7 +5054,7 @@ class Logging(LiteLLMLoggingBaseClass): """Restores trace_id/session_id contextvars once this attempt's own failure logging (including any nested calls its callbacks trigger) is fully done.""" try: - with outside_mirrored_call(): + with outside_mirrored_call(), self._parent_otel_span_anchored(): return await self._async_failure_handler_body( exception=exception, traceback_exception=traceback_exception, diff --git a/litellm/redis/stores/router_usage/__init__.py b/litellm/redis/stores/router_usage/__init__.py index e9d46f8a68b..3fbae3d42b6 100644 --- a/litellm/redis/stores/router_usage/__init__.py +++ b/litellm/redis/stores/router_usage/__init__.py @@ -1,32 +1,52 @@ from __future__ import annotations -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Callable, Mapping, Sequence +from dataclasses import dataclass from typing import TYPE_CHECKING, Final, Protocol from litellm._internal_context import with_service_target +from litellm.caching.redis_batch import Batch +from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.redis.stores._route import post_call_request, ready, route from litellm.redis.stores.router_usage.handle import RouterUsageHandle from litellm.redis.stores.router_usage.types import ( + MAX_INCREMENT, ROUTER_USAGE_TARGET, + CounterIncrement, CounterKind, + CounterScope, UsageCounter, ) +from litellm.rust_bridge._native import NativeRedisStoreError if TYPE_CHECKING: from litellm.redis.stores.connection import RedisStores __all__ = [ + "MAX_INCREMENT", "ROUTER_USAGE_TARGET", + "CounterIncrement", "CounterKind", + "CounterScope", "RouterUsage", "RouterUsageStore", "UsageCounter", + "UsageRead", ] +@dataclass(frozen=True, slots=True) +class UsageRead: + """A declared usage read, and how to give back the throttle refreshes it reserved.""" + + result: Awaitable[Sequence[float | None]] + release: Callable[[], None] + + class RouterUsage(Protocol): def attach_redis(self, redis: RedisStores) -> None: ... def flush_local_counters(self) -> None: ... - async def increment(self, counter: UsageCounter, value: float, ttl: int) -> float: ... + async def increment(self, counter: UsageCounter, value: float, ttl: int) -> float | None: ... async def increment_deployment( self, deployment_id: str, @@ -36,7 +56,7 @@ class RouterUsage(Protocol): requests: int, ) -> None: ... async def get_many(self, counters: Sequence[UsageCounter]) -> Sequence[float | None]: ... - def increment_sync(self, counter: UsageCounter, value: float, ttl: int) -> float: ... + def increment_sync(self, counter: UsageCounter, value: float, ttl: int) -> float | None: ... def increment_deployment_sync( self, deployment_id: str, @@ -46,6 +66,19 @@ class RouterUsage(Protocol): requests: int, ) -> None: ... def get_many_sync(self, counters: Sequence[UsageCounter]) -> Sequence[float | None]: ... + @property + def has_redis(self) -> bool: ... + def usage(self, counters: Sequence[UsageCounter], *, batch: Batch) -> UsageRead: ... + async def increment_pipeline(self, increments: Sequence[CounterIncrement]) -> Sequence[float] | None: ... + async def increment_post_call(self, increments: Sequence[CounterIncrement]) -> None: ... + async def increment_remote(self, increments: Sequence[CounterIncrement]) -> Sequence[float] | None: ... + def increment_local(self, increment: CounterIncrement) -> float: ... + def get_local(self, counter: UsageCounter) -> float | None: ... + def get_local_many(self, counters: Sequence[UsageCounter]) -> Sequence[float | None]: ... + def set_local(self, counter: UsageCounter, value: float) -> None: ... + async def current(self, counter: UsageCounter, limit: float) -> float | None: ... + def current_sync(self, counter: UsageCounter, limit: float) -> float | None: ... + async def shared_or_local(self, counter: UsageCounter) -> float | None: ... def snapshot(self) -> Mapping[str, float]: ... @@ -66,7 +99,7 @@ class RouterUsageStore: self._handle.flush_local_counters() @with_service_target(ROUTER_USAGE_TARGET) - async def increment(self, counter: UsageCounter, value: float, ttl: int) -> float: + async def increment(self, counter: UsageCounter, value: float, ttl: int) -> float | None: """Without Redis every async method answers from memory without suspending, as ``DualCache``'s did.""" if not self.has_redis: return self.increment_sync(counter, value, ttl) @@ -93,7 +126,7 @@ class RouterUsageStore: return await self._handle.get_many(counters) @with_service_target(ROUTER_USAGE_TARGET) - def increment_sync(self, counter: UsageCounter, value: float, ttl: int) -> float: + def increment_sync(self, counter: UsageCounter, value: float, ttl: int) -> float | None: """``increment`` for a sync call, which has no event loop to await on.""" return self._handle.increment_sync(counter, value, ttl) @@ -112,5 +145,79 @@ class RouterUsageStore: def get_many_sync(self, counters: Sequence[UsageCounter]) -> Sequence[float | None]: return self._handle.get_many_sync(counters) + @with_service_target(ROUTER_USAGE_TARGET) + def usage(self, counters: Sequence[UsageCounter], *, batch: Batch) -> UsageRead: + """``get_many`` declared now on the batch ``batch`` picks, so a pick's read can share the request's + pre-call trip. Without Redis it is answered from memory without suspending.""" + if not self.has_redis: + return UsageRead(ready(self.get_many_sync(counters)), lambda: None) + target: Final = route(batch) + declared: Final = self._handle.usage(counters, target.send) + return UsageRead(target.result(declared.read), declared.release) + + @with_service_target(ROUTER_USAGE_TARGET) + async def increment_pipeline(self, increments: Sequence[CounterIncrement]) -> Sequence[float] | None: + """``DualCache.async_increment_cache_pipeline``, on its own trip.""" + if not self.has_redis: + return [self._handle.increment_local(increment) for increment in increments] + return await self._handle.increment_pipeline(increments) + + @with_service_target(ROUTER_USAGE_TARGET) + async def increment_post_call(self, increments: Sequence[CounterIncrement]) -> None: + """``DualCache.async_increment_cache_pipeline_post_call``: memory now, Redis on the request's post-call + pipeline, or on a trip of its own outside a request.""" + request: Final = post_call_request() + if request is None or not self.has_redis: + await self.increment_pipeline(increments) + return + self._handle.increment_post_call(request, increments) + + @with_service_target(ROUTER_USAGE_TARGET) + async def increment_remote(self, increments: Sequence[CounterIncrement]) -> Sequence[float] | None: + """``redis_cache.async_increment_pipeline``: Redis alone, None without it, raising what it raises there, + so ``log_redis_failure`` keeps an open breaker at DEBUG and throttles timeouts.""" + try: + return await self._handle.increment_remote(increments) + except NativeRedisStoreError as e: + match e.args[:1]: + case ("breaker_open",): + raise RedisCircuitBreakerOpenError(*e.args[1:]) from e + case ("timeout",): + raise TimeoutError(*e.args[1:]) from e + case _: + raise + + def increment_local(self, increment: CounterIncrement) -> float: + return self._handle.increment_local(increment) + + def get_local(self, counter: UsageCounter) -> float | None: + return self._handle.get_local(counter) + + def get_local_many(self, counters: Sequence[UsageCounter]) -> Sequence[float | None]: + return self._handle.get_local_many(counters) + + def set_local(self, counter: UsageCounter, value: float) -> None: + self._handle.set_local(counter, value) + + @with_service_target(ROUTER_USAGE_TARGET) + async def current(self, counter: UsageCounter, limit: float) -> float | None: + """Memory alone once it reaches ``limit``, else Redis's count. No Redis and an open breaker serve memory, + and any other Redis failure reads as None, as ``RedisCache.async_get_cache`` swallows it.""" + if not self.has_redis: + return self._handle.get_local(counter) + return await self._handle.current(counter, limit) + + @with_service_target(ROUTER_USAGE_TARGET) + def current_sync(self, counter: UsageCounter, limit: float) -> float | None: + return self._handle.current_sync(counter, limit) + + @with_service_target(ROUTER_USAGE_TARGET) + async def shared_or_local(self, counter: UsageCounter) -> float | None: + """What every pod counted in Redis, else this router's memory count when Redis has none, is absent, refuses + or fails.""" + if not self.has_redis: + return self._handle.get_local(counter) + return await self._handle.shared_or_local(counter) + def snapshot(self) -> Mapping[str, float]: return self._handle.snapshot() diff --git a/litellm/redis/stores/router_usage/handle.py b/litellm/redis/stores/router_usage/handle.py index d3e138e4cd1..1b11f750079 100644 --- a/litellm/redis/stores/router_usage/handle.py +++ b/litellm/redis/stores/router_usage/handle.py @@ -3,15 +3,24 @@ from __future__ import annotations from collections.abc import Awaitable, Mapping, Sequence from typing import Protocol -from litellm.redis.stores.router_usage.types import UsageCounter -from litellm.rust_bridge._native import NativeRedisClient +from litellm.caching.redis_batch import NativePending +from litellm.redis.stores.router_usage.types import CounterIncrement, UsageCounter +from litellm.rust_bridge._native import NativeRedisClient, NativeRequestRedis, NativeSend + + +class DeclaredUsage(Protocol): + """A usage read declared on a batch, with the throttle refreshes it reserved.""" + + @property + def read(self) -> NativePending[Sequence[float | None]]: ... + def release(self) -> None: ... class RouterUsageHandle(Protocol): @property def has_redis(self) -> bool: ... def attach(self, client: NativeRedisClient, /) -> bool: ... - def increment(self, counter: UsageCounter, value: float, ttl_seconds: int) -> Awaitable[float]: ... + def increment(self, counter: UsageCounter, value: float, ttl_seconds: int) -> Awaitable[float | None]: ... def increment_deployment( self, deployment_id: str, @@ -21,7 +30,7 @@ class RouterUsageHandle(Protocol): requests: int, ) -> Awaitable[None]: ... def get_many(self, counters: Sequence[UsageCounter]) -> Awaitable[Sequence[float | None]]: ... - def increment_sync(self, counter: UsageCounter, value: float, ttl_seconds: int) -> float: ... + def increment_sync(self, counter: UsageCounter, value: float, ttl_seconds: int) -> float | None: ... def increment_deployment_sync( self, deployment_id: str, @@ -31,5 +40,16 @@ class RouterUsageHandle(Protocol): requests: int, ) -> None: ... def get_many_sync(self, counters: Sequence[UsageCounter]) -> Sequence[float | None]: ... + def usage(self, counters: Sequence[UsageCounter], send: NativeSend) -> DeclaredUsage: ... + def increment_pipeline(self, increments: Sequence[CounterIncrement]) -> Awaitable[Sequence[float] | None]: ... + def increment_post_call(self, request: NativeRequestRedis, increments: Sequence[CounterIncrement]) -> None: ... + def increment_remote(self, increments: Sequence[CounterIncrement]) -> Awaitable[Sequence[float] | None]: ... + def increment_local(self, increment: CounterIncrement) -> float: ... + def get_local(self, counter: UsageCounter) -> float | None: ... + def get_local_many(self, counters: Sequence[UsageCounter]) -> Sequence[float | None]: ... + def set_local(self, counter: UsageCounter, value: float) -> None: ... + def current(self, counter: UsageCounter, limit: float) -> Awaitable[float | None]: ... + def current_sync(self, counter: UsageCounter, limit: float) -> float | None: ... + def shared_or_local(self, counter: UsageCounter) -> Awaitable[float | None]: ... def flush_local_counters(self) -> None: ... def snapshot(self) -> Mapping[str, float]: ... diff --git a/litellm/redis/stores/router_usage/types.py b/litellm/redis/stores/router_usage/types.py index 062ee0a084e..a73c2ac38c0 100644 --- a/litellm/redis/stores/router_usage/types.py +++ b/litellm/redis/stores/router_usage/types.py @@ -4,12 +4,33 @@ from dataclasses import dataclass from typing import Final, Literal, TypeAlias ROUTER_USAGE_TARGET: Final = "router_usage" -CounterKind: TypeAlias = Literal["tpm", "rpm"] +CounterKind: TypeAlias = Literal["tpm", "rpm", "itpm", "otpm", "tpd", "rpd"] +CounterScope: TypeAlias = Literal["global_router", "deployment"] +# The native store counts in i64, so a larger amount never crosses into it. +MAX_INCREMENT: Final = 2**63 - 1 @dataclass(frozen=True, slots=True) class UsageCounter: + """One deployment's count in one window: ``HH-MM`` for a per-minute kind, ``YYYY-MM-DD`` for a per-day one. + ``global_router`` counters are the router core's, ``deployment`` ones those usage-based-routing-v2 and the + model rate limit check share.""" + + scope: CounterScope deployment_id: str model: str | None - minute: str + window: str kind: CounterKind + + @staticmethod + def for_deployment(model_id: str, model: str, kind: CounterKind, window: str) -> UsageCounter: + """A ``deployment`` counter. Callers spell the id and model with ``f"{value}"``, as main's + ``f"{model_id}:{model}:..."`` keys did, so a missing one is ``"None"``.""" + return UsageCounter("deployment", model_id, model, window, kind) + + +@dataclass(frozen=True, slots=True) +class CounterIncrement: + counter: UsageCounter + amount: float + ttl: float | None diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index a689b05a485..810bb4d3230 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -34,7 +34,7 @@ from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias from pydantic import BaseModel, ConfigDict, Field, TypeAdapter, ValidationError, field_validator, model_validator from typing_extensions import assert_never -from litellm._internal_context import service_target +from litellm._internal_context import request_parent_otel_span, service_target from litellm.constants import DEFAULT_COOLDOWN_TIME_SECONDS from litellm.integrations.custom_guardrail import is_guardrail_intervention from litellm.litellm_core_utils.add_retry_fallback_headers import ( @@ -1372,6 +1372,10 @@ class _EndpointPort: def routing_cache(self) -> "DualCache": return self._core._cache # pyright: ignore[reportPrivateUsage] # see the class docstring + @property + def router_usage(self) -> RouterUsage: + return self._core._usage # pyright: ignore[reportPrivateUsage] # see the class docstring + @property def selectors(self) -> StrategySelectors: return self._core._selectors # pyright: ignore[reportPrivateUsage] # see the class docstring @@ -1497,7 +1501,7 @@ class RustCore( # document router-core last received for them. self._index_written: frozenset[str] = frozenset() self._index_sent: Mapping[str, object] | None = None - self._selectors: Final = StrategySelectors(host, self._cache) + self._selectors: Final = StrategySelectors(host, self._cache, self._usage) self._selectors.init_default(resolved.routing_strategy, resolved.routing_strategy_args) self._selectors.init_groups(_group_selectors(resolved.routing_groups)) self._port: Final = _EndpointPort(self) @@ -1544,7 +1548,7 @@ class RustCore( self._hosted._core_routing_plugins = list(resolved.routing_plugins) # pyright: ignore[reportPrivateUsage] # the core fills its own host registry self._rebuild() self._join_callback_holders() - self._checks: Final = OptionalChecks(self, host, self._cache, self._hosted, ctx, config.arguments) + self._checks: Final = OptionalChecks(self, host, self._cache, self._usage, self._hosted, ctx, config.arguments) self._hosted.checks = self._checks strategies.ensure_affinity = self._checks.ensure_deployment_affinity_callback strategies.hand_over_affinity(self._checks.adopt) @@ -2959,7 +2963,10 @@ class RustCore( overrides if passthrough is None else {**(overrides or {}), "passthrough_on_no_deployment": True} ) # `owning_picks`: least-busy frees every pick this request made but never started. - with owning_picks(self._selectors.reservations, call_id_of(CALL_IDENTITY.validate_python(request))): + with ( + request_parent_otel_span(get_parent_otel_span_from_kwargs(request)), + owning_picks(self._selectors.reservations, call_id_of(CALL_IDENTITY.validate_python(request))), + ): call, prepared = self._start(plan, request, routed, prepared_already=prepared_already) arm_compaction_state(self._port, plan, prepared) # pyright: ignore[reportArgumentType] # the call's own dict with ( @@ -3037,7 +3044,10 @@ class RustCore( self, plan: Plan, request: Mapping[str, object], overrides: Mapping[str, object] | None = None ) -> object: # The sync routed calls pick with `picks_reserve(True)` inside `owning_picks`. - with owning_picks(self._selectors.reservations, call_id_of(CALL_IDENTITY.validate_python(request))): + with ( + request_parent_otel_span(get_parent_otel_span_from_kwargs(request)), + owning_picks(self._selectors.reservations, call_id_of(CALL_IDENTITY.validate_python(request))), + ): call, prepared = self._start(plan, request, overrides) with picks_reserve(plan.endpoint not in _UNRESERVED_OPS): driver, step = self._run_sync(plan, call, prepared, fresh=True) @@ -5212,9 +5222,8 @@ class _Driver: 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 + core._usage, # 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, diff --git a/litellm/router_rust/invoke/endpoints.py b/litellm/router_rust/invoke/endpoints.py index 5b4a65a70b1..496f8093706 100644 --- a/litellm/router_rust/invoke/endpoints.py +++ b/litellm/router_rust/invoke/endpoints.py @@ -57,6 +57,7 @@ from litellm.vector_stores.search_query import store_search_params if TYPE_CHECKING: from litellm.caching.dual_cache import DualCache + from litellm.redis.stores.router_usage import RouterUsage from litellm.router_api import RouterRegistry from litellm.router_rust.core import RustCore from litellm.router_rust.steps.selection import StrategySelectors @@ -94,6 +95,9 @@ class EndpointCore(Protocol): @property def routing_cache(self) -> "DualCache": ... + @property + def router_usage(self) -> "RouterUsage": ... + @property def selectors(self) -> "StrategySelectors": ... @@ -893,7 +897,6 @@ def attempt_call( data: Final = _deployment_params(plan, deployment) model_name: Final = str(data.get("model")) _validate_decisions_connection(plan, data, request) - cache: Final = core.routing_cache caching: Final = shaping.cache_responses # A fallback hop routes the target group in a copy of the metadata (`run_async_fallback`). if fallback_level > 0: @@ -910,7 +913,9 @@ def attempt_call( start_silent_experiments(core, plan, silent_model, named.get("messages"), kwargs) if "messages" in named: kwargs.setdefault("messages", named["messages"]) - update_kwargs_with_deployment(shaping, deployment, kwargs, cache=cache, splits_stream_timeout=True) + update_kwargs_with_deployment( + shaping, deployment, kwargs, usage=core.router_usage, splits_stream_timeout=True + ) kwargs.pop("silent_model", None) final = without_router_only_kwargs( # rebind-ok: rebound as the Python router rebinds it { @@ -924,7 +929,7 @@ def attempt_call( return AttemptCall(plan.host_op, final, kwargs, str(params.get("model")), group) case "text" | "embedding" | "image" | "transcription" | "rerank" | "adapter": update_kwargs_with_deployment( - shaping, deployment, kwargs, cache=cache, splits_stream_timeout=plan.splits_stream_timeout + shaping, deployment, kwargs, usage=core.router_usage, splits_stream_timeout=plan.splits_stream_timeout ) refused: Final = embedding_batch_refusal(plan, deployment, named.get("input"), kwargs, requested) if refused is not None: @@ -938,7 +943,7 @@ def attempt_call( } return AttemptCall(plan.host_op, without_router_only_kwargs(final), kwargs, model_name, group) case "speech": - update_kwargs_with_deployment(shaping, deployment, kwargs, cache=cache) + update_kwargs_with_deployment(shaping, deployment, kwargs, usage=core.router_usage) voice = kwargs.pop("voice", None) # rebind-ok: rebound as the Python router rebinds it final = { # rebind-ok: rebound as the Python router rebinds it **data, @@ -952,7 +957,9 @@ def attempt_call( if plan.kind == "cancel_batch": data.update(_resolved_credentials(core, deployment, requested)) data.pop("litellm_credential_name", None) - update_kwargs_with_deployment(shaping, deployment, kwargs, cache=cache, function_name=plan.function_name) + update_kwargs_with_deployment( + shaping, deployment, kwargs, usage=core.router_usage, function_name=plan.function_name + ) final = { # rebind-ok: rebound as the Python router rebinds it **data, "custom_llm_provider": _provider(core, data), @@ -968,7 +975,7 @@ def attempt_call( shaping, deployment, kwargs, - cache=cache, + usage=core.router_usage, function_name="generic_api_call", splits_stream_timeout=invoked_handler_name(core, plan) in STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, ) @@ -984,7 +991,7 @@ def attempt_call( } return AttemptCall(plan.host_op, final, kwargs, model_name, group) case "moderation": - update_kwargs_with_deployment(shaping, deployment, kwargs, cache=cache) + update_kwargs_with_deployment(shaping, deployment, kwargs, usage=core.router_usage) kwargs["model"] = model_name kwargs.update(data) return AttemptCall(plan.host_op, kwargs, kwargs, model_name, group) @@ -1015,7 +1022,7 @@ def _generic_attempt( core.shaping, deployment, kwargs, - cache=core.routing_cache, + usage=core.router_usage, function_name="_ageneric_api_call_with_fallbacks", splits_stream_timeout=handler in STREAM_TIMEOUT_AWARE_GENERIC_HANDLERS, opens_responses_socket=socket, @@ -2265,7 +2272,7 @@ class RustEndpoints: shaping, deployment, new_kwargs, # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values - cache=self._core.routing_cache, + usage=self._core.router_usage, function_name="aretrieve_batch", ) bucket: Final = metadata_bucket(new_kwargs, metadata_variable_name) # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values @@ -2350,7 +2357,7 @@ class RustEndpoints: self._core.shaping, deployment, kwargs_copy, # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values - cache=self._core.routing_cache, + usage=self._core.router_usage, function_name="acreate_file", ) data: Final = as_dict(deployment.get("litellm_params")) diff --git a/litellm/router_rust/invoke/prompt_management.py b/litellm/router_rust/invoke/prompt_management.py index 74a7991b06c..cc66fc16b6a 100644 --- a/litellm/router_rust/invoke/prompt_management.py +++ b/litellm/router_rust/invoke/prompt_management.py @@ -76,7 +76,7 @@ async def prompt_management_acompletion( core.shaping, deployment, kwargs, # pyright: ignore[reportUnknownArgumentType] # the request kwargs - cache=core.routing_cache, + usage=core.router_usage, splits_stream_timeout=True, ) params: Final = as_dict(deployment.get("litellm_params")) diff --git a/litellm/router_rust/invoke/shaping.py b/litellm/router_rust/invoke/shaping.py index bf6410d8242..5fa6fc2543c 100644 --- a/litellm/router_rust/invoke/shaping.py +++ b/litellm/router_rust/invoke/shaping.py @@ -52,7 +52,7 @@ from litellm.types.router import Deployment, GenericLiteLLMParams, LiteLLM_Param from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN if TYPE_CHECKING: - from litellm.caching.dual_cache import DualCache + from litellm.redis.stores.router_usage import RouterUsage from litellm.router_python_host import PythonHost # The SDK-client and max_parallel_requests reads of the routing cache: process-local, so the family @@ -404,7 +404,7 @@ def update_kwargs_with_deployment( deployment: Mapping[str, object], kwargs: MutableMapping[str, object], # mutable-ok: the request kwargs the Python router mutates in place *, - cache: "DualCache | None", + usage: "RouterUsage | None", function_name: str | None = None, splits_stream_timeout: bool = False, opens_responses_socket: bool = False, @@ -467,8 +467,8 @@ def update_kwargs_with_deployment( if isinstance(other_bucket_metadata, Mapping) and KUBERNETES_POD_ROUTING_KEY in other_bucket_metadata: kwargs[other_bucket] = {**as_dict(other_bucket_metadata), KUBERNETES_POD_ROUTING_KEY: routing} # pyright: ignore[reportUnknownArgumentType] # litellm's untyped helpers and request values # rebind-ok: the request kwargs are shaped in place, as the Python router shapes them - if cache is not None: - refund_stale_reservation_before_retry(cache, kwargs) # pyright: ignore[reportArgumentType] # the request's kwargs dict + if usage is not None: + refund_stale_reservation_before_retry(usage, kwargs) # pyright: ignore[reportArgumentType] # the request's kwargs dict set_io_token_rate_limit_request_kwargs( kwargs, # pyright: ignore[reportArgumentType] # the request's kwargs dict store_in_context=deployment_has_io_token_limits(deployment), # pyright: ignore[reportArgumentType] # a deployment record diff --git a/litellm/router_rust/invoke/websocket.py b/litellm/router_rust/invoke/websocket.py index b38252503c7..a787c729641 100644 --- a/litellm/router_rust/invoke/websocket.py +++ b/litellm/router_rust/invoke/websocket.py @@ -14,6 +14,7 @@ from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Final +from litellm._internal_context import request_parent_otel_span from litellm._logging import verbose_router_logger from litellm.constants import USAGE_BASED_ROUTING_COUNTED_TOKENS_METADATA_KEY from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs, without @@ -115,18 +116,18 @@ async def _count_usage_based_routing_tokens( or not isinstance(deployment_name, str) ): return False - await count_before_callbacks( - deployment_metadata, - USAGE_BASED_ROUTING_COUNTED_TOKENS_METADATA_KEY, - total_tokens, - partial( - selector.async_increment_tpm, - deployment_id=str(deployment_id), - deployment_name=deployment_name, - tokens=total_tokens, - parent_otel_span=parent_otel_span, - ), - ) + with request_parent_otel_span(parent_otel_span): + await count_before_callbacks( + deployment_metadata, + USAGE_BASED_ROUTING_COUNTED_TOKENS_METADATA_KEY, + total_tokens, + partial( + selector.async_increment_tpm, + deployment_id=str(deployment_id), + deployment_name=deployment_name, + tokens=total_tokens, + ), + ) return True @@ -167,7 +168,7 @@ async def admit_deployment_turn( io_token_admission: Final[Mapping[str, object]] = ( MappingProxyType({}) if rate_limit_check is None - else await async_admit_io_token_turn(rate_limit_check.dual_cache, deployment, get_utc_datetime()) # pyright: ignore[reportArgumentType] # the configured deployment dict + else await async_admit_io_token_turn(rate_limit_check.usage, deployment, get_utc_datetime()) # pyright: ignore[reportArgumentType] # the configured deployment dict ) budget_limiter: Final[object] = _budget_limiter(core) if budget_limiter is not None: diff --git a/litellm/router_rust/routing_stores.py b/litellm/router_rust/routing_stores.py index 9d120c0fea4..f3a6a485a55 100644 --- a/litellm/router_rust/routing_stores.py +++ b/litellm/router_rust/routing_stores.py @@ -55,7 +55,7 @@ class UsageWindow(_Step): minute: str def counter(self, kind: CounterKind) -> UsageCounter: - return UsageCounter(self.deployment_id, self.model, self.minute, kind) + return UsageCounter("global_router", self.deployment_id, self.model, self.minute, kind) class CooldownsStep(_Step): diff --git a/litellm/router_rust/steps/callbacks.py b/litellm/router_rust/steps/callbacks.py index 48cb998df9b..a8a98fd44b3 100644 --- a/litellm/router_rust/steps/callbacks.py +++ b/litellm/router_rust/steps/callbacks.py @@ -51,6 +51,7 @@ from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.redis.stores.router_usage import RouterUsage from litellm.router_python_host import PythonHost from litellm.router_rust.steps.selection import StrategySelectors from litellm.router_rust.surface.registry import HostRegistry @@ -109,6 +110,7 @@ class OptionalChecks: owner: "GlobalCallbackHolder", host: "PythonHost", cache: DualCache, + usage: "RouterUsage", hosted: "HostRegistry", ctx: object, arguments: Mapping[str, object], @@ -116,6 +118,7 @@ class OptionalChecks: self._owner: Final = owner self._host: Final = host self._cache: Final = cache + self._usage: Final = usage self._hosted: Final = hosted self._ctx: Final = ctx self.optional_callbacks: list[object] | None = None # mutable-ok: the router's live list, as Python keeps it @@ -313,7 +316,7 @@ class OptionalChecks: ) self.router_budget_logger = callback elif name == "enforce_model_rate_limits": - callback = ModelRateLimitingCheck(dual_cache=self._cache) + callback = ModelRateLimitingCheck(usage=self._usage) if callback is None: continue if self.optional_callbacks is not None and any( diff --git a/litellm/router_rust/steps/selection.py b/litellm/router_rust/steps/selection.py index 183350bc05f..f7a982356b6 100644 --- a/litellm/router_rust/steps/selection.py +++ b/litellm/router_rust/steps/selection.py @@ -35,6 +35,7 @@ from litellm.types.router import RoutingGroup, RoutingStrategy if TYPE_CHECKING: from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache + from litellm.redis.stores.router_usage import RouterUsage from litellm.router_python_host import PythonHost Selector: TypeAlias = ( @@ -68,6 +69,7 @@ def build_selector( strategy: str | None, args: Mapping[str, object], cache: "DualCache", + usage: "RouterUsage", reservations: "InMemoryCache | None" = None, ) -> Selector | None: """``_build_strategy_selector`` without registering: ``None`` for simple-shuffle and every @@ -80,7 +82,7 @@ def build_selector( case "usage-based-routing": return LowestTPMLoggingHandler(router_cache=cache, routing_args=routing_args) case "usage-based-routing-v2": - return LowestTPMLoggingHandler_v2(router_cache=cache, routing_args=routing_args) + return LowestTPMLoggingHandler_v2(usage=usage, routing_args=routing_args) case "latency-based-routing": return LowestLatencyLoggingHandler(router_cache=cache, routing_args=routing_args) case "cost-based-routing": @@ -98,11 +100,12 @@ class StrategySelectors: """The selectors of one router: the default group's, each routing group's, and the lazily built per-request override ones.""" - def __init__(self, host: "PythonHost", cache: "DualCache") -> None: + def __init__(self, host: "PythonHost", cache: "DualCache", usage: "RouterUsage") -> None: from litellm.caching.in_memory_cache import InMemoryCache self._host: Final = host self._cache: Final = cache + self._usage: Final = usage # `_core_least_busy_reservations`: the picks least-busy reserved, shared by its selectors. self.reservations: Final = InMemoryCache(max_size_in_memory=SELECTION_RESERVATION_MAX_KEYS) self._strategy: str | None = None @@ -144,7 +147,7 @@ class StrategySelectors: self._strategy = normalize_strategy(strategy) if configured: self._configured = self._strategy - self._default = build_selector(self._strategy, args, self._cache, self.reservations) + self._default = build_selector(self._strategy, args, self._cache, self._usage, self.reservations) if self._default is not None: self._register(self._default) @@ -158,7 +161,7 @@ class StrategySelectors: if self._default is None or self._strategy != self._configured: return try: - rebuilt: Final = build_selector(self._strategy, args, self._cache, self.reservations) + rebuilt: Final = build_selector(self._strategy, args, self._cache, self._usage, self.reservations) except (TypeError, ValidationError): verbose_router_logger.exception( "Invalid routing_strategy_args %s for '%s'; keeping the previous ones", dict(args), self._strategy @@ -177,7 +180,7 @@ class StrategySelectors: for name, strategy, args in groups if ( selector := build_selector( - None if strategy == "priority" else strategy, args, self._cache, self.reservations + None if strategy == "priority" else strategy, args, self._cache, self._usage, self.reservations ) ) is not None @@ -239,7 +242,7 @@ class StrategySelectors: return self.default with self._lock: if strategy not in self._overrides: - self._overrides[strategy] = build_selector(strategy, {}, self._cache, self.reservations) + self._overrides[strategy] = build_selector(strategy, {}, self._cache, self._usage, self.reservations) selector: Final = self._overrides[strategy] if selector is not None and request is not None and strategy not in self._registered_strategies(): from litellm.litellm_core_utils.litellm_logging import Logging diff --git a/litellm/router_rust/surface/legacy_members.py b/litellm/router_rust/surface/legacy_members.py index 640fdf739c0..9b436d6ce3e 100644 --- a/litellm/router_rust/surface/legacy_members.py +++ b/litellm/router_rust/surface/legacy_members.py @@ -23,7 +23,7 @@ from typing import TYPE_CHECKING, Final from typing_extensions import NotRequired, ReadOnly, TypedDict -from litellm._internal_context import with_service_target +from litellm._internal_context import request_parent_otel_span, with_service_target from litellm._logging import verbose_router_logger from litellm.constants import ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger @@ -423,10 +423,11 @@ class RustLegacyMembers(RustRegistryMembers, RustOptionalChecks): awaitable: Final[Awaitable[object]] = response # pyright: ignore[reportAssignmentType] # the endpoint's awaitable response = await awaitable # rebind-ok: awaited here except Exception: - with contextlib.suppress(Exception): - await async_io_token_refund_failure( - self._cache, kwargs, parent_otel_span=get_parent_otel_span_from_kwargs(kwargs) - ) + with ( + contextlib.suppress(Exception), + request_parent_otel_span(get_parent_otel_span_from_kwargs(kwargs)), + ): + await async_io_token_refund_failure(self._usage, kwargs) raise await self.increment_deployment_usage_for_response(response=response, request_kwargs=kwargs) return await self.set_response_headers( @@ -513,7 +514,9 @@ class RustLegacyMembers(RustRegistryMembers, RustOptionalChecks): if model_group is None or deployment_id is None: return model: Final = None if deployment_name is None else str(deployment_name) - counter: Final = UsageCounter(str(deployment_id), model, get_utc_datetime().strftime("%H-%M"), "rpm") + counter: Final = UsageCounter( + "global_router", str(deployment_id), model, get_utc_datetime().strftime("%H-%M"), "rpm" + ) await self._usage.increment(counter=counter, value=1, ttl=_USAGE_TTL) # -- the plugin context ------------------------------------------------------------------------ diff --git a/litellm/router_rust/surface/model_info.py b/litellm/router_rust/surface/model_info.py index dd67e252370..7cf4d46edc3 100644 --- a/litellm/router_rust/surface/model_info.py +++ b/litellm/router_rust/surface/model_info.py @@ -7,8 +7,8 @@ host registry. router-core holds the rules: a deployment's model info (``get_router_model_info``, ``get_deployment_model_info``: ``router-core/src/model_info.rs``), a model group's info and the configured-value reads (``_set_model_group_info``, ``get_model_listing_info``, the configured token -limits, mode, display name and service tiers: ``router-core/src/group_info.rs``), and the usage -counters' keys (``RouterCoreHandle.usage_key``). What lives in this process is answered here as a +limits, mode, display name and service tiers: ``router-core/src/group_info.rs``); the usage +counters' keys are the ``router_usage`` store's. What lives in this process is answered here as a query asks for it, at query time: the cost map (``get_model_info``, ``litellm.model_cost``), provider resolution, the wildcard pattern routes, what the background refresh discovered, and a group deployment's validated params, supported params, endpoints, fast mode and reasoning efforts. @@ -33,7 +33,6 @@ from pydantic import BaseModel, TypeAdapter from typing_extensions import overload import litellm -from litellm._internal_context import service_target from litellm._logging import verbose_router_logger from litellm.constants import DEFAULT_MAX_LRU_CACHE_SIZE from litellm.litellm_core_utils.add_retry_fallback_headers import ( @@ -56,7 +55,7 @@ from litellm.litellm_core_utils.reasoning_effort_capability import ( from litellm.litellm_core_utils.routing_common_utils import team_may_use_deployment from litellm.llms.anthropic.common_utils import AnthropicModelInfo from litellm.llms.bedrock.common_utils import get_bedrock_base_model -from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, RouterUsage, UsageCounter +from litellm.redis.stores.router_usage import RouterUsage, UsageCounter from litellm.router_python_host import python_host from litellm.router_utils.model_group_pricing import deployment_pricing from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import build_io_token_rate_limit_headers @@ -811,7 +810,10 @@ class HostModelInfo: return cached[2] deployments: Final = tuple(pair for pair in members if pair[0] is not None and pair[1] is not None) counters: Final = tuple( - tuple(UsageCounter(id, litellm_model, current_minute, kind) for id, litellm_model in deployments) # pyright: ignore[reportArgumentType] # what the listing names, as the Python core passed it + tuple( + UsageCounter("global_router", id, litellm_model, current_minute, kind) # pyright: ignore[reportArgumentType] # what the listing names, as the Python core passed it + for id, litellm_model in deployments + ) for kind in ("tpm", "rpm") ) if len(self._group_counters_cache) >= DEFAULT_MAX_LRU_CACHE_SIZE: @@ -839,30 +841,31 @@ class HostModelInfo: """ Returns current ITPM/OTPM usage for a model group (sum across deployments). """ - dt: Final = get_utc_datetime() - current_minute: Final = dt.strftime("%H-%M") - itpm_keys: Final[list[str]] = [] # mutable-ok: ported from the Python core unchanged - otpm_keys: Final[list[str]] = [] # mutable-ok: ported from the Python core unchanged - + current_minute: Final = get_utc_datetime().strftime("%H-%M") model_list: Final = self._core_get_model_list(model_name=model_group) if model_list is None: return None, None - for model in model_list: - model_id: str | None = model.get("model_info", {}).get("id") - litellm_model: str | None = model["litellm_params"].get("model") - if model_id is None or litellm_model is None: - continue - itpm_keys.append(self._core_handle().usage_key(model_id, litellm_model, "itpm", current_minute)) - otpm_keys.append(self._core_handle().usage_key(model_id, litellm_model, "otpm", current_minute)) + named: Final[tuple[tuple[object, object], ...]] = tuple( + (model.get("model_info", {}).get("id"), model["litellm_params"].get("model")) for model in model_list + ) + deployments: Final = tuple( + (f"{model_id}", f"{litellm_model}") + for model_id, litellm_model in named + if model_id is not None and litellm_model is not None + ) + itpm_counters: Final = tuple( + UsageCounter("global_router", model_id, litellm_model, current_minute, "itpm") + for model_id, litellm_model in deployments + ) + otpm_counters: Final = tuple( + UsageCounter("global_router", model_id, litellm_model, current_minute, "otpm") + for model_id, litellm_model in deployments + ) + combined_values: Final = await self._core_usage_store.get_many((*itpm_counters, *otpm_counters)) - with service_target(ROUTER_USAGE_TARGET): - combined_values: Final = await self._core_cache.async_batch_get_cache(keys=itpm_keys + otpm_keys) - if combined_values is None: - return None, None - - itpm_values: Final = combined_values[: len(itpm_keys)] - otpm_values: Final = combined_values[len(itpm_keys) :] + itpm_values: Final = combined_values[: len(itpm_counters)] + otpm_values: Final = combined_values[len(itpm_counters) :] total_itpm: int | None = None # rebind-ok: ported from the Python core unchanged for value in itpm_values: diff --git a/litellm/router_strategy/base_routing_strategy.py b/litellm/router_strategy/base_routing_strategy.py index 43acae1f62f..f1d0373ef25 100644 --- a/litellm/router_strategy/base_routing_strategy.py +++ b/litellm/router_strategy/base_routing_strategy.py @@ -9,28 +9,27 @@ from typing import Final from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger -from litellm.caching.caching import DualCache -from litellm.caching.redis_cache import RedisPipelineIncrementOperation, log_redis_failure +from litellm.caching.redis_cache import log_redis_failure from litellm.constants import DEFAULT_REDIS_SYNC_INTERVAL -from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET +from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, CounterIncrement, RouterUsage, UsageCounter class BaseRoutingStrategy(ABC): def __init__( self, - dual_cache: DualCache, + usage: RouterUsage, should_batch_redis_writes: bool, default_sync_interval: float | None, ): - self.dual_cache = dual_cache - self.redis_increment_operation_queue: list[RedisPipelineIncrementOperation] = [] + self.usage = usage + self.redis_increment_operation_queue: list[CounterIncrement] = [] self._sync_task: asyncio.Task[None] | None = None self._retired = False self._flush_after_retire: asyncio.Task[None] | None = None if should_batch_redis_writes: self.setup_sync_task(default_sync_interval) - self.in_memory_keys_to_update: set[str] = set() # Set with max size of 1000 keys + self.in_memory_keys_to_update: set[UsageCounter] = set() # Set with max size of 1000 keys def setup_sync_task(self, default_sync_interval: float | None): """Setup the sync task in a way that's compatible with FastAPI""" @@ -56,7 +55,7 @@ class BaseRoutingStrategy(ABC): def _schedule_flush_after_retire(self) -> None: if self._flush_after_retire is not None and not self._flush_after_retire.done(): return - if not self.redis_increment_operation_queue or self.dual_cache.redis_cache is None: + if not self.redis_increment_operation_queue or not self.usage.has_redis: return try: loop: Final = asyncio.get_running_loop() @@ -79,19 +78,7 @@ class BaseRoutingStrategy(ABC): except asyncio.CancelledError: pass - async def _increment_value_list_in_current_window( - self, increment_list: list[tuple[str, int]], ttl: int - ) -> list[float]: - """ - Increment a list of values in the current window - """ - results: Final = [] - for key, value in increment_list: - result = await self._increment_value_in_current_window(key=key, value=value, ttl=ttl) - results.append(result) - return results - - async def _increment_value_in_current_window(self, key: str, value: float, ttl: int): + async def _increment_value_in_current_window(self, counter: UsageCounter, value: float, ttl: int): """ Increment spend within existing budget window @@ -100,19 +87,11 @@ class BaseRoutingStrategy(ABC): - Increments the spend in memory cache (so spend instantly updated in memory) - Queues the increment operation to Redis Pipeline (using batched pipeline to optimize performance. Using Redis for multi instance environment of LiteLLM) """ - result: Final = await self.dual_cache.in_memory_cache.async_increment( - key=key, - value=value, - ttl=ttl, - ) - increment_op: Final = RedisPipelineIncrementOperation( - key=key, - increment_value=value, - ttl=ttl, - ) + increment_op: Final = CounterIncrement(counter=counter, amount=value, ttl=ttl) + result: Final = self.usage.increment_local(increment_op) self.redis_increment_operation_queue.append(increment_op) - self.add_to_in_memory_keys_to_update(key=key) + self.add_to_in_memory_keys_to_update(counter) if self._retired: self._schedule_flush_after_retire() return result @@ -148,35 +127,26 @@ class BaseRoutingStrategy(ABC): Only runs if Redis is initialized """ try: - if not self.dual_cache.redis_cache: + if not self.usage.has_redis: return # Redis is not initialized if len(self.redis_increment_operation_queue) > 0: - # Compress operations for the same key - compressed_ops: Final[dict[str, RedisPipelineIncrementOperation]] = {} - ops_to_remove: Final = [] - for idx, op in enumerate(self.redis_increment_operation_queue): - if op["key"] in compressed_ops: - # Add to existing increment - compressed_ops[op["key"]]["increment_value"] += op["increment_value"] - else: - compressed_ops[op["key"]] = op - - ops_to_remove.append(idx) - - # Convert back to list - compressed_queue: Final = list(compressed_ops.values()) - - increment_result: Final = await self.dual_cache.redis_cache.async_increment_pipeline( - increment_list=compressed_queue, + queued: Final = tuple(self.redis_increment_operation_queue) + compressed_queue: Final = tuple( + CounterIncrement( + counter=counter, + amount=sum(op.amount for op in queued if op.counter == counter), + ttl=next(op.ttl for op in queued if op.counter == counter), + ) + for counter in dict.fromkeys(op.counter for op in queued) ) - self.redis_increment_operation_queue = [ - op for idx, op in enumerate(self.redis_increment_operation_queue) if idx not in ops_to_remove - ] + increment_result: Final = await self.usage.increment_remote(compressed_queue) + + self.redis_increment_operation_queue = self.redis_increment_operation_queue[len(queued) :] if increment_result is not None: - return_result = {key["key"]: op for key, op in zip(compressed_queue, increment_result)} + return_result = {op.counter: total for op, total in zip(compressed_queue, increment_result)} else: return_result = {} return return_result @@ -185,19 +155,13 @@ class BaseRoutingStrategy(ABC): log_redis_failure(verbose_router_logger, logging.ERROR, "Error syncing in-memory cache with Redis", e) self.redis_increment_operation_queue = [] - def add_to_in_memory_keys_to_update(self, key: str): - self.in_memory_keys_to_update.add(key) + def add_to_in_memory_keys_to_update(self, counter: UsageCounter): + self.in_memory_keys_to_update.add(counter) - def get_key_pattern_to_sync(self) -> str | None: - """ - Get the key pattern to sync - """ - return None - - def get_in_memory_keys_to_update(self) -> set[str]: + def get_in_memory_keys_to_update(self) -> set[UsageCounter]: return self.in_memory_keys_to_update - def get_and_reset_in_memory_keys_to_update(self) -> set[str]: + def get_and_reset_in_memory_keys_to_update(self) -> set[UsageCounter]: """Atomic get and reset in-memory keys to update""" keys: Final = self.in_memory_keys_to_update self.in_memory_keys_to_update = set() @@ -221,7 +185,7 @@ class BaseRoutingStrategy(ABC): try: # No need to sync if Redis cache is not initialized - if self.dual_cache.redis_cache is None: + if not self.usage.has_redis: return # 2. Fetch all current provider spend from Redis to update in-memory cache @@ -232,10 +196,9 @@ class BaseRoutingStrategy(ABC): cache_keys_list: Final = list(cache_keys) # 1. Snapshot in-memory before - in_memory_before_dict: Final = {} - in_memory_before: Final = await self.dual_cache.in_memory_cache.async_batch_get_cache(keys=cache_keys_list) - for k, v in zip(cache_keys_list, in_memory_before): - in_memory_before_dict[k] = float(v or 0) + in_memory_before_dict: Final = { + k: float(v or 0) for k, v in zip(cache_keys_list, self.usage.get_local_many(cache_keys_list)) + } # 1. Push all provider spend increments to Redis redis_values: Final = await self._push_in_memory_increments_to_redis() @@ -246,7 +209,7 @@ class BaseRoutingStrategy(ABC): for key in cache_keys_list: redis_val = float(redis_values.get(key, 0) or 0) before = float(in_memory_before_dict.get(key, 0) or 0) - after = float(await self.dual_cache.in_memory_cache.async_get_cache(key=key) or 0) + after = float(self.usage.get_local(key) or 0) delta = after - before if after <= redis_val: merged = redis_val + delta @@ -259,7 +222,7 @@ class BaseRoutingStrategy(ABC): # import os # os._exit(1) # raise Exception(f"Redis is behind in-memory cache for key: {key}. This should not happen, since we should be updating redis with in-memory cache.") - await self.dual_cache.in_memory_cache.async_set_cache(key=key, value=merged) + self.usage.set_local(key, merged) except Exception as e: verbose_router_logger.exception("Error syncing in-memory cache with Redis: %s", e) diff --git a/litellm/router_strategy/lowest_tpm_rpm_v2.py b/litellm/router_strategy/lowest_tpm_rpm_v2.py index 8d2f45756e4..db024d47a53 100644 --- a/litellm/router_strategy/lowest_tpm_rpm_v2.py +++ b/litellm/router_strategy/lowest_tpm_rpm_v2.py @@ -15,14 +15,17 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._internal_context import with_service_target from litellm._logging import verbose_logger, verbose_router_logger -from litellm.caching.caching import DualCache from litellm.constants import USAGE_BASED_ROUTING_COUNTED_TOKENS_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type -from litellm.litellm_core_utils.core_helpers import get_parent_otel_span_from_kwargs from litellm.litellm_core_utils.deployment_pricing import deployment_backend_models from litellm.litellm_core_utils.token_counter import embedding_aware_token_count -from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET +from litellm.redis.stores.router_usage import ( + ROUTER_USAGE_TARGET, + CounterIncrement, + RouterUsage, + UsageCounter, +) from litellm.router_utils.router_callbacks.track_deployment_metrics import get_counted_usage_tokens from litellm.types.router import RouterErrors from litellm.types.utils import LiteLLMPydanticObjectBase, StandardLoggingPayload @@ -55,21 +58,21 @@ _active_prefetched_usage: Final[ContextVar["PrefetchedUsage | None"]] = ContextV @dataclass(frozen=True) class PrefetchedUsage: """ - tpm/rpm counter values another read of this request already fetched from the router cache. + tpm/rpm counter values another read of this request already fetched from the usage store. `values` is None when that read failed, which is what `async_batch_get_cache` returns on failure. """ - keys: frozenset[str] - values: Mapping[str, object] | None + counters: frozenset[UsageCounter] + values: Mapping[UsageCounter, float | None] | None - def covers(self, keys: Sequence[str]) -> bool: - return self.keys.issuperset(keys) + def covers(self, counters: Sequence[UsageCounter]) -> bool: + return self.counters.issuperset(counters) - def values_for(self, keys: Sequence[str]) -> list[object | None] | None: + def values_for(self, counters: Sequence[UsageCounter]) -> list[float | None] | None: if self.values is None: return None - return [self.values.get(key) for key in keys] + return [self.values.get(counter) for counter in counters] @staticmethod @contextmanager @@ -103,12 +106,11 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): logged_failure: int = 0 default_cache_time_seconds: int = 1 * 60 * 60 # 1 hour - def __init__(self, router_cache: DualCache, routing_args: dict = {}): - self.router_cache = router_cache + def __init__(self, usage: RouterUsage, routing_args: dict = {}): self.routing_args = RoutingArgs(**routing_args) BaseRoutingStrategy.__init__( self, - dual_cache=router_cache, + usage=usage, should_batch_redis_writes=True, default_sync_interval=0.1, ) @@ -129,11 +131,11 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): dt: Final = get_utc_datetime() current_minute: Final = dt.strftime("%H-%M") - model_id: Final = deployment.get("model_info", {}).get("id") - deployment_name: Final = deployment.get("litellm_params", {}).get("model") - rpm_key: Final = f"{model_id}:{deployment_name}:rpm:{current_minute}" + model_id: Final[object] = deployment.get("model_info", {}).get("id") + deployment_name: Final[object] = deployment.get("litellm_params", {}).get("model") + rpm_counter: Final = UsageCounter.for_deployment(f"{model_id}", f"{deployment_name}", "rpm", current_minute) - local_result: Final = self.router_cache.get_cache(key=rpm_key, local_only=True) # check local result first + local_result: Final = self.usage.get_local(rpm_counter) # check local result first deployment_rpm = None if deployment_rpm is None: @@ -168,7 +170,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): else: # if local result below limit, check redis ## prevent unnecessary redis checks - result: Final = self.router_cache.increment_cache(key=rpm_key, value=1, ttl=self.routing_args.ttl) + result: Final = self.usage.increment_sync(rpm_counter, 1, self.routing_args.ttl) if result is not None and result > deployment_rpm: raise litellm.RateLimitError( message=f"Deployment over defined rpm limit={deployment_rpm}. current usage={result}", @@ -208,13 +210,11 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): # ------------ dt: Final = get_utc_datetime() current_minute: Final = dt.strftime("%H-%M") - model_id: Final = deployment.get("model_info", {}).get("id") - deployment_name: Final = deployment.get("litellm_params", {}).get("model") + model_id: Final[object] = deployment.get("model_info", {}).get("id") + deployment_name: Final[object] = deployment.get("litellm_params", {}).get("model") - rpm_key: Final = f"{model_id}:{deployment_name}:rpm:{current_minute}" - local_result: Final = await self.router_cache.async_get_cache( - key=rpm_key, local_only=True - ) # check local result first + rpm_counter: Final = UsageCounter.for_deployment(f"{model_id}", f"{deployment_name}", "rpm", current_minute) + local_result: Final = self.usage.get_local(rpm_counter) # check local result first deployment_rpm = None if deployment_rpm is None: @@ -243,7 +243,7 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): ) else: # if local result below limit, check redis ## prevent unnecessary redis checks - result = await self._increment_value_in_current_window(key=rpm_key, value=1, ttl=self.routing_args.ttl) + result = await self._increment_value_in_current_window(rpm_counter, 1, self.routing_args.ttl) if result is not None and result > deployment_rpm: raise litellm.RateLimitError( message=f"Deployment over defined rpm limit={deployment_rpm}. current usage={result}", @@ -278,8 +278,8 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): if standard_logging_object is None: raise ValueError("standard_logging_object not passed in.") model_group: Final = standard_logging_object.get("model_group") - model: Final = standard_logging_object["hidden_params"].get("litellm_model_name") - id = standard_logging_object.get("model_id") + model: Final[object] = standard_logging_object["hidden_params"].get("litellm_model_name") + id: object = standard_logging_object.get("model_id") if model_group is None or id is None or model is None: return elif isinstance(id, int): @@ -293,14 +293,12 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): dt: Final = get_utc_datetime() current_minute: Final = dt.strftime("%H-%M") # use the same timezone regardless of system clock - tpm_key: Final = f"{id}:{model}:tpm:{current_minute}" - # ------------ - # Update usage - # ------------ - # update cache - ## TPM - self.router_cache.increment_cache(key=tpm_key, value=total_tokens, ttl=self.routing_args.ttl) + self.usage.increment_sync( + UsageCounter.for_deployment(f"{id}", f"{model}", "tpm", current_minute), + total_tokens, + self.routing_args.ttl, + ) ### TESTING ### if self.test_flag: self.logged_success += 1 @@ -334,7 +332,6 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): deployment_id=str(model_id), deployment_name=model if isinstance(model, str) else None, tokens=total_tokens if counted_tokens is None else max(0, total_tokens - counted_tokens), - parent_otel_span=get_parent_otel_span_from_kwargs(kwargs), ) ### TESTING ### @@ -346,16 +343,10 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): ) @with_service_target(ROUTER_USAGE_TARGET) - async def async_increment_tpm( - self, *, deployment_id: str, deployment_name: str | None, tokens: int, parent_otel_span: Span | None - ) -> None: + async def async_increment_tpm(self, *, deployment_id: str, deployment_name: str | None, tokens: int) -> None: current_minute: Final = get_utc_datetime().strftime("%H-%M") - await self.router_cache.async_increment_cache_post_call( - key=f"{deployment_id}:{deployment_name}:tpm:{current_minute}", - value=tokens, - ttl=self.routing_args.ttl, - parent_otel_span=parent_otel_span, - ) + counter: Final = UsageCounter.for_deployment(deployment_id, f"{deployment_name}", "tpm", current_minute) + await self.usage.increment_post_call((CounterIncrement(counter, tokens, self.routing_args.ttl),)) def _return_potential_deployments( self, @@ -414,10 +405,10 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): self, model_group: str, healthy_deployments: list, - tpm_keys: list, - tpm_values: Sequence | None, - rpm_keys: list, - rpm_values: Sequence | None, + tpm_counters: Sequence[UsageCounter], + tpm_values: Sequence[float | None] | None, + rpm_counters: Sequence[UsageCounter], + rpm_values: Sequence[float | None] | None, messages: list[dict[str, str]] | None = None, input: str | list | None = None, ) -> dict | None: @@ -428,13 +419,12 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): if tpm_values is None or rpm_values is None: return None - tpm_dict = {} # {model_id: 1, ..} - for idx, key in enumerate(tpm_keys): - tpm_dict[tpm_keys[idx].split(":")[0]] = tpm_values[idx] - - rpm_dict: Final = {} # {model_id: 1, ..} - for idx, key in enumerate(rpm_keys): - rpm_dict[rpm_keys[idx].split(":")[0]] = rpm_values[idx] + tpm_dict = { + counter.deployment_id.split(":")[0]: value for counter, value in zip(tpm_counters, tpm_values) + } # {model_id: 1, ..} + rpm_dict: Final = { + counter.deployment_id.split(":")[0]: value for counter, value in zip(rpm_counters, rpm_values) + } # {model_id: 1, ..} try: input_tokens = embedding_aware_token_count( @@ -473,17 +463,25 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): else: return None - def usage_counter_keys(self, healthy_deployments: list) -> tuple[list[str], list[str]]: - """The `::tpm:` and `::rpm:` counter keys selection reads.""" + def usage_counters( + self, healthy_deployments: Sequence[object] + ) -> tuple[tuple[UsageCounter, ...], tuple[UsageCounter, ...]]: + """The tpm and rpm counters selection reads, this minute's for each deployment.""" current_minute: Final = get_utc_datetime().strftime("%H-%M") - prefixes: Final = tuple( - f"{m.get('model_info', {}).get('id')}:{m.get('litellm_params', {}).get('model')}" + named: Final[tuple[tuple[object, object], ...]] = tuple( + (m.get("model_info", {}).get("id"), m.get("litellm_params", {}).get("model")) for m in healthy_deployments if isinstance(m, dict) ) return ( - [f"{prefix}:tpm:{current_minute}" for prefix in prefixes], - [f"{prefix}:rpm:{current_minute}" for prefix in prefixes], + tuple( + UsageCounter.for_deployment(f"{model_id}", f"{model}", "tpm", current_minute) + for model_id, model in named + ), + tuple( + UsageCounter.for_deployment(f"{model_id}", f"{model}", "rpm", current_minute) + for model_id, model in named + ), ) @with_service_target(ROUTER_USAGE_TARGET) @@ -508,20 +506,18 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): healthy_deployments, ) - tpm_keys, rpm_keys = self.usage_counter_keys(healthy_deployments) - combined_tpm_rpm_keys: Final = tpm_keys + rpm_keys + tpm_counters, rpm_counters = self.usage_counters(healthy_deployments) + combined_tpm_rpm_counters: Final = tpm_counters + rpm_counters prefetched_usage: Final = PrefetchedUsage.active() - if prefetched_usage is not None and prefetched_usage.covers(combined_tpm_rpm_keys): - combined_tpm_rpm_values = prefetched_usage.values_for(combined_tpm_rpm_keys) + if prefetched_usage is not None and prefetched_usage.covers(combined_tpm_rpm_counters): + combined_tpm_rpm_values = prefetched_usage.values_for(combined_tpm_rpm_counters) else: - combined_tpm_rpm_values = await self.router_cache.async_batch_get_cache( - keys=combined_tpm_rpm_keys - ) # [1, 2, None, ..] + combined_tpm_rpm_values = await self.usage.get_many(combined_tpm_rpm_counters) # [1, 2, None, ..] if combined_tpm_rpm_values is not None: - tpm_values = combined_tpm_rpm_values[: len(tpm_keys)] - rpm_values = combined_tpm_rpm_values[len(tpm_keys) :] + tpm_values = combined_tpm_rpm_values[: len(tpm_counters)] + rpm_values = combined_tpm_rpm_values[len(tpm_counters) :] else: tpm_values = None rpm_values = None @@ -529,9 +525,9 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): deployment: Final = self._common_checks_available_deployment( model_group=model_group, healthy_deployments=healthy_deployments, - tpm_keys=tpm_keys, + tpm_counters=tpm_counters, tpm_values=tpm_values, - rpm_keys=rpm_keys, + rpm_counters=rpm_counters, rpm_values=rpm_values, messages=messages, input=input, @@ -614,35 +610,16 @@ class LowestTPMLoggingHandler_v2(BaseRoutingStrategy, CustomLogger): healthy_deployments, ) - dt: Final = get_utc_datetime() - current_minute: Final = dt.strftime("%H-%M") - tpm_keys: Final = [] - rpm_keys: Final = [] - for m in healthy_deployments: - if isinstance(m, dict): - id = m.get("model_info", {}).get( - "id" - ) # a deployment should always have an 'id'. this is set in router.py - deployment_name = m.get("litellm_params", {}).get("model") - tpm_key = f"{id}:{deployment_name}:tpm:{current_minute}" - rpm_key = f"{id}:{deployment_name}:rpm:{current_minute}" - - tpm_keys.append(tpm_key) - rpm_keys.append(rpm_key) - - tpm_values: Final = self.router_cache.batch_get_cache( - keys=tpm_keys, parent_otel_span=parent_otel_span - ) # [1, 2, None, ..] - rpm_values: Final = self.router_cache.batch_get_cache( - keys=rpm_keys, parent_otel_span=parent_otel_span - ) # [1, 2, None, ..] + tpm_counters, rpm_counters = self.usage_counters(healthy_deployments) + tpm_values: Final = self.usage.get_many_sync(tpm_counters) # [1, 2, None, ..] + rpm_values: Final = self.usage.get_many_sync(rpm_counters) # [1, 2, None, ..] deployment: Final = self._common_checks_available_deployment( model_group=model_group, healthy_deployments=healthy_deployments, - tpm_keys=tpm_keys, + tpm_counters=tpm_counters, tpm_values=tpm_values, - rpm_keys=rpm_keys, + rpm_counters=rpm_counters, rpm_values=rpm_values, messages=messages, input=input, 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 4b618aba476..0e551de7842 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 @@ -18,21 +18,19 @@ import threading import time from collections import OrderedDict from collections.abc import Callable, Mapping, MutableMapping -from dataclasses import dataclass +from dataclasses import asdict, dataclass, replace from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final +from typing import Any, Final import httpx -from pydantic import TypeAdapter +from pydantic import TypeAdapter, ValidationError 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 @@ -40,25 +38,19 @@ 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_COUNTER_KEY, ITPM_RESERVED_KEY, - OTPM_CACHE_KEY, + OTPM_COUNTER_KEY, OTPM_RESERVED_KEY, ) from litellm.litellm_core_utils.token_counter import deployment_openai_model, offload_token_count from litellm.llms.azure.ptu_capacity import is_str_keyed_mapping -from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET -from litellm.types.router import RouterCacheEnum, RouterErrors +from litellm.redis.stores.router_usage import MAX_INCREMENT, ROUTER_USAGE_TARGET, RouterUsage, UsageCounter +from litellm.types.router import RouterErrors from litellm.utils import get_utc_datetime _MODEL_INFO_ADAPTER: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) - -if TYPE_CHECKING: - from opentelemetry.trace import Span as _Span - - Span = _Span -else: - Span = Any +_STASHED_COUNTER: Final[TypeAdapter[UsageCounter]] = TypeAdapter(UsageCounter) RoutingArgsTTL: Final = 60 @@ -177,16 +169,17 @@ def _without_io_token_limit_keys(section: object) -> object: return {key: value for key, value in section.items() if key not in _IO_TOKEN_LIMIT_KEYS} -def _get_cache_keys(deployment: dict, current_minute: str) -> tuple[str, str] | None: - model_id: Final = deployment.get("model_info", {}).get("id") - deployment_name: Final = deployment.get("litellm_params", {}).get("model") +def _get_counters(deployment: dict, current_minute: str) -> tuple[UsageCounter, UsageCounter] | None: + model_id: Final[object] = deployment.get("model_info", {}).get("id") + deployment_name: Final[object] = deployment.get("litellm_params", {}).get("model") # Without both a deployment id and model name the key would collapse to a # shared "None:None" bucket across misconfigured deployments, so bail out. if model_id is None or deployment_name is None: return None - itpm_key = RouterCacheEnum.ITPM.value.format(id=model_id, model=deployment_name, current_minute=current_minute) - otpm_key = RouterCacheEnum.OTPM.value.format(id=model_id, model=deployment_name, current_minute=current_minute) - return itpm_key, otpm_key + return ( + UsageCounter("global_router", f"{model_id}", f"{deployment_name}", current_minute, "itpm"), + UsageCounter("global_router", f"{model_id}", f"{deployment_name}", current_minute, "otpm"), + ) def _served_model(deployment: object) -> str: @@ -387,15 +380,16 @@ 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, int | str | None]: + itpm_reserved: int, otpm_reserved: int, itpm: UsageCounter | None, otpm: UsageCounter | None +) -> Mapping[str, object]: + """The reservation as plain fields, which survive the copies logging makes of the request metadata.""" 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, - OTPM_CACHE_KEY: otpm_cache_key, + ITPM_COUNTER_KEY: None if itpm is None else asdict(itpm), + OTPM_COUNTER_KEY: None if otpm is None else asdict(otpm), } ) @@ -407,33 +401,44 @@ def _stash_reservation_in_metadata(request: _IoTokenRateLimitRequest | None, res metadata[IO_TOKEN_ATTEMPT_KEY] = current_attempt_id() -def _key_for_this_minute(reserved_key: str) -> str: - counter, _, _ = reserved_key.rpartition(":") - return f"{counter}:{get_utc_datetime().strftime('%H-%M')}" +def _this_minute(counter: UsageCounter) -> UsageCounter: + return replace(counter, window=get_utc_datetime().strftime("%H-%M")) + + +_Reservation = tuple[int, int, UsageCounter | None, UsageCounter | None] def _reconcile_increments( - reservation: tuple[int, int, str | None, str | None], billable_input: int, completion_tokens: int -) -> tuple[tuple[str, int], ...]: - itpm_reserved, otpm_reserved, itpm_key, otpm_key = reservation - deltas: Final = ((itpm_key, billable_input - itpm_reserved), (otpm_key, completion_tokens - otpm_reserved)) + reservation: _Reservation, billable_input: int, completion_tokens: int +) -> tuple[tuple[UsageCounter, int], ...]: + itpm_reserved, otpm_reserved, itpm, otpm = reservation + deltas: Final = ((itpm, billable_input - itpm_reserved), (otpm, completion_tokens - otpm_reserved)) return tuple( - (key if delta < 0 else _key_for_this_minute(key), delta) for key, delta in deltas if key is not None and delta + (counter if delta < 0 else _this_minute(counter), delta) + for counter, delta in deltas + if counter is not None and delta ) +def _stashed_counter(value: object) -> UsageCounter | None: + if value is None: + return None + try: + return _STASHED_COUNTER.validate_python(value) + except ValidationError: + return 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) +def _extract_reservation(reservation: Mapping[str, object]) -> _Reservation: return ( _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, + _stashed_counter(reservation.get(ITPM_COUNTER_KEY)), + _stashed_counter(reservation.get(OTPM_COUNTER_KEY)), ) @@ -465,9 +470,7 @@ def _stashed_reservation(kwargs: Mapping[str, object] | None) -> Mapping[str, ob 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]: +def _claim_reservation_from_kwargs(kwargs: Mapping[str, object] | None, *, any_attempt: bool = False) -> _Reservation: """ 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 @@ -535,8 +538,8 @@ def _rate_limit_error(limit_label: str, limit: int, current: float) -> litellm.R @with_service_target(ROUTER_USAGE_TARGET) def _sync_increment_with_rollback( - dual_cache: DualCache, - key: str, + usage: RouterUsage, + counter: UsageCounter, value: int, limit: int | None, *, @@ -544,51 +547,36 @@ def _sync_increment_with_rollback( ) -> None: if value <= 0 or limit is None: return - current: Final = dual_cache.increment_cache( - key=key, - value=value, - ttl=RoutingArgsTTL, - ) + if value > min(limit, MAX_INCREMENT): + raise _rate_limit_error(limit_label, limit, value) + current: Final = usage.increment_sync(counter, value, RoutingArgsTTL) if current is not None and current > limit: - dual_cache.increment_cache( - key=key, - value=-value, - ttl=RoutingArgsTTL, - ) + usage.increment_sync(counter, -value, RoutingArgsTTL) raise _rate_limit_error(limit_label, limit, current) @with_service_target(ROUTER_USAGE_TARGET) async def _increment_with_rollback( - dual_cache: DualCache, - key: str, + usage: RouterUsage, + counter: UsageCounter, value: int, limit: int | None, *, - parent_otel_span: Span | None = None, limit_label: str, ) -> None: if value <= 0 or limit is None: return - current: Final = await dual_cache.async_increment_cache( - key=key, - value=value, - ttl=RoutingArgsTTL, - parent_otel_span=parent_otel_span, - ) + if value > min(limit, MAX_INCREMENT): + raise _rate_limit_error(limit_label, limit, value) + current: Final = await usage.increment(counter, value, RoutingArgsTTL) if current is not None and current > limit: - await dual_cache.async_increment_cache( - key=key, - value=-value, - ttl=RoutingArgsTTL, - parent_otel_span=parent_otel_span, - ) + await usage.increment(counter, -value, RoutingArgsTTL) raise _rate_limit_error(limit_label, limit, current) @with_service_target(ROUTER_USAGE_TARGET) def io_token_pre_call_check( - dual_cache: DualCache, + usage: RouterUsage, deployment: dict, ) -> dict | None: itpm_limit, otpm_limit = get_deployment_io_token_limits(deployment) @@ -605,10 +593,10 @@ def io_token_pre_call_check( dt: Final = get_utc_datetime() current_minute: Final = dt.strftime("%H-%M") - cache_keys: Final = _get_cache_keys(deployment, current_minute) - if cache_keys is None: + counters: Final = _get_counters(deployment, current_minute) + if counters is None: return deployment - itpm_key, otpm_key = cache_keys + itpm_counter, otpm_counter = counters itpm_reserved = 0 otpm_reserved = 0 @@ -616,8 +604,8 @@ def io_token_pre_call_check( if itpm_limit is not None: itpm_reserved = _reservation_value(estimated_input, itpm_limit) _sync_increment_with_rollback( - dual_cache, - itpm_key, + usage, + itpm_counter, itpm_reserved, itpm_limit, limit_label="ITPM", @@ -627,19 +615,15 @@ def io_token_pre_call_check( otpm_reserved = otpm_to_reserve try: _sync_increment_with_rollback( - dual_cache, - otpm_key, + usage, + otpm_counter, otpm_reserved, otpm_limit, limit_label="OTPM", ) except Exception: if itpm_reserved > 0: - dual_cache.increment_cache( - key=itpm_key, - value=-itpm_reserved, - ttl=RoutingArgsTTL, - ) + usage.increment_sync(itpm_counter, -itpm_reserved, RoutingArgsTTL) raise _stash_reservation_in_metadata( @@ -647,8 +631,8 @@ def io_token_pre_call_check( _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, + itpm_counter if itpm_limit is not None else None, + otpm_counter if otpm_limit is not None else None, ), ) return deployment @@ -656,9 +640,8 @@ def io_token_pre_call_check( @with_service_target(ROUTER_USAGE_TARGET) async def async_io_token_pre_call_check( - dual_cache: DualCache, + usage: RouterUsage, deployment: dict, - parent_otel_span: Span | None = None, ) -> dict | None: itpm_limit, otpm_limit = get_deployment_io_token_limits(deployment) if itpm_limit is None and otpm_limit is None: @@ -669,7 +652,7 @@ async def async_io_token_pre_call_check( if request is not None and request.opens_responses_socket: _stash_reservation_in_metadata( request, - await async_admit_io_token_turn(dual_cache, deployment, get_utc_datetime(), parent_otel_span), + await async_admit_io_token_turn(usage, deployment, get_utc_datetime()), ) return deployment _model, served_model = _tokenizer_and_served_models(deployment) @@ -680,72 +663,66 @@ async def async_io_token_pre_call_check( dt: Final = get_utc_datetime() current_minute: Final = dt.strftime("%H-%M") - cache_keys: Final = _get_cache_keys(deployment, current_minute) - if cache_keys is None: + counters: Final = _get_counters(deployment, current_minute) + if counters is None: return deployment - itpm_key, otpm_key = cache_keys + itpm_counter, otpm_counter = counters # 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, + usage, + itpm_counter, + otpm_counter, _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)) + _ = release_in_background(_refund_once_reserved(usage, reserving)) 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: +async def _refund_once_reserved(usage: RouterUsage, reserving: asyncio.Task[Mapping[str, object]]) -> 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) + await _async_increment_all(usage, _refunds(_extract_reservation(reservation))) async def _async_reserve( - dual_cache: DualCache, - itpm_key: str, - otpm_key: str, + usage: RouterUsage, + itpm_counter: UsageCounter, + otpm_counter: UsageCounter, itpm_reserved: int, otpm_reserved: int, itpm_limit: int | None, otpm_limit: int | None, - parent_otel_span: Span | None, -) -> Mapping[str, int | str | None]: +) -> Mapping[str, object]: if itpm_limit is not None: await _increment_with_rollback( - dual_cache, - itpm_key, + usage, + itpm_counter, itpm_reserved, itpm_limit, - parent_otel_span=parent_otel_span, limit_label="ITPM", ) if otpm_limit is not None: try: await _increment_with_rollback( - dual_cache, - otpm_key, + usage, + otpm_counter, otpm_reserved, otpm_limit, - parent_otel_span=parent_otel_span, limit_label="OTPM", ) except Exception: @@ -753,88 +730,64 @@ async def _async_reserve( # release the ITPM reservation already made, or it stays inflated # until the TTL expires. if itpm_reserved > 0: - await dual_cache.async_increment_cache( - key=itpm_key, - value=-itpm_reserved, - ttl=RoutingArgsTTL, - parent_otel_span=parent_otel_span, - ) + await usage.increment(itpm_counter, -itpm_reserved, RoutingArgsTTL) raise 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, + itpm_counter if itpm_limit is not None else None, + otpm_counter if otpm_limit is not None else None, ) -async def _async_redis_usage(redis_cache: RedisCache | None, key: str, parent_otel_span: Span | None) -> int | None: - if redis_cache is None: - return None - with contextlib.suppress(RedisCircuitBreakerOpenError): - return await redis_cache.async_get_cache(key=key, parent_otel_span=parent_otel_span) - return None - - -@with_service_target(ROUTER_USAGE_TARGET) -async def _async_shared_usage(dual_cache: DualCache, key: str, parent_otel_span: Span | None) -> int | None: - shared: Final = await _async_redis_usage(dual_cache.redis_cache, key, parent_otel_span) - if shared is not None: - return shared - return await dual_cache.async_get_cache(key=key, local_only=True) - - -async def _raise_if_at_limit( - dual_cache: DualCache, key: str, limit: int | None, limit_label: str, parent_otel_span: Span | None -) -> None: +async def _raise_if_at_limit(usage: RouterUsage, counter: UsageCounter, limit: int | None, limit_label: str) -> None: if limit is None: return - usage: Final = await _async_shared_usage(dual_cache, key, parent_otel_span) - if usage is not None and usage >= limit: - raise _rate_limit_error(limit_label, limit, usage) + count: Final = await usage.shared_or_local(counter) + if count is not None and count >= limit: + raise _rate_limit_error(limit_label, limit, count) @with_service_target(ROUTER_USAGE_TARGET) async def async_admit_io_token_turn( - dual_cache: DualCache, + usage: RouterUsage, deployment: Mapping[str, object], now: datetime, - parent_otel_span: Span | None = None, ) -> Mapping[str, object]: configured: Final = dict(deployment) itpm_limit, otpm_limit = get_deployment_io_token_limits(configured) if itpm_limit is None and otpm_limit is None: return MappingProxyType({}) - cache_keys: Final = _get_cache_keys(configured, now.strftime("%H-%M")) - if cache_keys is None: + counters: Final = _get_counters(configured, now.strftime("%H-%M")) + if counters is None: return MappingProxyType({}) - itpm_key, otpm_key = cache_keys - await _raise_if_at_limit(dual_cache, itpm_key, itpm_limit, "ITPM", parent_otel_span) - await _raise_if_at_limit(dual_cache, otpm_key, otpm_limit, "OTPM", parent_otel_span) + itpm_counter, otpm_counter = counters + await _raise_if_at_limit(usage, itpm_counter, itpm_limit, "ITPM") + await _raise_if_at_limit(usage, otpm_counter, otpm_limit, "OTPM") return _reservation( - 0, 0, itpm_key if itpm_limit is not None else None, otpm_key if otpm_limit is not None else None + 0, 0, itpm_counter if itpm_limit is not None else None, otpm_counter if otpm_limit is not None else None ) @with_service_target(ROUTER_USAGE_TARGET) def io_token_reconcile_success( - dual_cache: DualCache, + usage: RouterUsage, kwargs: Mapping[str, object] | None, response_obj: object, ) -> None: request_kwargs: Final[Mapping[str, object] | None] = kwargs response: Final[object] = response_obj 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: + itpm_reserved, otpm_reserved, itpm_counter, otpm_counter = reservation + if itpm_counter is None and otpm_counter is None: return billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(request_kwargs, response) if usage_resolved: - for key, delta in _reconcile_increments(reservation, billable_input, completion_tokens): - dual_cache.increment_cache(key=key, value=delta, ttl=RoutingArgsTTL) + for counter, delta in _reconcile_increments(reservation, billable_input, completion_tokens): + usage.increment_sync(counter, delta, RoutingArgsTTL) else: verbose_router_logger.debug( "[IO TOKEN LIMIT] usage missing; keeping reservation (itpm_reserved=%s, otpm_reserved=%s)", @@ -854,17 +807,15 @@ def io_token_reconcile_success( @with_service_target(ROUTER_USAGE_TARGET) async def async_io_token_reconcile_success( - dual_cache: DualCache, + usage: RouterUsage, kwargs: Mapping[str, object] | None, response_obj: object, - *, - parent_otel_span: Span | None = None, ) -> None: request_kwargs: Final[Mapping[str, object] | None] = kwargs response: Final[object] = response_obj 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: + itpm_reserved, otpm_reserved, itpm_counter, otpm_counter = reservation + if itpm_counter is None and otpm_counter is None: return billable_input, completion_tokens, usage_resolved = _resolve_reconcile_usage_tokens(request_kwargs, response) @@ -877,9 +828,8 @@ async def async_io_token_reconcile_success( ( release_in_background( _async_increment_all( - dual_cache, + usage, _reconcile_increments(reservation, billable_input, completion_tokens), - parent_otel_span, ) ), ) @@ -903,7 +853,7 @@ async def async_io_token_reconcile_success( @with_service_target(ROUTER_USAGE_TARGET) def io_token_refund_failure( - dual_cache: DualCache, + usage: RouterUsage, kwargs: Mapping[str, object] | None, *, any_attempt: bool = False, @@ -911,12 +861,12 @@ def io_token_refund_failure( request_kwargs: Final[Mapping[str, object] | None] = 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) + for counter, delta in _refunds(reservation): + usage.increment_sync(counter, delta, RoutingArgsTTL) verbose_router_logger.debug("[IO TOKEN LIMIT] refunded ITPM=%s OTPM=%s", itpm_reserved, otpm_reserved) -def refund_stale_reservation_before_retry(dual_cache: DualCache, kwargs: Mapping[str, object] | None) -> None: +def refund_stale_reservation_before_retry(usage: RouterUsage, kwargs: Mapping[str, object] | None) -> None: """ Synchronously refund and clear any reservation a previous deployment attempt stashed in ``kwargs``, before it's overwritten for the next @@ -940,48 +890,40 @@ def refund_stale_reservation_before_retry(dual_cache: DualCache, kwargs: Mapping """ if not kwargs: return - io_token_refund_failure(dual_cache, kwargs, any_attempt=True) + io_token_refund_failure(usage, kwargs, any_attempt=True) @with_service_target(ROUTER_USAGE_TARGET) async def async_io_token_refund_failure( - dual_cache: DualCache, + usage: RouterUsage, kwargs: Mapping[str, object] | None, - *, - parent_otel_span: Span | None = None, ) -> None: - refund: Final = release_io_token_reservation(dual_cache, kwargs, parent_otel_span=parent_otel_span) + refund: Final = release_io_token_reservation(usage, kwargs) if refund is not None: await released((refund,)) def release_io_token_reservation( - dual_cache: DualCache, + usage: RouterUsage, 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)) + return release_in_background(_async_increment_all(usage, refunds)) -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) +def _refunds(reservation: _Reservation) -> tuple[tuple[UsageCounter, int], ...]: + itpm_reserved, otpm_reserved, itpm_counter, otpm_counter = reservation + held: Final = ((itpm_counter, itpm_reserved), (otpm_counter, otpm_reserved)) + return tuple((counter, -reserved) for counter, reserved in held if counter 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=key, value=delta, ttl=RoutingArgsTTL, parent_otel_span=parent_otel_span - ) +async def _async_increment_all(usage: RouterUsage, increments: tuple[tuple[UsageCounter, int], ...]) -> None: + for counter, delta in increments: + await usage.increment(counter, delta, RoutingArgsTTL) def build_io_token_rate_limit_headers( diff --git a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py index 38ae7e0d951..10650def2f3 100644 --- a/litellm/router_utils/pre_call_checks/model_rate_limit_check.py +++ b/litellm/router_utils/pre_call_checks/model_rate_limit_check.py @@ -22,10 +22,8 @@ from pydantic import BaseModel, ConfigDict import litellm from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger -from litellm.caching.dual_cache import DualCache -from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.integrations.custom_logger import CustomLogger -from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET +from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, CounterIncrement, RouterUsage, UsageCounter from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( async_io_token_pre_call_check, async_io_token_reconcile_success, @@ -36,7 +34,6 @@ from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( io_token_reconcile_success, io_token_refund_failure, ) -from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.router import RouterErrors, parse_daily_reset_time from litellm.types.utils import StandardLoggingPayload from litellm.utils import get_utc_datetime @@ -97,12 +94,18 @@ class _DailyLimitDeployment(BaseModel): def reset_offset(self) -> timedelta: return self.litellm_params.reset_offset() - def cache_keys(self, now: datetime) -> tuple[str, str]: - prefix: Final = f"{self.model_info.id}:{self.litellm_params.model}" + def counters(self, now: datetime) -> tuple[UsageCounter, UsageCounter]: current_day: Final = daily_window_day(now, self.reset_offset()) - return f"{prefix}:tpd:{current_day}", f"{prefix}:rpd:{current_day}" + model_id: Final = self.model_info.id + model: Final = self.litellm_params.model + return ( + UsageCounter.for_deployment(f"{model_id}", f"{model}", "tpd", current_day), + UsageCounter.for_deployment(f"{model_id}", f"{model}", "rpd", current_day), + ) - def limit_error(self, kind: Literal["tpd", "rpd"], limit: int, usage: int, now: datetime) -> litellm.RateLimitError: + def limit_error( + self, kind: Literal["tpd", "rpd"], limit: int, usage: float, now: datetime + ) -> litellm.RateLimitError: return litellm.RateLimitError( message=f"Model rate limit exceeded. {kind.upper()} limit={limit}, current usage={usage}", llm_provider="", @@ -134,58 +137,44 @@ class ModelRateLimitingCheck(CustomLogger): this check actively enforces those limits across ALL routing strategies. """ - def __init__(self, dual_cache: DualCache, clock: Callable[[], datetime] = get_utc_datetime) -> None: - self.dual_cache = dual_cache + def __init__(self, usage: RouterUsage, clock: Callable[[], datetime] = get_utc_datetime) -> None: + self.usage = usage self.clock = clock # model_ids already warned about itpm/otpm + tpm/rpm on the same deployment, # so the warning is logged once per deployment rather than per request. self._io_token_conflict_warned_ids: set[str] = set() - def _increment(self, key: str, value: int, ttl: int) -> int: - return self.dual_cache.increment_cache(key=key, value=value, ttl=ttl) - - async def _async_increment( - self, key: str, value: int, ttl: int, parent_otel_span: Span | None = None - ) -> int | None: - count: Final = await self.dual_cache.async_increment_cache( - key=key, value=value, ttl=ttl, parent_otel_span=parent_otel_span - ) - return None if count is None else int(count) + def _increment(self, counter: UsageCounter, value: int, ttl: int) -> float | None: + return self.usage.increment_sync(counter, value, ttl) @staticmethod def _request_count_operations( - rpm_key: str | None, rpd_key: str | None, increment_value: int - ) -> tuple[RedisPipelineIncrementOperation, ...]: + rpm: UsageCounter | None, rpd: UsageCounter | None, increment_value: int + ) -> tuple[CounterIncrement, ...]: return tuple( - RedisPipelineIncrementOperation(key=key, increment_value=increment_value, ttl=ttl) - for key, ttl in ((rpm_key, RoutingArgs.ttl), (rpd_key, RoutingArgs.daily_ttl)) - if key is not None + CounterIncrement(counter=counter, amount=increment_value, ttl=ttl) + for counter, ttl in ((rpm, RoutingArgs.ttl), (rpd, RoutingArgs.daily_ttl)) + if counter is not None ) - async def _async_refund_request_counts( - self, rpm_key: str | None, rpd_key: str | None, parent_otel_span: Span | None - ) -> None: - operations: Final = self._request_count_operations(rpm_key, rpd_key, -1) + async def _async_refund_request_counts(self, rpm: UsageCounter | None, rpd: UsageCounter | None) -> None: + operations: Final = self._request_count_operations(rpm, rpd, -1) if operations: - await self.dual_cache.async_increment_cache_pipeline( - increment_list=list(operations), parent_otel_span=parent_otel_span - ) + await self.usage.increment_pipeline(operations) async def _async_increment_request_counts( - self, rpm_key: str | None, rpd_key: str | None, parent_otel_span: Span | None + self, rpm: UsageCounter | None, rpd: UsageCounter | None ) -> tuple[int | None, int | None]: - operations: Final = self._request_count_operations(rpm_key, rpd_key, 1) + operations: Final = self._request_count_operations(rpm, rpd, 1) if not operations: return None, None - counts: Final = await self.dual_cache.async_increment_cache_pipeline( - increment_list=list(operations), parent_otel_span=parent_otel_span - ) + counts: Final = await self.usage.increment_pipeline(operations) if counts is None or len(counts) != len(operations): return None, None - count_by_key: Final = {op["key"]: int(count) for op, count in zip(operations, counts)} + count_by_counter: Final = {op.counter: int(count) for op, count in zip(operations, counts)} return ( - None if rpm_key is None else count_by_key.get(rpm_key), - None if rpd_key is None else count_by_key.get(rpd_key), + None if rpm is None else count_by_counter.get(rpm), + None if rpd is None else count_by_counter.get(rpd), ) def _warn_io_token_and_tpm_rpm_coexist_once(self, deployment: dict) -> None: @@ -207,19 +196,12 @@ class ModelRateLimitingCheck(CustomLogger): def _refund_io_token_reservation_if_any(self) -> None: request_kwargs: Final = get_io_token_rate_limit_request_kwargs() if request_kwargs is not None: - io_token_refund_failure(self.dual_cache, request_kwargs) + io_token_refund_failure(self.usage, request_kwargs) - async def _async_refund_io_token_reservation_if_any( - self, - parent_otel_span: Span | None = None, - ) -> None: + async def _async_refund_io_token_reservation_if_any(self) -> None: request_kwargs: Final = get_io_token_rate_limit_request_kwargs() if request_kwargs is not None: - await async_io_token_refund_failure( - self.dual_cache, - request_kwargs, - parent_otel_span=parent_otel_span, - ) + await async_io_token_refund_failure(self.usage, request_kwargs) def _get_deployment_limits(self, deployment: dict) -> tuple[int | None, int | None]: """ @@ -255,48 +237,25 @@ class ModelRateLimitingCheck(CustomLogger): tpd_limit, rpd_limit = deployment.limits() if tpd_limit is None and rpd_limit is None: return - tpd_key, rpd_key = deployment.cache_keys(now) + tpd, rpd = deployment.counters(now) if tpd_limit is not None: - current_tpd: Final = self._get_current_tpm(tpd_key, tpd_limit) + current_tpd: Final = self.usage.current_sync(tpd, tpd_limit) if current_tpd is not None and current_tpd >= tpd_limit: raise deployment.limit_error("tpd", tpd_limit, current_tpd, now) if rpd_limit is not None: - current_rpd: Final = self._increment(rpd_key, 1, RoutingArgs.daily_ttl) - if current_rpd > rpd_limit: - self._increment(rpd_key, -1, RoutingArgs.daily_ttl) + current_rpd: Final = self._increment(rpd, 1, RoutingArgs.daily_ttl) + if current_rpd is not None and current_rpd > rpd_limit: + self._increment(rpd, -1, RoutingArgs.daily_ttl) raise deployment.limit_error("rpd", rpd_limit, current_rpd, now) - def _get_cache_keys(self, deployment: dict, current_minute: str) -> tuple[str, str]: - """Get the cache keys for TPM and RPM tracking.""" - model_id: Final = deployment.get("model_info", {}).get("id") - deployment_name: Final = deployment.get("litellm_params", {}).get("model") - - tpm_key: Final = f"{model_id}:{deployment_name}:tpm:{current_minute}" - rpm_key: Final = f"{model_id}:{deployment_name}:rpm:{current_minute}" - - return tpm_key, rpm_key - - @with_service_target(ROUTER_USAGE_TARGET) - def _get_current_tpm(self, tpm_key: str, tpm_limit: int) -> int | None: - local_tpm: Final = self.dual_cache.get_cache(key=tpm_key, local_only=True) - redis_cache: Final = self.dual_cache.redis_cache - if redis_cache is None or (local_tpm is not None and local_tpm >= tpm_limit): - return local_tpm - try: - return redis_cache.get_cache(key=tpm_key) - except RedisCircuitBreakerOpenError: - return local_tpm - - @with_service_target(ROUTER_USAGE_TARGET) - async def _async_get_current_tpm(self, tpm_key: str, tpm_limit: int, parent_otel_span: Span | None) -> int | None: - local_tpm: Final = await self.dual_cache.async_get_cache(key=tpm_key, local_only=True) - redis_cache: Final = self.dual_cache.redis_cache - if redis_cache is None or (local_tpm is not None and local_tpm >= tpm_limit): - return local_tpm - try: - return await redis_cache.async_get_cache(key=tpm_key, parent_otel_span=parent_otel_span) - except RedisCircuitBreakerOpenError: - return local_tpm + def _counters(self, deployment: dict, current_minute: str) -> tuple[UsageCounter, UsageCounter]: + """The counters TPM and RPM are tracked in.""" + model_id: Final[object] = deployment.get("model_info", {}).get("id") + deployment_name: Final[object] = deployment.get("litellm_params", {}).get("model") + return ( + UsageCounter.for_deployment(f"{model_id}", f"{deployment_name}", "tpm", current_minute), + UsageCounter.for_deployment(f"{model_id}", f"{deployment_name}", "rpm", current_minute), + ) @with_service_target(ROUTER_USAGE_TARGET) def pre_call_check(self, deployment: dict) -> dict | None: @@ -310,7 +269,7 @@ class ModelRateLimitingCheck(CustomLogger): if deployment_has_io_token_limits(deployment): self._warn_io_token_and_tpm_rpm_coexist_once(deployment) io_token_pre_call_check( - self.dual_cache, + self.usage, deployment, ) io_reservation_made = True @@ -320,7 +279,7 @@ class ModelRateLimitingCheck(CustomLogger): tpm_limit, rpm_limit = self._get_deployment_limits(deployment) current_minute: Final = dt.strftime("%H-%M") - tpm_key, rpm_key = self._get_cache_keys(deployment, current_minute) + tpm, rpm = self._counters(deployment, current_minute) model_id: Final = deployment.get("model_info", {}).get("id") model_name: Final = deployment.get("litellm_params", {}).get("model") @@ -328,7 +287,7 @@ class ModelRateLimitingCheck(CustomLogger): # Check TPM limit if tpm_limit is not None: - current_tpm: Final = self._get_current_tpm(tpm_key, tpm_limit) + current_tpm: Final = self.usage.current_sync(tpm, tpm_limit) if current_tpm is not None and current_tpm >= tpm_limit: raise litellm.RateLimitError( message=f"Model rate limit exceeded. TPM limit={tpm_limit}, current usage={current_tpm}", @@ -347,9 +306,9 @@ class ModelRateLimitingCheck(CustomLogger): # Check RPM limit (atomic increment-first to avoid race conditions) if rpm_limit is not None: - current_rpm: Final = self._increment(rpm_key, 1, RoutingArgs.ttl) + current_rpm: Final = self._increment(rpm, 1, RoutingArgs.ttl) if current_rpm is not None and current_rpm > rpm_limit: - self._increment(rpm_key, -1, RoutingArgs.ttl) + self._increment(rpm, -1, RoutingArgs.ttl) raise litellm.RateLimitError( message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", llm_provider="", @@ -369,7 +328,7 @@ class ModelRateLimitingCheck(CustomLogger): self._check_daily_limits(daily, dt) except litellm.RateLimitError: if rpm_limit is not None: - self._increment(rpm_key, -1, RoutingArgs.ttl) + self._increment(rpm, -1, RoutingArgs.ttl) raise return deployment @@ -393,11 +352,7 @@ class ModelRateLimitingCheck(CustomLogger): io_reservation_made = False if deployment_has_io_token_limits(deployment): self._warn_io_token_and_tpm_rpm_coexist_once(deployment) - await async_io_token_pre_call_check( - self.dual_cache, - deployment, - parent_otel_span=parent_otel_span, - ) + await async_io_token_pre_call_check(self.usage, deployment) io_reservation_made = True dt: Final = self.clock() @@ -406,8 +361,8 @@ class ModelRateLimitingCheck(CustomLogger): tpm_limit, rpm_limit = self._get_deployment_limits(deployment) current_minute: Final = dt.strftime("%H-%M") - tpm_key, rpm_key = self._get_cache_keys(deployment, current_minute) - tpd_key, rpd_key = daily.cache_keys(dt) + tpm, rpm = self._counters(deployment, current_minute) + tpd, rpd = daily.counters(dt) model_id: Final = deployment.get("model_info", {}).get("id") model_name: Final = deployment.get("litellm_params", {}).get("model") @@ -415,7 +370,7 @@ class ModelRateLimitingCheck(CustomLogger): # Check TPM limit if tpm_limit is not None: - current_tpm: Final = await self._async_get_current_tpm(tpm_key, tpm_limit, parent_otel_span) + current_tpm: Final = await self.usage.current(tpm, tpm_limit) if current_tpm is not None and current_tpm >= tpm_limit: raise litellm.RateLimitError( message=f"Model rate limit exceeded. TPM limit={tpm_limit}, current usage={current_tpm}", @@ -434,17 +389,15 @@ class ModelRateLimitingCheck(CustomLogger): ) if tpd_limit is not None: - current_tpd: Final = await self._async_get_current_tpm(tpd_key, tpd_limit, parent_otel_span) + current_tpd: Final = await self.usage.current(tpd, tpd_limit) if current_tpd is not None and current_tpd >= tpd_limit: raise daily.limit_error("tpd", tpd_limit, current_tpd, dt) - charged_rpm_key: Final = rpm_key if rpm_limit is not None else None - charged_rpd_key: Final = rpd_key if rpd_limit is not None else None - current_rpm, current_rpd = await self._async_increment_request_counts( - charged_rpm_key, charged_rpd_key, parent_otel_span - ) + charged_rpm: Final = rpm if rpm_limit is not None else None + charged_rpd: Final = rpd if rpd_limit is not None else None + current_rpm, current_rpd = await self._async_increment_request_counts(charged_rpm, charged_rpd) if rpm_limit is not None and current_rpm is not None and current_rpm > rpm_limit: - await self._async_refund_request_counts(charged_rpm_key, charged_rpd_key, parent_otel_span) + await self._async_refund_request_counts(charged_rpm, charged_rpd) raise litellm.RateLimitError( message=f"Model rate limit exceeded. RPM limit={rpm_limit}, current usage={current_rpm}", llm_provider="", @@ -462,14 +415,14 @@ class ModelRateLimitingCheck(CustomLogger): ) if rpd_limit is not None and current_rpd is not None and current_rpd > rpd_limit: - await self._async_refund_request_counts(charged_rpm_key, charged_rpd_key, parent_otel_span) + await self._async_refund_request_counts(charged_rpm, charged_rpd) raise daily.limit_error("rpd", rpd_limit, current_rpd, dt) return deployment except litellm.RateLimitError: if io_reservation_made: - await self._async_refund_io_token_reservation_if_any(parent_otel_span=parent_otel_span) + await self._async_refund_io_token_reservation_if_any() raise except Exception as e: verbose_router_logger.debug("Error in ModelRateLimitingCheck.async_pre_call_check: %s", e) @@ -478,28 +431,19 @@ class ModelRateLimitingCheck(CustomLogger): @with_service_target(ROUTER_USAGE_TARGET) async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): - from litellm.litellm_core_utils.core_helpers import ( - get_parent_otel_span_from_kwargs, - ) - try: standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") - await async_io_token_reconcile_success( - self.dual_cache, - kwargs, - response_obj, - parent_otel_span=get_parent_otel_span_from_kwargs(kwargs), - ) + await async_io_token_reconcile_success(self.usage, kwargs, response_obj) if standard_logging_object is None: return - model_id: Final = standard_logging_object.get("model_id") + model_id: Final[object] = standard_logging_object.get("model_id") if model_id is None: return total_tokens: Final = standard_logging_object.get("total_tokens", 0) - model: Final = standard_logging_object.get("hidden_params", {}).get("litellm_model_name") + model: Final[object] = standard_logging_object.get("hidden_params", {}).get("litellm_model_name") verbose_router_logger.debug( "[TPM TRACKING] model_id=%s, total_tokens=%s, model=%s", model_id, total_tokens, model @@ -509,32 +453,25 @@ class ModelRateLimitingCheck(CustomLogger): return dt: Final = self.clock() - current_minute: Final = dt.strftime("%H-%M") - tpm_key: Final = f"{model_id}:{model}:tpm:{current_minute}" + tpm: Final = UsageCounter.for_deployment(f"{model_id}", f"{model}", "tpm", dt.strftime("%H-%M")) reset_offset: Final = _reset_offset_from_kwargs(kwargs) # pyright: ignore[reportUnknownArgumentType] # CustomLogger passes an untyped kwargs dict - tpd_key: Final = f"{model_id}:{model}:tpd:{daily_window_day(dt, reset_offset)}" + tpd: Final = UsageCounter.for_deployment( + f"{model_id}", f"{model}", "tpd", daily_window_day(dt, reset_offset) + ) - verbose_router_logger.debug("[TPM TRACKING] Incrementing %s by %s", tpm_key, total_tokens) + verbose_router_logger.debug("[TPM TRACKING] Incrementing %s by %s", tpm, total_tokens) tokens: Final = int(total_tokens) - await self._async_increment(tpm_key, tokens, RoutingArgs.ttl) - await self._async_increment(tpd_key, tokens, RoutingArgs.daily_ttl) + await self.usage.increment(tpm, tokens, RoutingArgs.ttl) + await self.usage.increment(tpd, tokens, RoutingArgs.daily_ttl) except Exception as e: verbose_router_logger.debug("Error in ModelRateLimitingCheck.async_log_success_event: %s", e) async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): - from litellm.litellm_core_utils.core_helpers import ( - get_parent_otel_span_from_kwargs, - ) - # Never fail the primary logging pipeline over an io-token refund error. with contextlib.suppress(Exception): - await async_io_token_refund_failure( - self.dual_cache, - kwargs, - parent_otel_span=get_parent_otel_span_from_kwargs(kwargs), - ) + await async_io_token_refund_failure(self.usage, kwargs) @with_service_target(ROUTER_USAGE_TARGET) def log_success_event(self, kwargs, response_obj, start_time, end_time): @@ -544,30 +481,31 @@ class ModelRateLimitingCheck(CustomLogger): """ try: standard_logging_object: Final[StandardLoggingPayload | None] = kwargs.get("standard_logging_object") - io_token_reconcile_success(self.dual_cache, kwargs, response_obj) + io_token_reconcile_success(self.usage, kwargs, response_obj) if standard_logging_object is None: return - model_id: Final = standard_logging_object.get("model_id") + model_id: Final[object] = standard_logging_object.get("model_id") if model_id is None: return total_tokens: Final = standard_logging_object.get("total_tokens", 0) - model: Final = standard_logging_object.get("hidden_params", {}).get("litellm_model_name") + model: Final[object] = standard_logging_object.get("hidden_params", {}).get("litellm_model_name") if not model or not total_tokens: return dt: Final = self.clock() - current_minute: Final = dt.strftime("%H-%M") - tpm_key: Final = f"{model_id}:{model}:tpm:{current_minute}" + tpm: Final = UsageCounter.for_deployment(f"{model_id}", f"{model}", "tpm", dt.strftime("%H-%M")) reset_offset: Final = _reset_offset_from_kwargs(kwargs) # pyright: ignore[reportUnknownArgumentType] # CustomLogger passes an untyped kwargs dict - tpd_key: Final = f"{model_id}:{model}:tpd:{daily_window_day(dt, reset_offset)}" + tpd: Final = UsageCounter.for_deployment( + f"{model_id}", f"{model}", "tpd", daily_window_day(dt, reset_offset) + ) tokens: Final = int(total_tokens) - self._increment(tpm_key, tokens, RoutingArgs.ttl) - self._increment(tpd_key, tokens, RoutingArgs.daily_ttl) + self._increment(tpm, tokens, RoutingArgs.ttl) + self._increment(tpd, tokens, RoutingArgs.daily_ttl) except Exception as e: verbose_router_logger.debug("Error in ModelRateLimitingCheck.log_success_event: %s", e) @@ -575,6 +513,6 @@ class ModelRateLimitingCheck(CustomLogger): def log_failure_event(self, kwargs, response_obj, start_time, end_time): with contextlib.suppress(Exception): io_token_refund_failure( - self.dual_cache, + self.usage, kwargs, ) diff --git a/litellm/router_utils/routing_prefetch.py b/litellm/router_utils/routing_prefetch.py index 57d76220311..6d41eb867b5 100644 --- a/litellm/router_utils/routing_prefetch.py +++ b/litellm/router_utils/routing_prefetch.py @@ -1,7 +1,6 @@ from __future__ import annotations -import asyncio -from collections.abc import Awaitable, Mapping, Sequence +from collections.abc import Awaitable, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Final @@ -10,70 +9,26 @@ if TYPE_CHECKING: from litellm._internal_context import service_target from litellm._logging import verbose_router_logger -from litellm.caching.dual_cache import DualCache -from litellm.caching.redis_batch import Batch, BatchResult, active_request_redis_batches +from litellm.caching.redis_batch import Batch, active_request_redis_batches from litellm.redis.stores.router_cooldowns import ROUTER_COOLDOWNS_TARGET, ActiveCooldowns, RouterCooldowns -from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET +from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, RouterUsage, UsageCounter, UsageRead -_PREFETCH_SLOT: Final = "routing_read" +_USAGE_SLOT: Final = "routing_read" _COOLDOWNS_SLOT: Final = "routing_cooldowns" @dataclass(frozen=True, slots=True) -class TierRead: - cache: DualCache - keys: tuple[str, ...] +class ArmedUsage: + counters: tuple[UsageCounter, ...] + read: UsageRead -def selector_read(cache: DualCache, keys: Sequence[str]) -> TierRead: - return TierRead(cache, tuple(keys)) - - -@dataclass(frozen=True, slots=True) -class ArmedRead: - read: TierRead - fetched: frozenset[str] - result: BatchResult[Mapping[str, object]] - previous: Mapping[str, float | None] - - def release(self) -> None: - self.read.cache._rollback_redis_batch_key_reservations( # pyright: ignore[reportPrivateUsage] # the selector owns this tier - dict(self.previous) - ) - - async def settle(self, future: asyncio.Future[Mapping[str, object]]) -> None: - if future.cancelled() or future.exception() is not None: - self.release() - return - try: - pending: Final = await self.read.cache._prepare_batch_get( # pyright: ignore[reportPrivateUsage] # the selector owns this tier - list(self.read.keys), local_only=True - ) - values: Final = future.result() - missed: Final = { - key: values[key] - for key, local in zip(self.read.keys, pending.result) - if local is None and values.get(key) is not None - } - await self.read.cache._apply_batch_get(pending, missed) # pyright: ignore[reportPrivateUsage] # the selector owns this tier - except Exception: - self.release() - raise - - -def arm(read: TierRead) -> None: +def arm(store: RouterUsage, counters: Sequence[UsageCounter]) -> None: + """The usage read a usage-based-routing-v2 pick takes later, declared on the request's native pre-call batch.""" request: Final = active_request_redis_batches() - redis: Final = read.cache.redis_cache - if request is None or redis is None or _PREFETCH_SLOT in request.prefetched: + if request is None or not store.has_redis or _USAGE_SLOT in request.prefetched: return - due, previous = read.cache.reserve_redis_batch_reads(read.keys) - if not due: - return - with service_target(ROUTER_USAGE_TARGET): - result: Final = request.batch(redis).mget(tuple(due)) - armed_read: Final = ArmedRead(read=read, fetched=frozenset(due), result=result, previous=previous) - result.on_settled(armed_read.settle) - request.prefetched[_PREFETCH_SLOT] = armed_read + request.prefetched[_USAGE_SLOT] = ArmedUsage(tuple(counters), store.usage(counters, batch=Batch.PRE_CALL)) @dataclass(frozen=True, slots=True) @@ -92,7 +47,7 @@ def arm_cooldowns(store: RouterCooldowns, model_ids: Sequence[str]) -> None: def armed() -> bool: request: Final = active_request_redis_batches() - return request is not None and (_PREFETCH_SLOT in request.prefetched or _COOLDOWNS_SLOT in request.prefetched) + return request is not None and (_USAGE_SLOT in request.prefetched or _COOLDOWNS_SLOT in request.prefetched) async def _armed_cooldowns(model_ids: Sequence[str]) -> ActiveCooldowns | None: @@ -122,43 +77,29 @@ async def read_cooldowns( return await store.active(model_ids, batch=Batch.NONE) -def _take(read: TierRead) -> ArmedRead | None: +async def _armed_usage(counters: Sequence[UsageCounter]) -> Sequence[float | None] | None: request: Final = active_request_redis_batches() - if request is None: + taken: Final = None if request is None else request.prefetched.pop(_USAGE_SLOT, None) + if not isinstance(taken, ArmedUsage): return None - taken: Final = request.prefetched.pop(_PREFETCH_SLOT, None) - if isinstance(taken, ArmedRead) and set(taken.read.keys).issuperset(read.keys): - return taken - if isinstance(taken, ArmedRead): - taken.release() - return None - - -async def _read_prefetched(read: TierRead) -> Sequence[object | None] | None: - prefetch: Final = _take(read) - if prefetch is None: + if not frozenset(taken.counters).issuperset(counters): + taken.read.release() return None try: - values: Final = await prefetch.result - except Exception as e: # noqa: BLE001 # the read falls back to the tier's own Redis read + values: Final = await taken.read.result + except Exception as e: # noqa: BLE001 # the read falls back to the store's own Redis read verbose_router_logger.debug("routing prefetch failed, reading again: %s", e) return None - pending: Final = await read.cache._prepare_batch_get( # pyright: ignore[reportPrivateUsage] # the selector owns this tier - list(read.keys), local_only=True - ) - if any(key not in prefetch.fetched for key, value in zip(read.keys, pending.result) if value is None): - return None - return await read.cache._apply_batch_get( # pyright: ignore[reportPrivateUsage] # the selector owns this tier - pending, {key: values.get(key) for key, value in zip(read.keys, pending.result) if value is None} - ) + read: Final = dict(zip(taken.counters, values)) + return [read[counter] for counter in counters] -async def read(read: TierRead, parent_otel_span: Span | None) -> Sequence[object | None] | None: +async def read( + store: RouterUsage, counters: Sequence[UsageCounter], parent_otel_span: Span | None +) -> Sequence[float | None]: + """The armed usage read when it covers ``counters``, taken once; otherwise the store reads on its own.""" with service_target(ROUTER_USAGE_TARGET): - prefetched: Final = await _read_prefetched(read) - if prefetched is not None: - return prefetched - results: Final = await DualCache.async_batch_get_cache_shared( - ((read.cache, list(read.keys)),), parent_otel_span=parent_otel_span - ) - return results[0] + prefetched: Final = await _armed_usage(counters) + if prefetched is not None: + return prefetched + return await store.usage(counters, batch=Batch.NONE).result diff --git a/litellm/router_utils/routing_read_batch.py b/litellm/router_utils/routing_read_batch.py index 8f54a96e305..2d5aed65883 100644 --- a/litellm/router_utils/routing_read_batch.py +++ b/litellm/router_utils/routing_read_batch.py @@ -5,10 +5,10 @@ from contextvars import ContextVar from typing import TYPE_CHECKING, Final from litellm._logging import verbose_router_logger +from litellm.redis.stores.router_usage import UsageCounter from litellm.router_deprecated import renamed_router_param from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2, PrefetchedUsage from litellm.router_utils import routing_prefetch -from litellm.router_utils.routing_prefetch import TierRead, selector_read if TYPE_CHECKING: from opentelemetry.trace import Span @@ -18,13 +18,9 @@ if TYPE_CHECKING: __all__ = ("RoutingPrefetch", "RoutingReadBatch") -def _selector_keys(selector: LowestTPMLoggingHandler_v2, deployments: Sequence[object]) -> list[str]: - tpm_keys, rpm_keys = selector.usage_counter_keys(list(deployments)) - return tpm_keys + rpm_keys - - -def _selector_read(selector: LowestTPMLoggingHandler_v2 | None, deployments: Sequence[object]) -> TierRead | None: - return None if selector is None else selector_read(selector.router_cache, _selector_keys(selector, deployments)) +def _selector_counters(selector: LowestTPMLoggingHandler_v2, deployments: Sequence[object]) -> tuple[UsageCounter, ...]: + tpm_counters, rpm_counters = selector.usage_counters(deployments) + return tpm_counters + rpm_counters class RoutingPrefetch: @@ -34,9 +30,8 @@ class RoutingPrefetch: ctx: "RoutingContext", usage_selector: LowestTPMLoggingHandler_v2 | None, deployments: Sequence[object] ) -> None: RoutingPrefetch.arm_cooldowns(ctx) - read: Final = _selector_read(usage_selector, deployments) - if read is not None: - routing_prefetch.arm(read) + if usage_selector is not None: + routing_prefetch.arm(usage_selector.usage, _selector_counters(usage_selector, deployments)) @staticmethod def arm_cooldowns(ctx: "RoutingContext") -> None: @@ -92,11 +87,9 @@ class RoutingReadBatch: async def read_usage(self, deployments: Sequence[object], parent_otel_span: "Span | None") -> None: """The selector's tpm/rpm counters for ``deployments``, kept as ``prefetched_usage``.""" - usage_read: Final = _selector_read(self.usage_selector, deployments) - if usage_read is None: + selector: Final = self.usage_selector + if selector is None: return - values: Final = await routing_prefetch.read(usage_read, parent_otel_span) - self.prefetched_usage = PrefetchedUsage( - keys=frozenset(usage_read.keys), - values=None if values is None else dict(zip(usage_read.keys, values)), - ) + counters: Final = _selector_counters(selector, deployments) + values: Final = await routing_prefetch.read(selector.usage, counters, parent_otel_span) + self.prefetched_usage = PrefetchedUsage(counters=frozenset(counters), values=dict(zip(counters, values))) diff --git a/litellm/rust/crates/host-python/src/gil.rs b/litellm/rust/crates/host-python/src/gil.rs index 5f54fae81ed..886ed848399 100644 --- a/litellm/rust/crates/host-python/src/gil.rs +++ b/litellm/rust/crates/host-python/src/gil.rs @@ -22,6 +22,18 @@ impl PythonContext { ))) } + /// A context with no request-local values, for work no Python caller is waiting on. + pub fn empty(py: Python<'_>) -> PyResult { + Ok(Self(Arc::new( + py.import("contextvars")?.call_method0("Context")?.unbind(), + ))) + } + + /// A `contextvars.Context` the Python host built for work that outlives its caller. + pub fn built(context: Bound<'_, PyAny>) -> Self { + Self(Arc::new(context.unbind())) + } + /// Runs `f` inside a fresh copy of the captured context. The copy is what lets two blocking /// calls run concurrently: a `contextvars.Context` cannot be entered twice at once. pub fn enter(&self, py: Python<'_>, f: F) -> PyResult diff --git a/litellm/rust/crates/python-bridge/src/lib.rs b/litellm/rust/crates/python-bridge/src/lib.rs index 4a8f6917396..12795146ca9 100644 --- a/litellm/rust/crates/python-bridge/src/lib.rs +++ b/litellm/rust/crates/python-bridge/src/lib.rs @@ -62,7 +62,7 @@ mod _native { #[pymodule_export] use crate::redis::store::request::{NativeRequestRedis, NativeSend}; #[pymodule_export] - use crate::redis::store::router_usage::NativeRouterUsage; + use crate::redis::store::router_usage::{NativeRouterUsage, NativeUsageRead}; #[pymodule_export] use crate::redis::store::sensitive_route_pins::NativeSensitiveRoutePins; #[pymodule_export] @@ -166,6 +166,7 @@ mod tests { "NativePending", "NativeRouterCooldowns", "NativeRouterUsage", + "NativeUsageRead", "NativeSensitiveRoutePins", "NativeRedisStoreError", "NativeSpendCounters", diff --git a/litellm/rust/crates/python-bridge/src/redis/events.rs b/litellm/rust/crates/python-bridge/src/redis/events.rs index a281d78d5a6..b5d51442561 100644 --- a/litellm/rust/crates/python-bridge/src/redis/events.rs +++ b/litellm/rust/crates/python-bridge/src/redis/events.rs @@ -24,16 +24,52 @@ impl EventLoop { } } +/// The Python context a Redis event is reported in. Each constructor names whose context it is, +/// so no event is reported in the context of whoever happened to build the thing reporting it. +#[derive(Clone)] +pub(crate) struct ReportContext(PythonContext); + +impl ReportContext { + /// The caller of a store op, on the stack now: its `with_service_target` and the span its + /// request parents service events under. + pub(crate) fn of_caller(py: Python<'_>) -> PyResult { + Ok(Self(PythonContext::capture(py)?)) + } + + /// The request whose post-call batch its deadline sends, with no caller awaiting the flush. + fn of_deadline_flush(py: Python<'_>) -> PyResult { + let context = py + .import(ServiceLoggingEvents::MODULE)? + .call_method0("deadline_flush_context")?; + Ok(Self(PythonContext::built(context))) + } + + /// Nobody's, for breaker transitions and calls no store op scoped. + fn background(py: Python<'_>) -> PyResult { + Ok(Self(PythonContext::empty(py)?)) + } +} + /// The Python caller of one store operation. Its context carries `with_service_target` and the /// ambient OTel span, and its `call_type` is the label `RedisCache` reports, so dashboards keep /// their series. pub(crate) struct RedisCall { - pub(crate) context: PythonContext, + pub(crate) context: ReportContext, pub(crate) event_loop: Option, pub(crate) call_type: &'static str, } impl RedisCall { + /// The deadline's flush of a request's post-call batch. Built when a post-call declaration + /// arms the deadline, so it is the declaring caller's request the flush is reported under. + pub(crate) fn deadline_flush(py: Python<'_>) -> PyResult { + Ok(Self { + context: ReportContext::of_deadline_flush(py)?, + event_loop: EventLoop::running(py), + call_type: "post_call_redis_batch", + }) + } + pub(crate) fn tag(self) -> CallTag { Arc::new(self) } @@ -47,7 +83,7 @@ impl RedisCall { pub(crate) struct ServiceLoggingEvents { record_call: Arc>, record_breaker: Arc>, - context: PythonContext, + context: ReportContext, /// The loop the client was built on, for calls whose caller had none running. event_loop: Option, runtime: Handle, @@ -79,15 +115,15 @@ impl ServiceLoggingEvents { Ok(Self { record_call: Arc::new(module.getattr("record_redis_call")?.unbind()), record_breaker: Arc::new(module.getattr("record_redis_breaker")?.unbind()), - context: PythonContext::capture(py)?, + context: ReportContext::background(py)?, event_loop, runtime, }) } - fn report(&self, context: PythonContext, record: Arc>, args: ReportArgs) { + fn report(&self, context: ReportContext, record: Arc>, args: ReportArgs) { self.runtime.spawn(async move { - let reported = attach_blocking(context, move |py| match args { + let reported = attach_blocking(context.0, move |py| match args { ReportArgs::Call(call) => record .bind(py) .call1(( @@ -205,7 +241,6 @@ impl RedisEvents for ServiceLoggingEvents { mod tests { use std::{sync::Arc, time::Duration}; - use litellm_host_python::PythonContext; use litellm_redis_client::{CallEvent, FailureKind, Operation, Outcome, RedisEvents}; use pyo3::{ prelude::*, @@ -214,7 +249,7 @@ mod tests { use rstest::rstest; use tokio::runtime::Handle; - use super::{EventLoop, RedisCall, ServiceLoggingEvents}; + use super::{EventLoop, RedisCall, ReportContext, ServiceLoggingEvents}; const RECORDER: &str = " import contextvars, time @@ -272,7 +307,7 @@ def record_redis_breaker(state): .call_method1("set", ("spend_counter",)) .unwrap(); RedisCall { - context: PythonContext::capture(py).unwrap(), + context: ReportContext::of_caller(py).unwrap(), event_loop: caller_loop.map(|caller_loop| stand_in_loop(py, caller_loop)), call_type: "async_increment", } @@ -308,6 +343,41 @@ def record_redis_breaker(state): ); } + #[tokio::test(flavor = "multi_thread")] + async fn a_call_nobody_scoped_is_reported_outside_the_context_the_client_was_built_in() { + Python::initialize(); + let (events, namespace) = Python::attach(|py| { + let name = c"redis_events_recorder_built_inside_a_request"; + let module = + PyModule::from_code(py, &std::ffi::CString::new(RECORDER).unwrap(), name, name) + .unwrap(); + module + .getattr("target") + .unwrap() + .call_method1("set", ("the request that connected",)) + .unwrap(); + let events = ServiceLoggingEvents::from_module( + py, + &module, + Handle::current(), + Some(stand_in_loop(py, "client loop")), + ) + .unwrap(); + (events, module.unbind()) + }); + + events.call(&CallEvent { + operation: Operation::Command("INCRBYFLOAT"), + duration: Duration::from_millis(250), + outcome: Outcome::Succeeded, + }); + + assert_eq!( + seen(&namespace, 1).await, + "[('client loop', 'INCRBYFLOAT', 0.25, True, None, None)]" + ); + } + #[rstest] #[case::opened(litellm_circuit_breaker::State::Open { since: std::time::Instant::now() }, "open")] #[case::probing(litellm_circuit_breaker::State::HalfOpen, "half_open")] diff --git a/litellm/rust/crates/python-bridge/src/redis/store/mod.rs b/litellm/rust/crates/python-bridge/src/redis/store/mod.rs index 35eba7b924f..2700ae09593 100644 --- a/litellm/rust/crates/python-bridge/src/redis/store/mod.rs +++ b/litellm/rust/crates/python-bridge/src/redis/store/mod.rs @@ -17,7 +17,6 @@ use std::{ time::Duration, }; -use litellm_host_python::PythonContext; use litellm_python_compat::number::{SecondsError, seconds}; use litellm_redis_store::{Error, FromValue, IntoValue, Number, Pending, Remote}; use pyo3::{ @@ -30,14 +29,14 @@ use pyo3::{ use crate::execution::{run_async_value, run_sync_value}; use crate::redis::{ client::NativeRedisClient, - events::{EventLoop, RedisCall}, + events::{EventLoop, RedisCall, ReportContext}, }; /// The labels `RedisCache` reports for this operation, in the caller's context so its /// `service_target` reaches the service event, dispatched on the loop the caller runs on. fn call(py: Python<'_>, call_type: &'static str) -> PyResult { Ok(RedisCall { - context: PythonContext::capture(py)?, + context: ReportContext::of_caller(py)?, event_loop: EventLoop::running(py), call_type, }) diff --git a/litellm/rust/crates/python-bridge/src/redis/store/request.rs b/litellm/rust/crates/python-bridge/src/redis/store/request.rs index 3436b26a16c..57689ae83d5 100644 --- a/litellm/rust/crates/python-bridge/src/redis/store/request.rs +++ b/litellm/rust/crates/python-bridge/src/redis/store/request.rs @@ -1,8 +1,21 @@ +use std::sync::Arc; + +use litellm_redis_client::CallTag; use litellm_redis_store::{RequestRedis, Send}; use pyo3::prelude::*; use super::{call, duration}; -use crate::execution::run_async_value; +use crate::{execution::run_async_value, redis::events::RedisCall}; + +/// The tag of a deadline flush, asked on the thread whose post-call declaration arms the deadline. +fn deadline_flush_tag() -> Option { + Python::attach(RedisCall::deadline_flush) + .inspect_err(|error| { + tracing::warn!(%error, "a post-call deadline flush will be reported without its request"); + }) + .ok() + .map(RedisCall::tag) +} /// One request's native Redis batches. Python builds one per request scope and passes it to each /// request-aware store call. @@ -15,9 +28,8 @@ impl NativeRequestRedis { fn new(py: Python<'_>, post_call_deadline_seconds: f64) -> PyResult { let runtime = litellm_host_python::runtime_handle()?; let deadline = duration(post_call_deadline_seconds)?; - let deadline_tag = call(py, "post_call_redis_batch")?.tag(); Ok(Self(crate::logger::capture(py).scope(|| { - RequestRedis::with_deadline_tag(runtime, deadline, Some(deadline_tag)) + RequestRedis::with_deadline_tag(runtime, deadline, Arc::new(deadline_flush_tag)) }))) } diff --git a/litellm/rust/crates/python-bridge/src/redis/store/router_usage.rs b/litellm/rust/crates/python-bridge/src/redis/store/router_usage.rs index 48565028de3..dbe7ecfedb2 100644 --- a/litellm/rust/crates/python-bridge/src/redis/store/router_usage.rs +++ b/litellm/rust/crates/python-bridge/src/redis/store/router_usage.rs @@ -1,28 +1,40 @@ -use std::{collections::HashMap, sync::Arc, time::Duration}; +use std::{ + collections::HashMap, + sync::{Arc, Mutex, PoisonError}, + time::Duration, +}; -use litellm_redis_store::router_usage::{Counter, Kind, RouterUsage}; +use litellm_redis_store::router_usage::{ + Counter, Increment, Kind, RouterUsage, Scope, UsageReservation, +}; use pyo3::prelude::*; use super::{ - PyNumber, Spelled, kind_error, policy::router_tier_policies, redis_async, redis_sync, remote, + NativePending, PyNumber, Spelled, kind_error, pending, + policy::router_tier_policies, + redis_async, redis_sync, remote, + request::{NativeRequestRedis, NativeSend}, + ttl, }; use crate::redis::client::NativeRedisClient; /// `litellm.redis.stores.router_usage.UsageCounter`. #[derive(FromPyObject)] struct PyCounter { + scope: Spelled, deployment_id: String, model: Option, - minute: String, + window: String, kind: Spelled, } impl From for Counter { fn from(counter: PyCounter) -> Self { Self { + scope: counter.scope.0, deployment_id: counter.deployment_id, model: counter.model, - minute: counter.minute, + window: counter.window, kind: counter.kind.0, } } @@ -32,11 +44,66 @@ fn counters(counters: Vec) -> Vec { counters.into_iter().map(Counter::from).collect() } +/// `litellm.redis.stores.router_usage.CounterIncrement`. +#[derive(FromPyObject)] +struct PyIncrement { + counter: PyCounter, + amount: PyNumber, + ttl: Option, +} + +impl TryFrom for Increment { + type Error = PyErr; + + fn try_from(increment: PyIncrement) -> PyResult { + Ok(Self { + counter: increment.counter.into(), + amount: increment.amount.0, + ttl: ttl(increment.ttl)?, + }) + } +} + +fn increments(increments: Vec) -> PyResult> { + increments.into_iter().map(Increment::try_from).collect() +} + +fn numbers(read: Vec>) -> Vec> { + read.into_iter().map(|count| count.map(PyNumber)).collect() +} + +fn totals(totals: Option>) -> Option> { + totals.map(|totals| totals.into_iter().map(PyNumber).collect()) +} + /// One router's global TPM/RPM counters. The `_sync` methods block with the GIL released, for the /// Rust router's sync calls, which have no event loop to await on. #[pyclass(frozen, name = "NativeRouterUsage")] pub(crate) struct NativeRouterUsage(Arc); +/// A usage read declared on a batch, with the refreshes it reserved, given back at most once. +#[pyclass(frozen, name = "NativeUsageRead")] +pub(crate) struct NativeUsageRead { + #[pyo3(get)] + read: Py, + store: Arc, + reservation: Mutex>, +} + +#[pymethods] +impl NativeUsageRead { + fn release(&self) { + let reservation = self + .reservation + .lock() + .unwrap_or_else(PoisonError::into_inner) + .take(); + if let Some(reservation) = reservation { + self.store.release(reservation); + } + } +} + #[pymethods] impl NativeRouterUsage { #[new] @@ -166,6 +233,127 @@ impl NativeRouterUsage { }) } + /// `get_many` declared on the batch `send` picks, with the refreshes it reserved. + fn usage( + &self, + py: Python<'_>, + counters: Vec, + send: &NativeSend, + ) -> PyResult { + let declared = self + .0 + .declare_get_many(send.send(), &self::counters(counters)) + .map_err(kind_error)?; + Ok(NativeUsageRead { + read: Py::new( + py, + pending("async_batch_get_cache", kind_error, declared.read, numbers), + )?, + store: Arc::clone(&self.0), + reservation: Mutex::new(Some(declared.reservation)), + }) + } + + fn increment_pipeline<'py>( + &self, + py: Python<'py>, + increments: Vec, + ) -> PyResult> { + let store = self.0.clone(); + let increments = self::increments(increments)?; + redis_async(py, "async_increment_pipeline", kind_error, async move { + Ok(totals(store.increment_pipeline(&increments).await)) + }) + } + + /// Counts in memory now and declares the increments on the request's post-call batch. + fn increment_post_call( + &self, + request: &NativeRequestRedis, + increments: Vec, + ) -> PyResult<()> { + self.0 + .increment_post_call(&request.0, &self::increments(increments)?); + Ok(()) + } + + fn increment_remote<'py>( + &self, + py: Python<'py>, + increments: Vec, + ) -> PyResult> { + let store = self.0.clone(); + let increments = self::increments(increments)?; + redis_async(py, "async_increment_pipeline", kind_error, async move { + store.increment_remote(&increments).await.map(totals) + }) + } + + fn increment_local(&self, increment: PyIncrement) -> PyResult { + self.0 + .increment_local(&increment.try_into()?) + .map(PyNumber) + .map_err(kind_error) + } + + fn get_local(&self, counter: PyCounter) -> PyResult> { + self.0 + .get_local(&counter.into()) + .map(|count| count.map(PyNumber)) + .map_err(kind_error) + } + + fn get_local_many(&self, counters: Vec) -> PyResult>> { + self.0 + .get_local_many(&self::counters(counters)) + .map(numbers) + .map_err(kind_error) + } + + fn set_local(&self, counter: PyCounter, value: PyNumber) -> PyResult<()> { + self.0 + .set_local(&counter.into(), value.0) + .map_err(kind_error) + } + + fn current<'py>( + &self, + py: Python<'py>, + counter: PyCounter, + limit: PyNumber, + ) -> PyResult> { + let store = self.0.clone(); + let counter = Counter::from(counter); + redis_async(py, "async_get_cache", kind_error, async move { + Ok(store.current(&counter, limit.0).await?.map(PyNumber)) + }) + } + + fn shared_or_local<'py>( + &self, + py: Python<'py>, + counter: PyCounter, + ) -> PyResult> { + let store = self.0.clone(); + let counter = Counter::from(counter); + redis_async(py, "async_get_cache", kind_error, async move { + Ok(store.shared_or_local(&counter).await?.map(PyNumber)) + }) + } + + fn current_sync( + &self, + py: Python<'_>, + counter: PyCounter, + limit: PyNumber, + ) -> PyResult> { + let store = self.0.clone(); + let counter = Counter::from(counter); + redis_sync(py, "get_cache", kind_error, async move { + Ok(store.current(&counter, limit.0).await?.map(PyNumber)) + }) + } + fn flush_local_counters(&self) -> PyResult<()> { self.0.flush_local_counters().map_err(kind_error) } diff --git a/litellm/rust/crates/python-bridge/src/router/mod.rs b/litellm/rust/crates/python-bridge/src/router/mod.rs index 679194e012a..535717d8ed8 100644 --- a/litellm/rust/crates/python-bridge/src/router/mod.rs +++ b/litellm/rust/crates/python-bridge/src/router/mod.rs @@ -31,7 +31,7 @@ use pyo3::{ }; use serde_json::Value; -use litellm_redis_store::router_usage::{Counter, Kind}; +use litellm_redis_store::router_usage::{Counter, Kind, Scope}; use self::host::{ CacheOp, CacheOutcome, DriveError, Driver, Finish, HostCall, HostCallResponse, @@ -385,21 +385,6 @@ impl RouterCoreHandle { .collect() } - /// The key deployment `id` serving litellm `model` counts `metric` (`tpm`, `rpm`, `itpm` or - /// `otpm`) under in `minute` (`HH-MM`), as the usage store spells it. - fn usage_key(&self, id: &str, model: &str, metric: &str, minute: &str) -> PyResult { - let kind: Kind = metric - .parse() - .map_err(|_| PyValueError::new_err(format!("`{metric}` is not a usage metric")))?; - Ok(Counter { - deployment_id: id.to_owned(), - model: Some(model.to_owned()), - minute: minute.to_owned(), - kind, - } - .key()) - } - /// Keep the providers `other`'s build resolved for the deployments this router serves too. fn adopt_providers(&self, other: &Self) { self.router.adopt_providers(&other.router); @@ -1348,9 +1333,10 @@ fn cache_outcome(op: CacheOp, outcome: Option>) -> PyResult String { Counter { + scope: Scope::GlobalRouter, deployment_id: window.deployment_id, model: Some(window.model), - minute: window.minute, + window: window.minute, kind: Kind::Tpm, } .key() diff --git a/litellm/rust/crates/redis-store/src/dual_cache/batch_read.rs b/litellm/rust/crates/redis-store/src/dual_cache/batch_read.rs index 71dfda49233..abc49db428b 100644 --- a/litellm/rust/crates/redis-store/src/dual_cache/batch_read.rs +++ b/litellm/rust/crates/redis-store/src/dual_cache/batch_read.rs @@ -4,7 +4,7 @@ use std::{ time::Duration, }; -use litellm_redis_client::{Declared, Error as RedisError, commands}; +use litellm_redis_client::{Declared, Error as RedisError, RedisBatch, commands}; use super::{DualCache, MemorySettings}; use crate::{Error, FromValue, Send, codec, memory::lock, request::Misses}; @@ -66,9 +66,11 @@ impl DualCache = reservation.keys.iter().map(String::as_str).collect(); - let replies = send.batch(client).add(commands::mget::>(&keys)); + let batch = send.batch(client); + let replies = batch.add(commands::mget::>(&keys)); Fetch { reservation, + batch, replies, misses: send.misses(client), } @@ -81,24 +83,37 @@ impl DualCache Result>, Error> { + if let Some(batch) = declared.batch() { + batch.flush().await; + } + self.settle_batch_get(declared).await + } + + /// The second half of `batch_get` once the read's batch answers, without flushing it. Memory + /// is read again here, so a value written to memory since the declaration wins over what Redis + /// answered. The keys Redis lacks are noted as the request's misses. An undecodable value + /// serves memory for the whole read, where `main`'s `RedisCache` returns the raw bytes and + /// `DualCache` backfills memory with them. + pub(crate) async fn settle_batch_get( + &self, + declared: DeclaredRead, ) -> Result>, Error> { let DeclaredRead { keys, fetch } = declared; let Some(Fetch { reservation, + batch: _, replies, misses, }) = fetch else { return self.read_memory(&keys); }; - let fetched = replies.resolve().await; + let fetched = replies.settled().await; let local = self.read_memory(&keys)?; let fetched = match fetched { Ok(fetched) => fetched, @@ -174,7 +189,7 @@ impl DualCache, } +impl DeclaredRead { + /// The batch carrying the read, `None` when memory answers it alone. + pub(crate) fn batch(&self) -> Option { + self.fetch.as_ref().map(|fetch| fetch.batch.clone()) + } + + /// The refreshes the read reserved, `None` when memory answers it alone. + pub(crate) fn reservation(&self) -> Option { + self.fetch.as_ref().map(|fetch| fetch.reservation.clone()) + } +} + struct Fetch { reservation: Reservation, + batch: RedisBatch, replies: Declared>>>, misses: Option, } -struct Reservation { +#[derive(Clone)] +pub(crate) struct Reservation { keys: Vec, previous: Vec>, } diff --git a/litellm/rust/crates/redis-store/src/dual_cache/increment.rs b/litellm/rust/crates/redis-store/src/dual_cache/increment.rs index 705bc71b2c1..d565cdb91e0 100644 --- a/litellm/rust/crates/redis-store/src/dual_cache/increment.rs +++ b/litellm/rust/crates/redis-store/src/dual_cache/increment.rs @@ -1,7 +1,7 @@ use std::time::Duration; use futures_util::future::join_all; -use litellm_redis_client::{Declared, Expiry, RedisBatch, commands}; +use litellm_redis_client::{Declared, Error as RedisError, Expiry, RedisBatch, commands}; use super::DualCache; use crate::{Error, Number, RequestRedis, Send, failure_log::redis_failure, memory::Counter}; @@ -63,10 +63,26 @@ impl DualCache { .collect::, _>>() .map_err(|error| redis_failure!(&Error::from(error), PIPELINE_FAILED)) .ok()?; - let Some(client) = self.redis.client() else { - return Some(local); - }; - let pipeline = RedisBatch::pipelined(client.clone()); + match self.increment_pipeline_remote(increments).await { + None => Some(local), + Some(Ok(totals)) => Some(totals), + Some(Err(error)) => { + redis_failure!(&error, PIPELINE_FAILED); + Some(local) + } + } + } + + /// `RedisCache.async_increment_pipeline`, Redis alone: one pipeline of `INCRBYFLOAT` with + /// `EXPIRE` per increment. `None` without Redis or with nothing to increment. + pub(crate) async fn increment_pipeline_remote( + &self, + increments: &[PipelineIncrement], + ) -> Option, RedisError>> { + if increments.is_empty() { + return None; + } + let pipeline = RedisBatch::pipelined(self.redis.client()?.clone()); let totals = increments .iter() .map(|increment| { @@ -77,17 +93,13 @@ impl DualCache { )) }) .collect::>(); - match Declared::all(totals) - .await - .into_iter() - .collect::, _>>() - { - Ok(totals) => Some(totals.into_iter().map(Number::Float).collect()), - Err(error) => { - redis_failure!(&error, PIPELINE_FAILED); - Some(local) - } - } + Some( + Declared::all(totals) + .await + .into_iter() + .map(|total| total.map(Number::Float)) + .collect(), + ) } /// `DualCache.async_increment_cache_pipeline_post_call`: memory now, stopping at the first diff --git a/litellm/rust/crates/redis-store/src/dual_cache/mod.rs b/litellm/rust/crates/redis-store/src/dual_cache/mod.rs index b2c277a118e..10e62076d41 100644 --- a/litellm/rust/crates/redis-store/src/dual_cache/mod.rs +++ b/litellm/rust/crates/redis-store/src/dual_cache/mod.rs @@ -9,7 +9,10 @@ use std::{ }; pub use self::batch_read::TierPolicy; -pub(crate) use self::{batch_read::BatchReads, increment::PipelineIncrement}; +pub(crate) use self::{ + batch_read::{BatchReads, Reservation}, + increment::PipelineIncrement, +}; use crate::{ memory::{Clock, Memory, lock}, remote::RemoteSlot, diff --git a/litellm/rust/crates/redis-store/src/lib.rs b/litellm/rust/crates/redis-store/src/lib.rs index eaede611add..55b3abb88cc 100644 --- a/litellm/rust/crates/redis-store/src/lib.rs +++ b/litellm/rust/crates/redis-store/src/lib.rs @@ -21,5 +21,5 @@ pub use error::{Error, ErrorKind}; pub use memory::{Clock, system_clock}; pub use number::Number; pub use remote::Remote; -pub use request::{Pending, RequestRedis, Send}; +pub use request::{DeadlineTag, Pending, RequestRedis, Send}; pub use value::{FromValue, IntoValue}; diff --git a/litellm/rust/crates/redis-store/src/number.rs b/litellm/rust/crates/redis-store/src/number.rs index 7f3e9f4d10b..ae65461e542 100644 --- a/litellm/rust/crates/redis-store/src/number.rs +++ b/litellm/rust/crates/redis-store/src/number.rs @@ -70,13 +70,16 @@ impl Number { } } -/// Python's `+`: two ints stay an int, anything with a float is a float. +/// Python's `+`: two ints stay an int, anything with a float is a float. An int sum past `i64` +/// becomes the nearest float, which still compares above any limit the way Python's int does. impl Add for Number { type Output = Self; fn add(self, other: Self) -> Self { match (self, other) { - (Self::Int(left), Self::Int(right)) => Self::Int(left + right), + (Self::Int(left), Self::Int(right)) => left + .checked_add(right) + .map_or_else(|| Self::Float(left as f64 + right as f64), Self::Int), (left, right) => Self::Float(left.as_f64() + right.as_f64()), } } diff --git a/litellm/rust/crates/redis-store/src/request.rs b/litellm/rust/crates/redis-store/src/request.rs index 79497a5d503..93101ea01be 100644 --- a/litellm/rust/crates/redis-store/src/request.rs +++ b/litellm/rust/crates/redis-store/src/request.rs @@ -22,7 +22,7 @@ struct Shared { runtime: Handle, logs: Dispatch, post_call_deadline: Duration, - deadline_tag: Option, + deadline_tag: DeadlineTag, batches: Mutex>, follow_ups: Mutex>>, deadline: Mutex>>, @@ -60,15 +60,16 @@ impl RequestRedis { /// Python one that has no Tokio context. Work spawned for the request logs to the subscriber /// current here, since a declaring thread may have none. pub fn new(runtime: Handle, post_call_deadline: Duration) -> Self { - Self::with_deadline_tag(runtime, post_call_deadline, None) + Self::with_deadline_tag(runtime, post_call_deadline, Arc::new(|| None)) } - /// As `new`, with the deadline's flush of the post-call batch run inside `tagged(deadline_tag)`, - /// the scope an explicit `flush_post` runs in. + /// As `new`, with the deadline's flush of the post-call batch run inside the tag `deadline_tag` + /// gives for the declaration that armed it, as an explicit `flush_post` runs inside its + /// caller's. pub fn with_deadline_tag( runtime: Handle, post_call_deadline: Duration, - deadline_tag: Option, + deadline_tag: DeadlineTag, ) -> Self { Self(Arc::new(Shared { runtime, @@ -107,6 +108,16 @@ impl RequestRedis { lock(&self.0.follow_ups).push(follow_up); } + /// Runs `task` on the request's runtime without `flush_post` waiting for it, for work that + /// settles with the pre-call batch whoever awaits it. + pub(crate) fn detach(&self, task: impl Future + std::marker::Send + 'static) { + drop( + self.0 + .runtime + .spawn(self.logged(within(current_tag(), task))), + ); + } + fn logged(&self, task: F) -> impl Future + use { task.in_current_span().with_subscriber(self.0.logs.clone()) } @@ -151,6 +162,12 @@ impl RequestRedis { } fn arm_deadline(&self) { + if lock(&self.0.deadline).is_some() { + return; + } + // Asked before taking the lock: the host may answer under its interpreter lock, which + // another declaring thread can hold while it waits for this one. + let tag = (self.0.deadline_tag)(); let mut deadline = lock(&self.0.deadline); if deadline.is_some() { return; @@ -158,7 +175,6 @@ impl RequestRedis { let (disarm, disarmed) = oneshot::channel(); *deadline = Some(disarm); let request = self.clone(); - let tag = self.0.deadline_tag.clone(); self.0.runtime.spawn(self.logged(async move { tokio::select! { () = tokio::time::sleep(request.0.post_call_deadline) => within(tag, request.flush_post()).await, @@ -168,6 +184,10 @@ impl RequestRedis { } } +/// Names the scope a deadline flush runs in. It is asked when a post-call declaration arms the +/// deadline, on the declaring thread, so it can describe that declaration's caller. +pub type DeadlineTag = Arc Option + std::marker::Send + Sync>; + async fn within(tag: Option, task: F) -> F::Output { match tag { Some(tag) => tagged(tag, task).await, @@ -196,6 +216,14 @@ impl Send<'_> { } } + /// The request the op rides, `None` for one on its own. + pub(crate) fn request(&self) -> Option<&RequestRedis> { + match self { + Self::Now => None, + Self::PreCall(request) | Self::PostCall(request) => Some(request), + } + } + /// The request's misses on `client`; `None` for an op outside a request. pub(crate) fn misses(&self, client: &RedisClient) -> Option { match self { diff --git a/litellm/rust/crates/redis-store/src/router_usage/keys.rs b/litellm/rust/crates/redis-store/src/router_usage/keys.rs index 1482668ab2b..f897ed92020 100644 --- a/litellm/rust/crates/redis-store/src/router_usage/keys.rs +++ b/litellm/rust/crates/redis-store/src/router_usage/keys.rs @@ -1,5 +1,15 @@ +/// Who counts under a key. The Rust router core's counters carry a `global_router:` prefix, and +/// the counters usage-based-routing-v2 and the model rate limit check share carry none. +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, strum::EnumString, strum::IntoStaticStr)] +pub enum Scope { + #[strum(serialize = "global_router")] + GlobalRouter, + #[strum(serialize = "deployment")] + Deployment, +} + /// A counter's key segment, which is also how the Python router names it. -#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::EnumString, strum::IntoStaticStr)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, strum::EnumString, strum::IntoStaticStr)] pub enum Kind { #[strum(serialize = "tpm")] Tpm, @@ -11,25 +21,36 @@ pub enum Kind { /// Output tokens per minute. #[strum(serialize = "otpm")] Otpm, + /// Tokens per UTC day, whose window is `YYYY-MM-DD`. + #[strum(serialize = "tpd")] + Tpd, + /// Requests per UTC day. + #[strum(serialize = "rpd")] + Rpd, } -#[derive(Clone, Debug, PartialEq, Eq)] +#[derive(Clone, Debug, PartialEq, Eq, Hash)] pub struct Counter { + pub scope: Scope, pub deployment_id: String, pub model: Option, - pub minute: String, + /// `HH-MM` for a per-minute kind, `YYYY-MM-DD` for a per-day one. + pub window: String, pub kind: Kind, } impl Counter { - /// `global_router:{id}:{model}:{kind}:{HH-MM}`, where Python formats a missing model as `None`. - /// Public only for the hosts that hand a counter's key back to their callers. + /// `global_router:{id}:{model}:{kind}:{window}` or `{id}:{model}:{kind}:{window}`, where Python + /// formats a missing model as `None`. Public only for the hosts that hand a counter's key back + /// to their callers. pub fn key(&self) -> String { let kind: &str = self.kind.into(); let model = self.model.as_deref().unwrap_or("None"); - format!( - "global_router:{}:{model}:{kind}:{}", - self.deployment_id, self.minute - ) + let id = &self.deployment_id; + let window = &self.window; + match self.scope { + Scope::GlobalRouter => format!("global_router:{id}:{model}:{kind}:{window}"), + Scope::Deployment => format!("{id}:{model}:{kind}:{window}"), + } } } diff --git a/litellm/rust/crates/redis-store/src/router_usage/mod.rs b/litellm/rust/crates/redis-store/src/router_usage/mod.rs index 4b123a0f094..09cce2d3b61 100644 --- a/litellm/rust/crates/redis-store/src/router_usage/mod.rs +++ b/litellm/rust/crates/redis-store/src/router_usage/mod.rs @@ -2,10 +2,14 @@ mod keys; use std::{sync::Arc, time::Duration}; -pub use self::keys::{Counter, Kind}; +use litellm_redis_client::{Error as RedisError, commands}; +use tokio::sync::oneshot; + +pub use self::keys::{Counter, Kind, Scope}; use crate::{ - Error, Number, - dual_cache::{BatchReads, DualCache, PipelineIncrement, TierPolicy}, + Error, Number, Pending, RequestRedis, Send, codec, + dual_cache::{BatchReads, DualCache, PipelineIncrement, Reservation, TierPolicy}, + failure_log::redis_failure, memory::{Clock, system_clock}, remote::{Remote, RemoteSlot}, }; @@ -13,7 +17,40 @@ use crate::{ /// `RoutingArgs.ttl`, the window of a per-minute counter. pub const COUNTER_TTL: Duration = Duration::from_secs(60); -/// One router's global TPM/RPM counters: memory per router, Redis shared by every pod. +/// A `RedisPipelineIncrementOperation` on one counter. +#[derive(Clone, Debug, PartialEq)] +pub struct Increment { + pub counter: Counter, + pub amount: Number, + pub ttl: Option, +} + +impl From<&Increment> for PipelineIncrement { + fn from(increment: &Increment) -> Self { + Self { + key: increment.counter.key(), + amount: increment.amount, + ttl: increment.ttl, + } + } +} + +fn pipeline(increments: &[Increment]) -> Vec { + increments.iter().map(PipelineIncrement::from).collect() +} + +/// The throttle refreshes a declared usage read reserved, which `RouterUsage::release` gives back. +pub struct UsageReservation(Option); + +/// A usage read `RouterUsage::declare_get_many` declared, with the refreshes it reserved. +pub struct DeclaredUsage { + pub read: Pending>>, + pub reservation: UsageReservation, +} + +/// One router's usage counters: memory per router, Redis shared by every pod. The router core's +/// `global_router:` counters and the per-deployment ones its strategies and pre-call checks count +/// share one `DualCache`, as they shared the router's on `main`. pub struct RouterUsage { usage: DualCache, } @@ -61,9 +98,10 @@ impl RouterUsage { ) -> Result<(), Error> { let increment = |kind, amount| PipelineIncrement { key: Counter { + scope: Scope::GlobalRouter, deployment_id: deployment_id.to_owned(), model: model.map(str::to_owned), - minute: minute.to_owned(), + window: minute.to_owned(), kind, } .key(), @@ -94,8 +132,130 @@ impl RouterUsage { } pub async fn get_many(&self, counters: &[Counter]) -> Result>, Error> { - let keys = counters.iter().map(Counter::key).collect::>(); - self.usage.batch_get(&keys).await + self.usage.batch_get(&keys(counters)).await + } + + /// `get_many` declared on the batch `send` picks, so a pick's usage read can ride the request's + /// pre-call pipeline, with the refreshes it reserved. On a request it settles when its batch + /// answers, whoever awaits it, so a prefetch the pick never takes still backfills memory or gives + /// its reservation back, as the settled hook of `main`'s prefetch does. + pub fn declare_get_many( + self: &Arc, + send: Send<'_>, + counters: &[Counter], + ) -> Result { + let read = self.usage.declare_batch_get(send, &keys(counters))?; + let reservation = UsageReservation(read.reservation()); + let store = Arc::clone(self); + let (Some(request), Some(batch)) = (send.request(), read.batch()) else { + return Ok(DeclaredUsage { + read: Box::pin(async move { store.usage.resolve_batch_get(read).await }), + reservation, + }); + }; + let (settled, answer) = oneshot::channel(); + let settling = Arc::clone(self); + request.detach(async move { + drop(settled.send(settling.usage.settle_batch_get(read).await)); + }); + let counters = counters.to_vec(); + Ok(DeclaredUsage { + read: Box::pin(async move { + batch.flush().await; + match answer.await { + Ok(read) => read, + Err(_) => store.get_local_many(&counters), + } + }), + reservation, + }) + } + + /// `ArmedRead.release`: a prefetch a pick took but could not use gives its refreshes back, so the + /// pick's own read goes to Redis. + pub fn release(&self, reservation: UsageReservation) { + if let Some(reservation) = reservation.0 { + self.usage.release(reservation); + } + } + + /// `DualCache.async_increment_cache_pipeline`: memory, then one pipeline on its own round trip. + /// `None` when memory refuses an increment, and the memory totals when Redis fails. + pub async fn increment_pipeline(&self, increments: &[Increment]) -> Option> { + self.usage.increment_pipeline(&pipeline(increments)).await + } + + /// `DualCache.async_increment_cache_pipeline_post_call`: memory now, Redis on the request's + /// post-call batch. + pub fn increment_post_call(&self, request: &RequestRedis, increments: &[Increment]) { + self.usage + .increment_pipeline_post_call(request, &pipeline(increments)); + } + + /// `RedisCache.async_increment_pipeline`, which a strategy's periodic sync sends with the + /// increments it already counted in memory. `None` without Redis or increments, and a Redis failure is the + /// caller's to handle, as there. + pub async fn increment_remote( + &self, + increments: &[Increment], + ) -> Result>, Error> { + self.usage + .increment_pipeline_remote(&pipeline(increments)) + .await + .transpose() + .map_err(Error::from) + } + + /// `in_memory_cache.async_increment`. + pub fn increment_local(&self, increment: &Increment) -> Result { + Ok(self.usage.memory().increment( + &increment.counter.key(), + increment.amount, + increment.ttl, + )?) + } + + /// `get_cache(local_only=True)`. + pub fn get_local(&self, counter: &Counter) -> Result, Error> { + Ok(self.usage.memory().get(&counter.key())?) + } + + /// `in_memory_cache.async_batch_get_cache`. + pub fn get_local_many(&self, counters: &[Counter]) -> Result>, Error> { + counters + .iter() + .map(|counter| self.get_local(counter)) + .collect() + } + + /// `in_memory_cache.async_set_cache` without a TTL, which keeps `InMemoryCache`'s own. + pub fn set_local(&self, counter: &Counter, value: Number) -> Result<(), Error> { + Ok(self.usage.memory().set(&counter.key(), value, None)?) + } + + /// The model rate limit check's current usage: memory alone when it already reaches `limit`, + /// else `redis_cache.get_cache`. No Redis and an open breaker serve memory, and any other + /// failure reads as `None`, as `RedisCache` swallows it. + pub async fn current(&self, counter: &Counter, limit: Number) -> Result, Error> { + let local = self.get_local(counter)?; + if local.is_some_and(|local| local.as_f64() >= limit.as_f64()) { + return Ok(local); + } + Ok(match self.shared(counter).await { + Shared::Unasked => local, + Shared::Failed => None, + Shared::Read(count) => count, + }) + } + + /// The io-token check's usage: what every pod counted in Redis, else this router's memory + /// count when Redis has none, is absent, refuses or fails. + pub async fn shared_or_local(&self, counter: &Counter) -> Result, Error> { + let local = self.get_local(counter)?; + Ok(match self.shared(counter).await { + Shared::Read(Some(count)) => Some(count), + Shared::Unasked | Shared::Failed | Shared::Read(None) => local, + }) } pub fn flush_local_counters(&self) -> Result<(), Error> { @@ -107,4 +267,42 @@ impl RouterUsage { pub fn snapshot(&self) -> Result, Error> { Ok(self.usage.memory().entries()?) } + + /// `redis_cache.get_cache` of one counter, which `RedisCache` answers with `None` on a failure + /// and its breaker refuses. + async fn shared(&self, counter: &Counter) -> Shared { + let Some(client) = self.usage.redis().client() else { + return Shared::Unasked; + }; + let key = counter.key(); + match client.send(commands::get::>(&key)).await { + Err(RedisError::BreakerOpen) => Shared::Unasked, + Err(error) => { + redis_failure!( + &error, + "litellm.caching.caching: get() - Got exception from REDIS" + ); + Shared::Failed + } + Ok(raw) => match raw.as_deref().map(codec::decode).transpose() { + Ok(count) => Shared::Read(count), + Err(error) => { + tracing::warn!(%error, key, "a usage counter did not decode"); + Shared::Failed + } + }, + } + } +} + +/// What one counter's Redis read came back with. +enum Shared { + /// No Redis, or its breaker refused the read. + Unasked, + Failed, + Read(Option), +} + +fn keys(counters: &[Counter]) -> Vec { + counters.iter().map(Counter::key).collect() } diff --git a/litellm/rust/crates/redis-store/tests/request.rs b/litellm/rust/crates/redis-store/tests/request.rs index 1ee437d4d5a..5ecf39ee4c6 100644 --- a/litellm/rust/crates/redis-store/tests/request.rs +++ b/litellm/rust/crates/redis-store/tests/request.rs @@ -1,6 +1,9 @@ mod support; -use std::{sync::Arc, time::Duration}; +use std::{ + sync::{Arc, Mutex}, + time::Duration, +}; use litellm_redis_client::commands; use litellm_redis_client::{Amount, CallTag, Expiry, Operation, tagged}; @@ -192,16 +195,19 @@ async fn a_post_declaration_after_a_flush_arms_a_new_deadline( #[rstest] #[tokio::test] -async fn the_deadline_flush_runs_under_the_deadline_tag( +async fn the_deadline_flush_runs_under_the_tag_given_when_its_first_op_is_declared( #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, ) { let deployment = RedisDeployment::start(topology).await.unwrap(); let (client, tags) = connect_tagged(&deployment); + let caller = Arc::new(Mutex::new("the request before it has a span")); + let asked = caller.clone(); let request = RequestRedis::with_deadline_tag( Handle::current(), Duration::from_millis(100), - Some(Arc::new("post_call_redis_batch") as CallTag), + Arc::new(move || Some(Arc::new(*asked.lock().unwrap()) as CallTag)), ); + *caller.lock().unwrap() = "post_call_redis_batch"; let declared = tagged(Arc::new("declaring op"), async { request diff --git a/litellm/rust/crates/redis-store/tests/router_usage.rs b/litellm/rust/crates/redis-store/tests/router_usage.rs index cc327e9aa1d..d2140cec28b 100644 --- a/litellm/rust/crates/redis-store/tests/router_usage.rs +++ b/litellm/rust/crates/redis-store/tests/router_usage.rs @@ -11,30 +11,420 @@ use std::{ use litellm_redis_client::commands; use litellm_redis_client::{Amount, Expiry, Operation, Ttl}; use litellm_redis_store::{ - Number, TierPolicy, - router_usage::{COUNTER_TTL, Counter, Kind, RouterUsage}, + Number, RequestRedis, Send as RedisSend, TierPolicy, + router_usage::{COUNTER_TTL, Counter, Increment, Kind, RouterUsage, Scope}, }; use litellm_redis_testing::{RedisDeployment, Topology}; use rstest::rstest; -use support::{USAGE_POLICY as POLICY, connect, remote, usage, within}; +use support::{USAGE_POLICY as POLICY, connect, deployment_usage, remote, usage, within}; +use tokio::runtime::Handle; + +const DAY: Duration = Duration::from_secs(86_400); fn tpm_and_rpm(model: Option<&str>) -> [Counter; 2] { [Kind::Tpm, Kind::Rpm].map(|kind| Counter { + scope: Scope::GlobalRouter, deployment_id: "dep".to_owned(), model: model.map(str::to_owned), - minute: "12-34".to_owned(), + window: "12-34".to_owned(), kind, }) } +fn increment(counter: &Counter, amount: Number, ttl: Duration) -> Increment { + Increment { + counter: counter.clone(), + amount, + ttl: Some(ttl), + } +} + +async fn ttl_of(raw: &litellm_redis_client::RedisClient, key: &str) -> Duration { + let Ttl::Expires(remaining) = raw.send(commands::ttl(key)).await.unwrap() else { + panic!("{key} has no TTL"); + }; + remaining +} + #[rstest] #[case::tpm(Kind::Tpm, "tpm")] #[case::rpm(Kind::Rpm, "rpm")] +#[case::itpm(Kind::Itpm, "itpm")] +#[case::otpm(Kind::Otpm, "otpm")] +#[case::tpd(Kind::Tpd, "tpd")] +#[case::rpd(Kind::Rpd, "rpd")] fn a_counter_kind_round_trips_through_its_key_segment(#[case] kind: Kind, #[case] segment: &str) { assert_eq!(<&str>::from(kind), segment); assert_eq!(segment.parse::().unwrap(), kind); } +#[rstest] +#[case::global_router(Scope::GlobalRouter, "global_router")] +#[case::deployment(Scope::Deployment, "deployment")] +fn a_counter_scope_round_trips_through_its_spelling(#[case] scope: Scope, #[case] spelling: &str) { + assert_eq!(<&str>::from(scope), spelling); + assert_eq!(spelling.parse::().unwrap(), scope); +} + +#[rstest] +#[tokio::test] +async fn reads_deployment_counters_python_incremented_under_mains_keys( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY); + let cases = deployment_usage(); + for (case, delta) in cases.iter().zip([4.0, 1200.5]) { + raw.send(commands::incr_by_float( + case.key, + Amount::Float(delta), + Some(Expiry::Refresh(DAY)), + )) + .await + .unwrap(); + } + + assert_eq!( + store + .get_many(&cases.map(|case| case.counter)) + .await + .unwrap(), + vec![Some(Number::Int(4)), Some(Number::Float(1200.5))] + ); +} + +#[rstest] +#[tokio::test] +async fn a_declared_usage_read_rides_the_request_batch_and_backfills_memory( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let (client, trips) = support::connect_observed(&deployment); + let store = Arc::new(RouterUsage::new(Some(remote(client)), POLICY)); + let cases = deployment_usage(); + for (case, delta) in cases.iter().zip([2.0, 30.0]) { + raw.send(commands::incr_by_float( + case.key, + Amount::Float(delta), + Some(Expiry::Refresh(DAY)), + )) + .await + .unwrap(); + } + let counters = cases.map(|case| case.counter); + let request = RequestRedis::new(Handle::current(), Duration::from_secs(60)); + + let declared = store + .declare_get_many(RedisSend::PreCall(&request), &counters) + .unwrap() + .read; + assert!(trips.seen().is_empty()); + request.flush_pre().await; + let read = declared.await.unwrap(); + + let one_mget = match topology { + Topology::Standalone => Operation::Pipeline { ops: 1 }, + Topology::Cluster => Operation::Command("MGET"), + }; + assert_eq!(trips.seen(), [one_mget]); + assert_eq!(read, vec![Some(Number::Int(2)), Some(Number::Int(30))]); + assert_eq!( + store.get_local_many(&counters).unwrap(), + vec![Some(Number::Int(2)), Some(Number::Int(30))] + ); +} + +// main's ArmedRead settles on the batch's reply: a prefetch the pick never takes still backfills +// memory, and the throttle then keeps later reads off Redis +#[rstest] +#[tokio::test] +async fn an_abandoned_usage_prefetch_still_backfills_memory_and_holds_the_throttle( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let (client, trips) = support::connect_observed(&deployment); + let store = Arc::new(RouterUsage::new(Some(remote(client)), POLICY)); + let [rpm, tpd] = deployment_usage(); + raw.send(commands::incr_by_float( + rpm.key, + Amount::Float(6.0), + Some(Expiry::Refresh(COUNTER_TTL)), + )) + .await + .unwrap(); + let counters = [rpm.counter.clone(), tpd.counter.clone()]; + let request = RequestRedis::new(Handle::current(), Duration::from_secs(60)); + + drop( + store + .declare_get_many(RedisSend::PreCall(&request), &counters) + .unwrap(), + ); + request.flush_pre().await; + litellm_redis_testing::eventually(async || store.get_local(&rpm.counter).unwrap()).await; + let trips_after_prefetch = trips.seen().len(); + let read = store.get_many(&counters).await.unwrap(); + + assert_eq!(read, vec![Some(Number::Int(6)), None]); + assert_eq!(trips.seen().len(), trips_after_prefetch); +} + +// main's `_take` releases a prefetch that does not cover the pick taking it, so the counters it +// reserved are due again and the next read of them goes to Redis +#[rstest] +#[tokio::test] +async fn a_released_usage_prefetch_lets_the_next_read_go_to_redis( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let (client, trips) = support::connect_observed(&deployment); + let store = Arc::new(RouterUsage::new(Some(remote(client)), POLICY)); + let [rpm, tpd] = deployment_usage(); + let counters = [rpm.counter.clone(), tpd.counter.clone()]; + let request = RequestRedis::new(Handle::current(), Duration::from_secs(60)); + + let declared = store + .declare_get_many(RedisSend::PreCall(&request), &counters) + .unwrap(); + store.release(declared.reservation); + raw.send(commands::incr_by_float( + rpm.key, + Amount::Float(4.0), + Some(Expiry::Refresh(COUNTER_TTL)), + )) + .await + .unwrap(); + let read = store.get_many(&counters).await.unwrap(); + + assert_eq!(read, vec![Some(Number::Int(4)), None]); + assert_eq!(trips.seen().len(), 1); +} + +// The model rate limit check charges a request's rpm and rpd in one async_increment_cache_pipeline, +// whose RedisCache pipeline refreshes each window with EXPIRE +#[rstest] +#[tokio::test] +async fn a_pipeline_increment_is_one_round_trip_refreshing_each_window( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let (client, trips) = support::connect_observed(&deployment); + let store = RouterUsage::new(Some(remote(client)), POLICY); + let [rpm, tpd] = deployment_usage(); + raw.send(commands::incr_by_float( + rpm.key, + Amount::Float(3.0), + Some(Expiry::Refresh(Duration::from_secs(10))), + )) + .await + .unwrap(); + + let totals = store + .increment_pipeline(&[ + increment(&rpm.counter, Number::Int(1), COUNTER_TTL), + increment(&tpd.counter, Number::Int(1), DAY), + ]) + .await; + + assert_eq!(trips.seen(), [Operation::Pipeline { ops: 2 }]); + assert_eq!(totals, Some(vec![Number::Float(4.0), Number::Float(1.0)])); + assert!(within(ttl_of(&raw, rpm.key).await, 60)); + assert!(within(ttl_of(&raw, tpd.key).await, 86_400)); +} + +// LowestTPMLoggingHandler_v2.async_log_success_event's async_increment_cache_post_call +#[rstest] +#[tokio::test] +async fn a_post_call_increment_counts_in_memory_now_and_in_redis_at_flush_post( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY); + let [rpm, _] = deployment_usage(); + raw.send(commands::incr_by_float( + rpm.key, + Amount::Float(5.0), + Some(Expiry::Refresh(COUNTER_TTL)), + )) + .await + .unwrap(); + let request = RequestRedis::new(Handle::current(), Duration::from_secs(60)); + + store.increment_post_call( + &request, + &[increment(&rpm.counter, Number::Int(7), COUNTER_TTL)], + ); + let before_flush: Option = raw.send(commands::get(rpm.key)).await.unwrap(); + request.flush_post().await; + let after_flush: Option = raw.send(commands::get(rpm.key)).await.unwrap(); + + assert_eq!(store.get_local(&rpm.counter).unwrap(), Some(Number::Int(7))); + assert_eq!(before_flush, Some("5".to_owned())); + assert_eq!(after_flush, Some("12".to_owned())); +} + +// BaseRoutingStrategy's periodic sync counts in memory first, then sends the queued increments +// through redis_cache.async_increment_pipeline alone +#[rstest] +#[tokio::test] +async fn a_remote_increment_answers_the_redis_totals_and_leaves_memory_alone( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY); + let [rpm, _] = deployment_usage(); + let queued = increment(&rpm.counter, Number::Int(2), COUNTER_TTL); + raw.send(commands::incr_by_float( + rpm.key, + Amount::Float(10.0), + Some(Expiry::Refresh(COUNTER_TTL)), + )) + .await + .unwrap(); + + let counted = store.increment_local(&queued).unwrap(); + let totals = store + .increment_remote(std::slice::from_ref(&queued)) + .await + .unwrap(); + + assert_eq!(counted, Number::Int(2)); + assert_eq!(totals, Some(vec![Number::Float(12.0)])); + assert_eq!(store.get_local(&rpm.counter).unwrap(), Some(Number::Int(2))); +} + +#[rstest] +#[tokio::test] +async fn a_remote_increment_is_none_without_redis_and_an_error_when_redis_fails() { + let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let detached = RouterUsage::new(None, POLICY); + let attached = RouterUsage::new(Some(remote(connect(&deployment))), POLICY); + let [rpm, _] = deployment_usage(); + let queued = [increment(&rpm.counter, Number::Int(1), COUNTER_TTL)]; + deployment.stop().await.unwrap(); + + assert_eq!(detached.increment_remote(&queued).await.unwrap(), None); + assert!(attached.increment_remote(&queued).await.is_err()); +} + +#[rstest] +#[tokio::test] +async fn a_remote_increment_of_nothing_is_none_as_mains_pipeline_returns( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY); + + assert_eq!(store.increment_remote(&[]).await.unwrap(), None); +} + +#[rstest] +#[tokio::test] +async fn a_memory_count_past_i64_stays_above_the_limit_instead_of_wrapping_negative() { + let store = RouterUsage::new(None, POLICY); + let [tpm, _] = tpm_and_rpm(None); + store.increment(&tpm, Number::Int(1), COUNTER_TTL).await; + + assert_eq!( + store + .increment(&tpm, Number::Int(i64::MAX), COUNTER_TTL) + .await, + Some(Number::Float(9_223_372_036_854_775_808.0)) + ); +} + +#[rstest] +#[tokio::test] +async fn a_local_merge_replaces_the_memory_count_without_reaching_redis() { + let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let (client, trips) = support::connect_observed(&deployment); + let store = RouterUsage::new(Some(remote(client)), POLICY); + let [rpm, tpd] = deployment_usage(); + + store + .increment_local(&increment(&rpm.counter, Number::Int(1), COUNTER_TTL)) + .unwrap(); + store.set_local(&rpm.counter, Number::Float(9.0)).unwrap(); + + assert_eq!( + store.get_local_many(&[rpm.counter, tpd.counter]).unwrap(), + vec![Some(Number::Float(9.0)), None] + ); + assert!(trips.seen().is_empty()); +} + +#[rstest] +#[tokio::test] +async fn current_usage_skips_redis_once_memory_reaches_the_limit() { + let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let (client, trips) = support::connect_observed(&deployment); + let store = RouterUsage::new(Some(remote(client)), POLICY); + let [rpm, _] = deployment_usage(); + store + .increment_local(&increment(&rpm.counter, Number::Int(5), COUNTER_TTL)) + .unwrap(); + + let at_limit = store.current(&rpm.counter, Number::Int(5)).await.unwrap(); + let skipped = trips.seen(); + let below_limit = store.current(&rpm.counter, Number::Int(6)).await.unwrap(); + + assert_eq!(at_limit, Some(Number::Int(5))); + assert!(skipped.is_empty()); + assert_eq!(below_limit, None, "Redis holds no count, and Redis wins"); + assert_eq!(trips.seen(), [Operation::Command("GET")]); +} + +#[rstest] +#[tokio::test] +async fn current_usage_reads_none_on_a_failure_and_memory_once_the_breaker_opens() { + let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let breaker = litellm_circuit_breaker::Breaker::new( + litellm_circuit_breaker::Config { + failure_threshold: 1, + recovery_timeout: Duration::from_secs(60), + ..litellm_circuit_breaker::Config::default() + }, + Box::new(|_| {}), + ); + let store = RouterUsage::new( + Some(remote(support::connect_guarded(&deployment, breaker))), + POLICY, + ); + let [rpm, _] = deployment_usage(); + store + .increment_local(&increment(&rpm.counter, Number::Int(2), COUNTER_TTL)) + .unwrap(); + deployment.stop().await.unwrap(); + + let failed = store.current(&rpm.counter, Number::Int(10)).await.unwrap(); + let refused = store.current(&rpm.counter, Number::Int(10)).await.unwrap(); + + assert_eq!(failed, None); + assert_eq!(refused, Some(Number::Int(2))); +} + +#[rstest] +#[tokio::test] +async fn current_usage_without_redis_is_the_memory_count() { + let store = RouterUsage::new(None, POLICY); + let [rpm, _] = deployment_usage(); + store + .increment_local(&increment(&rpm.counter, Number::Int(3), COUNTER_TTL)) + .unwrap(); + + assert_eq!( + store.current(&rpm.counter, Number::Int(10)).await.unwrap(), + Some(Number::Int(3)) + ); +} + #[rstest] #[case::unknown("tokens")] #[case::wrong_case("TPM")] @@ -327,3 +717,69 @@ async fn a_failed_deployment_pipeline_keeps_the_memory_counts_as_integers() { [Some(Number::Int(5)), Some(Number::Int(1))] ); } + +// The io-token check stashes a counter's fields and the store spells the key, so the key must stay +// main's RouterCacheEnum.ITPM / OTPM spelling +#[rstest] +#[case::itpm(Kind::Itpm, "global_router:dep-1:bedrock/claude:itpm:12-34")] +#[case::otpm(Kind::Otpm, "global_router:dep-1:bedrock/claude:otpm:12-34")] +#[tokio::test] +async fn io_token_counters_increment_under_mains_keys( + #[values(Topology::Standalone, Topology::Cluster)] topology: Topology, + #[case] kind: Kind, + #[case] key: &str, +) { + let deployment = RedisDeployment::start(topology).await.unwrap(); + let raw = connect(&deployment); + let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY); + let counter = Counter { + scope: Scope::GlobalRouter, + deployment_id: "dep-1".to_owned(), + model: Some("bedrock/claude".to_owned()), + window: "12-34".to_owned(), + kind, + }; + + store + .increment(&counter, Number::Int(8), COUNTER_TTL) + .await + .unwrap(); + + let written: Option = raw.send(commands::get(key)).await.unwrap(); + assert_eq!(written, Some("8".to_owned())); +} + +#[rstest] +#[tokio::test] +async fn shared_usage_prefers_what_every_pod_counted_and_falls_back_to_memory() { + let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap(); + let raw = connect(&deployment); + let store = RouterUsage::new(Some(remote(connect(&deployment))), POLICY); + let [rpm, tpd] = deployment_usage(); + store + .increment_local(&increment(&rpm.counter, Number::Int(2), COUNTER_TTL)) + .unwrap(); + store + .increment_local(&increment(&tpd.counter, Number::Int(2), COUNTER_TTL)) + .unwrap(); + raw.send(commands::incr_by_float( + rpm.key, + Amount::Float(9.0), + Some(Expiry::Refresh(COUNTER_TTL)), + )) + .await + .unwrap(); + + let counted = store.shared_or_local(&rpm.counter).await.unwrap(); + let absent = store.shared_or_local(&tpd.counter).await.unwrap(); + deployment.stop().await.unwrap(); + let failed = store.shared_or_local(&rpm.counter).await.unwrap(); + + assert_eq!(counted, Some(Number::Int(9))); + assert_eq!( + absent, + Some(Number::Int(2)), + "Redis has no count, so memory answers" + ); + assert_eq!(failed, Some(Number::Int(2))); +} diff --git a/litellm/rust/crates/redis-store/tests/support/mod.rs b/litellm/rust/crates/redis-store/tests/support/mod.rs index f4b05c304e3..68966986d07 100644 --- a/litellm/rust/crates/redis-store/tests/support/mod.rs +++ b/litellm/rust/crates/redis-store/tests/support/mod.rs @@ -13,7 +13,7 @@ use litellm_redis_client::{ use litellm_redis_store::{ Number, Remote, TierPolicy, router_cooldowns::Cooldown, - router_usage::{Counter, Kind}, + router_usage::{Counter, Kind, Scope}, }; use litellm_redis_testing::RedisDeployment; @@ -212,18 +212,20 @@ pub fn usage() -> [UsageCase; 2] { [ UsageCase { counter: Counter { + scope: Scope::GlobalRouter, deployment_id: "dep".to_owned(), model: Some("anthropic/claude-sonnet-5-5".to_owned()), - minute: "12-34".to_owned(), + window: "12-34".to_owned(), kind: Kind::Tpm, }, key: "global_router:dep:anthropic/claude-sonnet-5-5:tpm:12-34", }, UsageCase { counter: Counter { + scope: Scope::GlobalRouter, deployment_id: "dep".to_owned(), model: None, - minute: "09-05".to_owned(), + window: "09-05".to_owned(), kind: Kind::Rpm, }, key: "global_router:dep:None:rpm:09-05", @@ -231,6 +233,32 @@ pub fn usage() -> [UsageCase; 2] { ] } +/// main's usage-based-routing-v2 and model rate limit check keys, which carry no prefix. +pub fn deployment_usage() -> [UsageCase; 2] { + [ + UsageCase { + counter: Counter { + scope: Scope::Deployment, + deployment_id: "dep".to_owned(), + model: Some("openai/gpt-5.4-nano".to_owned()), + window: "12-34".to_owned(), + kind: Kind::Rpm, + }, + key: "dep:openai/gpt-5.4-nano:rpm:12-34", + }, + UsageCase { + counter: Counter { + scope: Scope::Deployment, + deployment_id: "dep".to_owned(), + model: None, + window: "2026-11-10".to_owned(), + kind: Kind::Tpd, + }, + key: "dep:None:tpd:2026-11-10", + }, + ] +} + /// A TTL Redis reports for a key written `whole_seconds` ago at most a few seconds back. pub fn within(ttl: Duration, whole_seconds: u64) -> bool { ttl <= Duration::from_secs(whole_seconds) diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index baeb12476ed..e91d64ef0d6 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -230,14 +230,30 @@ class NativeSensitiveRoutePins: def set(self, tenant: str, session_id: str, model: str, ttl_seconds: float) -> Future[None]: ... class _UsageCounter(Protocol): + @property + def scope(self) -> Literal["global_router", "deployment"]: ... @property def deployment_id(self) -> str: ... @property def model(self) -> str | None: ... @property - def minute(self) -> str: ... + def window(self) -> str: ... @property - def kind(self) -> Literal["tpm", "rpm"]: ... + def kind(self) -> Literal["tpm", "rpm", "itpm", "otpm", "tpd", "rpd"]: ... + +class _CounterIncrement(Protocol): + @property + def counter(self) -> _UsageCounter: ... + @property + def amount(self) -> float: ... + @property + def ttl(self) -> float | None: ... + +@final +class NativeUsageRead: + @property + def read(self) -> NativePending[list[float | None]]: ... + def release(self) -> None: ... @final class NativeRouterUsage: @@ -245,7 +261,7 @@ class NativeRouterUsage: def attach(self, client: NativeRedisClient) -> bool: ... @property def has_redis(self) -> bool: ... - def increment(self, counter: _UsageCounter, value: float, ttl_seconds: int) -> Future[float]: ... + def increment(self, counter: _UsageCounter, value: float, ttl_seconds: int) -> Future[float | None]: ... def increment_deployment( self, deployment_id: str, @@ -255,7 +271,7 @@ class NativeRouterUsage: requests: int, ) -> Future[None]: ... def get_many(self, counters: Sequence[_UsageCounter]) -> Future[list[float | None]]: ... - def increment_sync(self, counter: _UsageCounter, value: float, ttl_seconds: int) -> float: ... + def increment_sync(self, counter: _UsageCounter, value: float, ttl_seconds: int) -> float | None: ... def increment_deployment_sync( self, deployment_id: str, @@ -265,6 +281,17 @@ class NativeRouterUsage: requests: int, ) -> None: ... def get_many_sync(self, counters: Sequence[_UsageCounter]) -> list[float | None]: ... + def usage(self, counters: Sequence[_UsageCounter], send: NativeSend) -> NativeUsageRead: ... + def increment_pipeline(self, increments: Sequence[_CounterIncrement]) -> Future[list[float] | None]: ... + def increment_post_call(self, request: NativeRequestRedis, increments: Sequence[_CounterIncrement]) -> None: ... + def increment_remote(self, increments: Sequence[_CounterIncrement]) -> Future[list[float] | None]: ... + def increment_local(self, increment: _CounterIncrement) -> float: ... + def get_local(self, counter: _UsageCounter) -> float | None: ... + def get_local_many(self, counters: Sequence[_UsageCounter]) -> list[float | None]: ... + def set_local(self, counter: _UsageCounter, value: float) -> None: ... + def current(self, counter: _UsageCounter, limit: float) -> Future[float | None]: ... + def current_sync(self, counter: _UsageCounter, limit: float) -> float | None: ... + def shared_or_local(self, counter: _UsageCounter) -> Future[float | None]: ... def flush_local_counters(self) -> None: ... def snapshot(self) -> dict[str, float]: ... @@ -559,6 +586,7 @@ def messages( def amessages( call: NativeCall, ) -> Execution: ... + @final class ResponsesWebSocketConnection: def __new__(cls, _uninstantiable: Never, /) -> Never: ... @@ -719,6 +747,7 @@ __all__ = [ "NativeSensitiveRoutePins", "NativeSpendCounters", "NativeTableSwapped", + "NativeUsageRead", "PostCallSettlement", "ProcessReservedForForking", "ResponsesWebSocketConnection", @@ -1197,7 +1226,6 @@ class RouterCoreHandle: @staticmethod def configured_service_tiers(model_infos: list[object]) -> list[object]: ... # The routing-cache key a deployment counts `metric` (tpm, rpm, itpm, otpm) under in `minute`. - def usage_key(self, id: str, model: str, metric: str, minute: str) -> str: ... def registry_version(self) -> int: ... # The version the deployments before the newest appends last changed at. def registry_rewritten(self) -> int: ... diff --git a/litellm/rust_bridge/host/redis_events.py b/litellm/rust_bridge/host/redis_events.py index ff94378be32..828d8936b71 100644 --- a/litellm/rust_bridge/host/redis_events.py +++ b/litellm/rust_bridge/host/redis_events.py @@ -4,13 +4,14 @@ from __future__ import annotations import asyncio from collections.abc import Callable, Coroutine +from contextvars import Context, copy_context from functools import lru_cache from typing import Final, Literal, TypeAlias from redis import exceptions as redis_exceptions from typing_extensions import assert_never -from litellm._internal_context import current_request_parent_otel_span +from litellm._internal_context import current_request_parent_otel_span, service_target from litellm._logging import verbose_logger from litellm._service_logger import ServiceLogging from litellm.types.services import ServiceTypes @@ -74,6 +75,14 @@ def _emit(event_loop: asyncio.AbstractEventLoop | None, hook: Hook) -> None: event_loop.call_soon_threadsafe(_spawn, hook) +def deadline_flush_context() -> Context: + """The context a request's post-call batch is reported in when its deadline sends it. It is the declaring caller's, + so the flush parents under that request's span, without the declaring store's target, since the flush that + closes the callbacks runs under none.""" + with service_target(None): + return copy_context() + + def record_redis_call( event_loop: asyncio.AbstractEventLoop | None, call_type: str, diff --git a/litellm/types/router.py b/litellm/types/router.py index 40b696887d7..6cc38a30d6f 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -1314,13 +1314,6 @@ class RoutingStrategy(enum.Enum): PROVIDER_BUDGET_LIMITING = "provider-budget-routing" -class RouterCacheEnum(enum.Enum): - TPM = "global_router:{id}:{model}:tpm:{current_minute}" - RPM = "global_router:{id}:{model}:rpm:{current_minute}" - ITPM = "global_router:{id}:{model}:itpm:{current_minute}" - OTPM = "global_router:{id}:{model}:otpm:{current_minute}" - - class GenericBudgetWindowDetails(LiteLLMBaseModel): """Details about a provider's budget window""" diff --git a/tests/unit/caching/test_request_redis_batch_post_call.py b/tests/unit/caching/test_request_redis_batch_post_call.py index 01feee0be54..e2ddda29944 100644 --- a/tests/unit/caching/test_request_redis_batch_post_call.py +++ b/tests/unit/caching/test_request_redis_batch_post_call.py @@ -36,10 +36,12 @@ from litellm.proxy.hooks.parallel_request_limiter_v3 import ( from litellm.proxy.spend_tracking.spend_counter_batch import PendingSpendIncrement from litellm.proxy.utils import InternalUsageCache from litellm.redis.stores.rate_limits import CounterIncrement, RateLimitScope +from litellm.redis.stores.router_usage import CounterIncrement as UsageIncrement from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import ModelResponse from tests.unit.redis.stores.fake_rate_limits import FakeRateLimits +from tests.unit.redis.stores.router_usage_double import FakeUsageHandle, usage_store from .test_redis_batch import FakeClient, FakeNatives, FakeRedisCache, fake_send @@ -132,10 +134,10 @@ def _response_cache(redis_cache: FakeRedisCache) -> Cache: return cache -def _tpm_router(redis_cache: FakeRedisCache) -> tuple[LowestTPMLoggingHandler_v2, DualCache]: - router_cache = DualCache() - router_cache.attach_redis_cache(redis_cache) - return LowestTPMLoggingHandler_v2(router_cache=router_cache, routing_args={"ttl": 60}), router_cache +def _tpm_selector(handle: FakeUsageHandle) -> LowestTPMLoggingHandler_v2: + selector: Final = LowestTPMLoggingHandler_v2(usage=usage_store(handle), routing_args={"ttl": 60}) + selector.cancel_sync_task() + return selector def _tpm_kwargs() -> Mapping[str, object]: @@ -151,19 +153,19 @@ def _tpm_kwargs() -> Mapping[str, object]: @pytest.mark.asyncio -async def test_post_call_owners_wait_for_the_callbacks_and_the_limiter_declares_on_the_native_post_call_batch(): +async def test_post_call_owners_declare_on_the_native_post_call_batch_until_the_callbacks_finish(): client = FakeClient(_ok_replies) redis_cache = PostCallFakeRedisCache(client) calls: list[tuple[str, list[str], list[object]]] = [] # mutable-ok: filled by the fake scripts fakes = FakeRateLimits(token_increment_script=_scripts(calls, [10.0]), parallel_release_script=_scripts(calls, [0])) limiter = _limiter(fakes) - tpm, router_cache = _tpm_router(redis_cache) + usage = FakeUsageHandle(redis={}) + tpm = _tpm_selector(usage) with request_redis_batch_scope(): await tpm.async_log_success_event(_tpm_kwargs(), None, None, None) await limiter.async_increment_tokens_with_ttl_preservation(_token_ops(_K1)) await limiter._release_stashed_parallel_slot(_slot_stash("slot-1", _K1), None, in_logging_callback=True) - assert client.pipelines == [] # nothing goes out while the callbacks are still declaring await flush_post_call_redis_batches() assert fakes.post_call == ["token_increment", "parallel_release"] @@ -171,15 +173,22 @@ async def test_post_call_owners_wait_for_the_callbacks_and_the_limiter_declares_ ("script", ["{api_key:k1}:tokens"], [10, 60]), ("script", ["{api_key:k1}:max_parallel_requests"], ["slot-1"]), ] - assert len(client.pipelines) == 1 - assert _names(client) == ["INCRBYFLOAT", "EXPIRE"] + tpm_counter = next(iter(usage.local)) + assert usage.calls == [("post_call", (UsageIncrement(tpm_counter, 42, 60),))] + assert (tpm_counter.deployment_id, tpm_counter.model, tpm_counter.kind) == ("dep-a", "openai/gpt-4o-mini", "tpm") + assert usage.local == {tpm_counter: 42} + assert client.pipelines == [] assert redis_cache.alone == [] - assert ( - await router_cache.in_memory_cache.async_get_cache( - next(k for k in router_cache.in_memory_cache.cache_dict if ":tpm:" in k) - ) - == 42 - ) + + +@pytest.mark.asyncio +async def test_outside_a_request_the_tpm_increment_goes_out_on_its_own_pipeline(): + usage = FakeUsageHandle(redis={}) + + await _tpm_selector(usage).async_log_success_event(_tpm_kwargs(), None, None, None) + + assert [name for name, _ in usage.calls] == ["pipeline"] + assert usage.local == usage.redis == {next(iter(usage.local)): 42} @pytest.mark.asyncio diff --git a/tests/unit/caching/test_request_redis_batch_pre_call.py b/tests/unit/caching/test_request_redis_batch_pre_call.py index b4268fe7bb0..b8e7658ef24 100644 --- a/tests/unit/caching/test_request_redis_batch_pre_call.py +++ b/tests/unit/caching/test_request_redis_batch_pre_call.py @@ -4,16 +4,14 @@ router's usage counter read, auth identity and spend counters all join the reque from __future__ import annotations import asyncio -import json from collections.abc import Callable, Sequence -from dataclasses import dataclass +from dataclasses import dataclass, field from datetime import datetime, timezone -from typing import Any, Final -from unittest.mock import AsyncMock, MagicMock +from typing import Any, Final, Literal +from unittest.mock import AsyncMock, MagicMock, patch import pytest -import litellm.caching.dual_cache as dual_cache_module from litellm import Router from litellm._internal_context import current_service_target from litellm.caching.dual_cache import DualCache @@ -30,12 +28,15 @@ from litellm.proxy.utils import InternalUsageCache from litellm.redis.stores._route import route from litellm.redis.stores.rate_limits import RateLimitScope from litellm.redis.stores.router_cooldowns import ActiveCooldowns, RouterCooldownStore -from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET +from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, RouterUsageStore, UsageCounter from litellm.router_api import RoutingContext -from litellm.router_utils.routing_prefetch import ArmedRead, TierRead, arm_cooldowns, read_cooldowns -from litellm.router_utils.routing_read_batch import RoutingPrefetch +from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2, PrefetchedUsage +from litellm.router_utils import routing_prefetch +from litellm.router_utils.routing_prefetch import arm_cooldowns, read_cooldowns +from litellm.router_utils.routing_read_batch import RoutingReadBatch from tests.unit.redis.stores.auth_cache_double import FakeAuthCache, shared_cache from tests.unit.redis.stores.fake_rate_limits import FakeRateLimits +from tests.unit.redis.stores.router_usage_double import FakeUsageRead from .test_redis_batch import ( FakeClient, @@ -43,6 +44,7 @@ from .test_redis_batch import ( FakeNativeRequest, FakeNatives, FakeRedisCache, + fake_pending, fake_send, ) @@ -323,43 +325,8 @@ def _router(redis_cache: FakeRedisCache, routing_strategy: str = "usage-based-ro return router -_PINNED_MINUTE: Final = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc) -_START: Final = 1_000_000.0 _PING: Final = [{"role": "user", "content": "ping"}] - - -@pytest.fixture -def pinned_clock(monkeypatch: pytest.MonkeyPatch) -> list[float]: - """Pins the usage counters' minute and the router cache's read throttle clock; tests move `now[0]`.""" - now: Final = [_START] # mutable-ok: the test advances the pinned clock - monkeypatch.setattr("litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime", lambda: _PINNED_MINUTE) - monkeypatch.setattr(dual_cache_module.time, "time", lambda: now[0]) - return now - - -def _counter_keys(router: Router, deployment_ids: tuple[str, ...] = ("dep-a", "dep-b")) -> frozenset[str]: - deployments: Final = [d for d in router.model_list if d["model_info"]["id"] in deployment_ids] - tpm, rpm = router.lowesttpm_logger_v2.usage_counter_keys(deployments) - return frozenset(tpm + rpm) - - -def _tpm_key(router: Router, deployment_id: str) -> str: - deployment: Final = next(d for d in router.model_list if d["model_info"]["id"] == deployment_id) - return router.lowesttpm_logger_v2.usage_counter_keys([deployment])[0][0] - - -def _counter_mgets(router: Router, client: FakeClient, redis_cache: FakeRedisCache) -> tuple[frozenset[str], ...]: - counters: Final = _counter_keys(router) - pipelined: Final = tuple( - frozenset(command[1:]) - for pipeline in client.pipelines - for command in pipeline.commands - if command[0] == "MGET" and counters.intersection(command[1:]) - ) - alone: Final = tuple( - frozenset(keys) for command, keys in redis_cache.alone if command == "MGET" and counters.intersection(keys) - ) - return pipelined + alone +_PINNED_MINUTE: Final = datetime(2026, 1, 1, 12, 0, 30, tzinfo=timezone.utc) def _python_cooldown_reads(client: FakeClient, redis_cache: FakeRedisCache) -> tuple[str, ...]: @@ -374,15 +341,6 @@ def _python_cooldown_reads(client: FakeClient, redis_cache: FakeRedisCache) -> t return tuple(key for key in pipelined + alone if key.endswith(":cooldown")) -def _replies_with(values: dict[str, object]) -> Callable[[tuple[Any, ...]], Any]: - def replies(command: tuple[Any, ...]) -> Any: - if command[0] == "MGET": - return [None if values.get(key) is None else json.dumps(values[key]) for key in command[1:]] - return _lua_ok_replies(command) - - return replies - - async def _route(router: Router) -> str: deployment: Final = await router.async_get_available_deployment( model=_MODEL_GROUP, messages=_PING, request_kwargs={} @@ -390,23 +348,139 @@ async def _route(router: Router) -> str: return deployment["model_info"]["id"] +def _counter(deployment_id: str, kind: Literal["tpm", "rpm"]) -> UsageCounter: + return UsageCounter.for_deployment(deployment_id, "anthropic/claude-x", kind, "12-00") + + +_DEPLOYMENT_COUNTERS: Final = tuple(_counter(d, kind) for kind in ("tpm", "rpm") for d in ("dep-a", "dep-b")) + + +@dataclass +class _FakeUsageHandle: + """A connected usage store whose ``usage`` reads land on the request's fake native batch, or alone for + ``NativeSend.Now()``, recording what each read asked for and under which service target.""" + + counts: dict[UsageCounter, float] + fail: Exception | None = None + reads: list[tuple[str, tuple[UsageCounter, ...], str | None]] = field(default_factory=list) + has_redis: bool = True + + def usage(self, counters: Sequence[UsageCounter], send: object) -> FakeUsageRead: + batched: Final = isinstance(send, FakeNativeRequest) + self.reads.append(("batch" if batched else "alone", tuple(counters), current_service_target())) + if batched: + send.events.append("declare usage") + + fail: Final = self.fail + + def answer() -> list[float | None]: + if fail is not None: + raise fail + return [self.counts.get(counter) for counter in counters] + + return FakeUsageRead( + fake_pending(send, answer), + lambda: self.reads.append(("release", tuple(counters), current_service_target())), + ) + + def get_many_sync(self, counters: Sequence[UsageCounter]) -> list[float | None]: + raise AssertionError("a connected store reads through usage()") + + +def _usage_store(handle: _FakeUsageHandle) -> RouterUsageStore: + return RouterUsageStore(handle) # pyright: ignore[reportArgumentType] # test double + + @pytest.mark.asyncio -@pytest.mark.usefixtures("pinned_clock") -async def test_armed_routing_read_is_one_pipeline_and_routing_issues_no_read_of_its_own(): - client = FakeClient(_lua_ok_replies) - redis_cache = FakeRedisCache(client) - router = _router(redis_cache) +async def test_the_armed_usage_read_and_the_armed_cooldown_read_go_out_in_one_native_pipeline(): + handle: Final = _FakeUsageHandle({_counter("dep-a", "tpm"): 40}) + usage: Final = _usage_store(handle) + cooldowns: Final = RouterCooldownStore(_FakeCooldownsHandle()) + natives: Final = FakeNatives() - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - deployment = await _route(router) + with request_redis_batch_scope(native_factory=natives, native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double + arm_cooldowns(cooldowns, ["dep-a", "dep-b"]) + routing_prefetch.arm(usage, _DEPLOYMENT_COUNTERS) + active, counts = await asyncio.gather( + read_cooldowns(cooldowns, ["dep-a", "dep-b"], None), + routing_prefetch.read(usage, _DEPLOYMENT_COUNTERS, None), + ) - assert deployment in {"dep-a", "dep-b"} - assert len(client.pipelines) == 1 - commands = client.pipelines[0].commands - assert [c[0] for c in commands] == ["MGET"] - assert frozenset(commands[0][1:]) == _counter_keys(router) - assert redis_cache.alone == [] + assert active == () + assert counts == [40, None, None, None] + assert _declared_before_the_first_flush(natives) == ["declare cooldowns", "declare usage"] + assert [kind for kind, _, _ in handle.reads] == ["batch"] + assert request.prefetched == {} + + +@pytest.mark.asyncio +async def test_the_armed_usage_read_is_declared_under_the_router_usage_family(): + handle: Final = _FakeUsageHandle({}) + + with request_redis_batch_scope(native_factory=FakeNatives(), native_send_factory=fake_send): # pyright: ignore[reportArgumentType] # test double + routing_prefetch.arm(_usage_store(handle), _DEPLOYMENT_COUNTERS) + + assert [target for _, _, target in handle.reads] == [ROUTER_USAGE_TARGET] + + +@pytest.mark.asyncio +async def test_a_usage_prefetch_that_does_not_cover_the_pick_gives_its_reservation_back_before_the_store_reads(): + handle: Final = _FakeUsageHandle({_counter("dep-b", "rpm"): 3}) + usage: Final = _usage_store(handle) + + with request_redis_batch_scope(native_factory=FakeNatives(), native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double + routing_prefetch.arm(usage, _DEPLOYMENT_COUNTERS[:1]) + counts: Final = await routing_prefetch.read(usage, _DEPLOYMENT_COUNTERS, None) + assert request.prefetched == {} + + assert counts == [None, None, None, 3] + assert [(kind, len(counters)) for kind, counters, _ in handle.reads] == [("batch", 1), ("release", 1), ("alone", 4)] + + +@pytest.mark.asyncio +async def test_a_failed_usage_prefetch_falls_back_to_the_stores_own_read(): + handle: Final = _FakeUsageHandle({_counter("dep-a", "tpm"): 9}, fail=ConnectionError("redis down")) + usage: Final = _usage_store(handle) + + with request_redis_batch_scope(native_factory=FakeNatives(), native_send_factory=fake_send): # pyright: ignore[reportArgumentType] # test double + routing_prefetch.arm(usage, _DEPLOYMENT_COUNTERS) + handle.fail = None + counts: Final = await routing_prefetch.read(usage, _DEPLOYMENT_COUNTERS, None) + + assert counts == [9, None, None, None] + assert [kind for kind, _, _ in handle.reads] == ["batch", "alone"] + + +@pytest.mark.asyncio +async def test_arming_a_store_without_redis_declares_nothing(): + handle: Final = _FakeUsageHandle({}, has_redis=False) + + with request_redis_batch_scope(native_factory=FakeNatives(), native_send_factory=fake_send) as request: # pyright: ignore[reportArgumentType] # test double + routing_prefetch.arm(_usage_store(handle), _DEPLOYMENT_COUNTERS) + + assert request.prefetched == {} + assert handle.reads == [] + + +@pytest.mark.asyncio +async def test_a_usage_based_pick_reads_the_counters_the_prefetch_armed(): + handle: Final = _FakeUsageHandle({_counter("dep-a", "tpm"): 100}) + selector: Final = LowestTPMLoggingHandler_v2(usage=_usage_store(handle)) + deployments: Final[list[object]] = [_deployment("dep-a"), _deployment("dep-b")] + read_batch: Final = RoutingReadBatch(usage_selector=selector) + + with ( + patch("litellm.router_strategy.lowest_tpm_rpm_v2.get_utc_datetime", return_value=_PINNED_MINUTE), + request_redis_batch_scope(native_factory=FakeNatives(), native_send_factory=fake_send), # pyright: ignore[reportArgumentType] # test double + ): + routing_prefetch.arm(selector.usage, _DEPLOYMENT_COUNTERS) + await read_batch.read_usage(deployments, None) + with PrefetchedUsage.scoped(read_batch.prefetched_usage): + picked: Final = await selector.async_get_available_deployments(_MODEL_GROUP, deployments, messages=_PING) + + selector.cancel_sync_task() + assert picked["model_info"]["id"] == "dep-b" + assert [kind for kind, _, _ in handle.reads] == ["batch"] @pytest.mark.asyncio @@ -431,158 +505,6 @@ async def test_a_cooldown_recorded_through_the_store_excludes_its_deployment_wit assert _python_cooldown_reads(client, redis_cache) == () -@pytest.mark.asyncio -@pytest.mark.usefixtures("pinned_clock") -async def test_a_prefetch_that_does_not_cover_the_routing_keys_is_ignored_and_routing_reads_itself(): - client = FakeClient(_lua_ok_replies) - redis_cache = FakeRedisCache(client) - router = _router(redis_cache) - - with request_redis_batch_scope() as request: - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - armed = request.prefetched["routing_read"] - assert isinstance(armed, ArmedRead) - request.prefetched["routing_read"] = ArmedRead( - read=TierRead(armed.read.cache, ("other",)), - fetched=armed.fetched, - result=armed.result, - previous=armed.previous, - ) - deployment = await _route(router) - assert request.prefetched == {} - - assert deployment in {"dep-a", "dep-b"} - assert [frozenset(keys) for command, keys in redis_cache.alone if command == "MGET"] == [_counter_keys(router)] - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("pinned_clock") -async def test_a_prefetch_with_incomplete_usage_keys_releases_its_reservations(): - client: Final = FakeClient(_lua_ok_replies) - redis_cache: Final = FakeRedisCache(client) - router: Final = _router(redis_cache) - - with request_redis_batch_scope(): - RoutingPrefetch.arm(RoutingContext(router), router.lowesttpm_logger_v2, router.model_list[:1]) - deployment: Final = await _route(router) - - assert deployment in {"dep-a", "dep-b"} - assert [frozenset(keys) for command, keys in redis_cache.alone if command == "MGET"] == [_counter_keys(router)] - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("pinned_clock") -async def test_a_failed_prefetch_falls_back_to_the_selectors_own_read(): - client = FakeClient(_lua_ok_replies, fail=ConnectionError("redis down")) - redis_cache = FakeRedisCache(client) - router = _router(redis_cache) - - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - deployment = await _route(router) - - assert deployment in {"dep-a", "dep-b"} - assert [frozenset(keys) for command, keys in redis_cache.alone if command == "MGET"] == [_counter_keys(router)] - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("pinned_clock") -async def test_the_armed_routing_read_is_declared_under_the_router_usage_family(): - """The prefetch is declared before routing runs under a target of its own, so the pipeline that - carries it renders ``redis.pipeline router_usage`` instead of a bare ``redis.pipeline``.""" - client = FakeClient(_lua_ok_replies) - redis_cache = FakeRedisCache(client) - router = _router(redis_cache) - pipeline_targets: list[str | None] = [] # mutable-ok: filled by the recording hook - - async def record(**kwargs: object) -> None: - pipeline_targets.append(current_service_target()) - - redis_cache.service_logger_obj.async_service_success_hook = record # pyright: ignore[reportAttributeAccessIssue] # fake, records the hook call - - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - await _route(router) - async with asyncio.timeout(5): - while not pipeline_targets: - await asyncio.sleep(0.01) - - assert pipeline_targets == [ROUTER_USAGE_TARGET] - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("pinned_clock") -async def test_an_abandoned_prefetch_still_backfills_the_counters_it_read(): - client: Final = FakeClient(_lua_ok_replies) - redis_cache: Final = FakeRedisCache(client) - router: Final = _router(redis_cache) - client.reply_for = _replies_with({_tpm_key(router, "dep-a"): 100}) - - with request_redis_batch_scope() as request: - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - await request.flush_all() - after_abandoned: Final = len(_counter_mgets(router, client, redis_cache)) - - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - deployment: Final = await _route(router) - - assert deployment == "dep-b" - assert after_abandoned == 1 - assert len(_counter_mgets(router, client, redis_cache)) == 1 - assert redis_cache.alone == [] - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("pinned_clock") -async def test_prefetch_settlement_keeps_newer_memory_values_and_backfills_misses(): - client: Final = FakeClient(_lua_ok_replies) - redis_cache: Final = FakeRedisCache(client) - router: Final = _router(redis_cache) - dep_a: Final = _tpm_key(router, "dep-a") - dep_b: Final = _tpm_key(router, "dep-b") - client.reply_for = _replies_with({dep_a: 5, dep_b: 7}) - usage_cache: Final = router.lowesttpm_logger_v2.router_cache - - with request_redis_batch_scope() as request: - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - usage_cache.in_memory_cache.set_cache(dep_a, 50) - await request.flush_all() - - assert await usage_cache.async_batch_get_cache([dep_a, dep_b], local_only=True) == [50, 7] - - -@pytest.mark.asyncio -async def test_an_abandoned_prefetch_whose_mget_fails_releases_its_reservation(pinned_clock: list[float]): - mget_replies: Final = iter((ConnectionError("redis down"), None)) - client: Final = FakeClient(_lua_ok_replies) - redis_cache: Final = FakeRedisCache(client) - router: Final = _router(redis_cache) - - def replies(command: tuple[Any, ...]) -> Any: - if command[0] == "MGET": - failure: Final = next(mget_replies) - if isinstance(failure, Exception): - return failure - return ["100" if key == _tpm_key(router, "dep-a") else None for key in command[1:]] - return _lua_ok_replies(command) - - client.reply_for = replies - - with request_redis_batch_scope() as request: - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - await request.flush_all() - after_failure: Final = len(_counter_mgets(router, client, redis_cache)) - - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - deployment: Final = await _route(router) - - assert after_failure == 1 - assert len(_counter_mgets(router, client, redis_cache)) == 2 - assert deployment == "dep-b" - - @pytest.mark.asyncio async def test_arming_outside_a_request_scope_is_a_no_op(): redis_cache = FakeRedisCache(FakeClient(_lua_ok_replies)) @@ -607,96 +529,6 @@ async def test_simple_shuffle_cooldown_prefetch_does_not_arm_usage_reads_or_read assert redis_cache.alone == [] -@pytest.mark.asyncio -@pytest.mark.usefixtures("pinned_clock") -async def test_requests_within_the_read_interval_read_the_counters_from_redis_once(): - client = FakeClient(_lua_ok_replies) - redis_cache = FakeRedisCache(client) - router = _router(redis_cache) - request_round_trips: list[tuple[int, int]] = [] # mutable-ok: one entry per simulated request - - for _ in range(3): - pipeline_count = len(client.pipelines) - alone_count = len(redis_cache.alone) - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - await _route(router) - request_round_trips.append((len(client.pipelines) - pipeline_count, len(redis_cache.alone) - alone_count)) - - assert len(_counter_mgets(router, client, redis_cache)) == 1 - assert request_round_trips[1:] == [(0, 0), (0, 0)] - - -@pytest.mark.asyncio -@pytest.mark.usefixtures("pinned_clock") -async def test_concurrent_requests_share_one_counter_read_per_interval(): - client: Final = FakeClient(_lua_ok_replies) - redis_cache: Final = FakeRedisCache(client) - router: Final = _router(redis_cache) - first_armed: Final = asyncio.Event() - both_armed: Final = asyncio.Event() - - async def route_after_both_requests_arm() -> str: - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - if first_armed.is_set(): - both_armed.set() - else: - first_armed.set() - await both_armed.wait() - return await _route(router) - - deployments: Final = await asyncio.gather(route_after_both_requests_arm(), route_after_both_requests_arm()) - - assert set(deployments) <= {"dep-a", "dep-b"} - assert len(_counter_mgets(router, client, redis_cache)) == 1 - - -@pytest.mark.asyncio -async def test_the_prefetch_reads_the_counters_again_once_the_read_interval_elapses(pinned_clock: list[float]): - client = FakeClient(_lua_ok_replies) - redis_cache = FakeRedisCache(client) - router = _router(redis_cache) - - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - await _route(router) - - client.reply_for = _replies_with({_tpm_key(router, "dep-a"): 100}) - pinned_clock[0] = _START + router.lowesttpm_logger_v2.router_cache.redis_batch_cache_expiry + 1 - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - deployment = await _route(router) - - assert len(_counter_mgets(router, client, redis_cache)) == 2 - assert deployment == "dep-b" - - -@pytest.mark.asyncio -async def test_the_prefetch_mget_carries_only_the_keys_whose_read_is_due(pinned_clock: list[float]): - client = FakeClient(_lua_ok_replies) - redis_cache = FakeRedisCache(client) - router = _router(redis_cache) - usage_cache = router.lowesttpm_logger_v2.router_cache - - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - await _route(router) - - for key in _counter_keys(router, ("dep-a",)): - usage_cache.in_memory_cache.set_cache(key, 1) - pinned_clock[0] = _START + usage_cache.redis_batch_cache_expiry + 1 - with request_redis_batch_scope(): - router.arm_routing_read_prefetch(_MODEL_GROUP, {}) - await _route(router) - - second_pipeline_mgets = tuple(command for command in client.pipelines[1].commands if command[0] == "MGET") - - assert len(client.pipelines) == 2 - assert len(second_pipeline_mgets) == 1 - assert frozenset(second_pipeline_mgets[0][1:]) == _counter_keys(router, ("dep-b",)) - - @pytest.mark.asyncio async def test_two_backends_flush_concurrently_one_pipeline_each(): a_client, b_client = FakeClient(_lua_ok_replies), FakeClient(_lua_ok_replies) diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index 41cd17b72db..81166caa094 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -26,7 +26,7 @@ from openai._legacy_response import HttpxBinaryResponseContent from pydantic import BaseModel import litellm -from litellm._internal_context import in_post_response_phase +from litellm._internal_context import current_request_parent_otel_span, in_post_response_phase from litellm._logging import session_id_var, trace_id_var, verbose_logger from litellm._service_logger import ServiceLogging from litellm.caching.caching import DualCache @@ -40,6 +40,7 @@ from litellm.constants import ( ) from litellm.cost_calculator import ocr_batch_cost from litellm.integrations.custom_logger import CustomLogger +from litellm.llms.custom_llm import CustomLLM, CustomLLMError from litellm.litellm_core_utils import litellm_logging from litellm.litellm_core_utils.litellm_logging import Logging as LitellmLogging from litellm.litellm_core_utils.token_counter import token_counter @@ -78,7 +79,7 @@ from litellm.types.utils import ( from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome import inspect from typing import List, Optional -from litellm.types.utils import LiteLLMCommonStrings +from litellm.types.utils import GenericStreamingChunk, LiteLLMCommonStrings @pytest.mark.parametrize("metadata_key", ("metadata", "litellm_metadata")) @@ -1763,7 +1764,9 @@ class TestRetrieveBatchPricesOnlyFinalBatches: assert batch.usage.total_tokens == 35 @pytest.mark.parametrize("batch_ignore_default_logging", [True, False]) - def test_sync_logger_gets_a_free_event_for_the_cost_poller_retrieve(self, batch_ignore_default_logging: bool) -> None: + def test_sync_logger_gets_a_free_event_for_the_cost_poller_retrieve( + self, batch_ignore_default_logging: bool + ) -> None: sync_events: list[object] = [] class SyncOnlyLogger(CustomLogger): @@ -7345,7 +7348,12 @@ _WEB_SEARCH_REPLY: Final = ( ( [ {"type": "mcp_tool_use", "id": "mcptoolu_1", "name": "lookup", "server_name": "docs", "input": {}}, - {"type": "mcp_tool_result", "tool_use_id": "mcptoolu_1", "is_error": False, "content": [{"type": "text", "text": "ok"}]}, + { + "type": "mcp_tool_result", + "tool_use_id": "mcptoolu_1", + "is_error": False, + "content": [{"type": "text", "text": "ok"}], + }, {"type": "text", "text": "Found it."}, ], "end_turn", @@ -12879,7 +12887,12 @@ def test_get_model_cost_information_maps_requested_model_when_response_names_an_ monkeypatch.setitem( litellm.model_cost, requested_model, - {"litellm_provider": "fireworks_ai", "mode": "chat", "input_cost_per_token": 3e-6, "output_cost_per_token": 7e-6}, + { + "litellm_provider": "fireworks_ai", + "mode": "chat", + "input_cost_per_token": 3e-6, + "output_cost_per_token": 7e-6, + }, ) response: Final = ModelResponse(model="Sol Display Max", choices=[]) response._hidden_params["provider_response_model"] = "Sol Display Max" @@ -12891,15 +12904,27 @@ def test_get_model_cost_information_maps_requested_model_when_response_names_an_ assert result["model_map_key"] == requested_model assert result["model_map_value"] is not None - assert result["model_map_value"]["input_cost_per_token"] == litellm.model_cost[requested_model]["input_cost_per_token"] + assert ( + result["model_map_value"]["input_cost_per_token"] == litellm.model_cost[requested_model]["input_cost_per_token"] + ) @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") def test_get_model_cost_information_keeps_a_priced_response_model_that_echoes_without_its_provider( monkeypatch: pytest.MonkeyPatch, ) -> None: - served_row: Final = {"litellm_provider": "mistral", "mode": "chat", "input_cost_per_token": 2e-6, "output_cost_per_token": 5e-6} - requested_row: Final = {"litellm_provider": "mistral", "mode": "chat", "input_cost_per_token": 9e-6, "output_cost_per_token": 9e-6} + served_row: Final = { + "litellm_provider": "mistral", + "mode": "chat", + "input_cost_per_token": 2e-6, + "output_cost_per_token": 5e-6, + } + requested_row: Final = { + "litellm_provider": "mistral", + "mode": "chat", + "input_cost_per_token": 9e-6, + "output_cost_per_token": 9e-6, + } monkeypatch.setitem(litellm.model_cost, "mistral/sol-snapshot-test", served_row) monkeypatch.setitem(litellm.model_cost, "mistral/sol-alias-test", requested_row) response: Final = ModelResponse(model="sol-snapshot-test", choices=[]) @@ -14787,7 +14812,9 @@ async def test_a_redacted_batch_retrieve_still_tells_the_spend_callbacks_whether ((kwargs, response_obj),) = recorder.events assert not isinstance(response_obj, LiteLLMBatch) assert defers_batch_cost(kwargs, response_obj) is deferred - assert (job_to_charge_once(kwargs, response_obj, 0.42) == ChargeOnceJob(kind="batch", job_id="batch_abc")) is charged + assert ( + job_to_charge_once(kwargs, response_obj, 0.42) == ChargeOnceJob(kind="batch", job_id="batch_abc") + ) is charged _FAILED_ATTEMPT: Final = (Usage(prompt_tokens=1000, completion_tokens=500, total_tokens=1500), 0.5) @@ -15062,3 +15089,109 @@ async def test_failed_chat_prompt_is_counted_only_when_an_opted_in_logger_receiv assert counted_prompts == ([messages] if expects_estimate else []) assert logging_obj.model_call_details.get(FAILED_PROMPT_TOKEN_ESTIMATE_KEY) == (7 if expects_estimate else None) + + +class _SpanSeenByCallbacks(CustomLogger): + """Records the request span a datastore op made inside each hook would report under.""" + + def __init__(self) -> None: + super().__init__() + self.seen: dict[str, object] = {} + self.logged = threading.Event() + + def _record(self, hook: str) -> None: + self.seen[hook] = current_request_parent_otel_span() + self.logged.set() + + def log_success_event(self, kwargs, response_obj, start_time, end_time): + self._record("success") + + async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + self._record("success") + + def log_failure_event(self, kwargs, response_obj, start_time, end_time): + self._record("failure") + + async def async_log_failure_event(self, kwargs, response_obj, start_time, end_time): + self._record("failure") + + +_FIRST_CHUNK: Final[GenericStreamingChunk] = { + "text": "hel", + "tool_use": None, + "is_finished": False, + "finish_reason": "", + "usage": None, + "index": 0, +} + + +class _BreaksMidStream(CustomLLM): + def streaming(self, *args, **kwargs) -> Iterator[GenericStreamingChunk]: + yield _FIRST_CHUNK + raise CustomLLMError(status_code=500, message="stream broke") + + async def astreaming(self, *args, **kwargs) -> AsyncIterator[GenericStreamingChunk]: + yield _FIRST_CHUNK + raise CustomLLMError(status_code=500, message="stream broke") + + +_SCENARIOS: Final = MappingProxyType( + { + "success": {"model": "gpt-4o-mini", "mock_response": "hello"}, + "stream_success": {"model": "gpt-4o-mini", "mock_response": "hello", "stream": True}, + "failure": {"model": "gpt-4o-mini", "mock_response": Exception("provider down")}, + "mid_stream_failure": {"model": "span-probe/broken", "stream": True}, + } +) + + +async def _sdk_call(*, asynchronous: bool, scenario: str, span: object) -> None: + request: Final = { + **_SCENARIOS[scenario], + "messages": [{"role": "user", "content": "hi"}], + "metadata": {"litellm_parent_otel_span": span}, + } + with contextlib.suppress(Exception): + if asynchronous: + response = await litellm.acompletion(**request) + if request.get("stream"): + async for _ in response: + pass + return + response = litellm.completion(**request) + if request.get("stream"): + for _ in response: + pass + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [True, False], ids=["async", "sync"]) +@pytest.mark.parametrize( + ("scenario", "hook"), + [ + ("success", "success"), + ("stream_success", "success"), + ("failure", "failure"), + ("mid_stream_failure", "failure"), + ], +) +async def test_every_callback_runs_under_the_span_an_sdk_caller_passed( + monkeypatch: pytest.MonkeyPatch, asynchronous: bool, scenario: str, hook: str +) -> None: + """A store op in any callback parents under the caller's span, with no Router scope around the call, after a + stream ends or breaks, and on the executor threads the sync stream paths dispatch to without the caller's context.""" + span: Final = object() + logger: Final = _SpanSeenByCallbacks() + monkeypatch.setattr(litellm, "callbacks", [logger]) + monkeypatch.setattr( + litellm, "custom_provider_map", [{"provider": "span-probe", "custom_handler": _BreaksMidStream()}] + ) + monkeypatch.setattr(litellm, "provider_list", list(litellm.provider_list)) + monkeypatch.setattr(litellm, "_custom_providers", []) + litellm.utils.custom_llm_setup() + + await _sdk_call(asynchronous=asynchronous, scenario=scenario, span=span) + await asyncio.wait_for(asyncio.to_thread(logger.logged.wait, 10), timeout=10) + + assert (logger.seen.get(hook), current_request_parent_otel_span()) == (span, None) diff --git a/tests/unit/redis/stores/router_usage_double.py b/tests/unit/redis/stores/router_usage_double.py new file mode 100644 index 00000000000..1cf472776a3 --- /dev/null +++ b/tests/unit/redis/stores/router_usage_double.py @@ -0,0 +1,151 @@ +"""A stand-in for ``NativeRouterUsage`` with the native store's semantics: memory in front of a Redis dict. +``redis`` None is no Redis, ``breaker_open`` refuses every Redis round trip and ``redis_down`` fails it, and every +call the store makes is logged in ``calls``.""" + +from __future__ import annotations + +import asyncio +from collections.abc import Awaitable, Callable, Sequence +from dataclasses import dataclass, field +from typing import Final, TypeVar + +from litellm.redis.stores.router_usage import CounterIncrement, RouterUsageStore, UsageCounter +from litellm.router_python_host import PythonHost +from litellm.rust_bridge._native import NativeRedisClient +from tests.unit.caching.test_redis_batch import fake_pending + +T = TypeVar("T") + + +def _done(value: T) -> asyncio.Future[T]: + future: Final[asyncio.Future[T]] = asyncio.get_running_loop().create_future() + future.set_result(value) + return future + + +@dataclass +class FakeUsageHandle: + local: dict[UsageCounter, float] = field(default_factory=dict) + redis: dict[UsageCounter, float] | None = None + breaker_open: bool = False + redis_down: bool = False + calls: list[tuple[str, tuple[object, ...]]] = field(default_factory=list) + + @property + def has_redis(self) -> bool: + return self.redis is not None + + def attach(self, client: NativeRedisClient, /) -> bool: + attached: Final = self.redis is None + self.redis = {} if self.redis is None else self.redis + return attached + + def _reachable(self) -> dict[UsageCounter, float] | None: + return None if self.breaker_open or self.redis_down else self.redis + + def _increment(self, increment: CounterIncrement) -> float: + self.local[increment.counter] = self.local.get(increment.counter, 0) + increment.amount + redis: Final = self._reachable() + if redis is None: + return self.local[increment.counter] + redis[increment.counter] = redis.get(increment.counter, 0) + increment.amount + return redis[increment.counter] + + def increment_sync(self, counter: UsageCounter, value: float, ttl_seconds: int) -> float: + self.calls.append(("increment", (counter, value, ttl_seconds))) + return self._increment(CounterIncrement(counter, value, ttl_seconds)) + + def increment(self, counter: UsageCounter, value: float, ttl_seconds: int) -> Awaitable[float]: + return _done(self.increment_sync(counter, value, ttl_seconds)) + + def increment_pipeline(self, increments: Sequence[CounterIncrement]) -> asyncio.Future[list[float]]: + self.calls.append(("pipeline", tuple(increments))) + return _done([self._increment(increment) for increment in increments]) + + def increment_post_call(self, request: object, increments: Sequence[CounterIncrement]) -> None: + """Memory now; the Redis half is the request's, so only the declaration is logged.""" + self.calls.append(("post_call", tuple(increments))) + for increment in increments: + self.local[increment.counter] = self.local.get(increment.counter, 0) + increment.amount + + def increment_local(self, increment: CounterIncrement) -> float: + self.local[increment.counter] = self.local.get(increment.counter, 0) + increment.amount + return self.local[increment.counter] + + def get_local(self, counter: UsageCounter) -> float | None: + return self.local.get(counter) + + def get_local_many(self, counters: Sequence[UsageCounter]) -> list[float | None]: + return [self.local.get(counter) for counter in counters] + + def set_local(self, counter: UsageCounter, value: float) -> None: + self.local[counter] = value + + def increment_remote(self, increments: Sequence[CounterIncrement]) -> asyncio.Future[list[float] | None]: + self.calls.append(("remote", tuple(increments))) + redis: Final = self._reachable() + if redis is None: + return _done(None) + totals: Final = [redis.get(i.counter, 0) + i.amount for i in increments] + redis.update({i.counter: total for i, total in zip(increments, totals)}) + return _done(totals) + + def current_sync(self, counter: UsageCounter, limit: float) -> float | None: + self.calls.append(("current", (counter, limit))) + local: Final = self.local.get(counter) + if self.redis is None or self.breaker_open or (local is not None and local >= limit): + return local + return None if self.redis_down else self.redis.get(counter) + + def current(self, counter: UsageCounter, limit: float) -> asyncio.Future[float | None]: + return _done(self.current_sync(counter, limit)) + + def shared_or_local(self, counter: UsageCounter) -> asyncio.Future[float | None]: + self.calls.append(("shared_or_local", (counter,))) + redis: Final = self._reachable() + shared: Final = None if redis is None else redis.get(counter) + return _done(self.local.get(counter) if shared is None else shared) + + def _read(self, counters: Sequence[UsageCounter]) -> list[float | None]: + redis: Final = self._reachable() + return [self.local.get(c, None if redis is None else redis.get(c)) for c in counters] + + def get_many_sync(self, counters: Sequence[UsageCounter]) -> list[float | None]: + self.calls.append(("get_many", tuple(counters))) + return self._read(counters) + + def get_many(self, counters: Sequence[UsageCounter]) -> asyncio.Future[list[float | None]]: + return _done(self.get_many_sync(counters)) + + def usage(self, counters: Sequence[UsageCounter], send: object) -> FakeUsageRead: + """Declared on the fake request's batch, or alone for ``NativeSend.Now()``.""" + self.calls.append(("usage", tuple(counters))) + return FakeUsageRead( + fake_pending(send, lambda: self._read(counters)), + lambda: self.calls.append(("release", tuple(counters))), + ) + + +@dataclass(frozen=True, slots=True) +class FakeUsageRead: + read: object + on_release: Callable[[], None] + + def release(self) -> None: + self.on_release() + + +def usage_store(handle: FakeUsageHandle) -> RouterUsageStore: + return RouterUsageStore(handle) # pyright: ignore[reportArgumentType] # implements what the tests call + + +class UsageHost(PythonHost): + """The Python host with a router's usage store counting in ``handle``.""" + + __slots__ = ("_usage",) + + def __init__(self, handle: FakeUsageHandle) -> None: + self._usage: Final = usage_store(handle) + + def router_usage(self, redis: object) -> RouterUsageStore: + return self._usage diff --git a/tests/unit/redis/stores/test_router_usage.py b/tests/unit/redis/stores/test_router_usage.py index 5b46f590c4d..54891c22a8f 100644 --- a/tests/unit/redis/stores/test_router_usage.py +++ b/tests/unit/redis/stores/test_router_usage.py @@ -10,7 +10,7 @@ from litellm._internal_context import current_service_target from litellm.caching.redis_batch import drain_post_call_redis_batches, request_redis_batch_scope from litellm.caching.redis_cache import RedisCache from litellm.redis.stores.connection import connect -from litellm.redis.stores.router_usage import RouterUsageStore, UsageCounter +from litellm.redis.stores.router_usage import CounterIncrement, RouterUsageStore, UsageCounter from litellm.rust_bridge._native import NativeRedisClient from litellm.types.utils import ModelResponse from tests.unit.caching.test_redis_batch import FakeNatives, fake_send @@ -52,8 +52,8 @@ class FakeUsage: requests: int, ) -> None: self.calls.append(("deployment", (deployment_id, model, minute, tokens, requests), current_service_target())) - tpm: Final = UsageCounter(deployment_id, model, minute, "tpm") - rpm: Final = UsageCounter(deployment_id, model, minute, "rpm") + tpm: Final = UsageCounter("global_router", deployment_id, model, minute, "tpm") + rpm: Final = UsageCounter("global_router", deployment_id, model, minute, "rpm") self.counts[tpm] = self.counts.get(tpm, 0) + tokens self.counts[rpm] = self.counts.get(rpm, 0) + requests @@ -84,8 +84,8 @@ class FakeUsage: @pytest.mark.asyncio async def test_counters_cross_the_bridge_as_domain_tuples_under_the_router_usage_target() -> None: handle: Final = FakeUsage() - store: Final = RouterUsageStore(handle) - counter: Final = UsageCounter("dep", None, "12-34", "rpm") + store: Final = RouterUsageStore(handle) # pyright: ignore[reportArgumentType] # implements what the test calls + counter: Final = UsageCounter("global_router", "dep", None, "12-34", "rpm") await store.increment(counter, 1, 60) await store.increment_deployment("dep", "model", "12-34", 7.5, 1) @@ -94,13 +94,35 @@ async def test_counters_cross_the_bridge_as_domain_tuples_under_the_router_usage ("increment", (counter, 1, 60), "router_usage"), ("deployment", ("dep", "model", "12-34", 7.5, 1), "router_usage"), ] - assert await store.get_many((counter, UsageCounter("dep", "model", "12-34", "tpm"))) == [1, 7.5] + assert await store.get_many((counter, UsageCounter("global_router", "dep", "model", "12-34", "tpm"))) == [1, 7.5] + + +@pytest.mark.parametrize( + ("counter", "key"), + [ + (UsageCounter("global_router", "dep", None, "09-05", "rpm"), "global_router:dep:None:rpm:09-05"), + ( + UsageCounter.for_deployment("dep", "openai/gpt-5.4-nano", "rpm", "12-34"), + "dep:openai/gpt-5.4-nano:rpm:12-34", + ), + (UsageCounter.for_deployment(f"{7}", f"{None}", "tpd", "2026-11-10"), "7:None:tpd:2026-11-10"), + ], +) +def test_counters_count_under_mains_keys(counter: UsageCounter, key: str) -> None: + store: Final = RouterUsageStore(connect(None).router_usage()) + + store.increment_local(CounterIncrement(counter, 3, 60)) + + assert store.snapshot() == {key: 3} @pytest.mark.asyncio async def test_late_attach_keeps_the_counts_this_router_already_holds() -> None: store: Final = RouterUsageStore(connect(None).router_usage()) - counters: Final = (UsageCounter("dep", "model", "12-34", "tpm"), UsageCounter("dep", "model", "12-34", "rpm")) + counters: Final = ( + UsageCounter("global_router", "dep", "model", "12-34", "tpm"), + UsageCounter("global_router", "dep", "model", "12-34", "rpm"), + ) await store.increment_deployment("dep", "model", "12-34", 7.5, 1) store.attach_redis(connect(RedisCache(host="redis.internal"))) assert await store.get_many(counters) == [7.5, 1] @@ -132,7 +154,7 @@ async def test_router_reset_clears_usage_and_failure_counters(operation: str) -> @pytest.mark.asyncio async def test_inside_a_request_deployment_increments_go_out_on_their_own_trip_like_main(): handle: Final = FakeUsage() - store: Final = RouterUsageStore(handle) + store: Final = RouterUsageStore(handle) # pyright: ignore[reportArgumentType] # implements what the test calls natives: Final = FakeNatives() await store.increment_deployment("dep", "model", "12-34", 7.5, 1) diff --git a/tests/unit/responses/test_responses_websocket_all_providers.py b/tests/unit/responses/test_responses_websocket_all_providers.py index fe4ee0c4bbe..cb7208f9b9b 100644 --- a/tests/unit/responses/test_responses_websocket_all_providers.py +++ b/tests/unit/responses/test_responses_websocket_all_providers.py @@ -48,10 +48,11 @@ from litellm.llms.volcengine.responses.transformation import ( VolcEngineResponsesAPIConfig, ) from litellm.llms.xai.responses.transformation import XAIResponsesAPIConfig +from litellm.redis.stores.connection import connect +from litellm.redis.stores.router_usage import CounterIncrement, RouterUsageStore, UsageCounter from litellm.types.utils import GenericStreamingChunk if TYPE_CHECKING: - from litellm.caching.dual_cache import DualCache from litellm.litellm_core_utils.litellm_logging import Logging from litellm.responses.streaming_iterator import ResponsesWsTurnAdmission from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck @@ -1704,8 +1705,8 @@ def _guard_recording(bill_refusal: MagicMock): return guard -_OPENED_ON: Final = "global_router:ws-dep:openai/gpt-5-mini:otpm:12-34" -_TURNED_ON: Final = "global_router:ws-dep:openai/gpt-5-mini:otpm:12-35" +_OPENED_ON: Final = UsageCounter("global_router", "ws-dep", "openai/gpt-5-mini", "12-34", "otpm") +_TURNED_ON: Final = UsageCounter("global_router", "ws-dep", "openai/gpt-5-mini", "12-35", "otpm") _OPENED_AT: Final = datetime(2026, 1, 1, 12, 34, 56, tzinfo=timezone.utc) _TURNED_AT: Final = datetime(2026, 1, 1, 12, 35, 5, tzinfo=timezone.utc) _WS_DEPLOYMENT: Final = { @@ -1731,12 +1732,12 @@ def _turn_completed() -> str: ) -async def _socket_opened_on(dual_cache: "DualCache") -> Mapping[str, object]: +async def _socket_opened_on(usage: RouterUsageStore) -> Mapping[str, object]: from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_admit_io_token_turn - await dual_cache.async_increment_cache(key=_OPENED_ON, value=_OTHER_USAGE, ttl=60) - await dual_cache.async_increment_cache(key=_TURNED_ON, value=_OTHER_USAGE, ttl=60) - return await async_admit_io_token_turn(dual_cache, _WS_DEPLOYMENT, _OPENED_AT) + usage.increment_local(CounterIncrement(_OPENED_ON, _OTHER_USAGE, 60)) + usage.increment_local(CounterIncrement(_TURNED_ON, _OTHER_USAGE, 60)) + return await async_admit_io_token_turn(usage, _WS_DEPLOYMENT, _OPENED_AT) def _turns_finish_at(monkeypatch: pytest.MonkeyPatch, finished_at: datetime) -> None: @@ -1745,13 +1746,13 @@ def _turns_finish_at(monkeypatch: pytest.MonkeyPatch, finished_at: datetime) -> ) -def _later_turns_admitted_through(dual_cache: "DualCache") -> "ResponsesWsTurnAdmission": +def _later_turns_admitted_through(usage: RouterUsageStore) -> "ResponsesWsTurnAdmission": from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_admit_io_token_turn - return partial(async_admit_io_token_turn, dual_cache, _WS_DEPLOYMENT, _TURNED_AT) + return partial(async_admit_io_token_turn, usage, _WS_DEPLOYMENT, _TURNED_AT) -_ITPM_OPENED_ON: Final = "global_router:ws-dep:openai/gpt-5-mini:itpm:12-34" +_ITPM_OPENED_ON: Final = UsageCounter("global_router", "ws-dep", "openai/gpt-5-mini", "12-34", "itpm") _ITPM_DEPLOYMENT: Final = { "litellm_params": {"model": "openai/gpt-5-mini", "itpm": 40_000}, "model_info": {"id": "ws-dep"}, @@ -1778,6 +1779,16 @@ def _turn_completed_reading_cache(cached_tokens: int) -> str: ) +def _memory_usage() -> RouterUsageStore: + return RouterUsageStore(connect(None).router_usage()) + + +def _rate_limit_check(usage: RouterUsageStore) -> "ModelRateLimitingCheck": + from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck + + return ModelRateLimitingCheck(usage=usage) + + def _connection_logged_through(rate_limits: "ModelRateLimitingCheck", opened_with: Mapping[str, object]) -> "Logging": import time @@ -2137,13 +2148,11 @@ class TestWebSocketTurnAdmission: ) -> None: import websockets.exceptions # noqa: F401 (lazy submodule must be importable) - from litellm.caching.dual_cache import DualCache from litellm.responses.streaming_iterator import drain_logging_tasks - from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck _turns_finish_at(monkeypatch, _TURNED_AT) - dual_cache: Final = DualCache() - opened_with: Final = await _socket_opened_on(dual_cache) + usage: Final = _memory_usage() + opened_with: Final = await _socket_opened_on(usage) first, second, third = _turn_frames() backend_ws: Final = MagicMock() backend_ws.send = AsyncMock() @@ -2154,18 +2163,18 @@ class TestWebSocketTurnAdmission: websocket=_ClientSocket(second, third), backend_ws=backend_ws, request_data={}, - logging_obj=_connection_logged_through(ModelRateLimitingCheck(dual_cache), opened_with), + logging_obj=_connection_logged_through(_rate_limit_check(usage), opened_with), first_message=first, - turn_admission=_later_turns_admitted_through(dual_cache), + turn_admission=_later_turns_admitted_through(usage), ) await handler.client_to_backend() await handler.backend_to_client() await drain_logging_tasks() - assert await dual_cache.async_get_cache(key=_OPENED_ON) == _OTHER_USAGE + assert usage.get_local(_OPENED_ON) == _OTHER_USAGE assert ( - await dual_cache.async_get_cache(key=_TURNED_ON) + usage.get_local(_TURNED_ON) == _OTHER_USAGE + first_turn_tokens + 2 * _OUTPUT_TOKENS_PER_TURN ) @@ -2182,14 +2191,12 @@ class TestWebSocketTurnAdmission: ) -> None: import websockets.exceptions # noqa: F401 (lazy submodule must be importable) - from litellm.caching.dual_cache import DualCache from litellm.responses.streaming_iterator import drain_logging_tasks from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import async_admit_io_token_turn - from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck _turns_finish_at(monkeypatch, _OPENED_AT) - dual_cache: Final = DualCache() - opened_with: Final = await async_admit_io_token_turn(dual_cache, _ITPM_DEPLOYMENT, _OPENED_AT) + usage: Final = _memory_usage() + opened_with: Final = await async_admit_io_token_turn(usage, _ITPM_DEPLOYMENT, _OPENED_AT) first, _, _ = _turn_frames() backend_ws: Final = MagicMock() backend_ws.send = AsyncMock() @@ -2198,7 +2205,7 @@ class TestWebSocketTurnAdmission: websocket=_ClientSocket(), backend_ws=backend_ws, request_data={}, - logging_obj=_connection_logged_through(ModelRateLimitingCheck(dual_cache), opened_with), + logging_obj=_connection_logged_through(_rate_limit_check(usage), opened_with), first_message=first, ) @@ -2206,7 +2213,7 @@ class TestWebSocketTurnAdmission: await handler.backend_to_client() await drain_logging_tasks() - assert await dual_cache.async_get_cache(key=_ITPM_OPENED_ON) == counted_input_tokens + assert usage.get_local(_ITPM_OPENED_ON) == counted_input_tokens @pytest.mark.asyncio @pytest.mark.parametrize( @@ -2219,25 +2226,23 @@ class TestWebSocketTurnAdmission: async def test_managed_relay_counts_every_turn_on_the_minute_it_finishes( self, monkeypatch: pytest.MonkeyPatch, failing_turns: frozenset[int], first_turn_tokens: int ) -> None: - from litellm.caching.dual_cache import DualCache - from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck _turns_finish_at(monkeypatch, _TURNED_AT) - dual_cache: Final = DualCache() - opened_with: Final = await _socket_opened_on(dual_cache) + usage: Final = _memory_usage() + opened_with: Final = await _socket_opened_on(usage) await _run_managed_relay( monkeypatch, _TwoOutputTokenTurns(failing_turns), - ModelRateLimitingCheck(dual_cache), + _rate_limit_check(usage), _turn_frames(), - _later_turns_admitted_through(dual_cache), + _later_turns_admitted_through(usage), opened_with, ) - assert await dual_cache.async_get_cache(key=_OPENED_ON) == _OTHER_USAGE + assert usage.get_local(_OPENED_ON) == _OTHER_USAGE assert ( - await dual_cache.async_get_cache(key=_TURNED_ON) + usage.get_local(_TURNED_ON) == _OTHER_USAGE + first_turn_tokens + 2 * _OUTPUT_TOKENS_PER_TURN ) @@ -2246,26 +2251,24 @@ class TestWebSocketTurnAdmission: self, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter ) -> None: import litellm - from litellm.caching.dual_cache import DualCache - from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) litellm.in_memory_llm_clients_cache.flush_cache() _turns_finish_at(monkeypatch, _TURNED_AT) provider: Final = respx_mock.post("https://api.openai.com/v1/responses").mock(side_effect=_provider_turn_stream) client_metadata: Final = {"session": "abc", "team": "search"} - dual_cache: Final = DualCache() + usage: Final = _memory_usage() await _relay_turns_to( monkeypatch, "openai/gpt-5-mini", - ModelRateLimitingCheck(dual_cache), + _rate_limit_check(usage), tuple( json.dumps({"type": "response.create", "input": word, "metadata": client_metadata}) for word in ("one", "two") ), - _later_turns_admitted_through(dual_cache), - await _socket_opened_on(dual_cache), + _later_turns_admitted_through(usage), + await _socket_opened_on(usage), api_key="sk-fake", ) @@ -2273,7 +2276,7 @@ class TestWebSocketTurnAdmission: client_metadata, client_metadata, ] - assert await dual_cache.async_get_cache(key=_TURNED_ON) == _OTHER_USAGE + 2 * _OUTPUT_TOKENS_PER_TURN + assert usage.get_local(_TURNED_ON) == _OTHER_USAGE + 2 * _OUTPUT_TOKENS_PER_TURN @pytest.mark.asyncio async def test_native_relay_sends_the_provider_only_the_metadata_the_socket_opened_with( diff --git a/tests/unit/router_strategy/test_base_routing_strategy.py b/tests/unit/router_strategy/test_base_routing_strategy.py index 19ff11dd3f6..21e39a06031 100644 --- a/tests/unit/router_strategy/test_base_routing_strategy.py +++ b/tests/unit/router_strategy/test_base_routing_strategy.py @@ -1,172 +1,108 @@ -import json +import asyncio import logging -from typing import Any, Dict, Final, List, Optional, Set, Union +import threading +from collections.abc import Sequence +from typing import Final import pytest - - -import asyncio -import threading -from unittest.mock import MagicMock, patch - from fakeredis import TcpFakeServer - -from litellm.caching.caching import DualCache -from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache -from litellm.caching.redis_cache import RedisCircuitBreakerOpenError, RedisPipelineIncrementOperation +from litellm.redis.stores.connection import connect +from litellm.redis.stores.router_usage import CounterIncrement, RouterUsageStore, UsageCounter from litellm.router_strategy.base_routing_strategy import BaseRoutingStrategy +from litellm.rust_bridge._native import NativeRedisStoreError +from tests.unit.redis.stores.router_usage_double import FakeUsageHandle, usage_store + +_A: Final = UsageCounter.for_deployment("dep-a", "openai/gpt-5.4-nano", "rpm", "12-34") +_B: Final = UsageCounter.for_deployment("dep-b", "openai/gpt-5.4-nano", "rpm", "12-34") -@pytest.fixture -async def mock_dual_cache(): - dual_cache = MagicMock(spec=DualCache) - dual_cache.in_memory_cache = MagicMock() - dual_cache.redis_cache = MagicMock() - - # Set up async method mocks to return coroutines - future1: asyncio.Future[None] = asyncio.Future() - future1.set_result(None) - dual_cache.in_memory_cache.async_increment.return_value = future1 - - future2: asyncio.Future[None] = asyncio.Future() - future2.set_result(None) - dual_cache.redis_cache.async_increment_pipeline.return_value = future2 - - future3: asyncio.Future[None] = asyncio.Future() - future3.set_result(None) - dual_cache.in_memory_cache.async_set_cache.return_value = future3 - - # Fix for async_batch_get_cache - batch_future: asyncio.Future[Dict[str, str]] = asyncio.Future() - batch_future.set_result({"key1": "10.0", "key2": "20.0"}) - dual_cache.redis_cache.async_batch_get_cache.return_value = batch_future - - return dual_cache - - -@pytest.fixture -async def base_strategy(mock_dual_cache): - return BaseRoutingStrategy( - dual_cache=mock_dual_cache, - should_batch_redis_writes=False, - default_sync_interval=1, - ) +def _strategy(handle: FakeUsageHandle) -> BaseRoutingStrategy: + return BaseRoutingStrategy(usage=usage_store(handle), should_batch_redis_writes=False, default_sync_interval=1) @pytest.mark.asyncio -async def test_increment_value_in_current_window(base_strategy, mock_dual_cache): - # Test incrementing value in current window - key = "test_key" - value = 10.0 - ttl = 3600 +async def test_an_increment_counts_in_memory_now_and_queues_the_redis_write(): + handle: Final = FakeUsageHandle(redis={}) + strategy: Final = _strategy(handle) - await base_strategy._increment_value_in_current_window(key, value, ttl) + counted: Final = await strategy._increment_value_in_current_window(_A, 1, 60) # pyright: ignore[reportPrivateUsage] # a step of the sync loop the strategy runs on a timer - # Verify in-memory cache was incremented - mock_dual_cache.in_memory_cache.async_increment.assert_called_once_with( - key=key, value=value, ttl=ttl - ) - - # Verify operation was queued for Redis - assert len(base_strategy.redis_increment_operation_queue) == 1 - queued_op = base_strategy.redis_increment_operation_queue[0] - assert isinstance(queued_op, dict) - assert queued_op["key"] == key - assert queued_op["increment_value"] == value - assert queued_op["ttl"] == ttl + assert counted == 1 + assert handle.local == {_A: 1} + assert handle.redis == {} + assert strategy.redis_increment_operation_queue == [CounterIncrement(_A, 1, 60)] + assert strategy.get_in_memory_keys_to_update() == {_A} @pytest.mark.asyncio -async def test_push_in_memory_increments_to_redis(base_strategy, mock_dual_cache): - # Add some operations to the queue - base_strategy.redis_increment_operation_queue = [ - RedisPipelineIncrementOperation(key="key1", increment_value=10, ttl=3600), - RedisPipelineIncrementOperation(key="key2", increment_value=20, ttl=3600), - ] +async def test_a_push_sends_each_counters_queued_increments_summed_in_one_pipeline(): + handle: Final = FakeUsageHandle(redis={_A: 10}) + strategy: Final = _strategy(handle) + for counter in (_A, _B, _A): + await strategy._increment_value_in_current_window(counter, 1, 60) # pyright: ignore[reportPrivateUsage] # a step of the sync loop the strategy runs on a timer - await base_strategy._push_in_memory_increments_to_redis() + totals: Final = await strategy._push_in_memory_increments_to_redis() # pyright: ignore[reportPrivateUsage] # a step of the sync loop the strategy runs on a timer - # Verify Redis pipeline was called - mock_dual_cache.redis_cache.async_increment_pipeline.assert_called_once() - # Verify queue was cleared - assert len(base_strategy.redis_increment_operation_queue) == 0 + assert handle.calls == [("remote", (CounterIncrement(_A, 2, 60), CounterIncrement(_B, 1, 60)))] + assert totals == {_A: 12, _B: 1} + assert strategy.redis_increment_operation_queue == [] @pytest.mark.asyncio -async def test_sync_in_memory_spend_with_redis(base_strategy, mock_dual_cache): - from litellm.types.caching import RedisPipelineIncrementOperation +async def test_a_sync_moves_memory_up_to_what_every_pod_counted_in_redis(): + handle: Final = FakeUsageHandle(redis={_A: 20}) + strategy: Final = _strategy(handle) + await strategy._increment_value_in_current_window(_A, 3, 60) # pyright: ignore[reportPrivateUsage] # a step of the sync loop the strategy runs on a timer - # Setup test data - base_strategy.in_memory_keys_to_update = {"key1"} - base_strategy.redis_increment_operation_queue = [ - RedisPipelineIncrementOperation(key="key1", increment_value=10, ttl=3600), - ] + await strategy._sync_in_memory_spend_with_redis() # pyright: ignore[reportPrivateUsage] # a step of the sync loop the strategy runs on a timer - # Mock the in-memory cache batch get responses for before snapshot - in_memory_before_future: asyncio.Future[List[str]] = asyncio.Future() - in_memory_before_future.set_result(["5.0"]) # Initial values - mock_dual_cache.in_memory_cache.async_batch_get_cache.return_value = ( - in_memory_before_future - ) - - # Mock Redis batch get response - redis_future: asyncio.Future[Dict[str, str]] = asyncio.Future() - redis_future.set_result([15.0]) # Redis values - mock_dual_cache.redis_cache.async_increment_pipeline.return_value = redis_future - - # Mock in-memory get for after snapshot - in_memory_after_future: asyncio.Future[Optional[str]] = asyncio.Future() - in_memory_after_future.set_result("8.0") # Value after potential updates - mock_dual_cache.in_memory_cache.async_get_cache.return_value = ( - in_memory_after_future - ) - - await base_strategy._sync_in_memory_spend_with_redis() - - # Verify the final merged values - set_cache_calls = mock_dual_cache.in_memory_cache.async_set_cache.call_args_list - print(f"set_cache_calls: {set_cache_calls}") - assert any( - call.kwargs["key"] == "key1" and float(call.kwargs["value"]) == 18.0 - for call in set_cache_calls - ) - - # Verify cache keys still exist - assert len(base_strategy.in_memory_keys_to_update) == 1 + assert handle.redis == {_A: 23} + assert handle.local == {_A: 23} @pytest.mark.asyncio -async def test_cache_keys_management(base_strategy): - # Test adding and getting cache keys - base_strategy.add_to_in_memory_keys_to_update("key1") - base_strategy.add_to_in_memory_keys_to_update("key2") - base_strategy.add_to_in_memory_keys_to_update("key1") # Duplicate should be ignored +async def test_without_redis_a_sync_leaves_the_queue_and_memory_alone(): + handle: Final = FakeUsageHandle() + strategy: Final = _strategy(handle) + await strategy._increment_value_in_current_window(_A, 3, 60) # pyright: ignore[reportPrivateUsage] # a step of the sync loop the strategy runs on a timer - cache_keys = base_strategy.get_in_memory_keys_to_update() - assert len(cache_keys) == 2 - assert "key1" in cache_keys - assert "key2" in cache_keys + await strategy._sync_in_memory_spend_with_redis() # pyright: ignore[reportPrivateUsage] # a step of the sync loop the strategy runs on a timer - # Test resetting cache keys - base_strategy.reset_in_memory_keys_to_update() - assert len(base_strategy.get_in_memory_keys_to_update()) == 0 + assert handle.local == {_A: 3} + assert handle.calls == [] + assert strategy.redis_increment_operation_queue == [CounterIncrement(_A, 3, 60)] @pytest.mark.asyncio -async def test_push_refused_by_the_open_circuit_breaker_is_not_logged_as_an_error(base_strategy, mock_dual_cache, caplog): +async def test_cache_keys_management(): + strategy: Final = _strategy(FakeUsageHandle()) + strategy.add_to_in_memory_keys_to_update(_A) + strategy.add_to_in_memory_keys_to_update(_B) + strategy.add_to_in_memory_keys_to_update(_A) + + assert strategy.get_in_memory_keys_to_update() == {_A, _B} + strategy.reset_in_memory_keys_to_update() + assert strategy.get_in_memory_keys_to_update() == set() + + +class _RefusedHandle(FakeUsageHandle): + def increment_remote(self, increments: Sequence[CounterIncrement]) -> asyncio.Future[list[float] | None]: + raise NativeRedisStoreError("breaker_open", "the Redis circuit breaker is open") + + +@pytest.mark.asyncio +async def test_push_refused_by_the_open_circuit_breaker_is_not_logged_as_an_error(caplog): """The sync loop pushes every 100 ms under usage-based routing, so an open breaker must not add an error line per cycle.""" - mock_dual_cache.redis_cache.async_increment_pipeline.side_effect = RedisCircuitBreakerOpenError( - "Redis circuit breaker is open - skipping async_increment_pipeline" - ) - base_strategy.redis_increment_operation_queue = [{"key": "k", "increment_value": 1.0, "ttl": 60}] + strategy: Final = _strategy(_RefusedHandle(redis={})) + await strategy._increment_value_in_current_window(_A, 1, 60) # pyright: ignore[reportPrivateUsage] # a step of the sync loop the strategy runs on a timer with caplog.at_level(logging.ERROR): - await base_strategy._push_in_memory_increments_to_redis() + await strategy._push_in_memory_increments_to_redis() # pyright: ignore[reportPrivateUsage] # a step of the sync loop the strategy runs on a timer assert caplog.records == [] - assert base_strategy.redis_increment_operation_queue == [] + assert strategy.redis_increment_operation_queue == [] @pytest.mark.asyncio @@ -175,16 +111,16 @@ async def test_an_increment_counted_after_retire_still_reaches_redis(): threading.Thread(target=server.serve_forever, daemon=True).start() redis_cache: Final = RedisCache(host="127.0.0.1", port=server.server_address[1]) strategy: Final = BaseRoutingStrategy( - dual_cache=DualCache(redis_cache=redis_cache, in_memory_cache=InMemoryCache()), + usage=RouterUsageStore(connect(redis_cache).router_usage()), should_batch_redis_writes=True, default_sync_interval=60, ) try: strategy.retire() - await strategy._increment_value_in_current_window("counted-after-retire", 3, 60) + await strategy._increment_value_in_current_window(_A, 3, 60) # pyright: ignore[reportPrivateUsage] # a step of the sync loop the strategy runs on a timer await asyncio.sleep(0.2) - assert float(await redis_cache.async_get_cache("counted-after-retire")) == 3 + assert float(await redis_cache.async_get_cache("dep-a:openai/gpt-5.4-nano:rpm:12-34")) == 3 finally: server.shutdown() server.server_close() diff --git a/tests/unit/router_strategy/test_lowest_tpm_rpm.py b/tests/unit/router_strategy/test_lowest_tpm_rpm.py index f1c56b9d033..bc4ea36e2ae 100644 --- a/tests/unit/router_strategy/test_lowest_tpm_rpm.py +++ b/tests/unit/router_strategy/test_lowest_tpm_rpm.py @@ -1,14 +1,15 @@ import asyncio from datetime import datetime, timedelta from typing import Final -from unittest.mock import AsyncMock import pytest from litellm import Router, token_counter -from litellm.caching.dual_cache import DualCache +from litellm.redis.stores.router_usage import UsageCounter from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2, PrefetchedUsage from litellm.types.router import DeploymentTypedDict, LiteLLMParamsTypedDict, RouterErrors +from litellm.utils import get_utc_datetime +from tests.unit.redis.stores.router_usage_double import FakeUsageHandle, usage_store MODEL_GROUP: Final = "lowest-tpm-router" HIGH_USAGE_DEPLOYMENT_ID: Final = "highest-usage" @@ -87,29 +88,29 @@ def test_usage_based_routing_v1_skips_a_deployment_whose_tpm_limit_is_used_up() @pytest.mark.asyncio async def test_v2_async_selection_uses_prefetched_counters_only_when_they_cover_its_keys(): - router_cache = DualCache() - router_cache.async_batch_get_cache = AsyncMock(return_value=[100, 10, None, None]) # type: ignore[method-assign] - strategy = LowestTPMLoggingHandler_v2(router_cache=router_cache) deployments = [ {"model_name": "g", "litellm_params": {"model": "m"}, "model_info": {"id": "a"}}, {"model_name": "g", "litellm_params": {"model": "m"}, "model_info": {"id": "b"}}, ] - tpm_keys, rpm_keys = strategy.usage_counter_keys(deployments) - keys = tpm_keys + rpm_keys + minute = get_utc_datetime().strftime("%H-%M") + tpm_counters = tuple(UsageCounter.for_deployment(d, "m", "tpm", minute) for d in ("a", "b")) + counters = tpm_counters + tuple(UsageCounter.for_deployment(d, "m", "rpm", minute) for d in ("a", "b")) + handle = FakeUsageHandle(redis=dict(zip(tpm_counters, [100, 10]))) + strategy = LowestTPMLoggingHandler_v2(usage=usage_store(handle)) - covering = PrefetchedUsage(keys=frozenset(keys), values=dict(zip(keys, [10, 100, None, None]))) + covering = PrefetchedUsage(counters=frozenset(counters), values=dict(zip(counters, [10, 100, None, None]))) with PrefetchedUsage.scoped(covering): chosen: Final = await strategy.async_get_available_deployments(model_group="g", healthy_deployments=deployments) assert chosen["model_info"]["id"] == "a", "the prefetched counters say a is the lowest" - router_cache.async_batch_get_cache.assert_not_awaited() + assert handle.calls == [] - stale = PrefetchedUsage(keys=frozenset(keys[:1]), values={keys[0]: 10}) + stale = PrefetchedUsage(counters=frozenset(counters[:1]), values={counters[0]: 10}) with PrefetchedUsage.scoped(stale): chosen_stale: Final = await strategy.async_get_available_deployments( model_group="g", healthy_deployments=deployments ) assert chosen_stale["model_info"]["id"] == "b", "counters that do not cover this minute's keys are read again" - router_cache.async_batch_get_cache.assert_awaited_once_with(keys=keys) + assert handle.calls == [("get_many", counters)] @pytest.mark.asyncio @@ -133,11 +134,9 @@ async def test_v2_subclass_overriding_async_get_available_deployments_with_the_o model_list=[_deployment(HIGH_USAGE_DEPLOYMENT_ID), _deployment(LOW_USAGE_DEPLOYMENT_ID)], routing_strategy="usage-based-routing-v2", ) - router.lowesttpm_logger_v2 = OldSignatureV2(router_cache=router.cache, routing_args={}) + router.lowesttpm_logger_v2 = OldSignatureV2(usage=router.lowesttpm_logger_v2.usage, routing_args={}) - response: Final = await router.acompletion( - model=MODEL_GROUP, messages=[{"role": "user", "content": "x"}] - ) + response: Final = await router.acompletion(model=MODEL_GROUP, messages=[{"role": "user", "content": "x"}]) assert response.choices[0].message.content in { f"from {HIGH_USAGE_DEPLOYMENT_ID}", diff --git a/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py b/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py index 60cad628598..b09759a9f04 100644 --- a/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py +++ b/tests/unit/router_strategy/test_lowest_tpm_rpm_v2.py @@ -22,7 +22,8 @@ from PIL import Image import litellm from litellm import Router from litellm._logging import verbose_logger -from litellm.caching.caching import DualCache +from litellm.redis.stores.connection import connect +from litellm.redis.stores.router_usage import RouterUsageStore, UsageCounter from litellm.router_strategy.lowest_tpm_rpm_v2 import ( LowestTPMLoggingHandler_v2 as LowestTPMLoggingHandler, ) @@ -32,6 +33,11 @@ from litellm.utils import get_utc_datetime ### UNIT TESTS FOR TPM/RPM ROUTING ### + +def _memory_usage() -> RouterUsageStore: + return RouterUsageStore(connect(None).router_usage()) + + """ - Given 2 deployments, make sure it's shuffling deployments correctly. """ @@ -48,8 +54,8 @@ def restore_verbose_logger_level() -> Iterator[None]: def test_tpm_rpm_updated(): from tests.unit.create_mock_standard_logging_payload import create_standard_logging_payload - test_cache = DualCache() - lowest_tpm_logger = LowestTPMLoggingHandler(router_cache=test_cache) + test_cache = _memory_usage() + lowest_tpm_logger = LowestTPMLoggingHandler(usage=test_cache) model_group = "gpt-3.5-turbo" deployment_id = "1234" deployment = "azure/gpt-4.1-mini" @@ -89,11 +95,11 @@ def test_tpm_rpm_updated(): ) dt = get_utc_datetime() current_minute = dt.strftime("%H-%M") - tpm_count_api_key = f"{deployment_id}:{deployment}:tpm:{current_minute}" - rpm_count_api_key = f"{deployment_id}:{deployment}:rpm:{current_minute}" + tpm_counter = UsageCounter.for_deployment(deployment_id, deployment, "tpm", current_minute) + rpm_counter = UsageCounter.for_deployment(deployment_id, deployment, "rpm", current_minute) - assert response_obj["usage"]["total_tokens"] == test_cache.get_cache(key=tpm_count_api_key) - assert 1 == test_cache.get_cache(key=rpm_count_api_key) + assert response_obj["usage"]["total_tokens"] == test_cache.get_local(tpm_counter) + assert 1 == test_cache.get_local(rpm_counter) # test_tpm_rpm_updated() @@ -103,7 +109,7 @@ def test_tpm_rpm_updated(): def test_get_available_deployments(): from tests.unit.create_mock_standard_logging_payload import create_standard_logging_payload - test_cache = DualCache() + test_cache = _memory_usage() model_list = [ { "model_name": "gpt-3.5-turbo", @@ -116,7 +122,7 @@ def test_get_available_deployments(): "model_info": {"id": "5678"}, }, ] - lowest_tpm_logger = LowestTPMLoggingHandler(router_cache=test_cache) + lowest_tpm_logger = LowestTPMLoggingHandler(usage=test_cache) model_group = "gpt-3.5-turbo" ## DEPLOYMENT 1 ## total_tokens = 50 @@ -570,8 +576,8 @@ def test_return_potential_deployments(): Assert deployment at limit is filtered out """ - test_cache = DualCache() - lowest_tpm_logger = LowestTPMLoggingHandler(router_cache=test_cache) + test_cache = _memory_usage() + lowest_tpm_logger = LowestTPMLoggingHandler(usage=test_cache) args: Dict = { "healthy_deployments": [ diff --git a/tests/unit/router_strategy/test_router_routing_groups.py b/tests/unit/router_strategy/test_router_routing_groups.py index 8a0d5b9a1d6..08b045aea82 100644 --- a/tests/unit/router_strategy/test_router_routing_groups.py +++ b/tests/unit/router_strategy/test_router_routing_groups.py @@ -10,19 +10,21 @@ import datetime import uuid from collections.abc import Mapping from typing import Final, Literal -from unittest.mock import patch import pytest from pydantic import ValidationError import litellm from litellm import Router -from litellm.caching.redis_cache import RedisPipelineIncrementOperation from litellm.integrations.custom_logger import CustomLogger +from litellm.redis.stores.router_usage import CounterIncrement, UsageCounter +from litellm.router import router_arguments from litellm.router_api import RoutingContext +from litellm.router_config import build_router_config from litellm.router_strategy.least_busy import picks_reserve from litellm.types.router import DeploymentTypedDict, FallbackAccessCheck, RoutingGroup from litellm.utils import Rules, function_setup +from tests.unit.redis.stores.router_usage_double import FakeUsageHandle, UsageHost def _model_list(): @@ -1094,33 +1096,21 @@ async def test_update_settings_same_routing_strategy_args_does_not_leak_sync_tas pass -class _RecordingRedisCache: - def __init__(self) -> None: - self.increment_lists: list[list[RedisPipelineIncrementOperation]] = [] - - async def async_increment_pipeline(self, increment_list: list[RedisPipelineIncrementOperation]) -> list[float]: - self.increment_lists.append(list(increment_list)) - return [float(op["increment_value"]) for op in increment_list] - - @pytest.mark.asyncio async def test_update_settings_changed_routing_strategy_args_flushes_replaced_selector_queue( monkeypatch, ) -> None: monkeypatch.setattr(litellm, "callbacks", []) monkeypatch.setattr(litellm, "input_callback", []) - router: Final = Router( - model_list=_model_list(), - routing_strategy="usage-based-routing-v2", - routing_strategy_args={"ttl": 60}, + handle: Final = FakeUsageHandle(redis={}) + queued: Final = CounterIncrement(UsageCounter.for_deployment("dep", "m", "rpm", "12-34"), 3, 60) + arguments: Final = router_arguments( + model_list=_model_list(), routing_strategy="usage-based-routing-v2", routing_strategy_args={"ttl": 60} ) + router: Final = Router.from_config(build_router_config(arguments), UsageHost(handle)) try: - redis_cache: Final = _RecordingRedisCache() replaced: Final = router.lowesttpm_logger_v2 - replaced.dual_cache.redis_cache = redis_cache - replaced.redis_increment_operation_queue.append( - RedisPipelineIncrementOperation(key="rpm-key", increment_value=3, ttl=60) - ) + replaced.redis_increment_operation_queue.append(queued) router.update_settings(routing_strategy_args={"ttl": 120}) await asyncio.sleep(0) @@ -1134,9 +1124,7 @@ async def test_update_settings_changed_routing_strategy_args_flushes_replaced_se ) assert router.lowesttpm_logger_v2 is not replaced - assert redis_cache.increment_lists == [ - [RedisPipelineIncrementOperation(key="rpm-key", increment_value=3, ttl=60)] - ] + assert [args for name, args in handle.calls if name == "remote"] == [(queued,)] assert replaced.redis_increment_operation_queue == [] finally: for t in [ diff --git a/tests/unit/router_utils/test_routing_read_batch.py b/tests/unit/router_utils/test_routing_read_batch.py index c8f4b42924c..96f46bc9e02 100644 --- a/tests/unit/router_utils/test_routing_read_batch.py +++ b/tests/unit/router_utils/test_routing_read_batch.py @@ -1,23 +1,27 @@ """ The router's pre-call reads. Cooldowns go through the native store, and the tpm/rpm selector -counters (`LowestTPMLoggingHandler_v2`) ride one MGET on the Python Redis. +counters (`LowestTPMLoggingHandler_v2`) are one read of the native usage store. """ import time from collections.abc import Sequence from typing import Final -from unittest.mock import AsyncMock, MagicMock +from unittest.mock import AsyncMock import pytest -import litellm from litellm import Router from litellm.caching.redis_batch import request_redis_batch_scope from litellm.redis.stores.router_cooldowns import CooldownCacheValue, RouterCooldownStore +from litellm.redis.stores.router_usage import UsageCounter +from litellm.router import router_arguments +from litellm.router_config import build_router_config from litellm.router_strategy.lowest_tpm_rpm_v2 import LowestTPMLoggingHandler_v2 from litellm.router_utils import routing_prefetch +from litellm.utils import get_utc_datetime from tests.unit.caching.test_redis_batch import FakeNativePending, FakeNatives, ReadyPending, fake_pending, fake_send from tests.unit.conftest import redis_cache_double +from tests.unit.redis.stores.router_usage_double import FakeUsageHandle, UsageHost from tests.unit.redis.stores.test_router_cooldowns import FakeCooldowns _MODEL_GROUP = "claude" @@ -32,30 +36,31 @@ def _deployment(deployment_id: str) -> dict: } -def _redis_answering(values_by_key_prefix: dict[str, object]) -> MagicMock: - """A Redis double that answers each key from its minute-less prefix and records every MGET.""" - - def _mget(key_list, parent_otel_span=None): - return {key: values_by_key_prefix.get(key.rsplit(":", 1)[0], values_by_key_prefix.get(key)) for key in key_list} - - redis = redis_cache_double() - redis.async_batch_get_cache = AsyncMock(side_effect=_mget) - return redis - - -def _router(redis: MagicMock, routing_strategy: str) -> Router: - router = Router( - model_list=[_deployment("dep-a"), _deployment("dep-b")], - routing_strategy=routing_strategy, +def _router(routing_strategy: str, handle: FakeUsageHandle | None = None) -> Router: + arguments: Final = router_arguments( + model_list=[_deployment("dep-a"), _deployment("dep-b")], routing_strategy=routing_strategy ) - router.update_redis_cache(cache=redis) - return router + return Router.from_config(build_router_config(arguments), UsageHost(handle or FakeUsageHandle(redis={}))) -def _redis_key_families(redis: MagicMock) -> list[list[str]]: +def _counter(deployment_id: str, kind: str) -> UsageCounter: + return UsageCounter.for_deployment( + deployment_id, + "anthropic/claude-x", + kind, # pyright: ignore[reportArgumentType] # tpm or rpm + get_utc_datetime().strftime("%H-%M"), + ) + + +def _tpm(**tokens: float | None) -> dict[UsageCounter, float]: + return {_counter(deployment_id, "tpm"): count for deployment_id, count in tokens.items() if count is not None} + + +def _usage_reads(handle: FakeUsageHandle) -> list[tuple[str, ...]]: return [ - sorted(key.rsplit(":", 1)[0] if ":tpm:" in key or ":rpm:" in key else key for key in call.args[0]) - for call in redis.async_batch_get_cache.await_args_list + tuple(f"{c.deployment_id}:{c.kind}" for c in counters) # pyright: ignore[reportAttributeAccessIssue] # usage() args are counters + for name, counters in handle.calls + if name == "usage" ] @@ -66,21 +71,16 @@ def _cool_down(router: Router, model_id: str, started: float) -> None: @pytest.mark.asyncio -async def test_usage_based_routing_reads_the_counters_in_one_python_redis_round_trip(): - redis = _redis_answering({}) - router = _router(redis, "usage-based-routing-v2") +async def test_usage_based_routing_reads_the_counters_in_one_store_read(): + handle: Final = FakeUsageHandle(redis={}) + router: Final = _router("usage-based-routing-v2", handle) - deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES) + deployment: Final = await router.async_get_available_deployment( + model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES + ) assert deployment["model_info"]["id"] in {"dep-a", "dep-b"} - assert _redis_key_families(redis) == [ - [ - "dep-a:anthropic/claude-x:rpm", - "dep-a:anthropic/claude-x:tpm", - "dep-b:anthropic/claude-x:rpm", - "dep-b:anthropic/claude-x:tpm", - ] - ], "the usage counters must arrive in one MGET, without cooldown keys" + assert _usage_reads(handle) == [("dep-a:tpm", "dep-b:tpm", "dep-a:rpm", "dep-b:rpm")] @pytest.mark.asyncio @@ -100,35 +100,30 @@ async def test_usage_based_routing_still_batches_when_the_strategy_is_a_fixed_si input=input, ) - redis: Final = _redis_answering({}) - router: Final = _router(redis, "usage-based-routing-v2") - router.lowesttpm_logger_v2 = OldSignatureV2(router_cache=router.cache) - router.cache.async_batch_get_cache = AsyncMock(wraps=router.cache.async_batch_get_cache) + handle: Final = FakeUsageHandle(redis={}) + router: Final = _router("usage-based-routing-v2", handle) + router.lowesttpm_logger_v2.cancel_sync_task() + router.lowesttpm_logger_v2 = OldSignatureV2(usage=router.lowesttpm_logger_v2.usage) deployment: Final = await router.async_get_available_deployment( model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES ) assert deployment["model_info"]["id"] in {"dep-a", "dep-b"} - assert _redis_key_families(redis) == [ - [ - "dep-a:anthropic/claude-x:rpm", - "dep-a:anthropic/claude-x:tpm", - "dep-b:anthropic/claude-x:rpm", - "dep-b:anthropic/claude-x:tpm", - ] - ], "the subclassed strategy must still get the batched read, not a second MGET" - router.cache.async_batch_get_cache.assert_not_awaited() + assert _usage_reads(handle) == [("dep-a:tpm", "dep-b:tpm", "dep-a:rpm", "dep-b:rpm")], ( + "the subclassed strategy must still get the batched read, not a second one" + ) + assert [name for name, _ in handle.calls] == ["usage"] @pytest.mark.asyncio -async def test_simple_shuffle_reads_nothing_from_the_python_redis(): - redis = _redis_answering({}) - router = _router(redis, "simple-shuffle") +async def test_simple_shuffle_reads_no_usage(): + handle: Final = FakeUsageHandle(redis={}) + router: Final = _router("simple-shuffle", handle) await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES) - assert _redis_key_families(redis) == [] + assert handle.calls == [] @pytest.mark.asyncio @@ -137,14 +132,13 @@ async def test_simple_shuffle_reads_nothing_from_the_python_redis(): [(100, 10, "dep-b"), (10, 100, "dep-a"), (None, 10, "dep-a"), (10, None, "dep-b")], ) async def test_batched_counters_pick_the_deployment_the_strategy_picks_reading_alone(tpm_a, tpm_b, expected): - counters = {"dep-a:anthropic/claude-x:tpm": tpm_a, "dep-b:anthropic/claude-x:tpm": tpm_b} - routed = _router(_redis_answering(counters), "usage-based-routing-v2") - alone = _router(_redis_answering(counters), "usage-based-routing-v2") + routed: Final = _router("usage-based-routing-v2", FakeUsageHandle(redis=_tpm(**{"dep-a": tpm_a, "dep-b": tpm_b}))) + alone: Final = _router("usage-based-routing-v2", FakeUsageHandle(redis=_tpm(**{"dep-a": tpm_a, "dep-b": tpm_b}))) - routed_choice = await routed.async_get_available_deployment( + routed_choice: Final = await routed.async_get_available_deployment( model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES ) - alone_choice = await alone.lowesttpm_logger_v2.async_get_available_deployments( + alone_choice: Final = await alone.lowesttpm_logger_v2.async_get_available_deployments( model_group=_MODEL_GROUP, healthy_deployments=alone.model_list, messages=_MESSAGES ) @@ -153,50 +147,36 @@ async def test_batched_counters_pick_the_deployment_the_strategy_picks_reading_a @pytest.mark.asyncio async def test_batched_read_still_excludes_a_cooled_down_deployment(): - redis = _redis_answering({"dep-a:anthropic/claude-x:tpm": 100, "dep-b:anthropic/claude-x:tpm": 10}) - router = _router(redis, "usage-based-routing-v2") + handle: Final = FakeUsageHandle(redis=_tpm(**{"dep-a": 100, "dep-b": 10})) + router: Final = _router("usage-based-routing-v2", handle) _cool_down(router, "dep-b", started=time.time()) - deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES) + deployment: Final = await router.async_get_available_deployment( + model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES + ) assert deployment["model_info"]["id"] == "dep-a", "dep-b has the lowest tpm but is cooling down" - assert redis.async_batch_get_cache.await_count == 1 + assert len(_usage_reads(handle)) == 1 @pytest.mark.asyncio async def test_batched_read_ignores_an_expired_cooldown(): - redis = _redis_answering({"dep-a:anthropic/claude-x:tpm": 100, "dep-b:anthropic/claude-x:tpm": 10}) - router = _router(redis, "usage-based-routing-v2") + router: Final = _router("usage-based-routing-v2", FakeUsageHandle(redis=_tpm(**{"dep-a": 100, "dep-b": 10}))) _cool_down(router, "dep-b", started=time.time() - 61) - deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES) + deployment: Final = await router.async_get_available_deployment( + model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES + ) assert deployment["model_info"]["id"] == "dep-b" -@pytest.mark.asyncio -async def test_a_failed_batched_read_degrades_like_the_two_failed_reads_did(): - redis = redis_cache_double() - redis.async_batch_get_cache = AsyncMock(side_effect=ConnectionError("redis unavailable")) - routed = _router(redis, "usage-based-routing-v2") - alone = _router(redis, "usage-based-routing-v2") - - with pytest.raises(litellm.RateLimitError, match="No deployments available") as routed_error: - await routed.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES) - with pytest.raises(litellm.RateLimitError, match="No deployments available") as alone_error: - await alone.lowesttpm_logger_v2.async_get_available_deployments( - model_group=_MODEL_GROUP, healthy_deployments=alone.model_list, messages=_MESSAGES - ) - - assert str(routed_error.value) == str(alone_error.value) - assert len(routed.cache.last_redis_batch_access_time) == 0, "a failed read must not throttle the next one" - - @pytest.mark.asyncio async def test_a_failed_batched_read_leaves_simple_shuffle_routing(): redis = redis_cache_double() redis.async_batch_get_cache = AsyncMock(side_effect=ConnectionError("redis unavailable")) - router = _router(redis, "simple-shuffle") + router = _router("simple-shuffle") + router.update_redis_cache(cache=redis) deployment = await router.async_get_available_deployment(model=_MODEL_GROUP, request_kwargs={}, messages=_MESSAGES) diff --git a/tests/unit/rust_bridge/host/test_redis_events.py b/tests/unit/rust_bridge/host/test_redis_events.py index 42caf6e4179..ef35de014ef 100644 --- a/tests/unit/rust_bridge/host/test_redis_events.py +++ b/tests/unit/rust_bridge/host/test_redis_events.py @@ -2,6 +2,7 @@ import asyncio import queue import threading from collections.abc import Iterator +from dataclasses import asdict from typing import Final import pytest @@ -15,6 +16,7 @@ from prometheus_client import REGISTRY from redis import exceptions as redis_exceptions import litellm +from litellm import Router from litellm._internal_context import service_target, set_request_parent_otel_span from litellm.caching.redis_batch import request_redis_batch_scope from litellm.caching.redis_cache import RedisCache @@ -24,7 +26,15 @@ from litellm.integrations.otel.logger import OpenTelemetryV2 from litellm.integrations.otel.plumbing import providers from litellm.redis.stores.auth_cache import AUTH_CACHE_TARGET, AuthCacheSettings, AuthCacheStore from litellm.redis.stores.connection import connect +from litellm.redis.stores.router_usage import RouterUsageStore, UsageCounter from litellm.redis.stores.spend_counters import SpendCounterStore +from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( + ITPM_COUNTER_KEY, + ITPM_RESERVED_KEY, + OTPM_COUNTER_KEY, + OTPM_RESERVED_KEY, +) +from litellm.router_utils.pre_call_checks.model_rate_limit_check import ModelRateLimitingCheck from litellm.rust_bridge.host import redis_events @@ -223,6 +233,86 @@ async def test_a_native_read_in_a_legacy_otel_request_reports_its_redis_span_und assert await _spans_under(exporter, native_request) == reported_by_redis_cache +async def _redis_calls_under(exporter: InMemorySpanExporter, parent: Span, polls: int = 100) -> frozenset[str]: + calls: Final = frozenset( + str(span.attributes["call_type"]) + for span in _children(exporter, parent) + if span.attributes is not None and span.attributes.get("service") == "redis" + ) + if calls or polls == 0: + return calls + await asyncio.sleep(0.05) + return await _redis_calls_under(exporter, parent, polls - 1) + + +async def test_a_deadline_flush_reports_under_the_request_span_auth_set_after_the_scope_opened( + fake_redis_port: int, legacy_otel: tuple[TracerProvider, InMemorySpanExporter] +): + provider, exporter = legacy_otel + store: Final = SpendCounterStore(connect(RedisCache(host="127.0.0.1", port=fake_redis_port)).spend_counters()) + request: Final = provider.get_tracer(__name__).start_span("request") + + async def proxy_request() -> frozenset[str]: + with request_redis_batch_scope(post_call_deadline=0.05): + set_request_parent_otel_span(request) + assert store.increment_many_post_call((("spend:key:k", 1.0),)) + return await _redis_calls_under(exporter, request) + + assert await asyncio.create_task(proxy_request()) == {"post_call_redis_batch"} + + +async def test_a_failed_sdk_calls_io_token_refund_reports_its_redis_span_under_the_span_its_caller_passed( + fake_redis_port: int, legacy_otel: tuple[TracerProvider, InMemorySpanExporter], monkeypatch: pytest.MonkeyPatch +): + provider, exporter = legacy_otel + usage: Final = RouterUsageStore(connect(RedisCache(host="127.0.0.1", port=fake_redis_port)).router_usage()) + monkeypatch.setattr(litellm, "callbacks", [ModelRateLimitingCheck(usage)]) + caller: Final = provider.get_tracer(__name__).start_span("sdk caller") + itpm: Final = UsageCounter("global_router", "d1", "openai/gpt-x", "12-00", "itpm") + otpm: Final = UsageCounter("global_router", "d1", "openai/gpt-x", "12-00", "otpm") + + with pytest.raises(Exception, match="provider down"): + await litellm.acompletion( + model="gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + mock_response=Exception("provider down"), + metadata={ + "litellm_parent_otel_span": caller, + ITPM_RESERVED_KEY: 5, + OTPM_RESERVED_KEY: 7, + ITPM_COUNTER_KEY: asdict(itpm), + OTPM_COUNTER_KEY: asdict(otpm), + }, + ) + + assert await _redis_calls_under(exporter, caller) == {"async_increment"} + + +async def test_an_sdk_router_call_reports_its_usage_read_under_the_span_its_caller_passed( + fake_redis_port: int, legacy_otel: tuple[TracerProvider, InMemorySpanExporter] +): + provider, exporter = legacy_otel + caller: Final = provider.get_tracer(__name__).start_span("sdk caller") + router: Final = Router( + model_list=[ + { + "model_name": "m", + "litellm_params": {"model": "openai/gpt-x", "api_key": "sk-fake", "mock_response": "hi", "tpm": 1000}, + "model_info": {"id": "d1"}, + } + ], + routing_strategy="usage-based-routing-v2", + redis_host="127.0.0.1", + redis_port=fake_redis_port, + ) + + await router.acompletion( + model="m", messages=[{"role": "user", "content": "hi"}], metadata={"litellm_parent_otel_span": caller} + ) + + assert "async_batch_get_cache" in await _redis_calls_under(exporter, caller) + + class _EndedSpans(SpanProcessor): def __init__(self) -> None: self.ended: Final = queue.SimpleQueue[tuple[ReadableSpan, int]]() diff --git a/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index 537afda84aa..cba28f075ac 100644 --- a/tests/unit/test_internal_context.py +++ b/tests/unit/test_internal_context.py @@ -14,10 +14,13 @@ from typing import Final import pytest from litellm._internal_context import ( + current_request_parent_otel_span, current_service_caller, current_service_target, + request_parent_otel_span, service_caller, service_target, + set_request_parent_otel_span, with_service_target, ) @@ -82,6 +85,7 @@ _IN_MEMORY_ONLY_CALLERS: Final = frozenset( "litellm/proxy/spend_tracking/key_metadata_recovery.py", "litellm/proxy/ui_crud_endpoints/latest_release_endpoints.py", "litellm/responses/litellm_completion_transformation/transformation.py", + "litellm/router_rust/surface/model_info.py", "litellm/router_utils/client_initalization_utils.py", "litellm/router_utils/pre_call_checks/prompt_caching_deployment_check.py", "litellm/router_utils/router_callbacks/track_deployment_metrics.py", @@ -195,6 +199,29 @@ class _Suspend: yield +def test_an_sdk_callers_span_anchors_the_block_and_the_tasks_it_starts_and_nothing_after() -> None: + sdk_span: Final = object() + + async def sdk_call() -> tuple[object, object, object]: + with request_parent_otel_span(sdk_span): # pyright: ignore[reportArgumentType] # any object stands in for a span + inside: Final = current_request_parent_otel_span() + background: Final = asyncio.create_task(asyncio.to_thread(current_request_parent_otel_span)) + return inside, await background, current_request_parent_otel_span() + + assert asyncio.run(sdk_call()) == (sdk_span, sdk_span, None) + + +def test_a_request_that_anchored_its_own_span_keeps_it_over_one_an_sdk_caller_passes() -> None: + proxy_span: Final = object() + + def proxy_request() -> object: + set_request_parent_otel_span(proxy_span) # pyright: ignore[reportArgumentType] # any object stands in for a span + with request_parent_otel_span(object()): # pyright: ignore[reportArgumentType] # any object stands in for a span + return current_request_parent_otel_span() + + assert contextvars.copy_context().run(proxy_request) is proxy_span + + def test_a_targeted_coroutine_closed_from_another_context_does_not_raise() -> None: @with_service_target("router_usage") async def sync_forever() -> None: diff --git a/tests/unit/test_router/core/test_registry_characterization_model_info.py b/tests/unit/test_router/core/test_registry_characterization_model_info.py index b613aed041d..3bf401ba9e5 100644 --- a/tests/unit/test_router/core/test_registry_characterization_model_info.py +++ b/tests/unit/test_router/core/test_registry_characterization_model_info.py @@ -4,13 +4,16 @@ limits/mode/display name/service tiers/listing info, group usage this minute, re the rate limit headers ``set_response_headers`` attaches. """ +from typing import Final, Literal + import pytest import litellm from litellm.litellm_core_utils.get_model_cost_map import GetModelCostMap from litellm.router_api import ChangeSet from litellm.router_core import RouterCore -from litellm.types.router import Deployment, RouterCacheEnum +from litellm.redis.stores.router_usage import UsageCounter +from litellm.types.router import Deployment from litellm.utils import get_utc_datetime from tests.unit.test_router.core.registry_helpers import MakeCore, typed, unique_id from tests.unit.test_router.core.scripted_host import completion_response, deployment @@ -46,9 +49,9 @@ async def _serve(core: RouterCore, deployment_id: str, model: str, total_tokens: ) -async def _count(core: RouterCore, kind: RouterCacheEnum, deployment_id: str, model: str, value: int) -> None: - key = kind.value.format(id=deployment_id, model=model, current_minute=_minute()) - await core.health().store.async_set_cache(key, value) +async def _count(core: RouterCore, kind: Literal["itpm", "otpm"], deployment_id: str, model: str, value: int) -> None: + usage: Final = core._usage # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # the store the core counts io tokens in + usage.set_local(UsageCounter("global_router", deployment_id, model, _minute(), kind), value) # -- model group info ----------------------------------------------------------------------------- @@ -500,7 +503,7 @@ async def test_group_usage_follows_deployments_added_and_removed_after_a_read(ma async def test_io_token_usage_sums_input_and_output_counters(make_core: MakeCore) -> None: core = _usage_core(make_core) assert await core.get_model_group_io_token_usage("io") == (None, None) # pyright: ignore[reportAttributeAccessIssue] # public on the core - await _count(core, RouterCacheEnum.ITPM, "i", "openai/gpt-4o", 7) + await _count(core, "itpm", "i", "openai/gpt-4o", 7) assert await core.get_model_group_io_token_usage("io") == (7, None) # pyright: ignore[reportAttributeAccessIssue] # public on the core @@ -525,7 +528,7 @@ async def test_remaining_usage_of_an_io_limited_group_reports_input_and_output_t make_core: MakeCore, ) -> None: core = _usage_core(make_core) - await _count(core, RouterCacheEnum.ITPM, "i", "openai/gpt-4o", 7) + await _count(core, "itpm", "i", "openai/gpt-4o", 7) remaining = await core.health().usage("io", kind="remaining") diff --git a/tests/unit/test_router/test_enforce_model_rate_limits.py b/tests/unit/test_router/test_enforce_model_rate_limits.py index ff3ffc0db0c..9edaf885e4f 100644 --- a/tests/unit/test_router/test_enforce_model_rate_limits.py +++ b/tests/unit/test_router/test_enforce_model_rate_limits.py @@ -6,20 +6,21 @@ regardless of the routing strategy being used. """ import asyncio +from collections.abc import Callable from datetime import datetime, timedelta, timezone from typing import Final -from unittest.mock import AsyncMock, MagicMock import pytest -from pydantic import TypeAdapter import litellm from litellm import Router -from litellm.caching.dual_cache import DualCache -from litellm.caching.redis_cache import RedisCircuitBreakerOpenError +from litellm.redis.stores.connection import connect +from litellm.redis.stores.router_usage import CounterIncrement, CounterKind, RouterUsage, RouterUsageStore, UsageCounter from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( ModelRateLimitingCheck, ) +from litellm.utils import get_utc_datetime +from tests.unit.redis.stores.router_usage_double import FakeUsageHandle, usage_store TPM_DEPLOYMENT = { "tpm": 1000, @@ -29,14 +30,49 @@ TPM_DEPLOYMENT = { } -def _dual_cache_with_local_tpm(local_tpm: int, redis_cache: MagicMock | None) -> DualCache: - dual_cache = DualCache(redis_cache=redis_cache) - check = ModelRateLimitingCheck(dual_cache=dual_cache) - now = litellm.utils.get_utc_datetime() - for minute in (now, now + timedelta(minutes=1)): - tpm_key, _ = check._get_cache_keys(TPM_DEPLOYMENT, minute.strftime("%H-%M")) - dual_cache.set_cache(key=tpm_key, value=local_tpm, local_only=True) - return dual_cache +def _memory_usage() -> RouterUsageStore: + return RouterUsageStore(connect(None).router_usage()) + + +def _check( + usage: RouterUsage | None = None, clock: Callable[[], datetime] = get_utc_datetime +) -> ModelRateLimitingCheck: + return ModelRateLimitingCheck(usage=usage or _memory_usage(), clock=clock) + + +def _counter(deployment: dict, kind: CounterKind, window: str) -> UsageCounter: + model_id: Final[object] = deployment["model_info"]["id"] + model: Final[object] = deployment["litellm_params"]["model"] + return UsageCounter.for_deployment(f"{model_id}", f"{model}", kind, window) + + +def _this_minute(deployment: dict, kind: CounterKind, value: float) -> dict[UsageCounter, float]: + """``value`` under this minute's counter and the next, so a test crossing a minute still finds it.""" + now: Final = get_utc_datetime() + return {_counter(deployment, kind, minute.strftime("%H-%M")): value for minute in (now, now + timedelta(minutes=1))} + + +RPM_DEPLOYMENT: Final = { + "rpm": 10, + "litellm_params": {"model": "gpt-4"}, + "model_info": {"id": "test-id"}, + "model_name": "test-model", +} + + +_SUCCESS_KWARGS: Final = { + "standard_logging_object": { + "model_id": "test-id", + "total_tokens": 50, + "hidden_params": {"litellm_model_name": "gpt-4"}, + } +} + + +def _tpm_increments(handle: FakeUsageHandle) -> tuple[tuple[object, object], ...]: + """The (tokens, ttl) each success charged to test-id's gpt-4 tpm counter this minute or the next.""" + tpm_counters: Final = frozenset(_this_minute(RPM_DEPLOYMENT, "tpm", 0)) + return tuple((args[1], args[2]) for name, args in handle.calls if name == "increment" and args[0] in tpm_counters) class TestModelRateLimitingCheck: @@ -44,7 +80,7 @@ class TestModelRateLimitingCheck: def test_get_deployment_limits_from_top_level(self): """Test extracting limits from top-level deployment config.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) + check = _check() deployment = { "tpm": 1000, @@ -59,7 +95,7 @@ class TestModelRateLimitingCheck: def test_get_deployment_limits_from_litellm_params(self): """Test extracting limits from litellm_params.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) + check = _check() deployment = { "litellm_params": {"model": "gpt-4", "tpm": 2000, "rpm": 20}, @@ -72,7 +108,7 @@ class TestModelRateLimitingCheck: def test_get_deployment_limits_from_model_info(self): """Test extracting limits from model_info.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) + check = _check() deployment = { "litellm_params": {"model": "gpt-4"}, @@ -85,7 +121,7 @@ class TestModelRateLimitingCheck: def test_get_deployment_limits_none_when_not_set(self): """Test that None is returned when limits are not set.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) + check = _check() deployment = { "litellm_params": {"model": "gpt-4"}, @@ -98,7 +134,7 @@ class TestModelRateLimitingCheck: def test_pre_call_check_allows_request_when_no_limits(self): """Test that requests are allowed when no limits are set.""" - check = ModelRateLimitingCheck(dual_cache=MagicMock()) + check = _check() deployment = { "litellm_params": {"model": "gpt-4"}, @@ -110,78 +146,50 @@ class TestModelRateLimitingCheck: def test_pre_call_check_raises_rate_limit_error_when_over_rpm(self): """Test that RateLimitError is raised when RPM limit is exceeded.""" - mock_cache = MagicMock() - mock_cache.increment_cache.return_value = 11 # Over limit after increment - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } + check = _check(usage_store(FakeUsageHandle(redis=_this_minute(RPM_DEPLOYMENT, "rpm", 10)))) with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(deployment) + check.pre_call_check(RPM_DEPLOYMENT) assert "RPM limit=10" in str(exc_info.value) assert "current usage=11" in str(exc_info.value) def test_pre_call_check_allows_request_under_limit(self): """Test that requests are allowed when under the limit.""" - mock_cache = MagicMock() - mock_cache.increment_cache.return_value = 6 + check = _check(usage_store(FakeUsageHandle(redis=_this_minute(RPM_DEPLOYMENT, "rpm", 5)))) - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - result = check.pre_call_check(deployment) - assert result == deployment + assert check.pre_call_check(RPM_DEPLOYMENT) == RPM_DEPLOYMENT def test_pre_call_check_raises_rate_limit_error_when_over_tpm(self): """Test that RateLimitError is raised when TPM limit is exceeded.""" - mock_cache = MagicMock() - mock_cache.get_cache.return_value = 1000 # Already at limit - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "tpm": 1000, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } + check = _check(usage_store(FakeUsageHandle(local=_this_minute(TPM_DEPLOYMENT, "tpm", 1000)))) with pytest.raises(litellm.RateLimitError) as exc_info: - check.pre_call_check(deployment) + check.pre_call_check(TPM_DEPLOYMENT) assert "TPM limit=1000" in str(exc_info.value) assert "current usage=1000" in str(exc_info.value) def test_pre_call_check_rejects_when_shared_tpm_is_over_limit_but_local_is_under(self): - redis_cache = MagicMock() - redis_cache.get_cache.return_value = 1000 - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + handle = FakeUsageHandle( + local=_this_minute(TPM_DEPLOYMENT, "tpm", 5), redis=_this_minute(TPM_DEPLOYMENT, "tpm", 1000) + ) + check = _check(usage_store(handle)) with pytest.raises(litellm.RateLimitError) as exc_info: check.pre_call_check(TPM_DEPLOYMENT) assert "current usage=1000" in str(exc_info.value) - @pytest.mark.parametrize( - "redis_get", [MagicMock(return_value=None), MagicMock(side_effect=RedisCircuitBreakerOpenError())] - ) - def test_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, redis_get): - redis_cache = MagicMock() - redis_cache.get_cache = redis_get - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache)) + @pytest.mark.parametrize("failure", ["redis_down", "breaker_open"]) + def test_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, failure: str): + handle = FakeUsageHandle( + local=_this_minute(TPM_DEPLOYMENT, "tpm", 1000), + redis={}, + breaker_open=failure == "breaker_open", + redis_down=failure == "redis_down", + ) + check = _check(usage_store(handle)) with pytest.raises(litellm.RateLimitError) as exc_info: check.pre_call_check(TPM_DEPLOYMENT) @@ -189,10 +197,12 @@ class TestModelRateLimitingCheck: assert "current usage=1000" in str(exc_info.value) def test_pre_call_check_falls_back_to_local_tpm_and_still_checks_rpm_when_redis_circuit_is_open(self): - redis_cache = MagicMock() - redis_cache.get_cache.side_effect = RedisCircuitBreakerOpenError() - redis_cache.increment_cache.return_value = 2 - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + handle = FakeUsageHandle( + local={**_this_minute(TPM_DEPLOYMENT, "tpm", 5), **_this_minute(TPM_DEPLOYMENT, "rpm", 1)}, + redis={}, + breaker_open=True, + ) + check = _check(usage_store(handle)) with pytest.raises(litellm.RateLimitError) as exc_info: check.pre_call_check({**TPM_DEPLOYMENT, "rpm": 1}) @@ -200,7 +210,7 @@ class TestModelRateLimitingCheck: assert "RPM limit=1" in str(exc_info.value) def test_pre_call_check_without_redis_enforces_local_tpm_and_rpm(self): - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None)) + check = _check(usage_store(FakeUsageHandle(local=_this_minute(TPM_DEPLOYMENT, "tpm", 5)))) deployment = {**TPM_DEPLOYMENT, "rpm": 1} assert check.pre_call_check(deployment) == deployment @@ -211,22 +221,12 @@ class TestModelRateLimitingCheck: def test_log_success_event_increments_cache(self): """Test that log_success_event correctly increments the cache.""" - mock_cache = MagicMock() - check = ModelRateLimitingCheck(dual_cache=mock_cache) + handle = FakeUsageHandle() + check = _check(usage_store(handle)) - kwargs = { - "standard_logging_object": { - "model_id": "test-id", - "total_tokens": 50, - "hidden_params": {"litellm_model_name": "gpt-4"}, - } - } + check.log_success_event(_SUCCESS_KWARGS, None, None, None) - check.log_success_event(kwargs, None, None, None) - - keys = tuple(call.kwargs["key"] for call in mock_cache.increment_cache.call_args_list) - assert any(key.startswith("test-id:gpt-4:tpm:") for key in keys), keys - assert all(call.kwargs["value"] == 50 for call in mock_cache.increment_cache.call_args_list) + assert _tpm_increments(handle) == ((50, 60),) class TestModelRateLimitingCheckAsync: @@ -235,10 +235,7 @@ class TestModelRateLimitingCheckAsync: @pytest.mark.asyncio async def test_async_pre_call_check_allows_request_when_no_limits(self): """Test that requests are allowed when no limits are set (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - - check = ModelRateLimitingCheck(dual_cache=mock_cache) + check = _check() deployment = { "litellm_params": {"model": "gpt-4"}, @@ -251,68 +248,36 @@ class TestModelRateLimitingCheckAsync: @pytest.mark.asyncio async def test_async_pre_call_check_raises_rate_limit_error_when_over_rpm(self): """Test that RateLimitError is raised when RPM limit is exceeded (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache_pipeline = AsyncMock(return_value=[11.0]) # Over limit - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } + check = _check(usage_store(FakeUsageHandle(redis=_this_minute(RPM_DEPLOYMENT, "rpm", 10)))) with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(deployment) + await check.async_pre_call_check(RPM_DEPLOYMENT) assert "RPM limit=10" in str(exc_info.value) @pytest.mark.asyncio async def test_async_pre_call_check_allows_request_under_limit(self): """Test that requests are allowed when under the limit (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache_pipeline = AsyncMock(return_value=[6.0]) + check = _check(usage_store(FakeUsageHandle(redis=_this_minute(RPM_DEPLOYMENT, "rpm", 5)))) - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "rpm": 10, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } - - result = await check.async_pre_call_check(deployment) - assert result == deployment + assert await check.async_pre_call_check(RPM_DEPLOYMENT) == RPM_DEPLOYMENT @pytest.mark.asyncio async def test_async_pre_call_check_raises_rate_limit_error_when_over_tpm(self): """Test that RateLimitError is raised when TPM limit is exceeded (async).""" - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=1000) # Already at limit - - check = ModelRateLimitingCheck(dual_cache=mock_cache) - - deployment = { - "tpm": 1000, - "litellm_params": {"model": "gpt-4"}, - "model_info": {"id": "test-id"}, - "model_name": "test-model", - } + check = _check(usage_store(FakeUsageHandle(local=_this_minute(TPM_DEPLOYMENT, "tpm", 1000)))) with pytest.raises(litellm.RateLimitError) as exc_info: - await check.async_pre_call_check(deployment) + await check.async_pre_call_check(TPM_DEPLOYMENT) assert "TPM limit=1000" in str(exc_info.value) @pytest.mark.asyncio async def test_async_pre_call_check_rejects_when_shared_tpm_is_over_limit_but_local_is_under(self): - redis_cache = MagicMock() - redis_cache.async_get_cache = AsyncMock(return_value=1000) - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + handle = FakeUsageHandle( + local=_this_minute(TPM_DEPLOYMENT, "tpm", 5), redis=_this_minute(TPM_DEPLOYMENT, "tpm", 1000) + ) + check = _check(usage_store(handle)) with pytest.raises(litellm.RateLimitError) as exc_info: await check.async_pre_call_check(TPM_DEPLOYMENT) @@ -320,13 +285,15 @@ class TestModelRateLimitingCheckAsync: assert "current usage=1000" in str(exc_info.value) @pytest.mark.asyncio - @pytest.mark.parametrize( - "redis_get", [AsyncMock(return_value=None), AsyncMock(side_effect=RedisCircuitBreakerOpenError())] - ) - async def test_async_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, redis_get): - redis_cache = MagicMock() - redis_cache.async_get_cache = redis_get - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(1000, redis_cache)) + @pytest.mark.parametrize("failure", ["redis_down", "breaker_open"]) + async def test_async_pre_call_check_keeps_rejecting_on_local_usage_when_redis_read_fails(self, failure: str): + handle = FakeUsageHandle( + local=_this_minute(TPM_DEPLOYMENT, "tpm", 1000), + redis={}, + breaker_open=failure == "breaker_open", + redis_down=failure == "redis_down", + ) + check = _check(usage_store(handle)) with pytest.raises(litellm.RateLimitError) as exc_info: await check.async_pre_call_check(TPM_DEPLOYMENT) @@ -337,10 +304,12 @@ class TestModelRateLimitingCheckAsync: async def test_async_pre_call_check_falls_back_to_local_tpm_and_still_checks_rpm_when_redis_circuit_is_open( self, ): - redis_cache = MagicMock() - redis_cache.async_get_cache = AsyncMock(side_effect=RedisCircuitBreakerOpenError()) - redis_cache.async_increment_pipeline = AsyncMock(return_value=[2.0]) - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, redis_cache)) + handle = FakeUsageHandle( + local={**_this_minute(TPM_DEPLOYMENT, "tpm", 5), **_this_minute(TPM_DEPLOYMENT, "rpm", 1)}, + redis={}, + breaker_open=True, + ) + check = _check(usage_store(handle)) with pytest.raises(litellm.RateLimitError) as exc_info: await check.async_pre_call_check({**TPM_DEPLOYMENT, "rpm": 1}) @@ -349,7 +318,7 @@ class TestModelRateLimitingCheckAsync: @pytest.mark.asyncio async def test_async_pre_call_check_without_redis_enforces_local_tpm_and_rpm(self): - check = ModelRateLimitingCheck(dual_cache=_dual_cache_with_local_tpm(5, None)) + check = _check(usage_store(FakeUsageHandle(local=_this_minute(TPM_DEPLOYMENT, "tpm", 5)))) deployment = {**TPM_DEPLOYMENT, "rpm": 1} assert await check.async_pre_call_check(deployment) == deployment @@ -361,23 +330,12 @@ class TestModelRateLimitingCheckAsync: @pytest.mark.asyncio async def test_async_log_success_event_increments_cache(self): """Test that async_log_success_event correctly increments the cache.""" - mock_cache = MagicMock() - mock_cache.async_increment_cache = AsyncMock() - check = ModelRateLimitingCheck(dual_cache=mock_cache) + handle = FakeUsageHandle(redis={}) + check = _check(usage_store(handle)) - kwargs = { - "standard_logging_object": { - "model_id": "test-id", - "total_tokens": 50, - "hidden_params": {"litellm_model_name": "gpt-4"}, - } - } + await check.async_log_success_event(_SUCCESS_KWARGS, None, None, None) - await check.async_log_success_event(kwargs, None, None, None) - - keys = tuple(call.kwargs["key"] for call in mock_cache.async_increment_cache.call_args_list) - assert any(key.startswith("test-id:gpt-4:tpm:") for key in keys), keys - assert all(call.kwargs["value"] == 50 for call in mock_cache.async_increment_cache.call_args_list) + assert _tpm_increments(handle) == ((50, 60),) class TestRouterWithEnforceModelRateLimits: @@ -441,8 +399,7 @@ class TestModelRateLimitConcurrency: the old check-then-increment pattern would let 3+ through due to a race condition on the local cache read. """ - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + check = _check() deployment = { "rpm": 2, @@ -473,6 +430,8 @@ RPD_DEPLOYMENT = { "model_name": "daily-model", } RPM_AND_RPD_DEPLOYMENT = {**RPD_DEPLOYMENT, "litellm_params": {"model": "gpt-4.1-mini", "rpm": 1, "rpd": 2}} +_DAY_ONE_RPM: Final = UsageCounter.for_deployment("daily-test-id", "gpt-4.1-mini", "rpm", "23-59") +_DAY_ONE_RPD: Final = UsageCounter.for_deployment("daily-test-id", "gpt-4.1-mini", "rpd", "2026-10-28") class Clock: @@ -497,7 +456,7 @@ async def _async_log_success(check: ModelRateLimitingCheck, kwargs: dict[str, ob class TestDailyLimits: def test_pre_call_check_rejects_the_request_over_rpd(self): - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) + check = _check(clock=Clock(DAY_ONE)) assert _pre_call(check, RPD_DEPLOYMENT) == RPD_DEPLOYMENT assert _pre_call(check, RPD_DEPLOYMENT) == RPD_DEPLOYMENT @@ -509,7 +468,7 @@ class TestDailyLimits: @pytest.mark.asyncio async def test_async_pre_call_check_rejects_the_request_over_rpd_and_lets_the_router_retry(self): - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) + check = _check(clock=Clock(DAY_ONE)) assert await _async_pre_call(check, RPD_DEPLOYMENT) == RPD_DEPLOYMENT assert await _async_pre_call(check, RPD_DEPLOYMENT) == RPD_DEPLOYMENT @@ -522,7 +481,7 @@ class TestDailyLimits: @pytest.mark.asyncio async def test_async_pre_call_check_rejects_once_the_day_tokens_reach_tpd(self): deployment = {**RPD_DEPLOYMENT, "litellm_params": {"model": "gpt-4.1-mini", "tpd": 100}} - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) + check = _check(clock=Clock(DAY_ONE)) kwargs = { "standard_logging_object": { "model_id": "daily-test-id", @@ -541,7 +500,7 @@ class TestDailyLimits: def test_pre_call_check_minute_rejected_request_does_not_spend_rpd(self): clock = Clock(DAY_ONE) - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + check = _check(clock=clock) assert _pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT with pytest.raises(litellm.RateLimitError, match="RPM limit=1"): @@ -556,7 +515,7 @@ class TestDailyLimits: @pytest.mark.asyncio async def test_async_pre_call_check_minute_rejected_request_does_not_spend_rpd(self): clock = Clock(DAY_ONE) - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + check = _check(clock=clock) assert await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT with pytest.raises(litellm.RateLimitError, match="RPM limit=1"): @@ -570,7 +529,7 @@ class TestDailyLimits: def test_pre_call_check_rejected_requests_hand_back_rpm_and_rpd(self): clock = Clock(DAY_ONE) - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + check = _check(clock=clock) assert _pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT with pytest.raises(litellm.RateLimitError, match="RPM limit=1, current usage=2"): @@ -589,7 +548,7 @@ class TestDailyLimits: @pytest.mark.asyncio async def test_async_pre_call_check_rejected_requests_hand_back_rpm_and_rpd(self): clock = Clock(DAY_ONE) - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + check = _check(clock=clock) assert await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT with pytest.raises(litellm.RateLimitError, match="RPM limit=1, current usage=2"): @@ -607,53 +566,43 @@ class TestDailyLimits: @pytest.mark.asyncio async def test_async_pre_call_check_refunds_rpm_and_rpd_in_one_cache_trip(self): - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache_pipeline = AsyncMock(side_effect=[[1.0, 3.0], [0.0, 2.0]]) - check = ModelRateLimitingCheck(dual_cache=mock_cache, clock=Clock(DAY_ONE)) + handle = FakeUsageHandle(redis={_DAY_ONE_RPD: 2}) + check = _check(usage_store(handle), clock=Clock(DAY_ONE)) with pytest.raises(litellm.RateLimitError, match="RPD limit=2, current usage=3"): await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) - assert mock_cache.async_increment_cache_pipeline.await_count == 2 - refund = mock_cache.async_increment_cache_pipeline.call_args_list[1].kwargs["increment_list"] - assert [(op["key"], op["increment_value"], op["ttl"]) for op in refund] == [ - ("daily-test-id:gpt-4.1-mini:rpm:23-59", -1, 60), - ("daily-test-id:gpt-4.1-mini:rpd:2026-10-28", -1, 86400), + assert handle.calls == [ + ("pipeline", (CounterIncrement(_DAY_ONE_RPM, 1, 60), CounterIncrement(_DAY_ONE_RPD, 1, 86400))), + ("pipeline", (CounterIncrement(_DAY_ONE_RPM, -1, 60), CounterIncrement(_DAY_ONE_RPD, -1, 86400))), ] @pytest.mark.asyncio async def test_async_pre_call_check_increments_rpm_and_rpd_in_one_cache_trip(self): - mock_cache = MagicMock() - mock_cache.async_get_cache = AsyncMock(return_value=None) - mock_cache.async_increment_cache_pipeline = AsyncMock(return_value=[1.0, 1.0]) - check = ModelRateLimitingCheck(dual_cache=mock_cache, clock=Clock(DAY_ONE)) + handle = FakeUsageHandle(redis={}) + check = _check(usage_store(handle), clock=Clock(DAY_ONE)) assert await _async_pre_call(check, RPM_AND_RPD_DEPLOYMENT) == RPM_AND_RPD_DEPLOYMENT - assert mock_cache.async_increment_cache_pipeline.await_count == 1 - operations = mock_cache.async_increment_cache_pipeline.call_args.kwargs["increment_list"] - assert [(op["key"], op["ttl"]) for op in operations] == [ - ("daily-test-id:gpt-4.1-mini:rpm:23-59", 60), - ("daily-test-id:gpt-4.1-mini:rpd:2026-10-28", 86400), + assert handle.calls == [ + ("pipeline", (CounterIncrement(_DAY_ONE_RPM, 1, 60), CounterIncrement(_DAY_ONE_RPD, 1, 86400))), ] - assert not mock_cache.async_increment_cache.called def test_pre_call_check_reads_the_shared_tpd_count_from_redis(self): - redis_cache = MagicMock() - redis_cache.get_cache.return_value = 100 + tpd: Final = UsageCounter.for_deployment("daily-test-id", "gpt-4.1-mini", "tpd", "2026-10-28") + handle = FakeUsageHandle(redis={tpd: 100}) deployment = {**RPD_DEPLOYMENT, "litellm_params": {"model": "gpt-4.1-mini", "tpd": 100}} - check = ModelRateLimitingCheck(dual_cache=DualCache(redis_cache=redis_cache), clock=Clock(DAY_ONE)) + check = _check(usage_store(handle), clock=Clock(DAY_ONE)) with pytest.raises(litellm.RateLimitError) as exc_info: _pre_call(check, deployment) assert "TPD limit=100, current usage=100" in str(exc_info.value) - assert redis_cache.get_cache.call_args.kwargs["key"] == "daily-test-id:gpt-4.1-mini:tpd:2026-10-28" + assert handle.calls[0] == ("current", (tpd, 100)) def test_rpd_count_starts_fresh_after_utc_midnight(self): clock = Clock(DAY_ONE) - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + check = _check(clock=clock) _pre_call(check, RPD_DEPLOYMENT) _pre_call(check, RPD_DEPLOYMENT) with pytest.raises(litellm.RateLimitError): @@ -677,7 +626,7 @@ class TestDailyLimits: ) def test_daily_429_retry_after_is_seconds_until_utc_midnight(self, now: datetime, seconds_to_midnight: int): deployment = {**RPD_DEPLOYMENT, "litellm_params": {"model": "gpt-4.1-mini", "rpd": 0}} - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(now)) + check = _check(clock=Clock(now)) with pytest.raises(litellm.RateLimitError) as exc_info: _pre_call(check, deployment) @@ -712,7 +661,7 @@ class TestDailyLimits: ], num_retries=1, ) - check = ModelRateLimitingCheck(dual_cache=router.cache, clock=Clock(DAY_ONE)) + check = _check(clock=Clock(DAY_ONE)) router.optional_callbacks = [check] litellm.logging_callback_manager.add_litellm_callback(check) try: @@ -728,18 +677,13 @@ class TestDailyLimits: assert served == ["a", "b", "b", "b"] -_LOGGED_TOKENS: Final = TypeAdapter[float | None](float | None) - - -async def _tpd_tokens_once_logged(cache: DualCache, key: str) -> int: +async def _tpd_tokens_once_logged(usage: RouterUsage, counter: UsageCounter) -> int: for _ in range(200): - logged: Final = _LOGGED_TOKENS.validate_python( - cache.in_memory_cache.get_cache(key) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # InMemoryCache.get_cache is untyped - ) + logged: Final = usage.get_local(counter) if logged is not None: return int(logged) await asyncio.sleep(0.01) - pytest.fail(f"the success callback never logged tpd tokens under {key}") + pytest.fail(f"the success callback never logged tpd tokens under {counter}") RESET_AT_EIGHT = datetime(2026, 10, 28, 8, 0, 0, tzinfo=timezone.utc) @@ -752,7 +696,7 @@ DAILY_RESET_DEPLOYMENT = { class TestDailyResetTime: def test_rpd_window_starts_over_at_daily_reset_time_not_midnight(self): clock = Clock(RESET_AT_EIGHT - timedelta(hours=9)) - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + check = _check(clock=clock) assert _pre_call(check, DAILY_RESET_DEPLOYMENT) == DAILY_RESET_DEPLOYMENT clock.now = RESET_AT_EIGHT - timedelta(hours=7) @@ -768,7 +712,7 @@ class TestDailyResetTime: @pytest.mark.asyncio async def test_async_rpd_window_starts_over_at_daily_reset_time_not_midnight(self): clock = Clock(RESET_AT_EIGHT - timedelta(hours=9)) - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + check = _check(clock=clock) assert await _async_pre_call(check, DAILY_RESET_DEPLOYMENT) == DAILY_RESET_DEPLOYMENT clock.now = RESET_AT_EIGHT - timedelta(hours=7) @@ -794,7 +738,7 @@ class TestDailyResetTime: **RPD_DEPLOYMENT, "litellm_params": {"model": "gpt-4.1-mini", "rpd": 0, "daily_reset_time": "08:00"}, } - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(now)) + check = _check(clock=Clock(now)) with pytest.raises(litellm.RateLimitError) as exc_info: _pre_call(check, deployment) @@ -808,7 +752,7 @@ class TestDailyResetTime: "litellm_params": {"model": "gpt-4.1-mini", "tpd": 100, "daily_reset_time": "08:00"}, } clock = Clock(RESET_AT_EIGHT - timedelta(minutes=1)) - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=clock) + check = _check(clock=clock) await _async_log_success( check, { @@ -834,7 +778,7 @@ class TestDailyResetTime: ], ) def test_rpd_outside_litellm_params_is_not_a_limit(self, deployment: dict[str, object]): - check = ModelRateLimitingCheck(dual_cache=DualCache(), clock=Clock(DAY_ONE)) + check = _check(clock=Clock(DAY_ONE)) assert _pre_call(check, deployment) == deployment @@ -866,19 +810,22 @@ class TestDailyResetTime: } ] ) - check = ModelRateLimitingCheck(dual_cache=router.cache, clock=Clock(RESET_AT_EIGHT - timedelta(minutes=1))) + usage: Final = _memory_usage() + check = _check(usage, clock=Clock(RESET_AT_EIGHT - timedelta(minutes=1))) router.optional_callbacks = [check] litellm.logging_callback_manager.add_litellm_callback(check) try: await router.acompletion( # pyright: ignore[reportUnknownMemberType] # Router.acompletion kwargs are untyped model="daily-group", messages=[{"role": "user", "content": "hi"}], daily_reset_time="08:00" ) - logged = await _tpd_tokens_once_logged(router.cache, "daily-test-id:openai/gpt-4.1-mini:tpd:2026-10-28") + logged = await _tpd_tokens_once_logged( + usage, UsageCounter.for_deployment("daily-test-id", "openai/gpt-4.1-mini", "tpd", "2026-10-28") + ) finally: router.discard() assert logged > 0 - day_before: Final = _LOGGED_TOKENS.validate_python( - router.cache.in_memory_cache.get_cache("daily-test-id:openai/gpt-4.1-mini:tpd:2026-10-27") # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # InMemoryCache.get_cache is untyped + day_before: Final = usage.get_local( + UsageCounter.for_deployment("daily-test-id", "openai/gpt-4.1-mini", "tpd", "2026-10-27") ) assert day_before is None 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 bb4813a273b..196a5fc767a 100644 --- a/tests/unit/test_router/test_io_token_rate_limits.py +++ b/tests/unit/test_router/test_io_token_rate_limits.py @@ -5,30 +5,27 @@ Tests for separate ITPM/OTPM deployment rate limits (enforce_model_rate_limits). import asyncio import contextvars from collections.abc import AsyncIterator, Awaitable, Callable, Mapping +from dataclasses import asdict from datetime import datetime, timezone -from typing import Final -from unittest.mock import AsyncMock, MagicMock +from typing import Final, Literal import httpx import pytest import respx -from fakeredis import FakeRedis, FakeServer -from fakeredis.aioredis import FakeRedis as FakeAsyncRedis import litellm from litellm import Router -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.redis.stores.connection import connect +from litellm.redis.stores.router_usage import CounterIncrement, RouterUsageStore, UsageCounter from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import ( - ITPM_CACHE_KEY, + ITPM_COUNTER_KEY, ITPM_RESERVED_KEY, - OTPM_CACHE_KEY, + OTPM_COUNTER_KEY, OTPM_RESERVED_KEY, _explicit_max_tokens, _reservation_value, @@ -47,9 +44,9 @@ from litellm.router_utils.pre_call_checks.model_rate_limit_check import ( ModelRateLimitingCheck, ) from litellm.types.llms.openai import InputTokensDetails, ResponseAPIUsage, ResponsesAPIResponse -from litellm.types.router import RouterCacheEnum from litellm.types.utils import ModelResponse, Usage from litellm.utils import get_utc_datetime +from tests.unit.redis.stores.router_usage_double import FakeUsageHandle, usage_store MANTLE_OPUS = "bedrock_mantle/anthropic.claude-opus-4-7" MANTLE_OPUS_MAX_OUTPUT = litellm.get_model_info(model=MANTLE_OPUS)["max_output_tokens"] @@ -59,6 +56,33 @@ AZURE_BASE_MODEL = "azure/gpt-5" AZURE_BASE_MODEL_MAX_OUTPUT = litellm.get_model_info(model=AZURE_BASE_MODEL)["max_output_tokens"] +def _memory_usage() -> RouterUsageStore: + return RouterUsageStore(connect(None).router_usage()) + + +def _io(deployment_id: str, model: str, kind: Literal["itpm", "otpm"], window: str) -> UsageCounter: + return UsageCounter("global_router", deployment_id, model, window, kind) + + +class _FailingKind(FakeUsageHandle): + """A usage store whose increments of one counter kind fail, as a transient Redis error would.""" + + def __init__(self, failing: str) -> None: + super().__init__() + self.failing: Final = failing + + def increment_sync(self, counter: UsageCounter, value: float, ttl_seconds: int) -> float: + if counter.kind == self.failing: + raise RuntimeError("transient cache error") + return super().increment_sync(counter, value, ttl_seconds) + + +def _counter_in(value: object) -> UsageCounter: + """The counter a reservation stashed in the request metadata.""" + assert isinstance(value, dict) + return UsageCounter(**value) + + class TestIOTokenRateLimitHelpers: def test_deployment_has_io_token_limits(self): assert deployment_has_io_token_limits({"litellm_params": {"itpm": 100, "otpm": 50}}) @@ -110,8 +134,8 @@ class TestModelRateLimitingCheckIOTokens: async def test_itpm_reservation_and_reconcile(self): from litellm.utils import get_utc_datetime - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) deployment = { "litellm_params": { "model": "bedrock_mantle/anthropic.claude-opus-4-7", @@ -131,8 +155,8 @@ class TestModelRateLimitingCheckIOTokens: await check.async_pre_call_check(deployment) minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-test-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" - otpm_key = f"global_router:io-test-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" + itpm_key = _io("io-test-id", "bedrock_mantle/anthropic.claude-opus-4-7", "itpm", minute) + otpm_key = _io("io-test-id", "bedrock_mantle/anthropic.claude-opus-4-7", "otpm", minute) kwargs = { "standard_logging_object": { @@ -154,16 +178,16 @@ class TestModelRateLimitingCheckIOTokens: ) await check.async_log_success_event(kwargs, response, None, None) - current_itpm = await dual_cache.async_get_cache(key=itpm_key) - current_otpm = await dual_cache.async_get_cache(key=otpm_key) + current_itpm = usage.get_local(itpm_key) + current_otpm = usage.get_local(otpm_key) # ITPM tracks input tokens only (billable prompt tokens), not output. assert current_itpm == 5 assert current_otpm == 3 @pytest.mark.asyncio async def test_itpm_limit_raises_429(self): - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) deployment = { "litellm_params": { "model": "bedrock_mantle/anthropic.claude-opus-4-7", @@ -198,8 +222,8 @@ class TestModelRateLimitingCheckIOTokens: async def test_otpm_atomic_reservation_no_overshoot_under_concurrency(self): from litellm.utils import get_utc_datetime - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) otpm_limit = 10 max_tokens = 4 deployment = { @@ -231,8 +255,8 @@ class TestModelRateLimitingCheckIOTokens: successes = sum(1 for r in results if r) minute = get_utc_datetime().strftime("%H-%M") - otpm_key = f"global_router:io-otpm-race-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" - current_otpm = await dual_cache.async_get_cache(key=otpm_key) + otpm_key = _io("io-otpm-race-id", "bedrock_mantle/anthropic.claude-opus-4-7", "otpm", minute) + current_otpm = usage.get_local(otpm_key) # Atomic reservation must never let concurrent requests overshoot the limit. assert current_otpm is not None @@ -252,8 +276,8 @@ class TestModelRateLimitingCheckIOTokens: """ from litellm.utils import get_utc_datetime - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) itpm_limit = 5 deployment = { "litellm_params": { @@ -284,8 +308,8 @@ class TestModelRateLimitingCheckIOTokens: successes = sum(1 for r in results if r) minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-itpm-estimate-fail-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" - current_itpm = await dual_cache.async_get_cache(key=itpm_key) + itpm_key = _io("io-itpm-estimate-fail-id", "bedrock_mantle/anthropic.claude-opus-4-7", "itpm", minute) + current_itpm = usage.get_local(itpm_key) # A minimal 1-token reservation per request lets itpm_limit concurrent # requests through, instead of a single request starving the rest. @@ -297,11 +321,11 @@ class TestModelRateLimitingCheckIOTokens: async def test_reservation_read_prefers_top_level_metadata_over_litellm_params(self): from litellm.utils import get_utc_datetime - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-lp-id:bedrock_mantle/test:itpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60) + itpm_key = _io("io-lp-id", "bedrock_mantle/test", "itpm", minute) + usage.increment_local(CounterIncrement(itpm_key, 20, 60)) # Production kwargs commonly carry litellm_params.metadata; the stashed # reservation lives in the top-level metadata and must still be found. @@ -311,20 +335,20 @@ class TestModelRateLimitingCheckIOTokens: "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, "metadata": {}, }, - "metadata": {ITPM_RESERVED_KEY: 20, ITPM_CACHE_KEY: itpm_key}, + "metadata": {ITPM_RESERVED_KEY: 20, ITPM_COUNTER_KEY: asdict(itpm_key)}, "litellm_params": {"metadata": {"user_api_key_hash": "abc123"}}, } await check.async_log_failure_event(kwargs, None, None, None) - current = await dual_cache.async_get_cache(key=itpm_key) + current = usage.get_local(itpm_key) assert current == 0 @pytest.mark.asyncio async def test_reconcile_tracks_actual_usage_when_estimate_zero(self): from litellm.utils import get_utc_datetime - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) deployment = { "litellm_params": { "model": "bedrock_mantle/anthropic.claude-opus-4-7", @@ -339,10 +363,10 @@ class TestModelRateLimitingCheckIOTokens: await check.async_pre_call_check(deployment) minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-zero-est-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + itpm_key = _io("io-zero-est-id", "bedrock_mantle/anthropic.claude-opus-4-7", "itpm", minute) # A failed/zero estimate reserves a minimal 1 token, not the full # itpm limit, so it doesn't starve concurrent requests. - assert await dual_cache.async_get_cache(key=itpm_key) == 1 + assert usage.get_local(itpm_key) == 1 kwargs = { "standard_logging_object": { @@ -358,7 +382,7 @@ class TestModelRateLimitingCheckIOTokens: ) await check.async_log_success_event(kwargs, response, None, None) - assert await dual_cache.async_get_cache(key=itpm_key) == 7 + assert usage.get_local(itpm_key) == 7 @pytest.mark.asyncio async def test_zero_estimate_reserves_minimal_capacity_before_reconcile(self): @@ -367,8 +391,8 @@ class TestModelRateLimitingCheckIOTokens: full itpm limit, so up to itpm_limit such calls are allowed concurrently instead of the first one claiming the entire bucket. """ - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) deployment = { "litellm_params": { "model": "bedrock_mantle/anthropic.claude-opus-4-7", @@ -392,8 +416,8 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio async def test_explicit_zero_max_tokens_does_not_reserve_otpm(self): - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) deployment = { "litellm_params": { "model": "bedrock_mantle/anthropic.claude-opus-4-7", @@ -410,9 +434,9 @@ class TestModelRateLimitingCheckIOTokens: set_io_token_rate_limit_request_kwargs(zero_output_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) - zero_output_otpm_key = zero_output_kwargs["metadata"][OTPM_CACHE_KEY] + zero_output_otpm_key = _counter_in(zero_output_kwargs["metadata"][OTPM_COUNTER_KEY]) assert zero_output_kwargs["metadata"][OTPM_RESERVED_KEY] == 0 - assert (await dual_cache.async_get_cache(key=zero_output_otpm_key) or 0) == 0 + assert (usage.get_local(zero_output_otpm_key) or 0) == 0 normal_output_kwargs = { "messages": [{"role": "user", "content": "hello"}], @@ -422,8 +446,8 @@ class TestModelRateLimitingCheckIOTokens: set_io_token_rate_limit_request_kwargs(normal_output_kwargs, metadata_channel="metadata") await check.async_pre_call_check(deployment) - normal_output_otpm_key = normal_output_kwargs["metadata"][OTPM_CACHE_KEY] - assert await dual_cache.async_get_cache(key=normal_output_otpm_key) == 5 + normal_output_otpm_key = _counter_in(normal_output_kwargs["metadata"][OTPM_COUNTER_KEY]) + assert usage.get_local(normal_output_otpm_key) == 5 @pytest.mark.asyncio @pytest.mark.parametrize("use_sync_check", [False, True]) @@ -439,8 +463,8 @@ class TestModelRateLimitingCheckIOTokens: async def test_request_without_max_tokens_is_admitted_below_model_max_output( self, model: str, otpm_limit: int, use_sync_check: bool ): - dual_cache: Final = DualCache() - check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + usage: Final = _memory_usage() + check: Final = ModelRateLimitingCheck(usage=usage) deployment: Final = { "litellm_params": {"model": model, "otpm": otpm_limit}, "model_info": {"id": f"io-no-max-tokens-{otpm_limit}"}, @@ -457,7 +481,7 @@ class TestModelRateLimitingCheckIOTokens: reserved: Final = request_kwargs["metadata"][OTPM_RESERVED_KEY] assert isinstance(reserved, int) assert 0 < reserved < otpm_limit - assert await dual_cache.async_get_cache(key=request_kwargs["metadata"][OTPM_CACHE_KEY]) == reserved + assert usage.get_local(_counter_in(request_kwargs["metadata"][OTPM_COUNTER_KEY])) == reserved @pytest.mark.asyncio @pytest.mark.parametrize( @@ -467,8 +491,8 @@ class TestModelRateLimitingCheckIOTokens: async def test_concurrent_requests_without_max_tokens_share_otpm_without_overshoot( self, model: str, otpm_limit: int ): - dual_cache: Final = DualCache() - check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + usage: Final = _memory_usage() + check: Final = ModelRateLimitingCheck(usage=usage) deployment: Final = { "litellm_params": {"model": model, "otpm": otpm_limit}, "model_info": {"id": "io-no-max-tokens-race-id"}, @@ -485,12 +509,12 @@ class TestModelRateLimitingCheckIOTokens: return request_kwargs["metadata"] admitted: Final = [m for m in await asyncio.gather(*[_attempt() for _ in range(8)]) if m is not None] - otpm_key: Final = admitted[0][OTPM_CACHE_KEY] + otpm_key: Final = _counter_in(admitted[0][OTPM_COUNTER_KEY]) assert len(admitted) == round(1 / NO_MAX_TOKENS_TPM_RESERVATION_FRACTION) - assert {m[OTPM_CACHE_KEY] for m in admitted} == {otpm_key} - assert await dual_cache.async_get_cache(key=otpm_key) == sum(m[OTPM_RESERVED_KEY] for m in admitted) - assert await dual_cache.async_get_cache(key=otpm_key) <= otpm_limit + assert {_counter_in(m[OTPM_COUNTER_KEY]) for m in admitted} == {otpm_key} + assert usage.get_local(otpm_key) == sum(m[OTPM_RESERVED_KEY] for m in admitted) + assert usage.get_local(otpm_key) <= otpm_limit @pytest.mark.asyncio @pytest.mark.parametrize( @@ -500,8 +524,8 @@ class TestModelRateLimitingCheckIOTokens: async def test_request_without_max_tokens_reserves_model_max_output_when_otpm_allows( self, model: str, model_max_output: int ): - dual_cache: Final = DualCache() - check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + usage: Final = _memory_usage() + check: Final = ModelRateLimitingCheck(usage=usage) deployment: Final = { "litellm_params": {"model": model, "otpm": model_max_output * 10}, "model_info": {"id": "io-no-max-tokens-roomy-id"}, @@ -520,7 +544,7 @@ class TestModelRateLimitingCheckIOTokens: catalog_max_output: Final = litellm.get_model_info(model=backend_model)["max_output_tokens"] assert isinstance(catalog_max_output, int) otpm_limit: Final = int(max(catalog_max_output, 987) / NO_MAX_TOKENS_TPM_RESERVATION_FRACTION) + 1 - check: Final = ModelRateLimitingCheck(dual_cache=DualCache()) + check: Final = ModelRateLimitingCheck(usage=_memory_usage()) async def _reserved(deployment_id: str, model_info: dict[str, object]) -> int: metadata: Final[dict[str, object]] = {} @@ -547,7 +571,7 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio @pytest.mark.parametrize("use_sync_check", [False, True]) async def test_base_model_does_not_change_the_default_otpm_reservation(self, use_sync_check: bool): - check: Final = ModelRateLimitingCheck(dual_cache=DualCache()) + check: Final = ModelRateLimitingCheck(usage=_memory_usage()) otpm_limit: Final = AZURE_BASE_MODEL_MAX_OUTPUT * 10 async def _reserved(model_info: dict[str, str]) -> int: @@ -572,8 +596,8 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio @pytest.mark.parametrize("use_sync_check", [False, True]) async def test_unparseable_max_tokens_leaves_itpm_uncharged(self, use_sync_check: bool): - dual_cache: Final = DualCache() - check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + usage: Final = _memory_usage() + check: Final = ModelRateLimitingCheck(usage=usage) deployment: Final = { "litellm_params": {"model": MANTLE_OPUS, "itpm": 2000, "otpm": 2000}, "model_info": {"id": "io-bad-max-tokens-id"}, @@ -595,13 +619,13 @@ class TestModelRateLimitingCheckIOTokens: } await _pre_call(valid_kwargs) - itpm_counter: Final = await dual_cache.async_get_cache(key=valid_kwargs["metadata"][ITPM_CACHE_KEY]) + itpm_counter: Final = usage.get_local(_counter_in(valid_kwargs["metadata"][ITPM_COUNTER_KEY])) assert itpm_counter == valid_kwargs["metadata"][ITPM_RESERVED_KEY] @pytest.mark.asyncio async def test_explicit_max_tokens_above_default_share_is_reserved_in_full(self): - dual_cache: Final = DualCache() - check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) + usage: Final = _memory_usage() + check: Final = ModelRateLimitingCheck(usage=usage) deployment: Final = { "litellm_params": {"model": MANTLE_OPUS, "otpm": 100}, "model_info": {"id": "io-explicit-max-tokens-id"}, @@ -622,7 +646,7 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio @pytest.mark.parametrize("use_sync_check", [False, True]) async def test_deployment_max_tokens_sets_the_reservation_unless_the_request_sets_one(self, use_sync_check: bool): - check: Final = ModelRateLimitingCheck(dual_cache=DualCache()) + check: Final = ModelRateLimitingCheck(usage=_memory_usage()) deployment: Final = { "litellm_params": {"model": MANTLE_OPUS, "otpm": 4000, "max_tokens": 100}, "model_info": {"id": "io-deployment-max-tokens-id"}, @@ -643,7 +667,7 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio async def test_deployment_cap_sets_the_reservation_only_through_the_param_its_api_sends(self): - check: Final = ModelRateLimitingCheck(dual_cache=DualCache()) + check: Final = ModelRateLimitingCheck(usage=_memory_usage()) chat: Final = {"messages": [{"role": "user", "content": "hello"}]} responses: Final = {"input": "hello"} @@ -673,7 +697,7 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio @pytest.mark.parametrize("use_sync_check", [False, True]) async def test_deployment_cap_above_the_default_reserves_the_default(self, use_sync_check: bool): - check: Final = ModelRateLimitingCheck(dual_cache=DualCache()) + check: Final = ModelRateLimitingCheck(usage=_memory_usage()) otpm_limit: Final = UNMAPPED_MODEL_DEFAULT_MAX_OUTPUT * 5 async def reserved(deployment_cap: Mapping[str, int]) -> int: @@ -698,8 +722,8 @@ class TestModelRateLimitingCheckIOTokens: def test_sync_io_pre_call_reserves_and_reconciles(self): from litellm.utils import get_utc_datetime - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) deployment = { "litellm_params": { "model": "bedrock_mantle/anthropic.claude-opus-4-7", @@ -718,8 +742,8 @@ class TestModelRateLimitingCheckIOTokens: check.pre_call_check(deployment) minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-sync-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" - otpm_key = f"global_router:io-sync-id:bedrock_mantle/anthropic.claude-opus-4-7:otpm:{minute}" + itpm_key = _io("io-sync-id", "bedrock_mantle/anthropic.claude-opus-4-7", "itpm", minute) + otpm_key = _io("io-sync-id", "bedrock_mantle/anthropic.claude-opus-4-7", "otpm", minute) kwargs = { "standard_logging_object": { "model_id": "io-sync-id", @@ -734,15 +758,15 @@ class TestModelRateLimitingCheckIOTokens: ) check.log_success_event(kwargs, response, None, None) - assert dual_cache.get_cache(key=itpm_key) == 5 - assert dual_cache.get_cache(key=otpm_key) == 3 + assert usage.get_local(itpm_key) == 5 + assert usage.get_local(otpm_key) == 3 @pytest.mark.asyncio async def test_reconcile_runs_via_success_event_without_model_id(self): - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - itpm_key = "global_router:io-noid:bedrock_mantle/test:itpm:12-34" - await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) + itpm_key = _io("io-noid", "bedrock_mantle/test", "itpm", "12-34") + usage.increment_local(CounterIncrement(itpm_key, 10, 60)) # standard_logging_object has no model_id (only the TPM path needs it); # IO reconciliation must still run off the stashed cache key. @@ -751,7 +775,7 @@ class TestModelRateLimitingCheckIOTokens: "hidden_params": {"litellm_model_name": "bedrock_mantle/test"}, "metadata": {}, }, - "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, + "metadata": {ITPM_RESERVED_KEY: 10, ITPM_COUNTER_KEY: asdict(itpm_key)}, } response = ModelResponse( choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], @@ -759,20 +783,20 @@ class TestModelRateLimitingCheckIOTokens: ) await check.async_log_success_event(kwargs, response, None, None) - assert await dual_cache.async_get_cache(key=itpm_key) == 3 + assert usage.get_local(itpm_key) == 3 @pytest.mark.asyncio async def test_failure_clears_reservation_so_retry_is_not_poisoned(self): from litellm.utils import get_utc_datetime - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-first:bedrock_mantle/test:itpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) + itpm_key = _io("io-first", "bedrock_mantle/test", "itpm", minute) + usage.increment_local(CounterIncrement(itpm_key, 8, 60)) # Shared request metadata carrying the first (IO) deployment's reservation. - metadata = {ITPM_RESERVED_KEY: 8, ITPM_CACHE_KEY: itpm_key} + metadata = {ITPM_RESERVED_KEY: 8, ITPM_COUNTER_KEY: asdict(itpm_key)} fail_kwargs = { "metadata": metadata, "standard_logging_object": { @@ -783,9 +807,9 @@ class TestModelRateLimitingCheckIOTokens: } await check.async_log_failure_event(fail_kwargs, None, None, None) - assert await dual_cache.async_get_cache(key=itpm_key) == 0 + assert usage.get_local(itpm_key) == 0 assert ITPM_RESERVED_KEY not in metadata - assert ITPM_CACHE_KEY not in metadata + assert ITPM_COUNTER_KEY not in metadata # Retry succeeds on a non-IO fallback deployment reusing the same metadata. retry_kwargs = { @@ -804,10 +828,10 @@ class TestModelRateLimitingCheckIOTokens: await check.async_log_success_event(retry_kwargs, response, None, None) # The first deployment's ITPM counter is not driven negative... - assert await dual_cache.async_get_cache(key=itpm_key) == 0 + assert usage.get_local(itpm_key) == 0 # ...and the non-IO deployment's TPM usage is tracked normally. - tpm_key = f"non-io-second:openai/gpt-4o-mini:tpm:{minute}" - assert await dual_cache.async_get_cache(key=tpm_key) == 12 + tpm = UsageCounter.for_deployment("non-io-second", "openai/gpt-4o-mini", "tpm", minute) + assert usage.get_local(tpm) == 12 @pytest.mark.asyncio async def test_stale_reservation_refunded_before_retry_overwrites_it(self): @@ -823,67 +847,67 @@ class TestModelRateLimitingCheckIOTokens: """ from litellm.utils import get_utc_datetime - dual_cache = DualCache() + usage = _memory_usage() minute = get_utc_datetime().strftime("%H-%M") - itpm_key_a = f"global_router:io-retry-a:bedrock_mantle/test-a:itpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key_a, value=9, ttl=60) + itpm_key_a = _io("io-retry-a", "bedrock_mantle/test-a", "itpm", minute) + usage.increment_local(CounterIncrement(itpm_key_a, 9, 60)) # Deployment A's still-unreconciled reservation, stashed on the shared # kwargs dict the retry loop reuses. - shared_kwargs = {"metadata": {ITPM_RESERVED_KEY: 9, ITPM_CACHE_KEY: itpm_key_a}} + shared_kwargs = {"metadata": {ITPM_RESERVED_KEY: 9, ITPM_COUNTER_KEY: asdict(itpm_key_a)}} # Router calls this before overwriting kwargs for deployment B's attempt - # simulating the fix landing ahead of set_io_token_rate_limit_request_kwargs. - refund_stale_reservation_before_retry(dual_cache, shared_kwargs) + refund_stale_reservation_before_retry(usage, shared_kwargs) # A's reservation is refunded immediately, not left stranded for a # background failure task that may run arbitrarily later (or never, # if the sentinels get cleared out from under it first). - assert await dual_cache.async_get_cache(key=itpm_key_a) == 0 + assert usage.get_local(itpm_key_a) == 0 assert ITPM_RESERVED_KEY not in shared_kwargs["metadata"] - assert ITPM_CACHE_KEY not in shared_kwargs["metadata"] + assert ITPM_COUNTER_KEY not in shared_kwargs["metadata"] # A's own (now-late) failure event finds nothing left to refund and # is a safe no-op, since the sentinels were already cleared above. - io_token_refund_failure(dual_cache, shared_kwargs) - assert await dual_cache.async_get_cache(key=itpm_key_a) == 0 + io_token_refund_failure(usage, shared_kwargs) + assert usage.get_local(itpm_key_a) == 0 # The retry proceeds to stash deployment B's own reservation on the # same dict; it starts clean, unaffected by A's cleared sentinels. set_io_token_rate_limit_request_kwargs(shared_kwargs, metadata_channel="metadata") - itpm_key_b = f"global_router:io-retry-b:bedrock_mantle/test-b:itpm:{minute}" + itpm_key_b = _io("io-retry-b", "bedrock_mantle/test-b", "itpm", minute) shared_kwargs["metadata"][ITPM_RESERVED_KEY] = 4 - shared_kwargs["metadata"][ITPM_CACHE_KEY] = itpm_key_b - await dual_cache.async_increment_cache(key=itpm_key_b, value=4, ttl=60) - assert await dual_cache.async_get_cache(key=itpm_key_b) == 4 + shared_kwargs["metadata"][ITPM_COUNTER_KEY] = asdict(itpm_key_b) + usage.increment_local(CounterIncrement(itpm_key_b, 4, 60)) + assert usage.get_local(itpm_key_b) == 4 @pytest.mark.asyncio async def test_client_supplied_reservation_keys_are_stripped(self): # metadata is caller-controlled; the server-only reservation sentinels # must be removed before the router captures the request kwargs. forged = { - "metadata": {ITPM_RESERVED_KEY: 999999, ITPM_CACHE_KEY: "attacker:key:itpm:00-00"}, + "metadata": {ITPM_RESERVED_KEY: 999999, ITPM_COUNTER_KEY: asdict(_io("attacker", "key", "itpm", "00-00"))}, "litellm_metadata": {OTPM_RESERVED_KEY: 7}, - "litellm_params": {"metadata": {OTPM_CACHE_KEY: "attacker:key:otpm:00-00"}}, + "litellm_params": {"metadata": {OTPM_COUNTER_KEY: asdict(_io("attacker", "key", "otpm", "00-00"))}}, } set_io_token_rate_limit_request_kwargs(forged, metadata_channel="metadata") stored = get_io_token_rate_limit_request_kwargs() assert ITPM_RESERVED_KEY not in stored["metadata"] - assert ITPM_CACHE_KEY not in stored["metadata"] + assert ITPM_COUNTER_KEY not in stored["metadata"] assert OTPM_RESERVED_KEY not in stored["litellm_metadata"] - assert OTPM_CACHE_KEY not in stored["litellm_params"]["metadata"] + assert OTPM_COUNTER_KEY not in stored["litellm_params"]["metadata"] @pytest.mark.asyncio async def test_forged_reservation_cannot_decrement_counter(self): - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) - victim_key = "global_router:victim:model:itpm:00-00" - await dual_cache.async_increment_cache(key=victim_key, value=100, ttl=60) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) + victim_key = _io("victim", "model", "itpm", "00-00") + usage.increment_local(CounterIncrement(victim_key, 100, 60)) # A caller forges a reservation pointing at another deployment's counter. kwargs = { - "metadata": {ITPM_RESERVED_KEY: 100, ITPM_CACHE_KEY: victim_key}, + "metadata": {ITPM_RESERVED_KEY: 100, ITPM_COUNTER_KEY: asdict(victim_key)}, "standard_logging_object": { "model_id": "m", "hidden_params": {"litellm_model_name": "model"}, @@ -900,19 +924,13 @@ class TestModelRateLimitingCheckIOTokens: await check.async_log_success_event(kwargs, response, None, None) # The forged reservation was stripped, so the victim counter is untouched. - assert await dual_cache.async_get_cache(key=victim_key) == 100 + assert usage.get_local(victim_key) == 100 @pytest.mark.asyncio async def test_otpm_reservation_error_rolls_back_itpm(self): from litellm.utils import get_utc_datetime - class _OtpmFailCache(DualCache): - async def async_increment_cache(self, key, **kwargs): - if ":otpm:" in key: - raise RuntimeError("transient cache error") - return await super().async_increment_cache(key=key, **kwargs) - - dual_cache = _OtpmFailCache() + usage = usage_store(_FailingKind("otpm")) deployment = { "litellm_params": { "model": "bedrock_mantle/anthropic.claude-opus-4-7", @@ -932,39 +950,33 @@ class TestModelRateLimitingCheckIOTokens: ) with pytest.raises(RuntimeError): - await async_io_token_pre_call_check(dual_cache, deployment) + await async_io_token_pre_call_check(usage, deployment) minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-rollback-id:bedrock_mantle/anthropic.claude-opus-4-7:itpm:{minute}" + itpm_key = _io("io-rollback-id", "bedrock_mantle/anthropic.claude-opus-4-7", "itpm", minute) # A transient OTPM error must release the ITPM reservation, not leak it. - assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0 + assert (usage.get_local(itpm_key) or 0) == 0 @pytest.mark.asyncio async def test_reconcile_clears_stash_even_when_increment_errors(self): - class _ItpmFailCache(DualCache): - async def async_increment_cache(self, key, **kwargs): - if ":itpm:" in key: - raise RuntimeError("transient cache error") - return await super().async_increment_cache(key=key, **kwargs) - - dual_cache = _ItpmFailCache() - metadata = {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: "global_router:x:model:itpm:00-00"} + usage = usage_store(_FailingKind("itpm")) + metadata = {ITPM_RESERVED_KEY: 5, ITPM_COUNTER_KEY: asdict(_io("x", "model", "itpm", "00-00"))} kwargs = {"metadata": metadata} response = ModelResponse( choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], usage=Usage(prompt_tokens=3, completion_tokens=0, total_tokens=3), ) - await async_io_token_reconcile_success(dual_cache, kwargs, response) + await async_io_token_reconcile_success(usage, kwargs, response) assert ITPM_RESERVED_KEY not in metadata - assert ITPM_CACHE_KEY not in metadata + assert ITPM_COUNTER_KEY not in metadata @pytest.mark.asyncio async def test_io_conflict_warning_not_collapsed_for_missing_model_id(self, caplog): import logging - check = ModelRateLimitingCheck(dual_cache=DualCache()) + check = ModelRateLimitingCheck(usage=_memory_usage()) deployment = { "litellm_params": {"model": "openai/gpt-4o-mini", "itpm": 100, "tpm": 1000}, "model_info": {}, @@ -979,7 +991,7 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio async def test_missing_deployment_id_skips_io_reservation(self): - dual_cache = DualCache() + usage = _memory_usage() deployment = { "litellm_params": {"model": "openai/gpt-4o-mini", "itpm": 100}, "model_info": {}, # no id -> cannot build a per-deployment cache key @@ -992,7 +1004,7 @@ class TestModelRateLimitingCheckIOTokens: } set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") - result = await async_io_token_pre_call_check(dual_cache, deployment) + result = await async_io_token_pre_call_check(usage, deployment) assert result is deployment # No reservation is stashed, so nothing lands in a shared None:None bucket. @@ -1000,32 +1012,33 @@ class TestModelRateLimitingCheckIOTokens: @pytest.mark.asyncio async def test_reconcile_uses_reservation_minute_key(self): - dual_cache = DualCache() + usage = _memory_usage() # Reservation was made on a fixed minute key; a call that finishes in a # later minute must reconcile against that same key, never a key built # from the response-time minute. - itpm_key = "global_router:io-min-id:bedrock_mantle/test:itpm:99-99" - await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) + itpm_key = _io("io-min-id", "bedrock_mantle/test", "itpm", "99-99") + usage.increment_local(CounterIncrement(itpm_key, 10, 60)) - kwargs = {"metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}} + kwargs = {"metadata": {ITPM_RESERVED_KEY: 10, ITPM_COUNTER_KEY: asdict(itpm_key)}} response = ModelResponse( choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], usage=Usage(prompt_tokens=4, completion_tokens=0, total_tokens=4), ) - await async_io_token_reconcile_success(dual_cache, kwargs, response) + await async_io_token_reconcile_success(usage, kwargs, response) - assert await dual_cache.async_get_cache(key=itpm_key) == 4 + assert usage.get_local(itpm_key) == 4 @pytest.mark.asyncio @pytest.mark.parametrize("use_sync_reconcile", [False, True]) async def test_reconcile_charges_usage_above_the_reservation_to_the_minute_it_finished_in( self, use_sync_reconcile: bool ): - dual_cache: Final = DualCache() - counter: Final = "global_router:io-late-id:bedrock_mantle/test:otpm" - reserved_minute_key: Final = f"{counter}:99-99" - await dual_cache.async_increment_cache(key=reserved_minute_key, value=10, ttl=60) - kwargs: Final = {"metadata": {ITPM_RESERVED_KEY: 0, OTPM_RESERVED_KEY: 10, OTPM_CACHE_KEY: reserved_minute_key}} + usage: Final = _memory_usage() + reserved_minute_key: Final = _io("io-late-id", "bedrock_mantle/test", "otpm", "99-99") + usage.increment_local(CounterIncrement(reserved_minute_key, 10, 60)) + kwargs: Final = { + "metadata": {ITPM_RESERVED_KEY: 0, OTPM_RESERVED_KEY: 10, OTPM_COUNTER_KEY: asdict(reserved_minute_key)} + } response: Final = ModelResponse( choices=[{"message": {"role": "assistant", "content": "hi"}, "index": 0, "finish_reason": "stop"}], usage=Usage(prompt_tokens=4, completion_tokens=25, total_tokens=29), @@ -1033,39 +1046,41 @@ class TestModelRateLimitingCheckIOTokens: minute_before: Final = get_utc_datetime().strftime("%H-%M") if use_sync_reconcile: - io_token_reconcile_success(dual_cache, kwargs, response) + io_token_reconcile_success(usage, kwargs, response) else: - await async_io_token_reconcile_success(dual_cache, kwargs, response) + await async_io_token_reconcile_success(usage, kwargs, response) finished_minutes: Final = {minute_before, get_utc_datetime().strftime("%H-%M")} - charged: Final = [await dual_cache.async_get_cache(key=f"{counter}:{minute}") for minute in finished_minutes] - assert await dual_cache.async_get_cache(key=reserved_minute_key) == 10 + charged: Final = [ + usage.get_local(_io("io-late-id", "bedrock_mantle/test", "otpm", minute)) for minute in finished_minutes + ] + assert usage.get_local(reserved_minute_key) == 10 assert sum(value or 0 for value in charged) == 25 - 10 @pytest.mark.asyncio async def test_reconcile_missing_usage_keeps_reservation(self): - dual_cache = DualCache() - itpm_key = "global_router:io-missing-usage:bedrock_mantle/test:itpm:00-00" - otpm_key = "global_router:io-missing-usage:bedrock_mantle/test:otpm:00-00" - await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) - await dual_cache.async_increment_cache(key=otpm_key, value=5, ttl=60) + usage = _memory_usage() + itpm_key = _io("io-missing-usage", "bedrock_mantle/test", "itpm", "00-00") + otpm_key = _io("io-missing-usage", "bedrock_mantle/test", "otpm", "00-00") + usage.increment_local(CounterIncrement(itpm_key, 8, 60)) + usage.increment_local(CounterIncrement(otpm_key, 5, 60)) kwargs = { "metadata": { ITPM_RESERVED_KEY: 8, OTPM_RESERVED_KEY: 5, - ITPM_CACHE_KEY: itpm_key, - OTPM_CACHE_KEY: otpm_key, + ITPM_COUNTER_KEY: asdict(itpm_key), + OTPM_COUNTER_KEY: asdict(otpm_key), } } response = ModelResponse( choices=[{"message": {"role": "assistant", "content": "ok"}, "index": 0, "finish_reason": "stop"}], ) - await async_io_token_reconcile_success(dual_cache, kwargs, response) + await async_io_token_reconcile_success(usage, kwargs, response) - assert await dual_cache.async_get_cache(key=itpm_key) == 8 - assert await dual_cache.async_get_cache(key=otpm_key) == 5 + assert usage.get_local(itpm_key) == 8 + assert usage.get_local(otpm_key) == 5 assert ITPM_RESERVED_KEY not in kwargs["metadata"] @pytest.mark.asyncio @@ -1076,50 +1091,50 @@ class TestModelRateLimitingCheckIOTokens: same as missing usage: keep the reservation instead of resolving to (0, 0) and refunding it in full. """ - dual_cache = DualCache() - itpm_key = "global_router:io-total-only:bedrock_mantle/test:itpm:00-00" - otpm_key = "global_router:io-total-only:bedrock_mantle/test:otpm:00-00" - await dual_cache.async_increment_cache(key=itpm_key, value=8, ttl=60) - await dual_cache.async_increment_cache(key=otpm_key, value=5, ttl=60) + usage = _memory_usage() + itpm_key = _io("io-total-only", "bedrock_mantle/test", "itpm", "00-00") + otpm_key = _io("io-total-only", "bedrock_mantle/test", "otpm", "00-00") + usage.increment_local(CounterIncrement(itpm_key, 8, 60)) + usage.increment_local(CounterIncrement(otpm_key, 5, 60)) kwargs = { "metadata": { ITPM_RESERVED_KEY: 8, OTPM_RESERVED_KEY: 5, - ITPM_CACHE_KEY: itpm_key, - OTPM_CACHE_KEY: otpm_key, + ITPM_COUNTER_KEY: asdict(itpm_key), + OTPM_COUNTER_KEY: asdict(otpm_key), } } response = {"type": "message", "usage": {"total_tokens": 13}} - await async_io_token_reconcile_success(dual_cache, kwargs, response) + await async_io_token_reconcile_success(usage, kwargs, response) - assert await dual_cache.async_get_cache(key=itpm_key) == 8 - assert await dual_cache.async_get_cache(key=otpm_key) == 5 + assert usage.get_local(itpm_key) == 8 + assert usage.get_local(otpm_key) == 5 def test_reconcile_standard_logging_total_tokens_only_keeps_reservation(self): - dual_cache = DualCache() - itpm_key = "global_router:io-slo-total-only:bedrock_mantle/test:itpm:00-00" - dual_cache.set_cache(key=itpm_key, value=10, ttl=60) + usage = _memory_usage() + itpm_key = _io("io-slo-total-only", "bedrock_mantle/test", "itpm", "00-00") + usage.set_local(itpm_key, 10) kwargs = { - "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, + "metadata": {ITPM_RESERVED_KEY: 10, ITPM_COUNTER_KEY: asdict(itpm_key)}, "standard_logging_object": {"total_tokens": 4}, } response = {"type": "message", "role": "assistant", "content": []} - io_token_reconcile_success(dual_cache, kwargs, response) + io_token_reconcile_success(usage, kwargs, response) - assert dual_cache.get_cache(key=itpm_key) == 10 + assert usage.get_local(itpm_key) == 10 @pytest.mark.asyncio async def test_reconcile_falls_back_to_standard_logging_object(self): - dual_cache = DualCache() - itpm_key = "global_router:io-slo-fallback:bedrock_mantle/test:itpm:00-00" - await dual_cache.async_increment_cache(key=itpm_key, value=10, ttl=60) + usage = _memory_usage() + itpm_key = _io("io-slo-fallback", "bedrock_mantle/test", "itpm", "00-00") + usage.increment_local(CounterIncrement(itpm_key, 10, 60)) kwargs = { - "metadata": {ITPM_RESERVED_KEY: 10, ITPM_CACHE_KEY: itpm_key}, + "metadata": {ITPM_RESERVED_KEY: 10, ITPM_COUNTER_KEY: asdict(itpm_key)}, "standard_logging_object": { "prompt_tokens": 4, "completion_tokens": 0, @@ -1128,23 +1143,23 @@ class TestModelRateLimitingCheckIOTokens: } response = {"type": "message", "role": "assistant", "content": []} - await async_io_token_reconcile_success(dual_cache, kwargs, response) + await async_io_token_reconcile_success(usage, kwargs, response) - assert await dual_cache.async_get_cache(key=itpm_key) == 4 + assert usage.get_local(itpm_key) == 4 def test_sync_reconcile_anthropic_dict_usage_counts_uncached_input_once(self): - dual_cache = DualCache() - itpm_key = "global_router:io-anthropic:bedrock_mantle/test:itpm:00-00" - otpm_key = "global_router:io-anthropic:bedrock_mantle/test:otpm:00-00" - dual_cache.set_cache(key=itpm_key, value=10, ttl=60) - dual_cache.set_cache(key=otpm_key, value=4, ttl=60) + usage = _memory_usage() + itpm_key = _io("io-anthropic", "bedrock_mantle/test", "itpm", "00-00") + otpm_key = _io("io-anthropic", "bedrock_mantle/test", "otpm", "00-00") + usage.set_local(itpm_key, 10) + usage.set_local(otpm_key, 4) kwargs = { "metadata": { ITPM_RESERVED_KEY: 10, OTPM_RESERVED_KEY: 4, - ITPM_CACHE_KEY: itpm_key, - OTPM_CACHE_KEY: otpm_key, + ITPM_COUNTER_KEY: asdict(itpm_key), + OTPM_COUNTER_KEY: asdict(otpm_key), } } response = { @@ -1157,10 +1172,10 @@ class TestModelRateLimitingCheckIOTokens: }, } - io_token_reconcile_success(dual_cache, kwargs, response) + io_token_reconcile_success(usage, kwargs, response) - assert dual_cache.get_cache(key=itpm_key) == 3 + 4 - assert dual_cache.get_cache(key=otpm_key) == 2 + assert usage.get_local(itpm_key) == 3 + 4 + assert usage.get_local(otpm_key) == 2 @pytest.mark.asyncio async def test_io_and_tpm_rpm_limits_both_enforced_with_warning(self, caplog): @@ -1168,8 +1183,8 @@ class TestModelRateLimitingCheckIOTokens: from litellm.utils import get_utc_datetime - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) model_id = "io-mixed-id" deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" deployment = { @@ -1183,9 +1198,9 @@ class TestModelRateLimitingCheckIOTokens: } minute = get_utc_datetime().strftime("%H-%M") - rpm_key = f"{model_id}:{deployment_name}:rpm:{minute}" - itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" - await dual_cache.async_increment_cache(key=rpm_key, value=5, ttl=60) + rpm = UsageCounter.for_deployment(model_id, deployment_name, "rpm", minute) + itpm_key = _io(model_id, deployment_name, "itpm", minute) + usage.increment_local(CounterIncrement(rpm, 5, 60)) request_kwargs = { "messages": [{"role": "user", "content": "hi"}], @@ -1198,8 +1213,8 @@ class TestModelRateLimitingCheckIOTokens: with pytest.raises(litellm.RateLimitError): await check.async_pre_call_check(deployment) - assert await dual_cache.async_get_cache(key=rpm_key) == 5, "the rejected request hands its rpm count back" - assert (await dual_cache.async_get_cache(key=itpm_key) or 0) == 0 + assert usage.get_local(rpm) == 5, "the rejected request hands its rpm count back" + assert (usage.get_local(itpm_key) or 0) == 0 assert ITPM_RESERVED_KEY not in request_kwargs["metadata"] assert any("both limit types are enforced" in record.message for record in caplog.records) @@ -1212,17 +1227,17 @@ class TestModelRateLimitingCheckIOTokens: """ from litellm.utils import get_utc_datetime - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) model_id = "io-tpm-mixed-id" deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" - tpm_key = f"{model_id}:{deployment_name}:tpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key, value=5, ttl=60) + itpm_key = _io(model_id, deployment_name, "itpm", minute) + tpm = UsageCounter.for_deployment(model_id, deployment_name, "tpm", minute) + usage.increment_local(CounterIncrement(itpm_key, 5, 60)) kwargs = { - "metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_key}, + "metadata": {ITPM_RESERVED_KEY: 5, ITPM_COUNTER_KEY: asdict(itpm_key)}, "standard_logging_object": { "model_id": model_id, "total_tokens": 7, @@ -1237,24 +1252,24 @@ class TestModelRateLimitingCheckIOTokens: await check.async_log_success_event(kwargs, response, None, None) # ITPM reconciled down from the 5-token reservation to actual usage (3). - assert await dual_cache.async_get_cache(key=itpm_key) == 3 + assert usage.get_local(itpm_key) == 3 # TPM tracking must still run so the tpm/rpm pre-call path can enforce it. - assert await dual_cache.async_get_cache(key=tpm_key) == 7 + assert usage.get_local(tpm) == 7 def test_io_success_still_tracks_tpm_for_mixed_deployment_sync(self): from litellm.utils import get_utc_datetime - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) model_id = "io-tpm-mixed-sync-id" deployment_name = "bedrock_mantle/anthropic.claude-opus-4-7" minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:{model_id}:{deployment_name}:itpm:{minute}" - tpm_key = f"{model_id}:{deployment_name}:tpm:{minute}" - dual_cache.set_cache(key=itpm_key, value=5, ttl=60) + itpm_key = _io(model_id, deployment_name, "itpm", minute) + tpm = UsageCounter.for_deployment(model_id, deployment_name, "tpm", minute) + usage.set_local(itpm_key, 5) kwargs = { - "metadata": {ITPM_RESERVED_KEY: 5, ITPM_CACHE_KEY: itpm_key}, + "metadata": {ITPM_RESERVED_KEY: 5, ITPM_COUNTER_KEY: asdict(itpm_key)}, "standard_logging_object": { "model_id": model_id, "total_tokens": 7, @@ -1268,20 +1283,20 @@ class TestModelRateLimitingCheckIOTokens: check.log_success_event(kwargs, response, None, None) - assert dual_cache.get_cache(key=itpm_key) == 3 - assert dual_cache.get_cache(key=tpm_key) == 7 + assert usage.get_local(itpm_key) == 3 + assert usage.get_local(tpm) == 7 @pytest.mark.asyncio async def test_failure_refunds_itpm_reservation(self): from litellm.utils import get_utc_datetime - dual_cache = DualCache() - check = ModelRateLimitingCheck(dual_cache=dual_cache) + usage = _memory_usage() + check = ModelRateLimitingCheck(usage=usage) minute = get_utc_datetime().strftime("%H-%M") - itpm_key = f"global_router:io-refund-id:bedrock_mantle/test:itpm:{minute}" - await dual_cache.async_increment_cache(key=itpm_key, value=20, ttl=60) + itpm_key = _io("io-refund-id", "bedrock_mantle/test", "itpm", minute) + usage.increment_local(CounterIncrement(itpm_key, 20, 60)) - reservation = {ITPM_RESERVED_KEY: 20, ITPM_CACHE_KEY: itpm_key} + reservation = {ITPM_RESERVED_KEY: 20, ITPM_COUNTER_KEY: asdict(itpm_key)} kwargs = { "standard_logging_object": { "model_id": "io-refund-id", @@ -1292,7 +1307,7 @@ class TestModelRateLimitingCheckIOTokens: } await check.async_log_failure_event(kwargs, None, None, None) - current = await dual_cache.async_get_cache(key=itpm_key) + current = usage.get_local(itpm_key) assert current == 0 @@ -1317,17 +1332,14 @@ def _current_minute() -> str: return get_utc_datetime().strftime("%H-%M") -def _counter(router: Router, template: RouterCacheEnum, minutes: frozenset[str]) -> int: +def _counter(router: Router, kind: Literal["itpm", "otpm"], minutes: frozenset[str]) -> int: deployment: Final = router.model_list[0] + usage: Final = router._core._usage # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # the store the router counts io tokens in return sum( - router.cache.get_cache( - key=template.value.format( - id=deployment["model_info"]["id"], - model=deployment["litellm_params"]["model"], - current_minute=minute, - ) + int( + usage.get_local(_io(deployment["model_info"]["id"], deployment["litellm_params"]["model"], kind, minute)) + or 0 ) - or 0 for minute in minutes ) @@ -1336,7 +1348,7 @@ async def _io_counters_since(router: Router, first_minute: str) -> tuple[int, in await asyncio.sleep(0) await asyncio.wait_for(GLOBAL_LOGGING_WORKER.flush(), timeout=10) minutes: Final = frozenset({first_minute, _current_minute()}) - return _counter(router, RouterCacheEnum.ITPM, minutes), _counter(router, RouterCacheEnum.OTPM, minutes) + return _counter(router, "itpm", minutes), _counter(router, "otpm", minutes) async def _sdk_chat(router: Router, mock_response: str | Exception) -> tuple[int, int]: @@ -1537,12 +1549,12 @@ class TestContextSlotRetention: """ kwargs = { "messages": [{"role": "user", "content": "x" * 1000}], - "metadata": {ITPM_RESERVED_KEY: 999, ITPM_CACHE_KEY: "forged"}, + "metadata": {ITPM_RESERVED_KEY: 999, ITPM_COUNTER_KEY: asdict(_io("forged", "model", "itpm", "00-00"))}, } set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=False, metadata_channel="metadata") assert get_io_token_rate_limit_request_kwargs() is None assert ITPM_RESERVED_KEY not in kwargs["metadata"] - assert ITPM_CACHE_KEY not in kwargs["metadata"] + assert ITPM_COUNTER_KEY not in kwargs["metadata"] set_io_token_rate_limit_request_kwargs(kwargs, store_in_context=True, metadata_channel="metadata") assert get_io_token_rate_limit_request_kwargs() is kwargs @@ -1561,7 +1573,7 @@ async def test_custom_azure_deployment_name_reserves_itpm_on_its_base_models_tok "metadata": metadata, } set_io_token_rate_limit_request_kwargs(request_kwargs, metadata_channel="metadata") - await async_io_token_pre_call_check(DualCache(), deployment) + await async_io_token_pre_call_check(_memory_usage(), deployment) return metadata[ITPM_RESERVED_KEY] named = await reserved_input( @@ -1620,21 +1632,26 @@ async def _deliver_success( async def test_success_settles_a_reservation_held_only_in_litellm_params_metadata_exactly_once( deliveries: tuple[str, ...], ) -> None: - dual_cache: Final = DualCache() - check: Final = ModelRateLimitingCheck(dual_cache=dual_cache) - itpm_key: Final = "global_router:proxy-shape:openai/gpt-5-mini:itpm:12-34" - otpm_key: Final = "global_router:proxy-shape:openai/gpt-5-mini:otpm:12-34" - await dual_cache.async_increment_cache(key=itpm_key, value=40, ttl=60) - await dual_cache.async_increment_cache(key=otpm_key, value=4096, ttl=60) + usage: Final = _memory_usage() + check: Final = ModelRateLimitingCheck(usage=usage) + itpm_key: Final = _io("proxy-shape", "openai/gpt-5-mini", "itpm", "12-34") + otpm_key: Final = _io("proxy-shape", "openai/gpt-5-mini", "otpm", "12-34") + usage.increment_local(CounterIncrement(itpm_key, 40, 60)) + usage.increment_local(CounterIncrement(otpm_key, 4096, 60)) kwargs: Final = _proxy_success_kwargs( - {ITPM_RESERVED_KEY: 40, OTPM_RESERVED_KEY: 4096, ITPM_CACHE_KEY: itpm_key, OTPM_CACHE_KEY: otpm_key} + { + ITPM_RESERVED_KEY: 40, + OTPM_RESERVED_KEY: 4096, + ITPM_COUNTER_KEY: asdict(itpm_key), + OTPM_COUNTER_KEY: asdict(otpm_key), + } ) for delivery in deliveries: await _deliver_success(check, delivery, kwargs, _reply_using(prompt_tokens=7, completion_tokens=3)) - assert await dual_cache.async_get_cache(key=itpm_key) == 7 - assert await dual_cache.async_get_cache(key=otpm_key) == 3 + assert usage.get_local(itpm_key) == 7 + assert usage.get_local(otpm_key) == 3 _TURN_AT: Final = datetime(2026, 1, 1, 12, 34, 56, tzinfo=timezone.utc) @@ -1645,8 +1662,8 @@ def _turn_deployment(kind: str) -> dict[str, object]: return {"litellm_params": {"model": "openai/gpt-5-mini", kind: 100}, "model_info": {"id": f"ws-turn-{kind}"}} -def _turn_key(kind: str, minute: str = "12-34") -> str: - return f"global_router:ws-turn-{kind}:openai/gpt-5-mini:{kind}:{minute}" +def _turn_key(kind: Literal["itpm", "otpm"], minute: str = "12-34") -> UsageCounter: + return _io(f"ws-turn-{kind}", "openai/gpt-5-mini", kind, minute) @pytest.fixture @@ -1660,11 +1677,11 @@ def settled_on_the_turn_s_minute(monkeypatch: pytest.MonkeyPatch) -> None: @pytest.mark.parametrize("kind", ["itpm", "otpm"]) @pytest.mark.parametrize("used", [100, 101, 250]) async def test_turn_admission_refuses_once_the_minutes_tokens_reach_the_limit(kind: str, used: int) -> None: - dual_cache: Final = DualCache() - await dual_cache.async_increment_cache(key=_turn_key(kind), value=used, ttl=60) + usage: Final = _memory_usage() + usage.increment_local(CounterIncrement(_turn_key(kind), used, 60)) with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=100, current usage={used}"): - await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + await async_admit_io_token_turn(usage, _turn_deployment(kind), _TURN_AT) @pytest.mark.asyncio @@ -1672,19 +1689,18 @@ async def test_turn_admission_refuses_once_the_minutes_tokens_reach_the_limit(ki @pytest.mark.parametrize("kind", ["itpm", "otpm"]) @pytest.mark.parametrize("used", [0, 1, 99]) async def test_an_admitted_turn_reserves_nothing_and_counts_its_actual_tokens(kind: str, used: int) -> None: - dual_cache: Final = DualCache() - await dual_cache.async_increment_cache(key=_turn_key(kind), value=used, ttl=60) + usage: Final = _memory_usage() + usage.increment_local(CounterIncrement(_turn_key(kind), used, 60)) actual: Final = _ACTUAL_TOKENS - reservation: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) - counter_key: Final = reservation[ITPM_CACHE_KEY if kind == "itpm" else OTPM_CACHE_KEY] - assert isinstance(counter_key, str) - assert await dual_cache.async_get_cache(key=counter_key) == used + reservation: Final = await async_admit_io_token_turn(usage, _turn_deployment(kind), _TURN_AT) + counter_key: Final = _counter_in(reservation[ITPM_COUNTER_KEY if kind == "itpm" else OTPM_COUNTER_KEY]) + assert usage.get_local(counter_key) == used await async_io_token_reconcile_success( - dual_cache, _proxy_success_kwargs(reservation), _reply_using(actual["itpm"], actual["otpm"]) + usage, _proxy_success_kwargs(reservation), _reply_using(actual["itpm"], actual["otpm"]) ) - assert await dual_cache.async_get_cache(key=counter_key) == used + actual[kind] + assert usage.get_local(counter_key) == used + actual[kind] @pytest.mark.asyncio @@ -1698,7 +1714,7 @@ async def test_an_admitted_turn_reserves_nothing_and_counts_its_actual_tokens(ki async def test_a_turn_without_io_token_counters_is_admitted_with_nothing_to_settle( deployment: dict[str, object], ) -> None: - assert await async_admit_io_token_turn(DualCache(), deployment, _TURN_AT) == {} + assert await async_admit_io_token_turn(_memory_usage(), deployment, _TURN_AT) == {} @pytest.mark.asyncio @@ -1708,44 +1724,32 @@ async def test_a_turn_without_io_token_counters_is_admitted_with_nothing_to_sett async def test_a_socket_open_reserves_nothing_and_its_first_turn_counts_its_actual_tokens_on_its_own_minute( kind: str, used_by_others: int ) -> None: - dual_cache: Final = DualCache() - await dual_cache.async_increment_cache(key=_turn_key(kind), value=used_by_others, ttl=60) + usage: Final = _memory_usage() + usage.increment_local(CounterIncrement(_turn_key(kind), used_by_others, 60)) socket_open: Final[dict[str, object]] = {"metadata": {}} set_io_token_rate_limit_request_kwargs(socket_open, opens_responses_socket=True, metadata_channel="metadata") - assert await async_io_token_pre_call_check(dual_cache, _turn_deployment(kind)) == _turn_deployment(kind) - assert await dual_cache.async_get_cache(key=_turn_key(kind)) == used_by_others + assert await async_io_token_pre_call_check(usage, _turn_deployment(kind)) == _turn_deployment(kind) + assert usage.get_local(_turn_key(kind)) == used_by_others opened_with: Final = socket_open["metadata"] assert isinstance(opened_with, dict) await async_io_token_reconcile_success( - dual_cache, _proxy_success_kwargs(opened_with), _reply_using(_ACTUAL_TOKENS["itpm"], _ACTUAL_TOKENS["otpm"]) + usage, _proxy_success_kwargs(opened_with), _reply_using(_ACTUAL_TOKENS["itpm"], _ACTUAL_TOKENS["otpm"]) ) - assert await dual_cache.async_get_cache(key=_turn_key(kind)) == used_by_others + _ACTUAL_TOKENS[kind] - - -class _RedisCacheOnFakeServer(RedisCache): - def __init__(self, server: FakeServer) -> None: - super().__init__(host="io-token-turn-redis", connection_pool=FakeRedis(server=server).connection_pool) - self._fake_async_client: Final = FakeAsyncRedis(server=server) - - def init_async_client(self) -> FakeAsyncRedis: - return self._fake_async_client + assert usage.get_local(_turn_key(kind)) == used_by_others + _ACTUAL_TOKENS[kind] @pytest.mark.asyncio @pytest.mark.parametrize("kind", ["itpm", "otpm"]) async def test_a_turn_is_refused_on_this_instance_s_count_when_reading_the_shared_counter_fails( - kind: str, + kind: Literal["itpm", "otpm"], ) -> None: - server: Final = FakeServer() - dual_cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=_RedisCacheOnFakeServer(server)) - server.connected = False - await dual_cache.in_memory_cache.async_increment(key=_turn_key(kind), value=100) + usage: Final = usage_store(FakeUsageHandle(local={_turn_key(kind): 100}, redis={}, redis_down=True)) with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=100, current usage=100"): - await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + await async_admit_io_token_turn(usage, _turn_deployment(kind), _TURN_AT) @pytest.mark.asyncio @@ -1758,17 +1762,17 @@ async def test_a_turn_finishing_in_the_next_minute_counts_its_actual_tokens_on_t "litellm.router_utils.pre_call_checks.io_token_rate_limit_check.get_utc_datetime", lambda: datetime(2026, 1, 1, 12, 35, 10, tzinfo=timezone.utc), ) - dual_cache: Final = DualCache() - reservation: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + usage: Final = _memory_usage() + reservation: Final = await async_admit_io_token_turn(usage, _turn_deployment(kind), _TURN_AT) await _deliver_success( - ModelRateLimitingCheck(dual_cache=dual_cache), + ModelRateLimitingCheck(usage=usage), delivery, _proxy_success_kwargs(reservation), _reply_using(_ACTUAL_TOKENS["itpm"], _ACTUAL_TOKENS["otpm"]), ) - assert [await dual_cache.async_get_cache(key=_turn_key(kind, minute)) for minute in ("12-34", "12-35")] == [ + assert [usage.get_local(_turn_key(kind, minute)) for minute in ("12-34", "12-35")] == [ None, _ACTUAL_TOKENS[kind], ] @@ -1779,14 +1783,11 @@ async def test_a_turn_finishing_in_the_next_minute_counts_its_actual_tokens_on_t async def test_a_turn_is_admitted_once_the_shared_counter_drops_below_the_limit_this_instance_last_read( kind: str, ) -> None: - redis_cache: Final = _RedisCacheOnFakeServer(FakeServer()) - dual_cache: Final = DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache) - await dual_cache.in_memory_cache.async_increment(key=_turn_key(kind), value=100) - await redis_cache.async_increment(key=_turn_key(kind), value=99) + usage: Final = usage_store(FakeUsageHandle(local={_turn_key(kind): 100}, redis={_turn_key(kind): 99})) - admitted: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + admitted: Final = await async_admit_io_token_turn(usage, _turn_deployment(kind), _TURN_AT) - assert admitted[ITPM_CACHE_KEY if kind == "itpm" else OTPM_CACHE_KEY] == _turn_key(kind) + assert _counter_in(admitted[ITPM_COUNTER_KEY if kind == "itpm" else OTPM_COUNTER_KEY]) == _turn_key(kind) @pytest.mark.asyncio @@ -1822,26 +1823,22 @@ async def test_a_turn_is_admitted_once_the_shared_counter_drops_below_the_limit_ ) @pytest.mark.usefixtures("settled_on_the_turn_s_minute") async def test_a_turn_counts_only_its_uncached_responses_input_toward_itpm(responses_reply: object) -> None: - dual_cache: Final = DualCache() - await dual_cache.async_increment_cache(key=_turn_key("itpm"), value=10, ttl=60) + usage: Final = _memory_usage() + usage.increment_local(CounterIncrement(_turn_key("itpm"), 10, 60)) - reservation: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment("itpm"), _TURN_AT) - await async_io_token_reconcile_success(dual_cache, _proxy_success_kwargs(reservation), responses_reply) + reservation: Final = await async_admit_io_token_turn(usage, _turn_deployment("itpm"), _TURN_AT) + await async_io_token_reconcile_success(usage, _proxy_success_kwargs(reservation), responses_reply) - assert await dual_cache.async_get_cache(key=_turn_key("itpm")) == 10 + 50 - 30 - - -def _shared_counter_answering(answer: int | Exception) -> DualCache: - redis_cache: Final = MagicMock(spec=RedisCache) - redis_cache.async_get_cache = AsyncMock(side_effect=[answer]) - return DualCache(in_memory_cache=InMemoryCache(), redis_cache=redis_cache) + assert usage.get_local(_turn_key("itpm")) == 10 + 50 - 30 @pytest.mark.asyncio @pytest.mark.parametrize("kind", ["itpm", "otpm"]) async def test_a_later_turn_is_refused_when_another_instance_filled_the_shared_minute(kind: str) -> None: with pytest.raises(litellm.RateLimitError, match=f"{kind.upper()} limit=100, current usage=100"): - await async_admit_io_token_turn(_shared_counter_answering(100), _turn_deployment(kind), _TURN_AT) + await async_admit_io_token_turn( + usage_store(FakeUsageHandle(redis={_turn_key(kind): 100})), _turn_deployment(kind), _TURN_AT + ) @pytest.mark.asyncio @@ -1849,12 +1846,11 @@ async def test_a_later_turn_is_refused_when_another_instance_filled_the_shared_m async def test_a_later_turn_is_admitted_on_this_instance_s_count_when_the_shared_counter_is_unreachable( kind: str, ) -> None: - dual_cache: Final = _shared_counter_answering(RedisCircuitBreakerOpenError("circuit open")) - await dual_cache.in_memory_cache.async_increment(key=_turn_key(kind), value=99) + usage: Final = usage_store(FakeUsageHandle(local={_turn_key(kind): 99}, redis={}, breaker_open=True)) - admitted: Final = await async_admit_io_token_turn(dual_cache, _turn_deployment(kind), _TURN_AT) + admitted: Final = await async_admit_io_token_turn(usage, _turn_deployment(kind), _TURN_AT) - assert admitted[ITPM_CACHE_KEY if kind == "itpm" else OTPM_CACHE_KEY] == _turn_key(kind) + assert _counter_in(admitted[ITPM_COUNTER_KEY if kind == "itpm" else OTPM_COUNTER_KEY]) == _turn_key(kind) _HELD_DEPLOYMENT: Final = { @@ -1864,58 +1860,59 @@ _HELD_DEPLOYMENT: Final = { } -def _held_keys() -> tuple[str, str]: +def _held_keys() -> tuple[UsageCounter, UsageCounter]: 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}", + _io("held-id", ROUTER_IO_MODEL, "itpm", minute), + _io("held-id", ROUTER_IO_MODEL, "otpm", minute), ) -class _HeldIncrement(DualCache): - """Holds the first increment matching ``holds`` until ``release`` is set.""" +class _HeldIncrement(FakeUsageHandle): + """A connected usage store holding the first increment matching ``holds`` until ``release`` is set.""" - def __init__(self, holds: Callable[[str, int], bool]) -> None: - super().__init__() + def __init__(self, holds: Callable[[UsageCounter, float], bool]) -> None: + super().__init__(redis={}) 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)): + async def increment(self, counter: UsageCounter, value: float, ttl_seconds: int) -> float: + if not self.holding.is_set() and self._holds(counter, value): self.holding.set() await self.release.wait() - return await super().async_increment_cache(key=key, value=value, **kwargs) + return self.increment_sync(counter, value, ttl_seconds) -async def _counts(cache: DualCache) -> tuple[int, int]: +async def _counts(usage: RouterUsageStore) -> 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) + itpm: Final = usage.get_local(itpm_key) + otpm: Final = usage.get_local(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() +async def _cancel_twice_while_held(held_store: _HeldIncrement, held: "asyncio.Task[object]") -> None: + await held_store.holding.wait() held.cancel() held.cancel() with pytest.raises(asyncio.CancelledError): await held - cache.release.set() + held_store.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) + held_store: Final = _HeldIncrement(lambda counter, value: counter.kind == "otpm" and value > 0) + cache: Final = usage_store(held_store) 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())) + await _cancel_twice_while_held(held_store, asyncio.create_task(reserve())) assert await _counts(cache) == (0, 0) assert ITPM_RESERVED_KEY not in request["metadata"] @@ -1927,40 +1924,46 @@ def _stashed(itpm_reserved: int, otpm_reserved: int) -> dict[str, object]: "metadata": { ITPM_RESERVED_KEY: itpm_reserved, OTPM_RESERVED_KEY: otpm_reserved, - ITPM_CACHE_KEY: itpm_key, - OTPM_CACHE_KEY: otpm_key, + ITPM_COUNTER_KEY: asdict(itpm_key), + OTPM_COUNTER_KEY: asdict(otpm_key), } } -async def _reserved(cache: DualCache, itpm_reserved: int, otpm_reserved: int) -> dict[str, object]: +async def _reserved(cache: RouterUsageStore, 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) + await cache.increment(itpm_key, itpm_reserved, 60) + await cache.increment(otpm_key, otpm_reserved, 60) 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) + held_store: Final = _HeldIncrement(lambda counter, value: counter.kind == "itpm" and value < 0) + cache: Final = usage_store(held_store) 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)) + held_store, 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: +class _YieldingIncrement(FakeUsageHandle): + """A connected usage store whose increments yield to the event loop first.""" + + def __init__(self) -> None: + super().__init__(redis={}) + + async def increment(self, counter: UsageCounter, value: float, ttl_seconds: int) -> float: await asyncio.sleep(0) - return await super().async_increment_cache(key=key, value=value, **kwargs) + return self.increment_sync(counter, value, ttl_seconds) @pytest.mark.asyncio async def test_competing_failure_callbacks_hand_a_reservation_back_once() -> None: - cache: Final = _YieldingIncrement() + cache: Final = usage_store(_YieldingIncrement()) kwargs: Final = await _reserved(cache, 9, 100) check: Final = ModelRateLimitingCheck(cache) @@ -1971,7 +1974,7 @@ async def test_competing_failure_callbacks_hand_a_reservation_back_once() -> Non @pytest.mark.asyncio async def test_a_late_failure_callback_of_one_attempt_leaves_the_next_attempts_reservation_held() -> None: - cache: Final = DualCache() + cache: Final = _memory_usage() check: Final = ModelRateLimitingCheck(cache) request: Final[dict[str, object]] = {"messages": PROXY_MESSAGES, "max_tokens": 100, "metadata": {}} @@ -2056,7 +2059,7 @@ async def test_a_cancelled_request_sharing_one_logging_object_hands_back_every_r @pytest.mark.asyncio async def test_a_reservation_refunded_through_a_logged_copy_of_its_metadata_is_refunded_once() -> None: - cache: Final = DualCache() + cache: Final = _memory_usage() check: Final = ModelRateLimitingCheck(cache) request: Final[dict[str, object]] = {"messages": PROXY_MESSAGES, "max_tokens": 100, "litellm_metadata": {}} @@ -2077,7 +2080,7 @@ async def test_a_reservation_refunded_through_a_logged_copy_of_its_metadata_is_r @pytest.mark.asyncio async def test_a_retry_refunds_the_reservation_its_previous_attempt_left_behind() -> None: - cache: Final = DualCache() + cache: Final = _memory_usage() request: Final[dict[str, object]] = {"messages": PROXY_MESSAGES, "max_tokens": 100, "metadata": {}} async def first_attempt() -> None: @@ -2129,11 +2132,11 @@ def _stashed_as(reservation_id: str, itpm_reserved: int, otpm_reserved: int) -> @pytest.mark.asyncio async def test_a_settled_reservation_stays_settled_however_many_others_settle_after_it() -> None: - cache: Final = DualCache() + cache: Final = _memory_usage() 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) + await cache.increment(itpm_key, 9, 60) + await cache.increment(otpm_key, 100, 60) 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) @@ -2184,3 +2187,68 @@ async def test_a_stream_cancelled_before_its_first_chunk_hands_back_its_reservat await call assert await _io_counters_since(router, first_minute) == (0, 0) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("settled_on_the_turn_s_minute") +async def test_a_stashed_reservation_counts_under_mains_redis_keys() -> None: + usage: Final = _memory_usage() + deployment: Final = { + "litellm_params": {"model": "bedrock/claude", "itpm": 100, "otpm": 100}, + "model_info": {"id": "dep-1"}, + } + + reservation: Final = await async_admit_io_token_turn(usage, deployment, _TURN_AT) + await async_io_token_reconcile_success(usage, _proxy_success_kwargs(reservation), _reply_using(5, 2)) + + assert usage.snapshot() == { + "global_router:dep-1:bedrock/claude:itpm:12-34": 5, + "global_router:dep-1:bedrock/claude:otpm:12-34": 2, + } + + +def _otpm_limited_deployment() -> dict[str, object]: + return { + "litellm_params": {"model": MANTLE_OPUS, "otpm": 50}, + "model_info": {"id": "oversized-cap-id"}, + "model_name": "opus", + } + + +@pytest.mark.asyncio +async def test_an_output_cap_too_large_for_the_store_to_count_is_refused_instead_of_admitted() -> None: + check: Final = ModelRateLimitingCheck(usage=_memory_usage()) + set_io_token_rate_limit_request_kwargs({"max_tokens": 2**63, "metadata": {}}, metadata_channel="metadata") + + with pytest.raises(litellm.RateLimitError, match="OTPM limit=50"): + await check.async_pre_call_check(_otpm_limited_deployment()) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("asynchronous", [True, False], ids=["async", "sync"]) +async def test_an_output_cap_past_the_limit_is_refused_without_touching_what_was_counted(asynchronous: bool) -> None: + usage: Final = _memory_usage() + check: Final = ModelRateLimitingCheck(usage=usage) + + async def pre_call(max_tokens: int) -> None: + set_io_token_rate_limit_request_kwargs({"max_tokens": max_tokens, "metadata": {}}, metadata_channel="metadata") + if asynchronous: + await check.async_pre_call_check(_otpm_limited_deployment()) + return + check.pre_call_check(_otpm_limited_deployment()) + + await pre_call(40) + counted: Final = usage.snapshot() + + with pytest.raises(litellm.RateLimitError, match="OTPM limit=50"): + await pre_call(2**63 - 1) + + assert usage.snapshot() == counted + + +def test_a_sync_output_cap_too_large_for_the_store_to_count_is_refused_instead_of_admitted() -> None: + check: Final = ModelRateLimitingCheck(usage=_memory_usage()) + set_io_token_rate_limit_request_kwargs({"max_tokens": 2**63, "metadata": {}}, metadata_channel="metadata") + + with pytest.raises(litellm.RateLimitError, match="OTPM limit=50"): + check.pre_call_check(_otpm_limited_deployment()) diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index d2b5d2b911b..1d9a07a4e3f 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -4,7 +4,7 @@ import json import os import time import weakref -from collections.abc import Awaitable, Callable, Iterator, Mapping +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Final @@ -17,16 +17,17 @@ from pydantic import BaseModel import litellm from litellm import Router -from litellm.caching.caching import DualCache, InMemoryCache from litellm.constants import DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS, ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.litellm_core_utils.get_litellm_params import get_litellm_params from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.redis.stores.connection import RedisStores +from litellm.redis.stores.router_usage import CounterIncrement, RouterUsageStore, UsageCounter from litellm.responses.streaming_iterator import ResponsesWsTurnAdmission, ResponsesWsTurnUsageRecorder from litellm.router import router_arguments -from litellm.router_config import RouterConfig, build_router_config +from litellm.router_config import build_router_config from litellm.router_python_host import PythonHost from litellm.router_strategy.adaptive_router.hooks import AdaptiveRouterPostCallHook from litellm.types.litellm_params import RoutingStrategyName @@ -40,11 +41,9 @@ from litellm.types.utils import ( StandardLoggingPayload, ) - if TYPE_CHECKING: - from opentelemetry.trace import Span - from tests.unit.conftest import RouterRedis + pass def create_standard_logging_payload() -> StandardLoggingPayload: @@ -475,25 +474,24 @@ async def test_usage_based_routing_counts_an_http_call_once(monkeypatch: pytest. await asyncio.sleep(0) await GLOBAL_LOGGING_WORKER.flush() - tpm_keys, _ = router.lowesttpm_logger_v2.usage_counter_keys(router.get_model_list(model_name="http-tpm-v2")) - assert await router.cache.async_get_cache(key=tpm_keys[0]) == response.usage.total_tokens + selector: Final = router.lowesttpm_logger_v2 + tpm_counters, _ = selector.usage_counters(router.get_model_list(model_name="http-tpm-v2")) + assert selector.usage.get_local(tpm_counters[0]) == response.usage.total_tokens -class _FirstTpmCountFails(DualCache): +class _FirstTpmCountFails(RouterUsageStore): failed: bool = False - async def async_increment_cache_post_call( - self, key: str, value: float, ttl: int | None, parent_otel_span: "Span | None" = None - ) -> None: - if self.failed or ":tpm:" not in key: - return await super().async_increment_cache_post_call(key, value, ttl, parent_otel_span) + async def increment_post_call(self, increments: Sequence[CounterIncrement]) -> None: + if self.failed or all(increment.counter.kind != "tpm" for increment in increments): + return await super().increment_post_call(increments) self.failed = True raise RuntimeError("routing cache unavailable") class _HostWhoseFirstTpmCountFails(PythonHost): - def routing_cache(self, config: RouterConfig) -> DualCache: - return _FirstTpmCountFails(in_memory_cache=InMemoryCache()) + def router_usage(self, redis: RedisStores) -> RouterUsageStore: + return _FirstTpmCountFails(redis.router_usage()) @pytest.mark.asyncio @@ -617,8 +615,9 @@ def _io_limited_router(kind: str, **litellm_params: object) -> Router: ) -def _minutes_counter_key(kind: str, minute: datetime, model: str = "openai/ws-io-unpriced") -> str: - return f"global_router:ws-io-dep:{model}:{kind}:{minute.strftime('%H-%M')}" +def _minutes_count(router: Router, kind: str, minute: datetime, model: str = "openai/ws-io-unpriced") -> object: + usage: Final = router._core._usage # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # the store the router counts io tokens in + return usage.get_local(UsageCounter("global_router", "ws-io-dep", model, minute.strftime("%H-%M"), kind)) _PROVIDER_RESPONSE: Final = { @@ -674,7 +673,7 @@ async def test_responses_call_sends_the_provider_only_the_client_metadata_and_se assert [json.loads(call.request.read()).get("metadata") for call in provider.calls] == [client_metadata] assert [ - await router.cache.async_get_cache(key=_minutes_counter_key(kind, router_minute_pinned)) + _minutes_count(router, kind, router_minute_pinned) for kind in ("itpm", "otpm") ] == [10, 5] @@ -781,7 +780,7 @@ async def test_retried_call_settles_deployment_io_tokens_to_the_successful_attem assert provider.call_count == 2 assert [ - await router.cache.async_get_cache(key=_minutes_counter_key(kind, router_minute_pinned, model)) + _minutes_count(router, kind, router_minute_pinned, model) for kind in ("itpm", "otpm") ] == [10, 5] @@ -1199,7 +1198,6 @@ async def test_get_model_group_io_token_usage_sums_across_deployments(): in the model group (not just the first), reading the same per-deployment cache keys the pre-call reservation writes to. """ - from litellm.types.router import RouterCacheEnum from litellm.utils import get_utc_datetime router = Router( @@ -1226,26 +1224,14 @@ async def test_get_model_group_io_token_usage_sums_across_deployments(): ) minute = get_utc_datetime().strftime("%H-%M") - keys_and_values = [ - ( - RouterCacheEnum.ITPM.value.format(id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute), - 30, - ), - ( - RouterCacheEnum.OTPM.value.format(id="io-usage-dep-1", model="openai/gpt-4o-mini", current_minute=minute), - 10, - ), - ( - RouterCacheEnum.ITPM.value.format(id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute), - 70, - ), - ( - RouterCacheEnum.OTPM.value.format(id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute), - 20, - ), - ] - for key, value in keys_and_values: - await router.cache.async_increment_cache(key=key, value=value, ttl=60) + usage = router._core._usage # pyright: ignore[reportPrivateUsage, reportAttributeAccessIssue] # the store the router counts io tokens in + for deployment_id, model, kind, value in ( + ("io-usage-dep-1", "openai/gpt-4o-mini", "itpm", 30), + ("io-usage-dep-1", "openai/gpt-4o-mini", "otpm", 10), + ("io-usage-dep-2", "openai/gpt-4o", "itpm", 70), + ("io-usage-dep-2", "openai/gpt-4o", "otpm", 20), + ): + usage.set_local(UsageCounter("global_router", deployment_id, model, minute, kind), value) current_itpm, current_otpm = await router.get_model_group_io_token_usage("opus")