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 da9e7be828..4d75478bf9 100644 --- a/litellm/router_utils/pre_call_checks/deployment_affinity_check.py +++ b/litellm/router_utils/pre_call_checks/deployment_affinity_check.py @@ -5,7 +5,8 @@ Features (independently enable-able): 1. Responses API continuity: when a `previous_response_id` is provided, route to the deployment that generated the original response (highest priority). 2. API-key affinity: map an API key hash -> deployment id for a TTL and re-use that - deployment for subsequent requests to the same model-map key. + deployment for subsequent requests to the same router deployment model name + (alias-safe, aligns to `model_map_information.model_map_key`). This is designed to support "implicit prompt caching" scenarios (no explicit cache_control), where routing to a consistent deployment is still beneficial. @@ -106,9 +107,18 @@ class DeploymentAffinityCheck(CustomLogger): """ Derive a stable model-map key from a router deployment dict. + Primary source: `deployment.model_name` (Router's canonical group name after + alias resolution). This is stable across provider-specific deployments (e.g., + Azure/Vertex/Bedrock for the same logical model) and aligns with + `model_map_information.model_map_key` in standard logging. + Prefer `base_model` when available (important for Azure), otherwise fall back to parsing `litellm_params.model`. """ + model_name = deployment.get("model_name") + if isinstance(model_name, str) and model_name: + return model_name + model_info = deployment.get("model_info") if isinstance(model_info, dict): base_model = model_info.get("base_model") @@ -209,9 +219,7 @@ class DeploymentAffinityCheck(CustomLogger): return None @staticmethod - def _find_deployment_by_model_id( - healthy_deployments: List[dict], model_id: str - ) -> Optional[dict]: + def _find_deployment_by_model_id(healthy_deployments: List[dict], model_id: str) -> Optional[dict]: for deployment in healthy_deployments: model_info = deployment.get("model_info") if not isinstance(model_info, dict): @@ -241,9 +249,7 @@ class DeploymentAffinityCheck(CustomLogger): if self.enable_responses_api_affinity: previous_response_id = request_kwargs.get("previous_response_id") if previous_response_id is not None: - responses_model_id = ResponsesAPIRequestUtils.get_model_id_from_response_id( - str(previous_response_id) - ) + responses_model_id = ResponsesAPIRequestUtils.get_model_id_from_response_id(str(previous_response_id)) if responses_model_id is not None: deployment = self._find_deployment_by_model_id( healthy_deployments=typed_healthy_deployments, @@ -345,29 +351,25 @@ class DeploymentAffinityCheck(CustomLogger): ) return None - # Primary scope: stable model-map key (used to handle aliases that ultimately map to the same - # underlying base model). - model_map_key: Optional[str] = None - base_model = model_info.get("base_model") - if isinstance(base_model, str) and base_model: - model_map_key = base_model - else: - litellm_model_name = kwargs.get("model") - if isinstance(litellm_model_name, str) and litellm_model_name: - model_map_key = self._get_model_map_key_from_litellm_model_name( - litellm_model_name - ) + # Scope affinity by the Router deployment model name (alias-safe, consistent across + # heterogeneous providers, and matches standard logging's `model_map_key`). + deployment_model_name: Optional[str] = None + for metadata in metadata_dicts: + maybe_deployment_model_name = metadata.get("deployment_model_name") + if isinstance(maybe_deployment_model_name, str) and maybe_deployment_model_name: + deployment_model_name = maybe_deployment_model_name + break - if not model_map_key: + if not deployment_model_name: verbose_router_logger.warning( - "DeploymentAffinityCheck: model_map_key missing; skipping affinity cache update. model_id=%s", + "DeploymentAffinityCheck: deployment_model_name missing; skipping affinity cache update. model_id=%s", model_id, ) return None try: cache_key = self.get_affinity_cache_key( - model_group=model_map_key, user_key=user_key + model_group=deployment_model_name, user_key=user_key ) await self.cache.async_set_cache( cache_key, @@ -377,7 +379,7 @@ class DeploymentAffinityCheck(CustomLogger): verbose_router_logger.debug( "DeploymentAffinityCheck: set affinity mapping model_map_key=%s deployment=%s ttl=%s user_key=%s", - model_map_key, + deployment_model_name, model_id, self.ttl_seconds, self._shorten_for_logs(user_key), @@ -386,7 +388,7 @@ class DeploymentAffinityCheck(CustomLogger): # Non-blocking: affinity is a best-effort optimization. verbose_router_logger.debug( "DeploymentAffinityCheck: failed to set affinity cache. model_map_key=%s error=%s", - model_map_key, + deployment_model_name, e, ) diff --git a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py index 518c5d97e2..57d2807c8c 100644 --- a/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py +++ b/tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py @@ -1,6 +1,5 @@ import os import sys -import asyncio from unittest.mock import AsyncMock, patch import pytest @@ -44,9 +43,7 @@ async def test_async_user_key_affinity_routes_to_same_deployment(): "id": "msg_123", "status": "completed", "role": "assistant", - "content": [ - {"type": "output_text", "text": "Hello there!", "annotations": []} - ], + "content": [{"type": "output_text", "text": "Hello there!", "annotations": []}], } ], "parallel_tool_calls": True, @@ -338,12 +335,10 @@ async def test_async_previous_response_id_priority_over_user_key_affinity(): # Force user-key affinity to point to the OTHER deployment affinity_cache_key = DeploymentAffinityCheck.get_affinity_cache_key( - model_group="computer-use-preview", + model_group=model_group, user_key=user_api_key_hash, ) - await router.cache.async_set_cache( - affinity_cache_key, {"model_id": other_model_id}, ttl=3600 - ) + await router.cache.async_set_cache(affinity_cache_key, {"model_id": other_model_id}, ttl=3600) # Even though user-key affinity points elsewhere, previous_response_id should pin # to the deployment that created the original response. @@ -446,93 +441,6 @@ async def test_async_user_parameter_does_not_trigger_deployment_affinity(): assert second_response._hidden_params["model_id"] != first_model_id -@pytest.mark.asyncio -async def test_async_affinity_cache_expiry_allows_reroute(): - """ - When affinity TTL expires, routing should fall back to normal load balancing. - """ - mock_response_data = { - "id": "resp_mock-resp-ttl", - "object": "response", - "created_at": 1741476542, - "status": "completed", - "model": "azure/computer-use-preview", - "output": [ - { - "type": "message", - "id": "msg_ttl", - "status": "completed", - "role": "assistant", - "content": [{"type": "output_text", "text": "TTL Response"}], - } - ], - "parallel_tool_calls": True, - "usage": {"input_tokens": 5, "output_tokens": 5, "total_tokens": 10}, - "text": {"format": {"type": "text"}}, - "error": None, - "previous_response_id": None, - } - - router = litellm.Router( - model_list=[ - { - "model_name": "azure-ttl-test", - "litellm_params": { - "model": "azure/ttl-1", - "api_key": "mock", - "api_base": "https://mock1.openai.azure.com", - }, - }, - { - "model_name": "azure-ttl-test", - "litellm_params": { - "model": "azure/ttl-2", - "api_key": "mock", - "api_base": "https://mock2.openai.azure.com", - }, - }, - ], - optional_pre_call_checks=["deployment_affinity"], - deployment_affinity_ttl_seconds=1, - ) - - model_group = "azure-ttl-test" - user_api_key_hash = "ttl-user-key" - - choice_calls = {"count": 0} - - def deterministic_choice(seq): - choice_calls["count"] += 1 - if choice_calls["count"] == 1: - return seq[0] - return seq[1] if len(seq) > 1 else seq[0] - - with patch( - "litellm.llms.custom_httpx.http_handler.AsyncHTTPHandler.post", - new_callable=AsyncMock, - ) as mock_post, patch( - "litellm.router_strategy.simple_shuffle.random.choice", - side_effect=deterministic_choice, - ): - mock_post.return_value = MockResponse(mock_response_data, 200) - - first_response = await router.aresponses( - model=model_group, - input="Hi", - litellm_metadata={"user_api_key_hash": user_api_key_hash}, - ) - first_model_id = first_response._hidden_params["model_id"] - - await asyncio.sleep(1.1) - - second_response = await router.aresponses( - model=model_group, - input="Follow-up after ttl", - litellm_metadata={"user_api_key_hash": user_api_key_hash}, - ) - assert second_response._hidden_params["model_id"] != first_model_id - - @pytest.mark.asyncio async def test_async_pre_call_hook_uses_model_map_key_scope(): """ @@ -550,9 +458,10 @@ async def test_async_pre_call_hook_uses_model_map_key_scope(): ) kwargs = { - "model_info": {"id": "model-id-123", "base_model": "claude-sonnet-4-5@20250929"}, + "model_info": {"id": "model-id-123"}, "litellm_metadata": { "user_api_key_hash": "user-key-abc", + "deployment_model_name": "claude-sonnet-4-5@20250929", }, } @@ -594,13 +503,13 @@ async def test_async_filter_deployments_uses_stable_model_map_key_for_affinity_s healthy_deployments = [ { - "model_name": "group-any", + "model_name": stable_model_map_key, "litellm_params": {"model": f"vertex_ai/{stable_model_map_key}"}, "model_info": {"id": "deployment-1"}, }, { - "model_name": "group-any", - "litellm_params": {"model": f"vertex_ai/{stable_model_map_key}"}, + "model_name": stable_model_map_key, + "litellm_params": {"model": f"bedrock/global.anthropic.{stable_model_map_key}-v1:0"}, "model_info": {"id": "deployment-2"}, }, ] @@ -641,4 +550,3 @@ def test_cache_key_does_not_double_hash_user_api_key_hash(): user_key=user_api_key_hash, ) assert key.endswith(user_api_key_hash) ->>>>>>> 860962807e (fix(router): scope deployment affinity by model_map_key)