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