From 9a1acee4a765549702e7f816cfd6aafd027de43b Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 13 Nov 2026 12:22:28 +0900 Subject: [PATCH] feat(redis): serve the router's session pins from a native store (#50870) Deployment affinity and the complexity router's session and tier pins move off the router's DualCache onto RouterSessionPinStore, backed by the native redis-store. The read and claim Lua scripts run byte-identical to main, the pin read rides the request's pre-call batch, and keys and values keep main's format. --- litellm/caching/affinity_cache.py | 178 ++------ litellm/redis/stores/connection.py | 4 + .../stores/router_session_pins/__init__.py | 133 ++++++ .../stores/router_session_pins/handle.py | 37 ++ .../redis/stores/router_session_pins/keys.py | 37 ++ litellm/router_api.py | 9 + litellm/router_rust/core.py | 15 +- litellm/router_rust/steps/callbacks.py | 5 +- litellm/router_rust/steps/prerouting.py | 12 +- .../complexity_router/complexity_router.py | 49 +- .../deployment_affinity_check.py | 72 ++- litellm/rust/Cargo.lock | 1 + litellm/rust/crates/python-bridge/src/lib.rs | 3 + .../python-bridge/src/redis/store/mod.rs | 5 + .../src/redis/store/router_session_pins.rs | 234 ++++++++++ .../rust/crates/python-compat/src/number.rs | 59 ++- litellm/rust/crates/redis-store/Cargo.toml | 1 + litellm/rust/crates/redis-store/src/lib.rs | 1 + .../src/router_session_pins/claim.rs | 192 ++++++++ .../src/router_session_pins/keys.rs | 181 ++++++++ .../src/router_session_pins/lua.rs | 51 +++ .../src/router_session_pins/mod.rs | 232 ++++++++++ .../src/router_session_pins/types.rs | 9 + .../redis-store/tests/router_session_pins.rs | 368 +++++++++++++++ litellm/rust_bridge/_native.pyi | 36 +- .../integration/sdk/test_dual_cache_redis.py | 183 -------- .../sdk/test_router_session_pins_redis.py | 395 ++++++++++++++++ .../stores/router_session_pins_memory.py | 13 + .../redis/stores/test_router_session_pins.py | 116 +++++ .../router_strategy/test_complexity_router.py | 401 +++++------------ tests/unit/router_strategy/test_llm_v2.py | 4 +- .../test_deployment_affinity_check.py | 85 ++-- .../test_encrypted_content_affinity_check.py | 6 +- .../test_session_id_affinity.py | 422 ++++-------------- .../core/test_fallbacks_walk_parity.py | 2 +- ...election_characterization_session_repin.py | 7 +- tests/unit/test_router/test_router.py | 6 +- 37 files changed, 2482 insertions(+), 1082 deletions(-) create mode 100644 litellm/redis/stores/router_session_pins/__init__.py create mode 100644 litellm/redis/stores/router_session_pins/handle.py create mode 100644 litellm/redis/stores/router_session_pins/keys.py create mode 100644 litellm/rust/crates/python-bridge/src/redis/store/router_session_pins.rs create mode 100644 litellm/rust/crates/redis-store/src/router_session_pins/claim.rs create mode 100644 litellm/rust/crates/redis-store/src/router_session_pins/keys.rs create mode 100644 litellm/rust/crates/redis-store/src/router_session_pins/lua.rs create mode 100644 litellm/rust/crates/redis-store/src/router_session_pins/mod.rs create mode 100644 litellm/rust/crates/redis-store/src/router_session_pins/types.rs create mode 100644 litellm/rust/crates/redis-store/tests/router_session_pins.rs create mode 100644 tests/integration/sdk/test_router_session_pins_redis.py create mode 100644 tests/unit/redis/stores/router_session_pins_memory.py create mode 100644 tests/unit/redis/stores/test_router_session_pins.py diff --git a/litellm/caching/affinity_cache.py b/litellm/caching/affinity_cache.py index 2b8dc67da7a..0b05e248a98 100644 --- a/litellm/caching/affinity_cache.py +++ b/litellm/caching/affinity_cache.py @@ -1,151 +1,42 @@ """Atomic affinity claims shared by deployment and tier-model selection.""" -import json from collections.abc import Mapping -from typing import ( - Final, - cast, # noqa: TID251 # Heterogeneous cache results are narrowed only to object, then validated -) - -from pydantic import JsonValue, TypeAdapter, ValidationError +from typing import Final from litellm._logging import verbose_router_logger -from litellm.caching.dual_cache import DualCache -from litellm.caching.redis_batch import active_request_redis_batch +from litellm.caching.redis_batch import Batch +from litellm.redis.stores.router_session_pins import ROUTER_SESSION_PINS_TARGET, PinKey, RouterSessionPinStore -ROUTER_SESSION_PINS_TARGET: Final = "router_session_pins" - -_PIN_JSON_ADAPTER: Final = TypeAdapter[JsonValue](JsonValue) - -_READ_PIN_SCRIPT: Final = """ -local current = redis.call('GET', KEYS[1]) -if current ~= false then - redis.call('EXPIRE', KEYS[1], ARGV[1]) -end -return current -""" - -_CLAIM_PIN_SCRIPT: Final = """ -local current = redis.call('GET', KEYS[1]) -if current == false then - redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) - return ARGV[1] -end -if ARGV[3] then - local decoded, stored = pcall(cjson.decode, current) - if ARGV[4] then - local legacy = current - if decoded and type(stored) == 'string' then legacy = stored end - for _, eligible in ipairs(cjson.decode(ARGV[3])) do - if eligible['model_id'] == legacy then - redis.call('EXPIRE', KEYS[1], ARGV[2]) - return current - end - end - end - if decoded and type(stored) == 'table' then - for _, eligible in ipairs(cjson.decode(ARGV[3])) do - local matches = true - for key, value in pairs(eligible) do - if stored[key] ~= value then matches = false; break end - end - for key, _ in pairs(stored) do - if eligible[key] == nil then matches = false; break end - end - if matches then - redis.call('EXPIRE', KEYS[1], ARGV[2]) - return current - end - end - end - redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) - return ARGV[1] -end -if current == ARGV[1] then - redis.call('EXPIRE', KEYS[1], ARGV[2]) -end -return current -""" +__all__ = [ + "ROUTER_SESSION_PINS_TARGET", + "claim_affinity_pin", + "read_affinity_pin", + "set_local_affinity_pin", +] -def set_local_affinity_pin(cache: DualCache, cache_key: str, value: object, ttl_seconds: int) -> None: - """Replace the entry because InMemoryCache.set_cache preserves a live key's expiry.""" - cache.in_memory_cache.delete_cache(cache_key) - cache.in_memory_cache.set_cache(cache_key, value, ttl=ttl_seconds) +def set_local_affinity_pin(pins: RouterSessionPinStore, cache_key: PinKey, value: object, ttl_seconds: int) -> None: + pins.set_local(cache_key, value, ttl_seconds) -def _legacy_pin_matches(stored: object, pin_value: Mapping[str, str]) -> bool: - if isinstance(stored, dict): - return all(stored.get(key) is not None and str(stored[key]) == value for key, value in pin_value.items()) - return isinstance(stored, str) and len(pin_value) == 1 and stored in pin_value.values() - - -def claim_affinity_pin_in_memory( - cache: DualCache, - cache_key: str, - pin_value: Mapping[str, str], - ttl_seconds: int, - *, - eligible_values: tuple[Mapping[str, str], ...] | None = None, - preserve_legacy_model_ids: bool = False, -) -> object: - """No await between read and write, so same-loop claims agree during a Redis outage.""" - existing: Final[object] = cache.in_memory_cache.get_cache(cache_key) - if existing is not None and eligible_values is None: - if _legacy_pin_matches(existing, pin_value): - set_local_affinity_pin(cache, cache_key, pin_value, ttl_seconds) - return existing - eligible_existing: Final = existing in (eligible_values or ()) or ( - preserve_legacy_model_ids - and isinstance(existing, str) - and any("model_id" in value and _legacy_pin_matches(existing, value) for value in (eligible_values or ())) - ) - winner: Final = existing if existing is not None and eligible_existing else pin_value - set_local_affinity_pin(cache, cache_key, winner, ttl_seconds) - return winner - - -def _decode_pin(value: str) -> object: - try: - return _PIN_JSON_ADAPTER.validate_json(value) - except ValidationError: - return value - - -async def read_affinity_pin(cache: DualCache, cache_key: str, ttl_seconds: int) -> object: +async def read_affinity_pin(pins: RouterSessionPinStore, cache_key: PinKey, ttl_seconds: int) -> object: """Refresh the authoritative pin without writing a previously observed value back.""" - redis_cache: Final = cache.redis_cache - if redis_cache is not None: + if pins.has_redis: try: - read_script: Final = redis_cache.async_register_script(_READ_PIN_SCRIPT) - batch: Final = active_request_redis_batch(redis_cache) - raw: Final = cast( # cast-ok: Redis script returns a serialized pin or a miss; validate before use - object, - await batch.script(_READ_PIN_SCRIPT, read_script, (cache_key,), (int(ttl_seconds),)) - if batch is not None and batch.pending - else await read_script(keys=(cache_key,), args=(int(ttl_seconds),)), - ) - decoded: Final = raw.decode("utf-8") if isinstance(raw, bytes) else raw - value: Final = _decode_pin(decoded) if isinstance(decoded, str) else None - if value is None: - cache.in_memory_cache.delete_cache(cache_key) - else: - set_local_affinity_pin(cache, cache_key, value, ttl_seconds) + value: Final = await pins.read(cache_key, ttl_seconds, batch=Batch.PRE_CALL) except Exception as error: # noqa: BLE001 # Redis faults retain pod-local continuity, unlike a successful miss verbose_router_logger.debug("Affinity cache: Redis read failed, using pod-local pin. error=%s", error) else: return value - local_value: Final = cast( # cast-ok: the heterogeneous local cache is parsed by the affinity caller - object, cache.in_memory_cache.get_cache(cache_key) - ) + local_value: Final = pins.get_local(cache_key) if local_value is not None: - set_local_affinity_pin(cache, cache_key, local_value, ttl_seconds) + set_local_affinity_pin(pins, cache_key, local_value, ttl_seconds) return local_value async def claim_affinity_pin( - cache: DualCache, - cache_key: str, + pins: RouterSessionPinStore, + cache_key: PinKey, pin_value: Mapping[str, str], ttl_seconds: int, *, @@ -157,33 +48,20 @@ async def claim_affinity_pin( Eligible claims refresh the returned winner. Legacy deployment claims only refresh a matching candidate. Resolve Redis per call because the proxy attaches it lazily. """ - redis_cache: Final = cache.redis_cache - if redis_cache is not None: + if pins.has_redis: try: - claim_script: Final = redis_cache.async_register_script(_CLAIM_PIN_SCRIPT) - args: Final = ( - json.dumps(dict(pin_value)), - int(ttl_seconds), - *( - (json.dumps(tuple(dict(value) for value in eligible_values)),) - if eligible_values is not None - else () - ), - *(("1",) if preserve_legacy_model_ids and eligible_values is not None else ()), + winner: Final = await pins.claim( + cache_key, + pin_value, + ttl_seconds, + eligible_values=eligible_values, + preserve_legacy_model_ids=preserve_legacy_model_ids, ) - raw: Final = cast( # cast-ok: Redis scripts return heterogeneous values; only object is asserted here - object, await claim_script(keys=(cache_key,), args=args) - ) - decoded: Final = raw.decode("utf-8") if isinstance(raw, bytes) else raw - if not isinstance(decoded, str): - return pin_value - winner: Final = _decode_pin(decoded) - set_local_affinity_pin(cache, cache_key, winner, ttl_seconds) - return winner except Exception as error: # noqa: BLE001 # Redis/Lua faults retain same-pod affinity through local claims verbose_router_logger.debug("Affinity cache: Redis claim failed, using pod-local claim. error=%s", error) - return claim_affinity_pin_in_memory( - cache, + else: + return pin_value if winner is None else winner + return pins.claim_local( cache_key, pin_value, ttl_seconds, diff --git a/litellm/redis/stores/connection.py b/litellm/redis/stores/connection.py index e33ec80bba1..d2be3493e4d 100644 --- a/litellm/redis/stores/connection.py +++ b/litellm/redis/stores/connection.py @@ -9,6 +9,7 @@ from litellm.rust_bridge._native import ( NativeRateLimits, NativeRedisClient, NativeRouterCooldowns, + NativeRouterSessionPins, NativeRouterUsage, NativeSensitiveRoutePins, NativeSpendCounters, @@ -48,6 +49,9 @@ class RedisStores: def deployment_health(self) -> NativeDeploymentHealth: return NativeDeploymentHealth(self.client) + def router_session_pins(self) -> NativeRouterSessionPins: + return NativeRouterSessionPins(self.client) + def connect(redis_cache: RedisCache | None) -> RedisStores: settings: Final = None if redis_cache is None else redis_cache.store_settings diff --git a/litellm/redis/stores/router_session_pins/__init__.py b/litellm/redis/stores/router_session_pins/__init__.py new file mode 100644 index 00000000000..d8ea1dd8bff --- /dev/null +++ b/litellm/redis/stores/router_session_pins/__init__.py @@ -0,0 +1,133 @@ +from __future__ import annotations + +from collections.abc import Awaitable, Mapping, Sequence +from typing import TYPE_CHECKING, Final + +from litellm._internal_context import with_service_target +from litellm.caching.redis_batch import Batch +from litellm.redis.stores._route import ready, route +from litellm.redis.stores.router_session_pins.handle import RouterSessionPinsHandle +from litellm.redis.stores.router_session_pins.keys import ( + DeploymentPin, + DeploymentSessionPin, + PinKey, + TierModelPin, + TierSessionPin, +) + +if TYPE_CHECKING: + from litellm.redis.stores.connection import RedisStores + +__all__ = [ + "ROUTER_SESSION_PINS_TARGET", + "DeploymentPin", + "DeploymentSessionPin", + "PinKey", + "RouterSessionPinStore", + "TierModelPin", + "TierSessionPin", +] + +ROUTER_SESSION_PINS_TARGET: Final = "router_session_pins" + + +class RouterSessionPinStore: + """The router's session affinity pins (deployment affinity, the complexity router's session and tier pins): + memory per router, Redis shared by every pod. The store owns the claim rule, in Redis and in memory.""" + + def __init__(self, handle: RouterSessionPinsHandle) -> None: + self._handle: Final = handle + + @property + def has_redis(self) -> bool: + return self._handle.has_redis + + def attach_redis(self, redis: RedisStores) -> None: + """Adds the Redis tier once, keeping this router's memory pins, as ``DualCache.attach_redis_cache`` does.""" + if redis.client is not None: + self._handle.attach(redis.client) + + @with_service_target(ROUTER_SESSION_PINS_TARGET) + async def get(self, key: PinKey) -> object: + """``DualCache.async_get_cache``: memory, else Redis written back to memory.""" + if not self.has_redis: + return self._handle.get_local(key) + return await self._handle.get(key) + + def get_local(self, key: PinKey) -> object: + return self._handle.get_local(key) + + def set_local(self, key: PinKey, value: object, ttl: float) -> None: + """Replaces the memory entry, since ``InMemoryCache.set_cache`` keeps a live key's expiry.""" + self._handle.set_local(key, _plain(value), ttl) + + def delete_local(self, key: PinKey) -> None: + self._handle.delete_local(key) + + @with_service_target(ROUTER_SESSION_PINS_TARGET) + async def set(self, key: PinKey, value: object, ttl: float) -> None: + """``DualCache.async_set_cache``: memory, then a ``SET`` of ``json.dumps(value)``.""" + if not self.has_redis: + self.set_local(key, value, ttl) + return + await self._handle.set(key, _plain(value), ttl) + + @with_service_target(ROUTER_SESSION_PINS_TARGET) + async def get_shared(self, key: PinKey) -> object: + """``redis_cache.async_get_cache(raise_on_error=True)``: Redis alone, raising what it raises.""" + return await self._handle.get_shared(key) + + @with_service_target(ROUTER_SESSION_PINS_TARGET) + def read(self, key: PinKey, ttl: float, *, batch: Batch) -> Awaitable[object]: + """The read script on the batch ``batch`` picks: refreshes the pin's TTL, then memory takes what Redis holds + or drops a missing pin. A Redis failure raises, and the caller keeps its memory pin.""" + if not self.has_redis: + return ready(self._handle.get_local(key)) + target: Final = route(batch) + return target.result(self._handle.read(key, ttl, target.send)) + + @with_service_target(ROUTER_SESSION_PINS_TARGET) + async def claim( + self, + key: PinKey, + value: Mapping[str, str], + ttl: float, + *, + eligible_values: Sequence[Mapping[str, str]] | None, + preserve_legacy_model_ids: bool, + ) -> object: + """The claim script: the first writer keeps the key unless it is no longer eligible, and memory takes the + winner. None when Redis answered no pin. A Redis failure raises, and the caller settles with + ``claim_local``.""" + return await self._handle.claim( + key, dict(value), ttl, *_eligibility(eligible_values, preserve_legacy_model_ids) + ) + + def claim_local( + self, + key: PinKey, + value: Mapping[str, str], + ttl: float, + *, + eligible_values: Sequence[Mapping[str, str]] | None, + preserve_legacy_model_ids: bool, + ) -> object: + """``claim`` against this router's memory alone, for a router without Redis or one whose Redis claim + failed. Nothing is awaited between its read and its write, so claims on one loop agree.""" + return self._handle.claim_local( + key, dict(value), ttl, *_eligibility(eligible_values, preserve_legacy_model_ids) + ) + + +def _eligibility( + eligible_values: Sequence[Mapping[str, str]] | None, preserve_legacy_model_ids: bool +) -> tuple[Sequence[object] | None, bool]: + return ( + None if eligible_values is None else [dict(eligible) for eligible in eligible_values], + preserve_legacy_model_ids and eligible_values is not None, + ) + + +def _plain(value: object) -> object: + """A pin as the plain value the native store takes: a mapping becomes the dict it holds.""" + return dict(value) if isinstance(value, Mapping) else value diff --git a/litellm/redis/stores/router_session_pins/handle.py b/litellm/redis/stores/router_session_pins/handle.py new file mode 100644 index 00000000000..5ec2480aea7 --- /dev/null +++ b/litellm/redis/stores/router_session_pins/handle.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from collections.abc import Awaitable, Sequence +from typing import Protocol + +from litellm.caching.redis_batch import NativePending +from litellm.redis.stores.router_session_pins.keys import PinKey +from litellm.rust_bridge._native import NativeRedisClient, NativeSend + + +class RouterSessionPinsHandle(Protocol): + @property + def has_redis(self) -> bool: ... + def attach(self, client: NativeRedisClient, /) -> bool: ... + def get(self, key: PinKey) -> Awaitable[object]: ... + def get_local(self, key: PinKey) -> object: ... + def set_local(self, key: PinKey, value: object, ttl_seconds: float) -> None: ... + def delete_local(self, key: PinKey) -> None: ... + def set(self, key: PinKey, value: object, ttl_seconds: float) -> Awaitable[None]: ... + def get_shared(self, key: PinKey) -> Awaitable[object]: ... + def read(self, key: PinKey, ttl_seconds: float, send: NativeSend) -> NativePending[object]: ... + def claim( + self, + key: PinKey, + value: object, + ttl_seconds: float, + eligible_values: Sequence[object] | None, + preserve_legacy_model_ids: bool, + ) -> Awaitable[object]: ... + def claim_local( + self, + key: PinKey, + value: object, + ttl_seconds: float, + eligible_values: Sequence[object] | None, + preserve_legacy_model_ids: bool, + ) -> object: ... diff --git a/litellm/redis/stores/router_session_pins/keys.py b/litellm/redis/stores/router_session_pins/keys.py new file mode 100644 index 00000000000..00d17b0e8c7 --- /dev/null +++ b/litellm/redis/stores/router_session_pins/keys.py @@ -0,0 +1,37 @@ +from __future__ import annotations + +from dataclasses import dataclass +from typing import TypeAlias + + +@dataclass(frozen=True, slots=True) +class DeploymentPin: + model_group: str + caller: str + + +@dataclass(frozen=True, slots=True) +class DeploymentSessionPin: + model_group: str + caller: str | None + session_id: str + + +@dataclass(frozen=True, slots=True) +class TierSessionPin: + router: str + caller: str | None + session_id: str + + +@dataclass(frozen=True, slots=True) +class TierModelPin: + router: str + tiers: str + caller: str | None + session_id: str + tier: str + replay_fallback: str | None + + +PinKey: TypeAlias = DeploymentPin | DeploymentSessionPin | TierSessionPin | TierModelPin diff --git a/litellm/router_api.py b/litellm/router_api.py index 38ed81570c9..e46337c37ee 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -58,6 +58,7 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.messages.dispatch import MessagesResult + from litellm.redis.stores.router_session_pins import RouterSessionPinStore from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.router import Router from litellm.router_rust.surface.model_list import HostedModelList @@ -1548,6 +1549,9 @@ class _HealthHost(Protocol): @property def _core_cooldown_cache(self) -> "CooldownCache": ... + @property + def _core_session_pins(self) -> "RouterSessionPinStore": ... + @property def _core_health_state_cache(self) -> "DeploymentHealthCache": ... @@ -1601,6 +1605,11 @@ class RouterHealth: """The cooldown store (was ``cooldown_cache``).""" return self._router._core_cooldown_cache + @property + def session_pins(self) -> "RouterSessionPinStore": + """The session affinity pins deployment affinity and the complexity router claim and read.""" + return self._router._core_session_pins + @property def health_checks(self) -> "DeploymentHealthCache": """Background health check results (was ``health_state_cache``).""" diff --git a/litellm/router_rust/core.py b/litellm/router_rust/core.py index 810bb4d3230..f419812947e 100644 --- a/litellm/router_rust/core.py +++ b/litellm/router_rust/core.py @@ -78,6 +78,7 @@ from litellm.litellm_core_utils.total_deadline import ( from litellm.redis.stores.connection import connect from litellm.redis.stores.deployment_health import DeploymentHealthStore from litellm.redis.stores.router_cooldowns import RouterCooldownStore +from litellm.redis.stores.router_session_pins import RouterSessionPinStore from litellm.redis.stores.router_usage import RouterUsage from litellm.router_api import ( ROUTER_INIT_ARGS, @@ -1165,6 +1166,10 @@ class RustHealth(RouterHealth): def cooldowns(self) -> "CooldownCache": return self._core._cooldowns # pyright: ignore[reportPrivateUsage] # see above + @property + def session_pins(self) -> RouterSessionPinStore: + return self._core._session_pins # pyright: ignore[reportPrivateUsage] # see above + def failures_this_minute(self, deployment_id: str) -> int: count: Final[object] = self.store.get_cache(f"{deployment_id}:fails", local_only=True) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped return int(count) if isinstance(count, int | float) else 0 @@ -1268,6 +1273,7 @@ _REBOUND_STATE: Final = MappingProxyType( { "_core_cache": "_cache", "_core_cooldown_cache": "_cooldowns", + "_core_session_pins": "_session_pins", "_core_health_state_cache": "_health_checks", "_core_slack_alerting_logger": "_slack_alerting_logger", } @@ -1278,6 +1284,7 @@ _HEALTH_SERVED: Final = frozenset( { "store", "cooldowns", + "session_pins", "failures_this_minute", "record_failure", "usage", @@ -1494,6 +1501,7 @@ class RustCore( self._health_checks: DeploymentHealthCache = resolved.health_routing.health_checks( DeploymentHealthStore(redis.deployment_health()) ) + self._session_pins: Final = RouterSessionPinStore(redis.router_session_pins()) self._version: int = 1 self._settings: RustSettings | None = None self._tag_filtering_models: tuple[int, bool] | None = None @@ -1522,7 +1530,7 @@ class RustCore( ctx: Final = RoutingContext(self) # pyright: ignore[reportArgumentType] # the strategies read registry(), settings(), health() and select() self._ctx: Final = ctx strategies: Final = PreRoutingStrategies( - self._hosted, host, ctx, config.get("auto_router_capability_limit"), self._cache + self._hosted, host, ctx, config.get("auto_router_capability_limit"), self._session_pins ) self._hosted.strategies = strategies self._prerouter: Final = PreRouter(strategies, self._hosted, host, ctx, self._cache) @@ -1548,7 +1556,9 @@ class RustCore( self._hosted._core_routing_plugins = list(resolved.routing_plugins) # pyright: ignore[reportPrivateUsage] # the core fills its own host registry self._rebuild() self._join_callback_holders() - self._checks: Final = OptionalChecks(self, host, self._cache, self._usage, self._hosted, ctx, config.arguments) + self._checks: Final = OptionalChecks( + self, host, self._cache, self._usage, self._session_pins, self._hosted, ctx, config.arguments + ) self._hosted.checks = self._checks strategies.ensure_affinity = self._checks.ensure_deployment_affinity_callback strategies.hand_over_affinity(self._checks.adopt) @@ -2135,6 +2145,7 @@ class RustCore( self._usage.attach_redis(redis) self._cooldowns.cooldown_store.attach_redis(redis) self._health_checks.store.attach_redis(redis) + self._session_pins.attach_redis(redis) self._checks.prompt_caching_pin_cache.attach_redis_cache(changes.redis_cache) self._prerouter.attach_redis_cache(changes.redis_cache) return ApplyResult( diff --git a/litellm/router_rust/steps/callbacks.py b/litellm/router_rust/steps/callbacks.py index a8a98fd44b3..e924193db78 100644 --- a/litellm/router_rust/steps/callbacks.py +++ b/litellm/router_rust/steps/callbacks.py @@ -51,6 +51,7 @@ from litellm.router_utils.pre_call_checks.prompt_caching_deployment_check import if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.redis.stores.router_session_pins import RouterSessionPinStore from litellm.redis.stores.router_usage import RouterUsage from litellm.router_python_host import PythonHost from litellm.router_rust.steps.selection import StrategySelectors @@ -111,6 +112,7 @@ class OptionalChecks: host: "PythonHost", cache: DualCache, usage: "RouterUsage", + pins: "RouterSessionPinStore", hosted: "HostRegistry", ctx: object, arguments: Mapping[str, object], @@ -119,6 +121,7 @@ class OptionalChecks: self._host: Final = host self._cache: Final = cache self._usage: Final = usage + self._pins: Final = pins self._hosted: Final = hosted self._ctx: Final = ctx self.optional_callbacks: list[object] | None = None # mutable-ok: the router's live list, as Python keeps it @@ -212,7 +215,7 @@ class OptionalChecks: def _affinity_check(self, user_key: bool, responses_api: bool, session_id: bool) -> DeploymentAffinityCheck: return DeploymentAffinityCheck( - cache=self._cache, + pins=self._pins, ttl_seconds=self.deployment_affinity_ttl_seconds, # pyright: ignore[reportArgumentType] # the argument as passed enable_user_key_affinity=user_key, enable_responses_api_affinity=responses_api, diff --git a/litellm/router_rust/steps/prerouting.py b/litellm/router_rust/steps/prerouting.py index 2953a87ed92..44a4800af8d 100644 --- a/litellm/router_rust/steps/prerouting.py +++ b/litellm/router_rust/steps/prerouting.py @@ -66,6 +66,7 @@ from litellm.utils import provider_rejectable_params if TYPE_CHECKING: from litellm.caching.redis_cache import RedisCache + from litellm.redis.stores.router_session_pins import RouterSessionPinStore from litellm.router_python_host import PythonHost from litellm.router_rust.surface.registry import HostRegistry from litellm.router_strategy.auto_router.auto_router import PreRoutingHookResponse @@ -172,13 +173,18 @@ class PreRoutingStrategies: reads ``_core_auto_routers`` and its siblings there).""" def __init__( - self, hosted: "HostRegistry", host: "PythonHost", ctx: object, capability_limit: object, cache: DualCache + self, + hosted: "HostRegistry", + host: "PythonHost", + ctx: object, + capability_limit: object, + pins: "RouterSessionPinStore", ) -> None: self._hosted: Final = hosted self._host: Final = host self._ctx: Final = ctx self._capability_limit: object = capability_limit - self._cache: Final = cache + self._pins: Final = pins self._affinity: object = None # The router's own `_ensure_deployment_affinity_callback` (its optional checks), when set. self.ensure_affinity: Callable[[], None] | None = None @@ -454,7 +460,7 @@ class PreRoutingStrategies: return group is not None and group.routing_strategy == "priority" affinity: Final = DeploymentAffinityCheck( - cache=self._cache, + pins=self._pins, ttl_seconds=3600, enable_user_key_affinity=False, enable_responses_api_affinity=False, diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 2ce1bda173a..79dfd64f3e8 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -65,6 +65,7 @@ from litellm.llms.anthropic.common_utils import is_claude_code_user_agent from litellm.llms.base_llm.base_utils import type_to_response_format_param from litellm.llms.custom_httpx.http_handler import get_async_httpx_client from litellm.llms.databricks.decisions.transformation import DATABRICKS_DECISIONS_CONFIG +from litellm.redis.stores.router_session_pins import TierModelPin, TierSessionPin from litellm.router_deprecated import RenamedRouterAttribute, renamed_router_param from litellm.router_strategy.adaptive_router.classifier import classify_prompt from litellm.router_strategy.complexity_router.context_compaction import compaction_pending @@ -3182,17 +3183,16 @@ class ComplexityRouter(CustomLogger): session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs) if session_id is None: return model - caller: Final = DeploymentAffinityCheck.get_user_key_from_request_kwargs(request_kwargs) - identity: Final = (self.model_name, self._tier_affinity_config, caller, session_id, _tier_name(tier)) - cache_identity: Final = ( - (*identity, ("replay_fallback", retained_model)) if retained_model is not None else identity - ) - cache_key: Final = ( - "complexity_router_tier_model_affinity:v1:" - + hashlib.sha256(json.dumps(cache_identity).encode()).hexdigest() + cache_key: Final = TierModelPin( + router=self.model_name, + tiers=self._tier_affinity_config, + caller=DeploymentAffinityCheck.get_user_key_from_request_kwargs(request_kwargs), + session_id=session_id, + tier=_tier_name(tier), + replay_fallback=retained_model, ) winner: Final = await claim_affinity_pin( - self.ctx.health().store, + self.ctx.health().session_pins, cache_key, MappingProxyType({"model": model}), self.config.session_affinity_ttl_seconds, @@ -4506,13 +4506,12 @@ class ComplexityRouter(CustomLogger): return str(user_key) return None - def _get_session_affinity_cache_key(self, session_id: str, request_kwargs: dict) -> str: - # Namespace by the caller's API key hash so two different callers reusing the - # same client-supplied session_id can't poison each other's routing pin. Falls - # back to "unscoped" only when there's no authenticated caller to scope by - # (e.g. direct Router usage without the proxy layer). - caller_scope: Final = self._get_user_api_key_hash_from_request_kwargs(request_kwargs) or "unscoped" - return f"complexity_router_session_affinity:v1:{self.model_name}:{caller_scope}:{session_id}" + def _get_session_affinity_cache_key(self, session_id: str, request_kwargs: dict) -> TierSessionPin: + return TierSessionPin( + router=self.model_name, + caller=self._get_user_api_key_hash_from_request_kwargs(request_kwargs), + session_id=session_id, + ) @property def _uses_tier_pin(self) -> bool: @@ -4594,7 +4593,7 @@ class ComplexityRouter(CustomLogger): if cache_key is not None and pin_replay_allowed: pinned_value: Final = await read_affinity_pin( - self.ctx.health().store, cache_key, self.config.session_affinity_ttl_seconds + self.ctx.health().session_pins, cache_key, self.config.session_affinity_ttl_seconds ) pinned_pin: Final = _parse_session_affinity_pin(pinned_value, self.config.tier_names()) if pinned_pin is not None: @@ -4693,10 +4692,10 @@ class ComplexityRouter(CustomLogger): else floor_model ) if escalated: - await self.ctx.health().store.async_set_cache( - key=cache_key, - value=_session_affinity_cache_value(session_model, resolved_pin_tier), - ttl=self.config.session_affinity_ttl_seconds, + await self.ctx.health().session_pins.set( + cache_key, + _session_affinity_cache_value(session_model, resolved_pin_tier), + self.config.session_affinity_ttl_seconds, ) if self.config.adaptive: from litellm.router_strategy.adaptive_router.config import ( @@ -4802,13 +4801,13 @@ class ComplexityRouter(CustomLogger): and self._matched_plan_mode_signal(request_kwargs, resolved_messages) is None ) if pinnable and cache_key is not None and response is not None: - await self.ctx.health().store.async_set_cache( - key=cache_key, - value=_session_affinity_cache_value( + await self.ctx.health().session_pins.set( + cache_key, + _session_affinity_cache_value( response.model, response.routing_decision.get("tier") if response.routing_decision is not None else None, ), - ttl=self.config.session_affinity_ttl_seconds, + self.config.session_affinity_ttl_seconds, ) return self._with_session_deployment_affinity(response) diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index dbf471d875c..19d80f0c0db 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -24,13 +24,17 @@ from litellm._logging import verbose_router_logger from litellm.caching.affinity_cache import ( ROUTER_SESSION_PINS_TARGET, claim_affinity_pin, - claim_affinity_pin_in_memory, set_local_affinity_pin, ) -from litellm.caching.dual_cache import DualCache from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger, Span from litellm.litellm_core_utils.routing_common_utils import deployment_ids_failed_in_request +from litellm.redis.stores.router_session_pins import ( + DeploymentPin, + DeploymentSessionPin, + PinKey, + RouterSessionPinStore, +) from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes @@ -97,12 +101,11 @@ class DeploymentAffinityCheck(CustomLogger): `Router(optional_pre_call_checks=[...])`. """ - CACHE_KEY_PREFIX = "deployment_affinity:v1" USER_ID_AFFINITY_PREFIX: Final = "user_id:" def __init__( self, - cache: DualCache, + pins: RouterSessionPinStore, ttl_seconds: int, enable_user_key_affinity: bool, enable_responses_api_affinity: bool, @@ -112,7 +115,7 @@ class DeploymentAffinityCheck(CustomLogger): session_affinity_repin_on_failover: bool = False, ): super().__init__() - self.cache = cache + self.pins: Final = pins self.ttl_seconds = ttl_seconds self.enable_user_key_affinity = enable_user_key_affinity self.enable_responses_api_affinity = enable_responses_api_affinity @@ -261,18 +264,16 @@ class DeploymentAffinityCheck(CustomLogger): return f"{value[:keep]}..." @classmethod - def get_affinity_cache_key(cls, model_group: str, user_key: str) -> str: - hashed_user_key: Final = cls._hash_user_key(user_key=user_key) - return f"{cls.CACHE_KEY_PREFIX}:{model_group}:{hashed_user_key}" + def get_affinity_cache_key(cls, model_group: str, user_key: str) -> DeploymentPin: + return DeploymentPin(model_group=model_group, caller=user_key) @classmethod - def get_session_affinity_cache_key(cls, model_group: str, session_id: str, user_key: str | None) -> str: - """Session pins are scoped by the caller's hashed API key so two callers reusing - the same client-supplied session_id cannot read or steer each other's pin. - `"unscoped"` covers direct Router usage with no authenticated caller, matching - the complexity router's own session pin key.""" - hashed_user_key: Final = cls._hash_user_key(user_key) if user_key is not None else "unscoped" - return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{hashed_user_key}:{session_id}" + def get_session_affinity_cache_key( + cls, model_group: str, session_id: str, user_key: str | None + ) -> DeploymentSessionPin: + """Session pins are scoped by the caller so two callers reusing the same client-supplied session_id cannot + read or steer each other's pin.""" + return DeploymentSessionPin(model_group=model_group, caller=user_key, session_id=session_id) @staticmethod def _get_session_id_from_metadata_dict(metadata: Mapping[object, object]) -> str | None: @@ -352,17 +353,13 @@ class DeploymentAffinityCheck(CustomLogger): return stored return None - def _set_local_pin(self, cache_key: str, value: object, ttl_seconds: int) -> None: - set_local_affinity_pin(self.cache, cache_key, value, ttl_seconds) + def _set_local_pin(self, cache_key: PinKey, value: object, ttl_seconds: int) -> None: + set_local_affinity_pin(self.pins, cache_key, value, ttl_seconds) - async def _claim_pin(self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int) -> str | None: - winner: Final = await claim_affinity_pin(self.cache, cache_key, pin_value, ttl_seconds) - return self._pinned_model_id(winner) - - def _claim_pin_in_memory( - self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int + async def _claim_pin( + self, cache_key: PinKey, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int ) -> str | None: - winner: Final = claim_affinity_pin_in_memory(self.cache, cache_key, pin_value, ttl_seconds) + winner: Final = await claim_affinity_pin(self.pins, cache_key, pin_value, ttl_seconds) return self._pinned_model_id(winner) @staticmethod @@ -453,7 +450,7 @@ class DeploymentAffinityCheck(CustomLogger): parent_otel_span, ) if self.session_affinity_repin_on_failover - else await self.cache.async_get_cache(key=session_cache_key) + else await self.pins.get(session_cache_key) ) session_model_id: str | None = None @@ -488,7 +485,7 @@ class DeploymentAffinityCheck(CustomLogger): return typed_healthy_deployments cache_key: Final = self.get_affinity_cache_key(model_group=stable_model_map_key, user_key=user_key) - cache_result: Final = await self.cache.async_get_cache(key=cache_key) + cache_result: Final = await self.pins.get(cache_key) model_id: str | None = None if isinstance(cache_result, dict): @@ -518,21 +515,20 @@ class DeploymentAffinityCheck(CustomLogger): ) return [deployment] - async def _read_shared_session_pin(self, cache_key: str, ttl_seconds: int, parent_otel_span: Span | None) -> object: - redis_cache: Final = self.cache.redis_cache - if redis_cache is None: - return TypeAdapter(object).validate_python(await self.cache.async_get_cache(key=cache_key)) - local_before: Final = TypeAdapter(object).validate_python(self.cache.in_memory_cache.get_cache(cache_key)) + async def _read_shared_session_pin( + self, cache_key: PinKey, ttl_seconds: int, parent_otel_span: Span | None + ) -> object: + if not self.pins.has_redis: + return await self.pins.get(cache_key) + local_before: Final = self.pins.get_local(cache_key) try: - shared: Final[object] = await redis_cache.async_get_cache( - key=cache_key, parent_otel_span=parent_otel_span, raise_on_error=True - ) + shared: Final = await self.pins.get_shared(cache_key) except Exception: # noqa: BLE001 # Optional affinity falls back locally on Redis admission or read failure verbose_router_logger.debug("Session pin read failed, using pod-local affinity") - return TypeAdapter(object).validate_python(await self.cache.async_get_cache(key=cache_key, local_only=True)) - if self.cache.in_memory_cache.get_cache(cache_key) is local_before: + return self.pins.get_local(cache_key) + if self.pins.get_local(cache_key) == local_before: if shared is None: - self.cache.in_memory_cache.delete_cache(cache_key) + self.pins.delete_local(cache_key) else: self._set_local_pin(cache_key, shared, ttl_seconds) return shared @@ -575,7 +571,7 @@ class DeploymentAffinityCheck(CustomLogger): model_group, session_id, self.get_user_key_from_request_kwargs(request_kwargs) ) await claim_affinity_pin( - self.cache, + self.pins, cache_key, DeploymentAffinityCacheValue(model_id=str(model_id)), marker_ttl if marker_ttl is not None else self.ttl_seconds, diff --git a/litellm/rust/Cargo.lock b/litellm/rust/Cargo.lock index d3cbb960521..ea790ecfb61 100644 --- a/litellm/rust/Cargo.lock +++ b/litellm/rust/Cargo.lock @@ -4553,6 +4553,7 @@ dependencies = [ "litellm-secret-value", "litellm-tracing", "rstest", + "sha2 0.10.9", "strum", "thiserror 2.0.19", "tokio", diff --git a/litellm/rust/crates/python-bridge/src/lib.rs b/litellm/rust/crates/python-bridge/src/lib.rs index 12795146ca9..a7acb35abb4 100644 --- a/litellm/rust/crates/python-bridge/src/lib.rs +++ b/litellm/rust/crates/python-bridge/src/lib.rs @@ -62,6 +62,8 @@ mod _native { #[pymodule_export] use crate::redis::store::request::{NativeRequestRedis, NativeSend}; #[pymodule_export] + use crate::redis::store::router_session_pins::NativeRouterSessionPins; + #[pymodule_export] use crate::redis::store::router_usage::{NativeRouterUsage, NativeUsageRead}; #[pymodule_export] use crate::redis::store::sensitive_route_pins::NativeSensitiveRoutePins; @@ -156,6 +158,7 @@ mod tests { "NativeSend", "NativeAuthCache", "NativeDeploymentHealth", + "NativeRouterSessionPins", "NativeRateLimits", "WindowCount", "AtomicApplied", diff --git a/litellm/rust/crates/python-bridge/src/redis/store/mod.rs b/litellm/rust/crates/python-bridge/src/redis/store/mod.rs index 2700ae09593..c788c92caa7 100644 --- a/litellm/rust/crates/python-bridge/src/redis/store/mod.rs +++ b/litellm/rust/crates/python-bridge/src/redis/store/mod.rs @@ -4,6 +4,7 @@ mod policy; pub(crate) mod rate_limits; pub(crate) mod request; pub(crate) mod router_cooldowns; +pub(crate) mod router_session_pins; pub(crate) mod router_usage; pub(crate) mod sensitive_route_pins; pub(crate) mod spend_counters; @@ -236,6 +237,10 @@ impl<'py, T: IntoValue> IntoPyObject<'py> for PyData { } } +fn unwrap_values(values: Vec) -> Vec { + values.into_iter().map(|value| value.0).collect() +} + fn unwrap_all(values: Vec>) -> Vec { values.into_iter().map(|value| value.0).collect() } diff --git a/litellm/rust/crates/python-bridge/src/redis/store/router_session_pins.rs b/litellm/rust/crates/python-bridge/src/redis/store/router_session_pins.rs new file mode 100644 index 00000000000..59dbcf853cb --- /dev/null +++ b/litellm/rust/crates/python-bridge/src/redis/store/router_session_pins.rs @@ -0,0 +1,234 @@ +use std::{sync::Arc, time::Duration}; + +use litellm_redis_store::{ + Clock, + router_session_pins::{Eligible, PinKey, RouterSessionPins}, + system_clock, +}; +use pyo3::prelude::*; + +use super::{ + NativePending, duration, kind_error, pending, policy::router_tier_policies, redis_async, + remote, request::NativeSend, unwrap_values, value::PyValue, +}; +use crate::redis::client::NativeRedisClient; + +fn pin(value: Option) -> Option { + value.map(PyValue) +} + +/// The `PinKey` dataclasses in `litellm.redis.stores.router_session_pins.keys`. A key whose fields +/// are a subset of another's comes after it, so extraction never stops at the smaller one. +#[derive(FromPyObject)] +enum PyPinKey { + TierModel { + router: String, + tiers: String, + caller: Option, + session_id: String, + tier: String, + replay_fallback: Option, + }, + TierSession { + router: String, + caller: Option, + session_id: String, + }, + DeploymentSession { + model_group: String, + caller: Option, + session_id: String, + }, + Deployment { + model_group: String, + caller: String, + }, +} + +impl From for PinKey { + fn from(key: PyPinKey) -> Self { + match key { + PyPinKey::TierModel { + router, + tiers, + caller, + session_id, + tier, + replay_fallback, + } => Self::TierModel { + router, + tiers, + caller, + session_id, + tier, + replay_fallback, + }, + PyPinKey::TierSession { + router, + caller, + session_id, + } => Self::TierSession { + router, + caller, + session_id, + }, + PyPinKey::DeploymentSession { + model_group, + caller, + session_id, + } => Self::DeploymentSession { + model_group, + caller, + session_id, + }, + PyPinKey::Deployment { + model_group, + caller, + } => Self::Deployment { + model_group, + caller, + }, + } + } +} + +/// `InMemoryCache(clock=...)`: a callable returning epoch seconds, which tests use to move time. +fn python_clock(clock: Py) -> Clock { + Arc::new(move || { + Python::attach(|py| clock.call0(py)?.extract::(py)) + .ok() + .and_then(|seconds| Duration::try_from_secs_f64(seconds).ok()) + .unwrap_or_default() + }) +} + +/// The router's session affinity pins, behind `RouterSessionPinStore`. +#[pyclass(frozen, name = "NativeRouterSessionPins")] +pub(crate) struct NativeRouterSessionPins(Arc); + +#[pymethods] +impl NativeRouterSessionPins { + #[new] + #[pyo3(signature = (client, clock=None))] + fn new( + py: Python<'_>, + client: Option>, + clock: Option>, + ) -> PyResult { + Ok(Self(Arc::new(RouterSessionPins::with_clock( + client.map(|client| remote(&client)), + router_tier_policies(py)?.default_in_memory_ttl, + clock.map_or_else(system_clock, python_clock), + )))) + } + + /// Adds the Redis tier and keeps memory; `false` when one was already attached. + fn attach(&self, client: PyRef<'_, NativeRedisClient>) -> bool { + self.0.attach(remote(&client)) + } + + #[getter] + fn has_redis(&self) -> bool { + self.0.has_redis() + } + + fn get<'py>(&self, py: Python<'py>, key: PyPinKey) -> PyResult> { + let store = self.0.clone(); + let key = PinKey::from(key); + redis_async(py, "async_get_cache", kind_error, async move { + store.get(&key).await.map(pin) + }) + } + + fn get_local(&self, key: PyPinKey) -> PyResult> { + self.0.get_local(&key.into()).map(pin).map_err(kind_error) + } + + fn set_local(&self, key: PyPinKey, value: PyValue, ttl_seconds: f64) -> PyResult<()> { + self.0 + .set_local(&key.into(), value.0, duration(ttl_seconds)?) + .map_err(kind_error) + } + + fn delete_local(&self, key: PyPinKey) -> PyResult<()> { + self.0.delete_local(&key.into()).map_err(kind_error) + } + + fn set<'py>( + &self, + py: Python<'py>, + key: PyPinKey, + value: PyValue, + ttl_seconds: f64, + ) -> PyResult> { + let store = self.0.clone(); + let key = PinKey::from(key); + let ttl = duration(ttl_seconds)?; + redis_async(py, "async_set_cache", kind_error, async move { + store.set(&key, value.0, ttl).await + }) + } + + fn get_shared<'py>(&self, py: Python<'py>, key: PyPinKey) -> PyResult> { + let store = self.0.clone(); + let key = PinKey::from(key); + redis_async(py, "async_get_cache", kind_error, async move { + store.get_shared(&key).await.map(pin) + }) + } + + /// The read script, declared on the batch the caller's `NativeSend` picked. + fn read(&self, key: PyPinKey, ttl_seconds: f64, send: &NativeSend) -> PyResult { + let declared = self + .0 + .read(send.send(), &key.into(), duration(ttl_seconds)?) + .map_err(kind_error)?; + Ok(pending("async_register_script", kind_error, declared, pin)) + } + + #[pyo3(signature = (key, value, ttl_seconds, eligible_values, preserve_legacy_model_ids))] + fn claim<'py>( + &self, + py: Python<'py>, + key: PyPinKey, + value: PyValue, + ttl_seconds: f64, + eligible_values: Option>, + preserve_legacy_model_ids: bool, + ) -> PyResult> { + let store = self.0.clone(); + let key = PinKey::from(key); + let ttl = duration(ttl_seconds)?; + let eligible = eligible(eligible_values, preserve_legacy_model_ids); + redis_async(py, "async_register_script", kind_error, async move { + store.claim(&key, &value.0, ttl, &eligible).await.map(pin) + }) + } + + #[pyo3(signature = (key, value, ttl_seconds, eligible_values, preserve_legacy_model_ids))] + fn claim_local( + &self, + key: PyPinKey, + value: PyValue, + ttl_seconds: f64, + eligible_values: Option>, + preserve_legacy_model_ids: bool, + ) -> PyResult { + self.0 + .claim_local( + &key.into(), + &value.0, + duration(ttl_seconds)?, + &eligible(eligible_values, preserve_legacy_model_ids), + ) + .map(PyValue) + .map_err(kind_error) + } +} + +fn eligible(values: Option>, preserve_legacy_model_ids: bool) -> Eligible { + Eligible { + values: values.map(unwrap_values), + preserve_legacy_model_ids, + } +} diff --git a/litellm/rust/crates/python-compat/src/number.rs b/litellm/rust/crates/python-compat/src/number.rs index 6e2f067ae25..c84e8dd613f 100644 --- a/litellm/rust/crates/python-compat/src/number.rs +++ b/litellm/rust/crates/python-compat/src/number.rs @@ -7,6 +7,36 @@ pub fn int(text: &str) -> Option { parse(text) } +/// `int(text, 16)`, or `None` where Python raises `ValueError`. +pub fn hex_int(text: &str) -> Option { + let text = text.trim(); + let (negative, unsigned) = match text.strip_prefix('-') { + Some(rest) => (true, rest), + None => (false, text.strip_prefix('+').unwrap_or(text)), + }; + let digits = match unsigned + .strip_prefix("0x") + .or_else(|| unsigned.strip_prefix("0X")) + { + Some(rest) => rest.strip_prefix('_').unwrap_or(rest), + None => unsigned, + }; + let hex_at = |index: Option| { + index + .and_then(|index| digits.as_bytes().get(index)) + .is_some_and(u8::is_ascii_hexdigit) + }; + let well_formed = digits.char_indices().all(|(index, c)| { + c.is_ascii_hexdigit() + || (c == '_' && hex_at(index.checked_sub(1)) && hex_at(Some(index + 1))) + }); + if digits.is_empty() || !well_formed { + return None; + } + let magnitude = BigInt::parse_bytes(digits.replace('_', "").as_bytes(), 16)?; + Some(if negative { -magnitude } else { magnitude }) +} + /// `float(text)`, or `None` where Python raises `ValueError`. pub fn float(text: &str) -> Option { parse(text) @@ -55,7 +85,7 @@ mod tests { use std::time::Duration; - use super::{SecondsError, float, int, seconds}; + use super::{SecondsError, float, hex_int, int, seconds}; // Expected values are CPython 3.14 int(text) #[rstest] @@ -78,6 +108,33 @@ mod tests { assert_eq!(int(text), expected.map(BigInt::from)); } + // Expected values are CPython 3.14 int(text, 16) + #[rstest] + #[case::lower("ff", Some(255))] + #[case::upper("FF", Some(255))] + #[case::surrounding_whitespace(" ff\n", Some(255))] + #[case::unicode_spaces("\u{a0}ff\u{2003}", Some(255))] + #[case::plus_sign("+ff", Some(255))] + #[case::minus_sign("-ff", Some(-255))] + #[case::prefix("0xff", Some(255))] + #[case::upper_prefix("0XFF", Some(255))] + #[case::underscore_after_prefix("0x_ff", Some(255))] + #[case::underscore_between_digits("f_f", Some(255))] + #[case::binary_looking("0b1", Some(177))] + #[case::double_underscore("f__f", None)] + #[case::leading_underscore("_ff", None)] + #[case::trailing_underscore("ff_", None)] + #[case::bare_prefix("0x", None)] + #[case::prefix_and_underscore("0x_", None)] + #[case::empty("", None)] + #[case::not_hex("g", None)] + #[case::information_separator("\x1cff", None)] + #[case::decimal_point("1.5", None)] + #[case::non_ascii_digit_is_a_known_limit("\u{667}", None)] + fn hex_int_parses_like_python(#[case] text: &str, #[case] expected: Option) { + assert_eq!(hex_int(text), expected.map(BigInt::from)); + } + // Expected values are CPython 3.14 float(text) #[rstest] #[case::underscore_and_newline(" 1_0.5\n", Some(10.5))] diff --git a/litellm/rust/crates/redis-store/Cargo.toml b/litellm/rust/crates/redis-store/Cargo.toml index 1dd61d76d3c..cba7e992834 100644 --- a/litellm/rust/crates/redis-store/Cargo.toml +++ b/litellm/rust/crates/redis-store/Cargo.toml @@ -11,6 +11,7 @@ litellm-cache.workspace = true litellm-cache-memory.workspace = true litellm-python-compat.workspace = true litellm-redis-client.workspace = true +sha2.workspace = true strum.workspace = true thiserror.workspace = true tokio = { workspace = true, features = ["macros", "rt", "sync", "time"] } diff --git a/litellm/rust/crates/redis-store/src/lib.rs b/litellm/rust/crates/redis-store/src/lib.rs index 55b3abb88cc..f1e3808fdce 100644 --- a/litellm/rust/crates/redis-store/src/lib.rs +++ b/litellm/rust/crates/redis-store/src/lib.rs @@ -10,6 +10,7 @@ pub mod rate_limits; mod remote; mod request; pub mod router_cooldowns; +pub mod router_session_pins; pub mod router_usage; pub mod sensitive_route_pins; pub mod spend_counters; diff --git a/litellm/rust/crates/redis-store/src/router_session_pins/claim.rs b/litellm/rust/crates/redis-store/src/router_session_pins/claim.rs new file mode 100644 index 00000000000..7b95af3c990 --- /dev/null +++ b/litellm/rust/crates/redis-store/src/router_session_pins/claim.rs @@ -0,0 +1,192 @@ +use litellm_python_compat::{Value, repr::to_str}; + +use super::types::Eligible; + +/// A claim settled against this router's memory alone: who holds the key, and what memory keeps +/// for it. A holder nobody may replace is answered without being written again. +#[derive(Debug, PartialEq)] +pub(super) struct LocalClaim { + pub(super) winner: Value, + pub(super) stored: Option, +} + +/// The claim script's rule for a router without Redis, or one whose Redis claim failed: the first +/// writer keeps the key unless `eligible` no longer lists it. +pub(super) fn local(existing: Option, pin: &Value, eligible: &Eligible) -> LocalClaim { + match (existing, &eligible.values) { + (Some(existing), None) => LocalClaim { + stored: legacy_matches(&existing, pin).then(|| pin.clone()), + winner: existing, + }, + (Some(existing), Some(values)) + if still_eligible(&existing, values, eligible.preserve_legacy_model_ids) => + { + LocalClaim { + stored: Some(existing.clone()), + winner: existing, + } + } + _ => LocalClaim { + stored: Some(pin.clone()), + winner: pin.clone(), + }, + } +} + +fn still_eligible(existing: &Value, eligible: &[Value], preserve_legacy_model_ids: bool) -> bool { + eligible.iter().any(|value| same_pin(existing, value)) + || (preserve_legacy_model_ids + && matches!(existing, Value::Str(_)) + && eligible.iter().any(|value| { + field(value, &Value::Str("model_id".to_owned())).is_some() + && legacy_matches(existing, value) + })) +} + +/// A stored mapping carries every field of `pin` as its text, or the bare string an older writer +/// left is the one value `pin` holds. +fn legacy_matches(stored: &Value, pin: &Value) -> bool { + let Value::Dict(fields) = pin else { + return false; + }; + match stored { + Value::Dict(_) => fields.iter().all(|(key, value)| { + field(stored, key) + .is_some_and(|held| *held != Value::None && Value::Str(to_str(held)) == *value) + }), + Value::Str(_) => matches!(fields.as_slice(), [(_, value)] if value == stored), + _ => false, + } +} + +/// Python's `==` between two pins: mappings compare by their fields, in whatever order. +fn same_pin(left: &Value, right: &Value) -> bool { + match (left, right) { + (Value::Dict(fields), Value::Dict(other)) => { + fields.len() == other.len() + && fields + .iter() + .all(|(key, value)| field(right, key) == Some(value)) + } + _ => left == right, + } +} + +fn field<'a>(mapping: &'a Value, key: &Value) -> Option<&'a Value> { + let Value::Dict(fields) = mapping else { + return None; + }; + fields + .iter() + .find(|(held, _)| held == key) + .map(|(_, value)| value) +} + +#[cfg(test)] +mod tests { + use litellm_python_compat::json; + use rstest::rstest; + + use super::*; + + fn value(text: &str) -> Value { + json::loads(text).unwrap() + } + + fn eligible(values: &[&str], preserve_legacy_model_ids: bool) -> Eligible { + Eligible { + values: Some(values.iter().copied().map(value).collect()), + preserve_legacy_model_ids, + } + } + + #[rstest] + #[case::nobody_holds_the_key( + None, + r#"{"model_id": "b"}"#, + r#"{"model_id": "b"}"#, + Some(r#"{"model_id": "b"}"#) + )] + #[case::the_first_writer_keeps_it( + Some(r#"{"model_id": "a"}"#), + r#"{"model_id": "b"}"#, + r#"{"model_id": "a"}"#, + None + )] + #[case::the_holder_claiming_again_is_refreshed( + Some(r#"{"model_id": "a"}"#), + r#"{"model_id": "a"}"#, + r#"{"model_id": "a"}"#, + Some(r#"{"model_id": "a"}"#) + )] + #[case::a_holder_with_a_number_matches_its_text( + Some(r#"{"model_id": 7}"#), + r#"{"model_id": "7"}"#, + r#"{"model_id": 7}"#, + Some(r#"{"model_id": "7"}"#) + )] + #[case::a_holder_with_a_null_field_is_not_refreshed( + Some(r#"{"model_id": null}"#), + r#"{"model_id": "None"}"#, + r#"{"model_id": null}"#, + None + )] + #[case::a_bare_id_an_older_writer_left_is_refreshed_as_the_pin( + Some(r#""a""#), + r#"{"model_id": "a"}"#, + r#""a""#, + Some(r#"{"model_id": "a"}"#) + )] + #[case::a_bare_id_does_not_match_a_pin_of_two_fields( + Some(r#""a""#), + r#"{"model_id": "a", "tier": "a"}"#, + r#""a""#, + None + )] + fn a_claim_naming_no_eligible_values_keeps_the_first_writer( + #[case] existing: Option<&str>, + #[case] pin: &str, + #[case] winner: &str, + #[case] stored: Option<&str>, + ) { + let settled = local(existing.map(value), &value(pin), &Eligible::default()); + + assert_eq!( + settled, + LocalClaim { + winner: value(winner), + stored: stored.map(value), + } + ); + } + + #[rstest] + #[case::nobody_holds_the_key(None, &[r#"{"model_id": "b"}"#], false, r#"{"model_id": "b"}"#)] + #[case::an_eligible_holder_keeps_it(Some(r#"{"model_id": "a"}"#), &[r#"{"model_id": "a"}"#, r#"{"model_id": "b"}"#], false, r#"{"model_id": "a"}"#)] + #[case::fields_match_in_any_order(Some(r#"{"model": "m", "tier": "t"}"#), &[r#"{"tier": "t", "model": "m"}"#], false, r#"{"model": "m", "tier": "t"}"#)] + #[case::a_holder_no_longer_eligible_is_replaced(Some(r#"{"model_id": "z"}"#), &[r#"{"model_id": "b"}"#], false, r#"{"model_id": "b"}"#)] + #[case::a_holder_with_an_extra_field_is_replaced(Some(r#"{"model_id": "b", "tier": "t"}"#), &[r#"{"model_id": "b"}"#], false, r#"{"model_id": "b"}"#)] + #[case::a_bare_id_is_kept_when_legacy_ids_are_preserved(Some(r#""a""#), &[r#"{"model_id": "a"}"#], true, r#""a""#)] + #[case::a_bare_id_is_replaced_when_they_are_not(Some(r#""a""#), &[r#"{"model_id": "a"}"#], false, r#"{"model_id": "b"}"#)] + #[case::a_bare_id_only_matches_a_model_id(Some(r#""a""#), &[r#"{"tier": "a"}"#], true, r#"{"model_id": "b"}"#)] + fn a_claim_naming_eligible_values_replaces_only_a_holder_they_do_not_list( + #[case] existing: Option<&str>, + #[case] listed: &[&str], + #[case] preserve_legacy_model_ids: bool, + #[case] winner: &str, + ) { + let settled = local( + existing.map(value), + &value(r#"{"model_id": "b"}"#), + &eligible(listed, preserve_legacy_model_ids), + ); + + assert_eq!( + settled, + LocalClaim { + winner: value(winner), + stored: Some(value(winner)), + } + ); + } +} diff --git a/litellm/rust/crates/redis-store/src/router_session_pins/keys.rs b/litellm/rust/crates/redis-store/src/router_session_pins/keys.rs new file mode 100644 index 00000000000..6e7db8811ce --- /dev/null +++ b/litellm/rust/crates/redis-store/src/router_session_pins/keys.rs @@ -0,0 +1,181 @@ +use litellm_python_compat::{Value, json, number::hex_int}; +use sha2::{Digest, Sha256}; + +use crate::Error; + +#[derive(Clone, Debug, PartialEq, Eq, Hash)] +pub enum PinKey { + /// Deployment affinity by caller. + Deployment { model_group: String, caller: String }, + /// Deployment affinity by session, scoped to the caller when there is one. + DeploymentSession { + model_group: String, + caller: Option, + session_id: String, + }, + /// The complexity router's session pin, whose caller is already the API key hash. + TierSession { + router: String, + caller: Option, + session_id: String, + }, + /// The complexity router's model within a tier. `tiers` hashes the router's tier config, and + /// `replay_fallback` is the model a replayed turn was pinned to. + TierModel { + router: String, + tiers: String, + caller: Option, + session_id: String, + tier: String, + replay_fallback: Option, + }, +} + +impl PinKey { + pub fn spelling(&self) -> Result { + Ok(match self { + Self::Deployment { + model_group, + caller, + } => format!("deployment_affinity:v1:{model_group}:{}", hashed(caller)), + Self::DeploymentSession { + model_group, + caller, + session_id, + } => { + let scope = caller + .as_deref() + .map_or_else(|| "unscoped".to_owned(), hashed); + format!("deployment_affinity:v1:session:{model_group}:{scope}:{session_id}") + } + Self::TierSession { + router, + caller, + session_id, + } => { + let scope = caller.as_deref().filter(|caller| !caller.is_empty()); + format!( + "complexity_router_session_affinity:v1:{router}:{}:{session_id}", + scope.unwrap_or("unscoped") + ) + } + Self::TierModel { + router, + tiers, + caller, + session_id, + tier, + replay_fallback, + } => { + let text = |value: &str| Value::Str(value.to_owned()); + let identity = [ + text(router), + text(tiers), + caller.as_deref().map_or(Value::None, text), + text(session_id), + text(tier), + ] + .into_iter() + .chain( + replay_fallback + .as_deref() + .map(|model| Value::List(vec![text("replay_fallback"), text(model)])), + ) + .collect(); + format!( + "complexity_router_tier_model_affinity:v1:{}", + sha256_hex(&json::dumps(&Value::List(identity))?) + ) + } + }) + } +} + +/// `_hash_user_key`: a caller that already reads as a SHA-256 hex digest to `int(text, 16)` keeps +/// its text lowercased, anything else is hashed. +fn hashed(caller: &str) -> String { + match caller.chars().count() == 64 && hex_int(caller).is_some() { + true => caller.to_lowercase(), + false => sha256_hex(caller), + } +} + +fn sha256_hex(text: &str) -> String { + format!("{:x}", Sha256::digest(text.as_bytes())) +} + +#[cfg(test)] +mod tests { + use rstest::rstest; + + use super::PinKey; + + fn text(value: &str) -> String { + value.to_owned() + } + + // Expected values are main's DeploymentAffinityCheck.get_affinity_cache_key and + // get_session_affinity_cache_key, and ComplexityRouter's session and tier model keys + #[rstest] + #[case::caller_hashed( + PinKey::Deployment { model_group: text("gpt-4o"), caller: text("sk-caller") }, + "deployment_affinity:v1:gpt-4o:8b6712790a2089c67aa97a2d80022df18cc65c7814350e33baebe79aab508891" + )] + #[case::caller_already_a_digest_is_lowercased( + PinKey::Deployment { model_group: text("gpt-4o"), caller: "A".repeat(64) }, + "deployment_affinity:v1:gpt-4o:aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa" + )] + #[case::caller_int_reads_as_hex_is_kept( + PinKey::Deployment { model_group: text("gpt-4o"), caller: format!("0x{}", "f".repeat(62)) }, + "deployment_affinity:v1:gpt-4o:0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff" + )] + #[case::user_id_caller( + PinKey::Deployment { model_group: text("gpt-4o"), caller: text("user_id:u-1") }, + "deployment_affinity:v1:gpt-4o:a8ae9346b4f3e606434b5bf77a2ba2ea81313df4283d69d85bcef9e93a78a4f8" + )] + #[case::session_without_caller( + PinKey::DeploymentSession { model_group: text("gpt-4o"), caller: None, session_id: text("sess-1") }, + "deployment_affinity:v1:session:gpt-4o:unscoped:sess-1" + )] + #[case::session_with_caller( + PinKey::DeploymentSession { model_group: text("gpt-4o"), caller: Some(text("sk-caller")), session_id: text("sess-1") }, + "deployment_affinity:v1:session:gpt-4o:8b6712790a2089c67aa97a2d80022df18cc65c7814350e33baebe79aab508891:sess-1" + )] + #[case::session_with_empty_caller_is_hashed( + PinKey::DeploymentSession { model_group: text("gpt-4o"), caller: Some(String::new()), session_id: text("sess-1") }, + "deployment_affinity:v1:session:gpt-4o:e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855:sess-1" + )] + #[case::tier_session( + PinKey::TierSession { router: text("smart"), caller: Some(text("key-hash")), session_id: text("sess-1") }, + "complexity_router_session_affinity:v1:smart:key-hash:sess-1" + )] + #[case::tier_session_with_empty_caller_is_unscoped( + PinKey::TierSession { router: text("smart"), caller: Some(String::new()), session_id: text("sess-1") }, + "complexity_router_session_affinity:v1:smart:unscoped:sess-1" + )] + #[case::tier_model( + PinKey::TierModel { + router: text("smart"), + tiers: text("cfg"), + caller: Some(text("sk-caller")), + session_id: text("sess-1"), + tier: text("SIMPLE"), + replay_fallback: None, + }, + "complexity_router_tier_model_affinity:v1:d93f3c0fe46c4ef59f1e45d1b2aa37455b2663b05deb14d4afb3339bc2094f51" + )] + #[case::tier_model_replay_without_caller( + PinKey::TierModel { + router: text("smart"), + tiers: text("cfg"), + caller: None, + session_id: text("séss"), + tier: text("COMPLEX"), + replay_fallback: Some(text("m-a")), + }, + "complexity_router_tier_model_affinity:v1:c83437032b0b0d7d8229f7aaf444096a83510ee630e5a510e4770b0f6f03a8d3" + )] + fn a_pin_key_is_spelled_as_main_spells_it(#[case] key: PinKey, #[case] expected: &str) { + assert_eq!(key.spelling().unwrap(), expected); + } +} diff --git a/litellm/rust/crates/redis-store/src/router_session_pins/lua.rs b/litellm/rust/crates/redis-store/src/router_session_pins/lua.rs new file mode 100644 index 00000000000..bb84cee3381 --- /dev/null +++ b/litellm/rust/crates/redis-store/src/router_session_pins/lua.rs @@ -0,0 +1,51 @@ +/// `_READ_PIN_SCRIPT`, byte for byte, so pods on either stack share one script SHA. +pub(super) const READ_PIN: &str = r#" +local current = redis.call('GET', KEYS[1]) +if current ~= false then + redis.call('EXPIRE', KEYS[1], ARGV[1]) +end +return current +"#; + +/// `_CLAIM_PIN_SCRIPT`, byte for byte. +pub(super) const CLAIM_PIN: &str = r#" +local current = redis.call('GET', KEYS[1]) +if current == false then + redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) + return ARGV[1] +end +if ARGV[3] then + local decoded, stored = pcall(cjson.decode, current) + if ARGV[4] then + local legacy = current + if decoded and type(stored) == 'string' then legacy = stored end + for _, eligible in ipairs(cjson.decode(ARGV[3])) do + if eligible['model_id'] == legacy then + redis.call('EXPIRE', KEYS[1], ARGV[2]) + return current + end + end + end + if decoded and type(stored) == 'table' then + for _, eligible in ipairs(cjson.decode(ARGV[3])) do + local matches = true + for key, value in pairs(eligible) do + if stored[key] ~= value then matches = false; break end + end + for key, _ in pairs(stored) do + if eligible[key] == nil then matches = false; break end + end + if matches then + redis.call('EXPIRE', KEYS[1], ARGV[2]) + return current + end + end + end + redis.call('SET', KEYS[1], ARGV[1], 'EX', ARGV[2]) + return ARGV[1] +end +if current == ARGV[1] then + redis.call('EXPIRE', KEYS[1], ARGV[2]) +end +return current +"#; diff --git a/litellm/rust/crates/redis-store/src/router_session_pins/mod.rs b/litellm/rust/crates/redis-store/src/router_session_pins/mod.rs new file mode 100644 index 00000000000..97e66d3a038 --- /dev/null +++ b/litellm/rust/crates/redis-store/src/router_session_pins/mod.rs @@ -0,0 +1,232 @@ +mod claim; +mod keys; +mod lua; +mod types; + +use std::{ + sync::{Arc, LazyLock}, + time::Duration, +}; + +use litellm_python_compat::{Value, json}; +use litellm_redis_client::{Script, commands}; + +use crate::{ + DecodeError, Error, Pending, Send, codec, + dual_cache::{DualCache, MemorySettings}, + memory::{Clock, system_clock}, + remote::{Remote, RemoteSlot}, + request::ready, +}; + +static READ_PIN: LazyLock