scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

feat(redis): serve the router's session pins from a native store (#50870)

Browse source 
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.
This commit is contained in:
yujonglee 2026-11-13 12:22:28 +09:00 • committed by GitHub
parent b88e4bff6b
commit 9a1acee4a7
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
37 changed files with 2482 additions and 1082 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

178
litellm/caching/affinity_cache.py
View file

@ -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,

4
litellm/redis/stores/connection.py
View file

@ -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

133
litellm/redis/stores/router_session_pins/__init__.py Normal file
View file

@ -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

37
litellm/redis/stores/router_session_pins/handle.py Normal file
View file

@ -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: ...

37
litellm/redis/stores/router_session_pins/keys.py Normal file
View file

@ -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

9
litellm/router_api.py
View file

@ -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``)."""

15
litellm/router_rust/core.py
View file

@ -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(

5
litellm/router_rust/steps/callbacks.py
View file

@ -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,

12
litellm/router_rust/steps/prerouting.py
View file

@ -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,

49
litellm/router_strategy/complexity_router/complexity_router.py
View file

@ -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)

72
litellm/router_utils/pre_call_checks/deployment_affinity_check.py
View file

@ -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,

1
litellm/rust/Cargo.lock generated
View file

@ -4553,6 +4553,7 @@ dependencies = [
"litellm-secret-value",
"litellm-tracing",
"rstest",
"sha2 0.10.9",
"strum",
"thiserror 2.0.19",
"tokio",

3
litellm/rust/crates/python-bridge/src/lib.rs
View file

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

5
litellm/rust/crates/python-bridge/src/redis/store/mod.rs
View file

@ -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<T> {
}
}
fn unwrap_values(values: Vec<PyValue>) -> Vec<litellm_python_compat::Value> {
values.into_iter().map(|value| value.0).collect()
}
fn unwrap_all<T>(values: Vec<PyData<T>>) -> Vec<T> {
values.into_iter().map(|value| value.0).collect()
}

234
litellm/rust/crates/python-bridge/src/redis/store/router_session_pins.rs Normal file
View file

@ -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<litellm_python_compat::Value>) -> Option<PyValue> {
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<String>,
session_id: String,
tier: String,
replay_fallback: Option<String>,
},
TierSession {
router: String,
caller: Option<String>,
session_id: String,
},
DeploymentSession {
model_group: String,
caller: Option<String>,
session_id: String,
},
Deployment {
model_group: String,
caller: String,
},
}
impl From<PyPinKey> 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<PyAny>) -> Clock {
Arc::new(move || {
Python::attach(|py| clock.call0(py)?.extract::<f64>(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<RouterSessionPins>);
#[pymethods]
impl NativeRouterSessionPins {
#[new]
#[pyo3(signature = (client, clock=None))]
fn new(
py: Python<'_>,
client: Option<PyRef<'_, NativeRedisClient>>,
clock: Option<Py<PyAny>>,
) -> PyResult<Self> {
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<Bound<'py, PyAny>> {
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<Option<PyValue>> {
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<Bound<'py, PyAny>> {
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<Bound<'py, PyAny>> {
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<NativePending> {
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<Vec<PyValue>>,
preserve_legacy_model_ids: bool,
) -> PyResult<Bound<'py, PyAny>> {
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<Vec<PyValue>>,
preserve_legacy_model_ids: bool,
) -> PyResult<PyValue> {
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<Vec<PyValue>>, preserve_legacy_model_ids: bool) -> Eligible {
Eligible {
values: values.map(unwrap_values),
preserve_legacy_model_ids,
}
}

59
litellm/rust/crates/python-compat/src/number.rs
View file

@ -7,6 +7,36 @@ pub fn int(text: &str) -> Option<BigInt> {
parse(text)
}
/// `int(text, 16)`, or `None` where Python raises `ValueError`.
pub fn hex_int(text: &str) -> Option<BigInt> {
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<usize>| {
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<f64> {
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<i128>) {
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))]

1
litellm/rust/crates/redis-store/Cargo.toml
View file

@ -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"] }

1
litellm/rust/crates/redis-store/src/lib.rs
View file

@ -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;

192
litellm/rust/crates/redis-store/src/router_session_pins/claim.rs Normal file
View file

@ -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<Value>,
}
/// 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<Value>, 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)),
}
);
}
}

181
litellm/rust/crates/redis-store/src/router_session_pins/keys.rs Normal file
View file

@ -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<String>,
session_id: String,
},
/// The complexity router's session pin, whose caller is already the API key hash.
TierSession {
router: String,
caller: Option<String>,
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<String>,
session_id: String,
tier: String,
replay_fallback: Option<String>,
},
}
impl PinKey {
pub fn spelling(&self) -> Result<String, Error> {
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);
}
}

51
litellm/rust/crates/redis-store/src/router_session_pins/lua.rs Normal file
View file

@ -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
"#;

232
litellm/rust/crates/redis-store/src/router_session_pins/mod.rs Normal file
View file

@ -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<Script> = LazyLock::new(|| Script::new(lua::READ_PIN));
static CLAIM_PIN: LazyLock<Script> = LazyLock::new(|| Script::new(lua::CLAIM_PIN));
pub use keys::PinKey;
pub use types::Eligible;
pub struct RouterSessionPins {
pins: DualCache<Value>,
}
impl RouterSessionPins {
pub fn new(remote: Option<Remote>, default_in_memory_ttl: Option<Duration>) -> Self {
Self::with_clock(remote, default_in_memory_ttl, system_clock())
}
pub fn with_clock(
remote: Option<Remote>,
default_in_memory_ttl: Option<Duration>,
now: Clock,
) -> Self {
Self {
pins: DualCache::new(
Arc::new(RemoteSlot::new(remote)),
MemorySettings {
entries: None,
ttl: default_in_memory_ttl,
},
now,
),
}
}
pub fn attach(&self, remote: Remote) -> bool {
self.pins.redis().attach(remote)
}
pub fn has_redis(&self) -> bool {
self.pins.redis().get().is_some()
}
pub async fn get(&self, key: &PinKey) -> Result<Option<Value>, Error> {
if let Some(local) = self.get_local(key)? {
return Ok(Some(local));
}
self.pins.get(None, &key.spelling()?, None).await
}
/// `in_memory_cache.get_cache`, which hands a stored string back as `json.loads` reads it.
pub fn get_local(&self, key: &PinKey) -> Result<Option<Value>, Error> {
Ok(self
.pins
.memory()
.get(&key.spelling()?)?
.map(|value| match value {
Value::Str(text) => json::loads(&text).unwrap_or(Value::Str(text)),
other => other,
}))
}
/// `set_local_affinity_pin`: the entry is replaced, since `InMemoryCache.set_cache` keeps a live
/// key's expiry.
pub fn set_local(&self, key: &PinKey, value: Value, ttl: Duration) -> Result<(), Error> {
self.replace_local(&key.spelling()?, value, ttl)
}
pub fn delete_local(&self, key: &PinKey) -> Result<(), Error> {
Ok(self.pins.memory().delete(&key.spelling()?)?)
}
pub async fn set(&self, key: &PinKey, value: Value, ttl: Duration) -> Result<(), Error> {
self.pins.set(&key.spelling()?, value, Some(ttl)).await
}
/// Unlike `get`, a Redis failure reaches the caller, which then keeps its memory pin.
pub async fn get_shared(&self, key: &PinKey) -> Result<Option<Value>, Error> {
let Some(client) = self.pins.redis().client() else {
return Ok(None);
};
let raw = client
.send(commands::get::<Vec<u8>>(&key.spelling()?))
.await?;
Ok(raw.as_deref().map(codec::decode).transpose()?)
}
/// `read_affinity_pin`'s Redis half: the read script refreshes the pin's TTL on the batch
/// `send` picks, then memory takes what Redis holds, or drops the key on a miss. A Redis failure
/// is the caller's, which keeps its memory pin. Without Redis, memory answers.
pub fn read(
self: &Arc<Self>,
send: Send<'_>,
key: &PinKey,
ttl: Duration,
) -> Result<Pending<Option<Value>>, Error> {
let Some(batch) = self.pins.redis().batch(send) else {
return Ok(ready(self.get_local(key)?));
};
let key = key.spelling()?;
let declared = batch.add(commands::eval::<Option<Vec<u8>>>(
&READ_PIN,
&[key.as_str()],
vec![ttl.as_secs().to_string()],
));
let store = Arc::clone(self);
Ok(Box::pin(async move {
let pin = decode_pin(declared.resolve().await?)?;
match &pin {
None => store.pins.memory().delete(&key)?,
Some(value) => store.replace_local(&key, value.clone(), ttl)?,
}
Ok(pin)
}))
}
/// `claim_affinity_pin`'s Redis half: the claim script keeps the first writer unless it is no
/// longer eligible, and memory takes the winner. `None` when the script answered something other
/// than a pin, or without Redis. A Redis failure is the caller's, which settles the claim with
/// `claim_local`.
pub async fn claim(
&self,
key: &PinKey,
pin: &Value,
ttl: Duration,
eligible: &Eligible,
) -> Result<Option<Value>, Error> {
let Some(client) = self.pins.redis().client() else {
return Ok(None);
};
let key = key.spelling()?;
let reply = client
.send(commands::eval::<Option<Vec<u8>>>(
&CLAIM_PIN,
&[key.as_str()],
claim_args(pin, ttl, eligible)?,
))
.await?;
let Some(winner) = decode_pin(reply)? else {
return Ok(None);
};
self.replace_local(&key, winner.clone(), ttl)?;
Ok(Some(winner))
}
/// `claim` against this router's memory alone, for a router without Redis or one whose Redis
/// claim failed.
pub fn claim_local(
&self,
key: &PinKey,
pin: &Value,
ttl: Duration,
eligible: &Eligible,
) -> Result<Value, Error> {
let settled = claim::local(self.get_local(key)?, pin, eligible);
if let Some(stored) = settled.stored {
self.set_local(key, stored, ttl)?;
}
Ok(settled.winner)
}
fn replace_local(&self, key: &str, value: Value, ttl: Duration) -> Result<(), Error> {
self.pins.memory().delete(key)?;
Ok(self.pins.memory().set(key, value, Some(ttl))?)
}
}
/// `json.dumps(pin)`, the TTL, then `json.dumps(eligible)` and `"1"` for legacy ids, each only when
/// given, as `claim_affinity_pin` builds `ARGV`.
fn claim_args(pin: &Value, ttl: Duration, eligible: &Eligible) -> Result<Vec<String>, Error> {
let eligible_args = match &eligible.values {
None => Vec::new(),
Some(values) => std::iter::once(json::dumps(&Value::List(values.clone()))?)
.chain(eligible.preserve_legacy_model_ids.then(|| "1".to_owned()))
.collect(),
};
Ok([json::dumps(pin)?, ttl.as_secs().to_string()]
.into_iter()
.chain(eligible_args)
.collect())
}
/// `_decode_pin` over a script's reply: JSON, else the text as written. A reply that is not UTF-8
/// fails, as its `decode("utf-8")` raises there.
fn decode_pin(raw: Option<Vec<u8>>) -> Result<Option<Value>, Error> {
let Some(raw) = raw else {
return Ok(None);
};
let text = String::from_utf8(raw).map_err(|_| DecodeError::NotUtf8)?;
Ok(Some(json::loads(&text).unwrap_or(Value::Str(text))))
}
#[cfg(test)]
mod tests {
use rstest::rstest;
use super::*;
#[rstest]
#[case::json(br#"{"model_id": "w"}"#.as_slice(), json::loads(r#"{"model_id": "w"}"#).unwrap())]
#[case::json_string(br#""w""#.as_slice(), Value::Str("w".into()))]
#[case::bare_text(b"w".as_slice(), Value::Str("w".into()))]
fn a_script_reply_decodes_as_json_else_as_its_text(
#[case] raw: &[u8],
#[case] expected: Value,
) {
assert_eq!(decode_pin(Some(raw.to_vec())).unwrap(), Some(expected));
}
#[test]
fn a_script_reply_that_is_not_utf8_fails() {
assert!(decode_pin(Some(vec![0xff])).is_err());
}
}

9
litellm/rust/crates/redis-store/src/router_session_pins/types.rs Normal file
View file

@ -0,0 +1,9 @@
use litellm_python_compat::Value;
/// What a claim asks for beyond the pin itself: the values that may keep holding the key, and
/// whether a bare deployment id an older writer left counts as one of them.
#[derive(Clone, Debug, Default, PartialEq)]
pub struct Eligible {
pub values: Option<Vec<Value>>,
pub preserve_legacy_model_ids: bool,
}

368
litellm/rust/crates/redis-store/tests/router_session_pins.rs Normal file
View file

@ -0,0 +1,368 @@
mod support;
use std::{sync::Arc, time::Duration};
use litellm_python_compat::{Value, json};
use litellm_redis_client::{Operation, Ttl, commands};
use litellm_redis_store::{
RequestRedis, Send as RedisSend,
router_session_pins::{Eligible, PinKey, RouterSessionPins},
};
use litellm_redis_testing::{RedisDeployment, Topology};
use rstest::rstest;
use support::{connect, remote, within};
use tokio::runtime::Handle;
const KEY: &str = "deployment_affinity:v1:session:group:unscoped:session-1";
fn key() -> PinKey {
PinKey::DeploymentSession {
model_group: "group".to_owned(),
caller: None,
session_id: "session-1".to_owned(),
}
}
const TTL: Duration = Duration::from_secs(600);
fn pin(model_id: &str) -> Value {
json::loads(&format!(r#"{{"model_id": "{model_id}"}}"#)).unwrap()
}
fn store(deployment: &RedisDeployment) -> Arc<RouterSessionPins> {
Arc::new(RouterSessionPins::new(
Some(remote(connect(deployment))),
None,
))
}
async fn stored(raw: &litellm_redis_client::RedisClient) -> Option<String> {
raw.send(commands::get(KEY)).await.unwrap()
}
async fn cli(deployment: &RedisDeployment, args: &[&str]) -> Vec<String> {
match deployment {
RedisDeployment::Standalone(server) => vec![server.cli(args).await.unwrap()],
RedisDeployment::Cluster(cluster) => cluster.cli_outputs(args).await.unwrap(),
}
}
enum Script {
Read,
Claim,
}
// sha1 of _READ_PIN_SCRIPT and _CLAIM_PIN_SCRIPT on main
#[rstest]
#[case::read_pin(Script::Read, "f5f07e272c68ee98239d41ee590bd6d17c2ed2c2")]
#[case::claim_pin(Script::Claim, "e3942e42b49307b19d1f79eaeff8f207bcc96934")]
#[tokio::test]
async fn each_script_runs_under_mains_sha(
#[case] script: Script,
#[case] sha: &str,
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let pins = store(&deployment);
match script {
Script::Read => {
pins.read(RedisSend::Now, &key(), TTL)
.unwrap()
.await
.unwrap();
}
Script::Claim => {
pins.claim(&key(), &pin("dep-a"), TTL, &Eligible::default())
.await
.unwrap();
}
}
let cached = cli(&deployment, &["SCRIPT", "EXISTS", sha]).await;
assert!(cached.iter().any(|reply| reply.trim() == "1"), "{cached:?}");
}
// claim_affinity_pin writes json.dumps(pin) with SET ... EX, and the first writer keeps the key
#[rstest]
#[tokio::test]
async fn the_first_claim_wins_and_is_stored_as_mains_json(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let first = store(&deployment);
let second = store(&deployment);
let won = first
.claim(&key(), &pin("dep-a"), TTL, &Eligible::default())
.await
.unwrap();
let lost = second
.claim(&key(), &pin("dep-b"), TTL, &Eligible::default())
.await
.unwrap();
assert_eq!(won, Some(pin("dep-a")));
assert_eq!(lost, Some(pin("dep-a")));
assert_eq!(second.get_local(&key()).unwrap(), Some(pin("dep-a")));
assert_eq!(
stored(&raw).await,
Some(r#"{"model_id": "dep-a"}"#.to_owned())
);
let Ttl::Expires(remaining) = raw.send(commands::ttl(KEY)).await.unwrap() else {
panic!("the pin has no TTL");
};
assert!(within(remaining, 600), "{remaining:?}");
}
#[rstest]
#[case::eligible_holder_kept(r#"{"model_id": "dep-a"}"#, false, "dep-a")]
#[case::ineligible_holder_replaced(r#"{"model_id": "dep-gone"}"#, false, "dep-b")]
#[case::legacy_id_kept(r#""dep-a""#, true, "dep-a")]
#[case::legacy_id_replaced_without_the_flag("dep-a", false, "dep-b")]
#[tokio::test]
async fn an_eligible_claim_replaces_only_a_holder_no_longer_eligible(
#[case] held: &str,
#[case] preserve_legacy_model_ids: bool,
#[case] winner: &str,
) {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let raw = connect(&deployment);
let pins = store(&deployment);
raw.send(commands::set(KEY, held, TTL)).await.unwrap();
let eligible = Eligible {
values: Some(vec![pin("dep-a"), pin("dep-b")]),
preserve_legacy_model_ids,
};
let claimed = pins
.claim(&key(), &pin("dep-b"), TTL, &eligible)
.await
.unwrap()
.unwrap();
let expected = if winner == "dep-b" {
pin("dep-b")
} else {
json::loads(held).unwrap()
};
assert_eq!(claimed, expected);
}
// read_affinity_pin's script rides the request's pre-call pipeline and refreshes the pin's TTL
#[rstest]
#[tokio::test]
async fn a_pin_read_rides_the_request_batch_and_refreshes_ttl_and_memory(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let (client, trips) = support::connect_observed(&deployment);
let pins = Arc::new(RouterSessionPins::new(Some(remote(client)), None));
raw.send(commands::set(
KEY,
r#"{"model_id": "dep-a"}"#,
Duration::from_secs(5),
))
.await
.unwrap();
let request = RequestRedis::new(Handle::current(), Duration::from_secs(60));
let declared = pins
.read(RedisSend::PreCall(&request), &key(), TTL)
.unwrap();
assert!(trips.seen().is_empty());
request.flush_pre().await;
let read = declared.await.unwrap();
let one_script = match topology {
Topology::Standalone => Operation::Pipeline { ops: 1 },
Topology::Cluster => Operation::Script,
};
assert_eq!(trips.seen(), [one_script]);
assert_eq!(read, Some(pin("dep-a")));
assert_eq!(pins.get_local(&key()).unwrap(), Some(pin("dep-a")));
let Ttl::Expires(remaining) = raw.send(commands::ttl(KEY)).await.unwrap() else {
panic!("the pin has no TTL");
};
assert!(within(remaining, 600), "{remaining:?}");
}
#[rstest]
#[tokio::test]
async fn a_pin_read_missing_in_redis_drops_the_memory_pin() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let pins = store(&deployment);
pins.set_local(&key(), pin("dep-a"), TTL).unwrap();
let read = pins
.read(RedisSend::Now, &key(), TTL)
.unwrap()
.await
.unwrap();
assert_eq!(read, None);
assert_eq!(pins.get_local(&key()).unwrap(), None);
}
#[rstest]
#[tokio::test]
async fn redis_failures_reach_the_caller_which_keeps_its_memory_pin() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let pins = store(&deployment);
pins.set_local(&key(), pin("dep-a"), TTL).unwrap();
deployment.stop().await.unwrap();
let read = pins.read(RedisSend::Now, &key(), TTL).unwrap().await;
let claim = pins
.claim(&key(), &pin("dep-b"), TTL, &Eligible::default())
.await;
let shared = pins.get_shared(&key()).await;
assert!(read.is_err());
assert!(claim.is_err());
assert!(shared.is_err());
assert_eq!(pins.get_local(&key()).unwrap(), Some(pin("dep-a")));
}
#[rstest]
#[tokio::test]
async fn without_redis_reads_answer_from_memory_and_claims_have_no_winner() {
let pins = Arc::new(RouterSessionPins::new(None, None));
pins.set_local(&key(), pin("dep-a"), TTL).unwrap();
assert_eq!(
pins.read(RedisSend::Now, &key(), TTL)
.unwrap()
.await
.unwrap(),
Some(pin("dep-a"))
);
assert_eq!(
pins.claim(&key(), &pin("dep-b"), TTL, &Eligible::default())
.await
.unwrap(),
None
);
assert_eq!(pins.get_shared(&key()).await.unwrap(), None);
}
// DualCache.async_set_cache writes json.dumps(value), which every pod's async_get_cache reads back
#[rstest]
#[tokio::test]
async fn a_set_pin_is_mains_json_and_another_router_reads_it_from_redis(
#[values(Topology::Standalone, Topology::Cluster)] topology: Topology,
) {
let deployment = RedisDeployment::start(topology).await.unwrap();
let raw = connect(&deployment);
let writer = store(&deployment);
let reader = store(&deployment);
let session_pin =
json::loads(r#"{"model": "anthropic/claude-sonnet-5-5", "tier": "COMPLEX"}"#).unwrap();
writer.set(&key(), session_pin.clone(), TTL).await.unwrap();
assert_eq!(
stored(&raw).await,
Some(r#"{"model": "anthropic/claude-sonnet-5-5", "tier": "COMPLEX"}"#.to_owned())
);
assert_eq!(reader.get(&key()).await.unwrap(), Some(session_pin.clone()));
assert_eq!(reader.get_local(&key()).unwrap(), Some(session_pin.clone()));
assert_eq!(reader.get_shared(&key()).await.unwrap(), Some(session_pin));
}
// set_local_affinity_pin deletes before setting, since InMemoryCache.set_cache keeps a live key's
// expiry: a re-pinned key lives for the new TTL
#[rstest]
#[tokio::test]
async fn a_local_pin_set_again_takes_the_new_ttl() {
let seconds = Arc::new(std::sync::atomic::AtomicU64::new(1_000));
let now = {
let seconds = seconds.clone();
Arc::new(move || Duration::from_secs(seconds.load(std::sync::atomic::Ordering::SeqCst)))
};
let pins = RouterSessionPins::with_clock(None, None, now);
pins.set_local(&key(), pin("dep-a"), Duration::from_secs(10))
.unwrap();
seconds.fetch_add(5, std::sync::atomic::Ordering::SeqCst);
pins.set_local(&key(), pin("dep-b"), TTL).unwrap();
seconds.fetch_add(30, std::sync::atomic::Ordering::SeqCst);
assert_eq!(pins.get_local(&key()).unwrap(), Some(pin("dep-b")));
}
// InMemoryCache.get_cache hands a stored string back as json.loads reads it
#[rstest]
#[case::json_text(r#"{ "model" : "first" }"#, r#"{"model": "first"}"#)]
#[case::bare_text("first", r#""first""#)]
#[tokio::test]
async fn a_memory_string_pin_reads_back_as_json_when_it_parses(
#[case] stored: &str,
#[case] expected: &str,
) {
let pins = RouterSessionPins::new(None, None);
pins.set_local(&key(), Value::Str(stored.to_owned()), TTL)
.unwrap();
assert_eq!(
pins.get_local(&key()).unwrap(),
Some(json::loads(expected).unwrap())
);
assert_eq!(
pins.get(&key()).await.unwrap(),
Some(json::loads(expected).unwrap())
);
}
// DualCache.async_get_cache keeps a pin it read from Redis for the router's default_in_memory_ttl
#[rstest]
#[tokio::test]
async fn a_pin_read_from_redis_stays_in_memory_for_the_default_in_memory_ttl() {
let deployment = RedisDeployment::start(Topology::Standalone).await.unwrap();
let raw = connect(&deployment);
raw.send(commands::set(KEY, r#"{"model_id": "dep-a"}"#, TTL))
.await
.unwrap();
let seconds = Arc::new(std::sync::atomic::AtomicU64::new(1_000));
let now = {
let seconds = seconds.clone();
Arc::new(move || Duration::from_secs(seconds.load(std::sync::atomic::Ordering::SeqCst)))
};
let pins = RouterSessionPins::with_clock(
Some(remote(connect(&deployment))),
Some(Duration::from_secs(5)),
now,
);
assert_eq!(pins.get(&key()).await.unwrap(), Some(pin("dep-a")));
seconds.fetch_add(4, std::sync::atomic::Ordering::SeqCst);
assert_eq!(pins.get_local(&key()).unwrap(), Some(pin("dep-a")));
seconds.fetch_add(2, std::sync::atomic::Ordering::SeqCst);
assert_eq!(pins.get_local(&key()).unwrap(), None);
}
#[rstest]
#[tokio::test]
async fn a_local_claim_keeps_the_first_writer_in_memory_until_it_is_no_longer_eligible() {
let store = RouterSessionPins::new(None, None);
let listing_b = Eligible {
values: Some(vec![pin("b")]),
preserve_legacy_model_ids: false,
};
let first = store
.claim_local(&key(), &pin("a"), TTL, &Eligible::default())
.unwrap();
let second = store
.claim_local(&key(), &pin("b"), TTL, &Eligible::default())
.unwrap();
let replaced = store
.claim_local(&key(), &pin("b"), TTL, &listing_b)
.unwrap();
assert_eq!(
(first, second, replaced, store.get_local(&key()).unwrap()),
(pin("a"), pin("a"), pin("b"), Some(pin("b")))
);
}

36
litellm/rust_bridge/_native.pyi
View file

@ -1,5 +1,5 @@
from asyncio import Future
from collections.abc import Mapping, Sequence
from collections.abc import Callable, Mapping, Sequence
from typing import Generic, Literal, Never, NotRequired, Protocol, TypedDict, TypeVar, final
import httpx
@ -7,6 +7,7 @@ from pydantic import JsonValue
from typing_extensions import disjoint_base
from litellm.redis.store_settings import RedisBreakerSettings, RedisStoreSettings
from litellm.redis.stores.router_session_pins.keys import PinKey
from litellm.rust_bridge.execution import Execution
from litellm.rust_bridge.public_call import NativeCall
@ -255,6 +256,38 @@ class NativeUsageRead:
def read(self) -> NativePending[list[float | None]]: ...
def release(self) -> None: ...
@final
class NativeRouterSessionPins:
def __new__(
cls, client: NativeRedisClient | None, clock: Callable[[], float] | None = None
) -> NativeRouterSessionPins: ...
def attach(self, client: NativeRedisClient) -> bool: ...
@property
def has_redis(self) -> bool: ...
def get(self, key: PinKey) -> Future[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) -> Future[None]: ...
def get_shared(self, key: PinKey) -> Future[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,
) -> Future[object]: ...
def claim_local(
self,
key: PinKey,
value: object,
ttl_seconds: float,
eligible_values: Sequence[object] | None,
preserve_legacy_model_ids: bool,
) -> object: ...
@final
class NativeRouterUsage:
def __new__(cls, client: NativeRedisClient | None) -> NativeRouterUsage: ...
@ -742,6 +775,7 @@ __all__ = [
"NativeRejected",
"NativeRequestRedis",
"NativeRouterCooldowns",
"NativeRouterSessionPins",
"NativeRouterUsage",
"NativeSend",
"NativeSensitiveRoutePins",

183
tests/integration/sdk/test_dual_cache_redis.py
View file

@ -6,7 +6,6 @@ from unittest.mock import patch
import pytest
from litellm.caching.affinity_cache import claim_affinity_pin, read_affinity_pin
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache
@ -93,185 +92,3 @@ async def test_sync_and_async_batch_reads_share_one_redis_without_sync_reads_goi
assert in_loop_results == [[expected[key]] for key in in_loop_keys]
assert await dual_cache.async_batch_get_cache(keys=[survivor_key]) == [expected[survivor_key]]
@pytest.mark.asyncio
async def test_concurrent_session_repin_claims_keep_the_shared_winner_after_recovery() -> None:
redis: Final = _redis_cache()
workers: Final = (DualCache(redis_cache=redis), DualCache(redis_cache=redis))
key: Final = f"session-repin-{uuid.uuid4().hex}"
choices: Final = ({"model_id": "replacement"}, {"model_id": "other"})
try:
await redis.async_set_cache(key, {"model_id": "original"}, ttl=60)
for worker in workers:
worker.in_memory_cache.set_cache(key, {"model_id": "original"}, ttl=60)
winners: Final = await asyncio.gather(
*(
claim_affinity_pin(worker, key, choice, 60, eligible_values=choices, preserve_legacy_model_ids=True)
for worker, choice in zip(workers, choices, strict=True)
)
)
assert winners[0] == winners[1] and winners[0] in choices
late_original: Final = await claim_affinity_pin(
workers[0],
key,
{"model_id": "original"},
60,
eligible_values=({"model_id": "original"}, *choices),
preserve_legacy_model_ids=True,
)
assert late_original == winners[0]
assert await redis.async_get_cache(key) == winners[0]
finally:
await redis.async_delete_cache(key)
@pytest.mark.asyncio
@pytest.mark.parametrize("stored", ["original", '"original"', '{"model_id":"original"}'])
async def test_session_repin_lua_preserves_healthy_legacy_pins_and_replaces_unavailable_ones(stored: str) -> None:
redis: Final = _redis_cache()
cache: Final = DualCache(redis_cache=redis)
key: Final = f"session-repin-legacy-{uuid.uuid4().hex}"
replacement: Final = {"model_id": "replacement"}
try:
redis.redis_client.set(redis.check_and_fix_namespace(key), stored, ex=60)
preserved: Final = await claim_affinity_pin(
cache,
key,
replacement,
60,
eligible_values=({"model_id": "original"}, replacement),
preserve_legacy_model_ids=True,
)
assert preserved in ("original", {"model_id": "original"})
moved: Final = await claim_affinity_pin(
cache,
key,
replacement,
60,
eligible_values=(replacement,),
preserve_legacy_model_ids=True,
)
assert moved == replacement
assert await redis.async_get_cache(key) == replacement
finally:
await redis.async_delete_cache(key)
@pytest.mark.asyncio
@pytest.mark.parametrize("stored", ['{"model":"current","tier":"REASONING"}', '"current"', "current"])
async def test_tier_pin_reads_shared_state_refreshes_ttl_and_forgets_deleted_pins(stored: str) -> None:
redis: Final = _redis_cache()
redis.namespace = f"tier-read-{uuid.uuid4().hex}"
workers: Final = (DualCache(redis_cache=redis), DualCache(redis_cache=redis))
key: Final = "session"
namespaced: Final = redis.check_and_fix_namespace(key)
try:
for worker in workers:
worker.in_memory_cache.set_cache(key, {"model": "old", "tier": "SIMPLE"}, ttl=60)
redis.redis_client.set(namespaced, stored, ex=10)
first: Final = await read_affinity_pin(workers[0], key, 60)
second: Final = await read_affinity_pin(workers[1], key, 60)
assert first == second and first in ("current", {"model": "current", "tier": "REASONING"})
assert all(worker.in_memory_cache.get_cache(key) == first for worker in workers)
assert 10 < redis.redis_client.ttl(namespaced) <= 60
replacement: Final = {"model": "newer", "tier": "SIMPLE"}
await redis.async_set_cache(key, replacement, ttl=10)
assert await read_affinity_pin(workers[0], key, 60) == replacement
assert await read_affinity_pin(workers[1], key, 60) == replacement
await redis.async_delete_cache(key)
for worker in workers:
assert await read_affinity_pin(worker, key, 60) is None
assert worker.in_memory_cache.get_cache(key) is None
assert redis.redis_client.exists(namespaced) == 0
finally:
await redis.async_delete_cache(key)
@pytest.mark.asyncio
@pytest.mark.parametrize("mode", ["joined", "separate", "script_error"])
async def test_tier_pin_joins_pending_reads_without_reusing_request_state(mode: str) -> None:
import json
from collections.abc import Iterable
import litellm
from redis.asyncio import BlockingConnectionPool, Connection, Redis
from litellm.caching.redis_batch import request_redis_batch_scope
sends: Final[list[bytes]] = []
class RecordingConnection(Connection):
async def send_packed_command(
self, command: bytes | str | Iterable[bytes], check_health: bool = True
) -> None:
packed: Final = command.encode() if isinstance(command, str) else (
command if isinstance(command, bytes) else b"".join(command)
)
sends.append(packed)
await super().send_packed_command(packed, check_health=check_health)
namespace: Final = f"tier-batch-{uuid.uuid4().hex}"
redis: Final = RedisCache(
host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]),
namespace=namespace, client_name=namespace,
)
pool: Final = BlockingConnectionPool(
host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]),
connection_class=RecordingConnection,
)
client: Final = Redis(connection_pool=pool)
client_key: Final = redis._get_async_client_cache_key()
litellm.in_memory_llm_clients_cache.flush_cache()
litellm.in_memory_llm_clients_cache.set_cache(key=client_key, value=client)
cache: Final = DualCache(redis_cache=redis)
key: Final = redis.check_and_fix_namespace("pin")
other_key: Final = redis.check_and_fix_namespace("other")
stale: Final = {"model": "old", "tier": "SIMPLE"}
current: Final = {"model": "current", "tier": "REASONING"}
newer: Final = {"model": "newer", "tier": "COMPLEX"}
try:
redis.redis_client.set(key, json.dumps(current), ex=10)
redis.redis_client.set(other_key, json.dumps("other-value"), ex=60)
assert await read_affinity_pin(cache, "pin", 60) == current
cache.in_memory_cache.set_cache("pin", stale, ttl=60)
assert cache.in_memory_cache.get_cache("pin") == stale
if mode == "script_error":
redis.redis_client.delete(key)
redis.redis_client.rpush(key, "wrong-type")
start: Final = len(sends)
with request_redis_batch_scope() as request:
batch: Final = request.batch(redis)
pending: Final = batch.mget(("other",))
if mode == "separate":
assert await pending == {"other": "other-value"}
assert await read_affinity_pin(cache, "pin", 60) == (stale if mode == "script_error" else current)
assert await pending == {"other": "other-value"}
wire: Final = tuple(data for data in sends[start:] if b"MGET" in data or b"EVALSHA" in data)
assert len(wire) == (2 if mode == "separate" else 1)
assert sum(b"MGET" in data for data in wire) == sum(b"EVALSHA" in data for data in wire) == 1
if mode == "script_error":
assert redis.redis_client.type(key) == b"list"
redis.redis_client.delete(key)
else:
assert 10 < redis.redis_client.ttl(key) <= 60
redis.redis_client.set(key, json.dumps(newer), ex=10)
later_pending: Final = batch.mget(("other",))
assert await read_affinity_pin(cache, "pin", 60) == newer
assert await later_pending == {"other": "other-value"}
assert cache.in_memory_cache.get_cache("pin") == newer
assert 10 < redis.redis_client.ttl(key) <= 60
assert redis.redis_client.get(key) == json.dumps(newer).encode()
redis.redis_client.delete(key)
assert await read_affinity_pin(cache, "pin", 60) is None
assert cache.in_memory_cache.get_cache("pin") is None
assert redis.redis_client.exists(key) == 0
finally:
redis.redis_client.delete(key, other_key)
for cached_key in tuple(litellm.in_memory_llm_clients_cache.cache_dict):
if client_key in cached_key:
litellm.in_memory_llm_clients_cache.delete_cache(cached_key)
await client.aclose()
await pool.aclose()
await redis.disconnect()
redis.redis_client.close()

395
tests/integration/sdk/test_router_session_pins_redis.py Normal file
View file

@ -0,0 +1,395 @@
import asyncio
import json
import os
import uuid
from collections.abc import Iterator, Mapping
from typing import Final
import pytest
from litellm.caching.affinity_cache import claim_affinity_pin, read_affinity_pin
from litellm.caching.redis_batch import request_redis_batch_scope
from litellm.caching.redis_cache import RedisCache
from litellm.redis.stores.connection import connect
from litellm.redis.stores.router_session_pins import (
DeploymentPin,
DeploymentSessionPin,
PinKey,
RouterSessionPinStore,
TierModelPin,
TierSessionPin,
)
from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck
_DEPLOYMENTS: Final = ("original", "replacement", "other")
@pytest.fixture
def redis() -> Iterator[RedisCache]:
cache: Final = RedisCache(host=os.environ["REDIS_HOST"], port=int(os.environ["REDIS_PORT"]))
yield cache
cache.redis_client.close()
@pytest.fixture
def session_id() -> str:
return uuid.uuid4().hex
def _pins(redis: RedisCache) -> RouterSessionPinStore:
return RouterSessionPinStore(connect(redis).router_session_pins())
def _session_key(session_id: str) -> DeploymentSessionPin:
return DeploymentSessionPin(model_group="repin-group", caller=None, session_id=session_id)
def _spelled(session_id: str) -> str:
return f"deployment_affinity:v1:session:repin-group:unscoped:{session_id}"
def _repin_check(pins: RouterSessionPinStore, enabled: bool = True) -> DeploymentAffinityCheck:
return DeploymentAffinityCheck(
pins=pins,
ttl_seconds=3600,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
enable_session_id_affinity=True,
session_affinity_repin_on_failover=enabled,
)
async def _selected(callback: DeploymentAffinityCheck, session_id: str) -> tuple[str, ...]:
selected: Final = await callback.async_filter_deployments(
model="repin-group",
healthy_deployments=[
{"model_name": "repin-group", "litellm_params": {"model": "openai/m"}, "model_info": {"id": deployment}}
for deployment in _DEPLOYMENTS
],
messages=None,
request_kwargs={"metadata": {"session_id": session_id}},
)
return tuple(deployment["model_info"]["id"] for deployment in selected)
def _break(redis: RedisCache, key: str) -> None:
redis.redis_client.delete(key)
redis.redis_client.rpush(key, "wrong-type")
# Keys main's Python spelled before they moved into the native store
@pytest.mark.asyncio
@pytest.mark.parametrize(
("key", "spelled"),
[
(
DeploymentPin(model_group="gpt-4o", caller="sk-caller"),
"deployment_affinity:v1:gpt-4o:8b6712790a2089c67aa97a2d80022df18cc65c7814350e33baebe79aab508891",
),
(
DeploymentSessionPin(model_group="gpt-4o", caller=None, session_id="sess-1"),
"deployment_affinity:v1:session:gpt-4o:unscoped:sess-1",
),
(
TierSessionPin(router="smart", caller="key-hash", session_id="sess-1"),
"complexity_router_session_affinity:v1:smart:key-hash:sess-1",
),
(
TierModelPin(
router="smart",
tiers="cfg",
caller=None,
session_id="séss",
tier="COMPLEX",
replay_fallback="m-a",
),
"complexity_router_tier_model_affinity:v1:c83437032b0b0d7d8229f7aaf444096a83510ee630e5a510e4770b0f6f03a8d3",
),
],
)
async def test_a_pin_key_reads_the_key_main_wrote(redis: RedisCache, key: PinKey, spelled: str) -> None:
try:
redis.redis_client.set(spelled, json.dumps({"model_id": "main-pod"}), ex=60)
assert await _pins(redis).get_shared(key) == {"model_id": "main-pod"}
finally:
redis.redis_client.delete(spelled)
@pytest.mark.asyncio
@pytest.mark.parametrize("enabled", [False, True])
async def test_session_selection_observes_another_workers_repin_only_when_enabled(
redis: RedisCache, session_id: str, enabled: bool
) -> None:
pins: Final = _pins(redis)
try:
pins.set_local(_session_key(session_id), {"model_id": "original"}, 3600)
await _pins(redis).set(_session_key(session_id), {"model_id": "replacement"}, 3600)
assert await _selected(_repin_check(pins, enabled), session_id) == (
"replacement" if enabled else "original",
)
finally:
redis.redis_client.delete(_spelled(session_id))
@pytest.mark.asyncio
@pytest.mark.parametrize("has_pin", [False, True])
async def test_redis_read_failures_preserve_session_routing(
redis: RedisCache, session_id: str, has_pin: bool
) -> None:
pins: Final = _pins(redis)
callback: Final = _repin_check(pins)
try:
if has_pin:
redis.redis_client.set(_spelled(session_id), json.dumps({"model_id": "replacement"}), ex=60)
pins.set_local(_session_key(session_id), {"model_id": "original"}, 3600)
assert await _selected(callback, session_id) == ("replacement",)
_break(redis, _spelled(session_id))
assert set(await _selected(callback, session_id)) == ({"replacement"} if has_pin else set(_DEPLOYMENTS))
finally:
redis.redis_client.delete(_spelled(session_id))
@pytest.mark.asyncio
async def test_confirmed_redis_miss_does_not_revive_a_local_pin_during_outage(
redis: RedisCache, session_id: str
) -> None:
pins: Final = _pins(redis)
callback: Final = _repin_check(pins)
try:
pins.set_local(_session_key(session_id), {"model_id": "original"}, 3600)
assert set(await _selected(callback, session_id)) == set(_DEPLOYMENTS)
_break(redis, _spelled(session_id))
assert set(await _selected(callback, session_id)) == set(_DEPLOYMENTS)
finally:
redis.redis_client.delete(_spelled(session_id))
@pytest.mark.asyncio
async def test_shared_session_pin_resumes_after_a_transient_redis_failure(redis: RedisCache, session_id: str) -> None:
pins: Final = _pins(redis)
callback: Final = _repin_check(pins)
try:
pins.set_local(_session_key(session_id), {"model_id": "original"}, 3600)
_break(redis, _spelled(session_id))
assert await _selected(callback, session_id) == ("original",)
redis.redis_client.delete(_spelled(session_id))
redis.redis_client.set(_spelled(session_id), json.dumps({"model_id": "replacement"}), ex=60)
assert await _selected(callback, session_id) == ("replacement",)
finally:
redis.redis_client.delete(_spelled(session_id))
@pytest.mark.asyncio
@pytest.mark.parametrize("shared_value", [None, {"model_id": "original"}])
async def test_a_shared_pin_read_in_flight_keeps_a_newer_local_write(
redis: RedisCache, session_id: str, shared_value: Mapping[str, str] | None
) -> None:
pins: Final = _pins(redis)
callback: Final = _repin_check(pins)
try:
if shared_value is not None:
redis.redis_client.set(_spelled(session_id), json.dumps(shared_value), ex=60)
pending: Final = asyncio.create_task(_selected(callback, session_id))
await asyncio.sleep(0)
assert not pending.done()
pins.set_local(_session_key(session_id), {"model_id": "replacement"}, 3600)
await pending
_break(redis, _spelled(session_id))
assert await _selected(callback, session_id) == ("replacement",)
finally:
redis.redis_client.delete(_spelled(session_id))
@pytest.mark.asyncio
async def test_a_claim_uses_redis_attached_after_construction(redis: RedisCache, session_id: str) -> None:
"""The proxy attaches Redis after the Router and this callback are built, so a claim resolves Redis per call."""
pins: Final = RouterSessionPinStore(connect(None).router_session_pins())
callback: Final = _repin_check(pins)
try:
redis.redis_client.set(_spelled(session_id), json.dumps({"model_id": "other-pod-winner"}), ex=60)
pins.attach_redis(connect(redis))
claimed: Final = await callback._claim_pin(_session_key(session_id), {"model_id": "ours"}, 777)
assert claimed == "other-pod-winner"
assert pins.get_local(_session_key(session_id)) == {"model_id": "other-pod-winner"}
finally:
redis.redis_client.delete(_spelled(session_id))
@pytest.mark.asyncio
async def test_a_claim_falls_back_to_pod_local_when_redis_fails(redis: RedisCache, session_id: str) -> None:
"""A Redis outage costs cross-pod agreement, never same-pod stickiness."""
pins: Final = _pins(redis)
callback: Final = _repin_check(pins)
try:
_break(redis, _spelled(session_id))
assert await callback._claim_pin(_session_key(session_id), {"model_id": "ours"}, 777) == "ours"
assert pins.get_local(_session_key(session_id)) == {"model_id": "ours"}
assert await callback._claim_pin(_session_key(session_id), {"model_id": "another"}, 777) == "ours"
finally:
redis.redis_client.delete(_spelled(session_id))
@pytest.mark.asyncio
@pytest.mark.parametrize(
("held", "expected", "stored"),
[
('{"model_id": "winner"}', "winner", {"model_id": "winner"}),
('"winner"', "winner", "winner"),
("winner", "winner", "winner"),
('{"model": "winner"}', None, {"model": "winner"}),
(None, "candidate", {"model_id": "candidate"}),
],
)
async def test_a_deployment_claim_reads_the_winner_in_any_shape_main_wrote(
redis: RedisCache, session_id: str, held: str | None, expected: str | None, stored: object
) -> None:
pins: Final = _pins(redis)
try:
if held is not None:
redis.redis_client.set(_spelled(session_id), held, ex=60)
winner: Final = await _repin_check(pins)._claim_pin(_session_key(session_id), {"model_id": "candidate"}, 30)
assert winner == expected
assert pins.get_local(_session_key(session_id)) == stored
finally:
redis.redis_client.delete(_spelled(session_id))
@pytest.mark.asyncio
async def test_concurrent_session_repin_claims_keep_the_shared_winner_after_recovery(
redis: RedisCache, session_id: str
) -> None:
workers: Final = (_pins(redis), _pins(redis))
key: Final = _session_key(session_id)
choices: Final = ({"model_id": "replacement"}, {"model_id": "other"})
try:
redis.redis_client.set(_spelled(session_id), json.dumps({"model_id": "original"}), ex=60)
for worker in workers:
worker.set_local(key, {"model_id": "original"}, 60)
winners: Final = await asyncio.gather(
*(
claim_affinity_pin(worker, key, choice, 60, eligible_values=choices, preserve_legacy_model_ids=True)
for worker, choice in zip(workers, choices, strict=True)
)
)
assert winners[0] == winners[1] and winners[0] in choices
late_original: Final = await claim_affinity_pin(
workers[0],
key,
{"model_id": "original"},
60,
eligible_values=({"model_id": "original"}, *choices),
preserve_legacy_model_ids=True,
)
assert late_original == winners[0]
assert json.loads(redis.redis_client.get(_spelled(session_id))) == winners[0]
finally:
redis.redis_client.delete(_spelled(session_id))
@pytest.mark.asyncio
@pytest.mark.parametrize("stored", ["original", '"original"', '{"model_id":"original"}'])
async def test_session_repin_lua_preserves_healthy_legacy_pins_and_replaces_unavailable_ones(
redis: RedisCache, session_id: str, stored: str
) -> None:
pins: Final = _pins(redis)
key: Final = _session_key(session_id)
replacement: Final = {"model_id": "replacement"}
try:
redis.redis_client.set(_spelled(session_id), stored, ex=60)
preserved: Final = await claim_affinity_pin(
pins,
key,
replacement,
60,
eligible_values=({"model_id": "original"}, replacement),
preserve_legacy_model_ids=True,
)
assert preserved in ("original", {"model_id": "original"})
moved: Final = await claim_affinity_pin(
pins, key, replacement, 60, eligible_values=(replacement,), preserve_legacy_model_ids=True
)
assert moved == replacement
assert json.loads(redis.redis_client.get(_spelled(session_id))) == replacement
finally:
redis.redis_client.delete(_spelled(session_id))
def _tier_key(session_id: str) -> TierSessionPin:
return TierSessionPin(router="smart", caller=None, session_id=session_id)
def _tier_spelled(session_id: str) -> str:
return f"complexity_router_session_affinity:v1:smart:unscoped:{session_id}"
@pytest.mark.asyncio
@pytest.mark.parametrize("stored", ['{"model":"current","tier":"REASONING"}', '"current"', "current"])
async def test_tier_pin_reads_shared_state_refreshes_ttl_and_forgets_deleted_pins(
redis: RedisCache, session_id: str, stored: str
) -> None:
workers: Final = (_pins(redis), _pins(redis))
key: Final = _tier_key(session_id)
spelled: Final = _tier_spelled(session_id)
try:
for worker in workers:
worker.set_local(key, {"model": "old", "tier": "SIMPLE"}, 60)
redis.redis_client.set(spelled, stored, ex=10)
first: Final = await read_affinity_pin(workers[0], key, 60)
second: Final = await read_affinity_pin(workers[1], key, 60)
assert first == second and first in ("current", {"model": "current", "tier": "REASONING"})
assert all(worker.get_local(key) == first for worker in workers)
assert 10 < redis.redis_client.ttl(spelled) <= 60
replacement: Final = {"model": "newer", "tier": "SIMPLE"}
redis.redis_client.set(spelled, json.dumps(replacement), ex=10)
assert await read_affinity_pin(workers[0], key, 60) == replacement
assert await read_affinity_pin(workers[1], key, 60) == replacement
redis.redis_client.delete(spelled)
for worker in workers:
assert await read_affinity_pin(worker, key, 60) is None
assert worker.get_local(key) is None
assert redis.redis_client.exists(spelled) == 0
finally:
redis.redis_client.delete(spelled)
@pytest.mark.asyncio
@pytest.mark.parametrize("script_error", [False, True])
async def test_tier_pin_read_inside_a_request_follows_redis_and_keeps_memory_when_the_script_fails(
redis: RedisCache, session_id: str, script_error: bool
) -> None:
pins: Final = _pins(redis)
key: Final = _tier_key(session_id)
spelled: Final = _tier_spelled(session_id)
stale: Final = {"model": "old", "tier": "SIMPLE"}
current: Final = {"model": "current", "tier": "REASONING"}
newer: Final = {"model": "newer", "tier": "COMPLEX"}
try:
redis.redis_client.set(spelled, json.dumps(current), ex=10)
assert await read_affinity_pin(pins, key, 60) == current
pins.set_local(key, stale, 60)
if script_error:
_break(redis, spelled)
with request_redis_batch_scope():
assert await read_affinity_pin(pins, key, 60) == (stale if script_error else current)
if script_error:
assert redis.redis_client.type(spelled) == b"list"
redis.redis_client.delete(spelled)
else:
assert 10 < redis.redis_client.ttl(spelled) <= 60
redis.redis_client.set(spelled, json.dumps(newer), ex=10)
assert await read_affinity_pin(pins, key, 60) == newer
assert pins.get_local(key) == newer
assert 10 < redis.redis_client.ttl(spelled) <= 60
redis.redis_client.delete(spelled)
assert await read_affinity_pin(pins, key, 60) is None
assert pins.get_local(key) is None
finally:
redis.redis_client.delete(spelled)

13
tests/unit/redis/stores/router_session_pins_memory.py Normal file
View file

@ -0,0 +1,13 @@
"""The real native session pin store without Redis, optionally on a clock the test moves, as ``InMemoryCache``'s
``clock`` argument does on main."""
from __future__ import annotations
from collections.abc import Callable
from litellm.redis.stores.router_session_pins import RouterSessionPinStore
from litellm.rust_bridge._native import NativeRouterSessionPins
def memory_pins(clock: Callable[[], float] | None = None) -> RouterSessionPinStore:
return RouterSessionPinStore(NativeRouterSessionPins(None, clock))

116
tests/unit/redis/stores/test_router_session_pins.py Normal file
View file

@ -0,0 +1,116 @@
import socket
from collections.abc import Awaitable, Callable, Iterator
from typing import Final
import pytest
from litellm.caching.affinity_cache import claim_affinity_pin, read_affinity_pin
from litellm.caching.redis_batch import Batch
from litellm.caching.redis_cache import RedisCache
from litellm.redis.stores.connection import RedisStores, connect
from litellm.redis.stores.router_session_pins import DeploymentSessionPin, RouterSessionPinStore
from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck
from litellm.rust_bridge._native import NativeRedisStoreError
from tests.unit.redis.stores.router_session_pins_memory import memory_pins
_KEY: Final = DeploymentSessionPin(model_group="group", caller=None, session_id="session")
@pytest.fixture
def unreachable_redis() -> Iterator[RedisStores]:
"""A real native client to a port nothing listens on, so every Redis call is refused at once."""
with socket.socket() as probe:
probe.bind(("127.0.0.1", 0))
port: Final = probe.getsockname()[1]
yield connect(RedisCache(host="127.0.0.1", port=port))
@pytest.fixture
def down_pins(unreachable_redis: RedisStores) -> RouterSessionPinStore:
return RouterSessionPinStore(unreachable_redis.router_session_pins())
async def test_attaching_redis_keeps_the_memory_pins(unreachable_redis: RedisStores) -> None:
pins: Final = memory_pins()
pins.set_local(_KEY, {"model_id": "dep-a"}, 60)
pins.attach_redis(unreachable_redis)
assert pins.has_redis
assert await pins.get(_KEY) == {"model_id": "dep-a"}
async def test_lookups_and_writes_serve_memory_while_redis_is_down(down_pins: RouterSessionPinStore) -> None:
assert await down_pins.get(_KEY) is None
await down_pins.set(_KEY, {"model_id": "dep-a"}, 60)
assert await down_pins.get(_KEY) == {"model_id": "dep-a"}
down_pins.delete_local(_KEY)
assert down_pins.get_local(_KEY) is None
@pytest.mark.parametrize(
"op",
[
lambda pins: pins.get_shared(_KEY),
lambda pins: pins.read(_KEY, 60, batch=Batch.NONE),
lambda pins: pins.claim(_KEY, {"model_id": "dep-a"}, 60, eligible_values=None, preserve_legacy_model_ids=False),
],
ids=["get_shared", "read", "claim"],
)
async def test_redis_only_ops_raise_while_redis_is_down(
down_pins: RouterSessionPinStore, op: Callable[[RouterSessionPinStore], Awaitable[object]]
) -> None:
with pytest.raises(NativeRedisStoreError, match="Redis did not answer"):
await op(down_pins)
async def test_a_read_without_redis_answers_from_memory() -> None:
pins: Final = memory_pins()
pins.set_local(_KEY, {"model_id": "dep-a"}, 60)
assert await pins.read(_KEY, 60, batch=Batch.PRE_CALL) == {"model_id": "dep-a"}
async def test_a_failed_pin_read_keeps_the_pod_local_pin(down_pins: RouterSessionPinStore) -> None:
down_pins.set_local(_KEY, {"model": "m", "tier": "SIMPLE"}, 60)
assert await read_affinity_pin(down_pins, _KEY, 60) == {"model": "m", "tier": "SIMPLE"}
assert await read_affinity_pin(down_pins, DeploymentSessionPin("group", None, "other"), 60) is None
async def test_a_failed_claim_falls_back_to_a_pod_local_first_writer(down_pins: RouterSessionPinStore) -> None:
assert await claim_affinity_pin(down_pins, _KEY, {"model_id": "dep-a"}, 60) == {"model_id": "dep-a"}
assert await claim_affinity_pin(down_pins, _KEY, {"model_id": "dep-b"}, 60) == {"model_id": "dep-a"}
@pytest.mark.parametrize("local", [None, "dep-a"])
async def test_session_routing_keeps_the_local_pin_when_the_shared_read_fails(
down_pins: RouterSessionPinStore, local: str | None
) -> None:
callback: Final = DeploymentAffinityCheck(
pins=down_pins,
ttl_seconds=60,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
enable_session_id_affinity=True,
session_affinity_repin_on_failover=True,
)
if local is not None:
down_pins.set_local(_KEY, {"model_id": local}, 60)
deployments: Final = [
{"model_name": "group", "litellm_params": {"model": "openai/m"}, "model_info": {"id": deployment}}
for deployment in ("dep-a", "dep-b")
]
selected: Final = await callback.async_filter_deployments(
model="group",
healthy_deployments=deployments,
messages=None,
request_kwargs={"metadata": {"session_id": "session"}},
)
assert [deployment["model_info"]["id"] for deployment in selected] == (
["dep-a"] if local is not None else ["dep-a", "dep-b"]
)

401
tests/unit/router_strategy/test_complexity_router.py
View file

@ -25,8 +25,7 @@ import litellm
from litellm import Router
from litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache
from litellm.redis.stores.router_session_pins import RouterSessionPinStore
from litellm.constants import (
CLIENT_OUTPUT_CEILING_METADATA_KEY,
OUTPUT_TOKEN_CEILING_PARAMS,
@ -36,6 +35,7 @@ from litellm.constants import (
from litellm.integrations.custom_logger import CustomLogger
from litellm.router import as_output_cap, stamped_client_ceiling
from litellm.router_api import RoutingContext
from tests.unit.redis.stores.router_session_pins_memory import memory_pins
from litellm.router_strategy.auto_router.auto_router_model_naming import (
CUSTOMIZATION_CAPABILITY,
GATED_AUTO_ROUTER_CAPABILITIES,
@ -3312,7 +3312,7 @@ class TestCapabilityClassifier:
bypass: Literal["literal_keyword_match", "session_affinity_pin", "housekeeping"],
) -> None:
mock_router_instance.acompletion = AsyncMock(return_value=_llm_response(_capability_reply(p_solve=0.8)))
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
router: Final = self._router(
mock_router_instance,
session_affinity=bypass == "session_affinity_pin",
@ -6720,7 +6720,7 @@ class TestTierModelAffinity:
deployment_affinity: bool = True,
plugins: bool = False,
) -> ComplexityRouter:
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
_answer_deployments_named(mock_router_instance, [])
return ComplexityRouter(
model_name="affinity-router",
@ -6865,8 +6865,8 @@ class TestTierModelAffinity:
)
return [] if unavailable else [{"model_name": model, "model_info": {"id": f"deployment-{model}"}}]
cache: Final = DualCache()
mock_router_instance.health.return_value.store = cache
cache: Final = memory_pins()
mock_router_instance.health.return_value.session_pins = cache
mock_router_instance.select = select
router: Final = TestModalityRouting._router(
mock_router_instance,
@ -6901,7 +6901,7 @@ class TestTierModelAffinity:
else ("modality_pin_override" if session_affinity else "user_turn_continuation")
)
cache_key: Final = router._get_session_affinity_cache_key("replay-session", {"metadata": metadata})
assert await cache.async_get_cache(cache_key) == {"model": "model-a", "tier": "SIMPLE"}
assert await cache.get(cache_key) == {"model": "model-a", "tier": "SIMPLE"}
next_ask: Final = await self._route(router, metadata, "model-b")
assert next_ask.model == "model-a"
@ -6915,7 +6915,7 @@ class TestTierModelAffinity:
self, mock_router_instance: MagicMock
) -> None:
clock: Final = MagicMock(return_value=100.0)
mock_router_instance.health.return_value.store = DualCache(in_memory_cache=InMemoryCache(clock=clock))
mock_router_instance.health.return_value.session_pins = memory_pins(clock)
router: Final = ComplexityRouter(
model_name="affinity-router",
ctx=RoutingContext(mock_router_instance),
@ -6951,8 +6951,8 @@ class TestTierModelAffinity:
async def test_session_escalation_keeps_the_selected_tier_when_models_overlap(
self, mock_router_instance: MagicMock
) -> None:
cache: Final = DualCache()
mock_router_instance.health.return_value.store = cache
cache: Final = memory_pins()
mock_router_instance.health.return_value.session_pins = cache
router: Final = ComplexityRouter(
model_name="affinity-router",
ctx=RoutingContext(mock_router_instance),
@ -6976,14 +6976,14 @@ class TestTierModelAffinity:
metadata: Final = {"session_id": "same-session"}
assert (await self._route(router, metadata, "higher", "visit_complex")).model == "higher"
cache_key: Final = router._get_session_affinity_cache_key("same-session", {"metadata": metadata})
await cache.async_set_cache(cache_key, {"model": "base", "tier": "SIMPLE"}, ttl=600)
await cache.set(cache_key, {"model": "base", "tier": "SIMPLE"}, 600)
result: Final = await self._route(router, metadata, "shared", "LITELLM ESCALATE")
assert result.model == "shared"
assert result.routing_decision["tier"] == "MEDIUM"
assert result.routing_decision["cause"] == "session_affinity_escalation"
assert result.litellm_params == {"temperature": 0.4}
assert await cache.async_get_cache(cache_key) == {"model": "shared", "tier": "MEDIUM"}
assert await cache.get(cache_key) == {"model": "shared", "tier": "MEDIUM"}
@pytest.mark.asyncio
@pytest.mark.parametrize(
@ -7004,8 +7004,8 @@ class TestTierModelAffinity:
expected_model: str,
expected_tier: str,
) -> None:
cache: Final = DualCache()
mock_router_instance.health.return_value.store = cache
cache: Final = memory_pins()
mock_router_instance.health.return_value.session_pins = cache
router: Final = ComplexityRouter(
model_name="affinity-router",
ctx=RoutingContext(mock_router_instance),
@ -7016,21 +7016,21 @@ class TestTierModelAffinity:
)
metadata: Final = {"session_id": "same-session"}
cache_key: Final = router._get_session_affinity_cache_key("same-session", {"metadata": metadata})
await cache.async_set_cache(cache_key, {"model": "model-a", "tier": stale_tier}, ttl=600)
await cache.set(cache_key, {"model": "model-a", "tier": stale_tier}, 600)
result: Final = await self._route(router, metadata, expected_model, prompt)
assert result.model == expected_model
assert result.routing_decision["tier"] == expected_tier
assert result.routing_decision["cause"] == "heuristic_scorer"
assert await cache.async_get_cache(cache_key) == {"model": expected_model, "tier": expected_tier}
assert await cache.get(cache_key) == {"model": expected_model, "tier": expected_tier}
@pytest.mark.asyncio
@pytest.mark.parametrize("classification_mode", ["every_request", "user_turn"])
async def test_custom_tier_keeps_its_own_model(
self, mock_router_instance: MagicMock, classification_mode: Literal["every_request", "user_turn"]
) -> None:
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
router: Final = ComplexityRouter(
model_name="affinity-router",
ctx=RoutingContext(mock_router_instance),
@ -7099,7 +7099,7 @@ class TestSessionAffinity:
"""The hook response's session_affinity_ttl_seconds is what the Router stamps as
the deployment-affinity marker, so both the classify path (turn 1) and the
session-pin path (turn 2) must carry the configured TTL."""
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
router = ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
@ -7139,7 +7139,7 @@ class TestSessionAffinity:
):
"""Deployment affinity retains a model per tier while classification continues.
Session affinity keeps the first tier too; plugins suppress both affinity policies."""
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
router = ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
@ -7166,7 +7166,7 @@ class TestSessionAffinity:
async def test_hook_response_has_no_session_affinity_ttl_when_disabled_or_plugins(
self, mock_router_instance, basic_config, session_affinity_config
):
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
disabled_router = ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
@ -7190,7 +7190,7 @@ class TestSessionAffinity:
async def test_disabled_by_default_reclassifies_every_turn(self, mock_router_instance, basic_config):
"""With session_affinity off, a shared session can move from REASONING to SIMPLE."""
assert "session_affinity" not in basic_config
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
router = ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
@ -7210,7 +7210,7 @@ class TestSessionAffinity:
async def test_proxy_generated_session_id_never_pins(self, mock_router_instance, session_affinity_config):
"""A session id the proxy generated for a request that had none is per request, so
it must not create a pin even with session_affinity enabled."""
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
router = ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
@ -7230,7 +7230,7 @@ class TestSessionAffinity:
async def test_can_be_enabled_to_pin_every_later_turn(self, mock_router_instance, session_affinity_config):
"""Regression: session_affinity=True is the opt-in, so a shared session_id reuses the
first turn's model instead of reclassifying."""
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
router = ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
@ -7248,7 +7248,7 @@ class TestSessionAffinity:
@pytest.mark.asyncio
async def test_pins_model_after_first_turn(self, mock_router_instance, session_affinity_config):
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
router = ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
@ -7274,7 +7274,7 @@ class TestSessionAffinity:
TTL, so a session whose only turn landed on the cooldown fallback must classify again once
the breaker closes instead of holding that fallback's model."""
now = 100.0
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
mock_router_instance.acompletion = AsyncMock(
side_effect=[TimeoutError("classifier timed out"), _llm_response('{"tier": "REASONING"}')]
)
@ -7309,7 +7309,7 @@ class TestSessionAffinity:
@pytest.mark.asyncio
async def test_a_pinned_turn_reports_the_tier_that_serves_it(self, mock_router_instance, session_affinity_config):
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
router = ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
@ -7326,7 +7326,7 @@ class TestSessionAffinity:
@pytest.mark.asyncio
async def test_different_sessions_classify_independently(self, mock_router_instance, session_affinity_config):
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
router = ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
@ -7347,8 +7347,8 @@ class TestSessionAffinity:
self, mock_router_instance: MagicMock, basic_config: dict[str, object], preexisting: bool
) -> None:
clock: Final = MagicMock(return_value=100.0)
cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock))
mock_router_instance.health.return_value.store = cache
cache: Final = memory_pins(clock)
mock_router_instance.health.return_value.session_pins = cache
router: Final = ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
@ -7361,7 +7361,7 @@ class TestSessionAffinity:
request_kwargs: Final = self._request_kwargs("session-1")
key: Final = router._get_session_affinity_cache_key("session-1", request_kwargs)
if preexisting:
cache.in_memory_cache.set_cache(key, "o1-preview", ttl=90)
cache.set_local(key, "o1-preview", 90)
first: Final = await router.async_pre_routing_hook(
model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE
)
@ -7372,145 +7372,17 @@ class TestSessionAffinity:
)
assert replay.model == first.model
clock.return_value = 200.0
assert cache.in_memory_cache.get_cache(key) == (
assert cache.get_local(key) == (
"o1-preview" if preexisting else {"model": "gpt-4o-mini", "tier": "SIMPLE"}
)
clock.return_value = 271.0
assert cache.in_memory_cache.get_cache(key) is None
@pytest.mark.asyncio
@pytest.mark.parametrize("batch_pending", [False, True])
@pytest.mark.parametrize("remote_state", ["latest", "interleaved", "outage"])
async def test_continuation_follows_redis_and_never_republishes_a_stale_pin(
self, mock_router_instance: MagicMock, basic_config: dict[str, object], remote_state: str, batch_pending: bool
) -> None:
from fakeredis import FakeRedis
from redis.asyncio import BlockingConnectionPool, Connection, Redis
from redis.asyncio.retry import Retry
from redis.backoff import NoBackoff
from litellm.caching.redis_batch import request_redis_batch_scope
clock: Final = MagicMock(return_value=100.0)
simple: Final = {"model": "gpt-4o-mini", "tier": "SIMPLE"}
latest: Final = {"model": "o1-preview", "tier": "REASONING"}
stored: Final = {"pin": json.dumps(simple if remote_state == "interleaved" else latest)}
available: Final = asyncio.Event()
read_started: Final = asyncio.Event()
release_read: Final = asyncio.Event()
if remote_state != "outage":
available.set()
if remote_state != "interleaved":
release_read.set()
async def execute(*args: object, **kwargs: object) -> object:
if args[0] in ("EVALSHA", "GET", "MGET"):
if not available.is_set():
raise ConnectionError("Redis unavailable")
if args[0] == "MGET":
return [json.dumps("other-value")]
observed: Final = stored["pin"]
read_started.set()
await release_read.wait()
return observed
if args[0] == "SET":
stored["pin"] = str(args[2])
return b"PONG" if args[0] == "PING" else b"OK"
pack_commands: Final = MagicMock(side_effect=tuple)
async def connection() -> MagicMock:
replies: Final[asyncio.Queue[object]] = asyncio.Queue()
async def send_command(*args: object) -> None:
await replies.put(await execute(*args))
async def send_packed_command(commands: Sequence[Sequence[object]]) -> None:
for command in commands:
await send_command(*command)
return MagicMock(
spec=Connection,
retry=Retry(NoBackoff(), 0),
send_command=AsyncMock(side_effect=send_command),
pack_commands=pack_commands,
send_packed_command=AsyncMock(side_effect=send_packed_command),
read_response=AsyncMock(side_effect=replies.get),
)
sync_client: Final = FakeRedis()
pool: Final = MagicMock(
spec=BlockingConnectionPool,
connection_kwargs={"protocol": 2},
get_encoder=sync_client.get_encoder,
get_connection=AsyncMock(side_effect=connection),
release=AsyncMock(),
)
client: Final = Redis(connection_pool=pool)
redis_cache: Final = RedisCache(host="tier-pin-fixture", connection_pool=sync_client.connection_pool)
client_key: Final = redis_cache._get_async_client_cache_key()
litellm.in_memory_llm_clients_cache.set_cache(key=client_key, value=client)
try:
cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock), redis_cache=redis_cache)
mock_router_instance.health.return_value.store = cache
router: Final = ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
complexity_router_config={
**basic_config,
"classification_mode": "user_turn",
"deployment_affinity": False,
"session_affinity_ttl_seconds": 10,
},
)
request: Final = self._request_kwargs("shared-session")
key: Final = router._get_session_affinity_cache_key("shared-session", request)
cache.in_memory_cache.set_cache(key, simple, ttl=10)
clock.return_value = 105.0
messages: Final = [*self.SIMPLE_MESSAGE, {"role": "tool", "content": "done"}]
with request_redis_batch_scope() as scope:
batch: Final = scope.batch(redis_cache)
unrelated: Final = batch.mget(("other",)) if batch_pending else None
pending: Final = asyncio.create_task(router.async_pre_routing_hook("m", request, messages))
if remote_state == "interleaved":
await asyncio.wait_for(read_started.wait(), timeout=1)
stored["pin"] = json.dumps(latest)
release_read.set()
result: Final = await pending
if unrelated is not None:
if remote_state == "outage":
with pytest.raises(ConnectionError, match="Redis unavailable"):
await unrelated
else:
assert await unrelated == {"other": "other-value"}
assert result.model == (latest["model"] if remote_state == "latest" else simple["model"])
assert result.routing_decision["cause"] == "user_turn_continuation"
assert json.loads(stored["pin"]) == latest
clock.return_value = 111.0
assert cache.in_memory_cache.get_cache(key) is not None
available.set()
recovery_read: Final = batch.mget(("other",)) if batch_pending else None
recovered: Final = await router.async_pre_routing_hook("m", request, messages)
if recovery_read is not None:
assert await recovery_read == {"other": "other-value"}
assert recovered.model == latest["model"]
assert cache.in_memory_cache.get_cache(key) == latest
assert read_started.is_set()
assert tuple(
tuple(args[0] for args in sent.args[0]) for sent in pack_commands.call_args_list
) == (("MGET", "EVALSHA"),) * (2 if batch_pending else 0)
finally:
for cached_key in tuple(litellm.in_memory_llm_clients_cache.cache_dict):
if client_key in cached_key:
litellm.in_memory_llm_clients_cache.delete_cache(cached_key)
await client.aclose()
sync_client.close()
assert cache.get_local(key) is None
@pytest.mark.asyncio
async def test_different_api_keys_do_not_share_pin(self, mock_router_instance, session_affinity_config):
"""A session_id is client-supplied and unauthenticated; two different callers
(API keys) reusing the same session_id must not poison each other's pin."""
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
router = ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
@ -7533,8 +7405,8 @@ class TestSessionAffinity:
@pytest.mark.asyncio
async def test_no_session_id_falls_back_to_reclassify(self, mock_router_instance, session_affinity_config):
cache = AsyncMock()
mock_router_instance.health.return_value.store = cache
pins: Final = MagicMock(spec=RouterSessionPinStore)
mock_router_instance.health.return_value.session_pins = pins
router = ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
@ -7544,15 +7416,14 @@ class TestSessionAffinity:
model="test-model", request_kwargs={}, messages=self.SIMPLE_MESSAGE
)
assert result.model == "gpt-4o-mini"
cache.async_get_cache.assert_not_called()
cache.async_set_cache.assert_not_called()
assert pins.method_calls == []
@pytest.mark.asyncio
async def test_adaptive_pinned_turn_still_stamps_chosen_model_metadata(self, mock_router_instance):
"""Regression: skipping classification on a pinned turn must not break the
adaptive bandit's reward-feedback loop, which only records a turn's outcome
when ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY is present in the request metadata."""
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
_answer_deployments_named(
mock_router_instance,
[
@ -7961,7 +7832,7 @@ class TestClassificationMode:
return {"metadata": {"session_id": session_id}}
def _router(self, mock_router_instance, config: dict) -> ComplexityRouter:
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
return ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
@ -8303,7 +8174,7 @@ class TestUserTurnClassifierFailureHoldsTheDefaultModel:
*classifier_replies: object, default_model: str = "gpt-4o", plan_mode_min_tier: str | None = None
) -> tuple[ComplexityRouter, MagicMock]:
router_instance: Final = MagicMock()
router_instance.health.return_value.store = DualCache()
router_instance.health.return_value.session_pins = memory_pins()
router_instance.acompletion = AsyncMock(side_effect=list(classifier_replies))
router: Final = ComplexityRouter(
model_name="test-router",
@ -8619,7 +8490,7 @@ class TestRoutingPlugins:
(e.g. a budget cap crossed mid-session) was only ever enforced on a session's
first turn. With plugins configured, every turn must go through
_classify_and_route (and therefore the plugin pipeline) again."""
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
class AllowAll:
async def run(self, context):
@ -9126,7 +8997,7 @@ class TestEscalationKeywords:
async def test_escalation_overrides_session_pin_and_persists(self, mock_router_instance, basic_config):
"""Mid-session escalation bumps relative to the pinned model (never below it) and
the bumped model persists for later turns."""
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
router = ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
@ -9189,7 +9060,7 @@ class TestEscalationKeywords:
survives across the plumbing turns after it, so reading escalation off it re-fires per turn and,
with the pin persisted, walks the session to the top tier. Escalation reads the newest turn's ask.
"""
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
router = ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
@ -9291,7 +9162,7 @@ class TestEscalationKeywords:
@pytest.mark.asyncio
async def test_session_escalation_at_ceiling_keeps_multi_model_pin(self, mock_router_instance):
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
router = ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
@ -9301,7 +9172,7 @@ class TestEscalationKeywords:
},
)
cache_key = router._get_session_affinity_cache_key("session-top", {})
await mock_router_instance.health().store.async_set_cache(key=cache_key, value="o1-b")
await mock_router_instance.health().session_pins.set(cache_key, "o1-b", 3600)
result = await router.async_pre_routing_hook(
model="test-model",
request_kwargs=self._request_kwargs("session-top"),
@ -9622,7 +9493,7 @@ class TestRoutingDecisionContents:
@pytest.mark.asyncio
async def test_session_pin_decision(self, mock_router_instance, basic_config):
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
router = ComplexityRouter(
model_name="test-complexity-router",
ctx=RoutingContext(mock_router_instance),
@ -9630,7 +9501,7 @@ class TestRoutingDecisionContents:
)
request_kwargs = {"metadata": {"session_id": "session-decision"}}
cache_key = router._get_session_affinity_cache_key("session-decision", request_kwargs)
await mock_router_instance.health().store.async_set_cache(key=cache_key, value="gpt-4o")
await mock_router_instance.health().session_pins.set(cache_key, "gpt-4o", 3600)
response = await router.async_pre_routing_hook(
model="test-complexity-router",
request_kwargs=request_kwargs,
@ -9995,7 +9866,7 @@ class TestEscalationIsRecordedConsistently:
@pytest.mark.asyncio
async def test_session_pin_at_ceiling_still_records_the_ask(self, mock_router_instance):
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
router = ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
@ -10003,7 +9874,7 @@ class TestEscalationIsRecordedConsistently:
)
request_kwargs = {"metadata": {"session_id": "session-ceiling"}}
cache_key = router._get_session_affinity_cache_key("session-ceiling", request_kwargs)
await mock_router_instance.health().store.async_set_cache(key=cache_key, value="o1-preview")
await mock_router_instance.health().session_pins.set(cache_key, "o1-preview", 3600)
response = await router.async_pre_routing_hook(
model="test-router",
@ -10020,7 +9891,7 @@ class TestEscalationIsRecordedConsistently:
@pytest.mark.asyncio
async def test_session_pin_below_ceiling_reports_the_bump(self, mock_router_instance):
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
router = ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
@ -10028,7 +9899,7 @@ class TestEscalationIsRecordedConsistently:
)
request_kwargs = {"metadata": {"session_id": "session-below"}}
cache_key = router._get_session_affinity_cache_key("session-below", request_kwargs)
await mock_router_instance.health().store.async_set_cache(key=cache_key, value="gpt-4o-mini")
await mock_router_instance.health().session_pins.set(cache_key, "gpt-4o-mini", 3600)
response = await router.async_pre_routing_hook(
model="test-router",
@ -10410,7 +10281,9 @@ class TestContextAwareClassifier:
)
router: Final = ComplexityRouter(
model_name="router",
ctx=RoutingContext(MagicMock(acompletion=completion, **{"health.return_value.store": DualCache()})),
ctx=RoutingContext(
MagicMock(acompletion=completion, **{"health.return_value.session_pins": memory_pins()})
),
complexity_router_config={
"tiers": {"SIMPLE": "simple-model", "COMPLEX": "task-model"},
"classifier_type": "llm",
@ -10507,7 +10380,9 @@ class TestContextAwareClassifier:
)
router: Final = ComplexityRouter(
model_name="router",
ctx=RoutingContext(MagicMock(acompletion=completion, **{"health.return_value.store": DualCache()})),
ctx=RoutingContext(
MagicMock(acompletion=completion, **{"health.return_value.session_pins": memory_pins()})
),
complexity_router_config={
"tiers": {"SIMPLE": "simple-model", "COMPLEX": "task-model"},
"classifier_type": "llm",
@ -12012,7 +11887,7 @@ class TestClientHousekeepingCalls:
An agent names the conversation on its first turn, so the cheapest tier would be the pin
every session starts with and the real work that follows would run there for the whole TTL.
"""
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "REASONING"}'))
router = ComplexityRouter(
model_name="test-router",
@ -12244,7 +12119,7 @@ class TestConversationShapeDiscriminator:
async def test_a_single_ask_is_a_first_turn(self, mock_router_instance, basic_config):
"""Nothing is cached for any model yet, so the baseline would have paid the same
cache write and the saving is the plain rate difference."""
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
result = await self._router(mock_router_instance, basic_config).async_pre_routing_hook(
model="test-model",
request_kwargs={"metadata": {}},
@ -12256,7 +12131,7 @@ class TestConversationShapeDiscriminator:
async def test_a_second_ask_means_the_baseline_was_already_warm(self, mock_router_instance, basic_config):
"""An earlier turn was served, so a single-model deployment wrote the prompt then
and would only read it now; this request's write is what switching cost."""
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
result = await self._router(mock_router_instance, basic_config).async_pre_routing_hook(
model="test-model",
request_kwargs={"metadata": {}},
@ -12272,7 +12147,7 @@ class TestConversationShapeDiscriminator:
async def test_it_needs_no_session_id(self, mock_router_instance, basic_config):
"""The whole point of reading the conversation rather than remembering it: a
caller that sends no session header is still classified correctly."""
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
router = self._router(mock_router_instance, basic_config)
first = await router.async_pre_routing_hook(
model="test-model", request_kwargs={}, messages=[{"role": "user", "content": "Hello!"}]
@ -12293,17 +12168,15 @@ class TestConversationShapeDiscriminator:
async def test_it_touches_no_cache(self, mock_router_instance, basic_config):
"""Reading the request instead of remembering it is what removes the routing-path
round-trip, and with it a cache failure that would read as a first turn."""
cache = AsyncMock()
cache.async_get_cache = AsyncMock(return_value=None)
mock_router_instance.health.return_value.store = cache
pins: Final = MagicMock(spec=RouterSessionPinStore)
mock_router_instance.health.return_value.session_pins = pins
result = await self._router(mock_router_instance, basic_config).async_pre_routing_hook(
model="test-model",
request_kwargs={"metadata": {}},
messages=[{"role": "user", "content": "Hello!"}],
)
assert result.routing_decision["conversation_continuing"] is False
assert cache.async_get_cache.await_count == 0
assert cache.async_set_cache.await_count == 0
assert pins.method_calls == []
@pytest.mark.parametrize(
"history",
@ -12733,7 +12606,7 @@ class TestClassifierFallbackChoice:
"session_affinity": True,
},
)
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
request_kwargs: Dict = {"metadata": {"session_id": "session-flaky"}}
mock_router_instance.acompletion = AsyncMock(side_effect=TimeoutError("classifier timed out"))
@ -12772,7 +12645,7 @@ class TestClassifierFallbackChoice:
"session_affinity": True,
},
)
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
request_kwargs: Dict = {"metadata": {"session_id": "session-steady"}}
mock_router_instance.acompletion = AsyncMock(return_value=_llm_response('{"tier": "REASONING"}'))
@ -13822,9 +13695,7 @@ class TestPlanModeTierFloor:
"""Mid-session shift+tab into plan mode: the plan turns route at the floor, but the
stored pin keeps the session's own model, so the first turn after plan mode exits
auto-routes back to it instead of staying premium."""
from litellm.caching.dual_cache import DualCache
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
config = {**basic_config, "plan_mode_min_tier": "COMPLEX", "session_affinity": True}
router = self._router(mock_router_instance, config)
session_kwargs = {"metadata": {"session_id": "plan-session"}}
@ -13863,9 +13734,7 @@ class TestPlanModeTierFloor:
async def test_plan_mode_first_turn_does_not_seed_the_session_pin(self, mock_router_instance, basic_config):
"""A session whose first turn is already in plan mode must not pin the floored model:
the first ordinary turn classifies and pins as if plan mode had never happened."""
from litellm.caching.dual_cache import DualCache
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
config = {**basic_config, "plan_mode_min_tier": "COMPLEX", "session_affinity": True}
router = self._router(mock_router_instance, config)
session_kwargs = {"metadata": {"session_id": "plan-first-session"}}
@ -13887,9 +13756,7 @@ class TestPlanModeTierFloor:
@pytest.mark.asyncio
async def test_pinned_session_at_or_above_floor_keeps_pin_cause(self, mock_router_instance, basic_config):
from litellm.caching.dual_cache import DualCache
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
config = {**basic_config, "plan_mode_min_tier": "MEDIUM", "session_affinity": True}
router = self._router(mock_router_instance, config)
session_kwargs = {"metadata": {"session_id": "premium-session"}}
@ -14070,9 +13937,7 @@ class TestPlanModeTierFloor:
"""A plan-mode turn routed at or above the floor keeps its ordinary cause, but it still
must not pin: on an adaptive router the hard floor shaped that pick, and any sentinel
turn's pin would carry plan mode past its exit."""
from litellm.caching.dual_cache import DualCache
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
config = {
**basic_config,
"plan_mode_min_tier": "MEDIUM",
@ -14278,10 +14143,11 @@ async def test_tier_params_mask_credentials_in_routing_decision(route, mock_rout
)
request_kwargs = {"metadata": {"session_id": "masked-params-session"}}
if route == "session":
mock_router_instance.health.return_value.store = DualCache()
await mock_router_instance.health().store.async_set_cache(
key=router._get_session_affinity_cache_key("masked-params-session", request_kwargs),
value={"model": "opus", "tier": "REASONING"},
mock_router_instance.health.return_value.session_pins = memory_pins()
await mock_router_instance.health().session_pins.set(
router._get_session_affinity_cache_key("masked-params-session", request_kwargs),
{"model": "opus", "tier": "REASONING"},
3600,
)
message = "reason carefully about this" if route == "keyword" else "hello"
@ -14302,7 +14168,7 @@ async def test_tier_params_mask_credentials_in_routing_decision(route, mock_rout
@pytest.mark.asyncio
async def test_session_pin_outside_tiers_does_not_inherit_medium_params(mock_router_instance):
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
router = ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
@ -14316,9 +14182,8 @@ async def test_session_pin_outside_tiers_does_not_inherit_medium_params(mock_rou
},
)
request_kwargs = {"metadata": {"session_id": "orphan-session"}}
await mock_router_instance.health().store.async_set_cache(
key=router._get_session_affinity_cache_key("orphan-session", request_kwargs),
value="orphan",
await mock_router_instance.health().session_pins.set(
router._get_session_affinity_cache_key("orphan-session", request_kwargs), "orphan", 3600
)
response = await router.async_pre_routing_hook(
@ -14334,7 +14199,7 @@ async def test_session_pin_outside_tiers_does_not_inherit_medium_params(mock_rou
@pytest.mark.asyncio
async def test_session_pin_uses_recorded_tier_when_model_is_in_multiple_tiers(mock_router_instance):
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
router = ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
@ -14347,9 +14212,10 @@ async def test_session_pin_uses_recorded_tier_when_model_is_in_multiple_tiers(mo
},
)
request_kwargs = {"metadata": {"session_id": "shared-session"}}
await mock_router_instance.health().store.async_set_cache(
key=router._get_session_affinity_cache_key("shared-session", request_kwargs),
value={"model": "shared", "tier": "SIMPLE"},
await mock_router_instance.health().session_pins.set(
router._get_session_affinity_cache_key("shared-session", request_kwargs),
{"model": "shared", "tier": "SIMPLE"},
3600,
)
response = await router.async_pre_routing_hook(
@ -14366,8 +14232,8 @@ async def test_session_pin_uses_recorded_tier_when_model_is_in_multiple_tiers(mo
@pytest.mark.asyncio
async def test_session_pin_survives_json_list_round_trip(mock_router_instance):
cache: Final = DualCache()
mock_router_instance.health.return_value.store = cache
cache: Final = memory_pins()
mock_router_instance.health.return_value.session_pins = cache
router = ComplexityRouter(
model_name="test-router",
ctx=RoutingContext(mock_router_instance),
@ -14381,7 +14247,7 @@ async def test_session_pin_survives_json_list_round_trip(mock_router_instance):
)
request_kwargs = {"metadata": {"session_id": "json-round-trip-session"}}
key: Final = router._get_session_affinity_cache_key("json-round-trip-session", request_kwargs)
cache.in_memory_cache.set_cache(key, ["shared", "SIMPLE"])
cache.set_local(key, ["shared", "SIMPLE"], 3600)
response = await router.async_pre_routing_hook(
model="test-router",
@ -14392,7 +14258,7 @@ async def test_session_pin_survives_json_list_round_trip(mock_router_instance):
assert response is not None
assert response.model == "shared"
assert response.litellm_params == {"reasoning_effort": "low"}
assert cache.in_memory_cache.get_cache(key) == ["shared", "SIMPLE"]
assert cache.get_local(key) == ["shared", "SIMPLE"]
HEURISTIC_FIRST_TIERS: dict[str, str] = {
@ -15531,8 +15397,8 @@ class TestModalityRouting:
{"role": "user", "content": [{"type": "text", "text": "quick lookup: what is this?"}, IMG_PART]}
]
elif path.startswith(("pin_kept", "pin_replacement", "pin_override")):
cache: Final = DualCache()
mock_router_instance.health.return_value.store = cache
cache: Final = memory_pins()
mock_router_instance.health.return_value.session_pins = cache
config["session_affinity"] = True
request_kwargs = {"metadata": {"session_id": "s1"}}
if path == "pin_replacement_gated":
@ -15553,9 +15419,10 @@ class TestModalityRouting:
_answer_deployments_named(mock_router_instance, [])
router = self._router(mock_router_instance, config, vision)
if path.startswith(("pin_kept", "pin_replacement", "pin_override")):
mock_router_instance.health().store.in_memory_cache.set_cache(
mock_router_instance.health().session_pins.set_local(
router._get_session_affinity_cache_key("s1", request_kwargs),
{"model": "text-cheap", "tier": "SIMPLE"},
3600,
)
result = await router.async_pre_routing_hook(model="m", request_kwargs=request_kwargs, messages=messages)
assert result.model == expected_model
@ -15658,7 +15525,7 @@ class TestModalityRouting:
async def test_continuation_turn_screenshot_escalates_past_the_held_model(self, mock_router_instance):
"""classification_mode user_turn replays the held model on continuation turns; a
continuation carrying a screenshot must still be re-placed when that model is text-only."""
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
config = {
"tiers": dict(self.BASE_TIERS),
"classification_mode": "user_turn",
@ -15726,8 +15593,8 @@ class TestModalityRouting:
@pytest.mark.asyncio
async def test_pin_override_serves_the_image_turn_without_repinning(self, mock_router_instance):
"""The override is for one request: the session keeps the model it was pinned to."""
cache: Final = DualCache()
mock_router_instance.health.return_value.store = cache
cache: Final = memory_pins()
mock_router_instance.health.return_value.session_pins = cache
router = self._router(
mock_router_instance,
{
@ -15740,7 +15607,7 @@ class TestModalityRouting:
)
request_kwargs = {"metadata": {"session_id": "s1"}}
key: Final = router._get_session_affinity_cache_key("s1", request_kwargs)
cache.in_memory_cache.set_cache(key, {"model": "text-cheap", "tier": "SIMPLE"})
cache.set_local(key, {"model": "text-cheap", "tier": "SIMPLE"}, 3600)
image_turn = await router.async_pre_routing_hook(
model="m", request_kwargs=request_kwargs, messages=self.IMAGE_MESSAGE
@ -15749,7 +15616,7 @@ class TestModalityRouting:
assert image_turn.routing_decision["cause"] == "modality_pin_override"
assert "modality_escalated_from:SIMPLE" in image_turn.routing_decision["signals"]
assert cache.in_memory_cache.get_cache(key) == {"model": "text-cheap", "tier": "SIMPLE"}
assert cache.get_local(key) == {"model": "text-cheap", "tier": "SIMPLE"}
text_turn = await router.async_pre_routing_hook(
model="m", request_kwargs={"metadata": {"session_id": "s1"}}, messages=[{"role": "user", "content": "hi"}]
@ -15760,8 +15627,8 @@ class TestModalityRouting:
@pytest.mark.asyncio
async def test_pin_override_with_no_capable_model_rejects_and_keeps_the_pin(self, mock_router_instance):
"""The clear 400 replaces the provider's, and a rejected turn must not cost the session its pin."""
cache: Final = DualCache()
mock_router_instance.health.return_value.store = cache
cache: Final = memory_pins()
mock_router_instance.health.return_value.session_pins = cache
router = self._router(
mock_router_instance,
{
@ -15773,12 +15640,12 @@ class TestModalityRouting:
{"text-cheap": False, "text-big": False},
)
key: Final = router._get_session_affinity_cache_key("s1", {"metadata": {"session_id": "s1"}})
cache.in_memory_cache.set_cache(key, {"model": "text-cheap", "tier": "SIMPLE"})
cache.set_local(key, {"model": "text-cheap", "tier": "SIMPLE"}, 3600)
with pytest.raises(litellm.BadRequestError, match="no model"):
await router.async_pre_routing_hook(
model="m", request_kwargs={"metadata": {"session_id": "s1"}}, messages=self.IMAGE_MESSAGE
)
assert cache.in_memory_cache.get_cache(key) == {"model": "text-cheap", "tier": "SIMPLE"}
assert cache.get_local(key) == {"model": "text-cheap", "tier": "SIMPLE"}
@pytest.mark.usefixtures("local_model_cost_map")
@ -16167,7 +16034,7 @@ class TestHealthFallbackDispatch:
assert [c.request.url.host for c in upstream.calls] == ["fallback.test"]
strategy: Final = router.complexity_routers["health-router"][0].strategy
key: Final = strategy._get_session_affinity_cache_key("outage", {})
assert await router.cache.async_get_cache(key=key) is None
assert await router.health().session_pins.get(key) is None
if source == "health":
router.health_state_cache.set_deployment_health_states(
{"primary-id": {"is_healthy": True, "timestamp": time.time()}}
@ -16194,7 +16061,7 @@ class TestHealthFallbackDispatch:
)
strategy: Final = router.complexity_routers["health-router"][0].strategy
key: Final = strategy._get_session_affinity_cache_key("precedence", {})
await router.cache.async_set_cache(key=key, value={"model": "primary", "tier": "SIMPLE"}, ttl=600)
await router.health().session_pins.set(key, {"model": "primary", "tier": "SIMPLE"}, 600)
with respx.mock(assert_all_mocked=True) as upstream:
upstream.post(host__regex=r"^(primary|peer|fallback)\.test$").mock(side_effect=self._http_response)
for model_id, expected, cause in (
@ -16206,7 +16073,7 @@ class TestHealthFallbackDispatch:
metadata: Final[dict[str, object]] = {"session_id": "precedence"}
assert await self._request(router, "chat", False, metadata) == expected
assert metadata["routing_decision"]["cause"] == cause
assert await router.cache.async_get_cache(key=key) == {"model": "primary", "tier": "SIMPLE"}
assert await router.health().session_pins.get(key) == {"model": "primary", "tier": "SIMPLE"}
assert [c.request.url.host for c in upstream.calls] == ["primary.test", "peer.test", "fallback.test"]
@pytest.mark.asyncio
@ -16559,7 +16426,7 @@ class TestTierHealthFailover:
mock_router_instance.select = select
mock_router_instance.probed_kwargs = probed_kwargs
mock_router_instance.probed_prompts = probed_prompts
mock_router_instance.health.return_value.store = DualCache()
mock_router_instance.health.return_value.session_pins = memory_pins()
return ComplexityRouter(
model_name="health-test-router",
ctx=RoutingContext(mock_router_instance),
@ -16587,9 +16454,7 @@ class TestTierHealthFailover:
)
# Seed the pin onto the dead group directly so the replay path is exercised.
key = router._get_session_affinity_cache_key("sess-dead", {})
await router.ctx.health().store.async_set_cache(
key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600
)
await router.ctx.health().session_pins.set(key, {"model": "dead-a", "tier": "SIMPLE"}, 600)
result = await router.async_pre_routing_hook(
model="m", request_kwargs={"metadata": {"session_id": "sess-dead"}}, messages=self.SIMPLE_MESSAGE
)
@ -16652,9 +16517,7 @@ class TestTierHealthFailover:
)
pinned = "only" if tiers else "dead-a"
key = router._get_session_affinity_cache_key("sess-open", {})
await router.ctx.health().store.async_set_cache(
key=key, value={"model": pinned, "tier": "SIMPLE"}, ttl=600
)
await router.ctx.health().session_pins.set(key, {"model": pinned, "tier": "SIMPLE"}, 600)
result = await router.async_pre_routing_hook(
model="m", request_kwargs={"metadata": {"session_id": "sess-open"}}, messages=self.SIMPLE_MESSAGE
)
@ -16682,9 +16545,7 @@ class TestTierHealthFailover:
request_kwargs={"metadata": {"session_id": session_id}},
messages=self.SIMPLE_MESSAGE,
)
return await router.ctx.health().store.async_get_cache(
key=router._get_session_affinity_cache_key(session_id, {})
)
return await router.ctx.health().session_pins.get(router._get_session_affinity_cache_key(session_id, {}))
stored = [await pin_after_session(turn) for turn in range(20)]
assert all(entry in (None, {"model": "live-b", "tier": "SIMPLE"}) for entry in stored)
@ -16707,9 +16568,7 @@ class TestTierHealthFailover:
messages=[{"role": "user", "content": TITLE_ASK}],
)
assert result.routing_decision["cause"] in ("housekeeping", "health_failover")
stored = await router.ctx.health().store.async_get_cache(
key=router._get_session_affinity_cache_key(session_id, {})
)
stored = await router.ctx.health().session_pins.get(router._get_session_affinity_cache_key(session_id, {}))
assert stored is None
@pytest.mark.asyncio
@ -16732,9 +16591,7 @@ class TestTierHealthFailover:
blocked=("id-b1",),
)
key = router._get_session_affinity_cache_key("sess-paused", {})
await router.ctx.health().store.async_set_cache(
key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600
)
await router.ctx.health().session_pins.set(key, {"model": "dead-a", "tier": "SIMPLE"}, 600)
results = [
await router.async_pre_routing_hook(
model="m", request_kwargs={"metadata": {"session_id": "sess-paused"}}, messages=self.SIMPLE_MESSAGE
@ -16774,9 +16631,7 @@ class TestTierHealthFailover:
cooling=("id-a1",),
)
key = router._get_session_affinity_cache_key("sess-adaptive", {})
await router.ctx.health().store.async_set_cache(
key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600
)
await router.ctx.health().session_pins.set(key, {"model": "dead-a", "tier": "SIMPLE"}, 600)
request_kwargs = {"metadata": {"session_id": "sess-adaptive", "adaptive_router_chosen_model": "dead-a"}}
result = await router.async_pre_routing_hook(
model="m", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE
@ -16820,9 +16675,7 @@ class TestTierHealthFailover:
cooling=("id-v1",),
)
key = router._get_session_affinity_cache_key("sess-image", {})
await router.ctx.health().store.async_set_cache(
key=key, value={"model": "dead-vision", "tier": "SIMPLE"}, ttl=600
)
await router.ctx.health().session_pins.set(key, {"model": "dead-vision", "tier": "SIMPLE"}, 600)
image_message = [
{
"role": "user",
@ -16868,9 +16721,7 @@ class TestTierHealthFailover:
},
)
key = router._get_session_affinity_cache_key("sess-window", {})
await router.ctx.health().store.async_set_cache(
key=key, value={"model": "dead-big", "tier": "SIMPLE"}, ttl=600
)
await router.ctx.health().session_pins.set(key, {"model": "dead-big", "tier": "SIMPLE"}, 600)
results = [
await router.async_pre_routing_hook(
model="m",
@ -16914,9 +16765,7 @@ class TestTierHealthFailover:
{"live-b": ["id-b1"]},
)
key = router._get_session_affinity_cache_key("sess-unknown", {})
await router.ctx.health().store.async_set_cache(
key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600
)
await router.ctx.health().session_pins.set(key, {"model": "dead-a", "tier": "SIMPLE"}, 600)
result = await router.async_pre_routing_hook(
model="m", request_kwargs={"metadata": {"session_id": "sess-unknown"}}, messages=self.SIMPLE_MESSAGE
)
@ -16943,9 +16792,7 @@ class TestTierHealthFailover:
excluded=("id-b1",),
)
key = router._get_session_affinity_cache_key("sess-tagged", {})
await router.ctx.health().store.async_set_cache(
key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600
)
await router.ctx.health().session_pins.set(key, {"model": "dead-a", "tier": "SIMPLE"}, 600)
results = [
await router.async_pre_routing_hook(
model="m", request_kwargs={"metadata": {"session_id": "sess-tagged"}}, messages=self.SIMPLE_MESSAGE
@ -16965,9 +16812,7 @@ class TestTierHealthFailover:
cooling=("id-a1",),
)
key = router._get_session_affinity_cache_key("sess-kwargs", {})
await router.ctx.health().store.async_set_cache(
key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600
)
await router.ctx.health().session_pins.set(key, {"model": "dead-a", "tier": "SIMPLE"}, 600)
request_kwargs = {
"metadata": {"session_id": "sess-kwargs"},
"_target_order": 1,
@ -17003,9 +16848,7 @@ class TestTierHealthFailover:
raises_for={"rpm-full-b": RouterRateLimitErrorBasic(model="rpm-full-b")},
)
key = router._get_session_affinity_cache_key("sess-rpm", {})
await router.ctx.health().store.async_set_cache(
key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600
)
await router.ctx.health().session_pins.set(key, {"model": "dead-a", "tier": "SIMPLE"}, 600)
results = [
await router.async_pre_routing_hook(
model="m", request_kwargs={"metadata": {"session_id": "sess-rpm"}}, messages=self.SIMPLE_MESSAGE
@ -17026,9 +16869,7 @@ class TestTierHealthFailover:
cooling=("id-a1",),
)
key = router._get_session_affinity_cache_key("sess-input", {})
await router.ctx.health().store.async_set_cache(
key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600
)
await router.ctx.health().session_pins.set(key, {"model": "dead-a", "tier": "SIMPLE"}, 600)
result = await router.async_pre_routing_hook(
model="m",
request_kwargs={"metadata": {"session_id": "sess-input"}},
@ -17076,10 +16917,8 @@ class TestTierHealthFailover:
sessions: Final = tuple(f"sess-exhausted-{sample}" for sample in range(20))
await asyncio.gather(
*(
router.ctx.health().store.async_set_cache(
key=router._get_session_affinity_cache_key(session_id, {}),
value={"model": "dead-a", "tier": "SIMPLE"},
ttl=600,
router.ctx.health().session_pins.set(
router._get_session_affinity_cache_key(session_id, {}), {"model": "dead-a", "tier": "SIMPLE"}, 600
)
for session_id in sessions
)
@ -17123,9 +16962,7 @@ class TestTierHealthFailover:
cooling=("id-a1",),
)
key = router._get_session_affinity_cache_key("sess-missing", {})
await router.ctx.health().store.async_set_cache(
key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600
)
await router.ctx.health().session_pins.set(key, {"model": "dead-a", "tier": "SIMPLE"}, 600)
results = [
await router.async_pre_routing_hook(
model="m", request_kwargs={"metadata": {"session_id": "sess-missing"}}, messages=self.SIMPLE_MESSAGE

4
tests/unit/router_strategy/test_llm_v2.py
View file

@ -8,7 +8,6 @@ from pydantic import ValidationError
import litellm
from litellm import ModelResponse, Router
from litellm.caching.dual_cache import DualCache
from litellm.router_api import RoutingContext
from litellm.router_strategy.auto_router.auto_router_model_naming import strategy_router_dependencies
from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter
@ -24,6 +23,7 @@ from litellm.router_strategy.complexity_router.llm_v2 import (
llm_v2_response_format,
)
from litellm.types.llms.openai import ResponsesAPIResponse
from tests.unit.redis.stores.router_session_pins_memory import memory_pins
def _config(**overrides: object) -> ComplexityRouterConfig:
@ -564,7 +564,7 @@ def test_response_schema_requires_both_model_forecasts() -> None:
@pytest.mark.asyncio
async def test_user_turn_mode_reuses_forecast_until_a_new_user_requirement() -> None:
router, client = _router(_verdict().model_dump_json(), _config(classification_mode="user_turn"))
client.health.return_value.store = DualCache()
client.health.return_value.session_pins = memory_pins()
initial: Final = [{"role": "user", "content": "Fix nested behavior"}]
first: Final = await router.async_pre_routing_hook(
model="v2-router", messages=initial, request_kwargs={"metadata": {"session_id": "v2-task"}}

85
tests/unit/router_utils/pre_call_checks/test_deployment_affinity_check.py
View file

@ -8,11 +8,11 @@ from unittest.mock import AsyncMock, patch
import pytest
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
DeploymentAffinityCheck,
)
from tests.unit.redis.stores.router_session_pins_memory import memory_pins
@pytest.fixture(autouse=True)
@ -371,7 +371,7 @@ async def test_async_previous_response_id_priority_over_user_key_affinity():
model_group=model_group,
user_key=user_api_key_hash,
)
await router.cache.async_set_cache(affinity_cache_key, {"model_id": other_model_id}, ttl=3600)
await router.health().session_pins.set(affinity_cache_key, {"model_id": other_model_id}, 3600)
# Even though user-key affinity points elsewhere, previous_response_id should pin
# to the deployment that created the original response.
@ -460,10 +460,10 @@ async def test_async_pre_call_hook_uses_model_map_key_scope():
Deployment affinity caching uses (user_api_key_hash, model_map_key) -> model_id.
"""
cache = DualCache()
pins = memory_pins()
callback = DeploymentAffinityCheck(
cache=cache,
pins=pins,
ttl_seconds=123,
enable_user_key_affinity=True,
enable_responses_api_affinity=False,
@ -483,7 +483,7 @@ async def test_async_pre_call_hook_uses_model_map_key_scope():
model_group="claude-sonnet-4-5@20250929",
user_key="user-key-abc",
)
assert await cache.async_get_cache(key=expected_cache_key) == {"model_id": "model-id-123"}
assert await pins.get(expected_cache_key) == {"model_id": "model-id-123"}
@pytest.mark.asyncio
@ -499,11 +499,15 @@ async def test_async_filter_deployments_uses_stable_model_map_key_for_affinity_s
user_key = "user-key-abc"
stable_model_map_key = "claude-sonnet-4-5@20250929"
cache = AsyncMock()
cache.async_get_cache = AsyncMock()
expected_cache_key = DeploymentAffinityCheck.get_affinity_cache_key(
model_group=stable_model_map_key,
user_key=user_key,
)
pins = memory_pins()
await pins.set(expected_cache_key, {"model_id": "deployment-2"}, 60)
callback = DeploymentAffinityCheck(
cache=cache,
pins=pins,
ttl_seconds=123,
enable_user_key_affinity=True,
enable_responses_api_affinity=False,
@ -522,18 +526,6 @@ async def test_async_filter_deployments_uses_stable_model_map_key_for_affinity_s
},
]
expected_cache_key = DeploymentAffinityCheck.get_affinity_cache_key(
model_group=stable_model_map_key,
user_key=user_key,
)
async def get_cache_side_effect(*, key: str):
if key == expected_cache_key:
return {"model_id": "deployment-2"}
return None
cache.async_get_cache.side_effect = get_cache_side_effect
filtered = await callback.async_filter_deployments(
model="some-router-model-group",
healthy_deployments=healthy_deployments,
@ -556,11 +548,15 @@ async def test_async_filter_deployments_falls_back_when_cached_deployment_is_unh
user_key = "user-key-unhealthy"
stable_model_map_key = "claude-sonnet-4-5@20250929"
cache = AsyncMock()
cache.async_get_cache = AsyncMock(return_value={"model_id": "stale-deployment"})
pins = memory_pins()
await pins.set(
DeploymentAffinityCheck.get_affinity_cache_key(model_group=stable_model_map_key, user_key=user_key),
{"model_id": "stale-deployment"},
60,
)
callback = DeploymentAffinityCheck(
cache=cache,
pins=pins,
ttl_seconds=123,
enable_user_key_affinity=True,
enable_responses_api_affinity=False,
@ -594,10 +590,14 @@ async def test_async_filter_deployments_falls_back_when_cached_deployment_is_unh
async def test_async_filter_deployments_does_not_pin_when_target_order_is_set():
user_key = "user-key-order-fallback"
stable_model_map_key = "claude-sonnet-4-5@20250929"
cache = AsyncMock()
cache.async_get_cache = AsyncMock(return_value={"model_id": "deployment-1"})
pins = memory_pins()
await pins.set(
DeploymentAffinityCheck.get_affinity_cache_key(model_group=stable_model_map_key, user_key=user_key),
{"model_id": "deployment-1"},
60,
)
callback = DeploymentAffinityCheck(
cache=cache,
pins=pins,
ttl_seconds=123,
enable_user_key_affinity=True,
enable_responses_api_affinity=False,
@ -624,7 +624,6 @@ async def test_async_filter_deployments_does_not_pin_when_target_order_is_set():
)
assert filtered == healthy_deployments
cache.async_get_cache.assert_not_called()
@pytest.mark.asyncio
@ -634,7 +633,7 @@ async def test_async_user_key_affinity_ttl_expiry_allows_reroute():
"""
callback = DeploymentAffinityCheck(
cache=DualCache(),
pins=memory_pins(),
ttl_seconds=1,
enable_user_key_affinity=True,
enable_responses_api_affinity=False,
@ -688,27 +687,13 @@ async def test_async_user_key_affinity_ttl_expiry_allows_reroute():
assert after_ttl_expiry == healthy_deployments
def test_cache_key_does_not_double_hash_user_api_key_hash():
"""
Proxy typically provides `metadata.user_api_key_hash` as a SHA-256 hex string.
The affinity cache key should not hash it again.
"""
user_api_key_hash = "b95b015b66dd02a1c14e1e0a8729211f8ee53ec962658764f4cf58546c2c68e1"
key = DeploymentAffinityCheck.get_affinity_cache_key(
model_group="any-model-group",
user_key=user_api_key_hash,
)
assert key.endswith(user_api_key_hash)
def test_get_effective_flags_returns_per_group_config():
"""
_get_effective_flags should return per-group flags when the model group has an entry
in model_group_affinity_config, and global flags otherwise.
"""
callback = DeploymentAffinityCheck(
cache=AsyncMock(),
pins=memory_pins(),
ttl_seconds=60,
enable_user_key_affinity=True,
enable_responses_api_affinity=True,
@ -848,7 +833,7 @@ async def test_model_group_affinity_config_falls_back_to_global():
unconfigured model groups should use the global settings.
"""
callback = DeploymentAffinityCheck(
cache=DualCache(),
pins=memory_pins(),
ttl_seconds=60,
enable_user_key_affinity=True,
enable_responses_api_affinity=False,
@ -905,7 +890,7 @@ async def test_model_group_affinity_config_overrides_global():
user-key affinity (from global config) should NOT apply to that group.
"""
callback = DeploymentAffinityCheck(
cache=DualCache(),
pins=memory_pins(),
ttl_seconds=60,
enable_user_key_affinity=True,
enable_responses_api_affinity=False,
@ -933,7 +918,7 @@ async def test_model_group_affinity_config_overrides_global():
# Set up user-key affinity cache for claude-3
cache_key = DeploymentAffinityCheck.get_affinity_cache_key(model_group=stable_model_map_key, user_key=user_key)
await callback.cache.async_set_cache(cache_key, {"model_id": "deployment-1"}, ttl=60)
await callback.pins.set(cache_key, {"model_id": "deployment-1"}, 60)
# claude-3 has per-group config (session_affinity only), so user-key affinity
# should NOT apply even though it's globally enabled
@ -1032,7 +1017,7 @@ async def test_proxy_jwt_auth_metadata_pins_per_user():
model_group = "gpt-5.4-mini"
healthy_deployments = _two_deployments(model_group)
callback = DeploymentAffinityCheck(
cache=DualCache(),
pins=memory_pins(),
ttl_seconds=60,
enable_user_key_affinity=True,
enable_responses_api_affinity=False,
@ -1086,7 +1071,7 @@ async def test_jwt_user_id_never_reads_a_virtual_key_pin():
model_group = "gpt-5.4-mini"
healthy_deployments = _two_deployments(model_group)
callback = DeploymentAffinityCheck(
cache=DualCache(),
pins=memory_pins(),
ttl_seconds=60,
enable_user_key_affinity=True,
enable_responses_api_affinity=False,
@ -1129,7 +1114,7 @@ async def test_virtual_key_hash_wins_over_user_id_for_affinity():
model_group = "gpt-5.4-mini"
healthy_deployments = _two_deployments(model_group)
callback = DeploymentAffinityCheck(
cache=DualCache(),
pins=memory_pins(),
ttl_seconds=60,
enable_user_key_affinity=True,
enable_responses_api_affinity=False,
@ -1221,7 +1206,7 @@ async def test_retry_reaches_healthy_sibling_when_affinity_pins_deployment_that_
@pytest.mark.asyncio
async def test_previous_response_id_pin_holds_even_after_its_deployment_failed():
callback: Final = DeploymentAffinityCheck(
cache=DualCache(),
pins=memory_pins(),
ttl_seconds=60,
enable_user_key_affinity=True,
enable_responses_api_affinity=True,

6
tests/unit/router_utils/pre_call_checks/test_encrypted_content_affinity_check.py
View file

@ -1833,11 +1833,7 @@ async def test_model_group_encrypted_content_affinity_overrides_global_deploymen
model_group=model_group,
user_key=user_api_key_hash,
)
await deployment_callback.cache.async_set_cache(
key=cache_key,
value={"model_id": "deployment-a"},
ttl=60,
)
await deployment_callback.pins.set(cache_key, {"model_id": "deployment-a"}, 60)
encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test")
request_kwargs = {
"input": [

422
tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py
View file

@ -1,29 +1,28 @@
import asyncio
import json
import uuid
from typing import Final
from collections.abc import AsyncIterator, Mapping
from collections.abc import Mapping
from datetime import datetime, timezone
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
from fakeredis import FakeRedis, FakeServer
from fakeredis.aioredis import FakeConnection
from redis.asyncio import BlockingConnectionPool, Redis
import litellm
from litellm.caching.affinity_cache import claim_affinity_pin
from litellm.caching.dual_cache import DualCache
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.caching.redis_cache import RedisCache
from litellm.redis.stores.router_session_pins import (
DeploymentPin,
DeploymentSessionPin,
RouterSessionPinStore,
TierModelPin,
)
from litellm.constants import (
REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
SESSION_ID_GENERATED_METADATA_KEY,
)
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
DeploymentAffinityCheck,
)
from tests.unit.redis.stores.router_session_pins_memory import memory_pins
@ -61,7 +60,11 @@ def _repin_params(deployment_id: str, caller: str = "caller-1") -> dict[str, obj
}
def _repin_key(caller: str = "caller-1") -> str:
def _pins(router: litellm.Router) -> RouterSessionPinStore:
return router.health().session_pins
def _repin_key(caller: str = "caller-1") -> DeploymentSessionPin:
return DeploymentAffinityCheck.get_session_affinity_cache_key("repin-group", "repin-session", caller)
@ -90,25 +93,25 @@ async def test_session_repin_requires_opt_in_unavailable_original_and_success(
) -> None:
router: Final = _repinning_router(enabled)
try:
await router.cache.async_set_cache(_repin_key(), {"model_id": "original"})
await router.cache.async_set_cache(_repin_key("caller-2"), {"model_id": "original"})
await _pins(router).set(_repin_key(), {"model_id": "original"}, 3600)
await _pins(router).set(_repin_key("caller-2"), {"model_id": "original"}, 3600)
user_key: Final = DeploymentAffinityCheck.get_affinity_cache_key("repin-group", "caller-1")
await router.cache.async_set_cache(user_key, {"model_id": "original"})
await _pins(router).set(user_key, {"model_id": "original"}, 3600)
if unavailable == "cooldown":
_cool_down_original(router)
elif unavailable == "removed":
router.delete_deployment("original")
callback: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck))
await callback.async_pre_call_deployment_hook(_repin_params("replacement"), None)
assert await router.cache.async_get_cache(_repin_key()) == {"model_id": "original"}
assert await _pins(router).get(_repin_key()) == {"model_id": "original"}
await _record_replacement_success(router)
expected: Final = "replacement" if enabled and unavailable != "temporary-filter" else "original"
assert await router.cache.async_get_cache(_repin_key()) == {"model_id": expected}
assert await router.cache.async_get_cache(_repin_key("caller-2")) == {"model_id": "original"}
assert await router.cache.async_get_cache(user_key) == {"model_id": "original"}
assert await _pins(router).get(_repin_key()) == {"model_id": expected}
assert await _pins(router).get(_repin_key("caller-2")) == {"model_id": "original"}
assert await _pins(router).get(user_key) == {"model_id": "original"}
router.cooldown_cache.cooldown_store.evict_local("original")
await _record_replacement_success(router, "other")
assert await router.cache.async_get_cache(_repin_key()) == {"model_id": expected}
assert await _pins(router).get(_repin_key()) == {"model_id": expected}
finally:
router.discard()
@ -118,11 +121,11 @@ async def test_session_repin_requires_opt_in_unavailable_original_and_success(
async def test_unsuccessful_responses_do_not_repin(status: str, wrapped: bool) -> None:
router: Final = _repinning_router()
try:
await router.cache.async_set_cache(_repin_key(), {"model_id": "original"})
await _pins(router).set(_repin_key(), {"model_id": "original"}, 3600)
_cool_down_original(router)
response: Final[Mapping[str, object]] = {"status": status}
await _record_replacement_success(router, response={"response": response} if wrapped else response)
assert await router.cache.async_get_cache(_repin_key()) == {"model_id": "original"}
assert await _pins(router).get(_repin_key()) == {"model_id": "original"}
finally:
router.discard()
@ -130,14 +133,14 @@ async def test_unsuccessful_responses_do_not_repin(status: str, wrapped: bool) -
async def test_concurrent_successful_replacements_keep_first_winner() -> None:
router: Final = _repinning_router()
try:
await router.cache.async_set_cache(_repin_key(), {"model_id": "original"})
await _pins(router).set(_repin_key(), {"model_id": "original"}, 3600)
_cool_down_original(router)
await asyncio.gather(_record_replacement_success(router), _record_replacement_success(router, "other"))
winner: Final = await router.cache.async_get_cache(_repin_key())
winner: Final = await _pins(router).get(_repin_key())
assert winner in ({"model_id": "replacement"}, {"model_id": "other"})
router.cooldown_cache.cooldown_store.evict_local("original")
await _record_replacement_success(router, "original")
assert await router.cache.async_get_cache(_repin_key()) == winner
assert await _pins(router).get(_repin_key()) == winner
finally:
router.discard()
@ -145,148 +148,13 @@ async def test_concurrent_successful_replacements_keep_first_winner() -> None:
async def test_healthy_legacy_string_pin_is_preserved_with_repinning_enabled() -> None:
router: Final = _repinning_router()
try:
await router.cache.async_set_cache(_repin_key(), "original")
await _pins(router).set(_repin_key(), "original", 3600)
await _record_replacement_success(router)
assert await router.cache.async_get_cache(_repin_key()) == "original"
assert await _pins(router).get(_repin_key()) == "original"
finally:
router.discard()
@pytest.mark.parametrize("enabled", [False, True])
async def test_session_selection_observes_another_workers_repin_only_when_enabled(enabled: bool) -> None:
server: Final = FakeServer()
sync_client: Final = FakeRedis(server=server)
pool: Final = BlockingConnectionPool(connection_class=FakeConnection, server=server)
router: Final = _repinning_router(enabled)
try:
with (
patch("redis.Redis.from_url", return_value=sync_client),
patch("redis.asyncio.BlockingConnectionPool.from_url", return_value=pool),
):
redis: Final = RedisCache(url=f"redis://repin-cache.example:6379/{int(enabled)}")
router.cache.attach_redis_cache(redis)
other_worker: Final = DualCache(redis_cache=redis)
router.cache.in_memory_cache.set_cache(_repin_key(), {"model_id": "original"}, ttl=3600)
await other_worker.async_set_cache(_repin_key(), {"model_id": "replacement"}, ttl=3600)
callback: Final = next(
cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck)
)
selected: Final = await callback.async_filter_deployments(
model="repin-group",
healthy_deployments=router.get_model_list(model_name="repin-group") or [],
messages=None,
request_kwargs=_repin_params("replacement"),
)
assert [deployment["model_info"]["id"] for deployment in selected] == [
"replacement" if enabled else "original"
]
finally:
router.discard()
await pool.aclose()
sync_client.close()
@pytest.fixture
async def shared_session_router() -> AsyncIterator[tuple[litellm.Router, FakeServer, RedisCache]]:
server: Final = FakeServer()
sync_client: Final = FakeRedis(server=server)
pool: Final = BlockingConnectionPool(connection_class=FakeConnection, server=server)
router: Final = _repinning_router()
try:
with (
patch("redis.Redis.from_url", return_value=sync_client),
patch("redis.asyncio.BlockingConnectionPool.from_url", return_value=pool),
):
redis: Final = RedisCache(url=f"redis://session-outage-{uuid.uuid4().hex}.example:6379/0")
router.cache.attach_redis_cache(redis)
yield router, server, redis
finally:
router.discard()
await pool.aclose()
sync_client.close()
async def _selected_session_ids(router: litellm.Router) -> tuple[str, ...]:
callback: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck))
selected: Final = await callback.async_filter_deployments(
model="repin-group",
healthy_deployments=router.get_model_list(model_name="repin-group") or [],
messages=None,
request_kwargs=_repin_params("replacement"),
)
return tuple(deployment["model_info"]["id"] for deployment in selected)
@pytest.mark.parametrize("has_pin", [False, True])
async def test_redis_read_failures_and_open_breaker_preserve_session_routing(
shared_session_router: tuple[litellm.Router, FakeServer, RedisCache], has_pin: bool
) -> None:
router, server, redis = shared_session_router
if has_pin:
await redis.async_set_cache(_repin_key(), {"model_id": "replacement"}, ttl=3600)
router.cache.in_memory_cache.set_cache(_repin_key(), {"model_id": "original"}, ttl=3600)
assert await _selected_session_ids(router) == ("replacement",)
server.connected = False
with patch("time.time", return_value=1790000000.0):
for _ in range(REDIS_CIRCUIT_BREAKER_FAILURE_THRESHOLD + 1):
assert set(await _selected_session_ids(router)) == (
{"replacement"} if has_pin else {"original", "replacement", "other"}
)
assert redis._health is not None and redis._health.breaker.state == "open"
async def test_confirmed_redis_miss_does_not_revive_a_local_pin_during_outage(
shared_session_router: tuple[litellm.Router, FakeServer, RedisCache],
) -> None:
router, server, _ = shared_session_router
router.cache.in_memory_cache.set_cache(_repin_key(), {"model_id": "original"}, ttl=3600)
assert set(await _selected_session_ids(router)) == {"original", "replacement", "other"}
server.connected = False
assert set(await _selected_session_ids(router)) == {"original", "replacement", "other"}
async def test_shared_session_pin_resumes_after_a_transient_redis_failure(
shared_session_router: tuple[litellm.Router, FakeServer, RedisCache],
) -> None:
router, server, redis = shared_session_router
await redis.async_set_cache(_repin_key(), {"model_id": "replacement"}, ttl=3600)
router.cache.in_memory_cache.set_cache(_repin_key(), {"model_id": "original"}, ttl=3600)
server.connected = False
assert await _selected_session_ids(router) == ("original",)
server.connected = True
assert await _selected_session_ids(router) == ("replacement",)
@pytest.mark.parametrize("shared_value", [None, {"model_id": "original"}])
async def test_delayed_shared_pin_read_preserves_a_newer_local_write(
shared_session_router: tuple[litellm.Router, FakeServer, RedisCache], shared_value: Mapping[str, str] | None
) -> None:
router, server, redis = shared_session_router
if shared_value is not None:
await redis.async_set_cache(_repin_key(), shared_value, ttl=3600)
observed: Final = asyncio.Event()
resume: Final = asyncio.Event()
class DelayedReader(Redis):
async def get(self, name: str) -> object:
value: Final = await super().get(name)
observed.set()
await resume.wait()
return value
reader: Final = DelayedReader(connection_pool=redis.init_async_client().connection_pool)
litellm.in_memory_llm_clients_cache.set_cache(key=redis._get_async_client_cache_key(), value=reader)
pending: Final = asyncio.create_task(_selected_session_ids(router))
try:
await observed.wait()
router.cache.in_memory_cache.set_cache(_repin_key(), {"model_id": "replacement"}, ttl=3600)
finally:
resume.set()
await pending
server.connected = False
assert await _selected_session_ids(router) == ("replacement",)
@pytest.mark.parametrize(
"response",
[
@ -298,10 +166,10 @@ async def test_delayed_shared_pin_read_preserves_a_newer_local_write(
async def test_unsuccessful_or_nested_response_envelopes_never_repin(response: Mapping[str, object]) -> None:
router: Final = _repinning_router()
try:
await router.cache.async_set_cache(_repin_key(), {"model_id": "original"})
await _pins(router).set(_repin_key(), {"model_id": "original"}, 3600)
_cool_down_original(router)
await _record_replacement_success(router, response=response)
assert await router.cache.async_get_cache(_repin_key()) == {"model_id": "original"}
assert await _pins(router).get(_repin_key()) == {"model_id": "original"}
finally:
router.discard()
@ -423,9 +291,9 @@ async def test_async_session_id_affinity_priority_over_user_key():
session_affinity should have priority. We test this by sending different
session ids for the same user.
"""
cache = DualCache()
pins = memory_pins()
callback = DeploymentAffinityCheck(
cache=cache,
pins=pins,
ttl_seconds=123,
enable_user_key_affinity=True,
enable_responses_api_affinity=False,
@ -445,14 +313,13 @@ async def test_async_session_id_affinity_priority_over_user_key():
},
]
await callback.cache.async_set_cache(
DeploymentAffinityCheck.get_affinity_cache_key("model_group", "user1"),
{"model_id": "deployment-1"},
await pins.set(
DeploymentAffinityCheck.get_affinity_cache_key("model_group", "user1"), {"model_id": "deployment-1"}, 60
)
await callback.cache.async_set_cache(
await pins.set(
DeploymentAffinityCheck.get_session_affinity_cache_key("model_group", "session1", user_key="user1"),
{"model_id": "deployment-2"},
60,
)
# Should use session mapping
@ -471,9 +338,9 @@ async def test_async_session_id_affinity_priority_over_user_key():
async def test_proxy_generated_session_id_does_not_pin_a_deployment():
"""A session id the proxy generated for a request that had none is per request, so a
pin stored under it must be ignored and none must be written."""
cache = DualCache()
pins = memory_pins()
callback = DeploymentAffinityCheck(
cache=cache,
pins=pins,
ttl_seconds=123,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
@ -483,9 +350,10 @@ async def test_proxy_generated_session_id_does_not_pin_a_deployment():
{"model_name": "model_group", "litellm_params": {"model": "model_1"}, "model_info": {"id": "deployment-1"}},
{"model_name": "model_group", "litellm_params": {"model": "model_2"}, "model_info": {"id": "deployment-2"}},
]
await cache.async_set_cache(
await pins.set(
DeploymentAffinityCheck.get_session_affinity_cache_key("model_group", "generated-1", user_key="user1"),
{"model_id": "deployment-2"},
60,
)
request_kwargs = {
"metadata": {"user_api_key_hash": "user1", "session_id": "generated-1", SESSION_ID_GENERATED_METADATA_KEY: True}
@ -503,7 +371,7 @@ async def test_proxy_generated_session_id_does_not_pin_a_deployment():
)
assert len(filtered) == 2
assert await cache.async_get_cache(
assert await pins.get(
DeploymentAffinityCheck.get_session_affinity_cache_key("model_group", "generated-1", user_key="user1")
) == {"model_id": "deployment-2"}
@ -569,6 +437,12 @@ def _smart_router(session_affinity=True, ttl_seconds=777, deployment_affinity=Tr
)
_TIER_PIN: Final = TierModelPin(
router="smart-router", tiers="tiers", caller=None, session_id="session", tier="SIMPLE", replay_fallback=None
)
_DEPLOYMENT_PIN: Final = DeploymentSessionPin(model_group="target-group", caller=None, session_id="session")
def _session_pin_key(session_id, user_key):
return DeploymentAffinityCheck.get_session_affinity_cache_key(
model_group="target-group", session_id=session_id, user_key=user_key
@ -609,14 +483,12 @@ async def test_auto_router_session_affinity_writes_scoped_pin_and_follows_it():
router = _smart_router()
try:
served = await _one_turn(router, "smart-router", "write-session", "key-1")
assert await router.cache.async_get_cache(key=_session_pin_key("write-session", "key-1")) == {
assert await _pins(router).get(_session_pin_key("write-session", "key-1")) == {
"model_id": served
}
assert await router.cache.async_get_cache(key=_session_pin_key("write-session", None)) is None
assert await _pins(router).get(_session_pin_key("write-session", None)) is None
await router.cache.async_set_cache(
key=_session_pin_key("read-session", "key-1"), value={"model_id": "deployment-2"}
)
await _pins(router).set(_session_pin_key("read-session", "key-1"), {"model_id": "deployment-2"}, 3600)
assert await _one_turn(router, "smart-router", "read-session", "key-1") == "deployment-2"
finally:
_cleanup_router_callbacks(router)
@ -639,9 +511,7 @@ async def test_seeded_session_pin_is_invisible_outside_its_scope(model, key_hash
for _ in range(20):
router = _smart_router()
try:
await router.cache.async_set_cache(
key=_session_pin_key("scoped-session", "key-1"), value={"model_id": "deployment-2"}
)
await _pins(router).set(_session_pin_key("scoped-session", "key-1"), {"model_id": "deployment-2"}, 3600)
served.add(await _one_turn(router, model, "scoped-session", key_hash))
finally:
_cleanup_router_callbacks(router)
@ -652,11 +522,10 @@ async def test_seeded_session_pin_is_invisible_outside_its_scope(model, key_hash
async def test_marker_write_uses_marker_ttl_and_writes_only_the_session_pin():
"""The write hook honors the marker's TTL over the callback default and writes
no user-key entry when only session affinity is active."""
import time as time_module
cache = DualCache()
clock = MagicMock(return_value=100.0)
pins = memory_pins(clock)
callback = DeploymentAffinityCheck(
cache=cache,
pins=pins,
ttl_seconds=3600,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
@ -676,16 +545,19 @@ async def test_marker_write_uses_marker_ttl_and_writes_only_the_session_pin():
)
session_key = _session_pin_key("ttl-session", "key-1")
assert cache.in_memory_cache.cache_dict == {session_key: {"model_id": "deployment-1"}}
assert cache.in_memory_cache.ttl_dict[session_key] == pytest.approx(time_module.time() + 777, abs=5)
assert pins.get_local(DeploymentPin(model_group="target-group", caller="key-1")) is None
clock.return_value = 876.0
assert pins.get_local(session_key) == {"model_id": "deployment-1"}
clock.return_value = 878.0
assert pins.get_local(session_key) is None
@pytest.mark.asyncio
@pytest.mark.parametrize("bad_marker", ["777", True, -5, 0, None])
async def test_malformed_marker_values_do_not_enable_session_affinity(bad_marker):
cache = DualCache()
pins = memory_pins()
callback = DeploymentAffinityCheck(
cache=cache,
pins=pins,
ttl_seconds=3600,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
@ -704,7 +576,8 @@ async def test_malformed_marker_values_do_not_enable_session_affinity(bad_marker
call_type=None,
)
assert cache.in_memory_cache.cache_dict == {}
assert pins.get_local(_session_pin_key("bad-marker-session", "key-1")) is None
assert pins.get_local(DeploymentPin(model_group="target-group", caller="key-1")) is None
@pytest.mark.asyncio
@ -713,9 +586,9 @@ async def test_concurrent_first_requests_never_flip_a_claimed_pin(enable_user_ke
"""Two overlapping first requests select different deployments before either
write lands. Pins are first-writer-wins claims, so the second write must leave
the stored pin unchanged instead of flipping it."""
cache = DualCache()
pins = memory_pins()
callback = DeploymentAffinityCheck(
cache=cache,
pins=pins,
ttl_seconds=3600,
enable_user_key_affinity=enable_user_key,
enable_responses_api_affinity=False,
@ -736,7 +609,7 @@ async def test_concurrent_first_requests_never_flip_a_claimed_pin(enable_user_ke
if enable_user_key
else _session_pin_key("racing-session", "key-1")
)
assert await cache.async_get_cache(key=pinned_key) == {"model_id": "deployment-1"}
assert await pins.get(pinned_key) == {"model_id": "deployment-1"}
@pytest.mark.asyncio
@ -751,102 +624,26 @@ async def test_in_memory_reclaim_slides_idle_window_only_for_the_stored_deployme
the value nor the expiry, so no-Redis setups keep stickiness across an active
session and ttl bounds idle time there too. Sameness is judged on the pinned
model id, so a legacy string pin written by the Redis branch slides the same."""
import time as time_module
cache = DualCache()
clock = MagicMock(return_value=100.0)
pins = memory_pins(clock)
callback = DeploymentAffinityCheck(
cache=cache,
pins=pins,
ttl_seconds=3600,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
)
pin_key = _session_pin_key("slide-session", "key-1")
cache.in_memory_cache.set_cache(pin_key, stored_pin, ttl=10)
first_expiry = cache.in_memory_cache.ttl_dict[pin_key]
pins.set_local(pin_key, stored_pin, 10)
reclaimed = await callback._claim_pin(cache_key=pin_key, pin_value={"model_id": "deployment-1"}, ttl_seconds=777)
assert reclaimed == "deployment-1"
assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5)
assert cache.in_memory_cache.ttl_dict[pin_key] > first_expiry
lost = await callback._claim_pin(cache_key=pin_key, pin_value={"model_id": "deployment-2"}, ttl_seconds=10)
assert lost == "deployment-1"
assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5)
@pytest.mark.asyncio
async def test_claim_pin_uses_redis_attached_after_construction():
"""The proxy attaches Redis via Router._update_redis_cache after the Router (and
this callback) are built. The claim must resolve the redis tier per call, or pins
silently stay pod-local and cross-pod first-writer-wins is lost."""
cache = DualCache()
callback = DeploymentAffinityCheck(
cache=cache,
ttl_seconds=3600,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
)
captured = {}
async def fake_runner(keys, args, client=None):
captured["keys"] = keys
captured["args"] = args
return b'{"model_id": "other-pod-winner"}'
late_redis = MagicMock()
late_redis.async_register_script = MagicMock(return_value=fake_runner)
cache.redis_cache = late_redis
import time as time_module
pin_key = _session_pin_key("late-redis-session", "key-1")
cache.in_memory_cache.set_cache(pin_key, {"model_id": "other-pod-winner"}, ttl=10)
claimed = await callback._claim_pin(
cache_key=pin_key,
pin_value={"model_id": "our-deployment"},
ttl_seconds=777,
)
assert claimed == "other-pod-winner"
assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5)
assert captured["keys"] == (pin_key,)
assert captured["args"] == ('{"model_id": "our-deployment"}', 777)
assert cache.in_memory_cache.get_cache(_session_pin_key("late-redis-session", "key-1")) == {
"model_id": "other-pod-winner"
}
@pytest.mark.asyncio
async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down():
"""A Redis outage must cost cross-pod agreement, never same-pod stickiness. The write
hook only logs this result, so an escaping error would leave the session unpinned and
reshuffle every turn for the whole outage. DualCache's write path, which this claim
replaced, wrote the in-memory tier before ever touching Redis."""
cache = DualCache()
callback = DeploymentAffinityCheck(
cache=cache,
ttl_seconds=3600,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
)
async def exploding_runner(keys, args, client=None):
raise ConnectionError("redis is down")
down_redis = MagicMock()
down_redis.async_register_script = MagicMock(return_value=exploding_runner)
cache.redis_cache = down_redis
key = _session_pin_key("outage-session", "key-1")
claimed = await callback._claim_pin(cache_key=key, pin_value={"model_id": "our-deployment"}, ttl_seconds=777)
assert claimed == "our-deployment"
assert cache.in_memory_cache.get_cache(key) == {"model_id": "our-deployment"}
second = await callback._claim_pin(cache_key=key, pin_value={"model_id": "another-deployment"}, ttl_seconds=777)
assert second == "our-deployment"
clock.return_value = 876.0
assert pins.get_local(pin_key) == {"model_id": "deployment-1"}
clock.return_value = 878.0
assert pins.get_local(pin_key) is None
@pytest.mark.asyncio
@ -854,7 +651,6 @@ async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down():
("stored", "expected"),
[
({"model": "first"}, {"model": "first"}),
('{ "model" : "first" }', {"model": "first"}),
({"model": "removed"}, {"model": "second"}),
({"model": "first", "extra": "stale"}, {"model": "second"}),
({"model_id": "first"}, {"model": "second"}),
@ -864,35 +660,35 @@ async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down():
)
async def test_eligible_affinity_claim_replaces_stale_pins_and_slides_ttl(stored: object, expected: object) -> None:
clock: Final = MagicMock(return_value=100.0)
cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock))
cache.in_memory_cache.set_cache("tier-pin", stored, ttl=10)
pins: Final = memory_pins(clock)
if stored is not None:
pins.set_local(_TIER_PIN, stored, 10)
clock.return_value = 105.0
winner: Final = await claim_affinity_pin(
cache,
"tier-pin",
pins,
_TIER_PIN,
{"model": "second"},
30,
eligible_values=({"model": "first"}, {"model": "second"}),
)
assert winner == expected
assert cache.in_memory_cache.ttl_dict["tier-pin"] == 135.0
clock.return_value = 111.0
assert cache.in_memory_cache.get_cache("tier-pin") == expected
assert pins.get_local(_TIER_PIN) == expected
clock.return_value = 136.0
assert cache.in_memory_cache.get_cache("tier-pin") is None
assert pins.get_local(_TIER_PIN) is None
@pytest.mark.asyncio
async def test_concurrent_eligible_claims_return_one_winner() -> None:
cache: Final = DualCache()
pins: Final = memory_pins()
candidates: Final = ({"model": "first"}, {"model": "second"})
winners: Final = await asyncio.gather(
*(
claim_affinity_pin(
cache,
"tier-pin",
pins,
_TIER_PIN,
candidates[index % 2],
30,
eligible_values=candidates,
@ -902,7 +698,7 @@ async def test_concurrent_eligible_claims_return_one_winner() -> None:
)
assert winners == [{"model": "first"}] * 20
assert cache.in_memory_cache.get_cache("tier-pin") == {"model": "first"}
assert pins.get_local(_TIER_PIN) == {"model": "first"}
@pytest.mark.asyncio
@ -919,54 +715,22 @@ async def test_legacy_deployment_claim_retains_decoder_and_keepalive(
stored: object, expected: str | None, refresh: bool
) -> None:
clock: Final = MagicMock(return_value=100.0)
cache: Final = DualCache(in_memory_cache=InMemoryCache(clock=clock))
pins: Final = memory_pins(clock)
callback: Final = DeploymentAffinityCheck(
cache=cache,
pins=pins,
ttl_seconds=30,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
)
cache.in_memory_cache.set_cache("deployment-pin", stored, ttl=10)
pins.set_local(_DEPLOYMENT_PIN, stored, 10)
clock.return_value = 105.0
winner: Final = await callback._claim_pin("deployment-pin", {"model_id": "7"}, 30)
winner: Final = await callback._claim_pin(_DEPLOYMENT_PIN, {"model_id": "7"}, 30)
assert winner == expected
assert cache.in_memory_cache.ttl_dict["deployment-pin"] == (135.0 if refresh else 110.0)
assert cache.in_memory_cache.get_cache("deployment-pin") == ({"model_id": "7"} if refresh else stored)
@pytest.mark.asyncio
@pytest.mark.parametrize(
("raw", "expected", "stored"),
[
(b'{"model_id": "winner"}', "winner", {"model_id": "winner"}),
('"winner"', "winner", "winner"),
("winner", "winner", "winner"),
(b"winner", "winner", "winner"),
('{"model": "winner"}', None, {"model": "winner"}),
(None, "candidate", None),
(123, "candidate", None),
({"model_id": "winner"}, "candidate", None),
],
)
async def test_redis_deployment_claim_preserves_legacy_result_decoding(
raw: object, expected: str | None, stored: object
) -> None:
redis: Final = MagicMock()
redis.async_register_script.return_value = AsyncMock(return_value=raw)
cache: Final = DualCache(redis_cache=redis)
callback: Final = DeploymentAffinityCheck(
cache=cache,
ttl_seconds=30,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
)
winner: Final = await callback._claim_pin("deployment-pin", {"model_id": "candidate"}, 30)
assert winner == expected
assert cache.in_memory_cache.get_cache("deployment-pin") == stored
assert pins.get_local(_DEPLOYMENT_PIN) == ({"model_id": "7"} if refresh else stored)
clock.return_value = 120.0
assert pins.get_local(_DEPLOYMENT_PIN) == ({"model_id": "7"} if refresh else None)
@pytest.mark.asyncio
@ -974,7 +738,7 @@ async def test_marker_session_affinity_read_and_write_agree_for_wildcard_groups(
"""Wildcard deployments keep the literal pattern as model_name on both the read
path and the write path, so the marker-gated pin round-trips through one key."""
callback = DeploymentAffinityCheck(
cache=DualCache(),
pins=memory_pins(),
ttl_seconds=3600,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,

2
tests/unit/test_router/core/test_fallbacks_walk_parity.py
View file

@ -269,7 +269,7 @@ async def test_weighted_failover_skips_an_affinity_pin_on_the_failed_deployment(
if isinstance(callback, DeploymentAffinityCheck)
)
key: Final = DeploymentAffinityCheck.get_session_affinity_cache_key("g", "s", "c")
await check.cache.async_set_cache(key, {"model_id": "a"}) # pyright: ignore[reportUnknownMemberType]
await check.pins.set(key, {"model_id": "a"}, 3600)
host.answer("openai/a", server_error())
response: Final = await complete(core, "g", metadata={"session_id": "s", "user_api_key_hash": "c"})

7
tests/unit/test_router/core/test_selection_characterization_session_repin.py
View file

@ -38,8 +38,7 @@ def _affinity_check(core: RouterCore) -> DeploymentAffinityCheck:
async def _pinned(check: DeploymentAffinityCheck) -> object:
key = DeploymentAffinityCheck.get_session_affinity_cache_key("g", _SESSION, _CALLER)
pinned: object = await check.cache.async_get_cache(key) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # DualCache's reads are untyped
return pinned # pyright: ignore[reportUnknownVariableType] # see above
return await check.pins.get(key)
async def _complete(core: RouterCore) -> object:
@ -62,7 +61,7 @@ async def test_a_failover_success_repins_the_session_only_with_repin_on_failover
)
check = _affinity_check(core)
key = DeploymentAffinityCheck.get_session_affinity_cache_key("g", _SESSION, _CALLER)
await check.cache.async_set_cache(key, {"model_id": "a"}) # pyright: ignore[reportUnknownMemberType]
await check.pins.set(key, {"model_id": "a"}, 3600)
host.answer("openai/a", rate_limit(), completion_response())
host.answer("openai/b", completion_response())
@ -98,7 +97,7 @@ async def test_a_failover_stream_repins_the_session_when_it_ends(
)
check = _affinity_check(core)
key = DeploymentAffinityCheck.get_session_affinity_cache_key("g", _SESSION, _CALLER)
await check.cache.async_set_cache(key, {"model_id": "a"}) # pyright: ignore[reportUnknownMemberType]
await check.pins.set(key, {"model_id": "a"}, 3600)
host.answer("openai/a", rate_limit())
host.answer(
"openai/b",

6
tests/unit/test_router/test_router.py
View file

@ -387,11 +387,7 @@ async def test_encrypted_content_affinity_takes_priority_over_user_key_affinity(
model_group=model_group,
user_key=user_api_key_hash,
)
await deployment_callback.cache.async_set_cache(
key=cache_key,
value={"model_id": "deployment-a"},
ttl=60,
)
await deployment_callback.pins.set(cache_key, {"model_id": "deployment-a"}, 60)
encoded_id = ResponsesAPIRequestUtils._build_encrypted_item_id("deployment-b", "rs_test")
request_kwargs = {
"input": [{"type": "reasoning", "id": encoded_id}],