diff --git a/litellm-rust/crates/router-core/src/lib.rs b/litellm-rust/crates/router-core/src/lib.rs index f9b71c689c9..e61d570c78f 100644 --- a/litellm-rust/crates/router-core/src/lib.rs +++ b/litellm-rust/crates/router-core/src/lib.rs @@ -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, 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(¤t.registry); for spec in installed { - next.push_built(spec, add)?; + next.install_built(spec, add, install)?; } next }; diff --git a/litellm-rust/crates/router-core/src/registry.rs b/litellm-rust/crates/router-core/src/registry.rs index 594aff5f6dd..2aff7f243ff 100644 --- a/litellm-rust/crates/router-core/src/registry.rs +++ b/litellm-rust/crates/router-core/src/registry.rs @@ -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( diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 38b814225ce..7846922bae2 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -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. diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 9f721e3fdcb..23d45e07dfa 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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 diff --git a/litellm/router_rust/steps/prerouting.py b/litellm/router_rust/steps/prerouting.py index 47bd1e69046..2953a87ed92 100644 --- a/litellm/router_rust/steps/prerouting.py +++ b/litellm/router_rust/steps/prerouting.py @@ -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 diff --git a/litellm/router_rust/surface/registry.py b/litellm/router_rust/surface/registry.py index 83008bdbb94..59d3d7389be 100644 --- a/litellm/router_rust/surface/registry.py +++ b/litellm/router_rust/surface/registry.py @@ -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() diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 2b269469f89..2f8456c6d0d 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -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 diff --git a/tests/unit/router_rust/test_atomic_upsert.py b/tests/unit/router_rust/test_atomic_upsert.py new file mode 100644 index 00000000000..51050b0002d --- /dev/null +++ b/tests/unit/router_rust/test_atomic_upsert.py @@ -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"] diff --git a/tests/unit/test_router/core/test_gaps_l11100_strategy.py b/tests/unit/test_router/core/test_gaps_l11100_strategy.py index 2459c577b9c..c311459c4e5 100644 --- a/tests/unit/test_router/core/test_gaps_l11100_strategy.py +++ b/tests/unit/test_router/core/test_gaps_l11100_strategy.py @@ -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 ------------------------------------------------------------------------------- diff --git a/tests/unit/test_router/core/test_registry_characterization_writes.py b/tests/unit/test_router/core/test_registry_characterization_writes.py index 23fa2a8e0e9..ecb6aa81832 100644 --- a/tests/unit/test_router/core/test_registry_characterization_writes.py +++ b/tests/unit/test_router/core/test_registry_characterization_writes.py @@ -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()