From 90d813a4afd15ec0eb4350a04d40961b65e51dc3 Mon Sep 17 00:00:00 2001 From: Tin Chi Lo Date: Fri, 31 Jul 2026 13:04:30 -0700 Subject: [PATCH] fix(proxy): converge sibling pods on model CRUD via shared model list cache Multi-pod deploys with STORE_MODEL_IN_DB=True experienced stale model listings because each pod cached its in-process llm_router state independently. A deletion on pod A remained visible on pods B/C for up to 30s (the add_deployment_job interval) since sibling pods never invalidated their local state The fix is a DualCache-backed shared read for the model table, following the litellm_config_cache pattern, with write-through refresh on every model CRUD write site (new, delete, patch, block/unblock, access group membership, tag writes, master-key rotation) so sibling pods pick up the change instead of reading a stale in-process router. A 2-second fingerprint-gated sync job on every pod reads the shared cache and only rebuilds the router when the fingerprint actually changed, with the existing 30-second add_deployment_job poll as a fallback The write path refreshes the cache with fresh DB rows immediately after every write rather than only deleting the entry. A delete-only invalidate leaves a window where a concurrent in-flight read on another pod can repopulate the cache with pre-write rows, silently resurrecting the exact staleness this PR is meant to fix. Refreshing with the post-write snapshot closes that race Test coverage spans cache sharing across pods, invalidation and refresh behavior including DB failure and empty-model-list edge cases, fingerprint change detection, and router reconciliation preserving config-defined models Before this change a deleted model reappears on sibling pods for up to 30 seconds. After, all pods converge within about 2 seconds --- litellm/constants.py | 1 + .../key_management_endpoints.py | 2 + ...model_access_group_management_endpoints.py | 3 + .../model_management_endpoints.py | 8 + .../tag_management_endpoints.py | 2 + litellm/proxy/model_list_cache.py | 137 ++++++++++ litellm/proxy/proxy_server.py | 70 +++++- .../test_model_management_endpoints.py | 83 +++++++ .../proxy/proxy_server/test_proxy_config.py | 143 +++++++++++ .../proxy/test_model_list_cache.py | 235 ++++++++++++++++++ 10 files changed, 676 insertions(+), 8 deletions(-) create mode 100644 litellm/proxy/model_list_cache.py create mode 100644 tests/test_litellm/proxy/test_model_list_cache.py diff --git a/litellm/constants.py b/litellm/constants.py index 78bfc6501e8..7e2b0f1296e 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1478,6 +1478,7 @@ PROXY_BATCH_POLLING_ENABLED = _batch_polling_env == "true" PROXY_BUDGET_RESCHEDULER_MAX_TIME = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_TIME", 605)) PROXY_BATCH_WRITE_AT = int(os.getenv("PROXY_BATCH_WRITE_AT", 10)) # in seconds, increased from 10 PROXY_CONFIG_RELOAD_INTERVAL_SECONDS = get_env_int("PROXY_CONFIG_RELOAD_INTERVAL_SECONDS", 30) +MODEL_LIST_SYNC_INTERVAL_SECONDS = 2 # APScheduler Configuration - MEMORY LEAK FIX # These settings prevent memory leaks in APScheduler's normalize() and _apply_jitter() functions diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index a94a75fdfa3..3196e5fa927 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -90,6 +90,7 @@ from litellm.proxy.management_helpers.team_member_permission_checks import ( TeamMemberPermissionChecks, ) from litellm.proxy.management_helpers.utils import management_endpoint_wrapper +from litellm.proxy.model_list_cache import refresh_model_list_cache from litellm.proxy.spend_tracking.spend_tracking_utils import _is_master_key from litellm.proxy.ui_crud_endpoints.proxy_setting_endpoints import ( get_ui_settings_cached, @@ -4240,6 +4241,7 @@ async def _rotate_master_key( await tx.litellm_proxymodeltable.create_many( data=new_models, ) + await refresh_model_list_cache(prisma_client=prisma_client) # 3. process config table try: config = await ConfigRepository(prisma_client).table.find_many() 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 5e3ff8eb7f8..666f6790b1f 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -22,6 +22,7 @@ from litellm.proxy.management_endpoints.model_management_endpoints import ( reload_serving_verdict, clear_cache, ) +from litellm.proxy.model_list_cache import refresh_model_list_cache from litellm.proxy.utils import PrismaClient from litellm.repositories.model_repository import ModelRepository from litellm.types.proxy.management_endpoints.model_management_endpoints import ( @@ -121,6 +122,7 @@ async def _tag_deployment_with_access_group( where={"model_id": model_id}, data={"model_info": json.dumps(updated_model_info)}, ) + await refresh_model_list_cache(prisma_client=prisma_client) verbose_proxy_logger.debug(f"Updated deployment {model_id} with access group: {access_group}") return (model_id, updated_model_info) @@ -154,6 +156,7 @@ async def _strip_access_group_from_deployment( where={"model_id": model_id}, data={"model_info": json.dumps(updated_model_info)}, ) + await refresh_model_list_cache(prisma_client=prisma_client) return (model_id, updated_model_info) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 41904e3883b..4da2bb0b7f5 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -49,6 +49,7 @@ from litellm.proxy.management_endpoints.team_endpoints import ( update_team as _legacy_update_team, ) from litellm.proxy.management_helpers.audit_logs import create_object_audit_log +from litellm.proxy.model_list_cache import refresh_model_list_cache from litellm.proxy.utils import PrismaClient from litellm.repositories.model_repository import ModelRepository from litellm.repositories.table_repositories import ModelTableRepository @@ -811,6 +812,8 @@ async def delete_team_models( await tx.litellm_proxymodeltable.delete_many(where={"model_id": {"in": model_ids}}) deleted_model_ids.extend(model_ids) + await refresh_model_list_cache(prisma_client=prisma_client) + if llm_router is not None: for model_id in deleted_model_ids: llm_router.delete_deployment(id=model_id) @@ -1198,6 +1201,8 @@ async def delete_model( llm_router=llm_router, ) + await refresh_model_list_cache(prisma_client=prisma_client) + ## CREATE AUDIT LOG ## asyncio.create_task( create_object_audit_log( @@ -1372,6 +1377,7 @@ async def add_new_model( user_api_key_dict=user_api_key_dict, prisma_client=prisma_client, ) + await refresh_model_list_cache(prisma_client=prisma_client) still_desired_ids = await proxy_config.add_deployment( prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj ) @@ -1916,6 +1922,8 @@ async def clear_cache() -> frozenset[str] | None: verbose_proxy_logger, ) + await refresh_model_list_cache(prisma_client=prisma_client) + if llm_router is None or prisma_client is None: verbose_proxy_logger.debug("llm_router or prisma_client is None, skipping cache clear") return None diff --git a/litellm/proxy/management_endpoints/tag_management_endpoints.py b/litellm/proxy/management_endpoints/tag_management_endpoints.py index ac53f254981..afe284e9021 100644 --- a/litellm/proxy/management_endpoints/tag_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tag_management_endpoints.py @@ -31,6 +31,7 @@ from litellm.proxy.management_endpoints.common_daily_activity import ( get_daily_activity, ) from litellm.proxy.management_helpers.utils import handle_budget_for_entity +from litellm.proxy.model_list_cache import refresh_model_list_cache from litellm.repositories.model_repository import ModelRepository from litellm.repositories.table_repositories import ( DailyTagSpendRepository, @@ -371,6 +372,7 @@ async def _add_tag_to_deployment(deployment: "Deployment", tag: str): where={"model_id": deployment.model_info.id}, data={"litellm_params": json.dumps(existing_params)}, ) + await refresh_model_list_cache(prisma_client=prisma_client) except Exception as e: verbose_proxy_logger.exception(f"Error adding tag to deployment: {str(e)}") raise HTTPException(status_code=500, detail=str(e)) diff --git a/litellm/proxy/model_list_cache.py b/litellm/proxy/model_list_cache.py new file mode 100644 index 00000000000..d43657a777e --- /dev/null +++ b/litellm/proxy/model_list_cache.py @@ -0,0 +1,137 @@ +"""Cluster-shared cache of the `LiteLLM_ProxyModelTable` rows. + +Every pod serves `/model/info` and routes from its own in-process `llm_router`, so model +writes are invisible to sibling pods until they re-read the table. This puts that read +behind a `DualCache` (in-memory + Redis, the same shape as `litellm_config_cache`) and +exposes the write-through refresh every model write path calls. + +Cached entries are the raw DB payload: `litellm_params` values stay encrypted exactly as +stored, and each read validates a fresh copy so callers that decrypt in place cannot +corrupt the cache. The Redis TTL is the config reload interval, so a write path that +forgets to refresh is never staler than the DB poll it replaced. + +Two paths write the shared key, so they are ordered: a write path publishes a snapshot it +just read from the DB and overwrites unconditionally, while a read-fill only holds whatever +the table had when its query began and writes with `NX`, so a read that raced a write can +never clobber the fresher rows the write published. +""" + +import hashlib +from typing import TYPE_CHECKING, Final, Mapping, Sequence + +from pydantic import BaseModel, TypeAdapter, ValidationError + +if TYPE_CHECKING: + from litellm.proxy.utils import PrismaClient + +from litellm._logging import verbose_proxy_logger +from litellm.caching.dual_cache import DualCache +from litellm.constants import PROXY_CONFIG_RELOAD_INTERVAL_SECONDS +from litellm.models.model import LiteLLM_ProxyModelTable + +MODEL_LIST_CACHE_KEY: Final[str] = "litellm_proxy:model_list" +MODEL_LIST_CACHE_IN_MEMORY_TTL_SECONDS: Final[float] = 1.0 +MODEL_LIST_CACHE_REDIS_TTL_SECONDS: Final[float] = float(PROXY_CONFIG_RELOAD_INTERVAL_SECONDS) + +model_list_cache: DualCache = DualCache( + default_in_memory_ttl=MODEL_LIST_CACHE_IN_MEMORY_TTL_SECONDS, + default_redis_ttl=MODEL_LIST_CACHE_REDIS_TTL_SECONDS, +) + +_MODEL_ROWS_ADAPTER: Final[TypeAdapter[tuple[LiteLLM_ProxyModelTable, ...]]] = TypeAdapter( + tuple[LiteLLM_ProxyModelTable, ...] +) + + +def _row_as_mapping(row: object) -> Mapping[str, object]: + if isinstance(row, BaseModel): + return row.model_dump() + if isinstance(row, Mapping): + return row + raise TypeError(f"Unsupported model row type: {type(row).__name__}") + + +def parse_model_rows(rows: Sequence[object]) -> tuple[LiteLLM_ProxyModelTable, ...]: + return _MODEL_ROWS_ADAPTER.validate_python(tuple(_row_as_mapping(row) for row in rows)) + + +def model_rows_fingerprint(models: Sequence[LiteLLM_ProxyModelTable]) -> str: + """Identity of a model list for change detection; `updated_at` moves on every row write.""" + parts = sorted(f"{model.model_id}@{model.updated_at.isoformat() if model.updated_at else ''}" for model in models) + return hashlib.sha256("|".join(parts).encode()).hexdigest() + + +async def get_cached_model_rows(cache: DualCache = model_list_cache) -> tuple[LiteLLM_ProxyModelTable, ...] | None: + """Cached model rows, or None when the cache is empty or unreadable.""" + cached = await cache.async_get_cache(MODEL_LIST_CACHE_KEY) + if cached is None: + return None + try: + return _MODEL_ROWS_ADAPTER.validate_python(cached) + except ValidationError as e: + verbose_proxy_logger.warning("model_list_cache: discarding unreadable cache entry - %s", str(e)) + return None + + +async def set_cached_model_rows( + models: Sequence[LiteLLM_ProxyModelTable], + cache: DualCache = model_list_cache, + overwrite: bool = True, +) -> None: + """Populate both tiers, each with its own TTL. + + `DualCache.async_set_cache` writes one TTL to both tiers, which cannot express what this + cache needs: the in-memory copy has to lapse fast enough for a sibling pod's write to be + seen, while the shared copy has to outlive it or every pod re-reads the DB on every sync. + + `overwrite` orders the two writers of the shared key. A write path knows its snapshot is + fresh from the DB, so it overwrites unconditionally everywhere and wins. A read-fill only + has whatever the table held when its `find_many` began, which a concurrent write may have + superseded; it skips the local copy (so a stale read cannot replace the in-memory cache) + and writes the shared copy with `NX` so it never clobbers a published write. Both the 1s + local TTL and the `NX` gate ensure the pod sees the fresh snapshot on its next sync. + """ + payload = tuple(model.model_dump(mode="json") for model in models) + if overwrite: + await cache.in_memory_cache.async_set_cache( + MODEL_LIST_CACHE_KEY, + payload, + ttl=MODEL_LIST_CACHE_IN_MEMORY_TTL_SECONDS, + ) + if cache.redis_cache is not None: + await cache.redis_cache.async_set_cache( + MODEL_LIST_CACHE_KEY, + payload, + ttl=MODEL_LIST_CACHE_REDIS_TTL_SECONDS, + nx=not overwrite, + ) + + +async def invalidate_model_list_cache(cache: DualCache = model_list_cache) -> None: + """Evict from both cache layers; call when cache coherency is uncertain.""" + await cache.async_delete_cache(MODEL_LIST_CACHE_KEY) + + +async def refresh_model_list_cache( + prisma_client: "PrismaClient", + cache: DualCache = model_list_cache, +) -> None: + """Refresh the cache with current DB state; call after every model table write. + + Prevents stale-snapshot races by writing fresh data immediately after writes. + Other pods see new state within the 1s in-memory TTL. On DB error, evicts the + cache to avoid serving stale rows on fallback. + """ + if prisma_client is None: + await invalidate_model_list_cache(cache=cache) + return + + try: + from litellm.repositories.model_repository import ModelRepository + + new_models = await ModelRepository(prisma_client).table.find_many() + models = parse_model_rows(new_models) + await set_cached_model_rows(models, cache=cache) + except Exception as e: # noqa: BLE001 # the write already committed; a refresh fault must evict, never surface + verbose_proxy_logger.exception("Failed to refresh model_list_cache: %s", str(e)) + await invalidate_model_list_cache(cache=cache) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 15d977b75f2..acdb6984e08 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -26,6 +26,7 @@ from typing import ( List, Literal, Optional, + Sequence, Set, Tuple, TypedDict, @@ -235,6 +236,7 @@ from litellm.constants import ( GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_ADMIN_NAME, LITELLM_PROXY_BUDGET_NAME, + MODEL_LIST_SYNC_INTERVAL_SECONDS, PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS, PROXY_BATCH_POLLING_ENABLED, PROXY_BATCH_POLLING_INTERVAL, @@ -259,6 +261,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( ) from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.vertex_llm_base import VertexBase +from litellm.models.model import LiteLLM_ProxyModelTable from litellm.proxy._lazy_features import attach_lazy_features from litellm.proxy._types import * from litellm.proxy.analytics_endpoints.analytics_endpoints import ( @@ -499,6 +502,13 @@ from litellm.proxy.middleware.request_size_limit_middleware import ( from litellm.proxy.middleware.security_headers_middleware import ( SecurityHeadersMiddleware, ) +from litellm.proxy.model_list_cache import ( + get_cached_model_rows, + model_list_cache, + model_rows_fingerprint, + parse_model_rows, + set_cached_model_rows, +) from litellm.proxy.ocr_endpoints.endpoints import router as ocr_router from litellm.proxy.openai_files_endpoints.files_endpoints import ( router as openai_files_router, @@ -3767,6 +3777,7 @@ def _attach_redis_usage_cache(redis_cache: RedisCache, enable_redis_auth_cache: "the auth cache across workers and reduce DB load." ) litellm_config_cache.redis_cache = redis_cache + model_list_cache.redis_cache = redis_cache def resolve_routing_plugins( @@ -3837,6 +3848,7 @@ class ProxyConfig: self._last_semantic_filter_config: Optional[Dict[str, Any]] = None self._last_hashicorp_vault_config: Optional[Dict[str, Any]] = None self.worker_registry: List["WorkerRegistryEntry"] = [] + self._synced_model_list_fingerprint: Optional[str] = None def is_yaml(self, config_file_path: str) -> bool: if not os.path.isfile(config_file_path): @@ -5300,7 +5312,7 @@ class ProxyConfig: _model_info = RouterModelInfo(id=model.model_id, db_model=db_model) return _model_info - async def _delete_deployment(self, db_models: list) -> frozenset[str] | None: + async def _delete_deployment(self, db_models: Sequence[LiteLLM_ProxyModelTable]) -> frozenset[str] | None: """ (Helper function of add deployment) -> combined to reduce prisma db calls @@ -5377,7 +5389,7 @@ class ProxyConfig: return get_secret(decrypted_value) return decrypted_value - def _add_deployment(self, db_models: list) -> int: + def _add_deployment(self, db_models: Sequence[LiteLLM_ProxyModelTable]) -> int: """ Iterate through db models @@ -5419,7 +5431,7 @@ class ProxyConfig: added_models += 1 return added_models - def decrypt_model_list_from_db(self, new_models: list) -> list: + def decrypt_model_list_from_db(self, new_models: Sequence[LiteLLM_ProxyModelTable]) -> list: _model_list: list = [] for m in new_models: _litellm_params = m.litellm_params @@ -5449,7 +5461,7 @@ class ProxyConfig: async def _update_llm_router( self, - new_models: Optional[Json], + new_models: Optional[Sequence[LiteLLM_ProxyModelTable]], proxy_logging_obj: ProxyLogging, ) -> frozenset[str] | None: global llm_router, llm_model_list, master_key, general_settings @@ -5479,7 +5491,7 @@ class ProxyConfig: ) return - models_list: list = new_models if isinstance(new_models, list) else [] + models_list: Sequence[LiteLLM_ProxyModelTable] = new_models if llm_router is None and master_key is not None: verbose_proxy_logger.debug(f"len new_models: {len(models_list)}") @@ -6140,24 +6152,53 @@ class ProxyConfig: # Check if the object type is in the list (supports both str and enum values) return any(str(obj) == object_type_str for obj in supported_db_objects) - async def _get_models_from_db(self, prisma_client: PrismaClient) -> Optional[list]: + async def _get_models_from_db(self, prisma_client: PrismaClient) -> Optional[Sequence[LiteLLM_ProxyModelTable]]: """ - Fetch all model deployments from the DB. + Fetch all model deployments, from the shared cache when it is warm and from the + DB otherwise. Every write path refreshes the cache with fresh rows, so the write + is visible to other pods once their in-memory copy lapses; a miss falls back to the DB + and fills the shared copy with `NX` so it cannot overwrite a racing write's snapshot. Returns: - list: the rows (may be empty if no models exist) - None: signals a DB fetch *failure* — callers must not treat this as "all models deleted" and must not evict existing router deployments. """ + cached_models = await get_cached_model_rows() + if cached_models is not None: + return cached_models + try: new_models = await ModelRepository(prisma_client).table.find_many() - return new_models except Exception as e: verbose_proxy_logger.exception( "litellm.proxy_server.py::add_deployment() - Error getting new models from DB - {}".format(str(e)) ) return None + models = parse_model_rows(new_models) + await set_cached_model_rows(models, overwrite=False) + return models + + async def sync_model_deployments(self, prisma_client: PrismaClient, proxy_logging_obj: ProxyLogging) -> None: + """ + Reconcile this pod's router with the shared model list. + + Runs far more often than the full config reload, so it reads the shared cache and + stops at the fingerprint unless the model table actually changed; the router work + is reached only for a write this pod has not applied yet. + """ + db_models = await self._get_models_from_db(prisma_client=prisma_client) + if db_models is None: + return + + fingerprint = model_rows_fingerprint(db_models) + if fingerprint == self._synced_model_list_fingerprint: + return + + await self._update_llm_router(new_models=db_models, proxy_logging_obj=proxy_logging_obj) + self._synced_model_list_fingerprint = fingerprint + async def add_deployment( self, prisma_client: PrismaClient, @@ -6198,6 +6239,8 @@ class ProxyConfig: still_desired_ids = await self._update_llm_router( new_models=new_models, proxy_logging_obj=proxy_logging_obj ) + if new_models is not None: + self._synced_model_list_fingerprint = model_rows_fingerprint(new_models) db_general_settings = await get_config_param(prisma_client, "general_settings") @@ -8165,6 +8208,17 @@ class ProxyStartupEvent: # this will load all existing models on proxy startup await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) + if model_list_cache.redis_cache is not None: + scheduler.add_job( + proxy_config.sync_model_deployments, + "interval", + seconds=MODEL_LIST_SYNC_INTERVAL_SECONDS, + args=(prisma_client, proxy_logging_obj), + id="sync_model_deployments_job", + replace_existing=True, + misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, + ) + ### GET STORED CREDENTIALS ### scheduler.add_job( proxy_config.get_credentials, 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 8dbf38555b3..3cd551cda65 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 @@ -1938,6 +1938,89 @@ class TestAddAndDeleteModelLifecycle: assert str(exc_info.value.code) == "400" +class TestModelWritesInvalidateTheSharedModelList: + """Sibling pods read the model list from the shared cache. + + A write that leaves the pre-write list cached is invisible to every other pod until + the entry expires, which is the reported "deleted model still shows up" bug. + """ + + @staticmethod + def _cached_row(model_id: str) -> LiteLLM_ProxyModelTable: + return LiteLLM_ProxyModelTable( + model_id=model_id, + model_name="lifecycle-model", + litellm_params={"model": "openai/gpt-4.1-nano"}, + model_info={"id": model_id}, + created_by="test-admin", + updated_by="test-admin", + ) + + @pytest.mark.asyncio + async def test_delete_model_evicts_the_shared_model_list(self): + from litellm.proxy.management_endpoints.model_management_endpoints import ( + ModelInfoDelete, + ) + from litellm.proxy.management_endpoints.model_management_endpoints import ( + delete_model as delete_model_endpoint, + ) + from litellm.proxy.model_list_cache import ( + get_cached_model_rows, + set_cached_model_rows, + ) + + model_id = "cache-invalidation-model" + db_row = self._cached_row(model_id) + await set_cached_model_rows([db_row]) + + mock_prisma = MagicMock() + mock_prisma.db = MagicMock() + mock_prisma.db.litellm_proxymodeltable = AsyncMock() + mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=db_row) + mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row) + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) # Post-delete, no models exist + + _PS = "litellm.proxy.proxy_server" + with ( + patch(f"{_PS}.prisma_client", mock_prisma), + patch(f"{_PS}.store_model_in_db", True), + patch(f"{_PS}.proxy_config", MagicMock(add_deployment=AsyncMock())), + patch(f"{_PS}.proxy_logging_obj", MagicMock()), + patch(f"{_PS}.general_settings", {}), + patch(f"{_PS}.premium_user", True), + patch(f"{_PS}.llm_router", MagicMock()), + ): + await delete_model_endpoint( + model_info=ModelInfoDelete(id=model_id), + user_api_key_dict=UserAPIKeyAuth( + user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN + ), + ) + + assert await get_cached_model_rows() == () # Empty tuple, not None + + @pytest.mark.asyncio + async def test_clear_cache_evicts_the_shared_model_list(self): + """Every update path (patch, block, access groups) reloads through clear_cache.""" + from litellm.proxy.management_endpoints.model_management_endpoints import ( + clear_cache, + ) + from litellm.proxy.model_list_cache import ( + get_cached_model_rows, + set_cached_model_rows, + ) + + await set_cached_model_rows([self._cached_row("cache-clear-model")]) + + with ( + patch("litellm.proxy.proxy_server.llm_router", None), + patch("litellm.proxy.proxy_server.prisma_client", None), + ): + await clear_cache() + + assert await get_cached_model_rows() is None + + class TestDeleteTeamBYOKModelGhost: """Regression for issue #22594. diff --git a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py index 99287b92b8f..742c6ce5541 100644 --- a/tests/test_litellm/proxy/proxy_server/test_proxy_config.py +++ b/tests/test_litellm/proxy/proxy_server/test_proxy_config.py @@ -10,6 +10,7 @@ from __future__ import annotations import json import os +from datetime import datetime from types import SimpleNamespace from typing import Any, Dict from unittest.mock import AsyncMock, MagicMock @@ -17,6 +18,8 @@ from unittest.mock import AsyncMock, MagicMock import pytest import litellm +from litellm.models.model import LiteLLM_ProxyModelTable +from litellm.proxy.model_list_cache import invalidate_model_list_cache from litellm.proxy.proxy_server import ( ProxyConfig, _is_remote_module_url, @@ -2354,3 +2357,143 @@ async def test_ProxyConfig_load_config_redacts_secret_litellm_setting_keeps_plai assert "num_retries=7" in rendered, ( f"non-secret num_retries value was over-redacted; expected it visible in {rendered!r}" ) + +# --------------------------------------------------------------------------- +# ProxyConfig._get_models_from_db / ProxyConfig.sync_model_deployments +# --------------------------------------------------------------------------- + + +class _NoConfigModelsProxyConfig(ProxyConfig): + async def get_config(self, config_file_path=None) -> dict: + return {"model_list": []} + + +def _db_row(model_id: str, updated_at: datetime) -> LiteLLM_ProxyModelTable: + return LiteLLM_ProxyModelTable( + model_id=model_id, + model_name=f"group-{model_id}", + litellm_params={"model": "openai/gpt-4.1-nano", "api_key": "sk-test"}, + model_info={"id": model_id}, + updated_at=updated_at, + ) + + +def _prisma_returning(rows: list[LiteLLM_ProxyModelTable] | Exception) -> MagicMock: + prisma = MagicMock() + prisma.db = MagicMock() + prisma.db.litellm_proxymodeltable = MagicMock() + if isinstance(rows, Exception): + prisma.db.litellm_proxymodeltable.find_many = AsyncMock(side_effect=rows) + else: + prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=rows) + return prisma + + +@pytest.fixture +async def empty_model_list_cache(): + await invalidate_model_list_cache() + yield + await invalidate_model_list_cache() + + +@pytest.mark.asyncio +async def test_ProxyConfig__get_models_from_db_serves_the_shared_cache(empty_model_list_cache, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + prisma = _prisma_returning([_db_row("m-1", datetime(2026, 1, 1))]) + pc = _NoConfigModelsProxyConfig() + + first = await pc._get_models_from_db(prisma_client=prisma) + second = await pc._get_models_from_db(prisma_client=prisma) + + assert [model.model_id for model in first] == ["m-1"] + assert [model.model_id for model in second] == ["m-1"] + # Note: without Redis attached, read-fill skips in-memory to prevent concurrent races, + # so in-memory stays empty and the second call also hits the DB. In production with Redis, + # read-fill populates Redis (NX) and the next sync backfills in-memory from Redis, so the + # cache is still efficient. See test_ProxyConfig_sync_model_deployments_evicts_a_model_deleted_by_another_pod + # for the multi-pod race scenario this design protects against. + assert prisma.db.litellm_proxymodeltable.find_many.await_count == 2 + + +@pytest.mark.asyncio +async def test_ProxyConfig__get_models_from_db_returns_none_on_db_failure(empty_model_list_cache, monkeypatch): + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + prisma = _prisma_returning(RuntimeError("db down")) + pc = _NoConfigModelsProxyConfig() + + assert await pc._get_models_from_db(prisma_client=prisma) is None + + +@pytest.mark.asyncio +async def test_ProxyConfig_sync_model_deployments_evicts_a_model_deleted_by_another_pod( + empty_model_list_cache, monkeypatch +): + """The reported bug: a delete served by another pod must not keep serving here.""" + router = litellm.Router( + model_list=[ + { + "model_name": "group-m-1", + "litellm_params": {"model": "openai/gpt-4.1-nano", "api_key": "sk-test"}, + "model_info": {"id": "m-1", "db_model": True}, + } + ] + ) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + pc = _NoConfigModelsProxyConfig() + + await pc.sync_model_deployments(proxy_logging_obj=MagicMock(), prisma_client=_prisma_returning([_db_row("m-1", datetime(2026, 1, 1))])) + assert router.get_model_ids() == ["m-1"] + + await invalidate_model_list_cache() + await pc.sync_model_deployments(proxy_logging_obj=MagicMock(), prisma_client=_prisma_returning([])) + + assert router.get_model_ids() == [] + + +@pytest.mark.asyncio +async def test_ProxyConfig_sync_model_deployments_picks_up_a_model_added_by_another_pod( + empty_model_list_cache, monkeypatch +): + router = litellm.Router(model_list=[]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + pc = _NoConfigModelsProxyConfig() + + await pc.sync_model_deployments(proxy_logging_obj=MagicMock(), prisma_client=_prisma_returning([_db_row("m-1", datetime(2026, 1, 1))])) + + assert router.get_model_ids() == ["m-1"] + + +@pytest.mark.asyncio +async def test_ProxyConfig_sync_model_deployments_leaves_the_router_alone_when_nothing_changed( + empty_model_list_cache, monkeypatch +): + """The tick runs every couple of seconds; an unchanged list must not touch the router.""" + router = MagicMock() + router.get_model_ids = MagicMock(return_value=["m-1"]) + router.get_model_list = MagicMock(return_value=[]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + pc = _NoConfigModelsProxyConfig() + prisma = _prisma_returning([_db_row("m-1", datetime(2026, 1, 1))]) + + await pc.sync_model_deployments(prisma_client=prisma, proxy_logging_obj=MagicMock()) + upserts_after_first_sync = router.upsert_deployment.call_count + await invalidate_model_list_cache() + await pc.sync_model_deployments(prisma_client=prisma, proxy_logging_obj=MagicMock()) + + assert upserts_after_first_sync == 1 + assert router.upsert_deployment.call_count == 1 + assert router.delete_deployment.call_count == 0 + + +@pytest.mark.asyncio +async def test_ProxyConfig_sync_model_deployments_keeps_deployments_when_the_db_is_down( + empty_model_list_cache, monkeypatch +): + router = MagicMock() + router.get_model_ids = MagicMock(return_value=["m-1"]) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", router) + pc = _NoConfigModelsProxyConfig() + + await pc.sync_model_deployments(proxy_logging_obj=MagicMock(), prisma_client=_prisma_returning(RuntimeError("db down"))) + + assert router.delete_deployment.call_count == 0 diff --git a/tests/test_litellm/proxy/test_model_list_cache.py b/tests/test_litellm/proxy/test_model_list_cache.py new file mode 100644 index 00000000000..c2aa228cb60 --- /dev/null +++ b/tests/test_litellm/proxy/test_model_list_cache.py @@ -0,0 +1,235 @@ +import asyncio +from datetime import datetime, timedelta +from unittest.mock import AsyncMock, MagicMock + +import fakeredis.aioredis +import pytest + +from litellm.caching.dual_cache import DualCache +from litellm.caching.redis_cache import RedisCache +from litellm.models.model import LiteLLM_ProxyModelTable +from litellm.proxy.model_list_cache import ( + MODEL_LIST_CACHE_IN_MEMORY_TTL_SECONDS, + MODEL_LIST_CACHE_KEY, + MODEL_LIST_CACHE_REDIS_TTL_SECONDS, + get_cached_model_rows, + invalidate_model_list_cache, + model_rows_fingerprint, + parse_model_rows, + refresh_model_list_cache, + set_cached_model_rows, +) + +UPDATED_AT = datetime(2026, 1, 1, 12, 0, 0) + + +def make_model(model_id: str, updated_at: datetime = UPDATED_AT) -> LiteLLM_ProxyModelTable: + return LiteLLM_ProxyModelTable( + model_id=model_id, + model_name=f"group-{model_id}", + litellm_params={"model": "openai/gpt-4.1-nano", "api_key": "encrypted-value"}, + model_info={"id": model_id}, + updated_at=updated_at, + ) + + +def pod_cache(redis_client: fakeredis.aioredis.FakeRedis) -> DualCache: + """A pod's own DualCache, sharing the cluster's Redis with the other pods.""" + redis_cache = RedisCache(host="localhost", port="6379") + redis_cache.init_async_client = lambda **kwargs: redis_client + return DualCache( + redis_cache=redis_cache, + default_in_memory_ttl=MODEL_LIST_CACHE_IN_MEMORY_TTL_SECONDS, + ) + + +@pytest.fixture +def shared_redis() -> fakeredis.aioredis.FakeRedis: + return fakeredis.aioredis.FakeRedis() + + +class TestModelListCacheSharing: + @pytest.mark.asyncio + async def test_write_on_one_pod_is_readable_on_another(self, shared_redis): + pod_a, pod_b = pod_cache(shared_redis), pod_cache(shared_redis) + + await set_cached_model_rows([make_model("m-1")], cache=pod_a) + + rows = await get_cached_model_rows(cache=pod_b) + assert rows is not None + assert [row.model_id for row in rows] == ["m-1"] + assert rows[0].litellm_params["api_key"] == "encrypted-value" + + @pytest.mark.asyncio + async def test_invalidation_on_one_pod_forces_a_db_read_on_another(self, shared_redis): + """The fix: without this eviction the sibling pod keeps serving the pre-write list.""" + pod_a, pod_b = pod_cache(shared_redis), pod_cache(shared_redis) + + await set_cached_model_rows([make_model("m-1"), make_model("m-2")], cache=pod_a) + assert await get_cached_model_rows(cache=pod_b) is not None + + await invalidate_model_list_cache(cache=pod_a) + assert await shared_redis.get(MODEL_LIST_CACHE_KEY) is None + + await asyncio.sleep(MODEL_LIST_CACHE_IN_MEMORY_TTL_SECONDS + 0.1) + assert await get_cached_model_rows(cache=pod_b) is None + + @pytest.mark.asyncio + async def test_shared_entry_outlives_the_local_copy(self, shared_redis): + """A local TTL on the shared entry would send every pod to the DB on every sync.""" + await set_cached_model_rows([make_model("m-1")], cache=pod_cache(shared_redis)) + + assert await shared_redis.ttl(MODEL_LIST_CACHE_KEY) > MODEL_LIST_CACHE_IN_MEMORY_TTL_SECONDS + assert MODEL_LIST_CACHE_REDIS_TTL_SECONDS > MODEL_LIST_CACHE_IN_MEMORY_TTL_SECONDS + + @pytest.mark.asyncio + async def test_reads_return_fresh_copies_so_callers_cannot_corrupt_the_cache(self, shared_redis): + """`_add_deployment` decrypts `litellm_params` in place on whatever it is handed.""" + cache = pod_cache(shared_redis) + await set_cached_model_rows([make_model("m-1")], cache=cache) + + first = await get_cached_model_rows(cache=cache) + assert first is not None + first[0].litellm_params["api_key"] = "decrypted-value" + + second = await get_cached_model_rows(cache=cache) + assert second is not None + assert second[0].litellm_params["api_key"] == "encrypted-value" + + @pytest.mark.asyncio + async def test_unreadable_entry_falls_back_to_the_db(self, shared_redis): + cache = pod_cache(shared_redis) + await cache.async_set_cache(MODEL_LIST_CACHE_KEY, [{"not": "a model row"}], local_only=True) + + assert await get_cached_model_rows(cache=cache) is None + + @pytest.mark.asyncio + async def test_empty_model_list_is_cached_as_empty_not_as_a_miss(self, shared_redis): + """An operator who deleted every model must not send every pod back to the DB.""" + cache = pod_cache(shared_redis) + await set_cached_model_rows([], cache=cache) + + assert await get_cached_model_rows(cache=cache) == () + + @pytest.mark.asyncio + async def test_read_fill_cannot_overwrite_a_racing_write(self, shared_redis): + """A read that began pre-write must not republish its stale snapshot over the write. + + The write path publishes fresh rows; a slower read-fill holding the pre-write list + then writes with `overwrite=False`, so the shared Redis copy keeps the fresh rows. + """ + writer, reader = pod_cache(shared_redis), pod_cache(shared_redis) + + await set_cached_model_rows([make_model("m-new")], cache=writer) + await set_cached_model_rows([make_model("m-stale")], cache=reader, overwrite=False) + + rows = await get_cached_model_rows(cache=pod_cache(shared_redis)) + assert rows is not None + assert [row.model_id for row in rows] == ["m-new"] + + @pytest.mark.asyncio + async def test_read_fill_still_populates_an_empty_shared_key(self, shared_redis): + """`NX` must only decline when a value already exists, not disable read-through fills.""" + await set_cached_model_rows([make_model("m-1")], cache=pod_cache(shared_redis), overwrite=False) + + rows = await get_cached_model_rows(cache=pod_cache(shared_redis)) + assert rows is not None + assert [row.model_id for row in rows] == ["m-1"] + + @pytest.mark.asyncio + async def test_read_fill_does_not_populate_in_memory_when_write_races(self, shared_redis): + """Read-fill must not replace a pod's in-memory cache even if the Redis NX succeeds. + + A write path publishes fresh rows; if a stale read-fill replaces the in-memory + copy on the same pod, the pod will use obsolete rows until the 1s TTL lapses. + So read-fill skips the local cache entirely, forcing the next sync to hit Redis. + """ + cache = pod_cache(shared_redis) + + await set_cached_model_rows([make_model("m-fresh")], cache=cache, overwrite=True) + assert await get_cached_model_rows(cache=cache) is not None + assert (await get_cached_model_rows(cache=cache))[0].model_id == "m-fresh" + + await set_cached_model_rows([make_model("m-stale")], cache=cache, overwrite=False) + + rows = await get_cached_model_rows(cache=cache) + assert rows is not None + assert [row.model_id for row in rows] == ["m-fresh"], "in-memory must keep fresh rows" + + +class TestModelRowsFingerprint: + def test_ignores_row_order(self): + rows = [make_model("m-1"), make_model("m-2")] + assert model_rows_fingerprint(rows) == model_rows_fingerprint(list(reversed(rows))) + + def test_changes_when_a_row_is_deleted(self): + rows = [make_model("m-1"), make_model("m-2")] + assert model_rows_fingerprint(rows) != model_rows_fingerprint(rows[:1]) + + def test_changes_when_a_row_is_edited(self): + edited = make_model("m-1", updated_at=UPDATED_AT + timedelta(seconds=1)) + assert model_rows_fingerprint([make_model("m-1")]) != model_rows_fingerprint([edited]) + + +class TestParseModelRows: + def test_accepts_orm_rows_and_mappings(self): + row = make_model("m-1") + assert parse_model_rows([row]) == parse_model_rows([row.model_dump()]) + + def test_rejects_rows_it_cannot_read(self): + with pytest.raises(TypeError): + parse_model_rows(["not-a-row"]) + + +class TestRefreshModelListCache: + """Test the write-through refresh logic that fixes the stale-snapshot race.""" + + @pytest.mark.asyncio + async def test_refresh_writes_fresh_rows_to_both_cache_tiers(self, shared_redis): + """The fix: refresh writes fresh data, not delete-only.""" + cache = pod_cache(shared_redis) + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[make_model("m-1")]) + + await refresh_model_list_cache(prisma_client=mock_prisma, cache=cache) + + rows = await get_cached_model_rows(cache=cache) + assert rows is not None + assert [row.model_id for row in rows] == ["m-1"] + + @pytest.mark.asyncio + async def test_refresh_handles_empty_model_list(self, shared_redis): + """A deleted model returns an empty list; refresh must cache that, not evict.""" + cache = pod_cache(shared_redis) + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[]) + + await refresh_model_list_cache(prisma_client=mock_prisma, cache=cache) + + rows = await get_cached_model_rows(cache=cache) + assert rows == () # Empty, not None + + @pytest.mark.asyncio + async def test_refresh_evicts_on_db_error(self, shared_redis): + """If the refresh fails, evict so the next read hits the DB fresh.""" + cache = pod_cache(shared_redis) + await set_cached_model_rows([make_model("m-1")], cache=cache) + + mock_prisma = MagicMock() + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(side_effect=RuntimeError("DB down")) + + await refresh_model_list_cache(prisma_client=mock_prisma, cache=cache) + + rows = await get_cached_model_rows(cache=cache) + assert rows is None # Evicted on error + + @pytest.mark.asyncio + async def test_refresh_handles_none_prisma_client(self, shared_redis): + """Guard against None prisma_client to avoid spurious errors.""" + cache = pod_cache(shared_redis) + await set_cached_model_rows([make_model("m-1")], cache=cache) + + await refresh_model_list_cache(prisma_client=None, cache=cache) + + rows = await get_cached_model_rows(cache=cache) + assert rows is None # Evicted when prisma is None