mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-13 08:52:02 +00:00
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:
parent
e20f07d178
commit
b88e4bff6b
61 changed files with 3267 additions and 2053 deletions
|
|
@ -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."""
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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]: ...
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"))
|
||||
|
|
|
|||
|
|
@ -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"))
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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):
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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 ------------------------------------------------------------------------
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
)
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)))
|
||||
|
|
|
|||
|
|
@ -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>
|
||||
|
|
|
|||
|
|
@ -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",
|
||||
|
|
|
|||
|
|
@ -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")]
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
})
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
})))
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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>>,
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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};
|
||||
|
|
|
|||
|
|
@ -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()),
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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 {
|
||||
|
|
|
|||
|
|
@ -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}"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)));
|
||||
}
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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: ...
|
||||
|
|
|
|||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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"""
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
151
tests/unit/redis/stores/router_usage_double.py
Normal file
151
tests/unit/redis/stores/router_usage_double.py
Normal 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
|
||||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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}",
|
||||
|
|
|
|||
|
|
@ -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": [
|
||||
|
|
|
|||
|
|
@ -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 [
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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]]()
|
||||
|
|
|
|||
|
|
@ -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:
|
||||
|
|
|
|||
|
|
@ -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")
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
|
@ -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")
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue