From b88e4bff6b8386eb7d9c6db362c7d3b7f17bd66a Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 13 Nov 2026 12:22:28 +0900 Subject: [PATCH] refactor(router): count the whole router_usage family on the native store (#50862) * refactor(router): count the whole router_usage family on the native store usage-based-routing-v2, the model rate limit check (tpm, rpm, tpd, rpd) and the io-token check (itpm, otpm) count through the native router_usage store, next to the router core's global_router tpm and rpm. The pick's usage read rides the native pre-call batch and settles on the batch reply, keeping DualCache's per-key read throttle, and the success tpm increment rides the native post-call batch after counting in memory. io-token reservations stash their counter's fields instead of a Redis key. Redis keys and values are unchanged. A Redis service event is reported in a context named for whose it is, either the calling store op's, the request whose post-call batch its deadline sends, or nobody's for breaker transitions and unscoped calls. The deadline's tag is built when a post-call declaration arms it, so that flush parents under the request's span even though auth sets the span after the batch scope opens. An SDK call anchors the span its caller passed, so its usage reads and io-token refunds parent under it without a span argument threaded through the family. * fix(redis): keep an int counter past i64 above the limit instead of wrapping Number::add summed two ints with plain i64 addition. With Redis off and the OTPM counter at 1, a request with max_tokens=2**63-1 passed the MAX_INCREMENT check, and the memory count wrapped negative in a release build, so it passed admission and let other requests past the cap. Python's int never overflows, so main sees a count above the limit, rolls it back and answers 429 An int sum past i64 now becomes the nearest float, which compares above any limit the way Python's int does. The rollback after such a refusal can leave the count one token low, since the float cannot hold 2**63 exactly * fix(logging): run every callback under the span an SDK caller passed Native Redis service spans take their parent from the request span contextvar. The proxy sets it at auth and the SDK router scopes it around the routed call, but a streamed response is iterated after that scope closes, so its callbacks saw no span. Three callbacks reopened the scope from kwargs by hand, and the sync stream path had no wrap at all because its success handler runs on an executor thread The four Logging handlers now anchor the span from model_call_details around their dispatch, so every CustomLogger callback sees it, streamed or not, sync or async, with or without a Router. The per-callback wraps in lowest_tpm_rpm_v2 and model_rate_limit_check are gone. A span already anchored, as on the proxy, is kept * test(logging): cover a mid-stream failure, and drop the span wrap in Router._release The callback span test now also breaks a stream after its first chunk, sync and async, through a custom provider, so the failure handlers the stream wrapper starts on a bare thread are covered too A probe through the Router showed Router._release only runs inside the routed attempt, where the router has already anchored the span, and never after a stream returns. Its own wrap was a no-op, so it is gone * fix(router): refuse an io-token reservation past its limit before counting it The overflow fix let a reservation near i64::MAX become a float, so the count still read above the limit and the request got its 429. But the rollback subtracted the same amount from that float, and with OTPM usage at 40 the count came back as 0.0, so later requests reused capacity already spent A reservation above its limit now gets the 429 before the store is touched, in the async and sync paths. On main the count went up and straight back down for such a reservation, so the answer and the count after it are the same as main's. Only the usage named in the error message differs, which now shows the reservation instead of the count plus the reservation * test(redis): send the failed SDK call's refund through a real completion The refund span test called ModelRateLimitingCheck.async_log_failure_event directly, so it relied on the span wrap that callback used to carry. The Logging failure handlers now anchor the span for every callback, so the test sends a failing litellm.acompletion with the reservation in its metadata and checks the refund's Redis span still lands under the caller's span --- litellm/_internal_context.py | 14 + litellm/caching/redis_batch.py | 11 +- .../internal_call_metadata.py | 8 +- litellm/litellm_core_utils/litellm_logging.py | 22 +- litellm/redis/stores/router_usage/__init__.py | 117 ++- litellm/redis/stores/router_usage/handle.py | 28 +- litellm/redis/stores/router_usage/types.py | 25 +- litellm/router_rust/core.py | 23 +- litellm/router_rust/invoke/endpoints.py | 27 +- .../router_rust/invoke/prompt_management.py | 2 +- litellm/router_rust/invoke/shaping.py | 8 +- litellm/router_rust/invoke/websocket.py | 27 +- litellm/router_rust/routing_stores.py | 2 +- litellm/router_rust/steps/callbacks.py | 5 +- litellm/router_rust/steps/selection.py | 15 +- litellm/router_rust/surface/legacy_members.py | 15 +- litellm/router_rust/surface/model_info.py | 51 +- .../router_strategy/base_routing_strategy.py | 105 +-- litellm/router_strategy/lowest_tpm_rpm_v2.py | 169 ++-- .../io_token_rate_limit_check.py | 310 +++----- .../pre_call_checks/model_rate_limit_check.py | 228 ++---- litellm/router_utils/routing_prefetch.py | 117 +-- litellm/router_utils/routing_read_batch.py | 29 +- litellm/rust/crates/host-python/src/gil.rs | 12 + litellm/rust/crates/python-bridge/src/lib.rs | 3 +- .../crates/python-bridge/src/redis/events.rs | 86 +- .../python-bridge/src/redis/store/mod.rs | 5 +- .../python-bridge/src/redis/store/request.rs | 18 +- .../src/redis/store/router_usage.rs | 198 ++++- .../crates/python-bridge/src/router/mod.rs | 20 +- .../redis-store/src/dual_cache/batch_read.rs | 47 +- .../redis-store/src/dual_cache/increment.rs | 44 +- .../crates/redis-store/src/dual_cache/mod.rs | 5 +- litellm/rust/crates/redis-store/src/lib.rs | 2 +- litellm/rust/crates/redis-store/src/number.rs | 7 +- .../rust/crates/redis-store/src/request.rs | 40 +- .../redis-store/src/router_usage/keys.rs | 39 +- .../redis-store/src/router_usage/mod.rs | 212 ++++- .../rust/crates/redis-store/tests/request.rs | 12 +- .../crates/redis-store/tests/router_usage.rs | 464 ++++++++++- .../crates/redis-store/tests/support/mod.rs | 34 +- litellm/rust_bridge/_native.pyi | 38 +- litellm/rust_bridge/host/redis_events.py | 11 +- litellm/types/router.py | 7 - .../test_request_redis_batch_post_call.py | 39 +- .../test_request_redis_batch_pre_call.py | 450 ++++------- .../test_litellm_logging.py | 151 +++- .../unit/redis/stores/router_usage_double.py | 151 ++++ tests/unit/redis/stores/test_router_usage.py | 38 +- .../test_responses_websocket_all_providers.py | 81 +- .../test_base_routing_strategy.py | 206 ++--- .../router_strategy/test_lowest_tpm_rpm.py | 29 +- .../router_strategy/test_lowest_tpm_rpm_v2.py | 28 +- .../test_router_routing_groups.py | 34 +- .../router_utils/test_routing_read_batch.py | 146 ++-- .../rust_bridge/host/test_redis_events.py | 90 +++ tests/unit/test_internal_context.py | 27 + ...st_registry_characterization_model_info.py | 15 +- .../test_enforce_model_rate_limits.py | 363 ++++----- .../test_router/test_io_token_rate_limits.py | 742 ++++++++++-------- .../test_router/test_router_helper_utils.py | 68 +- 61 files changed, 3267 insertions(+), 2053 deletions(-) create mode 100644 tests/unit/redis/stores/router_usage_double.py 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")