mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-14 23:21:35 +00:00
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
This commit is contained in:
parent
b4ff05be8e
commit
90d813a4af
10 changed files with 676 additions and 8 deletions
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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))
|
||||
|
|
|
|||
137
litellm/proxy/model_list_cache.py
Normal file
137
litellm/proxy/model_list_cache.py
Normal file
|
|
@ -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)
|
||||
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -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.
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
235
tests/test_litellm/proxy/test_model_list_cache.py
Normal file
235
tests/test_litellm/proxy/test_model_list_cache.py
Normal file
|
|
@ -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
|
||||
Loading…
Add table
Reference in a new issue