diff --git a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py index 1ca29eee89d..5e3ff8eb7f8 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -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" diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index bdaa69cd8fe..28c406edf76 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -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. diff --git a/litellm/router.py b/litellm/router.py index 487d6a31226..d3caa069b28 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -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) diff --git a/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py b/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py index 8eed696e77d..3240ad20edb 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py +++ b/tests/test_litellm/proxy/management_endpoints/test_access_group_management.py @@ -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() diff --git a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py index b14c9d1e490..3bd83ade20a 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_model_management_endpoints.py @@ -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 diff --git a/tests/test_litellm/test_model_block_unblock.py b/tests/test_litellm/test_model_block_unblock.py index 318b2f519c4..dc0098e405e 100644 --- a/tests/test_litellm/test_model_block_unblock.py +++ b/tests/test_litellm/test_model_block_unblock.py @@ -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", + ) diff --git a/tests/test_litellm/test_router.py b/tests/test_litellm/test_router.py index ad4e430c603..6db19d63a01 100644 --- a/tests/test_litellm/test_router.py +++ b/tests/test_litellm/test_router.py @@ -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"]})