diff --git a/litellm/caching/affinity_cache.py b/litellm/caching/affinity_cache.py new file mode 100644 index 00000000000..2712679b99b --- /dev/null +++ b/litellm/caching/affinity_cache.py @@ -0,0 +1,125 @@ +"""Atomic affinity claims shared by deployment and tier-model selection.""" + +import json +from collections.abc import Mapping +from typing import ( + Final, + cast, # noqa: TID251 # Redis script results are narrowed only to object, then validated +) + +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm._logging import verbose_router_logger +from litellm.caching.dual_cache import DualCache + +_PIN_JSON_ADAPTER: Final = TypeAdapter[JsonValue](JsonValue) + +_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 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 +""" + + +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 _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, +) -> 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 + winner: Final = existing if existing is not None and existing in (eligible_values or ()) 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 claim_affinity_pin( + cache: DualCache, + cache_key: str, + pin_value: Mapping[str, str], + ttl_seconds: int, + *, + eligible_values: tuple[Mapping[str, str], ...] | None = None, +) -> object: + """Return the authoritative first writer, replacing it only when it becomes ineligible. + + 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: + try: + claim_script: Final = redis_cache.async_register_script(_CLAIM_PIN_SCRIPT) + args: Final = ( + json.dumps(dict(pin_value)), # mutable-ok: JSON serialization requires dict, not a generic Mapping + int(ttl_seconds), + *( + (json.dumps(tuple(dict(value) for value in eligible_values)),) # mutable-ok: JSON requires dict + if eligible_values is not None + else () + ), + ) + 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, cache_key, pin_value, ttl_seconds, eligible_values=eligible_values) diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 9deccc9a468..1000f0f479f 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -16,6 +16,8 @@ Inspired by ClawRouter: https://github.com/BlockRunAI/ClawRouter from __future__ import annotations import asyncio +import hashlib +import json import random import re import time @@ -28,6 +30,7 @@ from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple, cast from pydantic import BaseModel, TypeAdapter, ValidationError, create_model from litellm._logging import verbose_router_logger +from litellm.caching.affinity_cache import claim_affinity_pin from litellm.constants import ( EMPTY_MAPPING, INTERNAL_CALL_ORIGIN_METADATA_KEY, @@ -55,6 +58,7 @@ from litellm.router_strategy.complexity_router.tier_predictor import ( TierSuccessPredictor, resolve_tier_artifact, ) +from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionImageObject, @@ -1119,10 +1123,10 @@ class _ContextWindowPlacement(NamedTuple): class _SessionAffinityPin(NamedTuple): model: str - tier: ComplexityTier | None + tier: ComplexityTier | str | None -def _parse_session_affinity_pin(value: object) -> _SessionAffinityPin | None: +def _parse_session_affinity_pin(value: object, active_tiers: tuple[str, ...]) -> _SessionAffinityPin | None: if isinstance(value, str): return _SessionAffinityPin(model=value, tier=None) parts: Final[tuple[object, object] | None] = ( @@ -1137,8 +1141,11 @@ def _parse_session_affinity_pin(value: object) -> _SessionAffinityPin | None: model, tier_value = parts if not isinstance(model, str): return None - tier: Final = ComplexityTier(tier_value) if isinstance(tier_value, str) else None - return _SessionAffinityPin(model=model, tier=tier) + if tier_value is None: + return _SessionAffinityPin(model=model, tier=None) + if not isinstance(tier_value, str) or tier_value not in active_tiers: + return None + return _SessionAffinityPin(model=model, tier=_built_in_tier_or_none(tier_value) or tier_value) def _session_affinity_cache_value(model: str, tier: ComplexityTier | str | None) -> Mapping[str, str | None]: @@ -1195,6 +1202,10 @@ class ComplexityRouter(CustomLogger): if default_model: self.config.default_model = default_model + self._tier_affinity_config = hashlib.sha256( + self.config.model_dump_json(include=MappingProxyType({"tiers": True, "tier_model_configs": True})).encode() + ).hexdigest() + # Checked here rather than on the config model because the deployment's # complexity_router_default_model arrives outside complexity_router_config and is # applied just above, so a validator on the model would reject a deployment that @@ -2259,6 +2270,51 @@ class ComplexityRouter(CustomLogger): def _tier_pools(self) -> dict[str, list[str]]: return {tier: (models if isinstance(models, list) else [models]) for tier, models in self.config.tiers.items()} + async def _pin_model_for_tier( + self, + tier: ComplexityTier | str, + model: str, + candidates: tuple[str, ...], + request_kwargs: dict[str, object], # mutable-ok: adaptive feedback metadata must follow the selected model + retained_pin: _SessionAffinityPin | None = None, + ) -> str: + if not self._uses_deployment_pin or model not in candidates: + return model + retained_model: Final = ( + retained_pin.model + if retained_pin is not None + and retained_pin.tier is not None + and _tier_name(retained_pin.tier) == _tier_name(tier) + else None + ) + if retained_model is not None and retained_model in candidates: + self._restamp_adaptive_choice(request_kwargs, model, retained_model) + return retained_model + 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() + ) + winner: Final = await claim_affinity_pin( + self.litellm_router_instance.cache, + cache_key, + MappingProxyType({"model": model}), + self.config.session_affinity_ttl_seconds, + eligible_values=tuple(MappingProxyType({"model": candidate}) for candidate in candidates), + ) + pinned: Final[object] = winner.get("model") if isinstance(winner, Mapping) else None + if not isinstance(pinned, str) or pinned not in candidates: + return model + self._restamp_adaptive_choice(request_kwargs, model, pinned) + return pinned + async def _pick_model_for_tier( self, tier: ComplexityTier | str, @@ -2266,11 +2322,18 @@ class ComplexityRouter(CustomLogger): resolved_messages: list[dict[str, Any]] | None, request_kwargs: dict, allowed_models: tuple[str, ...] | None = None, + retained_pin: _SessionAffinityPin | None = None, ) -> str: if not self.config.plugins: - if allowed_models is not None: - return self._pick_from_tier_value(allowed_models, _tier_name(tier)) - return self.get_model_for_tier(tier) + candidates: Final = ( + allowed_models if allowed_models is not None else tuple(self._tier_pools().get(_tier_name(tier), ())) + ) + selected: Final = ( + self._pick_from_tier_value(allowed_models, _tier_name(tier)) + if allowed_models is not None + else self.get_model_for_tier(tier) + ) + return await self._pin_model_for_tier(tier, selected, candidates, request_kwargs, retained_pin) from litellm.types.router import RoutingContext @@ -2369,6 +2432,40 @@ class ComplexityRouter(CustomLogger): self._adaptive_chosen_model_key = ADAPTIVE_ROUTER_CHOSEN_MODEL_KEY return self.adaptive_router + def _adaptive_candidate_models( + self, + classified_tier: ComplexityTier | str, + hard_floor: ComplexityTier | str | None = None, + hard_ceiling: ComplexityTier | str | None = None, + fit_filter: frozenset[str] | None = None, + ) -> tuple[str, ...]: + pools: Final = self._tier_pools() + candidates: Final = ( + tuple(pools.get(_tier_name(classified_tier), ())) + if self.config.adaptive_eligible == "classified_tier" + else tuple(dict.fromkeys(chain.from_iterable(pools.values()))) + ) + floor: Final = self._active_tier_severity(hard_floor) if hard_floor is not None else None + ceiling: Final = self._active_tier_severity(hard_ceiling) if hard_ceiling is not None else None + return tuple( + model + for model in _allowed(candidates, fit_filter) + if ( + floor is None + or any( + self._active_tier_severity(tier) >= floor + for tier in self._model_tiers.get(model, (classified_tier,)) + ) + ) + and ( + ceiling is None + or any( + self._active_tier_severity(tier) <= ceiling + for tier in self._model_tiers.get(model, (classified_tier,)) + ) + ) + ) + def _soft_floor_pick( self, classified_tier: ComplexityTier | str, @@ -2436,34 +2533,17 @@ class ComplexityRouter(CustomLogger): ], } return chosen_model - if self.config.adaptive_eligible == "classified_tier": - candidates = list(classified_candidates) - if not candidates: - return self._fitting_tier_fallback(classified_tier, fit_filter) - else: - candidates = list(_allowed(tuple(adaptive.config.available_models), fit_filter)) + candidates: Final = self._adaptive_candidate_models(classified_tier, fit_filter=fit_filter) all_costs: Final = [adaptive.model_to_cost.get(m, 0.0) for m in candidates] quality_weight: Final = self.config.adaptive_weights.quality cost_weight: Final = self.config.adaptive_weights.cost penalty_weight: Final = self.config.tier_distance_penalty - floor_severity: Final = self._active_tier_severity(hard_floor) if hard_floor is not None else None - ceiling_severity: Final = self._active_tier_severity(hard_ceiling) if hard_ceiling is not None else None best_model: str | None = None best_score = float("-inf") candidate_scores: Final[list[dict[str, object]]] = [] - for model in candidates: - if floor_severity is not None and all( - self._active_tier_severity(model_tier) < floor_severity - for model_tier in self._model_tiers.get(model, (classified_tier,)) - ): - continue - if ceiling_severity is not None and all( - self._active_tier_severity(model_tier) > ceiling_severity - for model_tier in self._model_tiers.get(model, (classified_tier,)) - ): - continue + for model in self._adaptive_candidate_models(classified_tier, hard_floor, hard_ceiling, fit_filter): cell = adaptive._cells[(request_type, model)] quality_sample = thompson_sample(cell) cost_score = normalized_cost(adaptive.model_to_cost.get(model, 0.0), all_costs) @@ -2644,8 +2724,6 @@ class ComplexityRouter(CustomLogger): """Prompt content the resolved message list never carries: the Responses API's `instructions`, the /v1/messages top-level `system` block, and tool definitions. A coding agent's context is dominated by these.""" - import json - instructions: Final = request_kwargs.get("instructions") proxy_request: Final = request_kwargs.get("proxy_server_request") body: Final = proxy_request.get("body") if isinstance(proxy_request, Mapping) else None @@ -2831,19 +2909,21 @@ class ComplexityRouter(CustomLogger): ) return higher_tiers[0] if higher_tiers else tier - def _escalated_pin(self, pinned_model: str) -> str | None: + def _escalated_pin(self, pinned_model: str, tier: ComplexityTier | str | None = None) -> _SessionAffinityPin | None: """Bump a session's pinned model to the next-higher configured tier. Returns None when the pin no longer maps to any configured tier, signalling a full reclassification instead. """ - pinned_tier: Final = self._tier_for_model(pinned_model) + pinned_tier: Final = tier if tier is not None else self._tier_for_model(pinned_model) if pinned_tier is None: return None escalated_tier: Final = self._escalate_tier(pinned_tier) if escalated_tier == pinned_tier: - return pinned_model - return self.get_model_for_tier(escalated_tier) + return _SessionAffinityPin(pinned_model, pinned_tier) + return _SessionAffinityPin( + self.get_model_for_tier(escalated_tier), _built_in_tier_or_none(_tier_name(escalated_tier)) + ) def _vision_verdicts(self, model_name: str) -> tuple[bool | None, ...]: """Declared vision support per deployment serving the name: True, False, or None when @@ -2907,6 +2987,7 @@ class ComplexityRouter(CustomLogger): resolved_messages: Sequence[Mapping[str, object]] | None, request_kwargs: dict, # mutable-ok: same shape the hook receives context_fit: _RequestContextFit | None = None, + retained_pin: _SessionAffinityPin | None = None, ) -> PreRoutingHookResponse: """Replace a routed model that cannot accept this request's image input. @@ -2955,6 +3036,7 @@ class ComplexityRouter(CustomLogger): repick_messages, # pyright: ignore[reportArgumentType] # hook-resolved message dicts; the pick only reads them request_kwargs, allowed_models=tuple(entry for entry in pools.get(capable, ()) if entry in eligible), + retained_pin=retained_pin, ) elif self._modality_default_model_usable(request_kwargs, resolved_messages, eligible): new_tier = None @@ -3098,6 +3180,7 @@ class ComplexityRouter(CustomLogger): resolved_messages: Sequence[Mapping[str, object]] | None, request_kwargs: dict, # mutable-ok: same shape the hook receives context_fit: _RequestContextFit | None = None, + retained_pin: _SessionAffinityPin | None = None, ) -> PreRoutingHookResponse: """Try compatible tier recovery before the default, preserving request policy and fit.""" decision: Final = response.routing_decision @@ -3155,6 +3238,7 @@ class ComplexityRouter(CustomLogger): repick_messages, # pyright: ignore[reportArgumentType] # hook-resolved message dicts; the pick only reads them request_kwargs, allowed_models=live, + retained_pin=retained_pin, ) except ValueError as exc: verbose_router_logger.debug( @@ -3247,8 +3331,13 @@ class ComplexityRouter(CustomLogger): """The adaptive feedback loop reads its chosen-model marker from request metadata; a gate rewrite must move the marker with the model or rewards land on the displaced one.""" metadata: Final = request_kwargs.get("metadata") - if isinstance(metadata, dict) and metadata.get("adaptive_router_chosen_model") == old_model: + if not isinstance(metadata, dict): + return + if metadata.get("adaptive_router_chosen_model") == old_model: metadata["adaptive_router_chosen_model"] = new_model + decision: Final = metadata.get("adaptive_router_decision") + if isinstance(decision, dict) and decision.get("chosen_model") == old_model: + decision["chosen_model"] = new_model def _lexical_tier_override(self, user_message: str) -> KeywordOverride | None: """When keyword_tier_rules match literally, the most-severe matched tier wins. @@ -3561,25 +3650,42 @@ class ComplexityRouter(CustomLogger): if cache_key is not None and pin_replay_allowed: pinned_value: Final = await self.litellm_router_instance.cache.async_get_cache(key=cache_key) - pinned_pin: Final = _parse_session_affinity_pin(pinned_value) + pinned_pin: Final = _parse_session_affinity_pin(pinned_value, self.config.tier_names()) if pinned_pin is not None: - routed_model: str | None = pinned_pin.model - pin_escalation_keyword: str | None = None - if self.escalation_keywords: - user_message: Final = ( - _newest_turn_ask(resolved_messages, marker_pairs) if resolved_messages else None + user_message: Final = _newest_turn_ask(resolved_messages, marker_pairs) if resolved_messages else None + pin_escalation_keyword: Final = ( + self._matched_escalation_keyword(user_message) if user_message is not None else None + ) + selected_pin: Final = ( + self._escalated_pin(pinned_pin.model, pinned_pin.tier) + if pin_escalation_keyword is not None + else _SessionAffinityPin( + pinned_pin.model, + pinned_pin.tier if pinned_pin.tier is not None else self._tier_for_model(pinned_pin.model), ) - if user_message is not None: - pin_escalation_keyword = self._matched_escalation_keyword(user_message) - if pin_escalation_keyword is not None: - routed_model = self._escalated_pin(pinned_pin.model) - if routed_model is not None: - escalated: Final = routed_model != pinned_pin.model - resolved_pin_tier: Final = ( - pinned_pin.tier - if not escalated and pinned_pin.tier is not None - else self._tier_for_model(routed_model) + ) + if selected_pin is not None: + escalated: Final = selected_pin.model != pinned_pin.model or ( + pin_escalation_keyword is not None + and pinned_pin.tier is not None + and selected_pin.tier != pinned_pin.tier ) + resolved_pin_tier: Final = selected_pin.tier + session_model: Final = ( + await self._pin_model_for_tier( + resolved_pin_tier, + selected_pin.model, + tuple(self._tier_pools().get(_tier_name(resolved_pin_tier), ())), + request_kwargs, + ) + if escalated and resolved_pin_tier is not None + else selected_pin.model + ) + retained_pin: Final = _SessionAffinityPin(session_model, resolved_pin_tier) + if resolved_pin_tier is not None: + await self._pin_model_for_tier( + resolved_pin_tier, session_model, (session_model,), request_kwargs + ) # The floor outranks the pin because plan mode is a transient state of the # session, not a request to move it: the turns carrying the sentinel route at # the floor, and the stored pin deliberately keeps the session's own model so @@ -3590,16 +3696,28 @@ class ComplexityRouter(CustomLogger): plan_floored: Final = ( pinned_tier is not None and self._apply_plan_mode_floor(pinned_tier) != pinned_tier ) - session_model: Final = routed_model - if plan_floored and pinned_tier is not None: - routed_model = self.get_model_for_tier(self._apply_plan_mode_floor(pinned_tier)) - pin_source_tier: Final = self._tier_for_model(routed_model) + floor_model: Final = ( + await self._pick_model_for_tier( + self._apply_plan_mode_floor(pinned_tier), + messages, + resolved_messages, + request_kwargs, + retained_pin=retained_pin, + ) + if plan_floored and pinned_tier is not None + else session_model + ) + pin_source_tier: Final = ( + self._apply_plan_mode_floor(pinned_tier) + if plan_floored and pinned_tier is not None + else resolved_pin_tier + ) pin_placement: Final = ( await self._context_window_placement( pin_source_tier, resolved_messages, request_kwargs, - pool_override=(routed_model,), + pool_override=(floor_model,), context_fit=context_fit, ) if pin_source_tier is not None @@ -3612,11 +3730,18 @@ class ComplexityRouter(CustomLogger): and _tier_name(pin_placement.tier) != _tier_name(pin_source_tier) else None ) - if pin_placement is not None and pin_context_original_tier is not None: - # The stored pin below keeps the session's own model on purpose. - routed_model = self._pick_from_tier_value( - pin_placement.allowed_models, _tier_name(pin_placement.tier) + routed_model: Final = ( + await self._pick_model_for_tier( + pin_placement.tier, + messages, + resolved_messages, + request_kwargs, + allowed_models=pin_placement.allowed_models, + retained_pin=retained_pin, ) + if pin_placement is not None and pin_context_original_tier is not None + else floor_model + ) # Refresh the TTL on every hit so an active session doesn't lose its # pin mid-conversation just because it outlives the original write. await self.litellm_router_instance.cache.async_set_cache( @@ -3644,7 +3769,7 @@ class ComplexityRouter(CustomLogger): routed_pin_tier: Final = ( pin_placement.tier if pin_placement is not None and pin_context_original_tier is not None - else (self._tier_for_model(routed_model) if plan_floored else resolved_pin_tier) + else pin_source_tier ) session_tier_litellm_params: Final = self._litellm_params_for_model(routed_pin_tier, routed_model) has_original_messages: Final = messages is not None and len(messages) > 0 @@ -3671,12 +3796,14 @@ class ComplexityRouter(CustomLogger): resolved_messages, request_kwargs, context_fit, + retained_pin, ), messages, input, resolved_messages, request_kwargs, context_fit, + retained_pin, ) ) @@ -3961,13 +4088,21 @@ class ComplexityRouter(CustomLogger): housekeeping_ceiling: Final = tier if outcome.cause == "housekeeping" else None # A context-escalated tier becomes the hard floor: a floor the bandit can slide # under is not a floor. - routed_model = self._soft_floor_pick( + adaptive_floor: Final = tier if context_original_tier is not None else plan_floor + adaptive_fit: Final = context_placement.holdable_models if context_placement is not None else None + sampled_model: Final = self._soft_floor_pick( tier, ask, request_kwargs, - hard_floor=tier if context_original_tier is not None else plan_floor, + hard_floor=adaptive_floor, hard_ceiling=housekeeping_ceiling, - fit_filter=context_placement.holdable_models if context_placement is not None else None, + fit_filter=adaptive_fit, + ) + routed_model = await self._pin_model_for_tier( # rebind-ok: reuse the eligible tier winner + tier, + sampled_model, + self._adaptive_candidate_models(tier, adaptive_floor, housekeeping_ceiling, adaptive_fit), + request_kwargs, ) adaptive: Final = self._ensure_adaptive_router() if adaptive is not None: diff --git a/litellm/router_strategy/complexity_router/config.py b/litellm/router_strategy/complexity_router/config.py index 1f1b5a5cc4b..bfc83f8dcab 100644 --- a/litellm/router_strategy/complexity_router/config.py +++ b/litellm/router_strategy/complexity_router/config.py @@ -1256,20 +1256,16 @@ class ComplexityRouterConfig(BaseModel): deployment_affinity: bool = Field( default=True, description=( - "When True and a session_id is resolvable on the request, pin the deployment chosen " - "inside each routed model group and reuse it whenever the session returns to that " - "group, without pinning which group the session routes to. Independent of " - "session_affinity, which pins the model group instead (and always carries this " - "deployment pin with it): with session_affinity off, " - "every turn is still classified on its own merits while a session that escalates to a " - "stronger tier and comes back still lands on the deployment it used before, which is " - "what keeps a provider prompt cache warm. Pins are held per model group, so switching " - "tiers does not disturb the pin left behind in the previous group. On by default " - "because re-shuffling a conversation across deployments of the same model discards " - "that cache for no benefit; set False to keep every turn load-balanced across the " - "group, which is what a deployment set with tight per-deployment rate limits wants. " - "Inert when no session_id is resolvable, since there is nothing to key a pin on, and " - "suppressed when plugins are configured, for the same reason session_affinity is." + "When True and a client session_id is resolvable, reuse the session's chosen model " + "for each classified tier and its deployment within each model group. With " + "session_affinity off, every turn is still classified: moving to another tier leaves " + "the previous tier's model pin intact for a later return. Pins yield to current " + "candidate, context, modality, and availability constraints. Adaptive selection chooses " + "the initial model from its eligible pool, then reuses that choice per tier. This " + "reduces avoidable provider prompt-cache misses; it does not guarantee cache hits. " + "Set False to select models and load-balance deployments on every turn, unless " + "session_affinity or user_turn classification requires a pin. Inert without a client " + "session_id and suppressed when plugins are configured." ), ) session_affinity_ttl_seconds: int = Field( @@ -1277,7 +1273,7 @@ class ComplexityRouterConfig(BaseModel): gt=0, description=( "TTL for the session affinity pin; refreshed on every cache hit. Bounds both the " - "session_affinity model pin and the deployment_affinity deployment pin, so it measures " + "session_affinity model pin and the deployment_affinity per-tier model and deployment pins, so it measures " "idle time for the session's routing decisions rather than total session length" ), ) diff --git a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py index c7eb46046ef..3b88ac2eb00 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -13,13 +13,13 @@ where routing to a consistent deployment is still beneficial. """ import hashlib -import json from collections.abc import Mapping, Sequence from typing import Any, Final, cast -from typing_extensions import TypedDict +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_router_logger +from litellm.caching.affinity_cache import 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 @@ -28,8 +28,8 @@ from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes -class DeploymentAffinityCacheValue(TypedDict): - model_id: str +class DeploymentAffinityCacheValue(TypedDict, closed=True): + model_id: ReadOnly[str] VALID_MODEL_GROUP_AFFINITY_FLAGS: Final = frozenset( @@ -60,19 +60,6 @@ def warn_on_unknown_model_group_affinity_flags(model_group_affinity_config: Mapp ) -_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 current == ARGV[1] then - redis.call('EXPIRE', KEYS[1], ARGV[2]) -end -return current -""" - - class DeploymentAffinityCheck(CustomLogger): """ Router deployment affinity callback. @@ -255,34 +242,33 @@ class DeploymentAffinityCheck(CustomLogger): return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{hashed_user_key}:{session_id}" @staticmethod - def _get_session_id_from_metadata_dict(metadata: dict) -> str | None: + def _get_session_id_from_metadata_dict(metadata: Mapping[object, object]) -> str | None: session_id: Final = metadata.get("session_id") if session_id is None or metadata.get(SESSION_ID_GENERATED_METADATA_KEY): return None return str(session_id) @staticmethod - def _iter_metadata_dicts(request_kwargs: dict) -> list[dict]: + def _iter_metadata_dicts(request_kwargs: Mapping[str, object]) -> tuple[Mapping[object, object], ...]: """ Return all metadata dicts available on the request. Depending on the endpoint, Router may populate `metadata` or `litellm_metadata`. Users may also send one or both, so we check both (rather than using `or`). """ - metadata_dicts: Final[list[dict]] = [] - for key in ("litellm_metadata", "metadata"): - md = request_kwargs.get(key) - if isinstance(md, dict): - metadata_dicts.append(md) - return metadata_dicts + return tuple( + cast(Mapping[object, object], metadata) # cast-ok: isinstance proves mapping shape; values remain opaque + for key in ("litellm_metadata", "metadata") + if isinstance(metadata := request_kwargs.get(key), dict) + ) @staticmethod - def _first_metadata_value(metadata_dicts: Sequence[dict], key: str) -> str | None: + def _first_metadata_value(metadata_dicts: Sequence[Mapping[object, object]], key: str) -> str | None: value: Final = next((metadata[key] for metadata in metadata_dicts if metadata.get(key) is not None), None) return None if value is None else str(value) @classmethod - def _get_user_key_from_request_kwargs(cls, request_kwargs: dict) -> str | None: + def get_user_key_from_request_kwargs(cls, request_kwargs: Mapping[str, object]) -> str | None: """ Extract a stable affinity key from request kwargs. @@ -334,74 +320,17 @@ class DeploymentAffinityCheck(CustomLogger): return None def _set_local_pin(self, cache_key: str, value: object, ttl_seconds: int) -> None: - """The one owner of authoritative local pin writes: a plain set keeps a live - key's original expiry (`allow_ttl_override`), so the entry is replaced to make - the TTL real. Every local pin write goes through here so the redis-winner sync - and the pod-local claim can never disagree about expiry again.""" - self.cache.in_memory_cache.delete_cache(cache_key) - self.cache.in_memory_cache.set_cache(cache_key, value, ttl=ttl_seconds) + set_local_affinity_pin(self.cache, cache_key, value, ttl_seconds) async def _claim_pin(self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int) -> str | None: - """First-writer-wins pin write: store `pin_value` only when the key is absent and - return the deployment id the key holds afterwards, so a caller learns whether it won - by comparing against its own id, and None when the stored value is one no reader can - interpret. Concurrent claimers converge on the - first write instead of the last. Re-claiming with the stored value refreshes its - TTL, the same keepalive the complexity router's model pin documents: an active - session must not lose its pin mid-conversation just because it outlives the - original write, so the affinity TTL (the Router's - `deployment_affinity_ttl_seconds`, or a pre-routing hook's per-request - `session_affinity_ttl_seconds` override) bounds idle time, not total - session length. On Redis one Lua script does the get-or-set-or-refresh - atomically (same registration seam the rate limiters use) and the in-memory - tier is synchronized to the winner; without Redis, and whenever Redis is - unreachable, the pod-local check-and-set below stands in and is atomic because it - runs synchronously on the event loop. Degrading to a pod-local claim rather than - propagating the fault is what keeps same-pod stickiness through a Redis blip: the - caller only logs this result, so an escaping error would leave the session with no - pin at all and reshuffle every turn for the outage, which is worse than losing - cross-pod agreement. The redis tier is - resolved per call because the proxy attaches it after Router construction - (`Router._update_redis_cache`); the compiled script is cached per event loop - underneath the registration seam. - """ - redis_cache: Final = self.cache.redis_cache - if redis_cache is not None: - try: - claim_script: Final = redis_cache.async_register_script(_CLAIM_PIN_SCRIPT) - raw: Final = await claim_script(keys=(cache_key,), args=(json.dumps(pin_value), int(ttl_seconds))) - decoded: Final = raw.decode("utf-8") if isinstance(raw, bytes) else raw - if not isinstance(decoded, str): - return pin_value["model_id"] - try: - winner: object = json.loads(decoded) - except json.JSONDecodeError: - winner = decoded - self._set_local_pin(cache_key=cache_key, value=winner, ttl_seconds=ttl_seconds) - return self._pinned_model_id(winner) - except Exception as e: # noqa: BLE001 # any Redis/Lua failure degrades to the pod-local claim, never unpins - verbose_router_logger.debug( - "DeploymentAffinityCheck: redis pin claim failed, falling back to pod-local claim. error=%s", e - ) - - return self._claim_pin_in_memory(cache_key=cache_key, pin_value=pin_value, ttl_seconds=ttl_seconds) + 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 ) -> str | None: - """Pod-local half of the claim, used when no Redis tier is attached and as the - fallback when the Redis claim fails. Mirrors the Lua script exactly, including - the keepalive: re-claiming with the stored value slides the idle window through - `_set_local_pin`. Both branches stay synchronous, hence atomic on the event - loop.""" - existing: Final = self.cache.in_memory_cache.get_cache(cache_key) - if existing is not None: - existing_model_id: Final = self._pinned_model_id(existing) - if existing_model_id == pin_value["model_id"]: - self._set_local_pin(cache_key=cache_key, value=pin_value, ttl_seconds=ttl_seconds) - return existing_model_id - self._set_local_pin(cache_key=cache_key, value=pin_value, ttl_seconds=ttl_seconds) - return pin_value["model_id"] + winner: Final = claim_affinity_pin_in_memory(self.cache, cache_key, pin_value, ttl_seconds) + return self._pinned_model_id(winner) @staticmethod def _find_deployment_by_model_id(healthy_deployments: list[dict], model_id: str) -> dict | None: @@ -465,7 +394,7 @@ class DeploymentAffinityCheck(CustomLogger): enable_session_id or self._get_marker_session_affinity_ttl(request_kwargs=request_kwargs) is not None ) user_key: Final = ( - self._get_user_key_from_request_kwargs(request_kwargs=request_kwargs) + self.get_user_key_from_request_kwargs(request_kwargs=request_kwargs) if (session_affinity_active or enable_user_key) else None ) @@ -580,7 +509,7 @@ class DeploymentAffinityCheck(CustomLogger): return None user_key: Final = ( - self._get_user_key_from_request_kwargs(request_kwargs=kwargs) + self.get_user_key_from_request_kwargs(request_kwargs=kwargs) if (enable_user_key or session_affinity_active) else None ) diff --git a/tests/test_litellm/router_strategy/test_complexity_router.py b/tests/test_litellm/router_strategy/test_complexity_router.py index dba44d1e2e8..cd72388ea21 100644 --- a/tests/test_litellm/router_strategy/test_complexity_router.py +++ b/tests/test_litellm/router_strategy/test_complexity_router.py @@ -9,7 +9,7 @@ import json import logging import sys import time -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Mapping, Sequence from copy import deepcopy from functools import partial from typing import Dict, Final, List, Literal @@ -31,6 +31,7 @@ from litellm.router_utils.auto_router_model_naming import ( ) from litellm._logging import verbose_router_logger from litellm.caching.dual_cache import DualCache +from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ( OUTPUT_TOKEN_CEILING_PARAMS, RETURN_RAW_MODEL_NAME_METADATA_KEY, @@ -70,6 +71,7 @@ from litellm.router_strategy.complexity_router.tier_predictor import ( from litellm.types.router import ( Deployment, LiteLLM_Params, + PreRoutingHookResponse, RouterErrors, TaggedPreRoutingStrategy, ) @@ -5574,6 +5576,387 @@ class TestRoutingDecisionCauseLogging: assert "cause=semantic_keyword_match" not in router_log_capture.text +class TestTierModelAffinity: + @staticmethod + async def _route( + router: ComplexityRouter, + metadata: Mapping[str, object], + proposed_model: str, + prompt: str = "compact", + messages: list[dict[str, object]] | None = None, + ) -> PreRoutingHookResponse: + def choose(candidates: Sequence[str]) -> str: + return proposed_model if proposed_model in candidates else candidates[0] + + request_metadata: Final = dict(metadata) + with patch( # test-quality-ok: [TQ008] alternate proposals make affinity reuse deterministic + "litellm.router_strategy.complexity_router.complexity_router.random.choice", + side_effect=choose, + ): + result: Final = await router.async_pre_routing_hook( + model="affinity-router", + request_kwargs={"metadata": request_metadata}, + messages=messages if messages is not None else [{"role": "user", "content": prompt}], + ) + assert result is not None + if router.config.adaptive: + assert request_metadata["adaptive_router_chosen_model"] == result.model + return result + + @staticmethod + def _router( + mock_router_instance: MagicMock, + adaptive: bool = False, + deployment_affinity: bool = True, + plugins: bool = False, + ) -> ComplexityRouter: + mock_router_instance.cache = DualCache() + mock_router_instance.model_list = [] + mock_router_instance.model_name_to_deployment_indices = {} + return ComplexityRouter( + model_name="affinity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": { + tier: [ + {"model_name": model, "litellm_params": {"temperature": temperature}} + for model in ("model-a", "model-b") + ] + for tier, temperature in (("SIMPLE", 0.1), ("REASONING", 0.9)) + }, + "adaptive": adaptive, + "deployment_affinity": deployment_affinity, + "session_affinity": False, + **({"plugins": [_DummyPlugin()]} if plugins else {}), + }, + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize("adaptive", [False, True]) + async def test_reuses_model_per_tier_without_pinning_classification( + self, mock_router_instance: MagicMock, adaptive: bool + ) -> None: + router: Final = self._router(mock_router_instance, adaptive=adaptive) + metadata: Final = {"session_id": "same-session"} + first: Final = await self._route(router, metadata, "model-a") + if adaptive: + from litellm.router_strategy.adaptive_router.bandit import BanditCell + from litellm.router_strategy.adaptive_router.classifier import classify_prompt + + bandit: Final = router._ensure_adaptive_router() + assert bandit is not None + bandit._cells[(classify_prompt("compact"), "model-a")] = BanditCell(alpha=5.0, beta=5.0) + repeated: Final = await self._route(router, metadata, "model-b") + reasoning: Final = await self._route( + router, metadata, "model-b", "Let's think step by step and reason through this problem carefully." + ) + returned: Final = await self._route(router, metadata, "model-b") + + assert (first.model, repeated.model, reasoning.model, returned.model) == ( + "model-a", "model-a", "model-b", "model-a" + ) + assert tuple(result.routing_decision["tier"] for result in (first, repeated, reasoning, returned)) == ( + "SIMPLE", "SIMPLE", "REASONING", "SIMPLE" + ) + assert returned.litellm_params == {"temperature": 0.1} + assert reasoning.litellm_params == {"temperature": 0.9} + + @pytest.mark.asyncio + @pytest.mark.parametrize("identity_key", ["user_api_key_hash", "user_api_key_user_id"]) + async def test_isolates_sessions_and_authenticated_callers( + self, mock_router_instance: MagicMock, identity_key: str + ) -> None: + router: Final = self._router(mock_router_instance) + first_caller: Final = {"session_id": "shared", identity_key: "caller-a"} + other_caller: Final = {"session_id": "shared", identity_key: "caller-b"} + other_session: Final = {"session_id": "separate", identity_key: "caller-a"} + + assert (await self._route(router, first_caller, "model-a")).model == "model-a" + assert (await self._route(router, other_caller, "model-b")).model == "model-b" + assert (await self._route(router, other_session, "model-b")).model == "model-b" + assert (await self._route(router, first_caller, "model-b")).model == "model-a" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "metadata,deployment_affinity,plugins", + [ + ({}, True, False), + ({"session_id": "generated", SESSION_ID_GENERATED_METADATA_KEY: True}, True, False), + ({"session_id": "provided"}, False, False), + ({"session_id": "provided"}, True, True), + ], + ids=["absent-session", "generated-session", "disabled", "plugin-policy"], + ) + async def test_does_not_pin_without_eligible_session( + self, + mock_router_instance: MagicMock, + metadata: Mapping[str, object], + deployment_affinity: bool, + plugins: bool, + ) -> None: + router: Final = self._router( + mock_router_instance, deployment_affinity=deployment_affinity, plugins=plugins + ) + assert (await self._route(router, metadata, "model-a")).model == "model-a" + assert (await self._route(router, metadata, "model-b")).model == "model-b" + + @pytest.mark.asyncio + @pytest.mark.parametrize("adaptive", [False, True]) + async def test_replaces_pin_outside_the_context_candidate_domain(self, adaptive: bool) -> None: + router: Final = ComplexityRouter( + model_name="affinity-router", + litellm_router_instance=_windowed_router(_SMALL, _BIG), + complexity_router_config={ + "tiers": {"SIMPLE": ["small-model", "big-model"]}, + "adaptive": adaptive, + "deployment_affinity": True, + "session_affinity": False, + }, + ) + metadata: Final = {"session_id": "growing-context"} + assert (await self._route(router, metadata, "small-model")).model == "small-model" + oversized: Final = await router.async_pre_routing_hook( + model="affinity-router", + request_kwargs={"metadata": dict(metadata)}, + messages=_OVERSIZED_TURNS, + ) + assert oversized is not None + assert oversized.model == "big-model" + assert oversized.routing_decision["tier"] == "SIMPLE" + assert (await self._route(router, metadata, "small-model")).model == "big-model" + + @pytest.mark.asyncio + @pytest.mark.parametrize("session_affinity", [False, True], ids=["user-turn", "session-affinity"]) + @pytest.mark.parametrize("gate", ["image", "health"]) + async def test_temporary_replay_gate_keeps_the_held_tiers_model_preference( + self, mock_router_instance: MagicMock, session_affinity: bool, gate: Literal["image", "health"] + ) -> None: + async def get_healthy_deployments( + model: str, + request_kwargs: Mapping[str, object], + messages: Sequence[Mapping[str, object]] | None = None, + input: object = None, + parent_otel_span: object = None, + health_check_probe: bool = False, + ) -> list[dict[str, object]]: + unavailable: Final = ( + gate == "health" + and model == "model-a" + and messages is not None + and bool(messages) + and messages[-1].get("role") == "tool" + ) + return [] if unavailable else [{"model_name": model, "model_info": {"id": f"deployment-{model}"}}] + + cache: Final = DualCache() + mock_router_instance.cache = cache + mock_router_instance.async_get_healthy_deployments = get_healthy_deployments + router: Final = TestModalityRouting._router( + mock_router_instance, + { + "tiers": {"SIMPLE": ["model-a", "model-b"]}, + "deployment_affinity": True, + "session_affinity": session_affinity, + "classification_mode": "every_request" if session_affinity else "user_turn", + "modality_routing": True, + "modality_pin_override": True, + }, + {"model-a": False, "model-b": True}, + ) + metadata: Final = {"session_id": "replay-session"} + continuation: Final[list[dict[str, object]]] = [ + {"role": "user", "content": "compact"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": [IMG_PART] if gate == "image" else "done"}, + ] + assert (await self._route(router, metadata, "model-a")).model == "model-a" + + replayed: Final = await self._route(router, metadata, "model-b", messages=continuation) + assert replayed.model == "model-b" + assert replayed.routing_decision["tier"] == "SIMPLE" + assert replayed.routing_decision["cause"] == ( + "health_failover" + if gate == "health" + 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"} + + next_ask: Final = await self._route(router, metadata, "model-b") + assert next_ask.model == "model-a" + assert next_ask.routing_decision["tier"] == "SIMPLE" + assert next_ask.routing_decision["cause"] == ( + "session_affinity_pin" if session_affinity else "heuristic_scorer" + ) + + @pytest.mark.asyncio + async def test_user_turn_replay_refreshes_the_model_used_within_its_tier( + self, mock_router_instance: MagicMock + ) -> None: + clock: Final = MagicMock(return_value=100.0) + mock_router_instance.cache = DualCache(in_memory_cache=InMemoryCache(clock=clock)) + router: Final = ComplexityRouter( + model_name="affinity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {"SIMPLE": ["model-a", "model-b"]}, + "classification_mode": "user_turn", + "session_affinity_ttl_seconds": 10, + }, + ) + metadata: Final = {"session_id": "same-session"} + continuation: Final[list[dict[str, object]]] = [ + {"role": "user", "content": "compact"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "done"}, + ] + assert (await self._route(router, metadata, "model-a")).model == "model-a" + clock.return_value = 105.0 + replayed: Final = await self._route(router, metadata, "model-b", messages=continuation) + assert replayed.model == "model-a" + assert replayed.routing_decision["cause"] == "user_turn_continuation" + + clock.return_value = 111.0 + next_ask: Final = await self._route(router, metadata, "model-b") + assert next_ask.model == "model-a" + assert next_ask.routing_decision["tier"] == "SIMPLE" + assert next_ask.routing_decision["cause"] == "heuristic_scorer" + + @pytest.mark.asyncio + async def test_session_escalation_keeps_the_selected_tier_when_models_overlap( + self, mock_router_instance: MagicMock + ) -> None: + cache: Final = DualCache() + mock_router_instance.cache = cache + router: Final = ComplexityRouter( + model_name="affinity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": { + "SIMPLE": "base", + **{ + tier: [ + {"model_name": model, "litellm_params": {"temperature": temperature}} + for model in models + ] + for tier, models, temperature in ( + ("MEDIUM", ("shared", "middle"), 0.4), + ("COMPLEX", ("shared", "higher"), 0.8), + ) + }, + }, + "session_affinity": True, + "keyword_tier_rules": [{"keywords": ["visit_complex"], "tier": "COMPLEX"}], + }, + ) + 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) + + 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"} + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "stale_tier", + ["NON_REASONING", "REMOVED_TIER", 7, []], + ids=["inactive-tier", "unknown-tier", "integer-tier", "list-tier"], + ) + @pytest.mark.parametrize( + "prompt,expected_model,expected_tier", + [("compact", "model-a", "SIMPLE"), ("LITELLM ESCALATE", "model-b", "MEDIUM")], + ids=["ordinary-replay", "escalation"], + ) + async def test_reclassifies_session_pin_outside_the_active_tier_ladder( + self, + mock_router_instance: MagicMock, + stale_tier: object, + prompt: str, + expected_model: str, + expected_tier: str, + ) -> None: + cache: Final = DualCache() + mock_router_instance.cache = cache + router: Final = ComplexityRouter( + model_name="affinity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config={ + "tiers": {"SIMPLE": "model-a", "MEDIUM": "model-b"}, + "session_affinity": True, + }, + ) + 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) + + 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} + + @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.cache = DualCache() + router: Final = ComplexityRouter( + model_name="affinity-router", + litellm_router_instance=mock_router_instance, + complexity_router_config=_custom_tier_config( + tiers={"SIMPLE": ["model-a", "model-b"], "SECURITY_REVIEW": ["model-a", "model-b"], "COMPLEX": "model-a"}, + deployment_affinity=True, + classification_mode=classification_mode, + keyword_tier_rules=[ + {"keywords": ["compact"], "tier": "SIMPLE"}, + {"keywords": ["audit"], "tier": "SECURITY_REVIEW"}, + ], + ), + ) + metadata: Final = {"session_id": "custom-session"} + assert (await self._route(router, metadata, "model-a")).model == "model-a" + assert (await self._route(router, metadata, "model-b", "audit")).model == "model-b" + assert (await self._route(router, metadata, "model-b")).model == "model-a" + retained: Final = await self._route(router, metadata, "model-a", "audit") + assert retained.model == "model-b" + assert retained.routing_decision["tier"] == "SECURITY_REVIEW" + if classification_mode == "user_turn": + continuation: Final[list[dict[str, object]]] = [ + {"role": "user", "content": "audit"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "done"}, + ] + replayed: Final = await self._route(router, metadata, "model-a", messages=continuation) + assert replayed.model == "model-b" + assert replayed.routing_decision["tier"] == "SECURITY_REVIEW" + assert replayed.routing_decision["cause"] == "user_turn_continuation" + + class TestSessionAffinity: """Test the session_affinity sticky-routing behavior (off by default).""" @@ -5638,11 +6021,8 @@ class TestSessionAffinity: tier_pinned, deployment_pinned, ): - """deployment_affinity pins the deployment inside each routed group without pinning which - group the session routes to, so with session_affinity off the tier must still reclassify - on every turn while the marker the Router stamps is still emitted. Turn 1 classifies - REASONING and turn 2 SIMPLE, so a reclassified turn 2 moves model while a tier-pinned one - does not. plugins suppress both pins, since a stale pin would bypass the plugin pipeline.""" + """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.cache = DualCache() router = ComplexityRouter( model_name="test-router", @@ -5692,8 +6072,7 @@ class TestSessionAffinity: @pytest.mark.asyncio async def test_disabled_by_default_reclassifies_every_turn(self, mock_router_instance, basic_config): - """Regression: session_affinity defaults to False, so a shared session_id must NOT - pin the first turn's model; every turn is classified on its own merits.""" + """With session_affinity off, a shared session can move from REASONING to SIMPLE.""" assert "session_affinity" not in basic_config mock_router_instance.cache = DualCache() router = ComplexityRouter( @@ -5848,7 +6227,7 @@ class TestSessionAffinity: @pytest.mark.asyncio async def test_respects_ttl_seconds(self, mock_router_instance, basic_config): - cache = AsyncMock() + cache: Final = AsyncMock(in_memory_cache=DualCache().in_memory_cache, redis_cache=None) cache.async_get_cache = AsyncMock(return_value=None) mock_router_instance.cache = cache router = ComplexityRouter( @@ -5872,7 +6251,7 @@ class TestSessionAffinity: async def test_ttl_refreshed_on_cache_hit(self, mock_router_instance, basic_config): """Regression: a pinned turn must refresh the TTL, not just the first write -- otherwise a session outliving session_affinity_ttl_seconds silently loses its pin.""" - cache = AsyncMock() + cache: Final = AsyncMock(in_memory_cache=DualCache().in_memory_cache, redis_cache=None) cache.async_get_cache = AsyncMock(return_value="o1-preview") mock_router_instance.cache = cache router = ComplexityRouter( @@ -7112,7 +7491,8 @@ class TestEscalationKeywords: complexity_router_config={"tiers": {"SIMPLE": "gpt-4o-mini", "REASONING": ["o1-a", "o1-b", "o1-c"]}}, ) for pinned in ("o1-a", "o1-b", "o1-c"): - assert router._escalated_pin(pinned) == pinned + escalated: Final = router._escalated_pin(pinned) + assert (escalated.model, escalated.tier) == (pinned, "REASONING") @pytest.mark.asyncio async def test_session_escalation_at_ceiling_keeps_multi_model_pin(self, mock_router_instance): @@ -12009,7 +12389,7 @@ 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 = AsyncMock() + cache: Final = AsyncMock(in_memory_cache=DualCache().in_memory_cache, redis_cache=None) cache.async_get_cache = AsyncMock(return_value=["shared", "SIMPLE"]) mock_router_instance.cache = cache router = ComplexityRouter( @@ -12988,7 +13368,7 @@ 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 = AsyncMock() + cache: Final = AsyncMock(in_memory_cache=DualCache().in_memory_cache, redis_cache=None) cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"}) mock_router_instance.cache = cache config["session_affinity"] = True @@ -13178,7 +13558,7 @@ 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 = AsyncMock() + cache: Final = AsyncMock(in_memory_cache=DualCache().in_memory_cache, redis_cache=None) cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"}) mock_router_instance.cache = cache router = self._router( @@ -13211,7 +13591,7 @@ 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 = AsyncMock() + cache: Final = AsyncMock(in_memory_cache=DualCache().in_memory_cache, redis_cache=None) cache.async_get_cache = AsyncMock(return_value={"model": "text-cheap", "tier": "SIMPLE"}) mock_router_instance.cache = cache router = self._router( @@ -14322,18 +14702,37 @@ class TestTierHealthFailover: cooling=("id-a1",), raises_for={"exhausted-b": raised}, ) - key = router._get_session_affinity_cache_key("sess-exhausted", {}) - await router.litellm_router_instance.cache.async_set_cache( - key=key, value={"model": "dead-a", "tier": "SIMPLE"}, ttl=600 - ) - results = [ - await router.async_pre_routing_hook( - model="m", request_kwargs={"metadata": {"session_id": "sess-exhausted"}}, messages=self.SIMPLE_MESSAGE + sessions: Final = tuple(f"sess-exhausted-{sample}" for sample in range(20)) + await asyncio.gather( + *( + router.litellm_router_instance.cache.async_set_cache( + key=router._get_session_affinity_cache_key(session_id, {}), + value={"model": "dead-a", "tier": "SIMPLE"}, + ttl=600, + ) + for session_id in sessions ) - for _ in range(20) + ) + results: Final = [ + await router.async_pre_routing_hook( + model="m", request_kwargs={"metadata": {"session_id": session_id}}, messages=self.SIMPLE_MESSAGE + ) + for session_id in sessions ] assert {r.model for r in results} == expected + def choose_other(candidates: Sequence[str]) -> str: + return next((model for model in candidates if model != results[0].model), candidates[0]) + + with patch( # test-quality-ok: [TQ008] an alternate healthy proposal proves retained affinity across failover + "litellm.router_strategy.complexity_router.complexity_router.random.choice", + side_effect=choose_other, + ): + retained: Final = await router.async_pre_routing_hook( + model="m", request_kwargs={"metadata": {"session_id": sessions[0]}}, messages=self.SIMPLE_MESSAGE + ) + assert retained.model == results[0].model + @pytest.mark.asyncio async def test_a_group_the_router_has_no_deployment_for_is_not_a_failover_target(self, mock_router_instance): """The owner answers an unconfigured group with BadRequestError. Reading that as live diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py index cf48888600e..780300bf9e1 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_session_id_affinity.py @@ -1,3 +1,5 @@ +import asyncio +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -6,7 +8,9 @@ import pytest import json 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.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, @@ -558,6 +562,124 @@ async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down(): assert second == "our-deployment" +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("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"}), + ("first", {"model": "second"}), + (None, {"model": "second"}), + ], +) +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) + clock.return_value = 105.0 + + winner: Final = await claim_affinity_pin( + cache, "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 + clock.return_value = 136.0 + assert cache.in_memory_cache.get_cache("tier-pin") is None + + +@pytest.mark.asyncio +async def test_concurrent_eligible_claims_return_one_winner() -> None: + cache: Final = DualCache() + candidates: Final = ({"model": "first"}, {"model": "second"}) + winners: Final = await asyncio.gather(*( + claim_affinity_pin( + cache, "tier-pin", candidates[index % 2], 30, + eligible_values=candidates, + ) + for index in range(20) + )) + + assert winners == [{"model": "first"}] * 20 + assert cache.in_memory_cache.get_cache("tier-pin") == {"model": "first"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("stored", "expected", "refresh"), + [ + ({"model_id": 7}, "7", True), + ({"model_id": "other"}, "other", False), + ({"model": "7"}, None, False), + (["7"], None, False), + ], +) +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)) + callback: Final = DeploymentAffinityCheck( + cache=cache, ttl_seconds=30, + enable_user_key_affinity=False, enable_responses_api_affinity=False, + ) + cache.in_memory_cache.set_cache("deployment-pin", stored, ttl=10) + clock.return_value = 105.0 + + 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 + + @pytest.mark.asyncio 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 diff --git a/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx b/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx index 9022d424369..325362ea177 100644 --- a/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AffinityControls.tsx @@ -28,13 +28,13 @@ export const AffinityControls: React.FC<{ onChange({ ...value, deployment_affinity: deploymentAffinity })} - aria-label="Pin a session to one deployment per model group" + aria-label="Pin one model deployment per tier" /> - Pin a session to one deployment per model group + Pin one model deployment per tier - Keeps a session on the same deployment within a group, so provider prompt caches stay warm. Turn off to - load-balance every turn. + Reuses the model chosen for each tier and its deployment when available. Requests can still move between tiers. + Turn off to select models and load-balance deployments every turn.