From 5c682ccd1e44cde3d823610bd6155fd75ba30a0e Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 21:56:48 +0000 Subject: [PATCH] refactor(redis): migrate deployment session affinity into a typed store --- .../stores/deployment_session_affinity.py | 135 +++++++++++++ .../deployment_affinity_check.py | 136 +++++++++---- tests/_support/fake_redis.py | 7 + .../test_deployment_session_affinity.py | 181 ++++++++++++++++++ .../test_session_id_affinity.py | 80 ++++---- 5 files changed, 458 insertions(+), 81 deletions(-) create mode 100644 litellm/redis/stores/deployment_session_affinity.py create mode 100644 tests/unit/redis/stores/test_deployment_session_affinity.py diff --git a/litellm/redis/stores/deployment_session_affinity.py b/litellm/redis/stores/deployment_session_affinity.py new file mode 100644 index 00000000000..94ca36683e3 --- /dev/null +++ b/litellm/redis/stores/deployment_session_affinity.py @@ -0,0 +1,135 @@ +import ast +import json +from dataclasses import dataclass +from typing import Final, Protocol + +from pydantic import JsonValue, TypeAdapter, ValidationError + +from litellm.redis.transport.contract import RedisFailure, RedisTransport + +_JSON: Final = TypeAdapter[JsonValue](JsonValue) +_OBJECT: Final = TypeAdapter(object) +_CLAIM: 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) + 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 + 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 affinity_key(model_group: str, caller_scope: str, session_id: str) -> str: + return f"deployment_affinity:v1:session:{model_group}:{caller_scope}:{session_id}" + + +@dataclass(frozen=True, slots=True) +class SessionPin: + value: object + + +@dataclass(frozen=True, slots=True) +class UnconfirmedSessionPin: + pass + + +class DeploymentSessionAffinityStore(Protocol): + async def read(self, model_group: str, caller_scope: str, session_id: str) -> SessionPin | RedisFailure: ... + + async def claim( + self, + model_group: str, + caller_scope: str, + session_id: str, + model_id: str, + *, + ttl_seconds: int, + eligible_model_ids: frozenset[str] | None = None, + ) -> SessionPin | UnconfirmedSessionPin | RedisFailure: ... + + +@dataclass(frozen=True, slots=True) +class TransportDeploymentSessionAffinityStore: + transport: RedisTransport + + async def read(self, model_group: str, caller_scope: str, session_id: str) -> SessionPin | RedisFailure: + raw: Final = await self.transport.get(affinity_key(model_group, caller_scope, session_id)) + if isinstance(raw, RedisFailure): + return raw + if raw is None: + return SessionPin(None) + try: + decoded: Final = raw.decode("utf-8") + try: + return SessionPin(_JSON.validate_json(decoded)) + except ValidationError: + return SessionPin(_OBJECT.validate_python(ast.literal_eval(decoded))) + except (UnicodeError, ValueError, SyntaxError) as error: + return RedisFailure(kind="decode", message=str(error)) + + async def claim( + self, + model_group: str, + caller_scope: str, + session_id: str, + model_id: str, + *, + ttl_seconds: int, + eligible_model_ids: frozenset[str] | None = None, + ) -> SessionPin | UnconfirmedSessionPin | RedisFailure: + raw: Final = await self.transport.eval( + _CLAIM, + (affinity_key(model_group, caller_scope, session_id),), + ( + json.dumps({"model_id": model_id}), + int(ttl_seconds), + *( + (json.dumps(tuple({"model_id": value} for value in eligible_model_ids)),) + if eligible_model_ids is not None + else () + ), + ), + ) + if isinstance(raw, RedisFailure): + return raw + if not isinstance(raw, bytes): + return UnconfirmedSessionPin() + try: + decoded: Final = raw.decode("utf-8") + except UnicodeError as error: + return RedisFailure(kind="decode", message=str(error)) + try: + return SessionPin(_JSON.validate_json(decoded)) + except ValidationError: + return SessionPin(decoded) 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 e8e12db5148..6a9e6a984a7 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -23,7 +23,6 @@ from litellm._internal_context import with_service_target from litellm._logging import verbose_router_logger from litellm.caching.affinity_cache import ( ROUTER_SESSION_PINS_TARGET, - claim_affinity_pin, claim_affinity_pin_in_memory, set_local_affinity_pin, ) @@ -31,6 +30,15 @@ from litellm.caching.dual_cache import DualCache from litellm.caching.redis_batch import active_request_redis_batch from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger, Span +from litellm.redis.stores.deployment_session_affinity import ( + DeploymentSessionAffinityStore, + SessionPin, + TransportDeploymentSessionAffinityStore, + UnconfirmedSessionPin, +) +from litellm.redis.stores.deployment_session_affinity import ( + affinity_key as session_affinity_key, +) from litellm.redis.stores.deployment_user_affinity import ( DeploymentUserAffinityStore, TransportDeploymentUserAffinityStore, @@ -280,7 +288,7 @@ class DeploymentAffinityCheck(CustomLogger): `"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}" + return session_affinity_key(model_group, hashed_user_key, session_id) @staticmethod def _get_session_id_from_metadata_dict(metadata: Mapping[object, object]) -> str | None: @@ -399,8 +407,45 @@ class DeploymentAffinityCheck(CustomLogger): return self._pinned_model_id(result.value) return self._claim_pin_in_memory(cache_key, DeploymentAffinityCacheValue(model_id=model_id), self.ttl_seconds) - async def _claim_pin(self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int) -> str | None: - winner: Final = await claim_affinity_pin(self.cache, cache_key, pin_value, ttl_seconds) + async def _claim_session_pin( + self, + model_group: str, + session_id: str, + user_key: str | None, + model_id: str, + ttl_seconds: int, + *, + eligible_model_ids: frozenset[str] | None = None, + ) -> str | None: + cache_key: Final = self.get_session_affinity_cache_key(model_group, session_id, user_key) + redis_cache: Final = self.cache.redis_cache + if redis_cache is not None: + store: Final[DeploymentSessionAffinityStore] = TransportDeploymentSessionAffinityStore( + redis_cache.transport() + ) + result: Final = await store.claim( + model_group, + self._hash_user_key(user_key) if user_key is not None else "unscoped", + session_id, + model_id, + ttl_seconds=ttl_seconds, + eligible_model_ids=eligible_model_ids, + ) + if isinstance(result, UnconfirmedSessionPin): + return model_id + if isinstance(result, SessionPin): + self._set_local_pin(cache_key, result.value, ttl_seconds) + return self._pinned_model_id(result.value) + winner: Final = claim_affinity_pin_in_memory( + self.cache, + cache_key, + DeploymentAffinityCacheValue(model_id=model_id), + ttl_seconds, + eligible_values=tuple(DeploymentAffinityCacheValue(model_id=value) for value in eligible_model_ids) + if eligible_model_ids is not None + else None, + preserve_legacy_model_ids=eligible_model_ids is not None, + ) return self._pinned_model_id(winner) def _claim_pin_in_memory( @@ -487,17 +532,12 @@ class DeploymentAffinityCheck(CustomLogger): 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, user_key=user_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_cache_result: Final = await self._read_session_pin( + stable_model_map_key, + session_id, + user_key, + self._get_marker_session_affinity_ttl(request_kwargs) or self.ttl_seconds, + parent_otel_span, ) session_model_id: str | None = None @@ -561,18 +601,36 @@ class DeploymentAffinityCheck(CustomLogger): ) return [deployment] - async def _read_shared_session_pin(self, cache_key: str, ttl_seconds: int, parent_otel_span: Span | None) -> object: + async def _read_session_pin( + self, + model_group: str, + session_id: str, + user_key: str | None, + ttl_seconds: int, + parent_otel_span: Span | None, + ) -> object: + cache_key: Final = self.get_session_affinity_cache_key(model_group, session_id, user_key) 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: Final = TypeAdapter(object).validate_python( + await self.cache.async_get_cache(key=cache_key, local_only=True) + ) + if redis_cache is None or (not self.session_affinity_repin_on_failover and local is not None): + return local 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") + batch: Final = active_request_redis_batch(redis_cache) + if not self.session_affinity_repin_on_failover and batch is not None and batch.read_as_missing(cache_key): + return None + store: Final[DeploymentSessionAffinityStore] = TransportDeploymentSessionAffinityStore(redis_cache.transport()) + result: Final = await store.read( + model_group, self._hash_user_key(user_key) if user_key is not None else "unscoped", session_id + ) + if not isinstance(result, SessionPin): return TypeAdapter(object).validate_python(await self.cache.async_get_cache(key=cache_key, local_only=True)) + shared: Final = result.value + if not self.session_affinity_repin_on_failover: + if shared is not None: + await self.cache.async_set_cache(key=cache_key, value=shared, local_only=True) + return shared 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) @@ -614,16 +672,13 @@ class DeploymentAffinityCheck(CustomLogger): 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)), + await self._claim_session_pin( + model_group, + session_id, + self.get_user_key_from_request_kwargs(request_kwargs), + 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, + eligible_model_ids=eligible_deployment_ids, ) async def async_pre_call_deployment_hook( @@ -733,19 +788,18 @@ class DeploymentAffinityCheck(CustomLogger): 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, user_key=user_key - ) - claimed_session_pin: Final = await self._claim_pin( - cache_key=session_cache_key, - pin_value=pin_value, + claimed_session_pin: Final = await self._claim_session_pin( + model_group=deployment_model_name, + session_id=session_id, + user_key=user_key, + model_id=str(model_id), 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, + pin_value["model_id"], session_affinity_ttl, session_id, ) @@ -754,7 +808,7 @@ class DeploymentAffinityCheck(CustomLogger): "DeploymentAffinityCheck: session pin already claimed model_map_key=%s existing=%s ours=%s session_id=%s", deployment_model_name, claimed_session_pin, - model_id, + pin_value["model_id"], session_id, ) except Exception as e: diff --git a/tests/_support/fake_redis.py b/tests/_support/fake_redis.py index 8b98f1f5f53..1b0328ef99a 100644 --- a/tests/_support/fake_redis.py +++ b/tests/_support/fake_redis.py @@ -3,6 +3,7 @@ from typing import Final, TypeVar import fakeredis from pydantic import TypeAdapter +from redis.asyncio import Redis import litellm from litellm._service_logger import ServiceLogging @@ -93,6 +94,12 @@ def async_client(cache: RedisCache) -> fakeredis.FakeAsyncRedis: return client +def raw_client(cache: RedisCache) -> "Redis[bytes]": + client: Final = cache.init_async_client() # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # legacy RedisCache API is untyped + assert isinstance(client, Redis) + return client # pyright: ignore[reportUnknownVariableType] # redis-py erases the reply type at runtime + + def pin_async_client(cache: RedisCache) -> Callable[[], None]: client: Final = async_client(cache) diff --git a/tests/unit/redis/stores/test_deployment_session_affinity.py b/tests/unit/redis/stores/test_deployment_session_affinity.py new file mode 100644 index 00000000000..ed4a5b71a70 --- /dev/null +++ b/tests/unit/redis/stores/test_deployment_session_affinity.py @@ -0,0 +1,181 @@ +import hashlib +import json +from collections.abc import Callable, Iterator +from typing import Final +from unittest.mock import patch + +import fakeredis +import pytest +from redis.exceptions import TimeoutError as RedisTimeoutError + +from litellm.caching.dual_cache import DualCache +from litellm.router_utils.pre_call_checks.deployment_affinity_check import DeploymentAffinityCheck +from tests._support.fake_redis import ( + fake_redis_cache, + keys_matching, + local_get, + local_set, + seed_async_client, +) + + +async def _pre_call(checker: DeploymentAffinityCheck, request: dict[str, object], call_type: None = None) -> None: + await checker.async_pre_call_deployment_hook(request, call_type) # pyright: ignore[reportUnknownMemberType] # legacy cache API is untyped + + +async def _read_session( + checker: DeploymentAffinityCheck, + model_group: str, + session_id: str, + user_key: str | None, + ttl_seconds: int, + parent_otel_span: None, +) -> object: + return await checker._read_session_pin(model_group, session_id, user_key, ttl_seconds, parent_otel_span) # pyright: ignore[reportPrivateUsage] # exercises the pin path without a full router request + + +@pytest.fixture(autouse=True) +def fixed_clock() -> Iterator[None]: + with patch("time.time", return_value=1700000000.0): + yield + + +def _checker(server: fakeredis.FakeServer, *, repin: bool = True) -> DeploymentAffinityCheck: + return DeploymentAffinityCheck( + DualCache(redis_cache=fake_redis_cache(server, namespace="session")), + 60, + False, + False, + enable_session_id_affinity=True, + session_affinity_repin_on_failover=repin, + ) + + +def _key(caller: str = "caller") -> str: + return "session:deployment_affinity:v1:session:group:" + hashlib.sha256(caller.encode()).hexdigest() + ":same" + + +def _request(model_id: str, caller: str = "caller") -> dict[str, object]: + return { + "model_info": {"id": model_id}, + "metadata": { + "deployment_model_name": "group", + "model_group": "group", + "session_id": "same", + "user_api_key_hash": caller, + }, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("repin", [False, True]) +@pytest.mark.parametrize("wire", [b'"chosen"', b'{"model_id": "chosen"}', b"{'model_id': 'chosen'}"]) +async def test_session_reads_legacy_values_without_refresh(wire: bytes, repin: bool) -> None: + server: Final = fakeredis.FakeServer() + checker: Final = _checker(server, repin=repin) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(_key(), wire, ex=20) + assert await _read_session(checker, "group", "same", "caller", 60, None) in ("chosen", {"model_id": "chosen"}) + assert await raw.get(_key()) == wire + assert await raw.pttl(_key()) == 20000 + + +@pytest.mark.asyncio +async def test_session_claim_bytes_scope_and_claim_time_ttl() -> None: + server: Final = fakeredis.FakeServer() + checker: Final = _checker(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await _pre_call(checker, _request("部署"), None) + assert await raw.get(_key()) == b'{"model_id": "\\u90e8\\u7f72"}' + assert await raw.pttl(_key()) == 60000 + await raw.pexpire(_key(), 20000) + await _pre_call(checker, _request("other"), None) + assert await raw.pttl(_key()) == 20000 + await _pre_call(checker, _request("部署"), None) + assert await raw.pttl(_key()) == 60000 + await _pre_call(checker, _request("other", "second"), None) + assert await raw.get(_key("second")) == b'{"model_id": "other"}' + assert await keys_matching(raw, "session:session:*") == () + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "wire,expected", + [ + (b'{"model_id": "chosen"}', b'{"model_id": "chosen"}'), + (b'"chosen"', b'"chosen"'), + (b"chosen", b"chosen"), + (b'{"model_id": "removed"}', b'{"model_id": "replacement"}'), + (b'{"model_id": "chosen", "extra": true}', b'{"model_id": "replacement"}'), + (b"invalid{", b'{"model_id": "replacement"}'), + ], +) +async def test_success_claim_preserves_only_eligible_complete_values(wire: bytes, expected: bytes) -> None: + server: Final = fakeredis.FakeServer() + checker: Final = _checker(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(_key(), wire, ex=20) + await checker.async_repin_session_after_success( + _request("replacement"), "group", frozenset(("chosen", "replacement")) + ) + assert await raw.get(_key()) == expected + assert await raw.pttl(_key()) == 60000 + + +@pytest.mark.asyncio +async def test_shared_miss_clears_local_while_outage_retains_it() -> None: + server: Final = fakeredis.FakeServer() + checker: Final = _checker(server) + local_key: Final = _key().removeprefix("session:") + local_set(checker.cache, local_key, {"model_id": "stale"}, ttl=60) + assert await _read_session(checker, "group", "same", "caller", 60, None) is None + assert local_get(checker.cache, local_key) is None + local_set(checker.cache, local_key, {"model_id": "local"}, ttl=60) + server.connected = False + assert await _read_session(checker, "group", "same", "caller", 60, None) == {"model_id": "local"} + + +class ConcurrentRead(fakeredis.FakeAsyncRedis): + after_read: Callable[[], None] + + async def get(self, name: str) -> object: # pyright: ignore[reportIncompatibleMethodOverride] # narrows fakeredis for the test + result = await super().get(name) # pyright: ignore[reportUnknownMemberType, reportUnknownVariableType] # fakeredis is untyped + self.after_read() + return result # pyright: ignore[reportUnknownVariableType] # fakeredis is untyped + + +@pytest.mark.asyncio +async def test_shared_read_does_not_overwrite_a_newer_local_claim() -> None: + server: Final = fakeredis.FakeServer() + checker: Final = _checker(server) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + await raw.set(_key(), json.dumps({"model_id": "shared"}), ex=20) + local_key: Final = _key().removeprefix("session:") + client: Final = ConcurrentRead(server=server) + client.after_read = lambda: local_set(checker.cache, local_key, {"model_id": "newer"}, ttl=60) + redis_cache: Final = checker.cache.redis_cache + assert redis_cache is not None + seed_async_client(redis_cache, client) + assert await _read_session(checker, "group", "same", "caller", 60, None) == {"model_id": "shared"} + assert local_get(checker.cache, local_key) == {"model_id": "newer"} + + +class LostClaimReply(fakeredis.FakeAsyncRedis): + async def eval(self, script: str, numkeys: int, *keys_and_args: str | bytes | int) -> object: + await super().eval(script, numkeys, *keys_and_args) # pyright: ignore[reportUnknownMemberType] # fakeredis is untyped + raise RedisTimeoutError("reply lost after apply") + + +@pytest.mark.asyncio +async def test_session_unknown_write_retains_local_continuity() -> None: + server: Final = fakeredis.FakeServer() + checker: Final = _checker(server) + redis_cache: Final = checker.cache.redis_cache + assert redis_cache is not None + seed_async_client(redis_cache, LostClaimReply(server=server)) + local_key: Final = _key().removeprefix("session:") + local_set(checker.cache, local_key, {"model_id": "local"}, ttl=60) + await _pre_call(checker, _request("shared"), None) + raw: Final = fakeredis.FakeAsyncRedis(server=server) + assert await raw.get(_key()) == b'{"model_id": "shared"}' + assert local_get(checker.cache, local_key) == {"model_id": "local"} 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 6b4ea74894a..d317f0d5741 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,8 +1,8 @@ import asyncio import json -from typing import Final from collections.abc import AsyncIterator, Mapping from datetime import datetime, timezone +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -19,6 +19,7 @@ from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESS from litellm.router_utils.pre_call_checks.deployment_affinity_check import ( DeploymentAffinityCheck, ) +from tests._support.fake_redis import async_client, fake_redis_cache, raw_client, seed_async_client def _repinning_router(enabled: bool | None = True) -> litellm.Router: @@ -121,6 +122,17 @@ async def test_unsuccessful_responses_do_not_repin(status: str, wrapped: bool) - router.discard() +async def _claim_session( + callback: DeploymentAffinityCheck, + model_group: str, + session_id: str, + user_key: str | None, + model_id: str, + ttl_seconds: int, +) -> object: + return await callback._claim_session_pin(model_group, session_id, user_key, model_id, ttl_seconds) # pyright: ignore[reportPrivateUsage] # exercises the pin path without a full router request + + async def test_concurrent_successful_replacements_keep_first_winner() -> None: router: Final = _repinning_router() try: @@ -288,8 +300,8 @@ async def test_delayed_shared_pin_read_preserves_a_newer_local_write( 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) + reader: Final = DelayedReader(connection_pool=raw_client(redis).connection_pool) + seed_async_client(redis, reader) pending: Final = asyncio.create_task(_selected_session_ids(router)) try: await observed.wait() @@ -817,12 +829,12 @@ async def test_in_memory_reclaim_slides_idle_window_only_for_the_stored_deployme 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) + reclaimed = await _claim_session(callback, "target-group", "slide-session", "key-1", "deployment-1", 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) + lost = await _claim_session(callback, "target-group", "slide-session", "key-1", "deployment-2", 10) assert lost == "deployment-1" assert cache.in_memory_cache.ttl_dict[pin_key] == pytest.approx(time_module.time() + 777, abs=5) @@ -840,32 +852,21 @@ async def test_claim_pin_uses_redis_attached_after_construction(): 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) + server = FakeServer() + late_redis = fake_redis_cache(server, namespace="late") cache.redis_cache = late_redis import time as time_module pin_key = _session_pin_key("late-redis-session", "key-1") + await late_redis.async_set_cache(pin_key, {"model_id": "other-pod-winner"}, ttl=10) 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, - ) + claimed = await _claim_session(callback, "target-group", "late-redis-session", "key-1", "our-deployment", 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 await late_redis.async_get_cache(pin_key) == {"model_id": "other-pod-winner"} assert cache.in_memory_cache.get_cache(_session_pin_key("late-redis-session", "key-1")) == { "model_id": "other-pod-winner" } @@ -885,20 +886,18 @@ async def test_claim_pin_falls_back_to_pod_local_when_redis_is_down(): 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) + server = FakeServer() + server.connected = False + down_redis = fake_redis_cache(server, namespace="down") 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) + claimed = await _claim_session(callback, "target-group", "outage-session", "key-1", "our-deployment", 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) + second = await _claim_session(callback, "target-group", "outage-session", "key-1", "another-deployment", 777) assert second == "our-deployment" @@ -979,14 +978,15 @@ async def test_legacy_deployment_claim_retains_decoder_and_keepalive( enable_user_key_affinity=False, enable_responses_api_affinity=False, ) - cache.in_memory_cache.set_cache("deployment-pin", stored, ttl=10) + pin_key: Final = callback.get_session_affinity_cache_key("target-group", "decoder", None) + cache.in_memory_cache.set_cache(pin_key, stored, ttl=10) clock.return_value = 105.0 - winner: Final = await callback._claim_pin("deployment-pin", {"model_id": "7"}, 30) + winner: Final = await _claim_session(callback, "target-group", "decoder", None, "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) + assert cache.in_memory_cache.ttl_dict[pin_key] == (135.0 if refresh else 110.0) + assert cache.in_memory_cache.get_cache(pin_key) == ({"model_id": "7"} if refresh else stored) @pytest.mark.asyncio @@ -998,16 +998,13 @@ async def test_legacy_deployment_claim_retains_decoder_and_keepalive( ("winner", "winner", "winner"), (b"winner", "winner", "winner"), ('{"model": "winner"}', None, {"model": "winner"}), - (None, "candidate", None), - (123, "candidate", None), - ({"model_id": "winner"}, "candidate", None), + (None, "candidate", {"model_id": "candidate"}), ], ) async def test_redis_deployment_claim_preserves_legacy_result_decoding( - raw: object, expected: str | None, stored: object + raw: str | bytes | None, expected: str | None, stored: object ) -> None: - redis: Final = MagicMock() - redis.async_register_script.return_value = AsyncMock(return_value=raw) + redis: Final = fake_redis_cache(FakeServer(), namespace="decode") cache: Final = DualCache(redis_cache=redis) callback: Final = DeploymentAffinityCheck( cache=cache, @@ -1016,10 +1013,13 @@ async def test_redis_deployment_claim_preserves_legacy_result_decoding( enable_responses_api_affinity=False, ) - winner: Final = await callback._claim_pin("deployment-pin", {"model_id": "candidate"}, 30) + pin_key: Final = callback.get_session_affinity_cache_key("target-group", "decoder", None) + if raw is not None: + await async_client(redis).set(redis.check_and_fix_namespace(pin_key), raw, ex=10) + winner: Final = await _claim_session(callback, "target-group", "decoder", None, "candidate", 30) assert winner == expected - assert cache.in_memory_cache.get_cache("deployment-pin") == stored + assert cache.in_memory_cache.get_cache(pin_key) == stored @pytest.mark.asyncio