From a8d0d21f3e075a8db44484200d8c72f11c6f3918 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 13 Nov 2026 12:22:27 +0900 Subject: [PATCH] refactor(router): move Redis key construction out of router-core into the stores (#50356) router-core's RouterCache now names the routing state it touches (cooldowns, allowed fails, per-minute usage, health states) instead of spelling Redis keys. The stores own every key, and the Python host answers each typed cache step from its store without parsing keys back into ids. --- litellm/router_rust/core.py | 193 ++---- litellm/router_rust/routing_stores.py | 337 +++++------ litellm/router_rust/surface/health.py | 5 +- litellm/router_rust/surface/health_routing.py | 2 - litellm/router_utils/health_state_cache.py | 29 +- .../track_deployment_metrics.py | 8 +- .../crates/python-bridge/src/router/host.rs | 569 ++++++++++-------- .../crates/python-bridge/src/router/mod.rs | 168 +++--- .../redis-store/src/router_usage/keys.rs | 11 +- litellm/rust/crates/router-core/src/cache.rs | 336 +++++++---- litellm/rust/crates/router-core/src/health.rs | 210 ++----- .../crates/router-core/src/health_routing.rs | 84 +-- litellm/rust/crates/router-core/src/lib.rs | 16 +- .../crates/router-core/tests/call_loop.rs | 50 +- litellm/rust_bridge/_native.pyi | 60 +- tests/unit/router_rust/test_routing_stores.py | 105 +++- tests/unit/test_internal_context.py | 4 +- 17 files changed, 1177 insertions(+), 1010 deletions(-) diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index fbb07ea7553..e151bae80fb 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -32,9 +32,10 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Literal, Protocol, TypeAlias, TypeVar, final, overload 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.constants import DEFAULT_COOLDOWN_TIME_SECONDS, HEALTH_CHECKS_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 ( add_fallback_headers_to_response, @@ -75,8 +76,8 @@ from litellm.litellm_core_utils.total_deadline import ( enforce_total_timeout_scope, ) from litellm.redis.stores.connection import connect -from litellm.redis.stores.router_cooldowns import ROUTER_COOLDOWNS_TARGET, RouterCooldownStore -from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, RouterUsage +from litellm.redis.stores.router_cooldowns import RouterCooldownStore +from litellm.redis.stores.router_usage import RouterUsage from litellm.router_api import ( ROUTER_INIT_ARGS, SETTING_FIELDS, @@ -127,7 +128,16 @@ from litellm.router_rust.invoke.streaming import ( ) from litellm.router_rust.invoke.walk import FallbackOutcomes, WalkMetadata, await_all, log_retry_into from litellm.router_rust.invoke.websocket import turns_counted_their_requests -from litellm.router_rust.routing_stores import RoutingStores +from litellm.router_rust.routing_stores import ( + CACHE_STEP, + CooldownsStep, + CountFailureStep, + CountSuccessStep, + LocalStep, + OutcomesStep, + RoutingStores, + service_target_of, +) from litellm.router_rust.steps.access import caller_fields, selection_refusal from litellm.router_rust.steps.callbacks import OptionalChecks, RustOptionalChecks from litellm.router_rust.steps.plugins import SYNC_PATH_MESSAGE, candidate_models, run_plugins @@ -170,7 +180,6 @@ from litellm.router_rust.surface.health import ( logged_call, ) from litellm.router_rust.surface.health_routing import ( - HEALTH_STATE_KEY, HealthRouting, failover_document, fully_unhealthy_model_names, @@ -240,7 +249,13 @@ from litellm.router_utils.kubernetes_pod_discovery import ( ) from litellm.router_utils.pre_call_checks.encrypted_content_affinity_check import EncryptedContentAffinityCheck from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import release_io_token_reservation -from litellm.router_utils.router_callbacks.track_deployment_metrics import find_deployment_metadata +from litellm.router_utils.router_callbacks.track_deployment_metrics import ( + find_deployment_metadata, + get_deployment_failures_for_current_minute, + get_deployment_successes_for_current_minute, + increment_deployment_failures_for_current_minute, + increment_deployment_successes_for_current_minute, +) from litellm.router_utils.routing_read_batch import RoutingReadBatch from litellm.types.router import ( AllowedFailsPolicy, @@ -500,13 +515,6 @@ class RustCoreUnsupported(AttributeError): """A member or request the Rust core does not serve. It is never forwarded to a Python core.""" -def _cache_target(key: str) -> str: - """The service target a routing cache read or write is attributed to, as the Python core's.""" - if key == HEALTH_STATE_KEY: - return HEALTH_CHECKS_TARGET - return ROUTER_COOLDOWNS_TARGET if key.endswith(":cooldown") else ROUTER_USAGE_TARGET - - # --- failures ----------------------------------------------------------------------------------- @@ -3564,80 +3572,43 @@ class RustCore( # -- cache steps ----------------------------------------------------------------------------- - async def _batch_read(self, keys: list[str]) -> Sequence[object]: # mutable-ok: the cache's key list - """One batch read: cooldowns and tpm/rpm usage from their stores, every other key from the routing cache.""" - stored: Final = await self._stored_with_usage(keys) - if stored is not None: - return stored - read: Final[object] = await self._cache.async_batch_get_cache(keys=keys) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped - return read if isinstance(read, list) else [None] * len(keys) # pyright: ignore[reportUnknownVariableType] # see above - - async def _stored_with_usage(self, keys: Sequence[str]) -> tuple[object, ...] | None: + async def _cooldowns_with_usage(self, step: CooldownsStep) -> object: """A usage-based-routing-v2 pick reads its tpm/rpm counters alongside its cooldowns (``_usage_read_scope``).""" usage_read: Final = _USAGE_READ.get() - if usage_read is None or not keys or not keys[0].endswith(":cooldown"): - return await self._stores.batch_get(keys) + if usage_read is None: + return await self._stores.answer(step) batch, deployments = usage_read - stored, _ = await asyncio.gather(self._stores.batch_get(keys), batch.read_usage(deployments, None)) + stored, _ = await asyncio.gather(self._stores.answer(step), batch.read_usage(deployments, None)) return stored - def _batch_read_sync(self, keys: list[str]) -> Sequence[object]: # mutable-ok: the cache's key list - """``_batch_read`` for a sync pick.""" - stored: Final = self._stores.batch_get_sync(keys) - if stored is not None: - return stored - read: Final[object] = self._cache.batch_get_cache(keys=keys) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped - return read if isinstance(read, list) else [None] * len(keys) # pyright: ignore[reportUnknownVariableType] # see above + async def _answer_cache(self, raw: Mapping[str, object]) -> object: + step: Final = CACHE_STEP.validate_python(raw) + with service_target(service_target_of(step)): + match step: + case CooldownsStep(): + return _plain(await self._cooldowns_with_usage(step)) + case OutcomesStep() | CountSuccessStep() | CountFailureStep(): + return self._answer_outcomes(step) + case _: + return _plain(await self._stores.answer(step)) - async def _increment_usage(self, keys: Sequence[object], amounts: Sequence[object]) -> None: - try: - counted: Final = await self._stores.increment_usage(keys, amounts) - except Exception as e: # noqa: BLE001 # usage that cannot be written never fails the call, as on the Python router - from litellm._logging import verbose_router_logger - - verbose_router_logger.debug("_increment_deployment_usage: usage was not recorded. Got - %s", e) - return - if not counted: - raise _not_usage(keys) - - async def _answer_cache(self, step: Mapping[str, object]) -> object: - with service_target(_cache_target(str(step.get("key")))): - return await self._answer_cache_in_target(step) - - async def _answer_cache_in_target(self, step: Mapping[str, object]) -> object: - key: Final = str(step.get("key")) - ttl: Final = _ttl(step) - local: Final = step.get("local") is True - stored: Final = None if local else await self._stores.answer(step.get("op"), key, step.get("value"), ttl) - if stored is not None: - return _plain(stored.value) - match step.get("op"): - case "get": - return _plain(await self._cache.async_get_cache(key, local_only=local)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # DualCache's reads are untyped - case "batch_get": - return _plain(await self._batch_read([str(k) for k in _as_list(step.get("keys"))])) - case "set": - if ttl is None: - await self._cache.async_set_cache(key, step.get("value")) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped - else: - await self._cache.async_set_cache(key, step.get("value"), ttl=ttl) # pyright: ignore[reportUnknownMemberType] # see above + def _answer_outcomes(self, step: LocalStep) -> object: + """This worker's per-minute outcomes, kept where the Python router's cooldown handlers keep them.""" + ctx: Final = self._core_ctx + match step: + case OutcomesStep(id=model_id): + return [ + get_deployment_successes_for_current_minute(ctx, model_id), + get_deployment_failures_for_current_minute(ctx, model_id), + ] + case CountSuccessStep(id=model_id): + _ = increment_deployment_successes_for_current_minute(ctx, model_id) return None - case "increment": - by: Final = _amount(step.get("value")) - total: Final = ( - await self._cache.async_increment_cache(key, by, local_only=local) # pyright: ignore[reportUnknownMemberType] # see above - if ttl is None - else await self._cache.async_increment_cache(key, by, local_only=local, ttl=ttl) # pyright: ignore[reportUnknownMemberType] # see above - ) - return by if total is None else total - case "increment_usage": - await self._increment_usage(_as_list(step.get("keys")), _as_list(step.get("value"))) + case CountFailureStep(id=model_id): + increment_deployment_failures_for_current_minute(ctx, model_id) return None - case "delete": - await self._cache.async_delete_cache(key) - return None - case other: - raise RuntimeError(f"router-core asked for an unknown cache op {other!r}") + case _: + assert_never(step) def run_routing_cache_sync(self, read_or_write: "Callable[[], Coroutine[object, object, object]]") -> object: """A routing-cache coroutine on a sync call, whose `make_call` runs on the one event loop @@ -3658,54 +3629,21 @@ class RustCore( @property def _stores(self) -> RoutingStores: """The stores router-core's cache steps reach, over the cooldowns as last assigned.""" - return RoutingStores(self._cooldowns, self._usage) + return RoutingStores(self._cooldowns, self._usage, self._health_checks) def io_token_usage_sync(self, model_group: str) -> tuple[int | None, int | None]: """``get_model_group_io_token_usage``, whose counters stay in the routing cache, on a sync call.""" read: Final = self.run_routing_cache_sync(lambda: self._hosted.get_model_group_io_token_usage(model_group)) return _IO_TOKEN_USAGE.validate_python(read) - def _answer_cache_sync(self, step: Mapping[str, object]) -> object: - with service_target(_cache_target(str(step.get("key")))): - return self._answer_cache_sync_in_target(step) - - def _answer_cache_sync_in_target(self, step: Mapping[str, object]) -> object: - key: Final = str(step.get("key")) - ttl: Final = _ttl(step) - local: Final = step.get("local") is True - op: Final = step.get("op") - stored: Final = None if local else self._stores.answer_sync(op, key, step.get("value"), ttl) - if stored is not None: - return _plain(stored.value) - match op: - case "get": - return _plain(self._cache.get_cache(key, local_only=local)) # pyright: ignore[reportUnknownMemberType, reportUnknownArgumentType] # DualCache's reads are untyped - case "batch_get": - return _plain(self._batch_read_sync([str(k) for k in _as_list(step.get("keys"))])) - case "set": - if ttl is None: - self._cache.set_cache(key, step.get("value")) # pyright: ignore[reportUnknownMemberType] # DualCache's writes are untyped - else: - self._cache.set_cache(key, step.get("value"), ttl=ttl) # pyright: ignore[reportUnknownMemberType] # see above - return None - case "increment": - by: Final = _amount(step.get("value")) - # DualCache.increment_cache is annotated int but adds floats as given. - return ( - self._cache.increment_cache(key, by, local_only=local) # pyright: ignore[reportUnknownMemberType, reportArgumentType] # see above - if ttl is None - else self._cache.increment_cache(key, by, local_only=local, ttl=ttl) # pyright: ignore[reportUnknownMemberType, reportArgumentType] # see above - ) - case "increment_usage": - keys: Final = _as_list(step.get("keys")) - if not self._stores.increment_usage_sync(keys, _as_list(step.get("value"))): - raise _not_usage(keys) - return None - case "delete": - self._cache.delete_cache(key) # pyright: ignore[reportUnknownMemberType] # DualCache's delete is untyped - return None - case other: - raise RuntimeError(f"router-core asked for an unknown cache op {other!r}") + def _answer_cache_sync(self, raw: Mapping[str, object]) -> object: + step: Final = CACHE_STEP.validate_python(raw) + with service_target(service_target_of(step)): + match step: + case OutcomesStep() | CountSuccessStep() | CountFailureStep(): + return self._answer_outcomes(step) + case _: + return _plain(self._stores.answer_sync(step)) # -- errors the core raises when no attempt failed ------------------------------------------ @@ -4074,19 +4012,6 @@ def _seconds(value: object, default: float) -> float | int: _IO_TOKEN_USAGE: Final = TypeAdapter(tuple[int | None, int | None]) -def _not_usage(keys: Sequence[object]) -> RuntimeError: - return RuntimeError(f"router-core asked to add usage to keys that are not one deployment's tpm/rpm: {keys!r}") - - -def _ttl(step: Mapping[str, object]) -> float | None: - ttl: Final = step.get("ttl") - return ttl if isinstance(ttl, int | float) else None - - -def _amount(value: object) -> float | int: - return value if isinstance(value, int | float) and not isinstance(value, bool) else 1 - - # The plans whose kwargs carry `original_generic_function`, which the fallback cooldown reads. _GENERIC_API_KINDS: Final = frozenset(("generic", "search")) diff --git a/litellm/router_rust/routing_stores.py b/litellm/router_rust/routing_stores.py index af3ef7cd375..706866a9d05 100644 --- a/litellm/router_rust/routing_stores.py +++ b/litellm/router_rust/routing_stores.py @@ -2,24 +2,39 @@ from __future__ import annotations from collections.abc import Sequence from dataclasses import dataclass -from typing import Final, TypeAlias +from typing import Annotated, Final, Literal, Protocol, TypeAlias -from pydantic import TypeAdapter +from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter +from typing_extensions import assert_never -from litellm._internal_context import service_target, with_service_target -from litellm.caching.redis_batch import Batch +from litellm._logging import verbose_router_logger +from litellm.constants import HEALTH_CHECKS_TARGET from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.redis.stores.router_cooldowns import ROUTER_COOLDOWNS_TARGET from litellm.redis.stores.router_usage import ROUTER_USAGE_TARGET, RouterUsage, UsageCounter +from litellm.redis.stores.router_usage.types import CounterKind from litellm.router_utils import routing_prefetch from litellm.router_utils.cooldown_cache import CooldownCache, CooldownCacheValue +from litellm.router_utils.health_state_cache import health_states_of +from litellm.router_utils.pre_call_checks.model_rate_limit_check import RoutingArgs -__all__ = ("AllowedFailsKey", "CooldownKey", "Handled", "RoutingStores", "UsageKey", "parse_key") +__all__ = ( + "CACHE_STEP", + "CacheStep", + "CooldownsStep", + "CountFailureStep", + "CountSuccessStep", + "LocalStep", + "OutcomesStep", + "RoutingStores", + "StoreStep", + "service_target_of", +) _COOLDOWN_VALUE: Final = TypeAdapter(CooldownCacheValue) -def _cooldown_value(value: object) -> CooldownCacheValue: +def _cooldown_value(value: JsonValue) -> CooldownCacheValue: """The cooldown router-core wrote, its whole-second ``cooldown_time`` kept an ``int`` as it serialized it.""" validated: Final = _COOLDOWN_VALUE.validate_python(value) written: Final = as_str_mapping(value) @@ -29,82 +44,104 @@ def _cooldown_value(value: object) -> CooldownCacheValue: return validated -@dataclass(frozen=True, slots=True) -class CooldownKey: - model_id: str +class _Step(BaseModel): + model_config = ConfigDict(frozen=True) -@dataclass(frozen=True, slots=True) -class AllowedFailsKey: - model_id: str +class UsageWindow(_Step): + deployment_id: str + model: str + minute: str + + def counter(self, kind: CounterKind) -> UsageCounter: + return UsageCounter(self.deployment_id, self.model, self.minute, kind) + + +class CooldownsStep(_Step): + op: Literal["cooldowns"] + ids: tuple[str, ...] + + +class PutCooldownStep(_Step): + op: Literal["put_cooldown"] + id: str + cooldown: JsonValue + ttl: float + + +class IncrementAllowedFailsStep(_Step): + op: Literal["increment_allowed_fails"] + id: str suffix: str | None + ttl: float -@dataclass(frozen=True, slots=True) -class UsageKey: - counter: UsageCounter +class AddUsageStep(_Step): + op: Literal["add_usage"] + window: UsageWindow + tokens: int | float + requests: int | float -RoutingKey: TypeAlias = CooldownKey | AllowedFailsKey | UsageKey +class AddRequestStep(_Step): + op: Literal["add_request"] + window: UsageWindow -@dataclass(frozen=True, slots=True) -class Handled: - value: object +class UsageStep(_Step): + op: Literal["usage"] + window: UsageWindow -def _deployment_key(rest: str) -> RoutingKey | None: - if rest.endswith(":cooldown"): - return CooldownKey(rest.removesuffix(":cooldown")) - model_id, marker, suffix = rest.partition(":allowed_fails") - if not marker or (suffix and not suffix.startswith(":")): - return None - return AllowedFailsKey(model_id, suffix.removeprefix(":") or None) +class HealthStatesStep(_Step): + op: Literal["health_states"] -def _usage_key(rest: str) -> UsageKey | None: - parts: Final = rest.rsplit(":", 2) - if len(parts) != 3: - return None - head, kind, minute = parts - model_id, separator, model = head.partition(":") - if not separator: - return None - match kind: - case "tpm" | "rpm": - return UsageKey(UsageCounter(model_id, model, minute, kind)) +class OutcomesStep(_Step): + op: Literal["outcomes"] + id: str + + +class CountSuccessStep(_Step): + op: Literal["count_success"] + id: str + + +class CountFailureStep(_Step): + op: Literal["count_failure"] + id: str + + +StoreStep: TypeAlias = ( + CooldownsStep + | PutCooldownStep + | IncrementAllowedFailsStep + | AddUsageStep + | AddRequestStep + | UsageStep + | HealthStatesStep +) +LocalStep: TypeAlias = OutcomesStep | CountSuccessStep | CountFailureStep +CacheStep: TypeAlias = Annotated[StoreStep | LocalStep, Field(discriminator="op")] +CACHE_STEP: Final[TypeAdapter[CacheStep]] = TypeAdapter(CacheStep) + + +def service_target_of(step: StoreStep | LocalStep) -> str: + """The service target a cache step is attributed to, as the Python core's.""" + match step: + case CooldownsStep() | PutCooldownStep() | IncrementAllowedFailsStep(): + return ROUTER_COOLDOWNS_TARGET + case HealthStatesStep(): + return HEALTH_CHECKS_TARGET + case AddUsageStep() | AddRequestStep() | UsageStep() | OutcomesStep() | CountSuccessStep() | CountFailureStep(): + return ROUTER_USAGE_TARGET case _: - return None + assert_never(step) -def parse_key(key: str) -> RoutingKey | None: - """The store family a router-core key belongs to. Each store rebuilds the same key from what is parsed - here, so an id or model holding a ``:`` still names the key router-core asked for.""" - if key.startswith("deployment:"): - return _deployment_key(key.removeprefix("deployment:")) - if key.startswith("global_router:"): - return _usage_key(key.removeprefix("global_router:")) - return None - - -def _service_target(key: RoutingKey | None) -> str | None: - if isinstance(key, CooldownKey | AllowedFailsKey): - return ROUTER_COOLDOWNS_TARGET - if isinstance(key, UsageKey): - return ROUTER_USAGE_TARGET - return None - - -def _cooldown_ids(keys: Sequence[str]) -> tuple[str, ...] | None: - ids: Final = tuple(parsed.model_id for parsed in map(parse_key, keys) if isinstance(parsed, CooldownKey)) - return ids if keys and len(ids) == len(keys) else None - - -def _usage_counters(keys: Sequence[object]) -> tuple[UsageCounter, ...] | None: - counters: Final = tuple( - parsed.counter for parsed in (parse_key(str(key)) for key in keys) if isinstance(parsed, UsageKey) - ) - return counters if keys and len(counters) == len(keys) else None +class HealthStates(Protocol): + async def async_get_states(self) -> object: ... + def get_states(self) -> object: ... def _aligned(model_ids: Sequence[str], active: Sequence[tuple[str, CooldownCacheValue]]) -> tuple[object, ...]: @@ -112,119 +149,67 @@ def _aligned(model_ids: Sequence[str], active: Sequence[tuple[str, CooldownCache return tuple(found.get(model_id) for model_id in model_ids) -def _one_deployment(tpm: UsageCounter, rpm: UsageCounter) -> bool: - return (tpm.kind, rpm.kind) == ("tpm", "rpm") and (tpm.deployment_id, tpm.model, tpm.minute) == ( - rpm.deployment_id, - rpm.model, - rpm.minute, - ) - - -def _deployment_increment(keys: Sequence[object], amounts: Sequence[object]) -> tuple[UsageCounter, float, int] | None: - """``_increment_deployment_usage``'s tpm then rpm increment of one deployment.""" - counters: Final = _usage_counters(keys) - if counters is None or len(counters) != 2 or len(amounts) != 2: - return None - tpm, rpm = counters - tokens, requests = amounts - if not _one_deployment(tpm, rpm) or not isinstance(tokens, int | float) or not isinstance(requests, int | float): - return None - return tpm, tokens, int(requests) - - -def _by(value: object) -> float: - return value if isinstance(value, int | float) and not isinstance(value, bool) else 1 +def _usage_counters(window: UsageWindow) -> tuple[UsageCounter, UsageCounter]: + return window.counter("tpm"), window.counter("rpm") +@dataclass(frozen=True, slots=True) class RoutingStores: - """Router-core's cache steps for cooldowns, allowed fails and tpm/rpm usage, answered by the native stores - the Python core uses. ``None`` leaves the step to the routing cache.""" + """Router-core's cache steps answered by the stores the Python core uses. The stores own every key.""" - def __init__(self, cooldowns: CooldownCache, usage: RouterUsage) -> None: - self._cooldowns: Final = cooldowns - self._usage: Final = usage + cooldowns: CooldownCache + usage: RouterUsage + health: HealthStates - def answer_sync(self, op: object, key: str, value: object, ttl: float | None) -> Handled | None: - store: Final = self._cooldowns.cooldown_store - parsed: Final = parse_key(key) - target: Final = _service_target(parsed) - if target is None: - return None - with service_target(target): - match op, parsed: - case "get", CooldownKey(model_id): - return Handled(_aligned((model_id,), store.active_sync((model_id,)))[0]) - case "set", CooldownKey(model_id): - store.put(model_id, _cooldown_value(value), ttl or 0.0) - return Handled(None) - case "delete", CooldownKey(model_id): - store.evict_local(model_id) - return Handled(None) - case "increment", AllowedFailsKey(model_id, suffix): - return Handled(store.increment_allowed_fails(model_id, suffix, ttl or 0.0)) - case "increment", UsageKey(counter): - return Handled(self._usage.increment_sync(counter, _by(value), int(ttl or 0))) - case _: - return None + async def answer(self, step: StoreStep) -> object: + store: Final = self.cooldowns.cooldown_store + match step: + case CooldownsStep(ids=ids): + return _aligned(ids, await routing_prefetch.read_cooldowns(store, ids, None)) + case PutCooldownStep(id=model_id, cooldown=cooldown, ttl=ttl): + await store.async_put(model_id, _cooldown_value(cooldown), ttl) + return None + case IncrementAllowedFailsStep(id=model_id, suffix=suffix, ttl=ttl): + return await store.async_increment_allowed_fails(model_id, suffix, ttl) + case AddUsageStep(window=window, tokens=tokens, requests=requests): + try: + await self.usage.increment_deployment( + window.deployment_id, window.model, window.minute, tokens, int(requests) + ) + except Exception as e: # noqa: BLE001 # usage that cannot be written never fails the call, as on the Python router + verbose_router_logger.debug("_increment_deployment_usage: usage was not recorded. Got - %s", e) + return None + case AddRequestStep(window=window): + await self.usage.increment(window.counter("rpm"), 1, RoutingArgs.ttl) + return None + case UsageStep(window=window): + return tuple(await self.usage.get_many(_usage_counters(window))) + case HealthStatesStep(): + return health_states_of(await self.health.async_get_states()) + case _: + assert_never(step) - async def answer(self, op: object, key: str, value: object, ttl: float | None) -> Handled | None: - parsed: Final = parse_key(key) - target: Final = _service_target(parsed) - if target is None: - return None - store: Final = self._cooldowns.cooldown_store - with service_target(target): - match op, parsed: - case "get", CooldownKey(model_id): - return Handled(_aligned((model_id,), await store.active((model_id,), batch=Batch.NONE))[0]) - case "set", CooldownKey(model_id): - await store.async_put(model_id, _cooldown_value(value), ttl or 0.0) - return Handled(None) - case "increment", AllowedFailsKey(model_id, suffix): - return Handled(await store.async_increment_allowed_fails(model_id, suffix, ttl or 0.0)) - case "increment", UsageKey(counter): - by: Final = _by(value) - return Handled(await self._usage.increment(counter, by, int(ttl or 0))) - case _: - return self.answer_sync(op, key, value, ttl) - - def batch_get_sync(self, keys: Sequence[str]) -> tuple[object, ...] | None: - model_ids: Final = _cooldown_ids(keys) - if model_ids is not None: - with service_target(ROUTER_COOLDOWNS_TARGET): - return _aligned(model_ids, self._cooldowns.cooldown_store.active_sync(model_ids)) - counters: Final = _usage_counters(keys) - if counters is None: - return None - with service_target(ROUTER_USAGE_TARGET): - return tuple(self._usage.get_many_sync(counters)) - - async def batch_get(self, keys: Sequence[str]) -> tuple[object, ...] | None: - model_ids: Final = _cooldown_ids(keys) - if model_ids is not None: - store: Final = self._cooldowns.cooldown_store - with service_target(ROUTER_COOLDOWNS_TARGET): - return _aligned(model_ids, await routing_prefetch.read_cooldowns(store, model_ids, None)) - counters: Final = _usage_counters(keys) - if counters is None: - return None - with service_target(ROUTER_USAGE_TARGET): - return tuple(await self._usage.get_many(counters)) - - @with_service_target(ROUTER_USAGE_TARGET) - def increment_usage_sync(self, keys: Sequence[object], amounts: Sequence[object]) -> bool: - found: Final = _deployment_increment(keys, amounts) - if found is None: - return False - tpm, tokens, requests = found - self._usage.increment_deployment_sync(tpm.deployment_id, tpm.model, tpm.minute, tokens, requests) - return True - - @with_service_target(ROUTER_USAGE_TARGET) - async def increment_usage(self, keys: Sequence[object], amounts: Sequence[object]) -> bool: - found: Final = _deployment_increment(keys, amounts) - if found is None: - return False - tpm, tokens, requests = found - await self._usage.increment_deployment(tpm.deployment_id, tpm.model, tpm.minute, tokens, requests) - return True + def answer_sync(self, step: StoreStep) -> object: + store: Final = self.cooldowns.cooldown_store + match step: + case CooldownsStep(ids=ids): + return _aligned(ids, store.active_sync(ids)) + case PutCooldownStep(id=model_id, cooldown=cooldown, ttl=ttl): + store.put(model_id, _cooldown_value(cooldown), ttl) + return None + case IncrementAllowedFailsStep(id=model_id, suffix=suffix, ttl=ttl): + return store.increment_allowed_fails(model_id, suffix, ttl) + case AddUsageStep(window=window, tokens=tokens, requests=requests): + self.usage.increment_deployment_sync( + window.deployment_id, window.model, window.minute, tokens, int(requests) + ) + return None + case AddRequestStep(window=window): + self.usage.increment_sync(window.counter("rpm"), 1, RoutingArgs.ttl) + return None + case UsageStep(window=window): + return tuple(self.usage.get_many_sync(_usage_counters(window))) + case HealthStatesStep(): + return health_states_of(self.health.get_states()) + case _: + assert_never(step) diff --git a/litellm/router_rust/surface/health.py b/litellm/router_rust/surface/health.py index f2ac72a262a..211ead61aeb 100644 --- a/litellm/router_rust/surface/health.py +++ b/litellm/router_rust/surface/health.py @@ -505,14 +505,15 @@ class RustHealthCallbacks: ) -> str | None: """``sync_deployment_callback_on_success``, for a sync stream that ended.""" from litellm.litellm_core_utils.batch_utils import is_batch_retrieve_call_type + from litellm.router_utils.router_callbacks.track_deployment_metrics import deployment_successes_key if is_batch_retrieve_call_type(kwargs.get("call_type")): return None target: Final = _success_target(kwargs) if target is None: return None - key: Final = self._drive_sync_steps(self._handle.count_success(target[0], callbacks)).get("key") - return key if isinstance(key, str) else None + counted: Final = self._drive_sync_steps(self._handle.count_success(target[0], callbacks)).get("counted") + return deployment_successes_key(target[0]) if counted is True else None async def repin_after_success(self, kwargs: Mapping[str, object], completion_response: object) -> None: """The session re-pin ``deployment_callback_on_success`` runs for a success logged against a diff --git a/litellm/router_rust/surface/health_routing.py b/litellm/router_rust/surface/health_routing.py index 46542bf9af2..44a045a0d5d 100644 --- a/litellm/router_rust/surface/health_routing.py +++ b/litellm/router_rust/surface/health_routing.py @@ -21,7 +21,6 @@ if TYPE_CHECKING: from litellm.router_config import RouterConfig from litellm.types.router import Deployment, DeploymentTypedDict -HEALTH_STATE_KEY: Final = DeploymentHealthCache.CACHE_KEY # The failover flags the Python core keeps as plain router attributes, with their defaults. FAILOVER_FLAGS: Final = ("enable_weighted_failover", "fallback_on_unsupported_params") @@ -149,7 +148,6 @@ async def fully_unhealthy_model_names( __all__ = [ "FAILOVER_FLAGS", - "HEALTH_STATE_KEY", "HealthRouting", "all_blocked", "failover_document", diff --git a/litellm/router_utils/health_state_cache.py b/litellm/router_utils/health_state_cache.py index c4cc9350cfa..47ac0abeaff 100644 --- a/litellm/router_utils/health_state_cache.py +++ b/litellm/router_utils/health_state_cache.py @@ -10,13 +10,14 @@ from collections.abc import Mapping from typing import TYPE_CHECKING, Any, Final from pydantic import TypeAdapter -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm import verbose_logger from litellm._internal_context import with_service_target from litellm.caching.caching import DualCache from litellm.caching.redis_cache import RedisCircuitBreakerOpenError from litellm.constants import HEALTH_CHECKS_TARGET +from litellm.litellm_core_utils.core_helpers import as_str_mapping if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -109,6 +110,16 @@ class DeploymentHealthCache: and (now - _timestamp_of(state)) < self.staleness_threshold # pyright: ignore[reportUnknownArgumentType] # a cached record } + @with_service_target(HEALTH_CHECKS_TARGET) + async def async_get_states(self) -> object: + """Every deployment's raw state, as the router core's health filter reads it; a failed read is None.""" + return await self.cache.async_get_cache(key=self.CACHE_KEY) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped + + @with_service_target(HEALTH_CHECKS_TARGET) + def get_states(self) -> object: + """``async_get_states`` for a sync call.""" + return self.cache.get_cache(key=self.CACHE_KEY) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped + @with_service_target(HEALTH_CHECKS_TARGET) async def async_get_unhealthy_deployment_ids(self, parent_otel_span: Span | None = None) -> set[str]: """Return set of deployment IDs currently marked unhealthy and not stale.""" @@ -141,3 +152,19 @@ def _timestamp_of(state: Mapping[str, object]) -> float: drops out alone instead of turning health routing off for every deployment.""" timestamp: Final = state.get("timestamp", 0) return float(timestamp) if isinstance(timestamp, int | float) and not isinstance(timestamp, bool) else 0.0 + + +class DeploymentHealthCheck(TypedDict): + is_healthy: ReadOnly[bool] + timestamp: ReadOnly[float] + + +def health_states_of(raw: object) -> Mapping[str, DeploymentHealthCheck]: + """The cached states as ``_extract_unhealthy_ids`` reads them: any falsy ``is_healthy`` is unhealthy, a + record without one is healthy, and what is no record is dropped.""" + records: Final = ((model_id, as_str_mapping(state)) for model_id, state in (as_str_mapping(raw) or {}).items()) + return { + model_id: {"is_healthy": bool(state.get("is_healthy", True)), "timestamp": _timestamp_of(state)} + for model_id, state in records + if state is not None + } diff --git a/litellm/router_utils/router_callbacks/track_deployment_metrics.py b/litellm/router_utils/router_callbacks/track_deployment_metrics.py index 2cf4e595296..0a4017c4dde 100644 --- a/litellm/router_utils/router_callbacks/track_deployment_metrics.py +++ b/litellm/router_utils/router_callbacks/track_deployment_metrics.py @@ -51,6 +51,10 @@ async def count_before_callbacks( raise +def deployment_successes_key(deployment_id: str) -> str: + return f"{deployment_id}:successes" + + @renamed_router_param("litellm_router_instance") def increment_deployment_successes_for_current_minute( ctx: "RoutingContext", @@ -59,7 +63,7 @@ def increment_deployment_successes_for_current_minute( """ In-Memory: Increments the number of successes for the current minute for a deployment_id """ - key: Final = f"{deployment_id}:successes" + key: Final = deployment_successes_key(deployment_id) ctx.health().store.increment_cache( local_only=True, key=key, @@ -96,7 +100,7 @@ def get_deployment_successes_for_current_minute( Returns 0 if no value found """ - key: Final = f"{deployment_id}:successes" + key: Final = deployment_successes_key(deployment_id) return ( ctx.health().store.get_cache( local_only=True, diff --git a/litellm/rust/crates/python-bridge/src/router/host.rs b/litellm/rust/crates/python-bridge/src/router/host.rs index 04d19e70edc..08169091e1b 100644 --- a/litellm/rust/crates/python-bridge/src/router/host.rs +++ b/litellm/rust/crates/python-bridge/src/router/host.rs @@ -6,11 +6,13 @@ //! current thread, so a host future finds its exchange without the router passing it through. //! //! The routing cache is a host operation too: [`StepCache`] parks every [`RouterCache`] op as a -//! `cache` step, and Python runs it against the router's own `DualCache` (memory plus Redis), so -//! Python and Rust read and write the same cooldowns, fail counters and usage. +//! `cache` step, and Python answers it from the router's stores, so Python and Rust read and write +//! the same cooldowns, fail counters and usage. A step names the state it touches, never a Redis +//! key: the keys are the stores'. use std::{ cell::RefCell, + collections::HashMap, future::{Future, poll_fn}, pin::Pin, sync::{Arc, Mutex, MutexGuard, PoisonError}, @@ -21,9 +23,10 @@ use std::{ use async_trait::async_trait; use litellm_router_core::{ Attempt, AttemptError, Blocked, BuildAsk, BuildPath, Built, CallbackFiltered, Deployment, - Error, Host, Output, Payload, PreRoute, PreRouted, Provided, Randomness, RouterCache, - RouterEvent, RouterRequest, RouterResponse, RoutingContext, SecretRef, SelectStage, - ThreadRandomness, build::BuiltDeployment, host::ModelInfoLookup, model_info::MergedModelInfo, + Error, HealthState, Host, Outcomes, Output, Payload, PreRoute, PreRouted, Provided, Randomness, + RouterCache, RouterEvent, RouterRequest, RouterResponse, RoutingContext, SecretRef, + SelectStage, ThreadRandomness, UsageCounts, UsageWindow, build::BuiltDeployment, + health::CooldownValue, host::ModelInfoLookup, model_info::MergedModelInfo, precall::PreCallFacts, }; use pyo3::{PyResult, Python, types::PyAnyMethods}; @@ -102,55 +105,123 @@ pub(crate) enum Stage { ProviderRequest, } -/// The `op` of a `cache` step: one per [`RouterCache`] method. +/// One [`RouterCache`] op, as a `cache` step carries it to Python. +#[derive(Clone, Debug, PartialEq, serde::Serialize)] +#[serde(tag = "op", rename_all = "snake_case")] +pub(crate) enum CacheCall { + /// Answered with one stored cooldown or `None` per id. + Cooldowns { ids: Vec }, + /// Answered with `None`. + PutCooldown { + id: String, + cooldown: CooldownValue, + ttl: f64, + }, + /// Answered with the count after this failure. + IncrementAllowedFails { + id: String, + suffix: Option, + ttl: f64, + }, + /// Answered with `None`. + AddUsage { + window: UsageWindow, + #[serde(serialize_with = "serialize_amount")] + tokens: f64, + #[serde(serialize_with = "serialize_amount")] + requests: f64, + }, + /// Answered with `None`. + AddRequest { window: UsageWindow }, + /// Answered with `[tokens, requests]`, each a number or `None`. + Usage { window: UsageWindow }, + /// Answered with each checked deployment's `{is_healthy, timestamp}` by id, or `None`. + HealthStates, + /// Answered with this worker's `[successes, fails]` of the current minute. + Outcomes { id: String }, + /// Answered with `None`. + CountSuccess { id: String }, + /// Answered with `None`. + CountFailure { id: String }, +} + +/// The `op` of a `cache` step. #[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display, strum::IntoStaticStr)] pub(crate) enum CacheOp { - /// `async_get_cache(key)`: answered with the stored JSON value, or `None`. - #[strum(serialize = "get")] - Get, - /// `async_set_cache(key, value, ttl)`: answered with `None`. - #[strum(serialize = "set")] - Set, - /// `async_increment(key, value, ttl)`: answered with the new number. The TTL applies when - /// the key is created. - #[strum(serialize = "increment")] - Increment, - /// `async_increment_usage(keys, value, ttl)`: `value` is one amount per key, added as - /// `_increment_deployment_usage` adds them. Answered with `None`. - #[strum(serialize = "increment_usage")] - IncrementUsage, - /// `async_batch_get_cache(keys)`: answered with one JSON value or `None` per key, in order. - #[strum(serialize = "batch_get")] - BatchGet, - /// `async_delete_cache(key)`: answered with `None`. - #[strum(serialize = "delete")] - Delete, + #[strum(serialize = "cooldowns")] + Cooldowns, + #[strum(serialize = "put_cooldown")] + PutCooldown, + #[strum(serialize = "increment_allowed_fails")] + IncrementAllowedFails, + #[strum(serialize = "add_usage")] + AddUsage, + #[strum(serialize = "add_request")] + AddRequest, + #[strum(serialize = "usage")] + Usage, + #[strum(serialize = "health_states")] + HealthStates, + #[strum(serialize = "outcomes")] + Outcomes, + #[strum(serialize = "count_success")] + CountSuccess, + #[strum(serialize = "count_failure")] + CountFailure, +} + +impl CacheCall { + pub(crate) fn op(&self) -> CacheOp { + match self { + Self::Cooldowns { .. } => CacheOp::Cooldowns, + Self::PutCooldown { .. } => CacheOp::PutCooldown, + Self::IncrementAllowedFails { .. } => CacheOp::IncrementAllowedFails, + Self::AddUsage { .. } => CacheOp::AddUsage, + Self::AddRequest { .. } => CacheOp::AddRequest, + Self::Usage { .. } => CacheOp::Usage, + Self::HealthStates => CacheOp::HealthStates, + Self::Outcomes { .. } => CacheOp::Outcomes, + Self::CountSuccess { .. } => CacheOp::CountSuccess, + Self::CountFailure { .. } => CacheOp::CountFailure, + } + } + + /// How many entries a `Values` answer holds. + fn values_len(&self) -> Option { + match self { + Self::Cooldowns { ids } => Some(ids.len()), + Self::Usage { .. } | Self::Outcomes { .. } => Some(2), + _ => None, + } + } } /// Python's outcome of a `cache` step. #[derive(Clone, Debug, PartialEq)] pub(crate) enum CacheOutcome { - /// `get`: the stored value (`None` when missing). Value(Option), - /// `batch_get`: one entry per key. Values(Vec>), - /// `increment`: the counter after the increment. Number(f64), - /// `set`, `delete` and `increment_usage`. Stored, } impl CacheOutcome { - /// The op this outcome answers (`Stored` answers both `set` and `delete`). fn answers(&self, op: CacheOp) -> bool { matches!( (self, op), - (Self::Value(_), CacheOp::Get) - | (Self::Values(_), CacheOp::BatchGet) - | (Self::Number(_), CacheOp::Increment) + (Self::Value(_), CacheOp::HealthStates) + | ( + Self::Values(_), + CacheOp::Cooldowns | CacheOp::Usage | CacheOp::Outcomes + ) + | (Self::Number(_), CacheOp::IncrementAllowedFails) | ( Self::Stored, - CacheOp::Set | CacheOp::Delete | CacheOp::IncrementUsage + CacheOp::PutCooldown + | CacheOp::AddUsage + | CacheOp::AddRequest + | CacheOp::CountSuccess + | CacheOp::CountFailure ) ) } @@ -219,18 +290,7 @@ pub(crate) enum HostCall { fallback_depth: u32, retry: u32, }, - /// One [`RouterCache`] op against the host's cache. `key` is set for every op but - /// `batch_get`, which sets `keys`; `value` is set for `set` (the JSON to store) and - /// `increment` (the amount); `ttl` is in seconds. - Cache { - op: CacheOp, - key: Option, - keys: Option>, - value: Option, - ttl: Option, - /// The op reads or writes this worker's in-memory tier only (`local_only=True`). - local: bool, - }, + Cache(CacheCall), /// The deployment build's `os.environ/` references: answered with one value per ref. Secrets { index: usize, @@ -263,7 +323,7 @@ impl HostCall { Self::RefuseFallbacks { .. } => StepKind::RefuseFallbacks, Self::Select { .. } => StepKind::Select, Self::PreRoute { .. } => StepKind::PreRoute, - Self::Cache { .. } => StepKind::Cache, + Self::Cache(_) => StepKind::Cache, Self::Secrets { .. } => StepKind::Secrets, Self::Build { kind, .. } => *kind, Self::ModelInfo { .. } => StepKind::ModelInfo, @@ -380,15 +440,17 @@ impl HostCallResponse { got: values.len(), }); } - if let (Self::Cache(outcome), HostCall::Cache { op, keys, .. }) = (self, call) { - if !outcome.answers(*op) { - return Err(DriveError::WrongCacheOutcome { op: *op }); + if let (Self::Cache(outcome), HostCall::Cache(cache)) = (self, call) { + let op = cache.op(); + if !outcome.answers(op) { + return Err(DriveError::WrongCacheOutcome { op }); } - if let (CacheOutcome::Values(values), Some(keys)) = (outcome, keys) - && values.len() != keys.len() + if let (CacheOutcome::Values(values), Some(expected)) = (outcome, cache.values_len()) + && values.len() != expected { return Err(DriveError::BatchLength { - expected: keys.len(), + op, + expected, got: values.len(), }); } @@ -418,8 +480,10 @@ pub(crate) enum Finish { }, /// `report_failure`: the cooldown applied, if the deployment was cooled down. Reported { cooldown: Option }, - /// `report_success` / `count_success`: the success counter's key, if one was counted. + /// `report_success`: the tokens counter's key, if usage was added. Counted { key: Option }, + /// `count_success`: whether a success was counted. + SuccessCounted { counted: bool }, /// `build`: the deployments it installed, in order. Built(Vec), /// `router_model_info`: the merged model info (`None` when the lookup answered `None`). @@ -454,8 +518,12 @@ pub(crate) enum DriveError { UnknownCandidate { id: String }, #[error("the outcome does not answer a `{op}` cache step")] WrongCacheOutcome { op: CacheOp }, - #[error("a `batch_get` cache step for {expected} keys was answered with {got} values")] - BatchLength { expected: usize, got: usize }, + #[error("a `{op}` cache step for {expected} entries was answered with {got} values")] + BatchLength { + op: CacheOp, + expected: usize, + got: usize, + }, #[error("a `secrets` step for {expected} references was answered with {got} values")] SecretsLength { expected: usize, got: usize }, #[error("the router call is pending on something other than a step")] @@ -541,12 +609,8 @@ async fn call_host(call: HostCall) -> HostCallResponse { const CHECKED: &str = "the driver checks each response against its host call"; -fn ttl_seconds(ttl: Option) -> Option { - ttl.map(|ttl| ttl.as_secs_f64()) -} - /// An increment amount as Python's counters store it: an int when it is whole. -fn amount_value(amount: f64) -> Value { +pub(crate) fn amount_value(amount: f64) -> Value { if amount.fract() == 0.0 && amount.abs() < i64::MAX as f64 { Value::from(amount as i64) } else { @@ -554,131 +618,151 @@ fn amount_value(amount: f64) -> Value { } } -/// The host's cache: every op is a `cache` step Python runs against the router's `DualCache`. -pub(crate) struct StepCache; - -async fn cache_step( - op: CacheOp, - key: Option<&str>, - keys: Option<&[String]>, - value: Option, - ttl: Option, -) -> CacheOutcome { - call_cache(op, key, keys, value, ttl, false).await +fn serialize_amount(amount: &f64, serializer: S) -> Result { + serde::Serialize::serialize(&amount_value(*amount), serializer) } -async fn call_cache( - op: CacheOp, - key: Option<&str>, - keys: Option<&[String]>, - value: Option, - ttl: Option, - local: bool, -) -> CacheOutcome { - let call = HostCall::Cache { - op, - key: key.map(str::to_owned), - keys: keys.map(<[String]>::to_vec), - value, - ttl: ttl_seconds(ttl), - local, - }; - let HostCallResponse::Cache(outcome) = call_host(call).await else { +/// The host's cache: every op is a `cache` step Python answers from the router's stores. +pub(crate) struct StepCache; + +async fn cache_step(call: CacheCall) -> CacheOutcome { + let HostCallResponse::Cache(outcome) = call_host(HostCall::Cache(call)).await else { unreachable!("{CHECKED}"); }; outcome } +async fn stored(call: CacheCall) -> Result<(), litellm_cache::Error> { + let CacheOutcome::Stored = cache_step(call).await else { + unreachable!("{CHECKED}"); + }; + Ok(()) +} + +async fn counted(call: CacheCall) -> Result { + let CacheOutcome::Number(next) = cache_step(call).await else { + unreachable!("{CHECKED}"); + }; + Ok(next) +} + +async fn value(call: CacheCall) -> Result, litellm_cache::Error> { + let CacheOutcome::Value(value) = cache_step(call).await else { + unreachable!("{CHECKED}"); + }; + Ok(value) +} + +async fn values(call: CacheCall) -> Result>, litellm_cache::Error> { + let CacheOutcome::Values(values) = cache_step(call).await else { + unreachable!("{CHECKED}"); + }; + Ok(values) +} + #[async_trait] impl RouterCache for StepCache { - async fn async_get_cache(&self, key: &str) -> Result, litellm_cache::Error> { - let CacheOutcome::Value(value) = - cache_step(CacheOp::Get, Some(key), None, None, None).await - else { - unreachable!("{CHECKED}"); - }; - Ok(value) + async fn cooldowns( + &self, + ids: &[String], + ) -> Result>, litellm_cache::Error> { + let stored = values(CacheCall::Cooldowns { ids: ids.to_vec() }).await?; + // A cooldown that does not read as one is no cooldown. + Ok(stored + .into_iter() + .map(|cooldown| serde_json::from_value(cooldown?).ok()) + .collect()) } - async fn async_set_cache( + async fn put_cooldown( &self, - key: &str, - value: Value, - ttl: Option, + id: &str, + cooldown: &CooldownValue, + ttl: Duration, ) -> Result<(), litellm_cache::Error> { - cache_step(CacheOp::Set, Some(key), None, Some(value), ttl).await; - Ok(()) + stored(CacheCall::PutCooldown { + id: id.to_owned(), + cooldown: cooldown.clone(), + ttl: ttl.as_secs_f64(), + }) + .await } - async fn async_increment( + async fn increment_allowed_fails( &self, - key: &str, - amount: f64, - ttl: Option, + id: &str, + suffix: Option<&str>, + ttl: Duration, ) -> Result { - let value = Some(amount_value(amount)); - let CacheOutcome::Number(next) = - cache_step(CacheOp::Increment, Some(key), None, value, ttl).await - else { - unreachable!("{CHECKED}"); - }; - Ok(next) + counted(CacheCall::IncrementAllowedFails { + id: id.to_owned(), + suffix: suffix.map(str::to_owned), + ttl: ttl.as_secs_f64(), + }) + .await } - async fn async_batch_get_cache( + async fn add_usage( &self, - keys: &[String], - ) -> Result>, litellm_cache::Error> { - let CacheOutcome::Values(values) = - cache_step(CacheOp::BatchGet, None, Some(keys), None, None).await - else { - unreachable!("{CHECKED}"); - }; - Ok(values) - } - - async fn async_delete_cache(&self, key: &str) -> Result<(), litellm_cache::Error> { - cache_step(CacheOp::Delete, Some(key), None, None, None).await; - Ok(()) - } - - async fn async_increment_usage( - &self, - increments: &[(String, f64)], - ttl: Option, + window: &UsageWindow, + tokens: f64, + requests: f64, ) -> Result<(), litellm_cache::Error> { - let keys: Vec = increments.iter().map(|(key, _)| key.clone()).collect(); - let amounts = increments - .iter() - .map(|(_, amount)| amount_value(*amount)) - .collect(); - let value = Some(Value::Array(amounts)); - cache_step(CacheOp::IncrementUsage, None, Some(&keys), value, ttl).await; - Ok(()) + stored(CacheCall::AddUsage { + window: window.clone(), + tokens, + requests, + }) + .await } - async fn async_get_local(&self, key: &str) -> Result, litellm_cache::Error> { - let CacheOutcome::Value(value) = - call_cache(CacheOp::Get, Some(key), None, None, None, true).await - else { - unreachable!("{CHECKED}"); - }; - Ok(value) + async fn add_request(&self, window: &UsageWindow) -> Result<(), litellm_cache::Error> { + stored(CacheCall::AddRequest { + window: window.clone(), + }) + .await } - async fn async_increment_local( - &self, - key: &str, - amount: f64, - ttl: Option, - ) -> Result { - let value = Some(amount_value(amount)); - let CacheOutcome::Number(next) = - call_cache(CacheOp::Increment, Some(key), None, value, ttl, true).await - else { - unreachable!("{CHECKED}"); + async fn usage(&self, window: &UsageWindow) -> Result { + let read = values(CacheCall::Usage { + window: window.clone(), + }) + .await?; + let count = |index: usize| read.get(index).cloned().flatten().and_then(|v| v.as_f64()); + Ok(UsageCounts { + tokens: count(0), + requests: count(1), + }) + } + + async fn health_states(&self) -> Result, litellm_cache::Error> { + let states = value(CacheCall::HealthStates).await?; + Ok(states + .and_then(|states| serde_json::from_value(states).ok()) + .unwrap_or_default()) + } + + async fn outcomes(&self, id: &str) -> Result { + let read = values(CacheCall::Outcomes { id: id.to_owned() }).await?; + let count = |index: usize| { + read.get(index) + .cloned() + .flatten() + .and_then(|count| count.as_f64()) + .map_or(0, |count| count as u64) }; - Ok(next) + Ok(Outcomes { + successes: count(0), + fails: count(1), + }) + } + + async fn count_success(&self, id: &str) -> Result<(), litellm_cache::Error> { + stored(CacheCall::CountSuccess { id: id.to_owned() }).await + } + + async fn count_failure(&self, id: &str) -> Result<(), litellm_cache::Error> { + stored(CacheCall::CountFailure { id: id.to_owned() }).await } } @@ -1151,13 +1235,12 @@ impl Driver { #[cfg(test)] mod tests { - use std::{ - collections::HashMap, - time::{SystemTime, UNIX_EPOCH}, - }; + use std::time::{SystemTime, UNIX_EPOCH}; + use futures_util::FutureExt; use litellm_router_core::{ - FailureKind, Op, Router, RouterConfig, RouterDynamicConfig, SeededRandomness, + FailureKind, InMemoryRouterCache, Op, Router, RouterConfig, RouterDynamicConfig, + SeededRandomness, }; use rstest::{fixture, rstest}; use serde_json::{Value, json}; @@ -1195,61 +1278,81 @@ mod tests { HostCallResponse::Attempt(Err(AttemptError::new(FailureKind::RateLimit, Some(429)))) } - /// A dict-backed stand-in for Python's `DualCache` that answers `cache` steps and keeps - /// every one it answered. + /// Router-core's in-memory cache standing in for Python's stores: it answers `cache` steps + /// and keeps every one it answered. #[derive(Default)] struct Store { - entries: HashMap, + cache: InMemoryRouterCache, calls: Vec, } impl Store { fn respond(&mut self, call: &HostCall) -> HostCallResponse { - let HostCall::Cache { - op, - key, - keys, - value, - .. - } = call - else { + let HostCall::Cache(cache_call) = call else { panic!("not a cache step: {call:?}"); }; self.calls.push(call.clone()); - let key = key.clone().unwrap_or_default(); - HostCallResponse::Cache(match op { - CacheOp::Get => CacheOutcome::Value(self.entries.get(&key).cloned()), - CacheOp::Set => { - self.entries.insert(key, value.clone().unwrap()); - CacheOutcome::Stored - } - CacheOp::Increment => { - let current = self.entries.get(&key).and_then(Value::as_f64); - let next = - current.unwrap_or(0.0) + value.as_ref().and_then(Value::as_f64).unwrap(); - self.entries.insert(key, amount_value(next)); - CacheOutcome::Number(next) - } - CacheOp::IncrementUsage => { - let amounts = value.as_ref().and_then(Value::as_array).unwrap(); - for (key, amount) in keys.iter().flatten().zip(amounts) { - let current = self.entries.get(key).and_then(Value::as_f64); - let next = current.unwrap_or(0.0) + amount.as_f64().unwrap(); - self.entries.insert(key.clone(), amount_value(next)); + let cache = &self.cache; + let seconds = Duration::from_secs_f64; + let outcome = async { + Ok::<_, litellm_cache::Error>(match cache_call { + CacheCall::Cooldowns { ids } => CacheOutcome::Values( + cache + .cooldowns(ids) + .await? + .iter() + .map(|cooldown| cooldown.as_ref().map(|cooldown| json!(cooldown))) + .collect(), + ), + CacheCall::PutCooldown { id, cooldown, ttl } => { + cache.put_cooldown(id, cooldown, seconds(*ttl)).await?; + CacheOutcome::Stored } - CacheOutcome::Stored - } - CacheOp::BatchGet => CacheOutcome::Values( - keys.iter() - .flatten() - .map(|key| self.entries.get(key).cloned()) - .collect(), - ), - CacheOp::Delete => { - self.entries.remove(&key); - CacheOutcome::Stored - } - }) + CacheCall::IncrementAllowedFails { id, suffix, ttl } => CacheOutcome::Number( + cache + .increment_allowed_fails(id, suffix.as_deref(), seconds(*ttl)) + .await?, + ), + CacheCall::AddUsage { + window, + tokens, + requests, + } => { + cache.add_usage(window, *tokens, *requests).await?; + CacheOutcome::Stored + } + CacheCall::AddRequest { window } => { + cache.add_request(window).await?; + CacheOutcome::Stored + } + CacheCall::Usage { window } => { + let counts = cache.usage(window).await?; + CacheOutcome::Values( + [counts.tokens, counts.requests] + .map(|count| count.map(Value::from)) + .to_vec(), + ) + } + CacheCall::HealthStates => CacheOutcome::Value(None), + CacheCall::Outcomes { id } => { + let outcomes = cache.outcomes(id).await?; + CacheOutcome::Values( + [outcomes.successes, outcomes.fails] + .map(|count| Some(Value::from(count))) + .to_vec(), + ) + } + CacheCall::CountSuccess { id } => { + cache.count_success(id).await?; + CacheOutcome::Stored + } + CacheCall::CountFailure { id } => { + cache.count_failure(id).await?; + CacheOutcome::Stored + } + }) + }; + HostCallResponse::Cache(outcome.now_or_never().unwrap().unwrap()) } /// Resume with `response`, answering every cache step on the way, up to the next other @@ -1305,18 +1408,10 @@ mod tests { fn cooldown_reads_are_cache_steps(pair: Router) { let mut driver = start(&pair, "primary"); let step = driver.resume(None).unwrap(); - let Step::HostCall(HostCall::Cache { - op: CacheOp::BatchGet, - key: None, - keys: Some(keys), - value: None, - ttl: None, - local: false, - }) = step - else { + let Step::HostCall(HostCall::Cache(CacheCall::Cooldowns { ids })) = step else { panic!("the call should read cooldowns first, got {step:?}"); }; - assert_eq!(keys, ["deployment:a1:cooldown", "deployment:a2:cooldown"]); + assert_eq!(ids, ["a1", "a2"]); } #[rstest] @@ -1365,10 +1460,21 @@ mod tests { .duration_since(UNIX_EPOCH) .unwrap() .as_secs_f64(); - store.entries.insert( - "deployment:a1:cooldown".into(), - json!({"exception_received": "slow down", "status_code": "429", "timestamp": now, "cooldown_time": 60}), - ); + store + .cache + .put_cooldown( + "a1", + &CooldownValue { + exception_received: "slow down".to_owned(), + status_code: "429".to_owned(), + timestamp: now, + cooldown_time: 60.0, + }, + Duration::from_secs(60), + ) + .now_or_never() + .unwrap() + .unwrap(); for _ in 0..8 { let mut call = start(&pair_without_retries, "primary"); let Step::HostCall(HostCall::Attempt { deployment_id, .. }) = @@ -1409,7 +1515,7 @@ mod tests { .calls .iter() .filter(|call| { - matches!(call, HostCall::Cache { op, .. } if !matches!(op, CacheOp::Get | CacheOp::BatchGet)) + matches!(call, HostCall::Cache(cache) if !matches!(cache, CacheCall::Cooldowns { .. } | CacheCall::Usage { .. } | CacheCall::HealthStates | CacheCall::Outcomes { .. })) }) .collect(); assert!(!writes.is_empty(), "the core wrote no outcome"); @@ -1504,13 +1610,7 @@ mod tests { let mut store = Store::default(); let mut driver = start(&pair, "primary"); let step = driver.resume(None).unwrap(); - let Step::HostCall( - call @ HostCall::Cache { - op: CacheOp::BatchGet, - .. - }, - ) = &step - else { + let Step::HostCall(call @ HostCall::Cache(CacheCall::Cooldowns { .. })) = &step else { panic!("expected the cooldown read, got {step:?}"); }; let error = driver.resume(Some(response)).unwrap_err(); @@ -1638,10 +1738,7 @@ mod tests { assert_eq!(model_group, "primary"); assert!(matches!( store.calls.as_slice(), - [HostCall::Cache { - op: CacheOp::BatchGet, - .. - }] + [HostCall::Cache(CacheCall::Cooldowns { .. })] )); } diff --git a/litellm/rust/crates/python-bridge/src/router/mod.rs b/litellm/rust/crates/python-bridge/src/router/mod.rs index 6b58d63a714..679194e012a 100644 --- a/litellm/rust/crates/python-bridge/src/router/mod.rs +++ b/litellm/rust/crates/python-bridge/src/router/mod.rs @@ -10,7 +10,7 @@ mod stream; use std::{ sync::{Arc, Mutex, PoisonError}, - time::{Duration, SystemTime}, + time::Duration, }; use litellm_error::Report; @@ -19,9 +19,8 @@ use litellm_router_core::lookup::{ListScope, Lookup, Row}; use litellm_router_core::{ AttemptError, Blocked, BuildPath, Built, Deployment, DeploymentSpec, Error, FailureKind, FailureReport, Op, OutcomeCallbacks, PreRouted, Provided, RetryCause, Router, RouterEvent, - SecretAt, SuccessReport, + SecretAt, SuccessReport, UsageWindow, build::Install, - health::{UsageMetric, current_minute, usage_key}, registry::{DeploymentFilter, ResolveKind, ResolveScope}, strategy::routing_context, }; @@ -32,6 +31,8 @@ use pyo3::{ }; use serde_json::Value; +use litellm_redis_store::router_usage::{Counter, Kind}; + use self::host::{ CacheOp, CacheOutcome, DriveError, Driver, Finish, HostCall, HostCallResponse, PythonRandomness, Stage, Step, StepHost, StepKind, @@ -384,13 +385,19 @@ impl RouterCoreHandle { .collect() } - /// The routing-cache key deployment `id` serving litellm `model` counts `metric` (`tpm`, - /// `rpm`, `itpm` or `otpm`) under in `minute` (`HH-MM`). + /// The key deployment `id` serving litellm `model` counts `metric` (`tpm`, `rpm`, `itpm` or + /// `otpm`) under in `minute` (`HH-MM`), as the usage store spells it. fn usage_key(&self, id: &str, model: &str, metric: &str, minute: &str) -> PyResult { - let metric: UsageMetric = metric + let kind: Kind = metric .parse() .map_err(|_| PyValueError::new_err(format!("`{metric}` is not a usage metric")))?; - Ok(usage_key(id, model, metric, minute)) + 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. @@ -457,7 +464,8 @@ impl RouterCoreHandle { } /// Record a success a logging callback saw, driven like `report_failure`; it ends with - /// `done` (`key`: the tokens counter's key when usage was added, else None). + /// `done` (`key`: the tokens counter's key when usage was added, else None, as + /// `deployment_callback_on_success` returns it). /// `usage_model` is the model its usage is keyed by, and `tracks_usage` the host's answer to /// whether it counts any (see `Output::tracks_usage`). `callbacks`: see `report_failure`. #[pyo3(signature = (id, total_tokens, requests, usage_model=None, tracks_usage=None, callbacks=None))] @@ -481,12 +489,14 @@ impl RouterCoreHandle { router .report_success(&id, total_tokens, requests, report) .await - .map(|key| Finish::Counted { key }) + .map(|window| Finish::Counted { + key: window.map(tokens_key), + }) })) } - /// Count a success against `id`, configured or not; ends with `done` (`key`). `callbacks`: - /// see `report_failure`. + /// Count a success against `id`, configured or not; ends with `done` (`counted`: whether it + /// counted one). `callbacks`: see `report_failure`. #[pyo3(signature = (id, callbacks=None))] fn count_success( &self, @@ -500,30 +510,10 @@ impl RouterCoreHandle { router .count_success(&id, callbacks) .await - .map(|key| Finish::Counted { key }) + .map(|counted| Finish::SuccessCounted { counted }) })) } - /// The usage keys of this minute for every deployment of `model_group`: - /// `[(id, tpm_key, rpm_key)]`, or None when the group has no deployment. - fn usage_keys(&self, model_group: &str) -> Option> { - let minute = current_minute(SystemTime::now()); - let group = self.router.registry().group(model_group)?; - Some( - group - .iter() - .map(|deployment| { - let model = &deployment.params.model; - ( - deployment.id.clone(), - usage_key(&deployment.id, model, UsageMetric::Tpm, &minute), - usage_key(&deployment.id, model, UsageMetric::Rpm, &minute), - ) - }) - .collect(), - ) - } - /// Every configured deployment, as `model_list` entries with their ids, in config order. fn deployments<'py>(&self, py: Python<'py>) -> PyResult> { let registry = self.router.registry(); @@ -1241,7 +1231,7 @@ fn host_call_response( result => HostCallResponse::Attempt(result), }) } - HostCall::Cache { op, .. } => cache_outcome(*op, outcome).map(HostCallResponse::Cache), + HostCall::Cache(call) => cache_outcome(call.op(), outcome).map(HostCallResponse::Cache), HostCall::Secrets { .. } => Ok(HostCallResponse::Secrets(from_py_argument( &outcome.ok_or_else(missing)?, )?)), @@ -1311,36 +1301,43 @@ fn pre_route_outcome(outcome: Option>) -> PyResult>) -> PyResult { let outcome = outcome.filter(|outcome| !outcome.is_none()); let json = |value: Option| value.filter(|value| !value.is_null()); match op { - CacheOp::Get => outcome + CacheOp::HealthStates => outcome .map(|outcome| from_py_argument::(&outcome)) .transpose() .map(|value| CacheOutcome::Value(json(value))), - CacheOp::BatchGet => { + CacheOp::Cooldowns | CacheOp::Usage | CacheOp::Outcomes => { let outcome = outcome.ok_or_else(|| { - PyValueError::new_err("a `batch_get` cache outcome must be a list, not None") + PyValueError::new_err(format!("a `{op}` cache outcome must be a list, not None")) })?; if !outcome.is_instance_of::() { - return Err(PyValueError::new_err( - "a `batch_get` cache outcome must be a list with one entry per key", - )); + return Err(PyValueError::new_err(format!( + "a `{op}` cache outcome must be a list with one entry per value asked for" + ))); } let values: Vec> = from_py_argument(&outcome)?; Ok(CacheOutcome::Values(values.into_iter().map(json).collect())) } - CacheOp::Increment => outcome + CacheOp::IncrementAllowedFails => outcome .filter(|outcome| !outcome.is_instance_of::()) .and_then(|outcome| outcome.extract::().ok()) .map(CacheOutcome::Number) .ok_or_else(|| { - PyValueError::new_err("an `increment` cache outcome must be the new int or float") + PyValueError::new_err(format!( + "an `{op}` cache outcome must be the new int or float" + )) }), - CacheOp::Set | CacheOp::Delete | CacheOp::IncrementUsage => match outcome { + CacheOp::PutCooldown + | CacheOp::AddUsage + | CacheOp::AddRequest + | CacheOp::CountSuccess + | CacheOp::CountFailure => match outcome { None => Ok(CacheOutcome::Stored), Some(_) => Err(PyValueError::new_err(format!( "a `{op}` cache step has no outcome, so pass None" @@ -1349,6 +1346,16 @@ fn cache_outcome(op: CacheOp, outcome: Option>) -> PyResult String { + Counter { + deployment_id: window.deployment_id, + model: Some(window.model), + minute: window.minute, + kind: Kind::Tpm, + } + .key() +} + /// `{"ok": True, "total_tokens": int}` (an attempt's may add `"blocked": str`, `"deferred": bool` /// and `"tracks_usage": bool`) or `{"ok": False, "kind": str, "status": int | None, /// "retry_after": float | None, "num_retries": int | None, "detail": str | None, @@ -1808,20 +1815,8 @@ fn host_call_items(py: Python<'_>, dict: &Bound<'_, PyDict>, call: HostCall) -> dict.set_item("fallback_depth", fallback_depth)?; dict.set_item("retry", retry) } - HostCall::Cache { - op, - key, - keys, - value, - ttl, - local, - } => { - dict.set_item("op", <&'static str>::from(op))?; - dict.set_item("key", key)?; - dict.set_item("keys", keys.map(|keys| PyList::new(py, keys)).transpose()?)?; - dict.set_item("value", value.map(|value| to_py(py, &value)).transpose()?)?; - dict.set_item("ttl", ttl)?; - dict.set_item("local", local) + HostCall::Cache(call) => { + dict.update(to_py(py, &call)?.bind(py).cast::()?.as_mapping()) } HostCall::Secrets { index, refs } => { dict.set_item("index", index)?; @@ -1914,6 +1909,7 @@ fn done_items(dict: &Bound<'_, PyDict>, finish: Finish) -> PyResult<()> { dict.set_item("built", PyList::new(py, records)?) } Finish::Counted { key } => dict.set_item("key", key), + Finish::SuccessCounted { counted } => dict.set_item("counted", counted), Finish::RouterModelInfo(merged) => { let py = dict.py(); match merged { @@ -1977,21 +1973,42 @@ import time store = {} asked = [] +def add(key, amount): + store[key] = store.get(key, 0) + amount + return store[key] + +def window(step): + w = step["window"] + return (w["deployment_id"], w["model"], w["minute"]) + def answer_cache(step): asked.append(step) - op, key = step["op"], step["key"] - if op == "get": - return store.get(key) - if op == "batch_get": - return [store.get(k) for k in step["keys"]] - if op == "set": - store[key] = step["value"] + op = step["op"] + if op == "cooldowns": + return [store.get(("cooldown", id)) for id in step["ids"]] + if op == "put_cooldown": + store[("cooldown", step["id"])] = step["cooldown"] return None - if op == "increment": - store[key] = store.get(key, 0) + step["value"] - return store[key] - if op == "delete": - store.pop(key, None) + if op == "increment_allowed_fails": + return add(("allowed_fails", step["id"], step["suffix"]), 1) + if op == "add_usage": + add(("tokens", *window(step)), step["tokens"]) + add(("requests", *window(step)), step["requests"]) + return None + if op == "add_request": + add(("requests", *window(step)), 1) + return None + if op == "usage": + return [store.get(("tokens", *window(step))), store.get(("requests", *window(step)))] + if op == "health_states": + return None + if op == "outcomes": + return [store.get(("successes", step["id"])), store.get(("fails", step["id"]))] + if op == "count_success": + add(("successes", step["id"]), 1) + return None + if op == "count_failure": + add(("fails", step["id"]), 1) return None raise AssertionError(op) @@ -2004,8 +2021,9 @@ def drive(call, outcome): handle = native.RouterCoreHandle(config, {"num_retries": 1}) picked = drive(handle.select("primary", {}), None) assert (picked["kind"], picked["deployment_id"], picked["model_group"]) == ("done", "a1", "primary"), picked -assert [step["op"] for step in asked] == ["batch_get"] -assert list(asked[0]) == ["kind", "op", "key", "keys", "value", "ttl", "local", "events"], list(asked[0]) +assert [step["op"] for step in asked] == ["cooldowns"] +assert list(asked[0]) == ["kind", "op", "ids", "events"], list(asked[0]) +assert asked[0]["ids"] == ["a1"], asked[0] call = handle.start("acompletion", '{"model": "primary"}') version = handle.registry_version() assert handle.apply({"delete": ["gone"]}) == version + 1 @@ -2027,20 +2045,20 @@ assert cooled.resolve("primary")["ids"] == ["a1", "a2"] assert cooled.resolve("nope") is None assert drive(call, {"ok": False, "kind": "RateLimitError", "status": 429})["kind"] == "failed" # The core records the failure itself: the 429 cools the deployment down in the host's store. -assert f"deployment:{tried}:cooldown" in store, store +assert ("cooldown", tried) in store, store other = drive(cooled.select("primary", {}), None) assert other["kind"] == "done" and other["deployment_id"] != tried, other call = handle.select("primary", {}) step = call.next(None) -assert step["op"] == "batch_get" +assert step["op"] == "cooldowns" for bad in ([], [None, None], "nope", None): try: call.next(bad) except ValueError: pass else: - raise AssertionError(f"{bad!r} should not answer a batch_get of one key") + raise AssertionError(f"{bad!r} should not answer a cooldowns read of one id") assert drive(call, [None])["kind"] == "done" "# ), diff --git a/litellm/rust/crates/redis-store/src/router_usage/keys.rs b/litellm/rust/crates/redis-store/src/router_usage/keys.rs index 261fbeb1d1d..1482668ab2b 100644 --- a/litellm/rust/crates/redis-store/src/router_usage/keys.rs +++ b/litellm/rust/crates/redis-store/src/router_usage/keys.rs @@ -5,6 +5,12 @@ pub enum Kind { Tpm, #[strum(serialize = "rpm")] Rpm, + /// Input tokens per minute, which the io-token rate limit check reads. + #[strum(serialize = "itpm")] + Itpm, + /// Output tokens per minute. + #[strum(serialize = "otpm")] + Otpm, } #[derive(Clone, Debug, PartialEq, Eq)] @@ -16,8 +22,9 @@ pub struct Counter { } impl Counter { - /// `global_router:{id}:{model}:{tpm|rpm}:{HH-MM}`, where Python formats a missing model as `None`. - pub(super) fn key(&self) -> String { + /// `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. + pub fn key(&self) -> String { let kind: &str = self.kind.into(); let model = self.model.as_deref().unwrap_or("None"); format!( diff --git a/litellm/rust/crates/router-core/src/cache.rs b/litellm/rust/crates/router-core/src/cache.rs index 9811e34a3e3..d32f4723101 100644 --- a/litellm/rust/crates/router-core/src/cache.rs +++ b/litellm/rust/crates/router-core/src/cache.rs @@ -1,9 +1,11 @@ -//! [`RouterCache`]: the DualCache-style store all routing state lives in, under the keys, -//! encodings and TTLs the Python router uses, so Python and Rust workers share it. Method names -//! follow `litellm-cache`'s `BaseCache`, `BatchCache`, `CounterCache` and `DeleteCache`. +//! [`RouterCache`]: the routing state router-core reads and writes, named by what it is (a +//! deployment's cooldown, its allowed-fails count, its per-minute usage, the health checks' +//! states). Where that state lives in Redis, and under which keys, is the host's store's business: +//! router-core never spells a Redis key. use std::{ collections::HashMap, + hash::Hash, sync::Arc, time::{Duration, Instant}, }; @@ -11,149 +13,261 @@ use std::{ use async_trait::async_trait; use litellm_cache::Error; use parking_lot::Mutex; -use serde_json::Value; +use serde::Deserialize; + +use crate::health::CooldownValue; + +/// `RoutingArgs.ttl`: a per-minute counter lives for its minute. +const MINUTE: Duration = Duration::from_secs(60); + +/// One deployment's per-minute usage counters: `model` is the model its usage is counted under, +/// and `minute` is `HH-MM` in UTC. +#[derive(Clone, Debug, PartialEq, Eq, Hash, serde::Serialize)] +pub struct UsageWindow { + pub deployment_id: String, + pub model: String, + pub minute: String, +} + +/// A [`UsageWindow`]'s counts, `None` for a counter nothing has written yet. +#[derive(Clone, Copy, Debug, Default, PartialEq)] +pub struct UsageCounts { + pub tokens: Option, + pub requests: Option, +} + +/// One background health check's result (`DeploymentHealthStateValue`): `timestamp` is when it +/// ran, in Unix seconds. +#[derive(Clone, Copy, Debug, PartialEq, Deserialize)] +pub struct HealthState { + pub is_healthy: bool, + pub timestamp: f64, +} + +/// What this worker saw of a deployment in the current minute. +#[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] +pub struct Outcomes { + pub successes: u64, + pub fails: u64, +} #[async_trait] pub trait RouterCache: Send + Sync { - async fn async_get_cache(&self, key: &str) -> Result, Error>; + /// The stored cooldown of each of `ids`, in order, `None` for one without. + async fn cooldowns(&self, ids: &[String]) -> Result>, Error>; - async fn async_set_cache( + async fn put_cooldown( &self, - key: &str, - value: Value, - ttl: Option, + id: &str, + cooldown: &CooldownValue, + ttl: Duration, ) -> Result<(), Error>; - /// Add `amount` to the number at `key` (0 when missing) and return the new value. The TTL - /// applies when the key is created, as Redis `INCRBYFLOAT` + `EXPIRE NX` does. - async fn async_increment( + /// Count one more failure against `id`'s allowed fails (per `suffix` when one is given) and + /// return the count. `ttl` applies when the counter is created. + async fn increment_allowed_fails( &self, - key: &str, - amount: f64, - ttl: Option, + id: &str, + suffix: Option<&str>, + ttl: Duration, ) -> Result; - async fn async_batch_get_cache(&self, keys: &[String]) -> Result>, Error>; - - async fn async_delete_cache(&self, key: &str) -> Result<(), Error>; - - /// Read this worker's in-memory tier only (`local_only=True`). - async fn async_get_local(&self, key: &str) -> Result, Error> { - self.async_get_cache(key).await - } - - /// Add to a counter in this worker's in-memory tier only (`local_only=True`). - async fn async_increment_local( + /// `_increment_deployment_usage`'s write: tokens and requests in one round trip. + async fn add_usage( &self, - key: &str, - amount: f64, - ttl: Option, - ) -> Result { - self.async_increment(key, amount, ttl).await - } + window: &UsageWindow, + tokens: f64, + requests: f64, + ) -> Result<(), Error>; - /// `_increment_deployment_usage`'s write: add each `(key, amount)` in one pipeline. With a - /// shared tier, this worker's in-memory tier then holds the shared totals, so its reads see - /// what other workers counted up to now. - async fn async_increment_usage( - &self, - increments: &[(String, f64)], - ttl: Option, - ) -> Result<(), Error> { - for (key, amount) in increments { - self.async_increment(key, *amount, ttl).await?; - } - Ok(()) + /// One more request on `window`. + async fn add_request(&self, window: &UsageWindow) -> Result<(), Error>; + + async fn usage(&self, window: &UsageWindow) -> Result; + + /// The background health checks' latest state of each deployment they checked, by id. + async fn health_states(&self) -> Result, Error>; + + /// `id`'s outcomes in this worker, which no other worker sees. + async fn outcomes(&self, id: &str) -> Result; + + async fn count_success(&self, id: &str) -> Result<(), Error>; + + async fn count_failure(&self, id: &str) -> Result<(), Error>; +} + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +enum Counter { + AllowedFails(String, Option), + Tokens(UsageWindow), + Requests(UsageWindow), + Successes(String), + Fails(String), +} + +struct Expiring { + value: T, + expires: Instant, +} + +fn live(entries: &mut HashMap>, key: &K) -> Option { + if entries + .get(key) + .is_some_and(|entry| entry.expires <= Instant::now()) + { + entries.remove(key); } + entries.get(key).map(|entry| entry.value.clone()) } #[derive(Default)] -struct Entry { - value: Value, - expires: Option, +struct State { + cooldowns: HashMap>, + counters: HashMap>, + health_states: HashMap, } -/// An in-process [`RouterCache`], for tests and for state Python keeps `local_only`. +/// An in-process [`RouterCache`], for tests. #[derive(Clone, Default)] pub struct InMemoryRouterCache { - entries: Arc>>, + state: Arc>, } impl InMemoryRouterCache { - fn live(entries: &mut HashMap, key: &str) -> Option { - let expired = entries - .get(key) - .is_some_and(|entry| entry.expires.is_some_and(|at| at <= Instant::now())); - if expired { - entries.remove(key); - } - entries.get(key).map(|entry| entry.value.clone()) + fn count(&self, counter: &Counter) -> Option { + live(&mut self.state.lock().counters, counter) } - pub fn get(&self, key: &str) -> Option { - Self::live(&mut self.entries.lock(), key) - } - - pub fn increment(&self, key: &str, amount: f64, ttl: Option) -> f64 { - let mut entries = self.entries.lock(); - let current = Self::live(&mut entries, key).and_then(|value| value.as_f64()); - let next = current.unwrap_or(0.0) + amount; - let value = serde_json::Number::from_f64(next) - .map(|number| { - if next.fract() == 0.0 && next.abs() < i64::MAX as f64 { - Value::from(next as i64) - } else { - Value::Number(number) - } - }) - .unwrap_or(Value::Null); - let expires = match (current, entries.get(key)) { - (Some(_), Some(entry)) => entry.expires, - _ => ttl.map(|ttl| Instant::now() + ttl), - }; - entries.insert(key.to_owned(), Entry { value, expires }); + fn increment(&self, counter: Counter, amount: f64, ttl: Duration) -> f64 { + let mut state = self.state.lock(); + let next = live(&mut state.counters, &counter).unwrap_or(0.0) + amount; + let expires = state + .counters + .get(&counter) + .map_or_else(|| Instant::now() + ttl, |entry| entry.expires); + state.counters.insert( + counter, + Expiring { + value: next, + expires, + }, + ); next } + + pub fn counted(&self, window: &UsageWindow) -> UsageCounts { + UsageCounts { + tokens: self.count(&Counter::Tokens(window.clone())), + requests: self.count(&Counter::Requests(window.clone())), + } + } + + pub fn set_health_states(&self, states: HashMap) { + self.state.lock().health_states = states; + } } #[async_trait] impl RouterCache for InMemoryRouterCache { - async fn async_get_cache(&self, key: &str) -> Result, Error> { - Ok(self.get(key)) - } - - async fn async_set_cache( - &self, - key: &str, - value: Value, - ttl: Option, - ) -> Result<(), Error> { - let expires = ttl.map(|ttl| Instant::now() + ttl); - self.entries - .lock() - .insert(key.to_owned(), Entry { value, expires }); - Ok(()) - } - - async fn async_increment( - &self, - key: &str, - amount: f64, - ttl: Option, - ) -> Result { - Ok(self.increment(key, amount, ttl)) - } - - async fn async_batch_get_cache(&self, keys: &[String]) -> Result>, Error> { - let mut entries = self.entries.lock(); - Ok(keys + async fn cooldowns(&self, ids: &[String]) -> Result>, Error> { + let mut state = self.state.lock(); + Ok(ids .iter() - .map(|key| Self::live(&mut entries, key)) + .map(|id| live(&mut state.cooldowns, id)) .collect()) } - async fn async_delete_cache(&self, key: &str) -> Result<(), Error> { - self.entries.lock().remove(key); + async fn put_cooldown( + &self, + id: &str, + cooldown: &CooldownValue, + ttl: Duration, + ) -> Result<(), Error> { + self.state.lock().cooldowns.insert( + id.to_owned(), + Expiring { + value: cooldown.clone(), + expires: Instant::now() + ttl, + }, + ); + Ok(()) + } + + async fn increment_allowed_fails( + &self, + id: &str, + suffix: Option<&str>, + ttl: Duration, + ) -> Result { + let counter = Counter::AllowedFails(id.to_owned(), suffix.map(str::to_owned)); + Ok(self.increment(counter, 1.0, ttl)) + } + + async fn add_usage( + &self, + window: &UsageWindow, + tokens: f64, + requests: f64, + ) -> Result<(), Error> { + self.increment(Counter::Tokens(window.clone()), tokens, MINUTE); + self.increment(Counter::Requests(window.clone()), requests, MINUTE); + Ok(()) + } + + async fn add_request(&self, window: &UsageWindow) -> Result<(), Error> { + self.increment(Counter::Requests(window.clone()), 1.0, MINUTE); + Ok(()) + } + + async fn usage(&self, window: &UsageWindow) -> Result { + Ok(self.counted(window)) + } + + async fn health_states(&self) -> Result, Error> { + Ok(self.state.lock().health_states.clone()) + } + + async fn outcomes(&self, id: &str) -> Result { + let count = |counter: Counter| self.count(&counter).map_or(0, |count| count as u64); + Ok(Outcomes { + successes: count(Counter::Successes(id.to_owned())), + fails: count(Counter::Fails(id.to_owned())), + }) + } + + async fn count_success(&self, id: &str) -> Result<(), Error> { + self.increment(Counter::Successes(id.to_owned()), 1.0, MINUTE); + Ok(()) + } + + async fn count_failure(&self, id: &str) -> Result<(), Error> { + self.increment(Counter::Fails(id.to_owned()), 1.0, MINUTE); Ok(()) } } + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::*; + + #[rstest] + #[tokio::test] + async fn usage_counters_expire_so_tomorrows_same_minute_starts_empty() { + let cache = InMemoryRouterCache::default(); + let window = UsageWindow { + deployment_id: "dep-1".to_owned(), + model: "gpt".to_owned(), + minute: "12-00".to_owned(), + }; + + cache.add_usage(&window, 90.0, 1.0).await.unwrap(); + + let state = cache.state.lock(); + let within_a_minute = Instant::now() + MINUTE; + let expires = |counter: Counter| state.counters[&counter].expires <= within_a_minute; + assert!(expires(Counter::Tokens(window.clone()))); + assert!(expires(Counter::Requests(window))); + } +} diff --git a/litellm/rust/crates/router-core/src/health.rs b/litellm/rust/crates/router-core/src/health.rs index eaa61dac7ba..5876a71bba5 100644 --- a/litellm/rust/crates/router-core/src/health.rs +++ b/litellm/rust/crates/router-core/src/health.rs @@ -1,17 +1,9 @@ -//! [`Health`]: cooldowns, allowed-fails counters and per-minute usage, kept in the shared -//! [`RouterCache`] under the Python router's keys, encodings and TTLs: -//! -//! | state | key | value | TTL | -//! |---|---|---|---| -//! | cooldown | `deployment:{id}:cooldown` | `{exception_received, status_code, timestamp, cooldown_time}` | cooldown time | -//! | allowed fails | `deployment:{id}:allowed_fails[:{suffix}]` | int counter | cooldown time | -//! | usage | `global_router:{id}:{model}:tpm:{HH-MM}` and `:rpm:` | int counter | 60s | -//! | outcomes (local only) | `{id}:successes`, `{id}:fails` | int counter | 60s | +//! [`Health`]: cooldowns, allowed-fails counters and per-minute usage, kept in the host's +//! [`RouterCache`] with the Python router's TTLs, and this worker's per-minute outcomes. //! //! From `router_utils/cooldown_cache.py`, `router_utils/cooldown_handlers.py`, -//! `router_utils/router_callbacks/track_deployment_metrics.py`, `RouterCacheEnum` in -//! `types/router.py` and `Router._increment_usage` / `deployment_callback_on_failure` in -//! `router.py`. +//! `router_utils/router_callbacks/track_deployment_metrics.py` and `Router._increment_usage` / +//! `deployment_callback_on_failure` in `router.py`. use std::{ sync::Arc, @@ -22,15 +14,13 @@ use serde::{Deserialize, Serialize}; use tracing::debug; use crate::{ - cache::RouterCache, + cache::{Outcomes, RouterCache, UsageWindow}, error::{AttemptError, Error, FailureKind}, outcome::{OutcomeCallbacks, OutcomeWrite}, registry::{Deployment, Registry}, settings::Settings, }; -/// `RoutingArgs.ttl`: per-minute usage keys expire after a minute. -pub const USAGE_TTL: Duration = Duration::from_secs(60); /// `DEFAULT_FAILURE_THRESHOLD_PERCENT`. pub const FAILURE_THRESHOLD_PERCENT: f64 = 0.5; /// `DEFAULT_FAILURE_THRESHOLD_MINIMUM_REQUESTS`. @@ -68,37 +58,6 @@ pub fn current_minute(time: SystemTime) -> String { format!("{:02}-{:02}", seconds / 3600, (seconds % 3600) / 60) } -pub fn cooldown_key(id: &str) -> String { - format!("deployment:{id}:cooldown") -} - -pub fn allowed_fails_key(id: &str, suffix: Option<&str>) -> String { - match suffix { - Some(suffix) => format!("deployment:{id}:allowed_fails:{suffix}"), - None => format!("deployment:{id}:allowed_fails"), - } -} - -#[derive(Clone, Copy, Debug, PartialEq, Eq, strum::Display, strum::EnumString)] -pub enum UsageMetric { - #[strum(serialize = "tpm")] - Tpm, - #[strum(serialize = "rpm")] - Rpm, - /// Input tokens per minute, which the io-token rate limit check counts. - #[strum(serialize = "itpm")] - Itpm, - /// Output tokens per minute. - #[strum(serialize = "otpm")] - Otpm, -} - -/// `RouterCacheEnum.TPM` / `RPM` / `ITPM` / `OTPM`: `model` is the deployment's -/// `litellm_params.model`. -pub fn usage_key(id: &str, model: &str, metric: UsageMetric, minute: &str) -> String { - format!("global_router:{id}:{model}:{metric}:{minute}") -} - /// `CooldownCacheValue`. #[derive(Clone, Debug, PartialEq, Serialize, Deserialize)] pub struct CooldownValue { @@ -260,14 +219,13 @@ impl Health { &self, ids: &[String], ) -> Result, Error> { - let keys: Vec = ids.iter().map(|id| cooldown_key(id)).collect(); - let values = self.cache.async_batch_get_cache(&keys).await?; + let values = self.cache.cooldowns(ids).await?; let now = self.now(); Ok(ids .iter() .zip(values) .filter_map(|(id, value)| { - let value: CooldownValue = serde_json::from_value(value?).ok()?; + let value = value?; (value.timestamp + value.cooldown_time - now > 0.0).then(|| (id.clone(), value)) }) .collect()) @@ -286,46 +244,23 @@ impl Health { timestamp: self.now(), cooldown_time: cooldown.as_secs_f64(), }; - let value = serde_json::to_value(value).unwrap_or_default(); - self.cache - .async_set_cache(&cooldown_key(id), value, Some(cooldown)) - .await?; + self.cache.put_cooldown(id, &value, cooldown).await?; Ok(()) } - async fn local_count(&self, key: &str) -> Result { - Ok(self - .cache - .async_get_local(key) - .await? - .and_then(|value| value.as_f64()) - .map_or(0, |count| count as u64)) + /// The successes and fails this worker saw this minute. + pub async fn outcomes_this_minute(&self, id: &str) -> Result { + Ok(self.cache.outcomes(id).await?) } - /// The per-minute (successes, fails) this worker saw (`{id}:successes`, `{id}:fails`, kept - /// `local_only`). - pub async fn outcomes_this_minute(&self, id: &str) -> Result<(u64, u64), Error> { - Ok(( - self.local_count(&format!("{id}:successes")).await?, - self.local_count(&format!("{id}:fails")).await?, - )) - } - - /// `increment_deployment_successes_for_current_minute`: returns the key it counted under. - pub async fn count_success(&self, id: &str) -> Result { - let key = format!("{id}:successes"); - self.cache - .async_increment_local(&key, 1.0, Some(USAGE_TTL)) - .await?; - Ok(key) + /// `increment_deployment_successes_for_current_minute`. + pub async fn count_success(&self, id: &str) -> Result<(), Error> { + Ok(self.cache.count_success(id).await?) } /// `increment_deployment_failures_for_current_minute`. pub async fn count_failure(&self, id: &str) -> Result<(), Error> { - self.cache - .async_increment_local(&format!("{id}:fails"), 1.0, Some(USAGE_TTL)) - .await?; - Ok(()) + Ok(self.cache.count_failure(id).await?) } /// Cooldown duration precedence from `deployment_callback_on_failure`: the deployment's @@ -499,7 +434,7 @@ impl Health { .await; } if settings.allowed_fails_policy.is_none() && !settings.allowed_fails_set() { - let (successes, fails) = self.outcomes_this_minute(&deployment.id).await?; + let Outcomes { successes, fails } = self.outcomes_this_minute(&deployment.id).await?; let total = successes + fails; let percent_fails = if total > 0 { fails as f64 / total as f64 @@ -534,10 +469,7 @@ impl Health { allowed: f64, ttl: Duration, ) -> Result { - let fails = self - .cache - .async_increment(&allowed_fails_key(id, suffix), 1.0, Some(ttl)) - .await?; + let fails = self.cache.increment_allowed_fails(id, suffix, ttl).await?; Ok(fails > allowed) } @@ -568,34 +500,22 @@ impl Health { /// `_increment_deployment_usage`, for a deployment that tracks usage: tokens and requests /// under `model` (metadata's `deployment`, which a wildcard route fills; else the - /// deployment's own), unless there is nothing to add. Returns the tokens counter's key when - /// it added to it. + /// deployment's own), unless there is nothing to add. Returns the window it added to. pub async fn add_usage( &self, deployment: &Deployment, model: Option<&str>, total_tokens: f64, requests: f64, - ) -> Result, Error> { + ) -> Result, Error> { if total_tokens <= 0.0 && requests <= 0.0 { return Ok(None); } - let minute = current_minute(self.clock.now()); - let model = model.unwrap_or(&deployment.params.model); - let tpm_key = usage_key(&deployment.id, model, UsageMetric::Tpm, &minute); + let window = self.usage_window(deployment, model); self.cache - .async_increment_usage( - &[ - (tpm_key.clone(), total_tokens), - ( - usage_key(&deployment.id, model, UsageMetric::Rpm, &minute), - requests, - ), - ], - Some(USAGE_TTL), - ) + .add_usage(&window, total_tokens, requests) .await?; - Ok(Some(tpm_key)) + Ok(Some(window)) } /// `async_deployment_callback_on_failure`: a failed request counts against the deployment's @@ -605,38 +525,27 @@ impl Health { deployment: &Deployment, model: Option<&str>, ) -> Result<(), Error> { - let minute = current_minute(self.clock.now()); - self.cache - .async_increment( - &usage_key( - &deployment.id, - model.unwrap_or(&deployment.params.model), - UsageMetric::Rpm, - &minute, - ), - 1.0, - Some(USAGE_TTL), - ) - .await?; + let window = self.usage_window(deployment, model); + self.cache.add_request(&window).await?; Ok(()) } /// This minute's (tpm, rpm) for a deployment. pub async fn usage(&self, deployment: &Deployment) -> Result<(u64, u64), Error> { - let minute = current_minute(self.clock.now()); - let model = &deployment.params.model; - let keys = [UsageMetric::Tpm, UsageMetric::Rpm] - .map(|metric| usage_key(&deployment.id, model, metric, &minute)); - let values = self.cache.async_batch_get_cache(&keys).await?; - let count = |index: usize| { - values - .get(index) - .cloned() - .flatten() - .and_then(|value| value.as_f64()) - .map_or(0, |count| count as u64) - }; - Ok((count(0), count(1))) + let counts = self + .cache + .usage(&self.usage_window(deployment, None)) + .await?; + let count = |value: Option| value.map_or(0, |count| count as u64); + Ok((count(counts.tokens), count(counts.requests))) + } + + fn usage_window(&self, deployment: &Deployment, model: Option<&str>) -> UsageWindow { + UsageWindow { + deployment_id: deployment.id.clone(), + model: model.unwrap_or(&deployment.params.model).to_owned(), + minute: current_minute(self.clock.now()), + } } } @@ -648,7 +557,7 @@ mod tests { use super::*; use crate::{ - cache::InMemoryRouterCache, + cache::{InMemoryRouterCache, UsageCounts}, config::{RouterConfig, RouterDynamicConfig}, registry::{DeploymentSpec, RoutingGroupSpec}, }; @@ -765,7 +674,10 @@ mod tests { .unwrap(); assert_eq!(cooldown.is_some(), cooled); - assert_eq!(health.outcomes_this_minute("a").await.unwrap().1, counted); + assert_eq!( + health.outcomes_this_minute("a").await.unwrap().fails, + counted + ); } #[rstest] @@ -798,7 +710,7 @@ mod tests { assert_eq!(cooldown, None); assert!(failure.is_a(base_class)); - assert_eq!(health.outcomes_this_minute("a").await.unwrap().1, 1); + assert_eq!(health.outcomes_this_minute("a").await.unwrap().fails, 1); } /// `litellm.allowed_fails` written after the build: the copied 3 starts counting instead of @@ -918,7 +830,7 @@ mod tests { .unwrap(); assert_eq!(cooldown.is_some(), cooled); - assert_eq!(health.outcomes_this_minute("a").await.unwrap().1, 1); + assert_eq!(health.outcomes_this_minute("a").await.unwrap().fails, 1); } #[rstest] @@ -936,29 +848,17 @@ mod tests { .await .unwrap(); - let minute = current_minute(clock.now()); + let window = UsageWindow { + deployment_id: "a".to_owned(), + model: keyed.to_owned(), + minute: current_minute(clock.now()), + }; assert_eq!( - ( - cache.get(&usage_key("a", keyed, UsageMetric::Tpm, &minute)), - cache.get(&usage_key("a", keyed, UsageMetric::Rpm, &minute)), - ), - (Some(serde_json::json!(15)), Some(serde_json::json!(1))), - ); - } - - #[rstest] - fn keys_match_python() { - assert_eq!( - ( - cooldown_key("abc"), - allowed_fails_key("abc", Some("generic")), - usage_key("abc", "openai/gpt-4o", UsageMetric::Tpm, "09-05"), - ), - ( - "deployment:abc:cooldown".to_owned(), - "deployment:abc:allowed_fails:generic".to_owned(), - "global_router:abc:openai/gpt-4o:tpm:09-05".to_owned(), - ) + cache.counted(&window), + UsageCounts { + tokens: Some(15.0), + requests: Some(1.0), + }, ); } } diff --git a/litellm/rust/crates/router-core/src/health_routing.rs b/litellm/rust/crates/router-core/src/health_routing.rs index bc3e43a6e92..d6b52d3c36f 100644 --- a/litellm/rust/crates/router-core/src/health_routing.rs +++ b/litellm/rust/crates/router-core/src/health_routing.rs @@ -1,44 +1,25 @@ -//! Health-check-driven routing: the background health checks' results, kept in the shared -//! [`RouterCache`](crate::cache::RouterCache) under one key as Python's `DeploymentHealthCache` -//! writes them, and the filter `_async_filter_health_check_unhealthy_deployments` / +//! Health-check-driven routing: the background health checks' results, read from the host's +//! [`RouterCache`](crate::cache::RouterCache) as Python's `DeploymentHealthCache` writes them, and +//! the filter `_async_filter_health_check_unhealthy_deployments` / //! `_filter_health_check_unhealthy_deployments` run before the cooldown filter. -use std::{collections::BTreeSet, sync::Arc, time::UNIX_EPOCH}; +use std::{ + collections::{BTreeSet, HashMap}, + sync::Arc, + time::UNIX_EPOCH, +}; -use serde_json::Value; - -use crate::{health::Health, registry::Deployment, settings::Settings}; - -/// `DeploymentHealthCache.CACHE_KEY`. -pub const HEALTH_STATE_KEY: &str = "litellm:health_check:deployment_health_state"; +use crate::{cache::HealthState, health::Health, registry::Deployment, settings::Settings}; /// `_extract_unhealthy_ids`: the ids marked unhealthy whose result is younger than `staleness`. -pub fn unhealthy_ids(raw: Option<&Value>, now: f64, staleness: f64) -> BTreeSet { - let Some(Value::Object(states)) = raw else { - return BTreeSet::new(); - }; +pub fn unhealthy_ids( + states: &HashMap, + now: f64, + staleness: f64, +) -> BTreeSet { states .iter() - .filter(|(_, state)| { - let Value::Object(state) = state else { - return false; - }; - // `not state.get("is_healthy", True)`: any falsy value is unhealthy. - let unhealthy = match state.get("is_healthy") { - None => false, - Some(Value::Null) => true, - Some(Value::Bool(value)) => !value, - Some(Value::Number(value)) => value.as_f64() == Some(0.0), - Some(Value::String(value)) => value.is_empty(), - Some(Value::Array(value)) => value.is_empty(), - Some(Value::Object(value)) => value.is_empty(), - }; - let timestamp = state - .get("timestamp") - .and_then(Value::as_f64) - .unwrap_or(0.0); - unhealthy && now - timestamp < staleness - }) + .filter(|(_, state)| !state.is_healthy && now - state.timestamp < staleness) .map(|(id, _)| id.clone()) .collect() } @@ -71,22 +52,13 @@ impl Health { return deployments; } // A read that fails counts as no health state. - let raw = self - .cache() - .async_get_cache(HEALTH_STATE_KEY) - .await - .ok() - .flatten(); + let states = self.cache().health_states().await.unwrap_or_default(); let now = self .now_time() .duration_since(UNIX_EPOCH) .map(|elapsed| elapsed.as_secs_f64()) .unwrap_or_default(); - let unhealthy = unhealthy_ids( - raw.as_ref(), - now, - settings.health_check_staleness.as_secs_f64(), - ); + let unhealthy = unhealthy_ids(&states, now, settings.health_check_staleness.as_secs_f64()); if unhealthy.is_empty() { return deployments; } @@ -109,18 +81,24 @@ impl Health { #[cfg(test)] mod tests { use rstest::rstest; - use serde_json::json; use super::*; #[rstest] - #[case::fresh_unhealthy(json!({"a": {"is_healthy": false, "timestamp": 95.0}}), vec!["a"])] - #[case::stale(json!({"a": {"is_healthy": false, "timestamp": 50.0}}), vec![])] - #[case::healthy(json!({"a": {"is_healthy": true, "timestamp": 95.0}}), vec![])] - #[case::no_flag(json!({"a": {"timestamp": 95.0}}), vec![])] - #[case::not_a_dict(json!({"a": 1}), vec![])] - fn unhealthy_ids_follow_python(#[case] raw: Value, #[case] expected: Vec<&str>) { - let ids = unhealthy_ids(Some(&raw), 100.0, 10.0); + #[case::fresh_unhealthy(false, 95.0, vec!["a"])] + #[case::stale(false, 50.0, vec![])] + #[case::healthy(true, 95.0, vec![])] + fn unhealthy_ids_follow_python( + #[case] is_healthy: bool, + #[case] timestamp: f64, + #[case] expected: Vec<&str>, + ) { + let state = HealthState { + is_healthy, + timestamp, + }; + let states = HashMap::from([("a".to_owned(), state)]); + let ids = unhealthy_ids(&states, 100.0, 10.0); assert_eq!(ids.iter().map(String::as_str).collect::>(), expected); } } diff --git a/litellm/rust/crates/router-core/src/lib.rs b/litellm/rust/crates/router-core/src/lib.rs index e61d570c78f..dc801bd14a2 100644 --- a/litellm/rust/crates/router-core/src/lib.rs +++ b/litellm/rust/crates/router-core/src/lib.rs @@ -118,7 +118,7 @@ use arc_swap::ArcSwap; use parking_lot::Mutex; pub use crate::{ - cache::{InMemoryRouterCache, RouterCache}, + cache::{HealthState, InMemoryRouterCache, Outcomes, RouterCache, UsageCounts, UsageWindow}, call::RouterResponse, config::{Op, RouterConfig, RouterDynamicConfig, RoutingStrategy}, error::{AttemptError, ConfigError, Error, FailureKind, Payload, Unsupported}, @@ -436,14 +436,14 @@ impl Router { /// A success a logging callback saw (`deployment_callback_on_success`): counted for a /// configured deployment, with its tokens and requests when it tracks usage. Returns the - /// tokens counter's key when it added usage (`_increment_deployment_usage`), else `None`. + /// window it added usage to (`_increment_deployment_usage`), else `None`. pub async fn report_success( &self, id: &str, total_tokens: f64, requests: f64, report: SuccessReport, - ) -> Result, Error> { + ) -> Result, Error> { report .callbacks .record( @@ -459,7 +459,7 @@ impl Router { total_tokens: f64, requests: f64, report: &SuccessReport, - ) -> Result, Error> { + ) -> Result, Error> { let snapshot = self.core.snapshot.load_full(); let Some(deployment) = snapshot.registry.deployment(id) else { return Ok(None); @@ -482,16 +482,16 @@ impl Router { .await } - /// `sync_deployment_callback_on_success`: count a success against `id`, configured or not, - /// and return the counter's key. + /// `sync_deployment_callback_on_success`: count a success against `id`, configured or not. + /// Returns whether it counted one. pub async fn count_success( &self, id: &str, callbacks: OutcomeCallbacks, - ) -> Result, Error> { + ) -> Result { callbacks .record(OutcomeWrite::Success, async { - self.core.health.count_success(id).await.map(Some) + self.core.health.count_success(id).await.map(|()| true) }) .await } diff --git a/litellm/rust/crates/router-core/tests/call_loop.rs b/litellm/rust/crates/router-core/tests/call_loop.rs index 732eb13371b..f95ff1f3994 100644 --- a/litellm/rust/crates/router-core/tests/call_loop.rs +++ b/litellm/rust/crates/router-core/tests/call_loop.rs @@ -8,12 +8,8 @@ use litellm_router_core::{ Attempt, AttemptError, CallbackFiltered, Deployment, Error, FailureKind, FailureReport, Host, InMemoryRouterCache, Op, Output, Payload, Router, RouterCache, RouterConfig, RouterDynamicConfig, RouterEvent, RouterRequest, RoutingContext, SeededRandomness, SelectStage, - SuccessReport, - call::StreamOrigin, - fallback::Fallbacks, - health::{UsageMetric, current_minute, usage_key}, - host::UncountedRetries, - select::PassThrough, + SuccessReport, UsageWindow, call::StreamOrigin, fallback::Fallbacks, health::current_minute, + host::UncountedRetries, select::PassThrough, }; use parking_lot::Mutex; use rstest::{fixture, rstest}; @@ -325,12 +321,12 @@ async fn a_rate_limit_cools_the_deployment_down(config: RouterConfig) { } #[rstest] -#[case::now(false, (Some(1.0), (10, 1)))] -#[case::deferred(true, (None, (10, 1)))] +#[case::now(false, (1, (10, 1)))] +#[case::deferred(true, (0, (10, 1)))] #[tokio::test] async fn a_deferred_attempt_counts_its_request_and_tokens_but_no_success( #[case] deferred: bool, - #[case] expected: (Option, (u64, u64)), + #[case] expected: (u64, (u64, u64)), ) { let config: RouterConfig = serde_json::from_value(json!({ "model_list": [ @@ -348,10 +344,7 @@ async fn a_deferred_attempt_counts_its_request_and_tokens_but_no_success( .await .unwrap(); let deployment = router.registry().deployment("a1").unwrap(); - let successes = host - .cache - .get("a1:successes") - .and_then(|count| count.as_f64()); + let successes = host.cache.outcomes("a1").await.unwrap().successes; let usage = router.health().usage(&deployment).await.unwrap(); assert_eq!((successes, usage), expected); } @@ -422,28 +415,33 @@ async fn a_reported_success_keeps_the_filled_wildcard_model() { }; router.report_success("w1", 7.0, 0.0, report).await.unwrap(); let minute = current_minute(std::time::SystemTime::now()); - let key = |model: &str| usage_key("w1", model, UsageMetric::Tpm, &minute); + let tokens = |model: &str| { + host.cache + .counted(&UsageWindow { + deployment_id: "w1".to_owned(), + model: model.to_owned(), + minute: minute.clone(), + }) + .tokens + }; assert_eq!( - ( - host.cache.get(&key("openai/gpt-4o")), - host.cache.get(&key("openai/*")), - ), - (Some(json!(7)), None), + (tokens("openai/gpt-4o"), tokens("openai/*")), + (Some(7.0), None), ); } /// `async_deployment_callback_on_failure`'s request is written only when it runs: a failure the /// host says ran no failure callback (`recorded`) writes none, nor does a sync call's. #[rstest] -#[case::the_first_logged_failure(false, false, 1, Some(1.0))] -#[case::a_later_failure_on_the_shared_logging_object(true, false, 0, None)] -#[case::a_sync_call(false, true, 0, Some(1.0))] +#[case::the_first_logged_failure(false, false, 1, 1)] +#[case::a_later_failure_on_the_shared_logging_object(true, false, 0, 0)] +#[case::a_sync_call(false, true, 0, 1)] #[tokio::test] async fn a_failed_attempt_counts_its_request_only_when_its_callback_ran( #[case] recorded: bool, #[case] uncounted_request: bool, #[case] requests: u64, - #[case] fails: Option, + #[case] fails: u64, ) { let config: RouterConfig = serde_json::from_value(json!({ "model_list": [ @@ -459,7 +457,7 @@ async fn a_failed_attempt_counts_its_request_only_when_its_callback_ran( let _ = router.call(Op::from("acompletion"), request()).await; let deployment = router.registry().deployment("a1").unwrap(); let usage = router.health().usage(&deployment).await.unwrap(); - let counted = host.cache.get("a1:fails").and_then(|count| count.as_f64()); + let counted = host.cache.outcomes("a1").await.unwrap().fails; assert_eq!((usage.1, counted), (requests, fails)); } @@ -499,9 +497,9 @@ async fn a_reported_failure_counts_its_request_when_asked( .await .unwrap(); let deployment = router.registry().deployment("a1").unwrap(); - let fails = host.cache.get("a1:fails").and_then(|count| count.as_f64()); + let fails = host.cache.outcomes("a1").await.unwrap().fails; let usage = router.health().usage(&deployment).await.unwrap(); - assert_eq!((fails, usage.1), (Some(1.0), requests)); + assert_eq!((fails, usage.1), (1, requests)); } #[rstest] diff --git a/litellm/rust_bridge/_native.pyi b/litellm/rust_bridge/_native.pyi index bde79aeac48..22714e7f325 100644 --- a/litellm/rust_bridge/_native.pyi +++ b/litellm/rust_bridge/_native.pyi @@ -890,20 +890,43 @@ class _PreRouteStep(_Step): fallback_depth: int retry: int -# One RouterCache op against the router's DualCache. `key` is set for every op but `batch_get`, -# and `increment_usage`, which set `keys`; `value` is the JSON to store (`set`), the amount -# (`increment`) or one amount per key (`increment_usage`); `ttl` is in seconds; `local` reads or -# writes this worker's in-memory tier only. Outcomes: get -> the JSON value or None; batch_get -> -# one JSON value or None per key, in order; increment -> the new int or float; set, delete and -# increment_usage -> None. +# One deployment's per-minute usage counters: `model` is the model its usage is counted under, +# `minute` is `HH-MM` in UTC. +class _UsageWindow(TypedDict): + deployment_id: str + model: str + minute: str + +# One RouterCache op, answered from the router's stores. It names the state it touches, never a +# Redis key. Fields by op: cooldowns (`ids`) -> one stored cooldown or None per id; put_cooldown +# (`id`, `cooldown`, `ttl`) -> None; increment_allowed_fails (`id`, `suffix`, `ttl`) -> the new +# count; add_usage (`window`, `tokens`, `requests`) -> None; add_request (`window`) -> None; +# usage (`window`) -> [tokens, requests], each a number or None; health_states -> each checked +# deployment's `{is_healthy: bool, timestamp: float}` by id, or None; outcomes (`id`) -> this +# worker's [successes, fails] of the current minute; count_success (`id`) and count_failure (`id`) +# -> None. `ttl` is in seconds. class _CacheStep(_Step): kind: Literal["cache"] - op: Literal["get", "set", "increment", "increment_usage", "batch_get", "delete"] - key: str | None - keys: list[str] | None - value: JsonValue - ttl: float | None - local: bool + op: Literal[ + "cooldowns", + "put_cooldown", + "increment_allowed_fails", + "add_usage", + "add_request", + "usage", + "health_states", + "outcomes", + "count_success", + "count_failure", + ] + ids: NotRequired[list[str]] + id: NotRequired[str] + cooldown: NotRequired[JsonValue] + suffix: NotRequired[str | None] + ttl: NotRequired[float] + window: NotRequired[_UsageWindow] + tokens: NotRequired[int | float] + requests: NotRequired[int | float] # One lookup of a model-info query (router-core `ModelInfoLookup`): its `lookup` name and that # lookup's arguments, answered as the lookup documents. @@ -1012,11 +1035,16 @@ class _ReportedStep(_Step): # The cooldown applied, in seconds; None when the deployment was not cooled down. cooldown_time: float | None -# The end of a `RouterCoreHandle.report_success` or `count_success`. +# The end of a `RouterCoreHandle.report_success`. class _CountedStep(_Step): kind: Literal["done"] key: str +# The end of a `RouterCoreHandle.count_success`. +class _SuccessCountedStep(_Step): + kind: Literal["done"] + counted: bool + # One `os.environ/` reference a `secrets` step asks for: a top-level param (`key`), or an # env-only `dataSources[source].parameters[key]` (`env_only`). class _SecretRef(TypedDict): @@ -1084,7 +1112,7 @@ _HealthySteps = _SelectStep | _PreRouteStep | _CacheStep | _HealthyStep | _Faile _BuildSteps = _SecretsStep | _BuildStep | _BuiltStep | _FailedStep _ModelInfoSteps = _ModelInfoStep | _ModelInfoDoneStep | _FailedStep _FailureReportSteps = _CacheStep | _ReportedStep | _FailedStep -_SuccessReportSteps = _CacheStep | _CountedStep | _FailedStep +_SuccessReportSteps = _CacheStep | _CountedStep | _SuccessCountedStep | _FailedStep class _AttemptSucceeded(TypedDict): ok: Literal[True] @@ -1181,12 +1209,10 @@ class RouterCoreHandle: tracks_usage: bool | None = None, callbacks: Mapping[str, bool] | None = None, ) -> RouterCall[_SuccessReportSteps]: ... - # Count a success against any id, driven like `report_failure`; `done` carries `key`. + # Count a success against any id, driven like `report_failure`; `done` carries `counted`. def count_success( self, id: str, callbacks: Mapping[str, bool] | None = None ) -> RouterCall[_SuccessReportSteps]: ... - # [(id, tpm_key, rpm_key)] for this minute, or None for a group with no deployment. - def usage_keys(self, model_group: str) -> list[tuple[str, str, str]] | None: ... # Every deployment as a `model_list` entry, in config order. def deployments(self) -> list[dict[str, JsonValue]]: ... def deployments_from(self, start: int) -> list[dict[str, JsonValue]]: ... diff --git a/tests/unit/router_rust/test_routing_stores.py b/tests/unit/router_rust/test_routing_stores.py index 6fc19cdcfa5..352e2819204 100644 --- a/tests/unit/router_rust/test_routing_stores.py +++ b/tests/unit/router_rust/test_routing_stores.py @@ -8,7 +8,16 @@ import pytest from litellm.redis.stores.connection import connect from litellm.redis.stores.router_cooldowns import CooldownCacheValue, RouterCooldownStore from litellm.redis.stores.router_usage import RouterUsageStore -from litellm.router_rust.routing_stores import Handled, RoutingStores +from litellm.router_rust.routing_stores import ( + CACHE_STEP, + AddRequestStep, + HealthStates, + HealthStatesStep, + IncrementAllowedFailsStep, + PutCooldownStep, + RoutingStores, + UsageStep, +) from litellm.router_utils.cooldown_cache import CooldownCache from tests.unit.redis.stores.test_router_cooldowns import FakeCooldowns @@ -49,33 +58,111 @@ class SlowRedis(FakeCooldowns): return self.increment_allowed_fails(model_id, suffix, ttl_seconds) +class NoHealthStates: + async def async_get_states(self) -> object: + return None + + def get_states(self) -> object: + return None + + +@dataclass(frozen=True) +class CachedHealthStates: + raw: object + + async def async_get_states(self) -> object: + return self.raw + + def get_states(self) -> object: + return self.raw + + +def _stores(health: HealthStates) -> RoutingStores: + cooldowns: Final = CooldownCache(store=RouterCooldownStore(FakeCooldowns()), default_cooldown_time=60) + return RoutingStores(cooldowns, RouterUsageStore(connect(None).router_usage()), health) + + async def _another_request(ran: asyncio.Event) -> None: ran.set() @pytest.mark.asyncio @pytest.mark.parametrize( - ("op", "key", "value", "answer"), + ("raw", "answer"), [ - ("set", "deployment:dep:cooldown", _ENTRY, None), - ("increment", "deployment:dep:allowed_fails:timeout", 1, 1), + ({"op": "put_cooldown", "id": "dep", "cooldown": _ENTRY, "ttl": 60.0}, None), + ({"op": "increment_allowed_fails", "id": "dep", "suffix": "timeout", "ttl": 60.0}, 1), ], ) async def test_an_async_cooldown_step_waiting_on_redis_lets_other_requests_run( - op: str, key: str, value: object, answer: object + raw: Mapping[str, object], answer: object ) -> None: """router-core's failure path writes cooldowns and allowed-fails counts; on an async call that Redis round trip must not hold the event loop, or every request on the worker stalls behind it.""" other_request: Final = asyncio.Event() handle: Final = SlowRedis(has_redis=True, other_request=other_request) cooldowns: Final = CooldownCache(store=RouterCooldownStore(handle), default_cooldown_time=60) - stores: Final = RoutingStores(cooldowns, RouterUsageStore(connect(None).router_usage())) + stores: Final = RoutingStores(cooldowns, RouterUsageStore(connect(None).router_usage()), NoHealthStates()) + step: Final = CACHE_STEP.validate_python(raw) + assert isinstance(step, PutCooldownStep | IncrementAllowedFailsStep) - step: Final = asyncio.create_task(stores.answer(op, key, value, 60.0)) + answering: Final = asyncio.create_task(stores.answer(step)) other: Final = asyncio.create_task(_another_request(other_request)) - handled: Final = await step + handled: Final = await answering await other - assert handled == Handled(answer) + assert handled == answer assert handle.overlapped is True assert handle.targets == ["router_cooldowns"] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("raw", "states"), + [ + pytest.param( + {"dep": {"is_healthy": False, "timestamp": 95.0, "reason": "timeout"}}, + {"dep": {"is_healthy": False, "timestamp": 95.0}}, + id="unhealthy", + ), + pytest.param({"dep": {"timestamp": 95.0}}, {"dep": {"is_healthy": True, "timestamp": 95.0}}, id="no_flag"), + pytest.param( + {"dep": {"is_healthy": 0, "timestamp": 95}}, + {"dep": {"is_healthy": False, "timestamp": 95.0}}, + id="falsy_flag", + ), + pytest.param( + {"dep": {"is_healthy": False, "timestamp": "soon"}}, + {"dep": {"is_healthy": False, "timestamp": 0.0}}, + id="timestamp_that_is_no_number", + ), + pytest.param({"dep": 1}, {}, id="not_a_record"), + pytest.param(None, {}, id="nothing_cached"), + ], +) +async def test_health_states_reach_router_core_as_the_python_filter_reads_them( + raw: object, states: Mapping[str, object] +) -> None: + """router-core takes each state as a strict bool and float, so the leniency of ``_extract_unhealthy_ids`` + (a falsy flag is unhealthy, a missing one healthy, a bad timestamp stale) has to be applied before it crosses.""" + stores: Final = _stores(CachedHealthStates(raw)) + step: Final = CACHE_STEP.validate_python({"op": "health_states"}) + assert isinstance(step, HealthStatesStep) + + assert await stores.answer(step) == states + assert stores.answer_sync(step) == states + + +@pytest.mark.asyncio +async def test_a_failed_request_counts_one_request_and_no_tokens() -> None: + stores: Final = _stores(NoHealthStates()) + window: Final = {"deployment_id": "dep", "model": "gpt", "minute": "09-05"} + counted: Final = CACHE_STEP.validate_python({"op": "add_request", "window": window}) + read: Final = CACHE_STEP.validate_python({"op": "usage", "window": window}) + assert isinstance(counted, AddRequestStep) + assert isinstance(read, UsageStep) + + await stores.answer(counted) + stores.answer_sync(counted) + + assert await stores.answer(read) == (None, 2) diff --git a/tests/unit/test_internal_context.py b/tests/unit/test_internal_context.py index d285d6fa619..537afda84aa 100644 --- a/tests/unit/test_internal_context.py +++ b/tests/unit/test_internal_context.py @@ -26,12 +26,14 @@ _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",) -# Helpers that act on a cache handed in by the declaring caller, or forward to the response-cache facade. +# Helpers that act on a cache handed in by the declaring caller, forward to the response-cache facade, or call a +# ``litellm.redis.stores`` store, which declares its own family. _CACHE_PARAMETER_HELPERS: Final = frozenset( { "litellm/llms/anthropic/prompt_cache_prediction.py", "litellm/proxy/common_utils/cache_coordinator.py", "litellm/proxy/common_utils/user_api_key_cache.py", + "litellm/router_rust/routing_stores.py", "litellm/router_utils/routing_prefetch.py", "litellm/utils.py", }