mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
32535987e8
* 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.
* fix(tests): return a ReconcileOutcome from the PTU test's add_deployment mock
test_ptu_model_settings.py stubs proxy_config.add_deployment with
AsyncMock(return_value=None). Now that add_deployment returns a
ReconcileOutcome, add_new_model reads .still_desired off that None and
the two PTU gate tests fail with "'NoneType' object has no attribute
'still_desired'".
Return ReconcileOutcome(still_desired=None, live_after=None), matching
the other reconcile mocks. Both fields None means no reconcile state was
captured, so the serving verdict falls back to reading the router live,
which is what the test's mock_router already drives — the PTU assertions
are unchanged.
Two sibling test files were updated for this in the parent commit; this
one was missed because the local env cannot collect four modules under
tests/test_litellm/proxy (prisma generate artifacts), so the full shard
only ran in CI.
Also applies ruff format to proxy_server.py: the new add_deployment
wrapper's single call fits on one line under the project's line length.
* fix(proxy): lock the delete evictions and stop clear_cache wiping deployments
Two follow-ups to MODEL_RECONCILE_LOCK, both found by review.
1. delete_model and delete_team_models evict from llm_router directly,
outside the lock. The db row is gone by then, but a reconcile that
snapshotted the db BEFORE the delete still lists that id as desired and
upserts the deployment straight back, so the pod keeps serving a model
the database no longer has until some later reconcile notices. Taking
the lock orders the eviction after any in-flight reconcile's re-add.
Both new tests fail without the lock ("did not wait for
MODEL_RECONCILE_LOCK") and pass with it.
2. clear_cache no longer wipes deployments. It used to delete_deployment()
every db model before the reload restored them, which left the router
serving ZERO db models for the entire width of the reload -- every
inference request landing in that window fell into a real hole, and
serializing reconciles made the aggregate outage additive rather than
overlapping. The wipe was also redundant: _delete_deployment evicts
exactly the ids the db no longer lists, and upsert_deployment
pops-and-re-adds a deployment whose params changed while no-opping one
that did not, so the reconcile converges to the same state on its own.
Every mutation is visible to that comparison (blocked, and updated_at
for premium, are written into model_info).
The auto-router pops are NOT redundant and stay: they are keyed by
model_name, which no deployment-id reconcile touches.
The new tests patch their own lock rather than contending the module-level
one: asyncio.Lock binds to the event loop of its first contended acquire
and raises on every other loop after that, which would poison the next
asyncio test in the process. The proxy has a single event loop for its
lifetime so this is test-only, but it is a trap worth naming for whoever
writes the next concurrency test here.
* fix(proxy): scope the clear_cache wipe to auto-router deployments
Review caught a regression in the previous commit. Dropping the wipe
entirely stranded every db-backed auto-router on the pod.
The strategy registries (auto_routers, complexity_routers,
adaptive_routers, quality_routers) are keyed by model_name, which no
deployment-id reconcile touches, so clear_cache pops them and relies on
the reload to rebuild them. But the rebuild only happens on the ADD path:
Router.upsert_deployment returns early when a deployment is unchanged and
never reaches add_deployment -> _add_deployment ->
init_auto_router_deployment, which is what repopulates them. With the wipe
gone the deployment was always unchanged, so the pop was permanent: ANY
unrelated model write -- a team admin patching one team-owned model --
left every db-backed auto, complexity, adaptive and quality router
unroutable across tenants until a restart.
Restore the wipe for exactly the auto_router/* db deployments, whose
strategy entries are the ones being popped. Deleting them forces upsert
down the add path so both the deployment and its strategy entry come back.
Ordinary db models stay un-wiped, which is the point of the previous
commit: wiping them un-served every db model for the width of the reload,
and the reconcile converges without it.
test_clear_cache_wipes_auto_routers_but_leaves_ordinary_db_models pins
both halves against each other, since fixing either one naively breaks the
other. Both clear_cache tests fail with the pop-without-delete version.
* refactor(clear_cache): fold auto-router wipe into the classification pass
The auto-router scoping added in 5deddfd introduced two new mutable-collection
constructions, pushing LIT002 five over its budget ceiling.
Rather than suppress, do the work in the single pass that already walks
current_models: detect and delete the auto_router/* db deployments while
classifying, accumulating names into a set that replaces the old
db_router_deployments comprehension. Net-zero LIT002, same behaviour.
Comment updated to describe where the wipe actually happens now.
253 lines
8.7 KiB
Python
253 lines
8.7 KiB
Python
from unittest.mock import AsyncMock, MagicMock
|
|
|
|
import pytest
|
|
|
|
import litellm
|
|
from litellm.proxy._types import (
|
|
BlockModelRequest,
|
|
LitellmUserRoles,
|
|
ProxyException,
|
|
ReconcileOutcome,
|
|
UserAPIKeyAuth,
|
|
)
|
|
from litellm.types.router import RouterRateLimitError
|
|
|
|
|
|
def _setup_model_block_mocks(monkeypatch, *, updated_blocked: bool):
|
|
model_id = "model-123"
|
|
|
|
existing_row = MagicMock()
|
|
existing_row.model_dump.return_value = {
|
|
"model_name": "gpt-4o",
|
|
"litellm_params": {"model": "openai/gpt-4o"},
|
|
"model_info": {"id": model_id},
|
|
}
|
|
|
|
updated_row = MagicMock()
|
|
updated_row.model_id = model_id
|
|
updated_row.blocked = updated_blocked
|
|
|
|
model_table = MagicMock()
|
|
model_table.find_unique = AsyncMock(return_value=existing_row)
|
|
model_table.update = AsyncMock(return_value=updated_row)
|
|
|
|
mock_prisma_client = MagicMock()
|
|
mock_prisma_client.db.litellm_proxymodeltable = model_table
|
|
|
|
mock_router = MagicMock()
|
|
mock_router.get_model_ids.return_value = [model_id]
|
|
|
|
# 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)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", True)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", mock_router)
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.litellm_proxy_admin_name", "admin")
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.management_endpoints.model_management_endpoints.clear_cache",
|
|
mock_clear_cache,
|
|
)
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.management_endpoints.model_management_endpoints.create_object_audit_log",
|
|
mock_audit_log,
|
|
)
|
|
|
|
return model_id, model_table, updated_row, mock_clear_cache, mock_audit_log
|
|
|
|
|
|
def _proxy_admin() -> UserAPIKeyAuth:
|
|
return UserAPIKeyAuth(
|
|
user_id="admin",
|
|
user_role=LitellmUserRoles.PROXY_ADMIN,
|
|
api_key="sk-admin",
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_model_block_endpoint_sets_blocked_true(monkeypatch):
|
|
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
|
block_model,
|
|
)
|
|
|
|
model_id, model_table, updated_row, mock_clear_cache, mock_audit_log = (
|
|
_setup_model_block_mocks(monkeypatch, updated_blocked=True)
|
|
)
|
|
|
|
result = await block_model(
|
|
data=BlockModelRequest(model_id=model_id),
|
|
http_request=MagicMock(),
|
|
user_api_key_dict=_proxy_admin(),
|
|
litellm_changed_by="operator@example.com",
|
|
)
|
|
|
|
assert result == updated_row
|
|
model_table.update.assert_awaited_once()
|
|
update_kwargs = model_table.update.await_args.kwargs
|
|
assert update_kwargs["where"] == {"model_id": model_id}
|
|
assert update_kwargs["data"]["blocked"] is True
|
|
assert update_kwargs["data"]["updated_by"] == "admin"
|
|
assert "updated_at" in update_kwargs["data"]
|
|
mock_clear_cache.assert_awaited_once_with()
|
|
assert mock_audit_log.call_args.kwargs["action"] == "blocked"
|
|
assert (
|
|
mock_audit_log.call_args.kwargs["litellm_changed_by"] == "operator@example.com"
|
|
)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_model_unblock_endpoint_sets_blocked_false(monkeypatch):
|
|
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
|
unblock_model,
|
|
)
|
|
|
|
model_id, model_table, updated_row, mock_clear_cache, mock_audit_log = (
|
|
_setup_model_block_mocks(monkeypatch, updated_blocked=False)
|
|
)
|
|
|
|
result = await unblock_model(
|
|
data=BlockModelRequest(model_id=model_id),
|
|
http_request=MagicMock(),
|
|
user_api_key_dict=_proxy_admin(),
|
|
litellm_changed_by=None,
|
|
)
|
|
|
|
assert result == updated_row
|
|
model_table.update.assert_awaited_once()
|
|
assert model_table.update.await_args.kwargs["data"]["blocked"] is False
|
|
mock_clear_cache.assert_awaited_once_with()
|
|
assert mock_audit_log.call_args.kwargs["action"] == "unblocked"
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_model_block_endpoint_requires_proxy_admin(monkeypatch):
|
|
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
|
block_model,
|
|
)
|
|
|
|
model_id, model_table, _, _, _ = _setup_model_block_mocks(
|
|
monkeypatch, updated_blocked=True
|
|
)
|
|
non_admin = UserAPIKeyAuth(
|
|
user_id="internal-user",
|
|
user_role=LitellmUserRoles.INTERNAL_USER,
|
|
api_key="sk-user",
|
|
)
|
|
|
|
with pytest.raises(ProxyException) as exc_info:
|
|
await block_model(
|
|
data=BlockModelRequest(model_id=model_id),
|
|
http_request=MagicMock(),
|
|
user_api_key_dict=non_admin,
|
|
litellm_changed_by=None,
|
|
)
|
|
|
|
assert exc_info.value.code == "403"
|
|
assert "Only proxy admins" in exc_info.value.message
|
|
model_table.update.assert_not_awaited()
|
|
|
|
|
|
def test_router_returns_no_healthy_deployment_when_model_is_fully_blocked():
|
|
router = litellm.Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "gpt-4o",
|
|
"litellm_params": {"model": "openai/gpt-4o-0"},
|
|
"model_info": {"id": "dep-0", "blocked": True},
|
|
},
|
|
{
|
|
"model_name": "gpt-4o",
|
|
"litellm_params": {"model": "openai/gpt-4o-1"},
|
|
"model_info": {"id": "dep-1", "blocked": True},
|
|
},
|
|
]
|
|
)
|
|
|
|
with pytest.raises(RouterRateLimitError) as exc_info:
|
|
router.get_available_deployment(model="gpt-4o", request_kwargs={})
|
|
|
|
assert "No deployments available for selected model" in str(exc_info.value)
|
|
assert "Passed model=gpt-4o" in str(exc_info.value)
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_route_request_returns_403_when_model_is_fully_blocked(monkeypatch):
|
|
from litellm.proxy.route_llm_request import route_request
|
|
|
|
router = litellm.Router(
|
|
model_list=[
|
|
{
|
|
"model_name": "gpt-4o",
|
|
"litellm_params": {"model": "openai/gpt-4o"},
|
|
"model_info": {"id": "dep-0", "blocked": True},
|
|
}
|
|
]
|
|
)
|
|
monkeypatch.setattr(
|
|
"litellm.proxy.route_llm_request.add_shared_session_to_data",
|
|
AsyncMock(return_value=None),
|
|
)
|
|
|
|
with pytest.raises(litellm.PermissionDeniedError) as exc_info:
|
|
await route_request(
|
|
data={"model": "gpt-4o"},
|
|
llm_router=router,
|
|
user_model=None,
|
|
route_type="acreate_eval",
|
|
)
|
|
|
|
assert exc_info.value.status_code == 403
|
|
assert "Model is blocked" in exc_info.value.message
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_model_block_surfaces_wholesale_reload_failure(monkeypatch):
|
|
"""The write endpoints owe the caller an error when the pod failed to reload at all;
|
|
the DB row is saved but this pod is not serving the change."""
|
|
from litellm.proxy._types import ProxyException
|
|
from litellm.proxy.management_endpoints.model_management_endpoints import block_model
|
|
|
|
model_id, model_table, updated_row, mock_clear_cache, mock_audit_log = _setup_model_block_mocks(
|
|
monkeypatch, updated_blocked=True
|
|
)
|
|
wiped_router = MagicMock()
|
|
wiped_router.get_model_ids.side_effect = [[model_id], []]
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", wiped_router)
|
|
|
|
with pytest.raises(ProxyException, match=model_id):
|
|
await block_model(
|
|
data=BlockModelRequest(model_id=model_id),
|
|
http_request=MagicMock(),
|
|
user_api_key_dict=_proxy_admin(),
|
|
litellm_changed_by="operator@example.com",
|
|
)
|
|
|
|
assert mock_audit_log.call_args.kwargs["object_id"] == model_id
|
|
|
|
|
|
@pytest.mark.asyncio
|
|
async def test_model_block_surfaces_model_dropped_by_reload(monkeypatch):
|
|
"""A reload that completes but drops the written model (ignore_invalid_deployments
|
|
swallowed its re-add) must not produce an unqualified success."""
|
|
from litellm.proxy._types import ProxyException
|
|
from litellm.proxy.management_endpoints.model_management_endpoints import block_model
|
|
|
|
model_id, model_table, updated_row, mock_clear_cache, _ = _setup_model_block_mocks(
|
|
monkeypatch, updated_blocked=True
|
|
)
|
|
dropped_router = MagicMock()
|
|
dropped_router.get_model_ids.return_value = []
|
|
monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", dropped_router)
|
|
|
|
with pytest.raises(ProxyException, match=model_id):
|
|
await block_model(
|
|
data=BlockModelRequest(model_id=model_id),
|
|
http_request=MagicMock(),
|
|
user_api_key_dict=_proxy_admin(),
|
|
litellm_changed_by="operator@example.com",
|
|
)
|