From f872b7db2cb45396f95b40c8b7c719a204552166 Mon Sep 17 00:00:00 2001 From: Akshat Shah Date: Wed, 7 Oct 2026 16:29:55 -0700 Subject: [PATCH 1/5] fix(router): re-pin stale session pin when pinned deployment is unhealthy A session pin naming a deployment that has since gone on cooldown was never refreshed. First-writer-wins kept the dead row for the rest of the TTL, so every later call in the session re-shuffled across the remaining deployments and the provider prompt cache was rebuilt from scratch each time. The filter now drops a session pin whose deployment is not in the healthy set, letting the post-call hook re-claim the pin for the deployment actually chosen. The claim stays atomic, so concurrent requests still converge on one winner. Fixes #45145 --- .../deployment_affinity_check.py | 21 ++++++ .../test_session_id_affinity.py | 73 +++++++++++++++++++ 2 files changed, 94 insertions(+) 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 432fe11dc2e..d302c391cb0 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -330,6 +330,21 @@ class DeploymentAffinityCheck(CustomLogger): def _set_local_pin(self, cache_key: str, value: object, ttl_seconds: int) -> None: set_local_affinity_pin(self.cache, cache_key, value, ttl_seconds) + async def _drop_stale_pin(self, cache_key: str) -> None: + """Best-effort invalidation of a pin whose deployment is no longer healthy. + + The post-call hook re-claims the pin for the deployment actually chosen + through an atomic first-writer-wins SET, so concurrent requests still + converge on a single winner. + """ + try: + await self.cache.async_delete_cache(cache_key) + except Exception as e: # noqa: BLE001 # affinity is best-effort; never break routing + verbose_router_logger.debug( + "DeploymentAffinityCheck: failed to drop stale session pin. error=%s", + e, + ) + async def _claim_pin(self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int) -> str | None: winner: Final = await claim_affinity_pin(self.cache, cache_key, pin_value, ttl_seconds) return self._pinned_model_id(winner) @@ -439,6 +454,12 @@ class DeploymentAffinityCheck(CustomLogger): ) return [session_deployment] else: + # The pinned deployment is not usable (e.g. on cooldown). + # Drop the stale pin so the post-call hook re-pins to the + # deployment actually chosen. Without this, first-writer-wins + # keeps the dead row for the rest of the TTL and every call + # re-shuffles (see #45145). + await self._drop_stale_pin(cache_key=session_cache_key) verbose_router_logger.debug( "DeploymentAffinityCheck: session-id pinned deployment=%s not found in healthy_deployments", session_model_id, 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..0065ddce5f3 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 @@ -771,3 +771,76 @@ def test_complexity_router_with_deployment_affinity_registers_affinity_callback( _cleanup_router_callbacks(enabled) _cleanup_router_callbacks(session_only) _cleanup_router_callbacks(disabled) + + +@pytest.mark.asyncio +async def test_stale_session_pin_is_dropped_and_reclaimed(): + """ + Regression test for https://github.com/BerriAI/litellm/issues/45145. + + A session pin naming a deployment that has since gone on cooldown must not + survive: the filter drops the stale pin, the post-call hook re-pins to the + deployment actually chosen, and subsequent calls stick to it instead of + re-shuffling on every request. + """ + cache = DualCache() + callback = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=True, + ) + + deployments = [ + { + "model_name": "group", + "litellm_params": {"model": "openai/model-a"}, + "model_info": {"id": "dep-a"}, + }, + { + "model_name": "group", + "litellm_params": {"model": "openai/model-b"}, + "model_info": {"id": "dep-b"}, + }, + { + "model_name": "group", + "litellm_params": {"model": "openai/model-c"}, + "model_info": {"id": "dep-c"}, + }, + ] + session_kwargs = {"metadata": {"session_id": "s1"}} + session_cache_key = DeploymentAffinityCheck.get_session_affinity_cache_key("group", "s1", user_key=None) + + # dep-a was pinned, then went on cooldown and left the healthy set. + await cache.async_set_cache(session_cache_key, {"model_id": "dep-a"}) + + filtered = await callback.async_filter_deployments( + model="group", + healthy_deployments=deployments[1:], + messages=[], + request_kwargs=session_kwargs, + ) + assert [d["model_info"]["id"] for d in filtered] == ["dep-b", "dep-c"] + + # The stale pin is gone, so the post-call hook can claim a fresh one. + assert await cache.async_get_cache(session_cache_key) is None + + # The router chose dep-b; the hook records the winner. + await callback.async_pre_call_deployment_hook( + kwargs={ + "model_info": {"id": "dep-b"}, + "metadata": {"session_id": "s1", "deployment_model_name": "group"}, + }, + call_type=None, + ) + assert await cache.async_get_cache(session_cache_key) == {"model_id": "dep-b"} + + # Later calls stick to dep-b instead of re-shuffling. + filtered = await callback.async_filter_deployments( + model="group", + healthy_deployments=deployments[1:], + messages=[], + request_kwargs=session_kwargs, + ) + assert [d["model_info"]["id"] for d in filtered] == ["dep-b"] From 89a040837124041c0af7f523b72d62cc9469e97a Mon Sep 17 00:00:00 2001 From: Akshat Shah Date: Wed, 7 Oct 2026 16:46:33 -0700 Subject: [PATCH 2/5] fix(router): address review feedback on stale session pin re-pinning - Compare the stored pin value before deleting it, so a worker that already re-pinned the session does not get its fresh pin erased by a lagging worker. - Preserve the pin when the request itself subsets deployments (web-search tools or team scoping): a request-only exclusion must not permanently migrate the session away from its warm deployment. - Test style: Final annotations, no rebinding, context in assertion messages. --- .../deployment_affinity_check.py | 45 ++++++- .../test_session_id_affinity.py | 120 ++++++++++++++---- 2 files changed, 133 insertions(+), 32 deletions(-) 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 d302c391cb0..a3feb6ae5de 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -30,6 +30,7 @@ from litellm.caching.dual_cache import DualCache from litellm.constants import SESSION_DEPLOYMENT_AFFINITY_TTL_METADATA_KEY, SESSION_ID_GENERATED_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger, Span from litellm.responses.utils import ResponsesAPIRequestUtils +from litellm.router_utils.common_utils import get_request_team_id from litellm.types.llms.openai import AllMessageValues from litellm.types.utils import CallTypes @@ -330,14 +331,20 @@ class DeploymentAffinityCheck(CustomLogger): def _set_local_pin(self, cache_key: str, value: object, ttl_seconds: int) -> None: set_local_affinity_pin(self.cache, cache_key, value, ttl_seconds) - async def _drop_stale_pin(self, cache_key: str) -> None: + async def _drop_stale_pin(self, cache_key: str, stale_model_id: str) -> None: """Best-effort invalidation of a pin whose deployment is no longer healthy. + Re-reads the pin and only deletes it when it still names the stale + deployment: a sibling worker may have already re-pinned the session + after our read, and deleting blindly would erase its fresh pin. The post-call hook re-claims the pin for the deployment actually chosen through an atomic first-writer-wins SET, so concurrent requests still converge on a single winner. """ try: + current: Final = await self.cache.async_get_cache(cache_key) + if self._pinned_model_id(current) != stale_model_id: + return await self.cache.async_delete_cache(cache_key) except Exception as e: # noqa: BLE001 # affinity is best-effort; never break routing verbose_router_logger.debug( @@ -345,6 +352,21 @@ class DeploymentAffinityCheck(CustomLogger): e, ) + @staticmethod + def _request_may_exclude_deployments(request_kwargs: dict) -> bool: + """True when router filtering may have removed healthy deployments for this request only. + + Mirrors the request-scoped filters in Router.async_get_available_deployment + (team-based and web-search): those exclusions must not be mistaken for the + pinned deployment going unhealthy, or one special turn would permanently + migrate the session away from its warm deployment. + """ + tools: Final = request_kwargs.get("tools") or [] + for tool in tools: + if isinstance(tool, dict) and tool.get("type") in ("web_search", "web_search_preview"): + return True + return get_request_team_id(request_kwargs) is not None + async def _claim_pin(self, cache_key: str, pin_value: DeploymentAffinityCacheValue, ttl_seconds: int) -> str | None: winner: Final = await claim_affinity_pin(self.cache, cache_key, pin_value, ttl_seconds) return self._pinned_model_id(winner) @@ -454,12 +476,21 @@ class DeploymentAffinityCheck(CustomLogger): ) return [session_deployment] else: - # The pinned deployment is not usable (e.g. on cooldown). - # Drop the stale pin so the post-call hook re-pins to the - # deployment actually chosen. Without this, first-writer-wins - # keeps the dead row for the rest of the TTL and every call - # re-shuffles (see #45145). - await self._drop_stale_pin(cache_key=session_cache_key) + # The pinned deployment is not in the healthy set. When the + # request itself subsets deployments (e.g. web-search or team + # scoping), the pin must be preserved: the deployment is fine, + # this turn just cannot use it, and deleting the pin would + # permanently migrate the session away from its warm + # deployment. Otherwise the pin is stale (e.g. the deployment + # went on cooldown): drop it so the post-call hook re-pins to + # the deployment actually chosen. Without this, + # first-writer-wins keeps the dead row for the rest of the + # TTL and every call re-shuffles (see #45145). + if not self._request_may_exclude_deployments(request_kwargs): + await self._drop_stale_pin( + cache_key=session_cache_key, + stale_model_id=session_model_id, + ) verbose_router_logger.debug( "DeploymentAffinityCheck: session-id pinned deployment=%s not found in healthy_deployments", session_model_id, 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 0065ddce5f3..c110014f181 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 @@ -775,16 +775,12 @@ def test_complexity_router_with_deployment_affinity_registers_affinity_callback( @pytest.mark.asyncio async def test_stale_session_pin_is_dropped_and_reclaimed(): - """ - Regression test for https://github.com/BerriAI/litellm/issues/45145. + """A session pin naming a deployment that went on cooldown is dropped and re-claimed. - A session pin naming a deployment that has since gone on cooldown must not - survive: the filter drops the stale pin, the post-call hook re-pins to the - deployment actually chosen, and subsequent calls stick to it instead of - re-shuffling on every request. + Regression test for https://github.com/BerriAI/litellm/issues/45145. """ - cache = DualCache() - callback = DeploymentAffinityCheck( + cache: Final = DualCache() + callback: Final = DeploymentAffinityCheck( cache=cache, ttl_seconds=600, enable_user_key_affinity=False, @@ -792,7 +788,7 @@ async def test_stale_session_pin_is_dropped_and_reclaimed(): enable_session_id_affinity=True, ) - deployments = [ + deployments: Final = ( { "model_name": "group", "litellm_params": {"model": "openai/model-a"}, @@ -808,25 +804,26 @@ async def test_stale_session_pin_is_dropped_and_reclaimed(): "litellm_params": {"model": "openai/model-c"}, "model_info": {"id": "dep-c"}, }, - ] - session_kwargs = {"metadata": {"session_id": "s1"}} - session_cache_key = DeploymentAffinityCheck.get_session_affinity_cache_key("group", "s1", user_key=None) + ) + healthy_without_dep_a: Final = [deployments[1], deployments[2]] + session_kwargs: Final = {"metadata": {"session_id": "s1"}} + session_cache_key: Final = DeploymentAffinityCheck.get_session_affinity_cache_key("group", "s1", user_key=None) - # dep-a was pinned, then went on cooldown and left the healthy set. await cache.async_set_cache(session_cache_key, {"model_id": "dep-a"}) - filtered = await callback.async_filter_deployments( + filtered_before_repin: Final = await callback.async_filter_deployments( model="group", - healthy_deployments=deployments[1:], + healthy_deployments=healthy_without_dep_a, messages=[], request_kwargs=session_kwargs, ) - assert [d["model_info"]["id"] for d in filtered] == ["dep-b", "dep-c"] + assert [d["model_info"]["id"] for d in filtered_before_repin] == ["dep-b", "dep-c"], ( + "with dep-a on cooldown the filter must fall through instead of narrowing to the stale pin" + ) + assert await cache.async_get_cache(session_cache_key) is None, ( + "the stale pin must be dropped so the post-call hook can claim a fresh one" + ) - # The stale pin is gone, so the post-call hook can claim a fresh one. - assert await cache.async_get_cache(session_cache_key) is None - - # The router chose dep-b; the hook records the winner. await callback.async_pre_call_deployment_hook( kwargs={ "model_info": {"id": "dep-b"}, @@ -834,13 +831,86 @@ async def test_stale_session_pin_is_dropped_and_reclaimed(): }, call_type=None, ) - assert await cache.async_get_cache(session_cache_key) == {"model_id": "dep-b"} + assert await cache.async_get_cache(session_cache_key) == {"model_id": "dep-b"}, ( + "the hook must re-pin the session to the deployment actually chosen" + ) - # Later calls stick to dep-b instead of re-shuffling. - filtered = await callback.async_filter_deployments( + filtered_after_repin: Final = await callback.async_filter_deployments( model="group", - healthy_deployments=deployments[1:], + healthy_deployments=healthy_without_dep_a, messages=[], request_kwargs=session_kwargs, ) - assert [d["model_info"]["id"] for d in filtered] == ["dep-b"] + assert [d["model_info"]["id"] for d in filtered_after_repin] == ["dep-b"], ( + "later calls must stick to the re-pinned deployment instead of re-shuffling" + ) + + +@pytest.mark.asyncio +async def test_stale_pin_drop_backs_off_when_sibling_worker_already_repinned(): + """Dropping a stale pin must not erase a fresh pin written by a sibling worker.""" + cache: Final = DualCache() + callback: Final = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=True, + ) + session_cache_key: Final = DeploymentAffinityCheck.get_session_affinity_cache_key("group", "s1", user_key=None) + + await cache.async_set_cache(session_cache_key, {"model_id": "dep-b"}) + await callback._drop_stale_pin(cache_key=session_cache_key, stale_model_id="dep-a") + assert await cache.async_get_cache(session_cache_key) == {"model_id": "dep-b"}, ( + "the drop must compare the stored value first and leave a re-pinned session alone" + ) + + await cache.async_set_cache(session_cache_key, {"model_id": "dep-a"}) + await callback._drop_stale_pin(cache_key=session_cache_key, stale_model_id="dep-a") + assert await cache.async_get_cache(session_cache_key) is None, ( + "the drop must still remove the pin when it names the stale deployment" + ) + + +@pytest.mark.asyncio +async def test_request_scoped_exclusion_preserves_session_pin(): + """A web-search turn that excludes the pinned deployment must not migrate the session.""" + cache: Final = DualCache() + callback: Final = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=True, + ) + deployments: Final = [ + { + "model_name": "group", + "litellm_params": {"model": "openai/model-a"}, + "model_info": {"id": "dep-a"}, + }, + { + "model_name": "group", + "litellm_params": {"model": "openai/model-b"}, + "model_info": {"id": "dep-b"}, + }, + ] + session_cache_key: Final = DeploymentAffinityCheck.get_session_affinity_cache_key("group", "s1", user_key=None) + await cache.async_set_cache(session_cache_key, {"model_id": "dep-a"}) + + web_search_kwargs: Final = { + "metadata": {"session_id": "s1"}, + "tools": [{"type": "web_search"}], + } + filtered: Final = await callback.async_filter_deployments( + model="group", + healthy_deployments=[deployments[1]], + messages=[], + request_kwargs=web_search_kwargs, + ) + assert [d["model_info"]["id"] for d in filtered] == ["dep-b"], ( + "the web-search turn still routes to a supporting deployment" + ) + assert await cache.async_get_cache(session_cache_key) == {"model_id": "dep-a"}, ( + "a request-only exclusion must not drop the pin, so ordinary turns keep the warm deployment" + ) From 6cccd47e734092b60c1d2f0b0f4a52210543b9d0 Mon Sep 17 00:00:00 2001 From: Akshat Shah Date: Wed, 7 Oct 2026 16:49:13 -0700 Subject: [PATCH 3/5] test(router): cover team-scoped exclusion and cache-failure paths of stale pin drop --- .../test_session_id_affinity.py | 63 +++++++++++++++++++ 1 file changed, 63 insertions(+) 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 c110014f181..e3277116173 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 @@ -914,3 +914,66 @@ async def test_request_scoped_exclusion_preserves_session_pin(): assert await cache.async_get_cache(session_cache_key) == {"model_id": "dep-a"}, ( "a request-only exclusion must not drop the pin, so ordinary turns keep the warm deployment" ) + + +@pytest.mark.asyncio +async def test_team_scoped_request_preserves_session_pin(): + """A team-scoped turn that excludes the pinned deployment must not migrate the session.""" + cache: Final = DualCache() + callback: Final = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=True, + ) + deployments: Final = [ + { + "model_name": "group", + "litellm_params": {"model": "openai/model-a"}, + "model_info": {"id": "dep-a"}, + }, + { + "model_name": "group", + "litellm_params": {"model": "openai/model-b"}, + "model_info": {"id": "dep-b"}, + }, + ] + session_cache_key: Final = DeploymentAffinityCheck.get_session_affinity_cache_key("group", "s1", user_key=None) + await cache.async_set_cache(session_cache_key, {"model_id": "dep-a"}) + + team_kwargs: Final = {"metadata": {"session_id": "s1", "user_api_key_team_id": "team-1"}} + filtered: Final = await callback.async_filter_deployments( + model="group", + healthy_deployments=[deployments[1]], + messages=[], + request_kwargs=team_kwargs, + ) + assert [d["model_info"]["id"] for d in filtered] == ["dep-b"], ( + "the team-scoped turn still routes to a deployment on its team" + ) + assert await cache.async_get_cache(session_cache_key) == {"model_id": "dep-a"}, ( + "a request-only exclusion must not drop the pin, so ordinary turns keep the warm deployment" + ) + + +@pytest.mark.asyncio +async def test_stale_pin_drop_never_breaks_routing_on_cache_failure(): + """A cache failure during stale-pin cleanup must not break the request.""" + cache: Final = DualCache() + callback: Final = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=True, + ) + session_cache_key: Final = DeploymentAffinityCheck.get_session_affinity_cache_key("group", "s1", user_key=None) + await cache.async_set_cache(session_cache_key, {"model_id": "dep-a"}) + + with patch.object(DualCache, "async_delete_cache", new=AsyncMock(side_effect=RuntimeError("cache down"))): + await callback._drop_stale_pin(cache_key=session_cache_key, stale_model_id="dep-a") + + assert await cache.async_get_cache(session_cache_key) == {"model_id": "dep-a"}, ( + "a failed cleanup must leave the pin untouched and never raise" + ) From 3500bd0b283111f85e10ee83cf264d3e5128fb34 Mon Sep 17 00:00:00 2001 From: Akshat Shah Date: Wed, 7 Oct 2026 18:13:54 -0700 Subject: [PATCH 4/5] fix(router): make stale session-pin delete an atomic compare-and-delete A sibling worker could re-pin the session between the read and the delete, erasing its fresh pin. The delete now runs as a single Lua script on Redis (compare model_id, then DEL) with a same-tick in-memory compare-and-delete fallback, so a fresh pin always survives. Add a regression test that replays the interleaving against a scripted Redis stand-in and asserts no Python-level read-then-delete remains. --- litellm/caching/affinity_cache.py | 56 ++++++++++++++++ .../deployment_affinity_check.py | 17 ++--- .../test_session_id_affinity.py | 64 ++++++++++++++++++- 3 files changed, 126 insertions(+), 11 deletions(-) diff --git a/litellm/caching/affinity_cache.py b/litellm/caching/affinity_cache.py index 2d27574b098..c6c3b2879a8 100644 --- a/litellm/caching/affinity_cache.py +++ b/litellm/caching/affinity_cache.py @@ -48,6 +48,26 @@ end return current """ +_DELETE_PIN_IF_MODEL_ID_SCRIPT: Final = """ +local current = redis.call('GET', KEYS[1]) +if current == false then + return 0 +end +local stored_id = current +local decoded, parsed = pcall(cjson.decode, current) +if decoded then + if type(parsed) == 'table' then + stored_id = parsed['model_id'] + else + stored_id = parsed + end +end +if tostring(stored_id) == ARGV[1] then + return redis.call('DEL', KEYS[1]) +end +return 0 +""" + def set_local_affinity_pin(cache: DualCache, cache_key: str, value: object, ttl_seconds: int) -> None: """Replace the entry because InMemoryCache.set_cache preserves a live key's expiry.""" @@ -87,6 +107,42 @@ def _decode_pin(value: str) -> object: return value +def _stored_pin_model_id(stored: object) -> str | None: + """Deployment id named by a stored pin, for both the dict shape and bare strings.""" + if isinstance(stored, dict): + model_id = 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 delete_affinity_pin_if_model_id_in_memory(cache: DualCache, cache_key: str, model_id: str) -> bool: + """Synchronous compare-and-delete; no await sits between the read and the write.""" + if _stored_pin_model_id(cache.in_memory_cache.get_cache(cache_key)) != model_id: + return False + cache.in_memory_cache.delete_cache(cache_key) + return True + + +async def delete_affinity_pin_if_model_id(cache: DualCache, cache_key: str, model_id: str) -> bool: + """Delete the pin only when it still names `model_id`; True when anything was deleted. + + The Redis compare-and-delete runs as a single Lua script, so a sibling worker + that re-pinned the session after our read keeps its fresh pin. Falls back to a + same-tick in-memory compare-and-delete when Redis is absent or fails. + """ + deleted: bool = False + redis_cache: Final = cache.redis_cache + if redis_cache is not None: + try: + delete_script: Final = redis_cache.async_register_script(_DELETE_PIN_IF_MODEL_ID_SCRIPT) + deleted = bool(await delete_script(keys=(cache_key,), args=(model_id,))) + except Exception as error: # noqa: BLE001 # Redis/Lua faults fall through to the local tier + verbose_router_logger.debug("Affinity cache: Redis pin delete failed. error=%s", error) + return delete_affinity_pin_if_model_id_in_memory(cache, cache_key, model_id) or deleted + + async def claim_affinity_pin( cache: DualCache, cache_key: str, 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 a3feb6ae5de..0099d0262a0 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -24,6 +24,7 @@ from litellm.caching.affinity_cache import ( ROUTER_SESSION_PINS_TARGET, claim_affinity_pin, claim_affinity_pin_in_memory, + delete_affinity_pin_if_model_id, set_local_affinity_pin, ) from litellm.caching.dual_cache import DualCache @@ -334,18 +335,14 @@ class DeploymentAffinityCheck(CustomLogger): async def _drop_stale_pin(self, cache_key: str, stale_model_id: str) -> None: """Best-effort invalidation of a pin whose deployment is no longer healthy. - Re-reads the pin and only deletes it when it still names the stale - deployment: a sibling worker may have already re-pinned the session - after our read, and deleting blindly would erase its fresh pin. - The post-call hook re-claims the pin for the deployment actually chosen - through an atomic first-writer-wins SET, so concurrent requests still - converge on a single winner. + The compare-and-delete is atomic (one Lua script on Redis, one synchronous + tick in memory): a sibling worker that re-pinned the session after our read + keeps its fresh pin. The post-call hook re-claims the pin for the deployment + actually chosen through an atomic first-writer-wins SET, so concurrent + requests still converge on a single winner. """ try: - current: Final = await self.cache.async_get_cache(cache_key) - if self._pinned_model_id(current) != stale_model_id: - return - await self.cache.async_delete_cache(cache_key) + await delete_affinity_pin_if_model_id(self.cache, cache_key=cache_key, model_id=stale_model_id) except Exception as e: # noqa: BLE001 # affinity is best-effort; never break routing verbose_router_logger.debug( "DeploymentAffinityCheck: failed to drop stale session pin. error=%s", 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 e3277116173..693927dc1dd 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 @@ -971,9 +971,71 @@ async def test_stale_pin_drop_never_breaks_routing_on_cache_failure(): session_cache_key: Final = DeploymentAffinityCheck.get_session_affinity_cache_key("group", "s1", user_key=None) await cache.async_set_cache(session_cache_key, {"model_id": "dep-a"}) - with patch.object(DualCache, "async_delete_cache", new=AsyncMock(side_effect=RuntimeError("cache down"))): + with patch.object(InMemoryCache, "delete_cache", side_effect=RuntimeError("cache down")): await callback._drop_stale_pin(cache_key=session_cache_key, stale_model_id="dep-a") assert await cache.async_get_cache(session_cache_key) == {"model_id": "dep-a"}, ( "a failed cleanup must leave the pin untouched and never raise" ) + + +@pytest.mark.asyncio +async def test_stale_pin_delete_is_atomic_compare_and_delete(): + """The stale-pin delete must be one atomic step, not read-then-delete. + + Regression test for review feedback on #45198: a sibling worker that re-pins + the session after our filter observed the stale pin must keep its fresh pin. + """ + redis_store: Final = {} + + async def fake_atomic_script(keys, args): + key: Final = keys[0] + expected: Final = args[0] + raw: Final = redis_store.get(key) + stored_id: Final = json.loads(raw)["model_id"] if raw is not None else None + if stored_id == expected: + del redis_store[key] + return 1 + return 0 + + registered: Final = {} + + def capture_script(source): + registered["source"] = source + return fake_atomic_script + + fake_redis: Final = MagicMock() + fake_redis.async_register_script = MagicMock(side_effect=capture_script) + cache: Final = DualCache() + cache.redis_cache = fake_redis + callback: Final = DeploymentAffinityCheck( + cache=cache, + ttl_seconds=600, + enable_user_key_affinity=False, + enable_responses_api_affinity=False, + enable_session_id_affinity=True, + ) + session_cache_key: Final = DeploymentAffinityCheck.get_session_affinity_cache_key("group", "s1", user_key=None) + + redis_store[session_cache_key] = json.dumps({"model_id": "dep-a"}) + # Sibling worker re-pins after our filter observed dep-a but before the delete runs. + redis_store[session_cache_key] = json.dumps({"model_id": "dep-b"}) + + with ( + patch.object(DualCache, "async_get_cache", new=AsyncMock()) as get_spy, + patch.object(DualCache, "async_delete_cache", new=AsyncMock()) as delete_spy, + ): + await callback._drop_stale_pin(cache_key=session_cache_key, stale_model_id="dep-a") + + assert json.loads(redis_store[session_cache_key]) == {"model_id": "dep-b"}, ( + "the atomic compare-and-delete must leave the sibling's fresh pin untouched" + ) + get_spy.assert_not_called() + delete_spy.assert_not_called() + script_source: Final = registered["source"] + assert "redis.call('GET', KEYS[1])" in script_source, ( + "the delete must run as a Lua script that compares before deleting" + ) + assert "redis.call('DEL', KEYS[1])" in script_source, ( + "the delete must run as a Lua script that compares before deleting" + ) From 98c6f786113e7c56f6f5db28832188d77b05d485 Mon Sep 17 00:00:00 2001 From: Akshat Shah Date: Wed, 7 Oct 2026 21:32:05 -0700 Subject: [PATCH 5/5] test(router): guard the Lua model-id predicate in the atomic pin-delete test The fake Redis script hardcodes the compare-and-delete contract, so the test now asserts the production Lua source keeps its model-id predicate, contains exactly one DEL, and keeps that DEL inside the predicate's conditional block. Verified via mutation: removing the predicate from the Lua turns the test red. --- .../pre_call_checks/test_session_id_affinity.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) 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 693927dc1dd..4a86ddd6f18 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 @@ -1033,9 +1033,14 @@ async def test_stale_pin_delete_is_atomic_compare_and_delete(): get_spy.assert_not_called() delete_spy.assert_not_called() script_source: Final = registered["source"] - assert "redis.call('GET', KEYS[1])" in script_source, ( - "the delete must run as a Lua script that compares before deleting" + predicate: Final = "if tostring(stored_id) == ARGV[1] then" + assert predicate in script_source, ( + "the Lua must compare the stored model id before deleting: the fake script " + "above hardcodes that contract, so only this assertion catches a Lua that " + "deletes unconditionally and would erase a sibling's fresh pin" ) - assert "redis.call('DEL', KEYS[1])" in script_source, ( - "the delete must run as a Lua script that compares before deleting" + assert script_source.count("redis.call('DEL'") == 1, ( + "the script must contain exactly one DEL, guarded by the predicate" ) + guarded_block: Final = script_source.split(predicate, 1)[1].split("end", 1)[0] + assert "redis.call('DEL', KEYS[1])" in guarded_block, "the DEL must sit inside the predicate's conditional block"