scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-05 07:37:50 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 360

refactor(redis): move global router usage behind a typed store

Browse source
This commit is contained in:
Devin AI 2026-10-30 20:53:38 +00:00
parent 371444b989
commit b82443abc7
8 changed files with 461 additions and 140 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

3
litellm/caching/dual_cache.py
View file

@ -336,6 +336,9 @@ class DualCache(BaseCache):
else:
self.last_redis_batch_access_time[key] = previous_time
def rollback_redis_batch_reads(self, previous_access_times: Mapping[str, float | None]) -> None:
self._rollback_redis_batch_key_reservations(dict(previous_access_times))
async def _prepare_batch_get(
self, keys: list[str], local_only: bool, throttle_redis: bool = True, **kwargs: object
) -> PendingBatchRead:

111
litellm/redis/stores/router_global_usage.py Normal file
View file

@ -0,0 +1,111 @@
import ast
from collections.abc import Sequence
from dataclasses import dataclass
from datetime import timedelta
from typing import Final, Literal, Protocol
from pydantic import JsonValue, TypeAdapter
from litellm.redis.transport.contract import RedisFailure, RedisTransport
TPM_KEY_TEMPLATE: Final = "global_router:{id}:{model}:tpm:{current_minute}"
RPM_KEY_TEMPLATE: Final = "global_router:{id}:{model}:rpm:{current_minute}"
_JSON: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue)
@dataclass(frozen=True, slots=True)
class GlobalUsageCounter:
deployment_id: str
model: str | None
minute: str
metric: Literal["tpm", "rpm"]
def usage_key(counter: GlobalUsageCounter) -> str:
template: Final = TPM_KEY_TEMPLATE if counter.metric == "tpm" else RPM_KEY_TEMPLATE
return template.format(id=counter.deployment_id, model=counter.model, current_minute=counter.minute)
def _decode(raw: bytes | None) -> JsonValue:
if raw is None:
return None
text: Final = raw.decode("utf-8")
try:
return _JSON.validate_json(text)
except ValueError:
return _JSON.validate_python(ast.literal_eval(text))
@dataclass(frozen=True, slots=True)
class GlobalUsageRead:
values: tuple[JsonValue | bytes, ...]
failure: RedisFailure | None = None
class RouterGlobalUsageStore(Protocol):
async def read_global_usage(self, counters: Sequence[GlobalUsageCounter]) -> GlobalUsageRead | RedisFailure: ...
async def record_global_usage(
self, deployment_id: str, model: str | None, minute: str, *, tokens: float, requests: int, ttl: timedelta
) -> tuple[float, float] | RedisFailure: ...
async def record_global_failure(
self, deployment_id: str, model: str | None, minute: str, *, ttl: timedelta
) -> float | RedisFailure: ...
@dataclass(frozen=True, slots=True)
class TransportRouterGlobalUsageStore:
transport: RedisTransport
async def read_global_usage(self, counters: Sequence[GlobalUsageCounter]) -> GlobalUsageRead | RedisFailure:
raw: Final = await self.transport.mget(tuple(usage_key(counter) for counter in counters))
if isinstance(raw, RedisFailure):
return raw
try:
return GlobalUsageRead(tuple(_decode(value) for value in raw))
except (ValueError, SyntaxError, UnicodeError) as error:
return GlobalUsageRead(raw, RedisFailure(kind="decode", message=str(error)))
async def record_global_usage(
self, deployment_id: str, model: str | None, minute: str, *, tokens: float, requests: int, ttl: timedelta
) -> tuple[float, float] | RedisFailure:
batch: Final = self.transport.post_call_batch() or self.transport.batch()
tpm: Final = usage_key(GlobalUsageCounter(deployment_id, model, minute, "tpm"))
rpm: Final = usage_key(GlobalUsageCounter(deployment_id, model, minute, "rpm"))
token_result: Final = batch.incrbyfloat(tpm, tokens)
token_expiry: Final = batch.expire_seconds(tpm, ttl)
request_result: Final = batch.incrbyfloat(rpm, requests)
request_expiry: Final = batch.expire_seconds(rpm, ttl)
token_value: Final = await token_result
token_expired: Final = await token_expiry
request_value: Final = await request_result
request_expired: Final = await request_expiry
if isinstance(token_value, RedisFailure):
return token_value
if isinstance(token_expired, RedisFailure):
return token_expired
if isinstance(request_value, RedisFailure):
return request_value
if isinstance(request_expired, RedisFailure):
return request_expired
return token_value, request_value
async def record_global_failure(
self, deployment_id: str, model: str | None, minute: str, *, ttl: timedelta
) -> float | RedisFailure:
key: Final = usage_key(GlobalUsageCounter(deployment_id, model, minute, "rpm"))
batch: Final = self.transport.batch()
increment: Final = batch.incrbyfloat(key, 1)
expiry: Final = batch.ttl(key)
value: Final = await increment
remaining: Final = await expiry
if isinstance(value, RedisFailure):
return value
if isinstance(remaining, RedisFailure):
return remaining
if remaining == -1:
expired: Final = await self.transport.expire_seconds(key, ttl)
if isinstance(expired, RedisFailure):
return expired
return value

159
litellm/router.py
View file

@ -58,7 +58,7 @@ from collections.abc import (
)
from contextvars import ContextVar
from dataclasses import dataclass
from datetime import datetime, timezone
from datetime import datetime, timedelta, timezone
from functools import lru_cache, partial
from itertools import chain
from types import MappingProxyType
@ -182,6 +182,13 @@ from litellm.llms.openai_like.model_info import (
MODEL_INFO_REFRESH_SECONDS,
get_openai_compatible_model_info,
)
from litellm.redis.stores.router_global_usage import (
GlobalUsageCounter,
RouterGlobalUsageStore,
TransportRouterGlobalUsageStore,
usage_key,
)
from litellm.redis.transport.contract import RedisFailure
from litellm.responses.bridge_check import keep_registered_mode
from litellm.responses.retrieve_stream_query import requests_streamed_retrieval, without_retrieve_stream_defaults
from litellm.router_api import (
@ -683,6 +690,7 @@ _DEPLOYMENT_PARAMS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | GenericLite
Mapping[str, object] | GenericLiteLLMParams
)
_OPAQUE_VALUE_ADAPTER: Final = TypeAdapter(object)
_GLOBAL_USAGE_LOCAL_VALUES: Final[TypeAdapter[list[object] | None]] = TypeAdapter(list[object] | None)
_DEPLOYMENT_SELECTED_EVENT: Final = "litellm.request.deployment_selected"
_CONTENT_BLOCKS_ADAPTER: Final = TypeAdapter(list[dict[str, object]])
_DEPLOYMENTS_ADAPTER: Final = TypeAdapter(Sequence[Mapping[str, object]])
@ -9861,25 +9869,34 @@ class Router(DeprecatedRouterMembers):
return None
current_minute: Final = get_utc_datetime().strftime("%H-%M") # use the same timezone regardless of system clock
tpm_key: Final = RouterCacheEnum.TPM.value.format(
id=deployment_id, current_minute=current_minute, model=deployment_name
)
rpm_key: Final = RouterCacheEnum.RPM.value.format(
id=deployment_id, current_minute=current_minute, model=deployment_name
)
tpm_key: Final = usage_key(GlobalUsageCounter(deployment_id, deployment_name, current_minute, "tpm"))
rpm_key: Final = usage_key(GlobalUsageCounter(deployment_id, deployment_name, current_minute, "rpm"))
pipeline_operations: Final[list[RedisPipelineIncrementOperation]] = [
RedisPipelineIncrementOperation(key=key, increment_value=increment_value, ttl=RoutingArgs.ttl.value)
for key, increment_value in ((tpm_key, total_tokens), (rpm_key, rpm_increment))
]
post_increment_values: Final = await self._core_cache.async_increment_cache_pipeline(
increment_list=pipeline_operations,
parent_otel_span=parent_otel_span,
local_values: Final = await self._core_cache.async_increment_cache_pipeline(
increment_list=pipeline_operations, local_only=True, parent_otel_span=parent_otel_span
)
if post_increment_values is not None and self._core_cache.redis_cache is not None:
for operation, value in zip(pipeline_operations, post_increment_values):
await self._core_cache.async_set_cache(
operation["key"], int(value), local_only=True, ttl=RoutingArgs.ttl.value
)
redis: Final = self._core_cache.redis_cache
if redis is not None:
store: Final[RouterGlobalUsageStore] = TransportRouterGlobalUsageStore(redis.transport())
shared_values: Final = await store.record_global_usage(
deployment_id,
deployment_name,
current_minute,
tokens=total_tokens,
requests=rpm_increment,
ttl=timedelta(seconds=RoutingArgs.ttl.value),
)
post_increment_values: Final = (
(None if shared_values.kind == "breaker_open" else local_values)
if isinstance(shared_values, RedisFailure)
else shared_values
)
if post_increment_values is not None:
for operation, value in zip(pipeline_operations, post_increment_values):
await self._set_global_usage_local(operation["key"], int(value), ttl=RoutingArgs.ttl.value)
return tpm_key
def _deployment_tracks_usage(self, deployment_id: str, model_group: str) -> bool:
@ -10098,15 +10115,26 @@ class Router(DeprecatedRouterMembers):
dt: Final = get_utc_datetime()
current_minute: Final = dt.strftime("%H-%M") # use the same timezone regardless of system clock
## RPM
rpm_key: Final = RouterCacheEnum.RPM.value.format(id=id, current_minute=current_minute, model=deployment_name)
counter: Final = GlobalUsageCounter(
str(_OPAQUE_VALUE_ADAPTER.validate_python(id)),
None if deployment_name is None else str(_OPAQUE_VALUE_ADAPTER.validate_python(deployment_name)),
current_minute,
"rpm",
)
with service_target(ROUTER_USAGE_TARGET):
await self._core_cache.async_increment_cache(
key=rpm_key,
key=usage_key(counter),
value=1,
local_only=True,
parent_otel_span=parent_otel_span,
ttl=RoutingArgs.ttl.value,
)
redis: Final = self._core_cache.redis_cache
if redis is not None:
store: Final[RouterGlobalUsageStore] = TransportRouterGlobalUsageStore(redis.transport())
await store.record_global_failure(
counter.deployment_id, counter.model, counter.minute, ttl=timedelta(seconds=RoutingArgs.ttl.value)
)
def _get_metadata_variable_name_from_kwargs(self, kwargs: dict) -> Literal["metadata", "litellm_metadata"]:
"""
@ -13204,60 +13232,59 @@ class Router(DeprecatedRouterMembers):
"""
dt: Final = get_utc_datetime()
current_minute: Final = dt.strftime("%H-%M") # use the same timezone regardless of system clock
tpm_keys: Final[list[str]] = []
rpm_keys: Final[list[str]] = []
model_list: Final = self._core_get_model_list(model_name=model_group)
if model_list is None: # no matching deployments
if model_list is None:
return None, None
for model in model_list:
id: str | None = model.get("model_info", {}).get("id")
litellm_model: str | None = model["litellm_params"].get(
"model"
) # USE THE MODEL SENT TO litellm.completion() - consistent with how global_router cache is written.
if id is None or litellm_model is None:
continue
tpm_keys.append(
RouterCacheEnum.TPM.value.format(
id=id,
model=litellm_model,
current_minute=current_minute,
)
identities: Final = tuple(
(
_OPAQUE_VALUE_ADAPTER.validate_python(model.get("model_info", {}).get("id")),
_OPAQUE_VALUE_ADAPTER.validate_python(model["litellm_params"].get("model")),
)
rpm_keys.append(
RouterCacheEnum.RPM.value.format(
id=id,
model=litellm_model,
current_minute=current_minute,
)
)
combined_tpm_rpm_keys: Final = tpm_keys + rpm_keys
combined_tpm_rpm_values: Final = await self._core_cache.async_batch_get_cache(keys=combined_tpm_rpm_keys)
if combined_tpm_rpm_values is None:
for model in model_list
)
tpm_counters: Final = tuple(
GlobalUsageCounter(str(deployment_id), model_name, current_minute, "tpm")
for deployment_id, model_name in identities
if deployment_id is not None and isinstance(model_name, str)
)
rpm_counters: Final = tuple(
GlobalUsageCounter(counter.deployment_id, counter.model, counter.minute, "rpm") for counter in tpm_counters
)
counters: Final = tpm_counters + rpm_counters
keys: Final = tuple(usage_key(counter) for counter in counters)
local: Final = _GLOBAL_USAGE_LOCAL_VALUES.validate_python(
await self._core_cache.async_batch_get_cache(keys=list(keys), local_only=True)
)
if local is None:
return None, None
remote: Final = await self._read_global_usage_remote(counters)
values: Final = tuple(remote.get(key, value) for key, value in zip(keys, local))
tpm_values: Final = tuple(value for value in values[: len(tpm_counters)] if isinstance(value, int))
rpm_values: Final = tuple(value for value in values[len(tpm_counters) :] if isinstance(value, int))
return (sum(tpm_values) if tpm_values else None, sum(rpm_values) if rpm_values else None)
tpm_usage_list: Final[list | None] = combined_tpm_rpm_values[: len(tpm_keys)]
rpm_usage_list: Final[list | None] = combined_tpm_rpm_values[len(tpm_keys) :]
async def _read_global_usage_remote(self, counters: Sequence[GlobalUsageCounter]) -> Mapping[str, object]:
keys: Final = tuple(usage_key(counter) for counter in counters)
due_keys, previous_access_times = self._core_cache.reserve_redis_batch_reads(keys)
due: Final = frozenset(due_keys)
redis: Final = self._core_cache.redis_cache
if redis is None or not due:
return {}
requested: Final = tuple(counter for counter in counters if usage_key(counter) in due)
store: Final[RouterGlobalUsageStore] = TransportRouterGlobalUsageStore(redis.transport())
shared: Final = await store.read_global_usage(requested)
if isinstance(shared, RedisFailure):
if shared.kind == "breaker_open":
self._core_cache.rollback_redis_batch_reads(previous_access_times)
return {}
remote: Final = MappingProxyType(dict(zip((usage_key(counter) for counter in requested), shared.values)))
for key, value in remote.items():
if value is not None:
await self._set_global_usage_local(key, value, ttl=self._core_cache.default_in_memory_ttl)
return remote
## TPM
tpm_usage: int | None = None
if tpm_usage_list is not None:
for t in tpm_usage_list:
if isinstance(t, int):
if tpm_usage is None:
tpm_usage = 0
tpm_usage += t
## RPM
rpm_usage: int | None = None
if rpm_usage_list is not None:
for t in rpm_usage_list:
if isinstance(t, int):
if rpm_usage is None:
rpm_usage = 0
rpm_usage += t
return tpm_usage, rpm_usage
async def _set_global_usage_local(self, key: str, value: object, *, ttl: float | None) -> None:
await self._core_cache.async_set_cache(key=key, value=value, local_only=True, ttl=ttl)
async def get_model_group_io_token_usage(self, model_group: str) -> tuple[int | None, int | None]:
"""

6
litellm/types/router.py
View file

@ -33,6 +33,8 @@ from litellm._uuid import uuid
from litellm.litellm_core_utils.completion_timeout import CompletionTimeout
from litellm.litellm_core_utils.core_helpers import normalize_drop_params
from litellm.litellm_core_utils.provider_affinity import validate_provider_affinity_header_name
from litellm.redis.stores.router_global_usage import RPM_KEY_TEMPLATE as _RPM_KEY_TEMPLATE
from litellm.redis.stores.router_global_usage import TPM_KEY_TEMPLATE as _TPM_KEY_TEMPLATE
from litellm.types.llms.base import LiteLLMBaseModel
from litellm.types.router_weights import RouterWeights
@ -1253,8 +1255,8 @@ class RoutingStrategy(enum.Enum):
class RouterCacheEnum(enum.Enum):
TPM = "global_router:{id}:{model}:tpm:{current_minute}"
RPM = "global_router:{id}:{model}:rpm:{current_minute}"
TPM = _TPM_KEY_TEMPLATE
RPM = _RPM_KEY_TEMPLATE
ITPM = "global_router:{id}:{model}:itpm:{current_minute}"
OTPM = "global_router:{id}:{model}:otpm:{current_minute}"

2
tests/code_coverage_tests/router_code_coverage.py
View file

@ -175,6 +175,8 @@ ignored_function_names = [
"_cannot_run_batches", # Tested through the Azure and OpenAI mixed-group acreate_batch tests in test_router.py
"_responses_continuation", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py
"_responses_continuation_ceilings", # Tested through the Responses mid-stream fallback tests in tests/unit/llms/base_llm/responses/test_transformation.py
"_read_global_usage_remote", # Tested through the shared Redis usage header tests in test_router_helper_utils.py
"_set_global_usage_local", # Tested through the shared Redis usage header tests in test_router_helper_utils.py
"_team_deployments_named", # Tested through the team lookups in the concurrent delete tests in test_router.py
"_served_openai_models", # Cached as served_openai_models, tested through the pre-call check token count tests in test_router.py
"priority_header_update", # Tested through the priority header tests in test_common_request_processing.py

215
tests/unit/redis/stores/test_router_global_usage.py Normal file
View file

@ -0,0 +1,215 @@
from collections.abc import Iterator, Mapping
from datetime import datetime, timedelta, timezone, tzinfo
from types import SimpleNamespace
from typing import Final, Protocol, cast
from unittest.mock import patch
import fakeredis
import pytest
from redis.asyncio.client import Pipeline
from redis.exceptions import ConnectionError as RedisConnectionError
from litellm import Router
from litellm.caching.redis_batch import request_redis_batch_scope
from litellm.caching.redis_cache import RedisCache
from litellm.redis.stores.router_global_usage import GlobalUsageCounter, TransportRouterGlobalUsageStore
from litellm.redis.transport.contract import RedisFailure
from tests._support.fake_redis import fake_redis_cache
TPM: Final = "global_router:deployment:openai/test-雪:tpm:12-34"
RPM: Final = "global_router:deployment:openai/test-雪:rpm:12-34"
class Clock(datetime):
@classmethod
def now(cls, tz: tzinfo | None = None) -> datetime:
return datetime(2026, 1, 1, 12, 34, 0, tzinfo=timezone.utc).astimezone(tz)
@pytest.fixture(autouse=True)
def clock() -> Iterator[None]:
with (
patch("litellm.utils.datetime", SimpleNamespace(datetime=Clock, timezone=timezone)),
patch("time.time", return_value=1767270840.0),
):
yield
class LegacyUsageCache(Protocol):
async def async_set_cache(self, key: str, value: object, *, ttl: int) -> bool | None: ...
async def async_get_cache(self, key: str) -> object: ...
class ExecutablePipeline(Protocol):
async def execute(self, raise_on_error: bool = True) -> object: ...
class UsageCallbacks(Protocol):
async def deployment_callback_on_success(
self, kwargs: Mapping[str, object], completion_response: object, start_time: object, end_time: object
) -> object: ...
async def async_deployment_callback_on_failure(
self, kwargs: Mapping[str, object], completion_response: object, start_time: object, end_time: object
) -> object: ...
def _router(cache: RedisCache | None) -> Router:
router: Final = Router(
model_list=[
{
"model_name": "group",
"litellm_params": {"model": "openai/test-雪", "api_key": "fake"},
"model_info": {"id": "deployment"},
"tpm": 1000,
"rpm": 100,
}
]
)
if cache is not None:
router.health().store.attach_redis_cache(cache)
return router
async def _record(router: Router, tokens: float = 3, *, failure: bool = False) -> None:
callbacks: Final = cast(UsageCallbacks, router) # cast-ok: typed view of the legacy callback signatures
metadata: Final = {"deployment": "openai/test-雪", "model_group": "group"}
kwargs: Final[Mapping[str, object]] = {
"call_type": "acompletion",
"standard_logging_object": {"total_tokens": tokens},
"litellm_params": {"metadata": metadata, "model_info": {"id": "deployment"}},
}
if failure:
await callbacks.async_deployment_callback_on_failure(kwargs, None, None, None)
else:
await callbacks.deployment_callback_on_success(kwargs, None, None, None)
@pytest.mark.asyncio
async def test_usage_old_new_bytes_namespace_and_post_result_reconciliation(server: fakeredis.FakeServer) -> None:
cache: Final = fake_redis_cache(server, namespace="tenant")
legacy: Final = cast(LegacyUsageCache, cache) # cast-ok: legacy writer/reader for rolling-deploy probes
raw: Final = fakeredis.FakeAsyncRedis(server=server)
await legacy.async_set_cache(TPM, 12, ttl=60)
await legacy.async_set_cache(RPM, 2, ttl=60)
router: Final = _router(cache)
assert await router.get_model_group_usage("group") == (12, 2)
with request_redis_batch_scope():
await _record(router, 3.5)
assert await raw.mget("tenant:" + TPM, "tenant:" + RPM) == [b"15.5", b"3"]
assert await raw.exists(TPM, RPM) == 0
assert await raw.pttl("tenant:" + TPM) == 60000
assert await raw.pttl("tenant:" + RPM) == 60000
assert await legacy.async_get_cache(TPM) == 15.5
assert await legacy.async_get_cache(RPM) == 3
assert await router.get_model_group_usage("group") == (15, 3)
assert await _router(cache).get_model_group_usage("group") == (None, 3)
@pytest.mark.asyncio
async def test_failure_preserves_window_but_success_refreshes_it(server: fakeredis.FakeServer) -> None:
cache: Final = fake_redis_cache(server)
raw: Final = fakeredis.FakeAsyncRedis(server=server)
router: Final = _router(cache)
await _record(router, failure=True)
assert await raw.get(RPM) == b"1"
assert await raw.pttl(RPM) == 60000
with patch("time.time", return_value=1767270850.0):
await _record(router, failure=True)
assert await raw.get(RPM) == b"2"
assert await raw.pttl(RPM) == 50000
await _record(router)
assert await raw.get(RPM) == b"3"
assert await raw.pttl(RPM) == 60000
@pytest.mark.asyncio
@pytest.mark.parametrize("seconds", [1.9, 0.9, 0, -0.9, -1.9])
@pytest.mark.parametrize("failure", [False, True])
async def test_usage_expiry_keeps_legacy_second_precision(
server: fakeredis.FakeServer, seconds: float, failure: bool
) -> None:
cache: Final = fake_redis_cache(server, namespace="tenant")
raw: Final = fakeredis.FakeAsyncRedis(server=server)
store: Final = TransportRouterGlobalUsageStore(cache.transport())
await raw.set("tenant:" + RPM, b"10")
if failure:
assert (
await store.record_global_failure("deployment", "openai/test-雪", "12-34", ttl=timedelta(seconds=seconds))
== 11
)
else:
assert await store.record_global_usage(
"deployment", "openai/test-雪", "12-34", tokens=3, requests=1, ttl=timedelta(seconds=seconds)
) == (3, 11)
with patch("time.time", return_value=1767270840.1):
assert await raw.get("tenant:" + RPM) == (b"11" if seconds >= 1 else None)
assert await raw.pttl("tenant:" + RPM) == (900 if seconds >= 1 else -2)
@pytest.mark.asyncio
async def test_missing_usage_is_throttled_and_late_attachment_does_not_backfill(server: fakeredis.FakeServer) -> None:
cache: Final = fake_redis_cache(server)
raw: Final = fakeredis.FakeAsyncRedis(server=server)
router: Final = _router(None)
await _record(router)
router.health().store.attach_redis_cache(cache)
assert await raw.exists(TPM, RPM) == 0
reader: Final = _router(cache)
assert await reader.get_model_group_usage("group") == (None, None)
await _record(router)
assert await raw.mget(TPM, RPM) == [b"3", b"1"]
assert await reader.get_model_group_usage("group") == (None, None)
with patch("time.time", return_value=1767270860.0):
assert await reader.get_model_group_usage("group") == (3, 1)
@pytest.mark.asyncio
async def test_malformed_batch_caches_raw_values_until_local_expiry(server: fakeredis.FakeServer) -> None:
cache: Final = fake_redis_cache(server)
raw: Final = fakeredis.FakeAsyncRedis(server=server)
await raw.set(TPM, b"bad-json")
await raw.set(RPM, b"5")
router: Final = _router(cache)
assert await router.get_model_group_usage("group") == (None, None)
await raw.set(TPM, b"10")
assert await router.get_model_group_usage("group") == (None, 5)
assert await _router(cache).get_model_group_usage("group") == (10, 5)
@pytest.mark.asyncio
async def test_outage_keeps_local_counts_and_shared_breaker() -> None:
server: Final = fakeredis.FakeServer()
cache: Final = fake_redis_cache(server)
router: Final = _router(cache)
server.connected = False
await _record(router)
assert await router.get_model_group_usage("group") == (3, 1)
store: Final = TransportRouterGlobalUsageStore(cache.transport())
for _ in range(5):
await store.read_global_usage((GlobalUsageCounter("deployment", "openai/test-雪", "12-34", "tpm"),))
result: Final = await store.record_global_failure(
"deployment", "openai/test-雪", "12-34", ttl=timedelta(seconds=60)
)
assert isinstance(result, RedisFailure) and result.kind == "breaker_open"
server.connected = True
@pytest.mark.asyncio
@pytest.mark.parametrize("failure", [False, True])
async def test_unknown_applied_usage_write_is_not_replayed(server: fakeredis.FakeServer, failure: bool) -> None:
cache: Final = fake_redis_cache(server)
raw: Final = fakeredis.FakeAsyncRedis(server=server)
router: Final = _router(cache)
original_execute: Final = cast(type[ExecutablePipeline], Pipeline).execute # cast-ok: typed redis-py execution seam
async def apply_then_disconnect(pipe: ExecutablePipeline, raise_on_error: bool = True) -> object:
await original_execute(pipe, raise_on_error=raise_on_error)
raise RedisConnectionError("response lost after pipeline applied")
with patch.object(Pipeline, "execute", new=apply_then_disconnect):
await _record(router, failure=failure)
assert await raw.mget(TPM, RPM) == [None if failure else b"3", b"1"]
assert await raw.pttl(RPM) == (-1 if failure else 60000)
assert await router.get_model_group_usage("group") == (None if failure else 3, 1)

2
tests/unit/test_internal_context.py
View file

@ -25,7 +25,7 @@ _REPO: Final = Path(__file__).resolve().parents[2]
_REDIS_PRODUCER_ROOTS: Final = ("litellm", "enterprise")
# The cache implementations and facades: they emit the service events, their callers declare the family.
_CACHE_LAYER_DIRS: Final = ("litellm/caching", "litellm/_v2/cache", "litellm/redis/transport")
_CACHE_LAYER_DIRS: Final = ("litellm/caching", "litellm/_v2/cache", "litellm/redis/transport", "litellm/redis/stores")
# Helpers that act on a cache handed in by the declaring caller, or forward to the response-cache facade.
_CACHE_PARAMETER_HELPERS: Final = frozenset(
{

103
tests/unit/test_router/test_router_helper_utils.py
View file

@ -10,6 +10,7 @@ from types import MethodType
from typing import Final
from unittest.mock import AsyncMock, MagicMock, patch
import fakeredis
import httpx
import pytest
import respx
@ -33,6 +34,7 @@ from litellm.types.utils import (
StandardLoggingModelInformation,
StandardLoggingPayload,
)
from tests._support.fake_redis import fake_redis_cache
def create_standard_logging_payload() -> StandardLoggingPayload:
@ -458,9 +460,7 @@ async def test_deployment_callback_on_success(sync_mode):
]
router = Router(model_list=model_list)
# Get the actual deployment ID that was generated
gpt_deployment = router.get_deployment_by_model_group_name(
model_group_name="gpt-5-mini"
)
gpt_deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-mini")
deployment_id = gpt_deployment["model_info"]["id"]
standard_logging_payload = create_standard_logging_payload()
@ -565,9 +565,16 @@ async def _finish_35_token_websocket_turn(routed: Mapping[str, object], call_id:
litellm_call_id=call_id,
function_id=call_id,
)
turn.update_environment_variables(litellm_params={"litellm_metadata": routed["litellm_metadata"]}, optional_params={})
turn.update_environment_variables(
litellm_params={"litellm_metadata": routed["litellm_metadata"]}, optional_params={}
)
await turn.dispatch_success_handlers(
[{"type": "response.completed", "response": {"usage": {"input_tokens": 20, "output_tokens": 15, "total_tokens": 35}}}],
[
{
"type": "response.completed",
"response": {"usage": {"input_tokens": 20, "output_tokens": 15, "total_tokens": 35}},
}
],
prefer_async_handlers=True,
)
@ -1327,9 +1334,7 @@ async def test_routing_strategy_pre_call_checks(model_list, sync_mode):
router = Router(model_list=model_list)
deployment = router.get_deployment_by_model_group_name(
model_group_name="gpt-5-mini"
)
deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-mini")
litellm_logging_obj = Logging(
model="gpt-5-mini",
@ -1344,9 +1349,7 @@ async def test_routing_strategy_pre_call_checks(model_list, sync_mode):
router.routing_strategy_pre_call_checks(deployment)
else:
## NO EXCEPTION
await router.async_routing_strategy_pre_call_checks(
deployment, litellm_logging_obj
)
await router.async_routing_strategy_pre_call_checks(deployment, litellm_logging_obj)
## WITH EXCEPTION - rate limit error
with patch.object(
@ -1361,27 +1364,19 @@ async def test_routing_strategy_pre_call_checks(model_list, sync_mode):
),
):
with pytest.raises(litellm.RateLimitError):
await router.async_routing_strategy_pre_call_checks(
deployment, litellm_logging_obj
)
await router.async_routing_strategy_pre_call_checks(deployment, litellm_logging_obj)
## WITH EXCEPTION - generic error
with patch.object(
callback, "async_pre_call_check", AsyncMock(side_effect=Exception("Error"))
):
with patch.object(callback, "async_pre_call_check", AsyncMock(side_effect=Exception("Error"))):
with pytest.raises(Exception, match="Error"):
await router.async_routing_strategy_pre_call_checks(
deployment, litellm_logging_obj
)
await router.async_routing_strategy_pre_call_checks(deployment, litellm_logging_obj)
@pytest.mark.parametrize(
"set_supported_environments, supported_environments, is_supported",
[(True, ["staging"], True), (False, None, True), (True, ["development"], False)],
)
def test_create_deployment(
model_list, set_supported_environments, supported_environments, is_supported
):
def test_create_deployment(model_list, set_supported_environments, supported_environments, is_supported):
"""Test if the '_create_deployment' function is working correctly"""
router = Router(model_list=model_list)
@ -1415,20 +1410,14 @@ def test_deployment_is_active_for_environment(
):
"""Test if the '_deployment_is_active_for_environment' function is working correctly"""
router = Router(model_list=model_list)
deployment = router.get_deployment_by_model_group_name(
model_group_name="gpt-5-mini"
)
deployment = router.get_deployment_by_model_group_name(model_group_name="gpt-5-mini")
if set_supported_environments:
os.environ["LITELLM_ENVIRONMENT"] = "staging"
deployment["model_info"]["supported_environments"] = supported_environments
if is_supported:
assert (
router.deployment_is_active_for_environment(deployment=deployment) is True
)
assert router.deployment_is_active_for_environment(deployment=deployment) is True
else:
assert (
router.deployment_is_active_for_environment(deployment=deployment) is False
)
assert router.deployment_is_active_for_environment(deployment=deployment) is False
def test_set_model_list(model_list):
@ -1755,31 +1744,14 @@ def test_deployment_tracks_usage_for_every_limit_source(
assert router._deployment_tracks_usage(deployment_id, "my-unmapped-model") is tracks
def _shared_redis_stub(store: dict) -> MagicMock:
from litellm.caching.redis_cache import RedisCache
async def increment_pipeline(increment_list, **kwargs):
for op in increment_list:
store[op["key"]] = store.get(op["key"], 0.0) + op["increment_value"]
return [store[op["key"]] for op in increment_list]
async def batch_get(keys, **kwargs):
return {key: store.get(key) for key in keys}
redis_stub = MagicMock(spec=RedisCache)
redis_stub.async_increment_pipeline = increment_pipeline
redis_stub.async_batch_get_cache = batch_get
return redis_stub
@pytest.mark.asyncio
@pytest.mark.usefixtures("router_minute_pinned")
async def test_headers_on_fresh_worker_reflect_shared_redis_usage():
store: dict = {}
server: Final = fakeredis.FakeServer()
worker_a = _rpm_tpm_router("lit-3058-workers")
worker_b = _rpm_tpm_router("lit-3058-workers")
worker_a.cache = DualCache(redis_cache=_shared_redis_stub(store), in_memory_cache=InMemoryCache())
worker_b.cache = DualCache(redis_cache=_shared_redis_stub(store), in_memory_cache=InMemoryCache())
worker_a.cache = DualCache(redis_cache=fake_redis_cache(server), in_memory_cache=InMemoryCache())
worker_b.cache = DualCache(redis_cache=fake_redis_cache(server), in_memory_cache=InMemoryCache())
messages = [{"role": "user", "content": "hi"}]
tokens_on_a = 0
@ -1840,27 +1812,19 @@ 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
),
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
),
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
),
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
),
RouterCacheEnum.OTPM.value.format(id="io-usage-dep-2", model="openai/gpt-4o", current_minute=minute),
20,
),
]
@ -2182,9 +2146,7 @@ def test_initialize_alerting(model_list):
from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting
from litellm.types.router import AlertingConfig
router = Router(
model_list=model_list, alerting_config=AlertingConfig(webhook_url="test")
)
router = Router(model_list=model_list, alerting_config=AlertingConfig(webhook_url="test"))
router._initialize_alerting()
callback_added = False
@ -2459,7 +2421,6 @@ async def test_discard_releases_per_request_override_selector():
assert selector_ref() is None
def test_discard_registers_the_selector_of_a_same_strategy_survivor_built_without_a_model_list():
retired: Final = Router(model_list=[], routing_strategy="latency-based-routing")
survivor: Final = Router(routing_strategy="latency-based-routing")
@ -2537,6 +2498,7 @@ async def test_discard_releases_its_router_budget_limiter():
assert budget_limiter_ref() is None
def test_initialize_assistants_endpoint(model_list):
"""Test if the 'initialize_assistants_endpoint' function is working correctly"""
router = Router(model_list=model_list)
@ -2612,7 +2574,6 @@ def test_deployments_by_pattern(model_list):
],
)
def test_pattern_match_deployment_set_model_name(user_request_model, model_name, litellm_model, expected_model):
from re import Match
from litellm.router_utils.pattern_match_deployments import PatternMatchRouter
@ -2667,7 +2628,6 @@ def test_add_optional_pre_call_checks(model_list):
@pytest.mark.asyncio
async def test_async_callback_filter_deployments(model_list):
from litellm.router_strategy.budget_limiter import RouterBudgetLimiting
router = Router(model_list=model_list)
@ -2699,7 +2659,7 @@ def test_cached_get_model_group_info(model_list):
def test_init_responses_api_endpoints(model_list):
"""Test if the '_init_responses_api_endpoints' function is working correctly"""
from typing import Callable
from collections.abc import Callable
router = Router(model_list=model_list)
@ -3554,6 +3514,7 @@ def test_resolve_model_name_from_model_id_keeps_a_wildcard_name_only_for_the_ids
router = Router(model_list=[{"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "sk-fake"}}])
assert router.resolve_model_name_from_model_id(model_id, custom_llm_provider) == expected
def test_get_valid_args():
"""Test get_valid_args static method returns valid Router.__init__ arguments"""