scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-12 23:01:41 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 88

fix(proxy): serialize model reconciles so concurrent writes stop evicting each other

Browse source 
A model write is a read-modify-write of the shared `llm_router` global: read the
db into a snapshot, then make the router match that snapshot. Nothing serialized
it, so two of them interleaving was not a lost update but an eviction --
_delete_deployment removes every live deployment absent from the snapshot it was
handed, so the request holding the older snapshot reconciles the newer request's
model straight back out of the router. The row survives in the db, which is what
makes it easy to miss: the pod simply stops serving a model it was told to serve
until some later reload happens to put it back.

clear_cache compounds it. It deletes every db model from the router before
reloading them, so for the width of that reload the pod serves none of them --
and any concurrent write sampling the router in that window sees the hole.

Fix is one lock (MODEL_RECONCILE_LOCK) held across both, so each reconcile reads
the db and applies it atomically and no stale snapshot can evict a newer model.
clear_cache holds it across wipe+reload and calls the already-locked
_add_deployment_locked, since asyncio.Lock is not reentrant and routing back
through the public add_deployment would deadlock the pod's whole model-write
path.

The verdict needed the same treatment. raise_if_reload_degraded_serving compared
a desired-set read during the reload against a router snapshot taken after it,
so a neighbouring reconcile's in-flight wipe was reported to the caller as
collateral damage from its own reload -- a 500 on a create that had in fact
succeeded. Reconciles now return a ReconcileOutcome carrying both the desired set
and the post-reconcile serving state, captured before the lock is released, and
the verdict judges against that. Omitting live_after keeps the old live re-read,
which stays correct for the no-reconcile-ran case.

Found by running the e2e suite with pytest-xdist at 8 workers: three unrelated
tests failed together on "Previously served model id(s) [...] are also no longer
being served by this pod", which is this. Serial runs concurrent enough to hit it
are rare, which is why 78 minutes of sequential e2e never surfaced it -- but any
customer provisioning models in parallel (terraform, CI) is in exactly this race.

test_reconciles_serialize_so_no_stale_snapshot_can_evict fails with 5 == 1
without the lock.
This commit is contained in:
Yuneng Jiang 2026-08-12 10:17:21 -07:00
parent 6a919aec6a
commit 6553bc8957
No known key found for this signature in database
5 changed files with 338 additions and 84 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

23
litellm/proxy/_types.py
View file

@ -3,7 +3,7 @@ import json
import os
from collections.abc import Callable
from datetime import datetime
from typing import TYPE_CHECKING, Any, Final, Literal
from typing import TYPE_CHECKING, Any, Final, Literal, NamedTuple
import httpx
from pydantic import (
@ -72,6 +72,27 @@ else:
Span = Any
class ReconcileOutcome(NamedTuple):
"""What a model reconcile observed, captured while it still held the reconcile
lock.
Both fields have to be read under that lock to be worth anything. ``live_after``
in particular is the router's serving state the instant this reconcile finished,
which is NOT the same as what a later snapshot would see: any other model write
admitted in between briefly un-serves every db model (see ``clear_cache``), so a
caller that re-snapshots at verdict time can observe that hole and blame its own
reload for it.
- ``still_desired``: the db + config ids the reconcile reconciled against, or None
when no reconcile ran and the desired set is therefore unknown.
- ``live_after``: the ids the router served immediately after the reconcile, or
None when no reconcile ran.
"""
still_desired: frozenset[str] | None
live_after: frozenset[str] | None
class SupportedDBObjectType(str, enum.Enum):
"""
Supported database object types for fine-grained DB storage control.

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

@ -34,6 +34,7 @@ from litellm.proxy._types import (
PrismaCompatibleUpdateDBModel,
ProxyErrorTypes,
ProxyException,
ReconcileOutcome,
TeamModelAddRequest,
TeamModelDeleteRequest,
UserAPIKeyAuth,
@ -402,7 +403,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: Final = live_model_ids_snapshot()
still_desired_ids: Final = await clear_cache()
reload_outcome: Final = await clear_cache()
## CREATE AUDIT LOG ##
asyncio.create_task(
@ -422,7 +423,8 @@ async def patch_model(
before=live_before_reload,
written_models=[(model_id, getattr(updated_model, "model_info", None))],
action="update",
still_desired=still_desired_ids,
still_desired=reload_outcome.still_desired,
live_after=reload_outcome.live_after,
)
return updated_model
@ -508,7 +510,7 @@ async def _set_model_blocked_status(
)
live_before_reload: Final = live_model_ids_snapshot()
still_desired_ids: Final = await clear_cache()
reload_outcome: Final = await clear_cache()
asyncio.create_task(
create_object_audit_log(
@ -529,7 +531,8 @@ async def _set_model_blocked_status(
before=live_before_reload,
written_models=[(data.model_id, getattr(updated_model, "model_info", None))],
action=action,
still_desired=still_desired_ids,
still_desired=reload_outcome.still_desired,
live_after=reload_outcome.live_after,
)
return updated_model
@ -1431,7 +1434,7 @@ async def add_new_model(
"""
live_before_reload: Final = live_model_ids_snapshot()
still_desired_ids: frozenset[str] | None = None
reload_outcome: ReconcileOutcome = ReconcileOutcome(still_desired=None, live_after=None)
try:
_original_litellm_model_name: Final = model_params.model_name
if model_params.model_info.team_id is None:
@ -1446,7 +1449,7 @@ async def add_new_model(
user_api_key_dict=user_api_key_dict,
prisma_client=prisma_client,
)
still_desired_ids = await proxy_config.add_deployment(
reload_outcome = await proxy_config.add_deployment(
prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj
)
# don't let failed slack alert block the /model/new response
@ -1493,7 +1496,8 @@ async def add_new_model(
before=live_before_reload,
written_models=[(model_response.model_id, getattr(model_response, "model_info", None))],
action="create",
still_desired=still_desired_ids,
still_desired=reload_outcome.still_desired,
live_after=reload_outcome.live_after,
)
return model_response
@ -1620,7 +1624,7 @@ async def update_model(
# Clear cache and reload models (uses config setting or defaults to preserving config models for DB updates)
live_before_reload: Final = live_model_ids_snapshot()
still_desired_ids: Final = await clear_cache()
reload_outcome: Final = await clear_cache()
## CREATE AUDIT LOG ##
asyncio.create_task(
create_object_audit_log(
@ -1647,7 +1651,8 @@ async def update_model(
before=live_before_reload,
written_models=[(_model_id, getattr(model_response, "model_info", None))],
action="update",
still_desired=still_desired_ids,
still_desired=reload_outcome.still_desired,
live_after=reload_outcome.live_after,
)
return model_response
@ -1952,6 +1957,7 @@ def reload_serving_verdict(
written_models: Sequence[tuple[str, object]],
written_must_serve: bool,
still_desired: frozenset[str] | None = None,
live_after: frozenset[str] | None = None,
) -> 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.
@ -1973,9 +1979,16 @@ def reload_serving_verdict(
yet polled, so the reload dropping it is the reconcile working rather than damage.
Without it (no reconcile ran) every drop is reported, which is the safe direction.
``live_after`` is the router's serving state captured by the reload itself, while it
still held MODEL_RECONCILE_LOCK. Pass it whenever the caller has it: re-reading the
router here instead means sampling it after the lock was released, where the NEXT
reconcile's leading wipe (clear_cache un-serves every db model before reloading
them) shows up as this reload having dropped them. Falling back to a fresh read is
only correct when no reconcile ran and there is nothing to be concurrent with.
Returns (written ids violating their obligation, collateral ids no longer served).
"""
now: Final = live_model_ids_snapshot()
now: Final = live_model_ids_snapshot() if live_after is None else live_after
written_ids: Final = frozenset(model_id for model_id, _ in written_models)
if written_must_serve:
missing = tuple(
@ -1995,16 +2008,23 @@ def raise_if_reload_degraded_serving(
written_models: Sequence[tuple[str, object]],
action: str,
still_desired: frozenset[str] | None = None,
live_after: frozenset[str] | None = None,
) -> 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."""
speaks only for the handling pod.
Callers hold a ReconcileOutcome from the reload; pass BOTH of its fields. Supplying
still_desired without live_after mixes a snapshot taken under the reconcile lock
with one taken after it was released, which is what makes a concurrent model write
look like collateral damage."""
missing, collateral = reload_serving_verdict(
before=before,
written_models=written_models,
written_must_serve=True,
still_desired=still_desired,
live_after=live_after,
)
if not missing and not collateral:
return
@ -2031,14 +2051,22 @@ def raise_if_reload_degraded_serving(
)
async def clear_cache() -> frozenset[str] | None:
async def clear_cache() -> ReconcileOutcome:
"""
Clear router caches and reload models.
Returns the db + config id set the reload reconciled against, or None when no
reload ran, so callers can pass it to raise_if_reload_degraded_serving.
Returns what the reload saw (see ReconcileOutcome) so callers can pass it to
raise_if_reload_degraded_serving.
Runs under MODEL_RECONCILE_LOCK for its whole extent, not just the reload at the
end. The wipe below un-serves EVERY db model before add_deployment puts them back,
so an unserialized concurrent model write would snapshot the router mid-hole and
report every one of them as collateral damage from its own reload. Holding the lock
across wipe+reload makes the pair atomic to any other reconcile; the inner call is
_add_deployment_locked because add_deployment would re-acquire and deadlock.
"""
from litellm.proxy.proxy_server import (
MODEL_RECONCILE_LOCK,
llm_router,
prisma_client,
proxy_config,
@ -2048,61 +2076,64 @@ async def clear_cache() -> frozenset[str] | None:
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
return ReconcileOutcome(still_desired=None, live_after=None)
try:
# Only clear DB models, preserve config models
verbose_proxy_logger.debug("Clearing only DB models, preserving config models")
async with MODEL_RECONCILE_LOCK:
try:
# Only clear DB models, preserve config models
verbose_proxy_logger.debug("Clearing only DB models, preserving config models")
# Get current models and filter out DB models
current_models: Final = llm_router.model_list.copy()
config_models: Final = []
db_model_ids: Final = []
# Get current models and filter out DB models
current_models: Final = llm_router.model_list.copy()
config_models: Final = []
db_model_ids: Final = []
for model in current_models:
model_info = model.get("model_info", {})
if model_info.get("db_model", False):
# This is a DB model, mark for deletion
db_model_ids.append(model_info.get("id"))
else:
# This is a config model, preserve it
config_models.append(model)
for model in current_models:
model_info = model.get("model_info", {})
if model_info.get("db_model", False):
# This is a DB model, mark for deletion
db_model_ids.append(model_info.get("id"))
else:
# This is a config model, preserve it
config_models.append(model)
# Clear only DB models
for model_id in db_model_ids:
llm_router.delete_deployment(id=model_id)
# Clear only DB models
for model_id in db_model_ids:
llm_router.delete_deployment(id=model_id)
# Clear only DB-backed auto-router-family entries, keyed by model_name, so the
# reload below rebuilds them fresh. A blanket .clear() would also drop config-defined
# routers, which are never re-added below (add_deployment only reloads DB models),
# leaving them permanently unroutable until a full proxy restart for every tenant.
# Restrict to deployments whose model is actually an auto_router/* so a config
# router that merely shares a model_name with a regular DB model isn't evicted. The
# auto_router/ prefix also covers quality_router/ and adaptive_router/, so pop the
# name from every router registry (no-op where absent); missing quality/adaptive
# entries would otherwise make init raise "already exists" on reload and abort it.
db_router_names: Final = {
model.get("model_name")
for model in current_models
if model.get("model_name") is not None
and model.get("model_info", {}).get("db_model", False)
and str(model.get("litellm_params", {}).get("model", "")).startswith("auto_router/")
}
for model_name in db_router_names:
llm_router.auto_routers.pop(model_name, None)
llm_router.complexity_routers.pop(model_name, None)
llm_router.adaptive_routers.pop(model_name, None)
llm_router.quality_routers.pop(model_name, None)
# Clear only DB-backed auto-router-family entries, keyed by model_name, so the
# reload below rebuilds them fresh. A blanket .clear() would also drop config-defined
# routers, which are never re-added below (add_deployment only reloads DB models),
# leaving them permanently unroutable until a full proxy restart for every tenant.
# Restrict to deployments whose model is actually an auto_router/* so a config
# router that merely shares a model_name with a regular DB model isn't evicted. The
# auto_router/ prefix also covers quality_router/ and adaptive_router/, so pop the
# name from every router registry (no-op where absent); missing quality/adaptive
# entries would otherwise make init raise "already exists" on reload and abort it.
db_router_names: Final = {
model.get("model_name")
for model in current_models
if model.get("model_name") is not None
and model.get("model_info", {}).get("db_model", False)
and str(model.get("litellm_params", {}).get("model", "")).startswith("auto_router/")
}
for model_name in db_router_names:
llm_router.auto_routers.pop(model_name, None)
llm_router.complexity_routers.pop(model_name, None)
llm_router.adaptive_routers.pop(model_name, None)
llm_router.quality_routers.pop(model_name, None)
# Reload only DB models
still_desired_ids: Final = await proxy_config.add_deployment(
prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj
)
# Reload only DB models. _add_deployment_locked, not add_deployment: this
# coroutine already holds MODEL_RECONCILE_LOCK and asyncio.Lock is not
# reentrant, so the public wrapper would deadlock against itself.
outcome: Final = await proxy_config._add_deployment_locked(
prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj
)
verbose_proxy_logger.debug(
"Cleared %s DB models, preserved %s config models", len(db_model_ids), len(config_models)
)
return still_desired_ids
except Exception as e:
verbose_proxy_logger.exception("Failed to clear cache and reload models. Due to error - %s", e)
return None
verbose_proxy_logger.debug(
"Cleared %s DB models, preserved %s config models", len(db_model_ids), len(config_models)
)
return outcome
except Exception as e:
verbose_proxy_logger.exception("Failed to clear cache and reload models. Due to error - %s", e)
return ReconcileOutcome(still_desired=None, live_after=None)

48
litellm/proxy/proxy_server.py
View file

@ -452,6 +452,7 @@ from litellm.proxy.management_endpoints.model_management_endpoints import (
_add_model_to_db,
_add_team_model_to_db,
_deduplicate_litellm_router_models,
live_model_ids_snapshot,
)
from litellm.proxy.management_endpoints.model_management_endpoints import (
router as model_management_router,
@ -2074,6 +2075,15 @@ experimental = False
#### GLOBAL VARIABLES ####
llm_router: Router | None = None
llm_model_list: list | None = None
# Serializes every model reconcile (ProxyConfig.add_deployment and clear_cache) so the
# read-modify-write of llm_router above is atomic. Without it, two concurrent model
# writes each reconcile the router against their OWN db snapshot, and the one holding
# the older snapshot evicts the deployment the newer one just added -- the db keeps the
# row, this pod stops serving it. Control-plane only (model create/update/delete and
# the config-sync tick), never on a completion path, so the serialization is free.
# Module-level rather than per-ProxyConfig because llm_router is a module global and a
# second ProxyConfig instance must not get its own independent lock over it.
MODEL_RECONCILE_LOCK: Final = asyncio.Lock()
general_settings: dict = {}
config_passthrough_endpoints: list[dict[str, Any]] | None = None
log_file: Final = "api_log.json"
@ -6353,16 +6363,39 @@ class ProxyConfig:
self,
prisma_client: PrismaClient,
proxy_logging_obj: ProxyLogging,
) -> frozenset[str] | None:
) -> ReconcileOutcome:
"""
- Check db for new models
- Check if model id's in router already
- If not, add to router
Returns the ids the db + config say should be served after the reconcile, or
None when no reconcile ran. Callers that judge their own reload need it to tell
a deliberate eviction from a deployment that went missing.
Serialized against every other model reconcile by MODEL_RECONCILE_LOCK, because
the work below is a read-modify-write of the shared ``llm_router`` global: it
reads the db into a snapshot and then makes the router match that snapshot. Two
of those interleaving is not a lost update but an eviction -- the request whose
snapshot predates the other's commit reconciles the newer model *out* of the
router, since _delete_deployment removes every live deployment absent from the
snapshot it was handed. The model stays in the db and this pod stops serving it
until some later reload puts it back.
Returns what the reconcile saw, captured before the lock is released so a
caller's verdict cannot be corrupted by the next reconcile's own in-flight
window. See ReconcileOutcome.
"""
async with MODEL_RECONCILE_LOCK:
return await self._add_deployment_locked(
prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj
)
async def _add_deployment_locked(
self,
prisma_client: PrismaClient,
proxy_logging_obj: ProxyLogging,
) -> ReconcileOutcome:
"""add_deployment's body, minus the locking. MODEL_RECONCILE_LOCK MUST already
be held. Split out for the one caller that has to hold the lock across more than
this reconcile -- clear_cache, which un-serves every db model before calling it
and would deadlock on a re-acquire."""
global llm_router, llm_model_list, master_key, general_settings
still_desired_ids: frozenset[str] | None = None
@ -6403,7 +6436,12 @@ class ProxyConfig:
except Exception as e:
verbose_proxy_logger.exception("litellm.proxy.proxy_server.py::ProxyConfig:add_deployment - %s", e)
return still_desired_ids
# Read while the lock is still held: once it is released the next reconcile can
# begin, and clear_cache's leading wipe would make this look like a mass drop.
return ReconcileOutcome(
still_desired=still_desired_ids,
live_after=None if still_desired_ids is None else live_model_ids_snapshot(),
)
def start_config_sync_subscriber(
self,

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

@ -18,6 +18,7 @@ from litellm.proxy._types import (
LiteLLM_TeamTable,
LitellmUserRoles,
Member,
ReconcileOutcome,
UserAPIKeyAuth,
)
from litellm.proxy.management_endpoints.model_management_endpoints import (
@ -103,11 +104,11 @@ class MockProxyConfig:
self.success = success
self.deployment_called = False
async def add_deployment(self, prisma_client, proxy_logging_obj):
async def _add_deployment_locked(self, prisma_client, proxy_logging_obj):
self.deployment_called = True
if not self.success:
raise Exception("Failed to add deployment")
return True
return ReconcileOutcome(still_desired=frozenset(), live_after=frozenset())
class TestModelManagementAuthChecks:
@ -409,7 +410,9 @@ class TestClearCache:
mock_router.model_list = ["openai/gpt-4o", "openai/gpt-4o-mini"]
mock_config = MagicMock()
mock_config.add_deployment = AsyncMock(return_value=True)
mock_config._add_deployment_locked = AsyncMock(
return_value=ReconcileOutcome(still_desired=frozenset(), live_after=frozenset())
)
mock_prisma = MagicMock()
mock_logging = MagicMock()
@ -463,7 +466,9 @@ class TestClearCache:
mock_router.complexity_routers = {"db-complexity-router": MagicMock(), "config-router": MagicMock()}
mock_config = MagicMock()
mock_config.add_deployment = AsyncMock(return_value=True)
mock_config._add_deployment_locked = AsyncMock(
return_value=ReconcileOutcome(still_desired=frozenset(), live_after=frozenset())
)
mock_prisma = MagicMock()
mock_logging = MagicMock()
@ -491,8 +496,8 @@ class TestClearCache:
assert "config-router" in mock_router.auto_routers
assert "config-router" in mock_router.complexity_routers
# Should have called add_deployment to reload DB models
mock_config.add_deployment.assert_called_once_with(
# Should have called the already-locked reload to restore DB models
mock_config._add_deployment_locked.assert_called_once_with(
prisma_client=mock_prisma, proxy_logging_obj=mock_logging
)
@ -534,7 +539,9 @@ class TestClearCachePreservesConfigRouters:
}
mock_config = MagicMock()
mock_config.add_deployment = AsyncMock(return_value=True)
mock_config._add_deployment_locked = AsyncMock(
return_value=ReconcileOutcome(still_desired=frozenset(), live_after=frozenset())
)
with (
patch("litellm.proxy.proxy_server.llm_router", mock_router),
@ -574,7 +581,9 @@ class TestClearCachePreservesConfigRouters:
mock_router.complexity_routers = {"shared-name": MagicMock()}
mock_config = MagicMock()
mock_config.add_deployment = AsyncMock(return_value=True)
mock_config._add_deployment_locked = AsyncMock(
return_value=ReconcileOutcome(still_desired=frozenset(), live_after=frozenset())
)
with (
patch("litellm.proxy.proxy_server.llm_router", mock_router),
@ -616,7 +625,9 @@ class TestClearCachePreservesConfigRouters:
mock_router.adaptive_routers = {"a1": MagicMock()}
mock_config = MagicMock()
mock_config.add_deployment = AsyncMock(return_value=True)
mock_config._add_deployment_locked = AsyncMock(
return_value=ReconcileOutcome(still_desired=frozenset(), live_after=frozenset())
)
with (
patch("litellm.proxy.proxy_server.llm_router", mock_router),
@ -839,7 +850,9 @@ class TestUpdateModel:
),
patch(
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
new=AsyncMock(return_value=None),
new=AsyncMock(
return_value=ReconcileOutcome(still_desired=None, live_after=None)
),
) as mock_clear_cache,
):
await update_model(
@ -1885,7 +1898,9 @@ class TestAddAndDeleteModelLifecycle:
mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row)
mock_proxy_config = MagicMock()
mock_proxy_config.add_deployment = AsyncMock()
mock_proxy_config._add_deployment_locked = AsyncMock(
return_value=ReconcileOutcome(still_desired=frozenset(), live_after=frozenset())
)
mock_router = MagicMock()
mock_router.delete_deployment = MagicMock()
@ -3223,7 +3238,9 @@ class TestPatchModelBlockedAuthGate:
),
patch(
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
new=AsyncMock(return_value=None),
new=AsyncMock(
return_value=ReconcileOutcome(still_desired=None, live_after=None)
),
),
):
result = await patch_model(
@ -3381,6 +3398,147 @@ class TestWriteSurfacesReloadDrop:
)
class TestConcurrentModelWritesDoNotEvictEachOther:
"""Two model writes racing on one pod must not un-serve each other's deployments,
and neither may report the other's in-flight reload as damage of its own.
The reconcile is a read-modify-write of the shared ``llm_router`` global: read the db
into a snapshot, then make the router match that snapshot. Unserialized, the request
holding the older snapshot deletes the deployment the newer one just added, because
_delete_deployment evicts every live id absent from the snapshot it was handed. The
row survives in the db, so the damage is invisible there -- the pod just stops
serving a model it was told to serve.
"""
@pytest.mark.asyncio
async def test_reconciles_serialize_so_no_stale_snapshot_can_evict(self, monkeypatch):
"""MODEL_RECONCILE_LOCK admits one reconcile at a time.
The fake body awaits, which is the whole point: without the lock the gather below
parks all five inside the critical section at that await and observed depth goes
to 5. Asserting depth never exceeds 1 is what pins the fix -- deleting the
`async with` makes this fail rather than merely getting slower.
"""
import asyncio
from litellm.proxy._types import ReconcileOutcome
from litellm.proxy.proxy_server import ProxyConfig
depth = 0
observed_max = 0
async def fake_locked(self, **kwargs):
nonlocal depth, observed_max
depth += 1
observed_max = max(observed_max, depth)
await asyncio.sleep(0)
depth -= 1
return ReconcileOutcome(still_desired=frozenset(), live_after=frozenset())
monkeypatch.setattr(ProxyConfig, "_add_deployment_locked", fake_locked)
config = ProxyConfig()
await asyncio.gather(
*[
config.add_deployment(prisma_client=MagicMock(), proxy_logging_obj=MagicMock())
for _ in range(5)
]
)
assert observed_max == 1
@pytest.mark.asyncio
async def test_clear_cache_reloads_under_the_lock_without_deadlocking(self, monkeypatch):
"""clear_cache un-serves every db model before reloading, so it has to hold the
lock across the pair -- and therefore must call the already-locked reload.
asyncio.Lock is not reentrant: routing this back through the public
add_deployment would block forever on a lock this coroutine already owns, taking
every model write on the pod down with it. The timeout is the assertion.
"""
import asyncio
import litellm
from litellm.proxy._types import ReconcileOutcome
from litellm.proxy.management_endpoints.model_management_endpoints import clear_cache
from litellm.proxy.proxy_server import ProxyConfig
live_router = litellm.Router(
model_list=[
{
"model_name": "gpt-4o",
"litellm_params": {"model": "gpt-4o"},
"model_info": {"id": "m-db", "db_model": True},
}
]
)
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", live_router)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", MagicMock())
async def fake_locked(self, **kwargs):
return ReconcileOutcome(still_desired=frozenset({"m-db"}), live_after=frozenset({"m-db"}))
monkeypatch.setattr(ProxyConfig, "_add_deployment_locked", fake_locked)
outcome = await asyncio.wait_for(clear_cache(), timeout=5)
assert outcome.still_desired == frozenset({"m-db"})
assert outcome.live_after == frozenset({"m-db"})
def test_verdict_trusts_the_lock_captured_snapshot_over_a_live_reread(self, monkeypatch):
"""Given live_after, the verdict judges the router as it stood when the reload
finished -- not as it stands now.
Re-reading here would sample the router after the lock was released, which is
exactly where the next writer's clear_cache has every db model deleted and not
yet re-added. That hole is another request's in-flight state; blaming this
request's reload for it is the 500 that made concurrent model creates fail.
"""
import litellm
from litellm.proxy._types import ProxyException
from litellm.proxy.management_endpoints.model_management_endpoints import (
raise_if_reload_degraded_serving,
reload_serving_verdict,
)
# The router as another writer's clear_cache leaves it mid-wipe: db models gone.
mid_wipe_router = litellm.Router(model_list=[])
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mid_wipe_router)
healthy_after_reload = frozenset({"m-live", "m-neighbour"})
_, collateral = reload_serving_verdict(
before=frozenset({"m-live", "m-neighbour"}),
written_models=[("m-live", None)],
written_must_serve=True,
still_desired=healthy_after_reload,
live_after=healthy_after_reload,
)
assert collateral == ()
assert (
raise_if_reload_degraded_serving(
before=frozenset({"m-live", "m-neighbour"}),
written_models=[("m-live", None)],
action="create",
still_desired=healthy_after_reload,
live_after=healthy_after_reload,
)
is None
)
# Same inputs, no lock-captured snapshot: the mid-wipe router is read live and
# the neighbour looks like collateral. This is the pre-fix behaviour, kept to
# show the parameter is what carries the difference.
with pytest.raises(ProxyException, match="m-neighbour"):
raise_if_reload_degraded_serving(
before=frozenset({"m-live", "m-neighbour"}),
written_models=[("m-live", None)],
action="create",
still_desired=healthy_after_reload,
)
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."""

8
tests/test_litellm/test_model_block_unblock.py
View file

@ -7,6 +7,7 @@ from litellm.proxy._types import (
BlockModelRequest,
LitellmUserRoles,
ProxyException,
ReconcileOutcome,
UserAPIKeyAuth,
)
from litellm.types.router import RouterRateLimitError
@ -36,7 +37,12 @@ def _setup_model_block_mocks(monkeypatch, *, updated_blocked: bool):
mock_router = MagicMock()
mock_router.get_model_ids.return_value = [model_id]
mock_clear_cache = AsyncMock(return_value=None)
# No reconcile ran in these tests, so both fields are None and the verdict falls
# back to reading the router live -- which is what the get_model_ids side_effects
# below drive.
mock_clear_cache = AsyncMock(
return_value=ReconcileOutcome(still_desired=None, live_after=None)
)
mock_audit_log = AsyncMock(return_value=None)
monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", mock_prisma_client)