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
This commit is contained in:
yujonglee 2026-11-13 12:22:28 +09:00 • committed by GitHub
parent e20f07d178
commit b88e4bff6b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
61 changed files with 3267 additions and 2053 deletions

View file

@ -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."""

View file

@ -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:

View file

@ -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,
}

View file

@ -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,

View file

@ -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()

View file

@ -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]: ...

View file

@ -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

View file

@ -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,

View file

@ -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"))

View file

@ -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"))

View file

@ -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

View file

@ -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:

View file

@ -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):

View file

@ -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(

View file

@ -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

View file

@ -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 ------------------------------------------------------------------------

View file

@ -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:

View file

@ -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)

View file

@ -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 `<id>:<model>:tpm:<HH-MM>` and `<id>:<model>:rpm:<HH-MM>` 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,

View file

@ -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(

View file

@ -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,
)

View file

@ -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

View file

@ -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)))

View file

@ -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<Self> {
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<T, F>(&self, py: Python<'_>, f: F) -> PyResult<T>

View file

@ -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",

View file

@ -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<Self> {
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<Self> {
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<Self> {
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<EventLoop>,
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<Self> {
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<Py<PyAny>>,
record_breaker: Arc<Py<PyAny>>,
context: PythonContext,
context: ReportContext,
/// The loop the client was built on, for calls whose caller had none running.
event_loop: Option<EventLoop>,
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<Py<PyAny>>, args: ReportArgs) {
fn report(&self, context: ReportContext, record: Arc<Py<PyAny>>, 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")]

View file

@ -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<RedisCall> {
Ok(RedisCall {
context: PythonContext::capture(py)?,
context: ReportContext::of_caller(py)?,
event_loop: EventLoop::running(py),
call_type,
})

View file

@ -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<CallTag> {
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<Self> {
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))
})))
}

View file

@ -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<Scope>,
deployment_id: String,
model: Option<String>,
minute: String,
window: String,
kind: Spelled<Kind>,
}
impl From<PyCounter> 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<PyCounter>) -> Vec<Counter> {
counters.into_iter().map(Counter::from).collect()
}
/// `litellm.redis.stores.router_usage.CounterIncrement`.
#[derive(FromPyObject)]
struct PyIncrement {
counter: PyCounter,
amount: PyNumber,
ttl: Option<f64>,
}
impl TryFrom<PyIncrement> for Increment {
type Error = PyErr;
fn try_from(increment: PyIncrement) -> PyResult<Self> {
Ok(Self {
counter: increment.counter.into(),
amount: increment.amount.0,
ttl: ttl(increment.ttl)?,
})
}
}
fn increments(increments: Vec<PyIncrement>) -> PyResult<Vec<Increment>> {
increments.into_iter().map(Increment::try_from).collect()
}
fn numbers(read: Vec<Option<litellm_redis_store::Number>>) -> Vec<Option<PyNumber>> {
read.into_iter().map(|count| count.map(PyNumber)).collect()
}
fn totals(totals: Option<Vec<litellm_redis_store::Number>>) -> Option<Vec<PyNumber>> {
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<RouterUsage>);
/// 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<NativePending>,
store: Arc<RouterUsage>,
reservation: Mutex<Option<UsageReservation>>,
}
#[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<PyCounter>,
send: &NativeSend,
) -> PyResult<NativeUsageRead> {
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<PyIncrement>,
) -> PyResult<Bound<'py, PyAny>> {
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<PyIncrement>,
) -> PyResult<()> {
self.0
.increment_post_call(&request.0, &self::increments(increments)?);
Ok(())
}
fn increment_remote<'py>(
&self,
py: Python<'py>,
increments: Vec<PyIncrement>,
) -> PyResult<Bound<'py, PyAny>> {
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<PyNumber> {
self.0
.increment_local(&increment.try_into()?)
.map(PyNumber)
.map_err(kind_error)
}
fn get_local(&self, counter: PyCounter) -> PyResult<Option<PyNumber>> {
self.0
.get_local(&counter.into())
.map(|count| count.map(PyNumber))
.map_err(kind_error)
}
fn get_local_many(&self, counters: Vec<PyCounter>) -> PyResult<Vec<Option<PyNumber>>> {
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<Bound<'py, PyAny>> {
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<Bound<'py, PyAny>> {
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<Option<PyNumber>> {
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)
}

View file

@ -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<String> {
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<Bound<'_, PyAny>>) -> PyResult<Cac
fn tokens_key(window: UsageWindow) -> String {
Counter {
scope: Scope::GlobalRouter,
deployment_id: window.deployment_id,
model: Some(window.model),
minute: window.minute,
window: window.minute,
kind: Kind::Tpm,
}
.key()

View file

@ -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<V: Clone + FromValue + std::marker::Send + Sync + 'static> DualCache<V, Bat
let reservation = self.reserve(keys, &self.read_memory(keys)?);
(!reservation.keys.is_empty()).then(|| {
let keys: Vec<&str> = reservation.keys.iter().map(String::as_str).collect();
let replies = send.batch(client).add(commands::mget::<Vec<u8>>(&keys));
let batch = send.batch(client);
let replies = batch.add(commands::mget::<Vec<u8>>(&keys));
Fetch {
reservation,
batch,
replies,
misses: send.misses(client),
}
@ -81,24 +83,37 @@ impl<V: Clone + FromValue + std::marker::Send + Sync + 'static> DualCache<V, Bat
})
}
/// The second half of `batch_get`. 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.
/// The second half of `batch_get`, flushing the read's batch if nobody has yet.
pub(crate) async fn resolve_batch_get(
&self,
declared: DeclaredRead,
) -> Result<Vec<Option<V>>, 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<Vec<Option<V>>, 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<V: Clone + FromValue + std::marker::Send + Sync + 'static> DualCache<V, Bat
Reservation { keys, previous }
}
fn release(&self, reservation: Reservation) {
pub(crate) fn release(&self, reservation: Reservation) {
let mut recent = lock(&self.reads.recent);
for (key, previous) in reservation.keys.into_iter().zip(reservation.previous) {
match previous {
@ -191,13 +206,27 @@ pub(crate) struct DeclaredRead {
fetch: Option<Fetch>,
}
impl DeclaredRead {
/// The batch carrying the read, `None` when memory answers it alone.
pub(crate) fn batch(&self) -> Option<RedisBatch> {
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<Reservation> {
self.fetch.as_ref().map(|fetch| fetch.reservation.clone())
}
}
struct Fetch {
reservation: Reservation,
batch: RedisBatch,
replies: Declared<Vec<Option<Vec<u8>>>>,
misses: Option<Misses>,
}
struct Reservation {
#[derive(Clone)]
pub(crate) struct Reservation {
keys: Vec<String>,
previous: Vec<Option<Duration>>,
}

View file

@ -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<V: Counter, R> DualCache<V, R> {
.collect::<Result<Vec<_>, _>>()
.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<Result<Vec<Number>, 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<V: Counter, R> DualCache<V, R> {
))
})
.collect::<Vec<_>>();
match Declared::all(totals)
.await
.into_iter()
.collect::<Result<Vec<_>, _>>()
{
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

View file

@ -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,

View file

@ -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};

View file

@ -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()),
}
}

View file

@ -22,7 +22,7 @@ struct Shared {
runtime: Handle,
logs: Dispatch,
post_call_deadline: Duration,
deadline_tag: Option<CallTag>,
deadline_tag: DeadlineTag,
batches: Mutex<Vec<ClientBatches>>,
follow_ups: Mutex<Vec<JoinHandle<()>>>,
deadline: Mutex<Option<oneshot::Sender<()>>>,
@ -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<CallTag>,
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<Output = ()> + std::marker::Send + 'static) {
drop(
self.0
.runtime
.spawn(self.logged(within(current_tag(), task))),
);
}
fn logged<F: Future>(&self, task: F) -> impl Future<Output = F::Output> + use<F> {
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<dyn Fn() -> Option<CallTag> + std::marker::Send + Sync>;
async fn within<F: Future>(tag: Option<CallTag>, 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<Misses> {
match self {

View file

@ -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<String>,
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}"),
}
}
}

View file

@ -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<Duration>,
}
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<PipelineIncrement> {
increments.iter().map(PipelineIncrement::from).collect()
}
/// The throttle refreshes a declared usage read reserved, which `RouterUsage::release` gives back.
pub struct UsageReservation(Option<Reservation>);
/// A usage read `RouterUsage::declare_get_many` declared, with the refreshes it reserved.
pub struct DeclaredUsage {
pub read: Pending<Vec<Option<Number>>>,
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<Number, BatchReads>,
}
@ -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<Vec<Option<Number>>, Error> {
let keys = counters.iter().map(Counter::key).collect::<Vec<_>>();
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<Self>,
send: Send<'_>,
counters: &[Counter],
) -> Result<DeclaredUsage, Error> {
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<Vec<Number>> {
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<Option<Vec<Number>>, 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<Number, Error> {
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<Option<Number>, Error> {
Ok(self.usage.memory().get(&counter.key())?)
}
/// `in_memory_cache.async_batch_get_cache`.
pub fn get_local_many(&self, counters: &[Counter]) -> Result<Vec<Option<Number>>, 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<Option<Number>, 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<Option<Number>, 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<Vec<(String, Number)>, 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::<Vec<u8>>(&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<Number>),
}
fn keys(counters: &[Counter]) -> Vec<String> {
counters.iter().map(Counter::key).collect()
}

View file

@ -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

View file

@ -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::<Kind>().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::<Scope>().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<String> = raw.send(commands::get(rpm.key)).await.unwrap();
request.flush_post().await;
let after_flush: Option<String> = 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<String> = 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)));
}

View file

@ -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)

View file

@ -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: ...

View file

@ -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,

View file

@ -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"""

View file

@ -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

View file

@ -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)

View file

@ -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)

View file

@ -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

View file

@ -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)

View file

@ -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(

View file

@ -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()

View file

@ -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}",

View file

@ -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": [

View file

@ -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 [

View file

@ -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)

View file

@ -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]]()

View file

@ -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:

View file

@ -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")

View file

@ -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

File diff suppressed because it is too large Load diff

View file

@ -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")