tiennm99/litellm
1
0
Fork 0
mirror of https://github.com/tiennm99/litellm.git synced 2026-08-05 08:23:17 +00:00
Code Issues Packages Projects Releases Wiki Activity

fix(team-routing): keep team model routing on public names

Browse Source 
Remove team model_alias rewrites and resolve team deployments by team_public_model_name with team_id so sibling deployments stay in the routing candidate pool, with explicit logs showing candidate selection before load balancing.

Made-with: Cursor
This commit is contained in:
Sameer Kankute
2026-03-27 20:11:27 +05:30
parent 5534b40ab3
commit aeb932d707
3 changed files with 171 additions and 110 deletions
Download Patch File Download Diff File Expand all files Collapse all files
litellm/proxy/management_endpoints/model_management_endpoints.py
+12 -38
View File
@@ -322,13 +322,9 @@ async def _add_team_model_to_db(
"""
If 'team_id' is provided,
- generate a deterministic 'model_name' for the model (e.g. 'model_name_{team_id}_{public_name}')
- store the model in the db with this shared group name
- store a team model alias mapping {"public_name": "model_name_{team_id}_{public_name}"}
Using a deterministic name (not UUID) ensures sibling deployments for the
same public model share a model_name, so the router treats them as a single
candidate pool for load balancing and failover.
- generate a unique 'model_name' for the model (e.g. 'model_name_{team_id}_{uuid})
- store the model in the db with the unique 'model_name'
- add the public model name to the team's allowed models list
"""
_team_id = model_params.model_info.team_id
if _team_id is None:
@@ -337,9 +333,9 @@ async def _add_team_model_to_db(
if original_model_name:
model_params.model_info.team_public_model_name = original_model_name
group_model_name = f"model_name_{_team_id}_{original_model_name}"
unique_model_name = f"model_name_{_team_id}_{uuid.uuid4()}"
model_params.model_name = group_model_name
model_params.model_name = unique_model_name
## CREATE MODEL IN DB ##
model_response = await _add_model_to_db(
@@ -348,17 +344,6 @@ async def _add_team_model_to_db(
prisma_client=prisma_client,
)
## CREATE MODEL ALIAS IN DB ##
await update_team(
data=UpdateTeamRequest(
team_id=_team_id,
model_aliases={original_model_name: group_model_name},
),
user_api_key_dict=user_api_key_dict,
http_request=Request(scope={"type": "http"}),
)
# add model to team object
await team_model_add(
data=TeamModelAddRequest(
team_id=_team_id,
@@ -457,18 +442,9 @@ async def _setup_new_team_model_assignment(
patch_data: updateDeployment,
user_api_key_dict: UserAPIKeyAuth,
) -> None:
"""Set up a new team model with deterministic name, alias, and team membership."""
group_model_name = f"model_name_{team_id}_{public_model_name}"
patch_data.model_name = group_model_name
await update_team(
data=UpdateTeamRequest(
team_id=team_id,
model_aliases={public_model_name: group_model_name},
),
user_api_key_dict=user_api_key_dict,
http_request=Request(scope={"type": "http"}),
)
"""Set up a new team model with unique name and team membership."""
unique_model_name = f"model_name_{team_id}_{uuid.uuid4()}"
patch_data.model_name = unique_model_name
await team_model_add(
data=TeamModelAddRequest(
@@ -492,18 +468,16 @@ async def _update_existing_team_model_assignment(
db_model.model_info.team_public_model_name if db_model.model_info else None
)
# Update alias only if public name changed
if old_public_name and public_model_name != old_public_name:
await update_team(
data=UpdateTeamRequest(
await team_model_add(
data=TeamModelAddRequest(
team_id=team_id,
model_aliases={public_model_name: db_model.model_name},
models=[public_model_name],
),
user_api_key_dict=user_api_key_dict,
http_request=Request(scope={"type": "http"}),
user_api_key_dict=user_api_key_dict,
)
# Keep existing unique model_name
patch_data.model_name = None
litellm/router.py
+66 -5
View File
@@ -8148,20 +8148,23 @@ class Router:
def map_team_model(self, team_model_name: str, team_id: str) -> Optional[str]:
"""
Map a team model name to a team-specific model name.
Check if team_model_name resolves to team-specific deployments.
Returns the public model name (unchanged) so the router can find all
sibling deployments via team_id filtering, instead of collapsing to a
single internal model_name.
Returns:
- deployment id: str - the deployment id of the team-specific model
- None: if no team-specific model name is found
- str: the team_model_name if team deployments exist for this team
- None: if no team-specific model is found
"""
models = self.get_model_list(model_name=team_model_name, team_id=team_id)
if not models:
return None
for model in models:
if model.get("model_info", {}).get("team_id") == team_id:
return model.get("model_name")
return team_model_name
## wildcard models
return None
def should_include_deployment(
@@ -8867,6 +8870,38 @@ class Router:
model = _model_from_alias
if model not in self.model_names:
# Check for team-specific deployments by team_public_model_name
if request_team_id is not None:
team_deployments = self._get_all_deployments(
model_name=model, team_id=request_team_id
)
if team_deployments:
candidate_details = []
for deployment in team_deployments:
deployment_info = deployment.get("model_info", {}) or {}
deployment_params = deployment.get("litellm_params", {}) or {}
candidate_details.append(
{
"model_name": deployment.get("model_name"),
"model_id": deployment_info.get("id"),
"team_public_model_name": deployment_info.get(
"team_public_model_name"
),
"api_base": deployment_params.get("api_base"),
}
)
verbose_router_logger.info(
"🔥 routing_candidates_before_lb "
f"model={model} count={len(team_deployments)} "
f"candidates={candidate_details}"
)
if len(team_deployments) > 1:
verbose_router_logger.info(
"🔥 load_balancer_candidate_pool "
f"model={model} candidate_count={len(team_deployments)}"
)
return model, team_deployments
# check if provider/ specific wildcard routing use pattern matching
pattern_deployments = self.pattern_router.get_deployments_by_pattern(
model=model,
@@ -8905,6 +8940,32 @@ class Router:
# check if the user sent in a deployment name instead
healthy_deployments = self._get_deployment_by_litellm_model(model=model)
if isinstance(healthy_deployments, list) and len(healthy_deployments) > 0:
candidate_details = []
for deployment in healthy_deployments:
deployment_info = deployment.get("model_info", {}) or {}
deployment_params = deployment.get("litellm_params", {}) or {}
candidate_details.append(
{
"model_name": deployment.get("model_name"),
"model_id": deployment_info.get("id"),
"team_public_model_name": deployment_info.get(
"team_public_model_name"
),
"api_base": deployment_params.get("api_base"),
}
)
verbose_router_logger.info(
"🔥 routing_candidates_before_lb "
f"model={model} count={len(healthy_deployments)} "
f"candidates={candidate_details}"
)
if len(healthy_deployments) > 1:
verbose_router_logger.info(
"🔥 load_balancer_candidate_pool "
f"model={model} candidate_count={len(healthy_deployments)}"
)
if verbose_router_logger.isEnabledFor(logging.DEBUG):
verbose_router_logger.debug(
f"initial list of deployments: {healthy_deployments}"
tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
+93 -67
View File
@@ -564,98 +564,124 @@ class TestUpdatePublicModelGroups:
litellm.public_model_groups_links = original_value
class TestTeamModelAliasSiblingOverwrite:
class TestTeamModelSiblingRouting:
"""
Verify that two sibling team deployments for the same public model name
produce the same deterministic internal model_name, so the alias write
is idempotent and the router groups both deployments together.
Verify that sibling team deployments (same public model name, different
api_base) are all reachable through routing — no alias overwrite, no
collapse to a single deployment.
"""
@pytest.mark.asyncio
async def test_sibling_team_models_share_deterministic_name(self):
async def test_no_model_aliases_written_for_team_models(self):
"""
_add_team_model_to_db must NOT write model_aliases (which caused
the second sibling to overwrite the first). It should only call
team_model_add to register the public name on the team's models list.
"""
from litellm.proxy.management_endpoints.model_management_endpoints import (
_add_team_model_to_db,
)
from litellm.types.router import ModelInfo
team_id = "team_alias_overwrite"
team_id = "team_no_alias"
public_name = "gpt-4.1-mini"
captured_alias_calls = []
async def mock_update_team(data, user_api_key_dict, http_request):
if data.model_aliases:
captured_alias_calls.append(dict(data.model_aliases))
mock_update_team = AsyncMock()
async def mock_add_model_to_db(model_params, user_api_key_dict, prisma_client):
return MagicMock(model_id=str(uuid.uuid4()))
async def mock_team_model_add(data, http_request, user_api_key_dict):
pass
mock_team_model_add = AsyncMock()
user = UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN)
prisma_client = MockPrismaClient(team_exists=True)
deployment_1 = Deployment(
model_name=public_name,
litellm_params=LiteLLM_Params(
model="azure/gpt-4o-mini",
api_key="key-1",
api_base="https://eastus.example.openai.azure.com",
),
model_info=ModelInfo(team_id=team_id),
)
deployment_2 = Deployment(
model_name=public_name,
litellm_params=LiteLLM_Params(
model="azure/gpt-4o-mini",
api_key="key-2",
api_base="https://westus.example.openai.azure.com",
),
model_info=ModelInfo(team_id=team_id),
)
with patch(
"litellm.proxy.management_endpoints.model_management_endpoints.update_team",
side_effect=mock_update_team,
), patch(
"litellm.proxy.management_endpoints.model_management_endpoints._add_model_to_db",
side_effect=mock_add_model_to_db,
), patch(
"litellm.proxy.management_endpoints.model_management_endpoints.team_model_add",
side_effect=mock_team_model_add,
):
await _add_team_model_to_db(
model_params=deployment_1,
user_api_key_dict=user,
prisma_client=prisma_client,
)
await _add_team_model_to_db(
model_params=deployment_2,
user_api_key_dict=user,
prisma_client=prisma_client,
for api_base in ["https://eastus.example.com", "https://westus.example.com"]:
dep = Deployment(
model_name=public_name,
litellm_params=LiteLLM_Params(
model="azure/gpt-4o-mini",
api_key="key",
api_base=api_base,
),
model_info=ModelInfo(team_id=team_id),
)
with patch(
"litellm.proxy.management_endpoints.model_management_endpoints.update_team",
mock_update_team,
), patch(
"litellm.proxy.management_endpoints.model_management_endpoints._add_model_to_db",
side_effect=mock_add_model_to_db,
), patch(
"litellm.proxy.management_endpoints.model_management_endpoints.team_model_add",
mock_team_model_add,
):
await _add_team_model_to_db(
model_params=dep,
user_api_key_dict=user,
prisma_client=prisma_client,
)
assert len(captured_alias_calls) == 2
mock_update_team.assert_not_called()
assert mock_team_model_add.call_count == 2
internal_name_1 = captured_alias_calls[0][public_name]
internal_name_2 = captured_alias_calls[1][public_name]
@pytest.mark.asyncio
async def test_router_finds_all_sibling_team_deployments(self):
"""
When two team deployments share team_public_model_name="gpt-4.1-mini",
the router's _common_checks_available_deployment must return BOTH as
healthy_deployments (not collapse to one).
"""
import litellm
expected_group_name = f"model_name_{team_id}_{public_name}"
team_id = "teamA"
public_name = "gpt-4.1-mini"
# Both sibling deployments get the same deterministic group name
assert internal_name_1 == expected_group_name
assert internal_name_2 == expected_group_name
assert internal_name_1 == internal_name_2, (
"Sibling deployments must share the same model_name so the "
"router treats them as a single candidate pool"
router = litellm.Router(
model_list=[
{
"model_name": f"model_name_{team_id}_uuid1",
"litellm_params": {
"model": "azure/gpt-4o-mini",
"api_key": "key-1",
"api_base": "https://eastus.openai.azure.com",
},
"model_info": {
"team_id": team_id,
"team_public_model_name": public_name,
},
},
{
"model_name": f"model_name_{team_id}_uuid2",
"litellm_params": {
"model": "azure/gpt-4o-mini",
"api_key": "key-2",
"api_base": "https://westus.openai.azure.com",
},
"model_info": {
"team_id": team_id,
"team_public_model_name": public_name,
},
},
],
)
# The second alias write is idempotent — same key, same value
final_aliases = {}
for alias_call in captured_alias_calls:
final_aliases.update(alias_call)
assert final_aliases == {public_name: expected_group_name}
# map_team_model should return the public name (not an internal UUID)
result = router.map_team_model(public_name, team_id)
assert result == public_name
# _common_checks_available_deployment should return both deployments
model, healthy = router._common_checks_available_deployment(
model=public_name,
request_kwargs={"metadata": {"user_api_key_team_id": team_id}},
)
assert isinstance(healthy, list)
assert len(healthy) == 2
api_bases = {d["litellm_params"]["api_base"] for d in healthy}
assert api_bases == {
"https://eastus.openai.azure.com",
"https://westus.openai.azure.com",
}
class TestTeamModelUpdate:
@@ -704,7 +730,7 @@ class TestTeamModelUpdate:
assert result.get("model_name", "").startswith("model_name_test_team_123_")
assert "team_public_model_name" in str(result.get("model_info", ""))
mock_update_team.assert_called_once()
mock_update_team.assert_not_called()
mock_team_model_add.assert_called_once()
@pytest.mark.asyncio