diff --git a/litellm/caching/affinity_cache.py b/litellm/caching/affinity_cache.py index 2d27574b098..d9a40c44d2b 100644 --- a/litellm/caching/affinity_cache.py +++ b/litellm/caching/affinity_cache.py @@ -24,6 +24,16 @@ if current == false then end if ARGV[3] then local decoded, stored = pcall(cjson.decode, current) + if ARGV[4] then + local legacy = current + if decoded and type(stored) == 'string' then legacy = stored end + for _, eligible in ipairs(cjson.decode(ARGV[3])) do + if eligible['model_id'] == legacy then + redis.call('EXPIRE', KEYS[1], ARGV[2]) + return current + end + end + end if decoded and type(stored) == 'table' then for _, eligible in ipairs(cjson.decode(ARGV[3])) do local matches = true @@ -68,6 +78,7 @@ def claim_affinity_pin_in_memory( ttl_seconds: int, *, eligible_values: tuple[Mapping[str, str], ...] | None = None, + preserve_legacy_model_ids: bool = False, ) -> object: """No await between read and write, so same-loop claims agree during a Redis outage.""" existing: Final[object] = cache.in_memory_cache.get_cache(cache_key) @@ -75,7 +86,12 @@ def claim_affinity_pin_in_memory( 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 + eligible_existing: Final = existing in (eligible_values or ()) or ( + preserve_legacy_model_ids + and isinstance(existing, str) + and any("model_id" in value and _legacy_pin_matches(existing, value) for value in (eligible_values or ())) + ) + winner: Final = existing if existing is not None and eligible_existing else pin_value set_local_affinity_pin(cache, cache_key, winner, ttl_seconds) return winner @@ -94,6 +110,7 @@ async def claim_affinity_pin( ttl_seconds: int, *, eligible_values: tuple[Mapping[str, str], ...] | None = None, + preserve_legacy_model_ids: bool = False, ) -> object: """Return the authoritative first writer, replacing it only when it becomes ineligible. @@ -112,6 +129,7 @@ async def claim_affinity_pin( if eligible_values is not None else () ), + *(("1",) if preserve_legacy_model_ids and 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) @@ -124,4 +142,11 @@ async def claim_affinity_pin( 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) + return claim_affinity_pin_in_memory( + cache, + cache_key, + pin_value, + ttl_seconds, + eligible_values=eligible_values, + preserve_legacy_model_ids=preserve_legacy_model_ids, + ) diff --git a/litellm/caching/redis_cache.py b/litellm/caching/redis_cache.py index eb17facfb12..afe6f8a1eab 100644 --- a/litellm/caching/redis_cache.py +++ b/litellm/caching/redis_cache.py @@ -1738,7 +1738,9 @@ class RedisCache(BaseCache): return key_value_dict @_redis_circuit_breaker_guard - async def async_get_cache(self, key, parent_otel_span: Span | None = None, **kwargs): + async def async_get_cache( + self, key, parent_otel_span: Span | None = None, *, raise_on_error: bool = False, **kwargs + ): from redis.asyncio import Redis _redis_client: Final[Redis] = self.init_async_client() @@ -1781,6 +1783,8 @@ class RedisCache(BaseCache): ) ) print_verbose(f"litellm.caching.caching: async get() - Got exception from REDIS: {e}") + if raise_on_error: + raise _record_swallowed_redis_failure(self._circuit_breaker, e) @_redis_circuit_breaker_guard diff --git a/litellm/router.py b/litellm/router.py index 25591a435f9..96102b678c3 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -1247,6 +1247,7 @@ class Router: fallback_access_check: FallbackAccessCheck | None = None, fallback_budget_check: FallbackBudgetCheck | None = None, auto_router_capability_limit: AutoRouterCapabilityLimit | None = None, + session_affinity_repin_on_failover: bool = False, ) -> None: """ Initialize the Router class with the given parameters for caching, reliability, and routing strategy. @@ -1283,6 +1284,8 @@ class Router: alerting_config (AlertingConfig): Slack alerting configuration. Defaults to None. provider_budget_config (ProviderBudgetConfig): Provider budget configuration. Use this to set llm_provider budget limits. example $100/day to OpenAI, $100/day to Azure, etc. Defaults to None. deployment_affinity_ttl_seconds (int): TTL for user-key -> deployment affinity mapping. Defaults to 3600. + session_affinity_repin_on_failover (bool): After a successful replacement response, move a session pin + whose deployment is in cooldown or removed. Async routes only. Defaults to False. Applied at startup. ignore_invalid_deployments (bool): Ignores invalid deployments, and continues with other deployments. Default is to raise an error. enable_weighted_failover (bool): When True and the routing strategy is "simple-shuffle", a retryable failure on one deployment causes the request to re-pick (weighted) across the other deployments in the same model group before any cross-group fallback runs. Bounded by `max_fallbacks`. Async-only: currently honored by `router.acompletion()` and other async entrypoints. The sync `router.completion()` path falls back to the regular fallback flow. Defaults to False. enable_mid_stream_fallback_continuation (bool): When True, a chat-completions stream that breaks after plain assistant text has been delivered continues on a fallback deployment via assistant prefill instead of surfacing the error. Only deployments whose model supports assistant prefill are eligible, so the partial text is continued, not regenerated; if none is, the original error is surfaced. Streams that emitted tool calls, thinking blocks, audio/images, or constrained (JSON / forced tool_choice) output are never continued. Async-only. Defaults to False. @@ -1467,6 +1470,7 @@ class Router: self._provider_unresolved_deployments: tuple[Callable[[], Deployment | None], ...] = () self.deployment_affinity_ttl_seconds = deployment_affinity_ttl_seconds + self.session_affinity_repin_on_failover = session_affinity_repin_on_failover self.model_group_affinity_config = model_group_affinity_config warn_on_unknown_model_group_affinity_flags(model_group_affinity_config) @@ -2862,6 +2866,7 @@ class Router: enable_session_id_affinity=False, model_group_affinity_config=self.model_group_affinity_config, is_priority_group=self._is_priority_routing_group, + session_affinity_repin_on_failover=self.session_affinity_repin_on_failover, ) self.optional_callbacks.append(affinity_callback) litellm.logging_callback_manager.add_litellm_callback(affinity_callback) @@ -2908,6 +2913,7 @@ class Router: enable_session_id_affinity=enable_session_id_affinity, model_group_affinity_config=self.model_group_affinity_config, is_priority_group=self._is_priority_routing_group, + session_affinity_repin_on_failover=self.session_affinity_repin_on_failover, ) self.optional_callbacks.append(affinity_callback) litellm.logging_callback_manager.add_litellm_callback(affinity_callback) @@ -9300,6 +9306,12 @@ class Router: if model_group is None or deployment_id is None or self.get_deployment(model_id=str(deployment_id)) is None: return + if self.session_affinity_repin_on_failover: + await self._repin_session_after_success( + _LITELLM_PARAMS_ADAPTER.validate_python(litellm_params), + _OPAQUE_VALUE_ADAPTER.validate_python(completion_response), + ) + # Always track deployment successes for cooldown logic, regardless of TPM/RPM limits increment_deployment_successes_for_current_minute( litellm_router_instance=self, @@ -9323,6 +9335,36 @@ class Router: "litellm.router.Router::deployment_callback_on_success(): Exception occured - %s", e ) + async def _repin_session_after_success(self, litellm_params: Mapping[str, object], response: object) -> None: + if not self.session_affinity_repin_on_failover: + return + try: + from litellm.router_utils.pre_call_checks.deployment_affinity_check import completed_affinity_response + + if not completed_affinity_response(response): + return + callback: Final = next( + (logger for logger in self._optional_loggers() if isinstance(logger, DeploymentAffinityCheck)), + None, + ) + if callback is None: + return + model_group: Final = callback.get_repin_model_group(litellm_params) + if model_group is None or self._is_priority_routing_group(model_group): + return + registered_ids: Final = self.get_candidate_model_ids_for_route( + model=model_group, team_id=get_request_team_id(litellm_params) + ) + cooldown_ids: Final = frozenset( + model_id + for model_id, _ in await self.cooldown_cache.async_get_active_cooldowns( + model_ids=list(registered_ids), parent_otel_span=None + ) + ) + await callback.async_repin_session_after_success(litellm_params, model_group, registered_ids - cooldown_ids) + except Exception as error: # noqa: BLE001 # Optional pin updates must not interrupt success accounting + verbose_router_logger.debug("Session re-pinning failed: %s", error) + async def increment_deployment_usage_for_response( self, response: object, 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 88fb8876175..8693eaa5151 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -16,6 +16,7 @@ import hashlib from collections.abc import Callable, Mapping, Sequence from typing import Final, cast +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict from litellm._internal_context import with_service_target @@ -38,6 +39,27 @@ class DeploymentAffinityCacheValue(TypedDict, closed=True): model_id: ReadOnly[str] +def _affinity_response_fields(response: object) -> Mapping[str, object]: + return ( + TypeAdapter(Mapping[str, object]).validate_python(response) + if isinstance(response, Mapping) + else {name: getattr(response, name, None) for name in ("response", "status", "error")} + ) + + +def completed_affinity_response(response: object) -> bool: + if response is None: + return False + envelope: Final = _affinity_response_fields(response) + if envelope.get("error") is not None or envelope.get("status") not in (None, "completed"): + return False + nested: Final = envelope.get("response") + fields: Final = _affinity_response_fields(nested) if nested is not None else envelope + return ( + fields.get("response") is None and fields.get("error") is None and fields.get("status") in (None, "completed") + ) + + VALID_MODEL_GROUP_AFFINITY_FLAGS: Final = frozenset( { "deployment_affinity", @@ -86,6 +108,7 @@ class DeploymentAffinityCheck(CustomLogger): enable_session_id_affinity: bool = False, model_group_affinity_config: dict[str, list[str]] | None = None, is_priority_group: Callable[[str], bool] | None = None, + session_affinity_repin_on_failover: bool = False, ): super().__init__() self.cache = cache @@ -95,6 +118,7 @@ class DeploymentAffinityCheck(CustomLogger): self.enable_session_id_affinity = enable_session_id_affinity self.model_group_affinity_config: dict[str, list[str]] = model_group_affinity_config or {} self.is_priority_group = is_priority_group + self.session_affinity_repin_on_failover = session_affinity_repin_on_failover def _get_effective_flags(self, model_group: str) -> tuple[bool, bool, bool]: """ @@ -297,7 +321,7 @@ class DeploymentAffinityCheck(CustomLogger): return f"{cls.USER_ID_AFFINITY_PREFIX}{user_id}" @staticmethod - def _get_session_id_from_request_kwargs(request_kwargs: dict) -> str | None: + def _get_session_id_from_request_kwargs(request_kwargs: Mapping[str, object]) -> str | None: for metadata in DeploymentAffinityCheck._iter_metadata_dicts(request_kwargs): session_id = DeploymentAffinityCheck._get_session_id_from_metadata_dict(metadata=metadata) if session_id is not None: @@ -305,7 +329,7 @@ class DeploymentAffinityCheck(CustomLogger): return None @staticmethod - def _get_marker_session_affinity_ttl(request_kwargs: dict) -> int | None: + def _get_marker_session_affinity_ttl(request_kwargs: Mapping[str, object]) -> 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 @@ -357,7 +381,7 @@ class DeploymentAffinityCheck(CustomLogger): model: str, healthy_deployments: list, messages: list[AllMessageValues] | None, - request_kwargs: dict | None = None, + request_kwargs: dict[str, object] | None = None, parent_otel_span: Span | None = None, ) -> list[dict]: """ @@ -418,7 +442,15 @@ class DeploymentAffinityCheck(CustomLogger): session_cache_key: Final = self.get_session_affinity_cache_key( 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) + session_cache_result: Final = ( + await self._read_shared_session_pin( + session_cache_key, + self._get_marker_session_affinity_ttl(request_kwargs) or self.ttl_seconds, + parent_otel_span, + ) + if self.session_affinity_repin_on_failover + else await self.cache.async_get_cache(key=session_cache_key) + ) session_model_id: str | None = None if isinstance(session_cache_result, dict): @@ -482,6 +514,71 @@ class DeploymentAffinityCheck(CustomLogger): ) return [deployment] + async def _read_shared_session_pin(self, cache_key: str, ttl_seconds: int, parent_otel_span: Span | None) -> object: + redis_cache: Final = self.cache.redis_cache + if redis_cache is None: + return TypeAdapter(object).validate_python(await self.cache.async_get_cache(key=cache_key)) + local_before: Final = TypeAdapter(object).validate_python(self.cache.in_memory_cache.get_cache(cache_key)) + try: + shared: Final[object] = await redis_cache.async_get_cache( + key=cache_key, parent_otel_span=parent_otel_span, raise_on_error=True + ) + except Exception: # noqa: BLE001 # Optional affinity falls back locally on Redis admission or read failure + verbose_router_logger.debug("Session pin read failed, using pod-local affinity") + return TypeAdapter(object).validate_python(await self.cache.async_get_cache(key=cache_key, local_only=True)) + if self.cache.in_memory_cache.get_cache(cache_key) is local_before: + if shared is None: + self.cache.in_memory_cache.delete_cache(cache_key) + else: + self._set_local_pin(cache_key, shared, ttl_seconds) + return shared + + def get_repin_model_group(self, request_kwargs: Mapping[str, object]) -> str | None: + model_group: Final = next( + ( + value + for metadata in self._iter_metadata_dicts(request_kwargs) + if isinstance(value := metadata.get("deployment_model_name"), str) and value + ), + None, + ) + if model_group is None or self._get_session_id_from_request_kwargs(request_kwargs) is None: + return None + if ( + not self._get_effective_flags(model_group)[2] + and self._get_marker_session_affinity_ttl(request_kwargs) is None + ): + return None + return model_group + + async def async_repin_session_after_success( + self, + request_kwargs: Mapping[str, object], + model_group: str, + eligible_deployment_ids: frozenset[str], + ) -> None: + marker_ttl: Final = self._get_marker_session_affinity_ttl(request_kwargs) + if not self._get_effective_flags(model_group)[2] and marker_ttl is None: + return + session_id: Final = self._get_session_id_from_request_kwargs(request_kwargs) + if session_id is None: + return + model_info: Final = TypeAdapter(Mapping[str, object]).validate_python(request_kwargs.get("model_info") or {}) + model_id: Final = model_info.get("id") + if model_id is None or str(model_id) not in eligible_deployment_ids: + return + cache_key: Final = self.get_session_affinity_cache_key( + model_group, session_id, self.get_user_key_from_request_kwargs(request_kwargs) + ) + await claim_affinity_pin( + self.cache, + cache_key, + DeploymentAffinityCacheValue(model_id=str(model_id)), + marker_ttl if marker_ttl is not None else self.ttl_seconds, + eligible_values=tuple(DeploymentAffinityCacheValue(model_id=value) for value in eligible_deployment_ids), + preserve_legacy_model_ids=True, + ) + async def async_pre_call_deployment_hook( self, kwargs: dict[str, object], call_type: CallTypes | None ) -> dict | None: diff --git a/tests/integration/sdk/test_dual_cache_redis.py b/tests/integration/sdk/test_dual_cache_redis.py index ddd01480d36..3527dbf91ee 100644 --- a/tests/integration/sdk/test_dual_cache_redis.py +++ b/tests/integration/sdk/test_dual_cache_redis.py @@ -6,6 +6,7 @@ from unittest.mock import patch import pytest +from litellm.caching.affinity_cache import claim_affinity_pin from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.caching.redis_cache import RedisCache @@ -92,3 +93,66 @@ async def test_sync_and_async_batch_reads_share_one_redis_without_sync_reads_goi assert in_loop_results == [[expected[key]] for key in in_loop_keys] assert await dual_cache.async_batch_get_cache(keys=[survivor_key]) == [expected[survivor_key]] + + +@pytest.mark.asyncio +async def test_concurrent_session_repin_claims_keep_the_shared_winner_after_recovery() -> None: + redis: Final = _redis_cache() + workers: Final = (DualCache(redis_cache=redis), DualCache(redis_cache=redis)) + key: Final = f"session-repin-{uuid.uuid4().hex}" + choices: Final = ({"model_id": "replacement"}, {"model_id": "other"}) + try: + await redis.async_set_cache(key, {"model_id": "original"}, ttl=60) + for worker in workers: + worker.in_memory_cache.set_cache(key, {"model_id": "original"}, ttl=60) + winners: Final = await asyncio.gather( + *( + claim_affinity_pin(worker, key, choice, 60, eligible_values=choices, preserve_legacy_model_ids=True) + for worker, choice in zip(workers, choices, strict=True) + ) + ) + assert winners[0] == winners[1] and winners[0] in choices + late_original: Final = await claim_affinity_pin( + workers[0], + key, + {"model_id": "original"}, + 60, + eligible_values=({"model_id": "original"}, *choices), + preserve_legacy_model_ids=True, + ) + assert late_original == winners[0] + assert await redis.async_get_cache(key) == winners[0] + finally: + await redis.async_delete_cache(key) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stored", ["original", '"original"', '{"model_id":"original"}']) +async def test_session_repin_lua_preserves_healthy_legacy_pins_and_replaces_unavailable_ones(stored: str) -> None: + redis: Final = _redis_cache() + cache: Final = DualCache(redis_cache=redis) + key: Final = f"session-repin-legacy-{uuid.uuid4().hex}" + replacement: Final = {"model_id": "replacement"} + try: + redis.redis_client.set(redis.check_and_fix_namespace(key), stored, ex=60) + preserved: Final = await claim_affinity_pin( + cache, + key, + replacement, + 60, + eligible_values=({"model_id": "original"}, replacement), + preserve_legacy_model_ids=True, + ) + assert preserved in ("original", {"model_id": "original"}) + moved: Final = await claim_affinity_pin( + cache, + key, + replacement, + 60, + eligible_values=(replacement,), + preserve_legacy_model_ids=True, + ) + assert moved == replacement + assert await redis.async_get_cache(key) == replacement + finally: + await redis.async_delete_cache(key) diff --git a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py index 9bbaed0ae1b..6b4ea74894a 100644 --- a/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py +++ b/tests/unit/router_utils/pre_call_checks/test_session_id_affinity.py @@ -1,20 +1,368 @@ import asyncio import json from typing import Final +from collections.abc import AsyncIterator, Mapping +from datetime import datetime, timezone from unittest.mock import AsyncMock, MagicMock, patch import pytest +from fakeredis import FakeRedis, FakeServer +from fakeredis.aioredis import FakeConnection +from redis.asyncio import BlockingConnectionPool, Redis import litellm from litellm.caching.affinity_cache import claim_affinity_pin from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache +from litellm.caching.redis_cache import RedisCache from litellm.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, ) +def _repinning_router(enabled: bool | None = True) -> litellm.Router: + return litellm.Router( + model_list=[ + { + "model_name": "repin-group", + "litellm_params": {"model": "openai/test-model", "api_key": "fake-api-key"}, + "model_info": {"id": deployment_id}, + } + for deployment_id in ("original", "replacement", "other") + ] + + [ + { + "model_name": "unrelated-group", + "litellm_params": {"model": "openai/test-model", "api_key": "fake-api-key"}, + "model_info": {"id": "unrelated"}, + } + ], + optional_pre_call_checks=["session_affinity", "deployment_affinity"], + **({"session_affinity_repin_on_failover": enabled} if enabled is not None else {}), + ) + + +def _repin_params(deployment_id: str, caller: str = "caller-1") -> dict[str, object]: + return { + "model_info": {"id": deployment_id}, + "metadata": { + "model_group": "repin-group", + "deployment_model_name": "repin-group", + "session_id": "repin-session", + "user_api_key_hash": caller, + }, + } + + +def _repin_key(caller: str = "caller-1") -> str: + return DeploymentAffinityCheck.get_session_affinity_cache_key("repin-group", "repin-session", caller) + + +def _cool_down_original(router: litellm.Router) -> None: + router.cooldown_cache.add_deployment_to_cooldown( + model_id="original", original_exception=RuntimeError("test outage"), exception_status=503, cooldown_time=3600 + ) + + +async def _record_replacement_success( + router: litellm.Router, deployment_id: str = "replacement", response: object = None +) -> None: + timestamp: Final = datetime(2026, 1, 1, tzinfo=timezone.utc) + await router.deployment_callback_on_success( + {"litellm_params": _repin_params(deployment_id), "standard_logging_object": {"total_tokens": 0}}, + {"choices": [{"message": {"content": "ok"}}]} if response is None else response, + timestamp, + timestamp, + ) + + +@pytest.mark.parametrize("enabled", [None, False, True], ids=["default", "disabled", "enabled"]) +@pytest.mark.parametrize("unavailable", ["cooldown", "removed", "temporary-filter"]) +async def test_session_repin_requires_opt_in_unavailable_original_and_success( + enabled: bool | None, unavailable: str +) -> None: + router: Final = _repinning_router(enabled) + try: + await router.cache.async_set_cache(_repin_key(), {"model_id": "original"}) + await router.cache.async_set_cache(_repin_key("caller-2"), {"model_id": "original"}) + user_key: Final = DeploymentAffinityCheck.get_affinity_cache_key("repin-group", "caller-1") + await router.cache.async_set_cache(user_key, {"model_id": "original"}) + if unavailable == "cooldown": + _cool_down_original(router) + elif unavailable == "removed": + router.delete_deployment("original") + callback: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck)) + await callback.async_pre_call_deployment_hook(_repin_params("replacement"), None) + assert await router.cache.async_get_cache(_repin_key()) == {"model_id": "original"} + await _record_replacement_success(router) + expected: Final = "replacement" if enabled and unavailable != "temporary-filter" else "original" + assert await router.cache.async_get_cache(_repin_key()) == {"model_id": expected} + assert await router.cache.async_get_cache(_repin_key("caller-2")) == {"model_id": "original"} + assert await router.cache.async_get_cache(user_key) == {"model_id": "original"} + router.cooldown_cache.cooldown_store.flush_cache() + await _record_replacement_success(router, "other") + assert await router.cache.async_get_cache(_repin_key()) == {"model_id": expected} + finally: + router.discard() + + +@pytest.mark.parametrize("status", ["failed", "incomplete", "in_progress", "cancelled"]) +@pytest.mark.parametrize("wrapped", [False, True]) +async def test_unsuccessful_responses_do_not_repin(status: str, wrapped: bool) -> None: + router: Final = _repinning_router() + try: + await router.cache.async_set_cache(_repin_key(), {"model_id": "original"}) + _cool_down_original(router) + response: Final[Mapping[str, object]] = {"status": status} + await _record_replacement_success(router, response={"response": response} if wrapped else response) + assert await router.cache.async_get_cache(_repin_key()) == {"model_id": "original"} + finally: + router.discard() + + +async def test_concurrent_successful_replacements_keep_first_winner() -> None: + router: Final = _repinning_router() + try: + await router.cache.async_set_cache(_repin_key(), {"model_id": "original"}) + _cool_down_original(router) + await asyncio.gather(_record_replacement_success(router), _record_replacement_success(router, "other")) + winner: Final = await router.cache.async_get_cache(_repin_key()) + assert winner in ({"model_id": "replacement"}, {"model_id": "other"}) + router.cooldown_cache.cooldown_store.flush_cache() + await _record_replacement_success(router, "original") + assert await router.cache.async_get_cache(_repin_key()) == winner + finally: + router.discard() + + +async def test_healthy_legacy_string_pin_is_preserved_with_repinning_enabled() -> None: + router: Final = _repinning_router() + try: + await router.cache.async_set_cache(_repin_key(), "original") + await _record_replacement_success(router) + assert await router.cache.async_get_cache(_repin_key()) == "original" + finally: + router.discard() + + +@pytest.mark.parametrize("route_bucket", ["metadata", "litellm_metadata"]) +async def test_repin_session_after_success_reads_split_metadata(route_bucket: str) -> None: + router: Final = _repinning_router() + try: + await router.cache.async_set_cache(_repin_key(), {"model_id": "original"}) + _cool_down_original(router) + session_bucket: Final = "litellm_metadata" if route_bucket == "metadata" else "metadata" + await router._repin_session_after_success( + { + "model_info": {"id": "replacement"}, + route_bucket: {"deployment_model_name": "repin-group"}, + session_bucket: {"session_id": "repin-session", "user_api_key_hash": "caller-1"}, + }, + {"status": "completed"}, + ) + assert await router.cache.async_get_cache(_repin_key()) == {"model_id": "replacement"} + finally: + router.discard() + + +@pytest.mark.parametrize("enabled", [False, True]) +async def test_session_selection_observes_another_workers_repin_only_when_enabled(enabled: bool) -> None: + server: Final = FakeServer() + sync_client: Final = FakeRedis(server=server) + pool: Final = BlockingConnectionPool(connection_class=FakeConnection, server=server) + router: Final = _repinning_router(enabled) + try: + with ( + patch("redis.Redis.from_url", return_value=sync_client), + patch("redis.asyncio.BlockingConnectionPool.from_url", return_value=pool), + ): + redis: Final = RedisCache(url=f"redis://repin-cache.example:6379/{int(enabled)}") + router.cache.attach_redis_cache(redis) + other_worker: Final = DualCache(redis_cache=redis) + router.cache.in_memory_cache.set_cache(_repin_key(), {"model_id": "original"}, ttl=3600) + await other_worker.async_set_cache(_repin_key(), {"model_id": "replacement"}, ttl=3600) + callback: Final = next( + cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck) + ) + selected: Final = await callback.async_filter_deployments( + model="repin-group", + healthy_deployments=router.get_model_list(model_name="repin-group") or [], + messages=None, + request_kwargs=_repin_params("replacement"), + ) + assert [deployment["model_info"]["id"] for deployment in selected] == [ + "replacement" if enabled else "original" + ] + finally: + router.discard() + await pool.aclose() + sync_client.close() + + +@pytest.fixture +async def shared_session_router() -> AsyncIterator[tuple[litellm.Router, FakeServer, RedisCache]]: + server: Final = FakeServer() + sync_client: Final = FakeRedis(server=server) + pool: Final = BlockingConnectionPool(connection_class=FakeConnection, server=server) + router: Final = _repinning_router() + try: + with ( + patch("redis.Redis.from_url", return_value=sync_client), + patch("redis.asyncio.BlockingConnectionPool.from_url", return_value=pool), + ): + redis: Final = RedisCache(url="redis://session-outage.example:6379/0") + router.cache.attach_redis_cache(redis) + yield router, server, redis + finally: + router.discard() + await pool.aclose() + sync_client.close() + + +async def _selected_session_ids(router: litellm.Router) -> tuple[str, ...]: + callback: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck)) + selected: Final = await callback.async_filter_deployments( + model="repin-group", + healthy_deployments=router.get_model_list(model_name="repin-group") or [], + messages=None, + request_kwargs=_repin_params("replacement"), + ) + return tuple(deployment["model_info"]["id"] for deployment in selected) + + +@pytest.mark.parametrize("has_pin", [False, True]) +async def test_redis_read_failures_and_open_breaker_preserve_session_routing( + shared_session_router: tuple[litellm.Router, FakeServer, RedisCache], has_pin: bool +) -> None: + router, server, redis = shared_session_router + if has_pin: + await redis.async_set_cache(_repin_key(), {"model_id": "replacement"}, ttl=3600) + router.cache.in_memory_cache.set_cache(_repin_key(), {"model_id": "original"}, ttl=3600) + assert await _selected_session_ids(router) == ("replacement",) + server.connected = False + with patch("time.time", return_value=1790000000.0): + for _ in range(redis._circuit_breaker.failure_threshold + 1): + assert set(await _selected_session_ids(router)) == ( + {"replacement"} if has_pin else {"original", "replacement", "other"} + ) + assert redis._circuit_breaker.is_open() + + +async def test_confirmed_redis_miss_does_not_revive_a_local_pin_during_outage( + shared_session_router: tuple[litellm.Router, FakeServer, RedisCache], +) -> None: + router, server, _ = shared_session_router + router.cache.in_memory_cache.set_cache(_repin_key(), {"model_id": "original"}, ttl=3600) + assert set(await _selected_session_ids(router)) == {"original", "replacement", "other"} + server.connected = False + assert set(await _selected_session_ids(router)) == {"original", "replacement", "other"} + + +async def test_shared_session_pin_resumes_after_a_transient_redis_failure( + shared_session_router: tuple[litellm.Router, FakeServer, RedisCache], +) -> None: + router, server, redis = shared_session_router + await redis.async_set_cache(_repin_key(), {"model_id": "replacement"}, ttl=3600) + router.cache.in_memory_cache.set_cache(_repin_key(), {"model_id": "original"}, ttl=3600) + server.connected = False + assert await _selected_session_ids(router) == ("original",) + server.connected = True + assert await _selected_session_ids(router) == ("replacement",) + + +@pytest.mark.parametrize("shared_value", [None, {"model_id": "original"}]) +async def test_delayed_shared_pin_read_preserves_a_newer_local_write( + shared_session_router: tuple[litellm.Router, FakeServer, RedisCache], shared_value: Mapping[str, str] | None +) -> None: + router, server, redis = shared_session_router + if shared_value is not None: + await redis.async_set_cache(_repin_key(), shared_value, ttl=3600) + observed: Final = asyncio.Event() + resume: Final = asyncio.Event() + + class DelayedReader(Redis): + async def get(self, name: str) -> object: + value: Final = await super().get(name) + observed.set() + await resume.wait() + return value + + reader: Final = DelayedReader(connection_pool=redis.init_async_client().connection_pool) + litellm.in_memory_llm_clients_cache.set_cache(key=redis._get_async_client_cache_key(), value=reader) + pending: Final = asyncio.create_task(_selected_session_ids(router)) + try: + await observed.wait() + router.cache.in_memory_cache.set_cache(_repin_key(), {"model_id": "replacement"}, ttl=3600) + finally: + resume.set() + await pending + server.connected = False + assert await _selected_session_ids(router) == ("replacement",) + + +@pytest.mark.parametrize( + "response", + [ + {"status": "failed", "response": {"status": "completed"}}, + {"error": {"message": "test failure"}, "response": {"status": "completed"}}, + {"response": {"response": {"status": "completed"}}}, + ], +) +async def test_unsuccessful_or_nested_response_envelopes_never_repin(response: Mapping[str, object]) -> None: + router: Final = _repinning_router() + try: + await router.cache.async_set_cache(_repin_key(), {"model_id": "original"}) + _cool_down_original(router) + await _record_replacement_success(router, response=response) + assert await router.cache.async_get_cache(_repin_key()) == {"model_id": "original"} + finally: + router.discard() + + +class _ReadCountingCache(InMemoryCache): + def __init__(self) -> None: + super().__init__() + self.read_keys: tuple[str, ...] = () + + def get_cache(self, key: str, **kwargs: object) -> object: + self.read_keys += (key,) + return super().get_cache(key, **kwargs) + + +@pytest.mark.parametrize("case", ["disabled", "no-session", "generated-session", "group-disabled", "active"]) +async def test_repin_cooldown_reads_are_limited_to_active_session_route(case: str) -> None: + router: Final = _repinning_router(case != "disabled") + reads: Final = _ReadCountingCache() + try: + router.cooldown_cache.cooldown_store.in_memory_cache = reads + callback: Final = next(cb for cb in router.optional_callbacks or () if isinstance(cb, DeploymentAffinityCheck)) + if case == "group-disabled": + callback.model_group_affinity_config = {"repin-group": []} + await router._repin_session_after_success( + { + "model_info": {"id": "replacement"}, + "metadata": { + "deployment_model_name": "repin-group", + **({"session_id": "repin-session"} if case != "no-session" else {}), + SESSION_ID_GENERATED_METADATA_KEY: case == "generated-session", + }, + }, + {"status": "completed"}, + ) + expected: Final = ( + { + router.cooldown_cache.get_cooldown_cache_key(model_id) + for model_id in ("original", "replacement", "other") + } + if case == "active" + else set() + ) + assert set(reads.read_keys) == expected + finally: + router.discard() + + class MockResponse: def __init__(self, json_data, status_code): self._json_data = json_data