scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-08-28 05:25:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 36

Merge pull request #34861 from BerriAI/litellm_lit4872_surface_reload_drop

Browse source 
fix(proxy): report when a model write does not survive the post-write reload
This commit is contained in:
tin-berri 2026-07-28 20:18:02 -07:00 • committed by GitHub
commit 1a6642ee2e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
7 changed files with 635 additions and 142 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

251
litellm/proxy/management_endpoints/model_access_group_management_endpoints.py
View file

@ -6,6 +6,7 @@ Endpoints here:
"""
import json
from collections.abc import Mapping, Sequence
from typing import Any, Dict, List, Tuple
from fastapi import APIRouter, Depends, HTTPException
@ -16,6 +17,9 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
# Clear cache and reload models to pick up the access group changes
from litellm.proxy.management_endpoints.model_management_endpoints import (
live_model_ids_snapshot,
model_info_as_mapping,
reload_serving_verdict,
clear_cache,
)
from litellm.proxy.utils import PrismaClient
@ -72,11 +76,92 @@ def add_access_group_to_deployment(model_info: Dict[str, Any], access_group: str
return model_info, True
def _raise_http_if_reload_degraded_serving(
before: frozenset[str],
written_models: Sequence[tuple[str, object]],
access_group: str,
) -> None:
"""Same verdict as the model-write endpoints, expressed through this file's
HTTPException error convention, with the metadata-only obligation: these writes
change group membership, not the models themselves, so a row that was already not
serving before the reload is never blamed here; only a model this reload stopped
serving is reported."""
missing, collateral = reload_serving_verdict(before=before, written_models=written_models, written_must_serve=False)
gone = tuple(dict.fromkeys((*missing, *collateral)))
if not gone:
return
raise HTTPException(
status_code=500,
detail={
"error": (
f"Access group '{access_group}' was saved to the database, but model id(s) {list(gone)} that "
"this pod was serving are no longer live after the reload it triggered. Other pods reload on "
"their own interval. Check server logs for 'Error upserting deployment' for the cause."
)
},
)
async def _tag_deployment_with_access_group(
model_id: str,
model_info: object,
access_group: str,
prisma_client: PrismaClient,
) -> tuple[str, Mapping[str, object]] | None:
"""Write `access_group` into one deployment's model_info; returns the
(model_id, updated model_info) pair when a write happened, None when the
deployment already carried the group."""
updated_model_info, was_modified = add_access_group_to_deployment(
model_info=dict(_readable_model_info_or_raise(model_id=model_id, model_info=model_info)),
access_group=access_group,
)
if not was_modified:
return None
await ModelRepository(prisma_client).table.update(
where={"model_id": model_id},
data={"model_info": json.dumps(updated_model_info)},
)
verbose_proxy_logger.debug(f"Updated deployment {model_id} with access group: {access_group}")
return (model_id, updated_model_info)
def _readable_model_info_or_raise(model_id: str, model_info: object) -> Mapping[str, object]:
"""These helpers rewrite the model_info column wholesale, so a present-but-unreadable
value must refuse loudly rather than be silently replaced with a fresh object; an
absent value stays a legitimate empty start."""
parsed = model_info_as_mapping(model_info)
if parsed is None and model_info is not None:
raise ValueError(f"model_info for deployment {model_id} is not a readable JSON object; refusing to rewrite it")
return parsed or {}
async def _strip_access_group_from_deployment(
model_id: str,
model_info: object,
access_group: str,
prisma_client: PrismaClient,
) -> tuple[str, Mapping[str, object]] | None:
"""Remove `access_group` from one deployment's model_info; returns the
(model_id, updated model_info) pair when a write happened, None when the
deployment did not carry the group."""
updated_model_info, was_modified = remove_access_group_from_deployment(
model_info=dict(_readable_model_info_or_raise(model_id=model_id, model_info=model_info)),
access_group=access_group,
)
if not was_modified:
return None
await ModelRepository(prisma_client).table.update(
where={"model_id": model_id},
data={"model_info": json.dumps(updated_model_info)},
)
return (model_id, updated_model_info)
async def update_deployments_with_access_group(
model_names: List[str],
access_group: str,
prisma_client: PrismaClient,
) -> int:
) -> tuple[tuple[str, Mapping[str, object]], ...]:
"""
Update all deployments for the given model names to include the access group.
@ -86,20 +171,15 @@ async def update_deployments_with_access_group(
prisma_client: Database client
Returns:
int: Number of deployments updated
The (model_id, updated model_info) pair of every deployment actually written,
so callers can verify each one survived the post-write reload
"""
models_updated = 0
deployments = await ModelRepository(prisma_client).table.find_many(where={"model_name": {"in": model_names}})
verbose_proxy_logger.debug(f"Found {len(deployments)} deployments for model_names: {model_names}")
found_names = {deployment.model_name for deployment in deployments}
for model_name in model_names:
verbose_proxy_logger.debug(f"Updating deployments for model_name: {model_name}")
# Get all deployments with this model_name
deployments = await ModelRepository(prisma_client).table.find_many(where={"model_name": model_name})
verbose_proxy_logger.debug(f"Found {len(deployments)} deployments for model_name: {model_name}")
# If no deployments found, this is a config model (not in DB)
if len(deployments) == 0:
if model_name not in found_names:
raise HTTPException(
status_code=400,
detail={
@ -107,65 +187,52 @@ async def update_deployments_with_access_group(
},
)
# Update each deployment
for deployment in deployments:
model_info = deployment.model_info or {}
# Add access group using helper
updated_model_info, was_modified = add_access_group_to_deployment(
model_info=model_info,
access_group=access_group,
)
# Only update in DB if modified
if was_modified:
await ModelRepository(prisma_client).table.update(
where={"model_id": deployment.model_id},
data={"model_info": json.dumps(updated_model_info)},
)
models_updated += 1
verbose_proxy_logger.debug(
f"Updated deployment {deployment.model_id} with access group: {access_group}"
)
return models_updated
tagged = [
await _tag_deployment_with_access_group(
model_id=deployment.model_id,
model_info=deployment.model_info,
access_group=access_group,
prisma_client=prisma_client,
)
for deployment in deployments
]
return tuple(pair for pair in tagged if pair is not None)
async def update_specific_deployments_with_access_group(
model_ids: List[str],
access_group: str,
prisma_client: PrismaClient,
) -> int:
) -> tuple[tuple[str, Mapping[str, object]], ...]:
"""
Update specific deployments (by model_id) to include the access group.
Unlike update_deployments_with_access_group which tags ALL deployments sharing
a model_name, this function only tags the specific deployments identified by
their unique model_id.
their unique model_id. Returns the (model_id, updated model_info) pair of every
deployment actually written.
"""
models_updated = 0
for model_id in model_ids:
verbose_proxy_logger.debug(f"Updating specific deployment model_id: {model_id}")
deployment = await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_id})
if deployment is None:
raise HTTPException(
status_code=400,
detail={"error": f"Deployment with model_id '{model_id}' not found in Database."},
)
model_info = deployment.model_info or {}
updated_model_info, was_modified = add_access_group_to_deployment(
model_info=model_info,
verbose_proxy_logger.debug(f"Updating specific deployment model_ids: {model_ids}")
tagged = [
await _tag_deployment_with_access_group(
model_id=model_id,
model_info=(await _find_deployment_or_400(model_id=model_id, prisma_client=prisma_client)),
access_group=access_group,
prisma_client=prisma_client,
)
if was_modified:
await ModelRepository(prisma_client).table.update(
where={"model_id": model_id},
data={"model_info": json.dumps(updated_model_info)},
)
models_updated += 1
verbose_proxy_logger.debug(f"Updated deployment {model_id} with access group: {access_group}")
return models_updated
for model_id in model_ids
]
return tuple(pair for pair in tagged if pair is not None)
async def _find_deployment_or_400(model_id: str, prisma_client: PrismaClient) -> Mapping[str, object] | None:
deployment = await ModelRepository(prisma_client).table.find_unique(where={"model_id": model_id})
if deployment is None:
raise HTTPException(
status_code=400,
detail={"error": f"Deployment with model_id '{model_id}' not found in Database."},
)
return deployment.model_info
def remove_access_group_from_deployment(model_info: Dict[str, Any], access_group: str) -> Tuple[Dict[str, Any], bool]:
@ -335,20 +402,28 @@ async def create_model_group(
# Update deployments using the appropriate method
if use_model_ids:
assert data.model_ids is not None
models_updated = await update_specific_deployments_with_access_group(
updated_pairs = await update_specific_deployments_with_access_group(
model_ids=data.model_ids,
access_group=data.access_group,
prisma_client=prisma_client,
)
else:
assert data.model_names is not None
models_updated = await update_deployments_with_access_group(
updated_pairs = await update_deployments_with_access_group(
model_names=data.model_names,
access_group=data.access_group,
prisma_client=prisma_client,
)
models_updated = len(updated_pairs)
live_before_reload = live_model_ids_snapshot()
await clear_cache()
_raise_http_if_reload_degraded_serving(
before=live_before_reload,
written_models=updated_pairs,
access_group=data.access_group,
)
verbose_proxy_logger.info(
f"Successfully created access group '{data.access_group}' with {models_updated} models updated"
@ -573,38 +648,42 @@ async def update_access_group(
# Step 1: Remove access group from ALL DB deployments (skip config models)
all_deployments = await ModelRepository(prisma_client).table.find_many()
for deployment in all_deployments:
model_info = deployment.model_info or {}
updated_model_info, was_modified = remove_access_group_from_deployment(
model_info=model_info,
stripped = [
await _strip_access_group_from_deployment(
model_id=deployment.model_id,
model_info=deployment.model_info,
access_group=access_group,
prisma_client=prisma_client,
)
if was_modified:
await ModelRepository(prisma_client).table.update(
where={"model_id": deployment.model_id},
data={"model_info": json.dumps(updated_model_info)},
)
for deployment in all_deployments
]
stripped_pairs = tuple(pair for pair in stripped if pair is not None)
# Step 2: Add access group using the appropriate method
if use_model_ids:
assert data.model_ids is not None
models_updated = await update_specific_deployments_with_access_group(
updated_pairs = await update_specific_deployments_with_access_group(
model_ids=data.model_ids,
access_group=access_group,
prisma_client=prisma_client,
)
else:
assert data.model_names is not None
models_updated = await update_deployments_with_access_group(
updated_pairs = await update_deployments_with_access_group(
model_names=data.model_names,
access_group=access_group,
prisma_client=prisma_client,
)
models_updated = len(updated_pairs)
# Clear cache and reload models to pick up the access group changes
live_before_reload = live_model_ids_snapshot()
await clear_cache()
_raise_http_if_reload_degraded_serving(
before=live_before_reload,
written_models=list({**dict(stripped_pairs), **dict(updated_pairs)}.items()),
access_group=access_group,
)
verbose_proxy_logger.info(
f"Successfully updated access group '{access_group}' with {models_updated} models updated"
@ -686,25 +765,27 @@ async def delete_access_group(
try:
# Remove access group from all DB deployments (skip config models)
all_deployments = await ModelRepository(prisma_client).table.find_many()
models_updated = 0
for deployment in all_deployments:
model_info = deployment.model_info or {}
updated_model_info, was_modified = remove_access_group_from_deployment(
model_info=model_info,
removed = [
await _strip_access_group_from_deployment(
model_id=deployment.model_id,
model_info=deployment.model_info,
access_group=access_group,
prisma_client=prisma_client,
)
if was_modified:
await ModelRepository(prisma_client).table.update(
where={"model_id": deployment.model_id},
data={"model_info": json.dumps(updated_model_info)},
)
models_updated += 1
for deployment in all_deployments
]
removed_pairs = tuple(pair for pair in removed if pair is not None)
models_updated = len(removed_pairs)
# Clear cache and reload models to pick up the access group changes
live_before_reload = live_model_ids_snapshot()
await clear_cache()
_raise_http_if_reload_degraded_serving(
before=live_before_reload,
written_models=removed_pairs,
access_group=access_group,
)
verbose_proxy_logger.info(
f"Successfully deleted access group '{access_group}' from {models_updated} deployments"

173
litellm/proxy/management_endpoints/model_management_endpoints.py
View file

@ -13,6 +13,7 @@ model/{model_id}/update - PATCH endpoint for model update.
import asyncio
import datetime
import json
from collections.abc import Mapping, Sequence
from typing import Any, Dict, List, Literal, Optional, Set, Tuple, Union, cast
from fastapi import APIRouter, Depends, HTTPException, Header, Request, status
@ -273,6 +274,7 @@ async def patch_model(
)
# Clear cache and reload models (uses config setting or defaults to preserving config models for DB updates)
live_before_reload = live_model_ids_snapshot()
await clear_cache()
## CREATE AUDIT LOG ##
@ -289,6 +291,12 @@ async def patch_model(
)
)
raise_if_reload_degraded_serving(
before=live_before_reload,
written_models=[(model_id, getattr(updated_model, "model_info", None))],
action="update",
)
return updated_model
except Exception as e:
@ -371,6 +379,7 @@ async def _set_model_blocked_status(
},
)
live_before_reload = live_model_ids_snapshot()
await clear_cache()
asyncio.create_task(
@ -388,6 +397,12 @@ async def _set_model_blocked_status(
)
)
raise_if_reload_degraded_serving(
before=live_before_reload,
written_models=[(data.model_id, getattr(updated_model, "model_info", None))],
action=action,
)
return updated_model
except Exception as e:
@ -714,13 +729,8 @@ async def _get_team_deployments(
# Confirm team_id in model_info (defensive check)
result = []
for row in response:
model_info = row.model_info
if isinstance(model_info, str):
try:
model_info = json.loads(model_info)
except (TypeError, ValueError):
continue
if isinstance(model_info, dict) and model_info.get("team_id") == team_id:
model_info = model_info_as_mapping(row.model_info)
if model_info is not None and model_info.get("team_id") == team_id:
result.append(row)
return result
@ -771,13 +781,8 @@ async def _get_team_public_model_names(
deployments = await _get_team_deployments(team_id, prisma_client)
public_names: Set[str] = set()
for row in deployments:
model_info = row.model_info
if isinstance(model_info, str):
try:
model_info = json.loads(model_info)
except (TypeError, ValueError):
continue
if isinstance(model_info, dict):
model_info = model_info_as_mapping(row.model_info)
if model_info is not None:
public_name = model_info.get("team_public_model_name")
if public_name:
public_names.add(public_name)
@ -879,18 +884,11 @@ async def _update_existing_team_model_assignment(
def _get_team_public_model_name(
model_info: Optional[Union[dict, str]],
) -> Optional[str]:
if isinstance(model_info, dict):
value = model_info.get("team_public_model_name")
return value if isinstance(value, str) else None
if isinstance(model_info, str):
try:
parsed = json.loads(model_info)
except (TypeError, ValueError):
return None
if isinstance(parsed, dict):
value = parsed.get("team_public_model_name")
return value if isinstance(value, str) else None
return None
parsed = model_info_as_mapping(model_info)
if parsed is None:
return None
value = parsed.get("team_public_model_name")
return value if isinstance(value, str) else None
old_public_name = db_model.model_info.team_public_model_name if db_model.model_info else None
@ -1302,6 +1300,7 @@ async def add_new_model(
- store keys separately
"""
live_before_reload = live_model_ids_snapshot()
try:
_original_litellm_model_name = model_params.model_name
if model_params.model_info.team_id is None:
@ -1357,6 +1356,12 @@ async def add_new_model(
)
)
raise_if_reload_degraded_serving(
before=live_before_reload,
written_models=[(model_response.model_id, getattr(model_response, "model_info", None))],
action="create",
)
return model_response
except Exception as e:
@ -1477,8 +1482,8 @@ async def update_model(
)
# Clear cache and reload models (uses config setting or defaults to preserving config models for DB updates)
live_before_reload = live_model_ids_snapshot()
await clear_cache()
## CREATE AUDIT LOG ##
asyncio.create_task(
create_object_audit_log(
@ -1501,6 +1506,12 @@ async def update_model(
)
)
raise_if_reload_degraded_serving(
before=live_before_reload,
written_models=[(_model_id, getattr(model_response, "model_info", None))],
action="update",
)
return model_response
except Exception as e:
verbose_proxy_logger.exception(
@ -1704,6 +1715,114 @@ def _deduplicate_litellm_router_models(models: List[Dict]) -> List[Dict]:
return unique_models
def model_info_as_mapping(model_info: object) -> Mapping[str, object] | None:
"""A DB row's model_info column arrives as a dict or as its JSON string depending on
the query path, and every consumer needs the mapping. Single owner of that parse:
returns None when no usable mapping exists (None, an unparseable string, or JSON
that is not an object), and callers choose what None means for them."""
if isinstance(model_info, Mapping):
return model_info
if not isinstance(model_info, str):
return None
try:
parsed = json.loads(model_info)
except (TypeError, ValueError):
return None
return parsed if isinstance(parsed, Mapping) else None
def _expects_liveness_on_this_pod(model_info: object) -> bool:
from litellm.router import model_info_is_active_for_environment
try:
return model_info_is_active_for_environment(model_info=model_info_as_mapping(model_info))
except ValueError:
return True
def live_model_ids_snapshot() -> frozenset[str]:
"""The ids this pod's router is currently serving, read fresh from the module global
because a reload can rebind it. The empirical ground truth every verdict below is
computed from; an absent router serves nothing."""
from litellm.proxy.proxy_server import llm_router
if llm_router is None:
return frozenset()
return frozenset(llm_router.get_model_ids())
def reload_serving_verdict(
before: frozenset[str],
written_models: Sequence[tuple[str, object]],
written_must_serve: bool,
) -> tuple[tuple[str, ...], tuple[str, ...]]:
"""Judge a write-triggered reload by diffing the router's serving state instead of
trusting any layer of the reload stack to report its own failure.
The full cell matrix, per id:
- written, must-serve (the write's purpose is this model's serving state): live now
is fine; not live is reported unless the row is deliberately inactive for this
pod's LITELLM_ENVIRONMENT; a row whose model_info cannot be read counts as
expecting to serve, so its drop is still reported
- written, metadata-only (must_not_degrade): live before and gone now is reported;
a row that was already not serving stays silent, because its deadness predates
this write and blaming it would block unrelated metadata fixes
- not written but live before and gone now: collateral degradation of this pod
caused by the reload this request triggered (a wholesale re-add failure, or a
newly introduced conflict), always reported
Returns (written ids violating their obligation, collateral ids no longer served).
Best effort under concurrent admin writes: the snapshot spans only this request.
"""
now = live_model_ids_snapshot()
written_ids = frozenset(model_id for model_id, _ in written_models)
if written_must_serve:
missing = tuple(
model_id
for model_id, model_info in written_models
if model_id not in now and _expects_liveness_on_this_pod(model_info)
)
else:
missing = tuple(model_id for model_id, _ in written_models if model_id in before and model_id not in now)
collateral = tuple(sorted(before - now - written_ids))
return (missing, collateral)
def raise_if_reload_degraded_serving(
before: frozenset[str],
written_models: Sequence[tuple[str, object]],
action: str,
) -> None:
"""The caller-visible error this pod's model-write endpoints owe their caller when
the model they wrote is not being served after the reload they triggered. The DB
write is durable either way and every other pod reloads on its own interval; this
speaks only for the handling pod."""
missing, collateral = reload_serving_verdict(before=before, written_models=written_models, written_must_serve=True)
if not missing and not collateral:
return
missing_clause = (
f"the model id(s) {list(missing)} are not live in this pod's router after the reload and are not "
"being served by this pod."
if missing
else "the reload it triggered degraded this pod's serving state."
)
collateral_clause = (
f" Previously served model id(s) {list(collateral)} are also no longer being served by this pod."
if collateral
else ""
)
raise ProxyException(
message=(
f"Model {action} was saved to the database, but {missing_clause}{collateral_clause} "
"Other pods reload on their own interval. Check server logs for 'Error upserting deployment' or "
"'Error creating deployment' for the cause."
),
type=ProxyErrorTypes.internal_server_error,
code=status.HTTP_500_INTERNAL_SERVER_ERROR,
param=None,
)
async def clear_cache():
"""
Clear router caches and reload models.

63
litellm/router.py
View file

@ -19,6 +19,7 @@ import threading
import time
import traceback
from collections import defaultdict
from collections.abc import Mapping
from functools import lru_cache
from typing import (
TYPE_CHECKING,
@ -270,6 +271,43 @@ def _cost_value_as_float(value: Union[str, int, float, None]) -> Optional[float]
return None
def model_info_is_active_for_environment(model_info: Mapping[str, object] | None) -> bool:
"""Single owner of the environment-gating rule: a deployment whose model_info names
`supported_environments` loads only on pods whose LITELLM_ENVIRONMENT is in that list.
`Router.deployment_is_active_for_environment` delegates here, and the model-write
endpoints consult the same rule to tell a deliberately inactive model from one that
was dropped by a failed reload."""
if model_info is None:
return True
supported_environments = model_info.get("supported_environments")
if supported_environments is None:
return True
if not isinstance(supported_environments, (list, tuple)):
raise ValueError(
f"supported_environments must be a list of {VALID_LITELLM_ENVIRONMENTS}. "
f"but set as: {supported_environments} for model_info: {model_info}"
)
litellm_environment = get_secret_str(secret_name="LITELLM_ENVIRONMENT")
if litellm_environment is None:
raise ValueError("Set 'supported_environments' for model but not 'LITELLM_ENVIRONMENT' set in .env")
if litellm_environment not in VALID_LITELLM_ENVIRONMENTS:
raise ValueError(
f"LITELLM_ENVIRONMENT must be one of {VALID_LITELLM_ENVIRONMENTS}. but set as: {litellm_environment}"
)
for _env in supported_environments:
if _env not in VALID_LITELLM_ENVIRONMENTS:
raise ValueError(
f"supported_environments must be one of {VALID_LITELLM_ENVIRONMENTS}. but set as: {_env} "
f"for model_info: {model_info}"
)
if litellm_environment in supported_environments:
return True
return False
_PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT")
@ -7982,30 +8020,7 @@ class Router:
- ValueError: If LITELLM_ENVIRONMENT is not set in .env or not one of the valid values
- ValueError: If supported_environments is not set in model_info or not one of the valid values
"""
if (
deployment.model_info is None
or "supported_environments" not in deployment.model_info
or deployment.model_info["supported_environments"] is None
):
return True
litellm_environment = get_secret_str(secret_name="LITELLM_ENVIRONMENT")
if litellm_environment is None:
raise ValueError("Set 'supported_environments' for model but not 'LITELLM_ENVIRONMENT' set in .env")
if litellm_environment not in VALID_LITELLM_ENVIRONMENTS:
raise ValueError(
f"LITELLM_ENVIRONMENT must be one of {VALID_LITELLM_ENVIRONMENTS}. but set as: {litellm_environment}"
)
for _env in deployment.model_info["supported_environments"]:
if _env not in VALID_LITELLM_ENVIRONMENTS:
raise ValueError(
f"supported_environments must be one of {VALID_LITELLM_ENVIRONMENTS}. but set as: {_env} for deployment: {deployment}"
)
if litellm_environment in deployment.model_info["supported_environments"]:
return True
return False
return model_info_is_active_for_environment(model_info=deployment.model_info)
def set_model_list(self, model_list: list):
original_model_list = copy.deepcopy(model_list)

112
tests/test_litellm/proxy/management_endpoints/test_access_group_management.py
View file

@ -160,7 +160,9 @@ async def test_create_access_group_with_model_names_tags_all_deployments():
{
"model_name": "gpt-4o",
"litellm_params": {"model": "gpt-4o", "api_key": "fake-key"},
"model_info": {"id": deployment_id, "db_model": True},
}
for deployment_id in ("deploy-A", "deploy-B", "deploy-C")
]
)
@ -318,3 +320,113 @@ async def test_create_access_group_invalid_model_id_returns_400():
await create_model_group(data=request_data, user_api_key_dict=mock_user)
assert exc_info.value.status_code == 400
assert "non-existent-id" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_create_access_group_surfaces_dropped_models():
"""An access-group write whose reload does not leave the tagged models live on this
pod must report the drop through this file's HTTPException contract, not a 200."""
from fastapi import HTTPException
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
create_model_group,
)
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
NewModelGroupRequest,
)
deploy_a = MagicMock(model_id="deploy-A", model_name="gpt-4o", model_info={})
mock_prisma = MagicMock()
mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[])
mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=deploy_a)
mock_prisma.db.litellm_proxymodeltable.update = AsyncMock()
mock_user = UserAPIKeyAuth(user_id="test_admin", user_role=LitellmUserRoles.PROXY_ADMIN)
wiped_router = MagicMock()
wiped_router.get_model_ids.side_effect = [["deploy-A"], []]
with (
patch("litellm.proxy.proxy_server.llm_router", wiped_router),
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch(
"litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
new=AsyncMock(return_value=None),
),
):
with pytest.raises(HTTPException) as exc_info:
await create_model_group(
data=NewModelGroupRequest(access_group="production-models", model_ids=["deploy-A"]),
user_api_key_dict=mock_user,
)
assert exc_info.value.status_code == 500
assert "deploy-A" in str(exc_info.value.detail)
@pytest.mark.asyncio
async def test_tag_deployment_parses_string_model_info_and_refuses_corrupt():
"""The model_info column can arrive as its JSON string; tagging must parse it rather
than crash, and must refuse to rewrite a present-but-unreadable value."""
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
_tag_deployment_with_access_group,
)
mock_prisma = MagicMock()
mock_prisma.db.litellm_proxymodeltable.update = AsyncMock()
pair = await _tag_deployment_with_access_group(
model_id="deploy-str",
model_info='{"access_groups": ["existing"]}',
access_group="new-group",
prisma_client=mock_prisma,
)
assert pair is not None
assert pair[0] == "deploy-str"
assert pair[1]["access_groups"] == ["existing", "new-group"]
with pytest.raises(ValueError, match="deploy-corrupt"):
await _tag_deployment_with_access_group(
model_id="deploy-corrupt",
model_info="{not json",
access_group="new-group",
prisma_client=mock_prisma,
)
@pytest.mark.asyncio
async def test_delete_access_group_ignores_models_that_were_already_dead():
"""A metadata-only strip over a model this pod never served must not fail the write;
the model's deadness predates the request, and blaming it here would make a broken
model block every access-group fix that touches it."""
deploy_broken = MagicMock(
model_id="deploy-broken", model_name="broken-model", model_info={"access_groups": ["doomed-group"]}
)
mock_prisma = MagicMock()
mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[deploy_broken])
mock_prisma.db.litellm_proxymodeltable.update = AsyncMock()
from litellm.proxy._types import LitellmUserRoles, UserAPIKeyAuth
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
delete_access_group,
)
never_served_router = MagicMock()
never_served_router.get_model_ids.return_value = []
with (
patch("litellm.proxy.proxy_server.llm_router", never_served_router),
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch(
"litellm.proxy.management_endpoints.model_access_group_management_endpoints.clear_cache",
new=AsyncMock(return_value=None),
),
):
response = await delete_access_group(
access_group="doomed-group",
user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN),
)
assert response.models_updated == 1
mock_prisma.db.litellm_proxymodeltable.update.assert_awaited_once()

107
tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py
View file

@ -819,6 +819,7 @@ class TestUpdateModel:
)
mock_router = MagicMock()
mock_router.get_model_ids.return_value = [model_id]
admin_user = UserAPIKeyAuth(
user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN
)
@ -838,7 +839,7 @@ class TestUpdateModel:
),
patch(
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
new=AsyncMock(return_value=None),
new=AsyncMock(return_value=True),
) as mock_clear_cache,
):
await update_model(
@ -1888,6 +1889,7 @@ class TestAddAndDeleteModelLifecycle:
mock_router = MagicMock()
mock_router.delete_deployment = MagicMock()
mock_router.get_model_ids.return_value = [model_id]
_PS = "litellm.proxy.proxy_server"
_ENCRYPT = "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper"
@ -3166,7 +3168,7 @@ class TestPatchModelBlockedAuthGate:
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch("litellm.proxy.proxy_server.llm_router", MagicMock()),
patch("litellm.proxy.proxy_server.llm_router", MagicMock(**{"get_model_ids.return_value": ["m1"]})),
patch("litellm.proxy.proxy_server.store_model_in_db", True),
patch("litellm.proxy.proxy_server.premium_user", True),
patch(
@ -3212,7 +3214,7 @@ class TestPatchModelBlockedAuthGate:
with (
patch("litellm.proxy.proxy_server.prisma_client", mock_prisma),
patch("litellm.proxy.proxy_server.llm_router", MagicMock()),
patch("litellm.proxy.proxy_server.llm_router", MagicMock(**{"get_model_ids.return_value": ["m1"]})),
patch("litellm.proxy.proxy_server.store_model_in_db", True),
patch("litellm.proxy.proxy_server.premium_user", True),
patch(
@ -3221,7 +3223,7 @@ class TestPatchModelBlockedAuthGate:
),
patch(
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
new=AsyncMock(return_value=None),
new=AsyncMock(return_value=True),
),
):
result = await patch_model(
@ -3231,3 +3233,100 @@ class TestPatchModelBlockedAuthGate:
)
assert result is updated_row
mock_prisma.db.litellm_proxymodeltable.update.assert_awaited_once()
class TestWriteSurfacesReloadDrop:
"""A model-write endpoint may report success only if every row it wrote is, after the
reload it triggered, live in this pod's router or deliberately environment-inactive."""
def test_reload_serving_verdict_matrix(self, monkeypatch):
import litellm
from litellm.proxy.management_endpoints.model_management_endpoints import (
reload_serving_verdict,
)
live_router = litellm.Router(
model_list=[
{
"model_name": "gpt-4o",
"litellm_params": {"model": "gpt-4o"},
"model_info": {"id": "m-live", "db_model": True},
}
]
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", live_router)
monkeypatch.setenv("LITELLM_ENVIRONMENT", "development")
written = [
("m-live", {"id": "m-live"}),
("m-gone", {"id": "m-gone"}),
("m-env", {"id": "m-env", "supported_environments": ["production"]}),
("m-env-str", '{"id": "m-env-str", "supported_environments": ["production"]}'),
("m-env-misconfigured", {"id": "m-env-misconfigured", "supported_environments": ["bogus"]}),
("m-corrupt", "{not json"),
]
missing, collateral = reload_serving_verdict(
before=frozenset({"m-live", "m-collateral"}), written_models=written, written_must_serve=True
)
assert missing == ("m-gone", "m-env-misconfigured", "m-corrupt")
assert collateral == ("m-collateral",)
missing, collateral = reload_serving_verdict(
before=frozenset({"m-live", "m-was-live"}),
written_models=[("m-live", None), ("m-was-live", None), ("m-never-lived", None)],
written_must_serve=False,
)
assert missing == ("m-was-live",)
assert collateral == ()
def test_raise_if_reload_degraded_serving_contract(self, monkeypatch):
import litellm
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.model_management_endpoints import (
raise_if_reload_degraded_serving,
)
live_router = litellm.Router(
model_list=[
{
"model_name": "gpt-4o",
"litellm_params": {"model": "gpt-4o"},
"model_info": {"id": "m-live", "db_model": True},
}
]
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", live_router)
assert (
raise_if_reload_degraded_serving(
before=frozenset({"m-live"}), written_models=[("m-live", None)], action="update"
)
is None
)
with pytest.raises(ProxyException, match="m-gone"):
raise_if_reload_degraded_serving(
before=frozenset(), written_models=[("m-gone", None)], action="update"
)
with pytest.raises(ProxyException, match="m-collateral"):
raise_if_reload_degraded_serving(
before=frozenset({"m-live", "m-collateral"}), written_models=[("m-live", None)], action="update"
)
class TestModelInfoAsMapping:
"""The model_info column reaches consumers as a dict or as its JSON string; this is
the single owner of that parse, and None means no usable mapping."""
def test_contract(self):
from litellm.proxy.management_endpoints.model_management_endpoints import (
model_info_as_mapping,
)
assert model_info_as_mapping({"id": "m1"}) == {"id": "m1"}
assert model_info_as_mapping('{"id": "m1"}') == {"id": "m1"}
assert model_info_as_mapping(None) is None
assert model_info_as_mapping("{not json") is None
assert model_info_as_mapping('["a", "b"]') is None
assert model_info_as_mapping(42) is None

52
tests/test_litellm/test_model_block_unblock.py
View file

@ -34,9 +34,9 @@ def _setup_model_block_mocks(monkeypatch, *, updated_blocked: bool):
mock_prisma_client.db.litellm_proxymodeltable = model_table
mock_router = MagicMock()
mock_router.get_deployment.return_value = None
mock_router.get_model_ids.return_value = [model_id]
mock_clear_cache = AsyncMock(return_value=None)
mock_clear_cache = AsyncMock(return_value=True)
mock_audit_log = AsyncMock(return_value=None)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)
@ -197,3 +197,51 @@ async def test_route_request_returns_403_when_model_is_fully_blocked(monkeypatch
assert exc_info.value.status_code == 403
assert "Model is blocked" in exc_info.value.message
@pytest.mark.asyncio
async def test_model_block_surfaces_wholesale_reload_failure(monkeypatch):
"""The write endpoints owe the caller an error when the pod failed to reload at all;
the DB row is saved but this pod is not serving the change."""
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.model_management_endpoints import block_model
model_id, model_table, updated_row, mock_clear_cache, mock_audit_log = _setup_model_block_mocks(
monkeypatch, updated_blocked=True
)
wiped_router = MagicMock()
wiped_router.get_model_ids.side_effect = [[model_id], []]
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", wiped_router)
with pytest.raises(ProxyException, match=model_id):
await block_model(
data=BlockModelRequest(model_id=model_id),
http_request=MagicMock(),
user_api_key_dict=_proxy_admin(),
litellm_changed_by="operator@example.com",
)
assert mock_audit_log.call_args.kwargs["object_id"] == model_id
@pytest.mark.asyncio
async def test_model_block_surfaces_model_dropped_by_reload(monkeypatch):
"""A reload that completes but drops the written model (ignore_invalid_deployments
swallowed its re-add) must not produce an unqualified success."""
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.model_management_endpoints import block_model
model_id, model_table, updated_row, mock_clear_cache, _ = _setup_model_block_mocks(
monkeypatch, updated_blocked=True
)
dropped_router = MagicMock()
dropped_router.get_model_ids.return_value = []
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", dropped_router)
with pytest.raises(ProxyException, match=model_id):
await block_model(
data=BlockModelRequest(model_id=model_id),
http_request=MagicMock(),
user_api_key_dict=_proxy_admin(),
litellm_changed_by="operator@example.com",
)

19
tests/test_litellm/test_router.py
View file

@ -6379,3 +6379,22 @@ class TestPreRoutingStrategyRegistryLifecycle:
litellm_params=LiteLLM_Params(**params)
)
assert actual is expected, params["model"]
def test_model_info_is_active_for_environment_matrix(monkeypatch):
"""The model-write endpoints consult this predicate to tell a deliberately
environment-inactive model from one dropped by a failed reload; the Router's own
deployment gate delegates to it, so the two can never diverge."""
from litellm.router import model_info_is_active_for_environment
assert model_info_is_active_for_environment(model_info=None) is True
assert model_info_is_active_for_environment(model_info={"id": "m1"}) is True
assert model_info_is_active_for_environment(model_info={"supported_environments": None}) is True
monkeypatch.setenv("LITELLM_ENVIRONMENT", "development")
assert model_info_is_active_for_environment(model_info={"supported_environments": ["development"]}) is True
assert model_info_is_active_for_environment(model_info={"supported_environments": ["production"]}) is False
monkeypatch.delenv("LITELLM_ENVIRONMENT")
with pytest.raises(ValueError, match="LITELLM_ENVIRONMENT"):
model_info_is_active_for_environment(model_info={"supported_environments": ["production"]})