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 88fb8876175..62c94ead0f4 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -24,12 +24,14 @@ 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 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,6 +332,38 @@ 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, stale_model_id: str) -> None: + """Best-effort invalidation of a pin whose deployment is no longer healthy. + + 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: + 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", + 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) @@ -439,6 +473,21 @@ class DeploymentAffinityCheck(CustomLogger): ) return [session_deployment] else: + # 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 9bbaed0ae1b..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 @@ -771,3 +771,276 @@ 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(): + """A session pin naming a deployment that went on cooldown is dropped and re-claimed. + + Regression test for https://github.com/BerriAI/litellm/issues/45145. + """ + 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"}, + }, + { + "model_name": "group", + "litellm_params": {"model": "openai/model-c"}, + "model_info": {"id": "dep-c"}, + }, + ) + 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) + + await cache.async_set_cache(session_cache_key, {"model_id": "dep-a"}) + + filtered_before_repin: Final = await callback.async_filter_deployments( + model="group", + healthy_deployments=healthy_without_dep_a, + messages=[], + request_kwargs=session_kwargs, + ) + 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" + ) + + 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"}, ( + "the hook must re-pin the session to the deployment actually chosen" + ) + + filtered_after_repin: Final = await callback.async_filter_deployments( + model="group", + healthy_deployments=healthy_without_dep_a, + messages=[], + request_kwargs=session_kwargs, + ) + 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" + ) + + +@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(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"] + 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 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"