From 6553bc89570b89020753e8aeed9f25df995661ae Mon Sep 17 00:00:00 2001 From: Yuneng Jiang Date: Wed, 12 Aug 2026 10:17:21 -0700 Subject: [PATCH] fix(proxy): serialize model reconciles so concurrent writes stop evicting each other 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. --- litellm/proxy/_types.py | 23 ++- .../model_management_endpoints.py | 161 +++++++++------- litellm/proxy/proxy_server.py | 48 ++++- .../test_model_management_endpoints.py | 182 ++++++++++++++++-- .../test_litellm/test_model_block_unblock.py | 8 +- 5 files changed, 338 insertions(+), 84 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index fa89df39c5f..ccd5bbb7b57 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -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. diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index c2087005863..ec756573abc 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -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) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 90daaaeae6b..ef431911c85 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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, 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 454849d6430..23510945542 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 @@ -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.""" diff --git a/tests/test_litellm/test_model_block_unblock.py b/tests/test_litellm/test_model_block_unblock.py index ff66bedf0dc..da63ed4a95a 100644 --- a/tests/test_litellm/test_model_block_unblock.py +++ b/tests/test_litellm/test_model_block_unblock.py @@ -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)