ai-workspace-services/litellm
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fix wrong keys being used for model sticky entry

Browse Source
This commit is contained in:
Emerson Gomes 2026-01-15 17:42:29 -06:00
parent 059e75ad88
commit f7726c8950
2 changed files with 34 additions and 124 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

50
litellm/router_utils/pre_call_checks/deployment_affinity_check.py
View File

@ -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,
)

108
tests/test_litellm/router_utils/pre_call_checks/test_deployment_affinity_check.py
View File

@ -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)