scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-11-13 08:52:02 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 174

fix(management): keep classifier deployments serving during reloads (#51111)

Browse source 
* fix(management): track self-heal for test_moving_the_workspace_under_a_burst_moves_the_classifier_and_logs_every_request_once

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(management): keep classifier deployments serving during reloads

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: mateo <277851410+mateo-berri@users.noreply.github.com>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-11-11 18:48:06 +00:00 • committed by GitHub
parent 3cd48a6ec7
commit 3a91fdfc07
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
10 changed files with 293 additions and 208 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

7
litellm-rust/crates/router-core/src/lib.rs
View file

@ -688,7 +688,7 @@ impl Router {
entries,
resolve_secrets,
|deployment| {
working.push_built(deployment.spec.clone(), add)?;
working.install_built(deployment.spec.clone(), add, install)?;
installed.push(deployment.spec.clone());
core.validate_deployment(working.last().expect("just installed"))
.map_err(Error::from)
@ -698,7 +698,7 @@ impl Router {
if built.is_ok() && replace && path == BuildPath::Create {
working.set_model_names(model_names);
}
self.serve_built(&base, working, installed, add)?;
self.serve_built(&base, working, installed, add, install)?;
built
}
@ -710,6 +710,7 @@ impl Router {
built: Registry,
installed: Vec<DeploymentSpec>,
add: bool,
install: build::Install,
) -> Result<(), ConfigError> {
let _writer = self.core.writer.lock();
let current = self.core.snapshot.load_full();
@ -718,7 +719,7 @@ impl Router {
} else {
let mut next = Registry::clone(&current.registry);
for spec in installed {
next.push_built(spec, add)?;
next.install_built(spec, add, install)?;
}
next
};

75
litellm-rust/crates/router-core/src/registry.rs
View file

@ -626,6 +626,27 @@ impl Registry {
}
}
pub fn install_built(
&mut self,
spec: DeploymentSpec,
add: bool,
install: crate::build::Install,
) -> Result<(), ConfigError> {
if install == crate::build::Install::Upsert {
let id = self.carried(spec.clone())?.id;
if self.by_id.contains_key(&id) {
*self = self.apply(
&ChangeSet {
remove: Some(id),
..ChangeSet::default()
},
None,
)?;
}
}
self.push_built(spec, add)
}
/// Add one deployment a build accepted at the end and index it in place (`add` is
/// `add_deployment`, which also names its group). A build installs into a registry of its
/// own this way, so a list of n deployments builds in O(n) rather than reindexing per
@ -1531,6 +1552,60 @@ mod tests {
assert_eq!((resolution.kind, resolution.group.as_str()), (kind, group));
}
#[rstest]
#[case::empty(false)]
#[case::replacement(true)]
fn installing_an_upsert_replaces_only_the_working_snapshot(#[case] populated: bool) {
let registry = Registry::build(
if populated {
vec![spec(
"a",
json!({"model": "openai/old"}),
json!({"id": "1"}),
)]
} else {
Vec::new()
},
BTreeMap::new(),
Vec::new(),
)
.unwrap();
let mut working = registry.clone();
working
.install_built(
spec("a", json!({"model": "openai/new"}), json!({"id": "1"})),
true,
crate::build::Install::Upsert,
)
.unwrap();
assert_eq!(registry.all().len(), usize::from(populated));
assert_eq!(working.all().len(), 1);
assert_eq!(working.all()[0].raw_params["model"], json!("openai/new"));
assert_eq!(working.host.position_of("1"), Some(0));
}
#[rstest]
fn an_invalid_upsert_keeps_the_previous_deployment() {
let mut registry = Registry::build(
vec![spec(
"a",
json!({"model": "openai/old"}),
json!({"id": "1"}),
)],
BTreeMap::new(),
Vec::new(),
)
.unwrap();
let invalid = spec("a", json!({"model": 42}), json!({"id": "1"}));
assert!(
registry
.install_built(invalid, true, crate::build::Install::Upsert)
.is_err()
);
assert_eq!(registry.all().len(), 1);
assert_eq!(registry.all()[0].raw_params["model"], json!("openai/old"));
}
#[rstest]
fn apply_moves_an_edited_deployment_to_the_end() {
let registry = Registry::build(

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

@ -4043,8 +4043,8 @@ async def clear_cache() -> ReconcileOutcome:
raise_if_reload_degraded_serving.
Runs under DB_RECONCILE_LOCK for its whole extent, not just the reload at the
end, so the auto-router reset and the reload that rebuilds those routers are atomic
to any other reconcile. The inner call is _add_deployment_locked because
end, so concurrent model writes cannot interleave their reconciles.
The inner call is _add_deployment_locked because
add_deployment would re-acquire the same non-reentrant lock and deadlock.
"""
from litellm.proxy.proxy_server import (
@ -4062,75 +4062,22 @@ async def clear_cache() -> ReconcileOutcome:
async with DB_RECONCILE_LOCK:
try:
# Only clear DB models, preserve config models
verbose_proxy_logger.debug("Clearing only DB models, preserving config models")
verbose_proxy_logger.debug("Reconciling DB models, preserving config models")
# Get current models and filter out DB models
current_models: Final = llm_router.registry().all().copy()
config_models: Final = []
db_model_ids: Final = []
db_router_names: Final = set()
for model in as_object_list(current_models) or ():
deployment = decoded_json_object(model)
model_info = decoded_json_object(deployment.get("model_info", {}))
if model_info.get("db_model", False):
db_model_ids.append(model_info.get("id"))
# Auto-router deployments (and only those) are wiped here, in the
# same pass, so the reload rebuilds them -- see the comment below.
deployment_name = deployment.get("model_name")
if isinstance(deployment_name, str) and str(
decoded_json_object(deployment.get("litellm_params", {})).get("model", "")
).startswith("auto_router/"):
db_router_names.add(deployment_name)
router_model_id = model_info.get("id")
if isinstance(router_model_id, str):
llm_router.apply(ChangeSet(delete=[router_model_id]))
else:
# This is a config model, preserved by the reconcile below
config_models.append(model)
# ORDINARY db deployments are deliberately NOT wiped. This used to
# delete_deployment() every db model before the reload put them back, which
# left the router serving ZERO db models for the whole width of the reload
# -- a real data-plane hole that every inference request landing in it fell
# into. It was also redundant for them: the reload's _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 on its own. Every mutation is
# visible to that comparison -- `blocked` and (for premium) `updated_at`
# are written into model_info.
#
# AUTO-ROUTER db deployments are the exception and ARE wiped -- in the
# classification pass above, together with the strategy entries popped
# just below. Their strategy registries are keyed
# by model_name, which no deployment-id reconcile touches, so they have to
# be popped and rebuilt here. 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 the registries.
# Popping without deleting would therefore strip every db-backed auto,
# complexity, adaptive and quality router on this pod and never put it back,
# so ANY unrelated model write would leave them unroutable until a restart.
# Deleting the deployment forces upsert down the add path, which rebuilds
# both the deployment and its strategy entry.
#
# That pass restricts the wipe 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 -- config routers are never re-added by the
# reload (it only reloads db models) and would be permanently unroutable.
# The auto_router/ prefix also covers quality_router/ and adaptive_router/,
# so pop the name from every registry (no-op where absent); a missing
# quality/adaptive entry would otherwise make init raise "already exists"
# on reload and abort it.
for model_name in db_router_names:
llm_router.registry().plugins("auto").pop(model_name, None)
llm_router.registry().plugins("complexity").pop(model_name, None)
llm_router.registry().plugins("adaptive").pop(model_name, None)
llm_router.registry().plugins("quality").pop(model_name, None)
llm_router.registry().plugins("switchyard").pop(model_name, None)
# Reload only DB models. _add_deployment_locked, not add_deployment: this
# coroutine already holds DB_RECONCILE_LOCK and asyncio.Lock is not
# reentrant, so the public wrapper would deadlock against itself.

5
litellm/proxy/proxy_server.py
View file

@ -9024,9 +9024,8 @@ class ProxyConfig:
proxy_logging_obj: ProxyLogging,
) -> ReconcileOutcome:
"""add_deployment's body, minus the locking. DB_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."""
be held. Split out for clear_cache, which already holds the lock and would
deadlock on a re-acquire."""
global llm_router, llm_model_list, master_key, general_settings
still_desired_ids: frozenset[str] | None = None

81
litellm/router_rust/steps/prerouting.py
View file

@ -133,9 +133,13 @@ def _routing_definition_fingerprint(deployment: Deployment) -> str | None:
return deployment_id(deployment.model_name, {"routing_definition": definition})
def _unregister(registry: StrategyRegistry, model_name: str, tags: tuple[str, ...]) -> bool:
def _unregister(
registry: StrategyRegistry, model_name: str, tags: tuple[str, ...], previous: Deployment | None = None
) -> bool:
existing: Final = registry.get(model_name, [])
remaining: Final = [entry for entry in existing if entry.tags != tags]
remaining: Final = [
entry for entry in existing if entry.tags != tags or (previous is not None and entry.deployment != previous)
]
if len(remaining) == len(existing):
return False
if remaining:
@ -217,30 +221,31 @@ class PreRoutingStrategies:
self._hosted._core_auto_routers = {} # pyright: ignore[reportPrivateUsage] # see above
self._hosted._core_switchyard_routers = {} # pyright: ignore[reportPrivateUsage] # see above
def register(self, deployment: Deployment) -> None:
def register(self, deployment: Deployment, *, replace_existing: bool = False) -> None:
"""The end of ``_add_deployment``: an auto, complexity or quality marker registers its router
(adaptive ones wait for ``finalize``)."""
match classify_strategy_router_model(deployment.litellm_params.model):
case "semantic":
self._init_auto_router(deployment)
self._init_auto_router(deployment, replace_existing=replace_existing)
case "complexity":
self._init_complexity_router(deployment)
self._init_complexity_router(deployment, replace_existing=replace_existing)
case "quality":
self._init_quality_router(deployment)
self._init_quality_router(deployment, replace_existing=replace_existing)
case "switchyard":
self._init_switchyard_router(deployment)
self._init_switchyard_router(deployment, replace_existing=replace_existing)
case _:
pass
def release(self, deployment: Deployment) -> None:
def release(self, deployment: Deployment, *, replaced: bool = False) -> None:
"""``_unregister_pre_routing_strategy_for_deployment``."""
if not deployment.litellm_params.model.startswith("auto_router/"):
return
model_name: Final = deployment.model_name
tags: Final = _deployment_tags(deployment)
previous: Final = deployment if replaced else None
for registry in (self.auto, self.complexity, self.quality, self.switchyard):
_ = _unregister(registry, model_name, tags)
if _unregister(self.adaptive, model_name, tags):
_ = _unregister(registry, model_name, tags, previous)
if _unregister(self.adaptive, model_name, tags, previous):
self._sync_adaptive_router_hooks()
def release_entry(self, entry: Mapping[str, object]) -> None:
@ -261,15 +266,6 @@ class PreRoutingStrategies:
):
self.finalize()
def without_capability_limit(self, action: Callable[[], object]) -> None:
"""Run ``action`` without judging the capability ceiling: a rollback re-admits a serving router."""
resolver: Final = self._capability_limit
self._capability_limit = None
try:
_ = action()
finally:
self._capability_limit = resolver
def finalize(self) -> None:
"""``_finalize_adaptive_router_if_configured``: an adaptive router for every adaptive marker,
and the adaptive companion of every complexity router configured with ``adaptive``."""
@ -328,16 +324,30 @@ class PreRoutingStrategies:
# -- building --------------------------------------------------------------------------------
def _register(
self, registry: StrategyRegistry, deployment: Deployment, strategy: object, strategy_label: str
self,
registry: StrategyRegistry,
deployment: Deployment,
strategy: object,
strategy_label: str,
*,
replace_existing: bool = False,
) -> None:
tags: Final = _deployment_tags(deployment)
if _has_registered_strategy(registry, deployment.model_name, tags):
existing: Final = registry.get(deployment.model_name, [])
remaining: Final = [
entry
for entry in existing
if not replace_existing
or entry.deployment is None
or entry.deployment.model_info.id != deployment.model_info.id
]
if any(entry.tags == tags for entry in remaining):
raise ValueError(
f"{strategy_label} deployment {deployment.model_name} with tags {list(tags)} already exists. "
"Please use a different model name or set different tags."
)
registry[deployment.model_name] = [ # rebind-ok: the live registry, edited in place as the Python router does
*registry.get(deployment.model_name, []),
*remaining,
TaggedPreRoutingStrategy(
tags=tags,
strategy=strategy,
@ -346,17 +356,22 @@ class PreRoutingStrategies:
),
]
def _capability_violation(self, capability: GatedAutoRouterCapability) -> str | None:
def _capability_violation(self, capability: GatedAutoRouterCapability, deployment: Deployment) -> str | None:
resolver: Final = self._capability_limit
# `self.auto_router_capability_limit() if ... is not None`: a value that cannot be called raises.
limit: Final = resolver() if resolver is not None else None # pyright: ignore[reportCallIssue] # called as given, as Python calls it
others: Final = count_capability_routers(
(entry for entry in self._hosted._core_model_list if isinstance(entry, Mapping)), # pyright: ignore[reportPrivateUsage, reportUnnecessaryIsInstance] # the live list
(
entry
for entry in self._hosted._core_model_list # pyright: ignore[reportPrivateUsage] # the live list
if isinstance(entry, Mapping) # pyright: ignore[reportUnnecessaryIsInstance] # stored entries are mappings
and as_dict(entry.get("model_info")).get("id") != deployment.model_info.id
),
capability=capability,
)
return capability_limit_violation(capability=capability, held=others + 1, limit=limit) # pyright: ignore[reportArgumentType] # the resolver's limit
def _init_auto_router(self, deployment: Deployment) -> None:
def _init_auto_router(self, deployment: Deployment, *, replace_existing: bool = False) -> None:
from litellm.router_strategy.auto_router.auto_router import AutoRouter
params: Final = deployment.litellm_params
@ -382,15 +397,15 @@ class PreRoutingStrategies:
ctx=self._ctx, # pyright: ignore[reportArgumentType] # the core's RoutingContext
max_input_chars=max_input_chars if max_input_chars is not None else DEFAULT_AUTO_ROUTER_MAX_INPUT_CHARS,
)
self._register(self.auto, deployment, router, "Auto-router")
self._register(self.auto, deployment, router, "Auto-router", replace_existing=replace_existing)
def _init_complexity_router(self, deployment: Deployment) -> None:
def _init_complexity_router(self, deployment: Deployment, *, replace_existing: bool = False) -> None:
from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter
from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig
config: Final = deployment.litellm_params.complexity_router_config
for capability in claimed_capabilities(config):
if (violation := self._capability_violation(capability)) is not None:
if (violation := self._capability_violation(capability, deployment)) is not None:
raise ValueError(violation)
default_model: Final = (
ComplexityRouterConfig.model_validate(config).resolve_default_model(
@ -410,7 +425,7 @@ class PreRoutingStrategies:
ctx=self._ctx, # pyright: ignore[reportArgumentType] # the core's RoutingContext
complexity_router_config=config,
)
self._register(self.complexity, deployment, router, "Complexity-router")
self._register(self.complexity, deployment, router, "Complexity-router", replace_existing=replace_existing)
if router._uses_deployment_pin: # pyright: ignore[reportPrivateUsage] # the Python router reads the same flag
self._ensure_deployment_affinity_callback()
@ -451,7 +466,7 @@ class PreRoutingStrategies:
self._affinity = affinity
self._host.callbacks.add_litellm_callback(affinity)
def _init_quality_router(self, deployment: Deployment) -> None:
def _init_quality_router(self, deployment: Deployment, *, replace_existing: bool = False) -> None:
from litellm.router_strategy.quality_router.quality_router import QualityRouter
config: Final = deployment.litellm_params.quality_router_config
@ -468,9 +483,9 @@ class PreRoutingStrategies:
ctx=self._ctx, # pyright: ignore[reportArgumentType] # the core's RoutingContext
quality_router_config=config,
)
self._register(self.quality, deployment, router, "Quality-router")
self._register(self.quality, deployment, router, "Quality-router", replace_existing=replace_existing)
def _init_switchyard_router(self, deployment: Deployment) -> None:
def _init_switchyard_router(self, deployment: Deployment, *, replace_existing: bool = False) -> None:
from litellm.router_strategy.switchyard_router.config import SwitchyardRouterConfig, build_algorithm
from litellm.router_strategy.switchyard_router.switchyard_router import SwitchyardRouter
from litellm.types.decisions import OpenAIDecisionResponse
@ -494,7 +509,7 @@ class PreRoutingStrategies:
return result
router: Final = SwitchyardRouter(deployment.model_name, config, build_algorithm(config), complete_judge, decide)
self._register(self.switchyard, deployment, router, "Switchyard")
self._register(self.switchyard, deployment, router, "Switchyard", replace_existing=replace_existing)
def _init_adaptive_router(self, deployment: Deployment) -> None:
from litellm.router_strategy.adaptive_router.adaptive_router import AdaptiveRouter

89
litellm/router_rust/surface/registry.py
View file

@ -215,6 +215,7 @@ class _Build:
entries: dict[int, _BuildEntry] = dataclasses.field(default_factory=dict) # mutable-ok: filled in step order
# `add`: the deployment to add, and the entry it is stored as.
deployment: Deployment | None = None
replace_existing: bool = False
stored: UntypedEntry | None = None
# The deployment the build is on, its provider and its data-source values.
pending: Deployment | None = None
@ -1007,12 +1008,13 @@ class HostRegistry(HostModelInfo):
raise RuntimeError("the host registry builds deployments through its Rust core")
for entry in entries:
refuse_wide_numbers(entry)
self._core_model_list.begin_build(replacing=install == "replace")
try:
self.build_runner(entries, build.path, resolve_secrets, install, partial(self._answer_shown, build))
finally:
self._core_model_list.end_build()
self._drop_removed_live_values()
with self._deployment_index_lock:
self._core_model_list.begin_build(replacing=install == "replace")
try:
self.build_runner(entries, build.path, resolve_secrets, install, partial(self._answer_shown, build))
finally:
self._core_model_list.end_build()
self._drop_removed_live_values()
if build.failure is not None:
# router-core served what it built before the entry that raised (`set_model_list`
# cleared the list first), so the router is left as the Python core leaves it.
@ -1220,7 +1222,11 @@ class HostRegistry(HostModelInfo):
deployment: Final = build.deployment
if deployment is None:
raise RuntimeError("an add build holds its deployment")
if deployment.model_info.id and self._core_has_model_id(deployment.model_info.id):
if (
not build.replace_existing
and deployment.model_info.id
and self._core_has_model_id(deployment.model_info.id)
):
return {"skip": True}
declared_params: Final = deployment.litellm_params.model_dump(exclude_none=True)
warn_on_provider_credential_mismatch(model_name=deployment.model_name, litellm_params=declared_params)
@ -1273,7 +1279,9 @@ class HostRegistry(HostModelInfo):
deployment: Final = build.pending
if deployment is None:
raise RuntimeError("an added step follows a validated deployment")
self._register_deployment(deployment, build.provider, build.data_source_values)
self._register_deployment(
deployment, build.provider, build.data_source_values, replace_existing=build.replace_existing
)
build.pending = None # rebind-ok: the build's state, which each step advances
build.data_source_values = () # rebind-ok: the build's state, which each step advances
self._invalidate_model_group_info_cache()
@ -1319,6 +1327,8 @@ class HostRegistry(HostModelInfo):
deployment: Deployment,
custom_llm_provider: str | None,
data_source_values: Sequence[tuple[int, str, str]],
*,
replace_existing: bool,
) -> None:
"""``_add_deployment`` after the provider: its data-source values, pass-through
credentials and pre-routing strategies (router-core registers its names and patterns)."""
@ -1366,7 +1376,7 @@ class HostRegistry(HostModelInfo):
if custom_llm_provider is not None:
self._initialize_deployment_for_pass_through(deployment, custom_llm_provider)
if self.strategies is not None:
self.strategies.register(deployment)
self.strategies.register(deployment, replace_existing=replace_existing)
def _initialize_deployment_for_pass_through(self, deployment: Deployment, custom_llm_provider: str) -> None:
"""Register vertex credentials for a pass-through deployment (``_initialize_deployment_for_pass_through``)."""
@ -1395,7 +1405,10 @@ class HostRegistry(HostModelInfo):
def add_deployment(self, deployment: Deployment) -> Deployment | None:
"""Add one deployment unless its id exists (``add_deployment``), built by router-core."""
build: Final = _Build(path="add", deployment=deployment)
return self._install_deployment(deployment, replace_existing=False)
def _install_deployment(self, deployment: Deployment, *, replace_existing: bool) -> Deployment | None:
build: Final = _Build(path="add", deployment=deployment, replace_existing=replace_existing)
self._run_build(build, [deployment.to_json(exclude_none=True)], resolve_secrets=False, install="upsert")
if build.result is None:
return None
@ -1410,62 +1423,38 @@ class HostRegistry(HostModelInfo):
on_router: Final = self._core_get_deployment(model_id=model_id)
served: Final = without_model_capability_flags(deployment)
try:
if (
on_router is not None
and deployment.model_name == on_router.model_name
and served.litellm_params == on_router.litellm_params
and deployment.model_info == on_router.model_info
):
return None
self._install_deployment(deployment, replace_existing=True)
if on_router is not None:
if (
deployment.model_name == on_router.model_name
and served.litellm_params == on_router.litellm_params
and deployment.model_info == on_router.model_info
):
return None
with self._deployment_index_lock:
removal_idx: Final = self.model_id_to_deployment_index_map.get(model_id)
if removal_idx is not None:
removed: Final = self._core_model_list[removal_idx] # pyright: ignore[reportArgumentType] # raises IndexError for a position past the list, as the pop does
self._remove(model_id)
withdraw_shared_capabilities(
self._host, removed, deployment.to_json(exclude_none=True), self._served_by_live_routers()
)
withdraw_shared_capabilities(
self._host,
on_router.to_json(exclude_none=True),
deployment.to_json(exclude_none=True),
self._served_by_live_routers(),
)
if self.strategies is not None:
self.strategies.release(on_router)
self._on_change()
self.add_deployment(deployment=deployment)
self.strategies.release(on_router, replaced=True)
if self.strategies is not None:
self.strategies.upserted(deployment, on_router)
except Exception as e:
if not self.ignore_invalid_deployments:
raise
verbose_router_logger.warning(
"Error upserting deployment %s (id=%s): %s. Dropping it and continuing with other deployments.",
"Error upserting deployment %s (id=%s): %s. Keeping its previous deployment, if any.",
deployment.model_name,
deployment.model_info.id,
e,
)
self._restore_deployment_after_failed_upsert(on_router, model_id)
return None
else:
return deployment
def _restore_deployment_after_failed_upsert(self, previous: Deployment | None, model_id: str) -> None:
if previous is None or self._core_has_model_id(model_id):
return
try:
if self.strategies is None:
self.add_deployment(deployment=previous)
else:
self.strategies.without_capability_limit(lambda: self.add_deployment(deployment=previous))
verbose_router_logger.info(
"Restored deployment %s (id=%s); it keeps serving its previous configuration.",
previous.model_name,
model_id,
)
except Exception as restore_error: # noqa: BLE001 # best-effort restore: a second failure must not abort the reload
verbose_router_logger.warning(
"Could not restore previously served deployment %s (id=%s) after the failed upsert: %s",
previous.model_name,
model_id,
restore_error,
)
def _live_handle(self) -> "RouterCoreHandle | None":
return None if self.handle_lookup is None else self.handle_lookup()

66
tests/unit/proxy/management_endpoints/test_model_management_endpoints.py
View file

@ -1113,7 +1113,7 @@ class TestClearCache:
@pytest.mark.asyncio
async def test_clear_cache_preserve_config_models(self):
"""
clear_cache resets DB-backed auto-router entries and delegates every deployment
clear_cache preserves DB-backed auto-router entries and delegates every deployment
change to the reload, leaving config models untouched. It must not wipe
deployments itself -- see the delete_deployment assertion below.
"""
@ -1122,8 +1122,7 @@ class TestClearCache:
)
# Create mock router with mixed DB and config router deployments. The two DB
# entries are auto_router/* deployments (so their router-map entries should be
# cleared for reload); the config-defined router is preserved.
# entries and the config-defined router must remain available during reload.
mock_router = MagicMock()
plugins: dict[str, dict] = {"auto": {}, "complexity": {}, "adaptive": {}, "quality": {}, "switchyard": {}}
mock_router.registry.return_value.plugins.side_effect = plugins.__getitem__
@ -1165,20 +1164,9 @@ class TestClearCache:
):
await clear_cache()
# clear_cache must wipe ONLY the db auto-router deployments -- the ones whose
# strategy entries are popped below and can only be rebuilt via the add path.
# Ordinary db models are left alone: wiping them un-served every db model for
# the width of the reload, and the reconcile converges without it.
assert mock_router.apply.call_count == 2
mock_router.apply.assert_any_call(ChangeSet(delete=["db-model-1"]))
mock_router.apply.assert_any_call(ChangeSet(delete=["db-model-2"]))
# DB-backed router entries are cleared so they can be re-populated by the
# reload below; the config-backed router must survive, since add_deployment()
# only reloads DB models and would otherwise leave it permanently unroutable
# (see TestClearCachePreservesConfigRouters).
assert "db-auto-router" not in plugins["auto"]
assert "db-complexity-router" not in plugins["complexity"]
mock_router.apply.assert_not_called()
assert "db-auto-router" in plugins["auto"]
assert "db-complexity-router" in plugins["complexity"]
assert "config-router" in plugins["auto"]
assert "config-router" in plugins["complexity"]
@ -1188,22 +1176,8 @@ class TestClearCache:
)
@pytest.mark.asyncio
async def test_clear_cache_wipes_auto_routers_but_leaves_ordinary_db_models(self):
"""An ordinary db model must survive clear_cache; a db auto-router must not.
Two separate hazards meet here, and fixing one naively breaks the other:
- Wiping ordinary db models un-serves EVERY db model for the width of the
reload. The reconcile converges without that, so the wipe is a pure
data-plane hole.
- NOT wiping a db auto-router strands it. Its strategy registries are keyed by
model_name and are popped here, but Router.upsert_deployment returns early
for an unchanged deployment and never reaches the add path that rebuilds
them. Any unrelated model write would then leave every db-backed auto,
complexity, adaptive and quality router unroutable until a restart.
So the wipe is scoped to exactly the auto-router deployments.
"""
async def test_clear_cache_preserves_auto_routers_and_ordinary_db_models(self):
"""Deployments and their strategies remain available while the reload runs."""
from litellm.proxy.management_endpoints.model_management_endpoints import (
clear_cache,
)
@ -1242,10 +1216,8 @@ class TestClearCache:
):
await clear_cache()
# The auto-router deployment is wiped so the reload takes the add path and
# rebuilds its strategy entry; the ordinary db model is never touched.
mock_router.apply.assert_called_once_with(ChangeSet(delete=["db-auto-1"]))
assert "db-auto-router" not in plugins["auto"]
mock_router.apply.assert_not_called()
assert "db-auto-router" in plugins["auto"]
class TestClearCachePreservesConfigRouters:
@ -1299,9 +1271,7 @@ class TestClearCachePreservesConfigRouters:
):
await clear_cache()
# The DB-backed router for the model that was actually updated is cleared
# so the reload below can re-populate it.
assert "team-a-db-router" not in plugins["complexity"]
assert "team-a-db-router" in plugins["complexity"]
# Config-defined routers for unrelated tenants must survive untouched.
assert "config-defined-complexity-router" in plugins["complexity"]
assert "config-semantic-router" in plugins["auto"]
@ -1346,11 +1316,8 @@ class TestClearCachePreservesConfigRouters:
assert "shared-name" in plugins["complexity"]
@pytest.mark.asyncio
async def test_db_quality_adaptive_and_switchyard_routers_are_evicted(self):
"""The auto_router/ prefix also covers quality_router/, adaptive_router/ and switchyard. Their
registry entries must be popped too, or reload's init raises 'already exists'
(quality) or leaves a stale entry (adaptive).
"""
async def test_db_quality_adaptive_and_switchyard_routers_are_preserved(self):
"""Every strategy remains available until the deployment reconcile replaces it."""
from litellm.proxy.management_endpoints.model_management_endpoints import clear_cache
mock_router = MagicMock()
@ -1393,9 +1360,9 @@ class TestClearCachePreservesConfigRouters:
):
await clear_cache()
assert "q1" not in plugins["quality"]
assert "a1" not in plugins["adaptive"]
assert "s1" not in plugins["switchyard"]
assert "q1" in plugins["quality"]
assert "a1" in plugins["adaptive"]
assert "s1" in plugins["switchyard"]
class TestDeleteModelClearsRouterRegistry:
@ -6136,8 +6103,7 @@ class TestConcurrentModelWritesDoNotEvictEachOther:
@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.
"""clear_cache holds the reconcile lock and 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

89
tests/unit/router_rust/test_atomic_upsert.py Normal file
View file

@ -0,0 +1,89 @@
from collections.abc import Callable
from types import SimpleNamespace
from typing import Final
import pytest
from litellm.proxy import proxy_server
from litellm.proxy._types import ReconcileOutcome
from litellm.proxy.management_endpoints.model_management_endpoints import clear_cache
from litellm.router_api import ChangeSet
from litellm.router_core import RouterCore
from litellm.router_python_host import Endpoint
from tests.unit.test_router.core.registry_helpers import typed
from tests.unit.test_router.core.scripted_host import ScriptedHost, deployment, rust_core
class ObservingHost(ScriptedHost):
def __init__(self) -> None:
super().__init__()
self.observe: Callable[[], None] | None = None
def invoke(self, op: str) -> Endpoint:
endpoint: Final = super().invoke(op)
if op != "get_llm_provider":
return endpoint
def provider(*args: object, **kwargs: object) -> object: # kwargs-ok: provider resolver arguments
if self.observe is not None:
self.observe()
return endpoint(*args, **kwargs)
return provider
@pytest.mark.asyncio
async def test_upsert_keeps_the_previous_deployment_serving_while_resolving_its_replacement() -> None:
host: Final = ObservingHost()
core: Final[RouterCore] = rust_core(host, model_list=[deployment("group", "only")])
seen: Final[list[tuple[str, ...]]] = [] # mutable-ok: observations made by the injected provider resolver
def observe() -> None:
seen.append(tuple(str(item["model_info"]["id"]) for item in core.registry().deployments("group")))
host.observe = observe
result: Final = core.apply(ChangeSet(upsert=[typed("group", "only", api_base="https://replacement.example")]))
assert seen and set(seen) == {("only",)}, seen
assert list(result.upserted) == [typed("group", "only", api_base="https://replacement.example")]
await core.acompletion(model="group", messages=[{"role": "user", "content": "after replacement"}]) # pyright: ignore[reportAttributeAccessIssue] # the core's completion endpoint
assert host.deployments_called("acompletion") == ["only"]
@pytest.mark.asyncio
async def test_clear_cache_keeps_a_db_classifier_serving_while_the_reload_waits(
monkeypatch: pytest.MonkeyPatch,
) -> None:
host: Final = ScriptedHost()
core: Final = rust_core(
host,
model_list=[
deployment("tier", "tier"),
deployment(
"smart",
"router",
model="auto_router/complexity_router",
complexity_router_config={
"classifier_type": "heuristic_v2",
"tiers": {"SIMPLE": "tier", "MEDIUM": "tier", "COMPLEX": "tier", "REASONING": "tier"},
},
model_info={"db_model": True},
),
],
)
seen: Final[list[tuple[str, ...]]] = [] # mutable-ok: snapshots at the reload boundary
async def reload(**kwargs: object) -> ReconcileOutcome: # kwargs-ok: the config reload's injected I/O boundary
seen.append(tuple(str(item["model_info"]["id"]) for item in core.registry().deployments("smart")))
await core.acompletion(model="smart", messages=[{"role": "user", "content": "hello"}]) # pyright: ignore[reportAttributeAccessIssue] # the core's completion endpoint
return ReconcileOutcome(still_desired=frozenset({"router"}), live_after=frozenset({"router"}))
monkeypatch.setattr(proxy_server, "llm_router", core)
monkeypatch.setattr(proxy_server, "prisma_client", object())
monkeypatch.setattr(proxy_server, "proxy_config", SimpleNamespace(add_deployment_locked=reload))
outcome: Final = await clear_cache()
assert seen == [("router",)]
assert outcome == ReconcileOutcome(still_desired=frozenset({"router"}), live_after=frozenset({"router"}))
assert host.deployments_called("acompletion") == ["tier"]

6
tests/unit/test_router/core/test_gaps_l11100_strategy.py
View file

@ -300,7 +300,7 @@ def test_a_failed_edit_rolls_back_to_the_serving_router_even_above_a_tightened_l
assert restored.litellm_params.tags is None
def test_a_failed_edit_without_ignore_invalid_deployments_raises_and_leaves_the_router_removed(
def test_a_failed_edit_without_ignore_invalid_deployments_raises_and_keeps_the_router_serving(
make_core: MakeCore,
) -> None:
limit: Final = {"value": 2}
@ -319,8 +319,8 @@ def test_a_failed_edit_without_ignore_invalid_deployments_raises_and_leaves_the_
ChangeSet(upsert=[Deployment(**complexity_marker("two", "m2", config=HEURISTIC_V2, tags=["edited"]))]) # pyright: ignore[reportArgumentType]
)
assert str(raised.value) == HEURISTIC_V2_LIMIT_MESSAGE
assert model_ids(core) == ["s", "b", "m1"]
assert registered_tags(core, "complexity") == {"one": [()]}
assert model_ids(core) == ["s", "b", "m1", "m2"]
assert registered_tags(core, "complexity") == {"one": [()], "two": [()]}
# --- the adaptive router -------------------------------------------------------------------------------

24
tests/unit/test_router/core/test_registry_characterization_writes.py
View file

@ -158,14 +158,13 @@ def test_an_upsert_that_empties_a_group_drops_the_group_name(make_core: MakeCore
assert core.registry().deployments("emb") == []
def test_a_failing_upsert_of_an_existing_deployment_raises_and_leaves_it_removed(make_core: MakeCore) -> None:
"""Odd but current: without ignore_invalid_deployments the old deployment is popped and not restored."""
def test_a_failing_upsert_of_an_existing_deployment_raises_and_keeps_it_serving(make_core: MakeCore) -> None:
core = _core(make_core)
with pytest.raises(litellm.BadRequestError):
core.apply(ChangeSet(upsert=[typed("gpt", "a", model="notaprovider/x")]))
assert core.registry().names(kind="ids") == ["b", "e", "w", "t1"]
assert core.registry().names(kind="ids") == ["a", "b", "e", "w", "t1"]
def test_a_failing_upsert_under_ignore_invalid_deployments_restores_the_previous_deployment(
@ -179,7 +178,7 @@ def test_a_failing_upsert_under_ignore_invalid_deployments_restores_the_previous
restored = core.registry().deployment("a")
assert restored is not None
assert restored.litellm_params.model == "openai/a"
assert core.registry().names(kind="ids") == ["b", "e", "w", "t1", "a"]
assert core.registry().names(kind="ids") == ["a", "b", "e", "w", "t1"]
def test_the_deprecated_upsert_deployment_warns_and_returns_the_deployment(make_core: MakeCore) -> None:
@ -540,7 +539,7 @@ def test_a_read_on_another_thread_waits_for_an_upsert_to_read_the_list_back(make
past the list nor, after a removal, the neighbouring group's entries."""
core: Final = _two_per_group(make_core)
hosted: Final = core._hosted # pyright: ignore[reportAttributeAccessIssue] # the registry under test
write: Final[Callable[[Mapping[str, object]], None]] = hosted.changes_runner
write: Final = hosted.build_runner
started: Final = threading.Event()
pool: Final = ThreadPoolExecutor(max_workers=1)
reads: Final[list[object]] = [] # mutable-ok: one per write
@ -549,25 +548,30 @@ def test_a_read_on_another_thread_waits_for_an_upsert_to_read_the_list_back(make
started.set()
return _group_ids(core)
def write_with_a_read_started(changes: Mapping[str, object]) -> None:
write(changes)
def write_with_a_read_started(
entries: Sequence[object],
path: str,
resolve_secrets: bool,
install: str,
answer: Callable[[Mapping[str, object]], object],
) -> None:
write(entries, path, resolve_secrets, install, answer)
started.clear()
reads.append(pool.submit(read))
assert started.wait(timeout=5)
# The reader runs while router-core holds the write and the mirror does not yet.
_ = threading.Event().wait(0.02)
hosted.changes_runner = write_with_a_read_started
hosted.build_runner = write_with_a_read_started
try:
_ = core.apply(ChangeSet(upsert=[typed(f"m{upserted % _GROUPS}", f"d{upserted}", timeout=11)]))
finally:
hosted.changes_runner = write
hosted.build_runner = write
assert reads
for future in reads:
seen = future.result(timeout=5) # pyright: ignore[reportAttributeAccessIssue] # a Future
for group in range(_GROUPS):
# Between the upsert's removal and its add, its group is served by the other deployment.
assert seen[group], (group, seen)
assert set(seen[group]) <= {f"d{group}", f"d{group + _GROUPS}"}, (group, seen)
pool.shutdown()