feat(router): independent, default-on deployment affinity for the auto-router (#36146)

This commit is contained in:
tin-berri 2026-08-08 13:02:29 -07:00 committed by GitHub
parent 554f065361
commit e35ee4e5fa
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
15 changed files with 964 additions and 111 deletions

View file

@ -1322,6 +1322,7 @@ X_LITELLM_DISABLE_CALLBACKS: Final = "x-litellm-disable-callbacks"
LITELLM_METADATA_FIELD: Final = "litellm_metadata"
OLD_LITELLM_METADATA_FIELD: Final = "metadata"
RETURN_RAW_MODEL_NAME_METADATA_KEY: Final = "_complexity_router_return_raw_model_name"
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: Final = "_session_deployment_affinity_ttl"
INTERNAL_CALL_ORIGIN_METADATA_KEY: Final = "internal_call_origin"
LITELLM_TRUNCATED_PAYLOAD_FIELD: Final = "litellm_truncated"
LITELLM_TRUNCATION_DB_SAFEGUARD_NOTE: Final = (

View file

@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Any, Final, Optional
import litellm
from litellm import get_secret
from litellm._logging import verbose_proxy_logger
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY
from litellm.constants import PRE_CALL_EXECUTED_GUARDRAILS_KEY, SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.core_helpers import (
get_metadata_variable_name_from_kwargs,
@ -425,6 +425,7 @@ LITELLM_PROXY_INTERNAL_METADATA_KEYS: Final = frozenset(
"_guardrail_pipelines",
"_pipeline_managed_guardrails",
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
"disable_global_guardrails",
"disable_global_guardrail",
"opted_out_global_guardrails",

View file

@ -18,6 +18,7 @@ from litellm.constants import (
INTERNAL_CALL_ORIGIN_METADATA_KEY,
LITELLM_PROXY_MASTER_KEY_ALIAS,
PRE_CALL_EXECUTED_GUARDRAILS_KEY,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
)
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
from litellm.litellm_core_utils.initialize_dynamic_callback_params import (
@ -226,6 +227,7 @@ _UNTRUSTED_METADATA_CONTROL_FIELDS: Final = (
"applied_policies",
"policy_sources",
"routing_decision",
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
INTERNAL_CALL_ORIGIN_METADATA_KEY,
"standard_logging_object",
"proxy_server_request",

View file

@ -46,6 +46,7 @@ from litellm.constants import (
DEFAULT_HEALTH_CHECK_INTERVAL,
DEFAULT_HEALTH_CHECK_STALENESS_MULTIPLIER,
DEFAULT_MAX_LRU_CACHE_SIZE,
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
)
from litellm.integrations.custom_logger import CustomLogger
from litellm.litellm_core_utils.asyncify import run_async_function
@ -135,6 +136,7 @@ from litellm.router_utils.handle_error import (
from litellm.router_utils.health_state_cache import DeploymentHealthCache
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
DeploymentAffinityCheck,
warn_on_unknown_model_group_affinity_flags,
)
from litellm.router_utils.pre_call_checks.io_token_rate_limit_check import (
build_io_token_rate_limit_headers,
@ -603,6 +605,10 @@ class Router:
# ``litellm.proxy.auth.auth_checks._is_model_cost_zero``.
self._zero_cost_cache: dict[str, bool] = {}
self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds
self.model_group_affinity_config = model_group_affinity_config
warn_on_unknown_model_group_affinity_flags(model_group_affinity_config)
if model_list is not None:
# set_model_list will build indices automatically
self.set_model_list(model_list)
@ -744,7 +750,6 @@ class Router:
litellm.failure_callback = [self.deployment_callback_on_failure]
self.routing_strategy_args = routing_strategy_args
self.provider_budget_config = provider_budget_config
self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds
self.router_budget_logger: RouterBudgetLimiting | None = None
if RouterBudgetLimiting.should_init_router_budget_limiter(
model_list=model_list, provider_budget_config=self.provider_budget_config
@ -766,7 +771,6 @@ class Router:
)
self.model_group_retry_policy: dict[str, RetryPolicy] | None = model_group_retry_policy
self.model_group_affinity_config: dict[str, list[str]] | None = model_group_affinity_config
self.allowed_fails_policy: AllowedFailsPolicy | None = None
if allowed_fails_policy is not None:
@ -789,21 +793,8 @@ class Router:
# If model_group_affinity_config is set but no global affinity checks were
# enabled, we still need the DeploymentAffinityCheck callback (with global
# flags all False) so per-group config can activate affinity per model group.
if self.model_group_affinity_config and not any(
isinstance(cb, DeploymentAffinityCheck) for cb in (self.optional_callbacks or [])
):
if self.optional_callbacks is None:
self.optional_callbacks = []
affinity_callback: Final = DeploymentAffinityCheck(
cache=self.cache,
ttl_seconds=self.deployment_affinity_ttl_seconds,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
enable_session_id_affinity=False,
model_group_affinity_config=self.model_group_affinity_config,
)
self.optional_callbacks.append(affinity_callback)
litellm.logging_callback_manager.add_litellm_callback(affinity_callback)
if self.model_group_affinity_config:
self._ensure_deployment_affinity_callback()
if self.alerting_config is not None:
self._initialize_alerting()
@ -1662,6 +1653,28 @@ class Router:
_move_before_deployment_affinity(self.optional_callbacks, ec_callback)
_move_before_deployment_affinity(litellm.callbacks, ec_callback)
def _ensure_deployment_affinity_callback(self) -> None:
"""Register the DeploymentAffinityCheck callback (global flags all False) if absent.
Needed when nothing enabled a global affinity flag but affinity can still
activate per request: per-group `model_group_affinity_config` entries, or the
session-affinity marker a complexity router stamps at pre-routing time.
"""
if any(isinstance(cb, DeploymentAffinityCheck) for cb in (self.optional_callbacks or [])):
return
if self.optional_callbacks is None:
self.optional_callbacks = []
affinity_callback: Final = DeploymentAffinityCheck(
cache=self.cache,
ttl_seconds=self.deployment_affinity_ttl_seconds,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
enable_session_id_affinity=False,
model_group_affinity_config=self.model_group_affinity_config,
)
self.optional_callbacks.append(affinity_callback)
litellm.logging_callback_manager.add_litellm_callback(affinity_callback)
def add_optional_pre_call_checks(self, optional_pre_call_checks: OptionalPreCallChecks | None):
if optional_pre_call_checks is None:
return
@ -7683,6 +7696,8 @@ class Router:
strategy=complexity_router,
strategy_label="Complexity-router",
)
if complexity_router._uses_deployment_pin:
self._ensure_deployment_affinity_callback()
def _is_adaptive_router_deployment(self, litellm_params: LiteLLM_Params) -> bool:
"""True when this deployment opts in via the `auto_router/adaptive_router` model prefix."""
@ -11190,6 +11205,9 @@ class Router:
router_strategy: Final = self._select_pre_routing_strategy(model=model, request_kwargs=request_kwargs)
if router_strategy is None:
self._record_routing_decision(request_kwargs=request_kwargs, routing_decision=None)
self._stamp_or_clear_metadata_key(
request_kwargs=request_kwargs, key=SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, value=None
)
return None
pre_routing_hook_response: Final = await router_strategy.async_pre_routing_hook(
@ -11203,6 +11221,11 @@ class Router:
request_kwargs=request_kwargs,
routing_decision=(pre_routing_hook_response.routing_decision if pre_routing_hook_response else None),
)
self._stamp_or_clear_metadata_key(
request_kwargs=request_kwargs,
key=SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY,
value=(pre_routing_hook_response.session_affinity_ttl_seconds if pre_routing_hook_response else None),
)
# `model` (the alias, e.g. "smart-router") is never the deployment actually
# called - apply the alias's own litellm_params (besides `model` itself,
@ -11234,21 +11257,40 @@ class Router:
to the deployment that actually served the request. Every attempt therefore
writes or clears, never just writes.
"""
if routing_decision is None:
Router._stamp_or_clear_metadata_key(
request_kwargs=request_kwargs,
key="routing_decision",
value=(
None
if routing_decision is None
else Router._redact_prompt_text_if_needed(
request_kwargs=request_kwargs, routing_decision=routing_decision
)
),
)
@staticmethod
def _stamp_or_clear_metadata_key(request_kwargs: dict, key: str, value: object | None) -> None:
"""Write a proxy-internal metadata key for THIS routing attempt, or clear it.
Fallbacks and retries re-enter the pre-routing hook with the same
`request_kwargs`, so every attempt must write or clear, never just write;
a value left behind by an earlier attempt would be attributed to this one.
`get_or_create_metadata_bucket` is the single owner of "which dict holds
proxy-internal metadata": it picks `litellm_metadata` when present (so the
value never lands in the `metadata` dict that routes like /v1/messages
forward to the provider) and replaces a non-dict value rather than silently
skipping the write. Clearing pops from BOTH buckets so a request whose
bucket resolution changed between attempts cannot resurrect a stale value.
"""
if value is None:
for bucket in (request_kwargs.get("metadata"), request_kwargs.get("litellm_metadata")):
if isinstance(bucket, dict):
bucket.pop("routing_decision", None)
bucket.pop(key, None)
return
# `get_or_create_metadata_bucket` is the single owner of "which dict holds
# proxy-internal metadata": it picks `litellm_metadata` when present (so the
# decision never lands in the `metadata` dict that routes like /v1/messages
# forward to the provider) and replaces a non-dict value rather than silently
# skipping the write.
_, metadata_bucket = get_or_create_metadata_bucket(request_kwargs)
metadata_bucket["routing_decision"] = Router._redact_prompt_text_if_needed(
request_kwargs=request_kwargs, routing_decision=routing_decision
)
metadata_bucket[key] = value
@staticmethod
def _redact_prompt_text_if_needed(

View file

@ -1651,6 +1651,28 @@ class ComplexityRouter(CustomLogger):
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}"
@property
def _uses_tier_pin(self) -> bool:
return bool(self.config.session_affinity and not self.config.plugins)
@property
def _uses_deployment_pin(self) -> bool:
"""session_affinity implies the deployment pin: a session frozen onto one model
group but load-balanced across its deployments would still go cache-cold, which
is the exact failure both flags exist to prevent."""
return bool((self.config.deployment_affinity or self.config.session_affinity) and not self.config.plugins)
def _with_session_deployment_affinity(
self, response: PreRoutingHookResponse | None
) -> PreRoutingHookResponse | None:
if response is None or not self._uses_deployment_pin:
return response
return response.model_copy(
update={ # mutable-ok: model_copy types update as a plain dict
"session_affinity_ttl_seconds": self.config.session_affinity_ttl_seconds
}
)
async def async_pre_routing_hook(
self,
model: str,
@ -1685,7 +1707,7 @@ class ComplexityRouter(CustomLogger):
resolved_messages: Final = self._resolve_messages(messages, request_kwargs)
conversation_continuing: Final = _conversation_is_continuing(resolved_messages)
use_session_affinity: Final = self.config.session_affinity and not self.config.plugins
use_session_affinity: Final = self._uses_tier_pin
session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs) if use_session_affinity else None
cache_key = self._get_session_affinity_cache_key(session_id, request_kwargs) if session_id is not None else None
@ -1724,17 +1746,19 @@ class ComplexityRouter(CustomLogger):
"ComplexityRouter: routing decision cause=%s, routed_model=%s", cause, routed_model
)
has_original_messages: Final = messages is not None and len(messages) > 0
return PreRoutingHookResponse(
model=routed_model,
messages=messages if has_original_messages else None,
routing_decision=self._build_routing_decision(
routed_model=routed_model,
cause=cause,
tier=self._tier_for_model(routed_model),
escalation_keyword=pin_escalation_keyword,
escalated=escalated,
conversation_continuing=conversation_continuing,
),
return self._with_session_deployment_affinity(
PreRoutingHookResponse(
model=routed_model,
messages=messages if has_original_messages else None,
routing_decision=self._build_routing_decision(
routed_model=routed_model,
cause=cause,
tier=self._tier_for_model(routed_model),
escalation_keyword=pin_escalation_keyword,
escalated=escalated,
conversation_continuing=conversation_continuing,
),
)
)
response: Final = await self._classify_and_route(
@ -1752,7 +1776,7 @@ class ComplexityRouter(CustomLogger):
value=response.model,
ttl=self.config.session_affinity_ttl_seconds,
)
return response
return self._with_session_deployment_affinity(response)
async def _classify_and_route(
self,

View file

@ -508,13 +508,39 @@ class ComplexityRouterConfig(BaseModel):
"session's first turn and reuse it for every later turn, skipping re-classification. "
"Off by default so every turn is classified on its own merits and routed to the cheapest "
"adequate tier. Set True to keep a multi-turn session on one model, which preserves "
"provider prompt caches and avoids cross-model conversation-history errors."
"provider prompt caches and avoids cross-model conversation-history errors. Always "
"implies the deployment pin regardless of deployment_affinity: the session sticks to "
"one deployment of the pinned model, since freezing the model while re-shuffling its "
"deployments would still go cache-cold."
),
)
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."
),
)
session_affinity_ttl_seconds: int = Field(
default=3600,
gt=0,
description="TTL for the session affinity pin; refreshed on every cache hit",
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 "
"idle time for the session's routing decisions rather than total session length"
),
)
plugins: list[RoutingPlugin] | None = Field(

View file

@ -13,12 +13,15 @@ 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 litellm._logging import verbose_router_logger
from litellm.caching.dual_cache import DualCache
from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY
from litellm.integrations.custom_logger import CustomLogger, Span
from litellm.responses.utils import ResponsesAPIRequestUtils
from litellm.types.llms.openai import AllMessageValues
@ -29,6 +32,47 @@ class DeploymentAffinityCacheValue(TypedDict):
model_id: str
VALID_MODEL_GROUP_AFFINITY_FLAGS: Final = frozenset(
{
"deployment_affinity",
"responses_api_deployment_check",
"session_affinity",
"encrypted_content_affinity",
}
)
def warn_on_unknown_model_group_affinity_flags(model_group_affinity_config: Mapping[str, Sequence[str]] | None) -> None:
"""`model_group_affinity_config` is one Router-level config consumed by two callbacks:
DeploymentAffinityCheck acts on three of the flags and EncryptedContentAffinityCheck
on the fourth, so typo detection lives here at the schema, not inside either consumer.
"""
if model_group_affinity_config is None:
return
for group, flags in model_group_affinity_config.items():
unknown = set(flags) - VALID_MODEL_GROUP_AFFINITY_FLAGS
if unknown:
verbose_router_logger.warning(
"model_group_affinity_config: unknown flag(s) %s for model group '%s'; will be ignored. Valid flags: %s",
unknown,
group,
VALID_MODEL_GROUP_AFFINITY_FLAGS,
)
_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.
@ -38,14 +82,6 @@ class DeploymentAffinityCheck(CustomLogger):
"""
CACHE_KEY_PREFIX = "deployment_affinity:v1"
VALID_FLAGS = frozenset(
{
"deployment_affinity",
"responses_api_deployment_check",
"session_affinity",
"encrypted_content_affinity",
}
)
def __init__(
self,
@ -63,15 +99,6 @@ class DeploymentAffinityCheck(CustomLogger):
self.enable_responses_api_affinity = enable_responses_api_affinity
self.enable_session_id_affinity = enable_session_id_affinity
self.model_group_affinity_config: dict[str, list[str]] = model_group_affinity_config or {}
for group, flags in self.model_group_affinity_config.items():
unknown = set(flags) - self.VALID_FLAGS
if unknown:
verbose_router_logger.warning(
"DeploymentAffinityCheck: unknown flag(s) %s for model group '%s'; will be ignored. Valid flags: %s",
unknown,
group,
self.VALID_FLAGS,
)
def _get_effective_flags(self, model_group: str) -> tuple[bool, bool, bool]:
"""
@ -218,8 +245,13 @@ class DeploymentAffinityCheck(CustomLogger):
return f"{cls.CACHE_KEY_PREFIX}:{model_group}:{hashed_user_key}"
@classmethod
def get_session_affinity_cache_key(cls, model_group: str, session_id: str) -> str:
return f"{cls.CACHE_KEY_PREFIX}:session:{model_group}:{session_id}"
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}"
@staticmethod
def _get_user_key_from_metadata_dict(metadata: dict) -> str | None:
@ -278,6 +310,97 @@ class DeploymentAffinityCheck(CustomLogger):
return session_id
return None
@staticmethod
def _get_marker_session_affinity_ttl(request_kwargs: dict) -> int | None:
"""TTL from the session-affinity marker the Router stamps at pre-routing time
when an auto-router routed this request with session_affinity enabled.
Marker presence enables session pinning for this request only; anything that
is not a positive int is treated as absent."""
for metadata in DeploymentAffinityCheck._iter_metadata_dicts(request_kwargs):
ttl = metadata.get(SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY)
if isinstance(ttl, int) and not isinstance(ttl, bool) and ttl > 0:
return ttl
return None
@staticmethod
def _pinned_model_id(stored: object) -> str | None:
"""Deployment id held by a stored pin, for both the dict shape this writes and the
bare string older writers left behind. None when the value is neither."""
if isinstance(stored, dict):
model_id: Final = stored.get("model_id")
return str(model_id) if model_id is not None else None
if isinstance(stored, str):
return stored
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)
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 `session_affinity_ttl_seconds` 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)
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"]
@staticmethod
def _find_deployment_by_model_id(healthy_deployments: list[dict], model_id: str) -> dict | None:
for deployment in healthy_deployments:
@ -334,12 +457,21 @@ class DeploymentAffinityCheck(CustomLogger):
if stable_model_map_key is None:
return typed_healthy_deployments
session_affinity_active: Final = (
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)
if (session_affinity_active or enable_user_key)
else None
)
# 2) Session-id -> deployment affinity
if enable_session_id:
if session_affinity_active:
session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs=request_kwargs)
if session_id is not None:
session_cache_key: Final = self.get_session_affinity_cache_key(
model_group=stable_model_map_key, session_id=session_id
model_group=stable_model_map_key, session_id=session_id, user_key=user_key
)
session_cache_result: Final = await self.cache.async_get_cache(key=session_cache_key)
@ -371,7 +503,6 @@ class DeploymentAffinityCheck(CustomLogger):
if not enable_user_key:
return typed_healthy_deployments
user_key: Final = self._get_user_key_from_request_kwargs(request_kwargs=request_kwargs)
if user_key is None:
return typed_healthy_deployments
@ -438,18 +569,22 @@ class DeploymentAffinityCheck(CustomLogger):
enable_session_id,
) = self._get_effective_flags(deployment_model_name)
if not enable_user_key and not enable_session_id:
marker_session_ttl: Final = self._get_marker_session_affinity_ttl(request_kwargs=kwargs)
session_affinity_active: Final = enable_session_id or marker_session_ttl is not None
if not enable_user_key and not session_affinity_active:
return None
user_key = None
if enable_user_key:
user_key = self._get_user_key_from_request_kwargs(request_kwargs=kwargs)
user_key: Final = (
self._get_user_key_from_request_kwargs(request_kwargs=kwargs)
if (enable_user_key or session_affinity_active)
else None
)
session_id: Final = (
self._get_session_id_from_request_kwargs(request_kwargs=kwargs) if session_affinity_active else None
)
session_id = None
if enable_session_id:
session_id = self._get_session_id_from_request_kwargs(request_kwargs=kwargs)
if user_key is None and session_id is None:
if not ((enable_user_key and user_key is not None) or session_id is not None):
return None
model_info = kwargs.get("model_info")
@ -473,22 +608,31 @@ class DeploymentAffinityCheck(CustomLogger):
verbose_router_logger.warning("DeploymentAffinityCheck: model_id missing; skipping affinity cache update.")
return None
if user_key is not None:
pin_value: Final = DeploymentAffinityCacheValue(model_id=str(model_id))
if enable_user_key and user_key is not None:
try:
cache_key: Final = self.get_affinity_cache_key(model_group=deployment_model_name, user_key=user_key)
await self.cache.async_set_cache(
cache_key,
DeploymentAffinityCacheValue(model_id=str(model_id)),
ttl=self.ttl_seconds,
)
verbose_router_logger.debug(
"DeploymentAffinityCheck: set affinity mapping model_map_key=%s deployment=%s ttl=%s user_key=%s",
deployment_model_name,
model_id,
self.ttl_seconds,
self._shorten_for_logs(user_key),
claimed_user_pin: Final = await self._claim_pin(
cache_key=cache_key,
pin_value=pin_value,
ttl_seconds=self.ttl_seconds,
)
if claimed_user_pin == pin_value["model_id"]:
verbose_router_logger.debug(
"DeploymentAffinityCheck: set affinity mapping model_map_key=%s deployment=%s ttl=%s user_key=%s",
deployment_model_name,
model_id,
self.ttl_seconds,
self._shorten_for_logs(user_key),
)
else:
verbose_router_logger.debug(
"DeploymentAffinityCheck: affinity pin already claimed model_map_key=%s existing=%s ours=%s",
deployment_model_name,
claimed_user_pin,
model_id,
)
except Exception as e:
# Non-blocking: affinity is a best-effort optimization.
verbose_router_logger.debug(
@ -500,21 +644,31 @@ class DeploymentAffinityCheck(CustomLogger):
# Also persist Session-ID affinity if enabled and session-id is provided
if session_id is not None:
try:
session_affinity_ttl: Final = marker_session_ttl if marker_session_ttl is not None else self.ttl_seconds
session_cache_key: Final = self.get_session_affinity_cache_key(
model_group=deployment_model_name, session_id=session_id
model_group=deployment_model_name, session_id=session_id, user_key=user_key
)
await self.cache.async_set_cache(
session_cache_key,
DeploymentAffinityCacheValue(model_id=str(model_id)),
ttl=self.ttl_seconds,
)
verbose_router_logger.debug(
"DeploymentAffinityCheck: set session affinity mapping model_map_key=%s deployment=%s ttl=%s session_id=%s",
deployment_model_name,
model_id,
self.ttl_seconds,
session_id,
claimed_session_pin: Final = await self._claim_pin(
cache_key=session_cache_key,
pin_value=pin_value,
ttl_seconds=session_affinity_ttl,
)
if claimed_session_pin == pin_value["model_id"]:
verbose_router_logger.debug(
"DeploymentAffinityCheck: set session affinity mapping model_map_key=%s deployment=%s ttl=%s session_id=%s",
deployment_model_name,
model_id,
session_affinity_ttl,
session_id,
)
else:
verbose_router_logger.debug(
"DeploymentAffinityCheck: session pin already claimed model_map_key=%s existing=%s ours=%s session_id=%s",
deployment_model_name,
claimed_session_pin,
model_id,
session_id,
)
except Exception as e:
verbose_router_logger.debug(
"DeploymentAffinityCheck: failed to set session affinity cache. model_map_key=%s error=%s",

View file

@ -816,6 +816,7 @@ class PreRoutingHookResponse(BaseModel):
model: str
messages: list[dict[str, Any]] | None
routing_decision: StandardLoggingRoutingDecision | None = None
session_affinity_ttl_seconds: int | None = None
_PreRoutingStrategyT_co = TypeVar("_PreRoutingStrategyT_co", covariant=True)

View file

@ -676,6 +676,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields():
"applied_policies": ["spoofed-policy"],
"policy_sources": {"spoofed-policy": "request"},
"routing_decision": {"cause": "forged", "routed_model": "spoofed"},
"_session_deployment_affinity_ttl": 999999,
"internal_call_origin": "autorouter_classifier",
"_guardrail_pipelines": [{"name": "spoofed"}],
"_pipeline_managed_guardrails": ["evaded"],
@ -719,6 +720,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields():
"applied_policies",
"policy_sources",
"routing_decision",
"_session_deployment_affinity_ttl",
"internal_call_origin",
"_guardrail_pipelines",
"_pipeline_managed_guardrails",

View file

@ -3414,6 +3414,103 @@ class TestSessionAffinity:
def _request_kwargs(session_id: str) -> Dict:
return {"metadata": {"session_id": session_id}}
@pytest.mark.asyncio
async def test_hook_response_carries_session_affinity_ttl_on_classify_and_pin_paths(
self, mock_router_instance, session_affinity_config
):
"""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.cache = DualCache()
router = ComplexityRouter(
model_name="test-router",
litellm_router_instance=mock_router_instance,
complexity_router_config={**session_affinity_config, "session_affinity_ttl_seconds": 321},
)
request_kwargs = self._request_kwargs("marker-session")
first = await router.async_pre_routing_hook(
model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE
)
second = await router.async_pre_routing_hook(
model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE
)
assert first.session_affinity_ttl_seconds == 321
assert second.session_affinity_ttl_seconds == 321
@pytest.mark.parametrize(
"session_affinity,deployment_affinity,plugins,tier_pinned,deployment_pinned",
[
(False, False, False, False, False),
(False, True, False, False, True),
(True, False, False, True, True),
(True, True, False, True, True),
(False, True, True, False, False),
(True, True, True, False, False),
],
)
@pytest.mark.asyncio
async def test_tier_pin_and_deployment_pin_are_independently_gated(
self,
mock_router_instance,
basic_config,
session_affinity,
deployment_affinity,
plugins,
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."""
mock_router_instance.cache = DualCache()
router = ComplexityRouter(
model_name="test-router",
litellm_router_instance=mock_router_instance,
complexity_router_config={
**basic_config,
"session_affinity": session_affinity,
"deployment_affinity": deployment_affinity,
**({"plugins": [_DummyPlugin()]} if plugins else {}),
},
)
request_kwargs = self._request_kwargs("matrix-session")
first = await router.async_pre_routing_hook(
model="test-model", request_kwargs=request_kwargs, messages=self.REASONING_MESSAGE
)
second = await router.async_pre_routing_hook(
model="test-model", request_kwargs=request_kwargs, messages=self.SIMPLE_MESSAGE
)
assert first.model == "o1-preview"
assert second.model == ("o1-preview" if tier_pinned else "gpt-4o-mini")
assert (first.session_affinity_ttl_seconds is not None) is deployment_pinned
assert (second.session_affinity_ttl_seconds is not None) is deployment_pinned
@pytest.mark.asyncio
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.cache = DualCache()
disabled_router = ComplexityRouter(
model_name="test-router",
litellm_router_instance=mock_router_instance,
complexity_router_config={**basic_config, "deployment_affinity": False},
)
plugin_router = ComplexityRouter(
model_name="test-router-plugins",
litellm_router_instance=mock_router_instance,
complexity_router_config={**session_affinity_config, "plugins": [_DummyPlugin()]},
)
disabled = await disabled_router.async_pre_routing_hook(
model="test-model", request_kwargs=self._request_kwargs("s-off"), messages=self.SIMPLE_MESSAGE
)
with_plugins = await plugin_router.async_pre_routing_hook(
model="test-model", request_kwargs=self._request_kwargs("s-plugins"), messages=self.SIMPLE_MESSAGE
)
assert disabled.session_affinity_ttl_seconds is None
assert with_plugins.session_affinity_ttl_seconds is None
@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

View file

@ -465,8 +465,7 @@ 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 = AsyncMock()
cache.async_set_cache = AsyncMock()
cache = DualCache()
callback = DeploymentAffinityCheck(
cache=cache,
@ -489,11 +488,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",
)
cache.async_set_cache.assert_called_once_with(
expected_cache_key,
{"model_id": "model-id-123"},
ttl=123,
)
assert await cache.async_get_cache(key=expected_cache_key) == {"model_id": "model-id-123"}
@pytest.mark.asyncio

View file

@ -1,6 +1,6 @@
import os
import sys
from unittest.mock import AsyncMock, patch
from unittest.mock import AsyncMock, MagicMock, patch
import pytest
@ -10,6 +10,7 @@ import json
import litellm
from litellm.caching.dual_cache import DualCache
from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
DeploymentAffinityCheck,
)
@ -163,7 +164,7 @@ async def test_async_session_id_affinity_priority_over_user_key():
await callback.cache.async_set_cache(
DeploymentAffinityCheck.get_session_affinity_cache_key(
"model_group", "session1"
"model_group", "session1", user_key="user1"
),
{"model_id": "deployment-2"},
)
@ -180,3 +181,439 @@ async def test_async_session_id_affinity_priority_over_user_key():
assert len(filtered) == 1
assert filtered[0]["model_info"]["id"] == "deployment-2"
MOCK_RESPONSES_API_RESPONSE = {
"id": "resp_mock-resp-456",
"object": "response",
"created_at": 1741476542,
"status": "completed",
"model": "azure/computer-use-preview",
"output": [],
"usage": {
"input_tokens": 5,
"output_tokens": 10,
"total_tokens": 15,
"output_tokens_details": {"reasoning_tokens": 0},
},
}
def _smart_router(session_affinity=True, ttl_seconds=777, deployment_affinity=True):
return litellm.Router(
model_list=[
{
"model_name": "smart-router",
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_default_model": "target-group",
"complexity_router_config": {
"session_affinity": session_affinity,
"deployment_affinity": deployment_affinity,
"session_affinity_ttl_seconds": ttl_seconds,
"tiers": {
"SIMPLE": "target-group",
"MEDIUM": "target-group",
"COMPLEX": "target-group",
"REASONING": "target-group",
},
},
},
},
{
"model_name": "target-group",
"litellm_params": {
"model": "azure/computer-use-preview-1",
"api_key": "mock-api-key-1",
"api_version": "mock-api-version",
"api_base": "https://mock-endpoint-1.openai.azure.com",
},
"model_info": {"id": "deployment-1", "base_model": "computer-use-preview"},
},
{
"model_name": "target-group",
"litellm_params": {
"model": "azure/computer-use-preview-2",
"api_key": "mock-api-key-2",
"api_version": "mock-api-version-2",
"api_base": "https://mock-endpoint-2.openai.azure.com",
},
"model_info": {"id": "deployment-2", "base_model": "computer-use-preview"},
},
],
)
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
)
def _cleanup_router_callbacks(router):
for callback in router.optional_callbacks or []:
litellm.logging_callback_manager.remove_callback_from_all_lists(callback)
async def _one_turn(router, model, session_id, key_hash):
"""One request with the shuffle forced to deployment-1, so any other landing
deployment can only come from a pin read."""
with (
patch(
"litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post",
new_callable=AsyncMock,
) as mock_post,
patch(
"litellm.router_strategy.simple_shuffle.random.choice",
side_effect=lambda seq: seq[0],
),
):
mock_post.return_value = MockResponse(MOCK_RESPONSES_API_RESPONSE, 200)
response = await router.aresponses(
model=model,
input=f"turn for {session_id} {key_hash}",
litellm_metadata={"session_id": session_id, "user_api_key_hash": key_hash},
)
return response._hidden_params["model_id"]
@pytest.mark.asyncio
async def test_auto_router_session_affinity_writes_scoped_pin_and_follows_it():
"""Turn 1 persists a key-scoped deployment pin; a pin seeded to the deployment
the shuffle would never pick is then followed, proving the read path."""
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")) == {
"model_id": served
}
assert await router.cache.async_get_cache(key=_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"}
)
assert await _one_turn(router, "smart-router", "read-session", "key-1") == "deployment-2"
finally:
_cleanup_router_callbacks(router)
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model,key_hash",
[
("target-group", "key-1"),
("smart-router", "key-2"),
],
ids=["direct-group-call", "different-api-key"],
)
async def test_seeded_session_pin_is_invisible_outside_its_scope(model, key_hash):
"""The pin binds (auto-routed request, api key, session): a direct call to the
group and a different key reusing the session id must both ignore it."""
router = _smart_router()
try:
await router.cache.async_set_cache(
key=_session_pin_key("scoped-session", "key-1"), value={"model_id": "deployment-2"}
)
assert await _one_turn(router, model, "scoped-session", key_hash) == "deployment-1"
finally:
_cleanup_router_callbacks(router)
@pytest.mark.asyncio
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()
callback = DeploymentAffinityCheck(
cache=cache,
ttl_seconds=3600,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
)
await callback.async_pre_call_deployment_hook(
kwargs={
"model_info": {"id": "deployment-1"},
"metadata": {
"deployment_model_name": "target-group",
"session_id": "ttl-session",
"user_api_key_hash": "key-1",
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: 777,
},
},
call_type=None,
)
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)
@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()
callback = DeploymentAffinityCheck(
cache=cache,
ttl_seconds=3600,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
)
await callback.async_pre_call_deployment_hook(
kwargs={
"model_info": {"id": "deployment-1"},
"metadata": {
"deployment_model_name": "target-group",
"session_id": "bad-marker-session",
"user_api_key_hash": "key-1",
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: bad_marker,
},
},
call_type=None,
)
assert cache.in_memory_cache.cache_dict == {}
@pytest.mark.asyncio
@pytest.mark.parametrize("enable_user_key", [False, True], ids=["session-pin", "user-key-pin"])
async def test_concurrent_first_requests_never_flip_a_claimed_pin(enable_user_key):
"""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()
callback = DeploymentAffinityCheck(
cache=cache,
ttl_seconds=3600,
enable_user_key_affinity=enable_user_key,
enable_responses_api_affinity=False,
)
def racing_kwargs(deployment_id):
metadata = {"deployment_model_name": "target-group", "user_api_key_hash": "key-1"}
if not enable_user_key:
metadata["session_id"] = "racing-session"
metadata[SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY] = 777
return {"model_info": {"id": deployment_id}, "metadata": metadata}
await callback.async_pre_call_deployment_hook(kwargs=racing_kwargs("deployment-1"), call_type=None)
await callback.async_pre_call_deployment_hook(kwargs=racing_kwargs("deployment-2"), call_type=None)
pinned_key = (
DeploymentAffinityCheck.get_affinity_cache_key(model_group="target-group", user_key="key-1")
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"}
@pytest.mark.asyncio
@pytest.mark.parametrize(
"stored_pin",
[{"model_id": "deployment-1"}, "deployment-1"],
ids=["dict-pin", "legacy-string-pin"],
)
async def test_in_memory_reclaim_slides_idle_window_only_for_the_stored_deployment(stored_pin):
"""The pod-local claim mirrors the Lua keepalive: the winning deployment's
re-claim extends the pin's expiry, a losing deployment's claim touches neither
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()
callback = DeploymentAffinityCheck(
cache=cache,
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]
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"
@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
path and the write path, so the marker-gated pin round-trips through one key."""
callback = DeploymentAffinityCheck(
cache=DualCache(),
ttl_seconds=3600,
enable_user_key_affinity=False,
enable_responses_api_affinity=False,
)
request_kwargs = {
"model_info": {"id": "wild-deployment-2"},
"metadata": {
"deployment_model_name": "openai/*",
"session_id": "wild-session",
"user_api_key_hash": "key-1",
SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: 777,
},
}
await callback.async_pre_call_deployment_hook(kwargs=request_kwargs, call_type=None)
filtered = await callback.async_filter_deployments(
model="openai/gpt-4o",
healthy_deployments=[
{
"model_name": "openai/*",
"litellm_params": {"model": "openai/gpt-4o"},
"model_info": {"id": f"wild-deployment-{i}"},
}
for i in (1, 2)
],
messages=[],
request_kwargs=request_kwargs,
)
assert [d["model_info"]["id"] for d in filtered] == ["wild-deployment-2"]
@pytest.mark.asyncio
@pytest.mark.parametrize(
"model,session_affinity,deployment_affinity,expect_marker",
[
("smart-router", False, True, True),
("smart-router", True, False, True),
("smart-router", False, False, False),
("target-group", False, True, False),
],
ids=[
"deployment-affinity-stamps",
"session-affinity-implies-deployment-pin",
"both-off-no-stamp",
"non-auto-routed-clears",
],
)
async def test_pre_routing_hook_stamps_or_clears_the_marker_per_attempt(
model, session_affinity, deployment_affinity, expect_marker
):
"""Every routing attempt writes or clears the marker, so a fallback from an
auto-routed group to a plain group cannot carry a stale marker. session_affinity
implies the deployment pin: a session frozen onto one group must not re-shuffle
across that group's deployments."""
router = _smart_router(session_affinity=session_affinity, deployment_affinity=deployment_affinity)
try:
request_kwargs = {
"metadata": {SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: 111},
"litellm_metadata": {SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY: 111},
}
await router.async_pre_routing_hook(
model=model,
request_kwargs=request_kwargs,
messages=[{"role": "user", "content": "Hello"}],
)
if expect_marker:
assert request_kwargs["litellm_metadata"][SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY] == 777
else:
assert SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY not in request_kwargs["metadata"]
assert SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY not in request_kwargs["litellm_metadata"]
finally:
_cleanup_router_callbacks(router)
def test_complexity_router_with_deployment_affinity_registers_affinity_callback():
enabled = _smart_router()
session_only = _smart_router(session_affinity=True, deployment_affinity=False)
disabled = _smart_router(session_affinity=False, deployment_affinity=False)
try:
assert [
(cb.enable_user_key_affinity, cb.enable_responses_api_affinity, cb.enable_session_id_affinity)
for cb in enabled.optional_callbacks or []
if isinstance(cb, DeploymentAffinityCheck)
] == [(False, False, False)]
assert any(isinstance(cb, DeploymentAffinityCheck) for cb in session_only.optional_callbacks or [])
assert not any(isinstance(cb, DeploymentAffinityCheck) for cb in disabled.optional_callbacks or [])
finally:
_cleanup_router_callbacks(enabled)
_cleanup_router_callbacks(session_only)
_cleanup_router_callbacks(disabled)

View file

@ -7550,3 +7550,68 @@ async def test_fallback_failure_detail_from_upstream_is_bounded():
assert capture.messages, "the fallback failure path did not log at ERROR"
assert huge_message not in "".join(capture.messages)
assert max(len(message) for message in capture.messages) < 5_000
def test_stamp_or_clear_metadata_key_writes_and_clears_both_buckets():
request_kwargs = {"metadata": {}}
litellm.Router._stamp_or_clear_metadata_key(request_kwargs=request_kwargs, key="probe", value=7)
assert request_kwargs["metadata"]["probe"] == 7
stale_kwargs = {"metadata": {"probe": 7}, "litellm_metadata": {"probe": 7}}
litellm.Router._stamp_or_clear_metadata_key(request_kwargs=stale_kwargs, key="probe", value=None)
assert "probe" not in stale_kwargs["metadata"]
assert "probe" not in stale_kwargs["litellm_metadata"]
@pytest.mark.parametrize(
"complexity_router_config,expect_callback",
[
({"tiers": {"SIMPLE": "gpt-4o"}}, True),
({"tiers": {"SIMPLE": "gpt-4o"}, "deployment_affinity": False}, False),
({"tiers": {"SIMPLE": "gpt-4o"}, "deployment_affinity": False, "session_affinity": True}, True),
],
)
def test_complexity_router_registers_affinity_callback_for_deployment_pin(complexity_router_config, expect_callback):
"""The marker the complexity router stamps is inert unless a DeploymentAffinityCheck is
registered to read it, so deployment_affinity has to pull the callback in, and its default-on
means a bare config registers one. Opting out must skip the callback entirely rather than
register a filter that can never fire, including when session_affinity is on, since the two
pins are independent."""
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
DeploymentAffinityCheck,
)
router = litellm.Router(
model_list=[
{"model_name": "gpt-4o", "litellm_params": {"model": "gpt-4o"}},
{
"model_name": "my-complexity-router",
"litellm_params": {
"model": "auto_router/complexity_router",
"complexity_router_config": complexity_router_config,
},
},
]
)
try:
registered = any(isinstance(cb, DeploymentAffinityCheck) for cb in router.optional_callbacks or [])
assert registered is expect_callback
finally:
for cb in router.optional_callbacks or []:
litellm.logging_callback_manager.remove_callback_from_all_lists(cb)
def test_ensure_deployment_affinity_callback_is_idempotent():
from litellm.router_utils.pre_call_checks.deployment_affinity_check import (
DeploymentAffinityCheck,
)
router = litellm.Router(model_list=[])
try:
router._ensure_deployment_affinity_callback()
router._ensure_deployment_affinity_callback()
affinity_callbacks = [
cb for cb in router.optional_callbacks or [] if isinstance(cb, DeploymentAffinityCheck)
]
assert len(affinity_callbacks) == 1
finally:
for cb in router.optional_callbacks or []:
litellm.logging_callback_manager.remove_callback_from_all_lists(cb)

View file

@ -27,7 +27,7 @@
"limit": 0
},
"LIT010": {
"limit": 16760
"limit": 16758
},
"LIT011": {
"limit": 5598

View file

@ -31695,6 +31695,12 @@ export interface components {
* @description Default model to use if tier cannot be determined
*/
default_model?: string | null;
/**
* Deployment Affinity
* @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.
* @default true
*/
deployment_affinity: boolean;
/**
* Dimension Weights
* @description Weights for each scoring dimension
@ -31752,13 +31758,13 @@ export interface components {
semantic_keyword_matching: boolean;
/**
* Session Affinity
* @description When True and a session_id is resolvable on the request, pin the model chosen on the session's first turn and reuse it for every later turn, skipping re-classification. Off by default so every turn is classified on its own merits and routed to the cheapest adequate tier. Set True to keep a multi-turn session on one model, which preserves provider prompt caches and avoids cross-model conversation-history errors.
* @description When True and a session_id is resolvable on the request, pin the model chosen on the session's first turn and reuse it for every later turn, skipping re-classification. Off by default so every turn is classified on its own merits and routed to the cheapest adequate tier. Set True to keep a multi-turn session on one model, which preserves provider prompt caches and avoids cross-model conversation-history errors. Always implies the deployment pin regardless of deployment_affinity: the session sticks to one deployment of the pinned model, since freezing the model while re-shuffling its deployments would still go cache-cold.
* @default false
*/
session_affinity: boolean;
/**
* Session Affinity Ttl Seconds
* @description TTL for the session affinity pin; refreshed on every cache hit
* @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 idle time for the session's routing decisions rather than total session length
* @default 3600
*/
session_affinity_ttl_seconds: number;