From b82443abc758cc6afb993bbd6904f1731d16d2e1 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 20:53:38 +0000 Subject: [PATCH] refactor(redis): move global router usage behind a typed store --- litellm/caching/dual_cache.py | 3 + litellm/redis/stores/router_global_usage.py | 111 +++++++++ litellm/router.py | 159 +++++++------ litellm/types/router.py | 6 +- .../router_code_coverage.py | 2 + .../redis/stores/test_router_global_usage.py | 215 ++++++++++++++++++ tests/unit/test_internal_context.py | 2 +- .../test_router/test_router_helper_utils.py | 103 +++------ 8 files changed, 461 insertions(+), 140 deletions(-) create mode 100644 litellm/redis/stores/router_global_usage.py create mode 100644 tests/unit/redis/stores/test_router_global_usage.py diff --git a/litellm/caching/dual_cache.py b/litellm/caching/dual_cache.py index e6214a3650c..d929953f3c1 100644 --- a/litellm/caching/dual_cache.py +++ b/litellm/caching/dual_cache.py @@ -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: diff --git a/litellm/redis/stores/router_global_usage.py b/litellm/redis/stores/router_global_usage.py new file mode 100644 index 00000000000..8471cba1a54 --- /dev/null +++ b/litellm/redis/stores/router_global_usage.py @@ -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 diff --git a/litellm/router.py b/litellm/router.py index 233e7581a40..1809f861ad8 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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]: """ diff --git a/litellm/types/router.py b/litellm/types/router.py index 726a328e913..d4aedfb45c6 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -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}" diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 7e05e98d917..791515f3b41 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -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 diff --git a/tests/unit/redis/stores/test_router_global_usage.py b/tests/unit/redis/stores/test_router_global_usage.py new file mode 100644 index 00000000000..cd3c1182731 --- /dev/null +++ b/tests/unit/redis/stores/test_router_global_usage.py @@ -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) diff --git a/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index 2155a75858d..6ead20c4cd5 100644 --- a/tests/unit/test_internal_context.py +++ b/tests/unit/test_internal_context.py @@ -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( { diff --git a/tests/unit/test_router/test_router_helper_utils.py b/tests/unit/test_router/test_router_helper_utils.py index 5f31ded626d..9084307d0c0 100644 --- a/tests/unit/test_router/test_router_helper_utils.py +++ b/tests/unit/test_router/test_router_helper_utils.py @@ -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"""