mirror of
https://github.com/BerriAI/litellm.git
synced 2026-08-28 05:25:59 +00:00
fix(router): clean pattern_router state on upsert/delete (#29601)
* fix(router): clean pattern_router state on upsert/delete PatternMatchRouter.add_pattern was append-only, and neither Router.upsert_deployment nor Router.delete_deployment removed the existing entry. Rotated-out api_keys stayed in the routing rotation for wildcard deployments (model_name with `*`) until proxy restart, silently defeating key rotation as an admin operation. The same leak applied to provider_default_deployment_ids and per-team pattern routers, and the patterns list grew unboundedly on every edit * test(router): direct unit tests for _remove_deployment_from_wildcard_state router_code_coverage.py greps test files for AST Call nodes and flagged the helper as untested because the existing coverage only exercised it transitively through upsert/delete. Adds two direct tests that pin the helper's contract (cleans across global pattern router, per-team routers with empty-router pop, and provider_default_deployment_ids; noop on falsy model_id) * fix(router): address Greptile review on pattern_router cleanup Widen PatternMatchRouter.remove_deployment annotation to Optional[str]; the implementation already handles None via the falsy guard and the unit test exercises it directly. Move _remove_deployment_from_wildcard_state up one level in upsert_deployment so it runs whenever the prior deployment is on the router, not only when the model_id is present in the fast-mapping index. The scenario is currently unreachable (get_deployment shares the same index), but the cleanup is idempotent so this is defensive against any future divergence between those code paths. * fix(router): widen _remove_deployment_from_wildcard_state to Optional[str] Moving the call out of the inner `deployment_id in deployment_fast_mapping` block in the previous commit lost mypy's narrowing of `deployment_id` from Optional[str] to str, tripping the lint CI. The helper already handles None via its falsy guard, so widening the annotation matches the actual contract. * fix(router): make delete_deployment wildcard cleanup symmetric with upsert After the previous commit moved _remove_deployment_from_wildcard_state out of the inner index-map guard in upsert_deployment, delete_deployment was still calling it only inside `if deployment_idx is not None`. Greptile flagged the asymmetry: under a desynced index_map, delete would silently leave the stale wildcard credential in pattern_router. Moves the cleanup call to the top of the try block, mirroring the upsert path. Cleanup is idempotent so the change is a no-op on the happy path. Adds a regression test that simulates the desync by removing the entry from model_id_to_deployment_index_map and asserts delete still clears pattern_router.
This commit is contained in:
Aarkin Karnik
GitHub
parent
b11e59f3a3
commit
ad4e6e2395
4 changed files with 380 additions and 0 deletions
|
|
@ -8855,6 +8855,8 @@ class Router:
|
|||
model_id=deployment_id, removal_idx=removal_idx
|
||||
)
|
||||
|
||||
self._remove_deployment_from_wildcard_state(model_id=deployment_id)
|
||||
|
||||
# if the model_id is not in router
|
||||
self.add_deployment(deployment=deployment)
|
||||
return deployment
|
||||
|
|
@ -8881,6 +8883,9 @@ class Router:
|
|||
deployment_idx = self.model_id_to_deployment_index_map[id]
|
||||
|
||||
try:
|
||||
# Idempotent and symmetric with upsert_deployment, so a desynced
|
||||
# index_map cannot leave stale wildcard credentials behind.
|
||||
self._remove_deployment_from_wildcard_state(model_id=id)
|
||||
if deployment_idx is not None:
|
||||
# Pop the item from the list first
|
||||
item = self.model_list.pop(deployment_idx)
|
||||
|
|
@ -8898,6 +8903,28 @@ class Router:
|
|||
except Exception:
|
||||
return None
|
||||
|
||||
def _remove_deployment_from_wildcard_state(self, model_id: Optional[str]) -> None:
|
||||
"""
|
||||
Drop every reference to model_id from the wildcard-routing data
|
||||
structures. Without this, upsert/delete leaves stale credentials in
|
||||
pattern_router which silently defeats key rotation for wildcard
|
||||
deployments.
|
||||
"""
|
||||
if not model_id:
|
||||
return
|
||||
self.pattern_router.remove_deployment(model_id)
|
||||
empty_team_ids: List[str] = []
|
||||
for team_id, team_router in self.team_pattern_routers.items():
|
||||
team_router.remove_deployment(model_id)
|
||||
if not team_router.patterns:
|
||||
empty_team_ids.append(team_id)
|
||||
for team_id in empty_team_ids:
|
||||
del self.team_pattern_routers[team_id]
|
||||
if model_id in self.provider_default_deployment_ids:
|
||||
self.provider_default_deployment_ids = [
|
||||
i for i in self.provider_default_deployment_ids if i != model_id
|
||||
]
|
||||
|
||||
def _get_router_deployment_budget_limiter(
|
||||
self,
|
||||
) -> Optional[RouterBudgetLimiting]:
|
||||
|
|
|
|||
|
|
@ -75,6 +75,34 @@ class PatternMatchRouter:
|
|||
self.patterns[regex] = []
|
||||
self.patterns[regex].append(llm_deployment)
|
||||
|
||||
def remove_deployment(self, model_id: Optional[str]) -> int:
|
||||
"""
|
||||
Remove every stored entry whose model_info.id equals model_id.
|
||||
|
||||
Returns the number of entries removed. A regex whose deployment list
|
||||
becomes empty is dropped so the patterns dict does not grow unboundedly.
|
||||
|
||||
Empty / falsy model_id is a no-op so callers cannot accidentally wipe
|
||||
entries whose model_info.id is missing.
|
||||
"""
|
||||
if not model_id:
|
||||
return 0
|
||||
removed_total = 0
|
||||
for regex in list(self.patterns.keys()):
|
||||
original = self.patterns[regex]
|
||||
filtered = [
|
||||
d for d in original if (d.get("model_info") or {}).get("id") != model_id
|
||||
]
|
||||
removed_here = len(original) - len(filtered)
|
||||
if removed_here == 0:
|
||||
continue
|
||||
removed_total += removed_here
|
||||
if filtered:
|
||||
self.patterns[regex] = filtered
|
||||
else:
|
||||
del self.patterns[regex]
|
||||
return removed_total
|
||||
|
||||
def _pattern_to_regex(self, pattern: str) -> str:
|
||||
"""
|
||||
Convert a wildcard pattern to a regex pattern
|
||||
|
|
|
|||
|
|
@ -395,3 +395,90 @@ def test_wildcard_priority_over_deployment_names():
|
|||
assert (
|
||||
deployments[0]["litellm_params"]["api_base"] == "http://localhost:8081/openai"
|
||||
), f"Expected '*' wildcard deployment (8081), got {deployments[0]['litellm_params']['api_base']}"
|
||||
|
||||
|
||||
def _make_wildcard_entry(model_id: str, api_key: str) -> dict:
|
||||
"""Build a dict shaped like Deployment(...).to_json(exclude_none=True)."""
|
||||
return Deployment(
|
||||
model_name="openai/*",
|
||||
litellm_params=LiteLLM_Params(model="openai/*", api_key=api_key),
|
||||
model_info=ModelInfo(id=model_id),
|
||||
).to_json(exclude_none=True)
|
||||
|
||||
|
||||
def test_remove_deployment_drops_matching_id_only():
|
||||
router = PatternMatchRouter()
|
||||
router.add_pattern("openai/*", _make_wildcard_entry("id-A", "key-A"))
|
||||
router.add_pattern("openai/*", _make_wildcard_entry("id-B", "key-B"))
|
||||
|
||||
removed = router.remove_deployment("id-A")
|
||||
|
||||
assert removed == 1
|
||||
survivors = router.patterns["openai/(.*)"]
|
||||
assert len(survivors) == 1
|
||||
assert survivors[0]["model_info"]["id"] == "id-B"
|
||||
assert survivors[0]["litellm_params"]["api_key"] == "key-B"
|
||||
|
||||
|
||||
def test_remove_deployment_drops_regex_key_when_list_empties():
|
||||
router = PatternMatchRouter()
|
||||
router.add_pattern("openai/*", _make_wildcard_entry("id-A", "key-A"))
|
||||
|
||||
router.remove_deployment("id-A")
|
||||
|
||||
assert router.patterns == {}
|
||||
|
||||
|
||||
def test_remove_deployment_spans_multiple_regexes():
|
||||
router = PatternMatchRouter()
|
||||
router.add_pattern("openai/*", _make_wildcard_entry("id-A", "key-A"))
|
||||
router.add_pattern("anthropic/*", _make_wildcard_entry("id-A", "key-A"))
|
||||
router.add_pattern("openai/*", _make_wildcard_entry("id-B", "key-B"))
|
||||
|
||||
removed = router.remove_deployment("id-A")
|
||||
|
||||
assert removed == 2
|
||||
assert "anthropic/(.*)" not in router.patterns
|
||||
assert [d["model_info"]["id"] for d in router.patterns["openai/(.*)"]] == ["id-B"]
|
||||
|
||||
|
||||
def test_remove_deployment_noop_for_unknown_id():
|
||||
router = PatternMatchRouter()
|
||||
router.add_pattern("openai/*", _make_wildcard_entry("id-A", "key-A"))
|
||||
before = {k: list(v) for k, v in router.patterns.items()}
|
||||
|
||||
removed = router.remove_deployment("id-DOES-NOT-EXIST")
|
||||
|
||||
assert removed == 0
|
||||
assert router.patterns == before
|
||||
|
||||
|
||||
def test_remove_deployment_with_falsy_id_is_noop_even_when_entries_have_no_id():
|
||||
"""
|
||||
An entry without model_info.id has '' / None as its id. remove_deployment('')
|
||||
must NOT match those entries; otherwise a stray empty-string call would
|
||||
wipe every id-less wildcard deployment in the router.
|
||||
"""
|
||||
router = PatternMatchRouter()
|
||||
router.patterns["openai/(.*)"] = [
|
||||
{"model_name": "openai/*", "litellm_params": {"api_key": "key-X"}}
|
||||
]
|
||||
|
||||
for falsy in ("", None):
|
||||
removed = router.remove_deployment(falsy)
|
||||
assert removed == 0
|
||||
assert len(router.patterns["openai/(.*)"]) == 1
|
||||
|
||||
|
||||
def test_remove_deployment_tolerates_missing_model_info():
|
||||
router = PatternMatchRouter()
|
||||
router.patterns["openai/(.*)"] = [
|
||||
{"model_name": "openai/*", "litellm_params": {"api_key": "X"}},
|
||||
_make_wildcard_entry("id-A", "key-A"),
|
||||
]
|
||||
|
||||
router.remove_deployment("id-A")
|
||||
|
||||
survivors = router.patterns["openai/(.*)"]
|
||||
assert len(survivors) == 1
|
||||
assert "model_info" not in survivors[0]
|
||||
|
|
|
|||
|
|
@ -4756,3 +4756,241 @@ def test_is_deployment_blocked_static_helper_reflects_blocked_flag():
|
|||
)
|
||||
is True
|
||||
)
|
||||
|
||||
|
||||
def _build_wildcard_router(api_key: str, model_id: str = "wildcard-deployment-1"):
|
||||
from litellm.router import Deployment
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "openai/*",
|
||||
"litellm_params": {"model": "openai/*", "api_key": api_key},
|
||||
"model_info": {"id": model_id},
|
||||
},
|
||||
],
|
||||
)
|
||||
return router, Deployment
|
||||
|
||||
|
||||
def test_upsert_wildcard_deployment_removes_stale_api_key():
|
||||
"""
|
||||
Regression: rotating the api_key on a wildcard deployment must NOT leave
|
||||
the previous credential in pattern_router. Before the fix, the old entry
|
||||
survived in pattern_router.patterns[regex] and was still selected on
|
||||
roughly half of wildcard traffic.
|
||||
"""
|
||||
router, Deployment = _build_wildcard_router(api_key="OLD-KEY")
|
||||
model_id = "wildcard-deployment-1"
|
||||
|
||||
router.upsert_deployment(
|
||||
Deployment(
|
||||
model_name="openai/*",
|
||||
litellm_params={"model": "openai/*", "api_key": "NEW-KEY"},
|
||||
model_info={"id": model_id},
|
||||
)
|
||||
)
|
||||
|
||||
entries = router.pattern_router.patterns["openai/(.*)"]
|
||||
assert (
|
||||
len(entries) == 1
|
||||
), f"pattern_router accumulated stale entries: {len(entries)} present"
|
||||
assert entries[0]["litellm_params"]["api_key"] == "NEW-KEY"
|
||||
assert all(
|
||||
e["litellm_params"]["api_key"] != "OLD-KEY" for e in entries
|
||||
), "Rotated-out api_key is still resident in pattern_router"
|
||||
|
||||
|
||||
def test_upsert_wildcard_deployment_is_idempotent_in_pattern_router():
|
||||
"""
|
||||
Repeated upserts of the same wildcard deployment must not grow the
|
||||
pattern_router unboundedly.
|
||||
"""
|
||||
router, Deployment = _build_wildcard_router(api_key="key-0")
|
||||
model_id = "wildcard-deployment-1"
|
||||
|
||||
for i in range(1, 6):
|
||||
router.upsert_deployment(
|
||||
Deployment(
|
||||
model_name="openai/*",
|
||||
litellm_params={"model": "openai/*", "api_key": f"key-{i}"},
|
||||
model_info={"id": model_id},
|
||||
)
|
||||
)
|
||||
|
||||
entries = router.pattern_router.patterns["openai/(.*)"]
|
||||
assert len(entries) == 1
|
||||
assert entries[0]["litellm_params"]["api_key"] == "key-5"
|
||||
assert router.provider_default_deployment_ids.count(model_id) == 1
|
||||
|
||||
|
||||
def test_delete_wildcard_deployment_clears_pattern_router():
|
||||
router, _ = _build_wildcard_router(api_key="OLD-KEY")
|
||||
model_id = "wildcard-deployment-1"
|
||||
|
||||
assert router.pattern_router.patterns # sanity
|
||||
|
||||
router.delete_deployment(id=model_id)
|
||||
|
||||
assert router.pattern_router.patterns == {}
|
||||
assert model_id not in router.provider_default_deployment_ids
|
||||
|
||||
|
||||
def test_upsert_wildcard_to_concrete_model_removes_old_wildcard_entry():
|
||||
"""
|
||||
Changing model_name from a wildcard to a concrete name must drop the
|
||||
old wildcard entry. Otherwise traffic to openai/<anything> still routes
|
||||
to a model the operator believes they renamed.
|
||||
"""
|
||||
router, Deployment = _build_wildcard_router(api_key="key-A")
|
||||
model_id = "wildcard-deployment-1"
|
||||
|
||||
router.upsert_deployment(
|
||||
Deployment(
|
||||
model_name="openai/gpt-4o",
|
||||
litellm_params={"model": "openai/gpt-4o", "api_key": "key-A"},
|
||||
model_info={"id": model_id},
|
||||
)
|
||||
)
|
||||
|
||||
assert router.pattern_router.patterns == {}
|
||||
assert model_id not in router.provider_default_deployment_ids
|
||||
|
||||
|
||||
def test_upsert_team_wildcard_deployment_does_not_leak_old_key():
|
||||
from litellm.router import Deployment
|
||||
|
||||
team_id = "team-acme"
|
||||
model_id = "team-wildcard-1"
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "openai/team-foo-*",
|
||||
"litellm_params": {"model": "openai/*", "api_key": "OLD-KEY"},
|
||||
"model_info": {
|
||||
"id": model_id,
|
||||
"team_id": team_id,
|
||||
"team_public_model_name": "team-foo-*",
|
||||
},
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
assert team_id in router.team_pattern_routers
|
||||
assert router.team_pattern_routers[team_id].patterns
|
||||
|
||||
router.upsert_deployment(
|
||||
Deployment(
|
||||
model_name="openai/team-foo-*",
|
||||
litellm_params={"model": "openai/*", "api_key": "NEW-KEY"},
|
||||
model_info={
|
||||
"id": model_id,
|
||||
"team_id": team_id,
|
||||
"team_public_model_name": "team-foo-*",
|
||||
},
|
||||
)
|
||||
)
|
||||
|
||||
team_router = router.team_pattern_routers[team_id]
|
||||
assert len(team_router.patterns) == 1
|
||||
team_entries = next(iter(team_router.patterns.values()))
|
||||
assert len(team_entries) == 1
|
||||
assert team_entries[0]["litellm_params"]["api_key"] == "NEW-KEY"
|
||||
|
||||
|
||||
def test_delete_team_wildcard_removes_empty_team_router():
|
||||
"""
|
||||
When the last wildcard deployment for a team is deleted, its dedicated
|
||||
team_pattern_router entry should be discarded too, not left as an empty
|
||||
PatternMatchRouter holding the team_id forever.
|
||||
"""
|
||||
team_id = "team-acme"
|
||||
model_id = "team-wildcard-1"
|
||||
|
||||
router = litellm.Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "openai/team-foo-*",
|
||||
"litellm_params": {"model": "openai/*", "api_key": "key-A"},
|
||||
"model_info": {
|
||||
"id": model_id,
|
||||
"team_id": team_id,
|
||||
"team_public_model_name": "team-foo-*",
|
||||
},
|
||||
},
|
||||
],
|
||||
)
|
||||
|
||||
router.delete_deployment(id=model_id)
|
||||
|
||||
assert team_id not in router.team_pattern_routers
|
||||
|
||||
|
||||
def test_remove_deployment_from_wildcard_state_cleans_all_three_structures():
|
||||
"""
|
||||
Pin the contract of Router._remove_deployment_from_wildcard_state: a single
|
||||
call must strip model_id from the global pattern router, from every per-team
|
||||
pattern router (dropping the team entry when its router empties out), and
|
||||
from provider_default_deployment_ids. Exercised directly so the helper's
|
||||
behavior is locked in independently of the upsert/delete code paths.
|
||||
"""
|
||||
from litellm.router_utils.pattern_match_deployments import PatternMatchRouter
|
||||
|
||||
router, _ = _build_wildcard_router(api_key="key-A")
|
||||
model_id = "wildcard-deployment-1"
|
||||
|
||||
router.team_pattern_routers["team-x"] = PatternMatchRouter()
|
||||
router.team_pattern_routers["team-x"].add_pattern(
|
||||
"openai/team-x-*",
|
||||
{
|
||||
"model_name": "openai/team-x-*",
|
||||
"litellm_params": {"model": "openai/*", "api_key": "key-A"},
|
||||
"model_info": {"id": model_id},
|
||||
},
|
||||
)
|
||||
assert router.pattern_router.patterns
|
||||
assert "team-x" in router.team_pattern_routers
|
||||
assert model_id in router.provider_default_deployment_ids
|
||||
|
||||
router._remove_deployment_from_wildcard_state(model_id=model_id)
|
||||
|
||||
assert router.pattern_router.patterns == {}
|
||||
assert "team-x" not in router.team_pattern_routers
|
||||
assert model_id not in router.provider_default_deployment_ids
|
||||
|
||||
|
||||
def test_remove_deployment_from_wildcard_state_is_noop_for_empty_id():
|
||||
"""
|
||||
A falsy model_id must not touch any wildcard-routing state; otherwise an
|
||||
accidental empty-string call could wipe deployments that lack model_info.id.
|
||||
"""
|
||||
router, _ = _build_wildcard_router(api_key="key-A")
|
||||
snapshot_patterns = {k: list(v) for k, v in router.pattern_router.patterns.items()}
|
||||
snapshot_ids = list(router.provider_default_deployment_ids)
|
||||
|
||||
router._remove_deployment_from_wildcard_state(model_id="")
|
||||
|
||||
assert router.pattern_router.patterns == snapshot_patterns
|
||||
assert router.provider_default_deployment_ids == snapshot_ids
|
||||
|
||||
|
||||
def test_delete_deployment_cleans_wildcard_state_even_when_index_is_desynced():
|
||||
"""
|
||||
Symmetric with the upsert path: delete_deployment must clean wildcard
|
||||
state from the model_id regardless of whether the fast-mapping index
|
||||
still knows about it. Simulates index corruption / partial-failure by
|
||||
removing the entry from model_id_to_deployment_index_map while leaving
|
||||
pattern_router intact, then calls delete and asserts pattern_router is
|
||||
cleaned.
|
||||
"""
|
||||
router, _ = _build_wildcard_router(api_key="key-A")
|
||||
model_id = "wildcard-deployment-1"
|
||||
|
||||
assert router.pattern_router.patterns
|
||||
del router.model_id_to_deployment_index_map[model_id]
|
||||
|
||||
router.delete_deployment(id=model_id)
|
||||
|
||||
assert router.pattern_router.patterns == {}
|
||||
assert model_id not in router.provider_default_deployment_ids
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue