diff --git a/litellm/proxy/common_utils/registry_read_through.py b/litellm/proxy/common_utils/registry_read_through.py index 487c6feb68d..d58c58fa1a5 100644 --- a/litellm/proxy/common_utils/registry_read_through.py +++ b/litellm/proxy/common_utils/registry_read_through.py @@ -133,7 +133,7 @@ async def _resync_model_deployments(model_name: str) -> bool: rows: Final = await _find_rows_by_name_or_id(model_name, ModelRepository(prisma_client).table) if not rows: return False - async with proxy_server.MODEL_RECONCILE_LOCK: + async with proxy_server.DB_RECONCILE_LOCK: outcome: Final = await _install_model_rows(rows, _always) if outcome in ("no_router", "auto_router"): await _reload_every_model() @@ -150,7 +150,7 @@ async def resync_model_deployments_by_ids(model_ids: tuple[str, ...]) -> None: prisma_client: Final = proxy_server.prisma_client assert prisma_client is not None writer_table: Final = ModelRepository(WriterPinnedClient(prisma_client.db)).table - async with proxy_server.MODEL_RECONCILE_LOCK: + async with proxy_server.DB_RECONCILE_LOCK: rows: Final = await _find_rows_by_ids(model_ids, writer_table) _ = await _install_model_rows(_without_auto_routers(rows), _names_an_unloaded_credential) diff --git a/litellm/proxy/credential_endpoints/endpoints.py b/litellm/proxy/credential_endpoints/endpoints.py index 86e716f19e3..3026f2706ac 100644 --- a/litellm/proxy/credential_endpoints/endpoints.py +++ b/litellm/proxy/credential_endpoints/endpoints.py @@ -324,7 +324,7 @@ async def create_credential( Stores credential in DB. Reloads credentials in memory. """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import DB_RECONCILE_LOCK, prisma_client try: if prisma_client is None: @@ -365,21 +365,22 @@ async def create_credential( credentials_dict_jsonified: Final = cast( # cast-ok: deep-copies a model_dump, so keys are str "dict[str, object]", jsonify_object(credentials_dict) ) - try: - await CredentialsRepository(prisma_client).create( - data={ - **credentials_dict_jsonified, - "created_by": user_api_key_dict.user_id, - "updated_by": user_api_key_dict.user_id, - } - ) - except Exception as e: - if not is_unique_violation(e): - raise - raise HTTPException(status_code=409, detail=_credential_exists_detail(credential.credential_name)) + async with DB_RECONCILE_LOCK: + try: + await CredentialsRepository(prisma_client).create( + data={ + **credentials_dict_jsonified, + "created_by": user_api_key_dict.user_id, + "updated_by": user_api_key_dict.user_id, + } + ) + except Exception as e: + if not is_unique_violation(e): + raise + raise HTTPException(status_code=409, detail=_credential_exists_detail(credential.credential_name)) - ## ADD TO LITELLM ## - CredentialAccessor.upsert_credentials([processed_credential]) + ## ADD TO LITELLM ## + CredentialAccessor.upsert_credentials([processed_credential]) return {"success": True, "message": "Credential created successfully"} except Exception as e: @@ -865,7 +866,7 @@ async def delete_credential( """ [BETA] endpoint. This might change unexpectedly. """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import DB_RECONCILE_LOCK, prisma_client try: if prisma_client is None: @@ -876,14 +877,17 @@ async def delete_credential( _reject_non_admin_wif_fields( await named_credential_wif_fields(credential_name, prisma_client), user_api_key_dict ) - deleted: Final = await CredentialsRepository(prisma_client).delete_by_name(credential_name) - if deleted is None: - raise _not_found_unless_config_defined( - credential_name, "Credential not found. Got credential name: " + credential_name - ) + async with DB_RECONCILE_LOCK: + deleted: Final = await CredentialsRepository(prisma_client).delete_by_name(credential_name) + if deleted is None: + raise _not_found_unless_config_defined( + credential_name, "Credential not found. Got credential name: " + credential_name + ) - ## DELETE FROM LITELLM ## - litellm.credential_list = [cred for cred in litellm.credential_list if cred.credential_name != credential_name] + ## DELETE FROM LITELLM ## + litellm.credential_list = [ + cred for cred in litellm.credential_list if cred.credential_name != credential_name + ] await _purge_user_connections_for_credential(credential_name) return {"success": True, "message": "Credential deleted successfully"} except Exception as e: @@ -983,7 +987,7 @@ async def update_credential( """ [BETA] endpoint. This might change unexpectedly. """ - from litellm.proxy.proxy_server import prisma_client + from litellm.proxy.proxy_server import DB_RECONCILE_LOCK, prisma_client try: if credential.credential_name and credential.credential_name != credential_name: @@ -1007,53 +1011,55 @@ async def update_credential( detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) credentials_repository: Final = CredentialsRepository(WriterPinnedClient(prisma_client.db)) - db_credential: Final = await credentials_repository.find_by_name(credential_name) - if db_credential is None: - raise _not_found_unless_config_defined(credential_name, "Credential not found in DB.") - _reject_non_admin_wif_fields(_stored_wif_fields(db_credential), user_api_key_dict) - decrypted_stored_credential: Final = decrypted_credential(db_credential) - decrypted_values: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( - decrypted_stored_credential.credential_values - ) - incoming_values: Final = _without_masked_read_backs(requested_values, _stored_objects(decrypted_values)) - patch: Final = CredentialItem( - credential_name=credential_name, - display_name=( - requested_display_name if "display_name" in credential.model_fields_set else db_credential.display_name - ), - credential_info=_CREDENTIAL_DICT_ADAPTER.validate_python(credential.credential_info), - credential_values=_CREDENTIAL_DICT_ADAPTER.validate_python(incoming_values), - credential_values_to_delete=credential.credential_values_to_delete, - ) - stored_info: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( - db_credential.credential_info # pyright: ignore[reportUnknownMemberType] # untyped credential field - ) - patch_info: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( - patch.credential_info # pyright: ignore[reportUnknownMemberType] # untyped credential field - ) - effective_info: Final = {**stored_info, **patch_info} - normalized_patch: Final = ( - patch.model_copy( - update={"credential_values": _normalized_github_copilot_values(incoming_values, effective_info)} + async with DB_RECONCILE_LOCK: + db_credential: Final = await credentials_repository.find_by_name(credential_name) + if db_credential is None: + raise _not_found_unless_config_defined(credential_name, "Credential not found in DB.") + _reject_non_admin_wif_fields(_stored_wif_fields(db_credential), user_api_key_dict) + decrypted_stored_credential: Final = decrypted_credential(db_credential) + decrypted_values: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( + decrypted_stored_credential.credential_values ) - if incoming_values.get(GITHUB_COPILOT_HOST_KEY) is not None - else patch - ) - merged_credential: Final = update_db_credential(db_credential, normalized_patch) - credential_object_jsonified: Final = cast( # cast-ok: deep-copies a model_dump, so keys are str - "dict[str, object]", jsonify_object(merged_credential.model_dump(exclude_none=True)) - ) - await credentials_repository.update_by_name( - credential_name, - data={ - **credential_object_jsonified, - "display_name": merged_credential.display_name, - "updated_by": user_api_key_dict.user_id, - }, - ) - - # Sync in-memory credential_list (skip if not in memory - e.g., proxy restarted) - _sync_in_memory_credential(normalized_patch, credential_name) + incoming_values: Final = _without_masked_read_backs(requested_values, _stored_objects(decrypted_values)) + patch: Final = CredentialItem( + credential_name=credential_name, + display_name=( + requested_display_name + if "display_name" in credential.model_fields_set + else db_credential.display_name + ), + credential_info=_CREDENTIAL_DICT_ADAPTER.validate_python(credential.credential_info), + credential_values=_CREDENTIAL_DICT_ADAPTER.validate_python(incoming_values), + credential_values_to_delete=credential.credential_values_to_delete, + ) + stored_info: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( + db_credential.credential_info # pyright: ignore[reportUnknownMemberType] # untyped credential field + ) + patch_info: Final = _CREDENTIAL_DICT_ADAPTER.validate_python( + patch.credential_info # pyright: ignore[reportUnknownMemberType] # untyped credential field + ) + effective_info: Final = {**stored_info, **patch_info} + normalized_patch: Final = ( + patch.model_copy( + update={"credential_values": _normalized_github_copilot_values(incoming_values, effective_info)} + ) + if incoming_values.get(GITHUB_COPILOT_HOST_KEY) is not None + else patch + ) + merged_credential: Final = update_db_credential(db_credential, normalized_patch) + credential_object_jsonified: Final = cast( # cast-ok: deep-copies a model_dump, so keys are str + "dict[str, object]", jsonify_object(merged_credential.model_dump(exclude_none=True)) + ) + await credentials_repository.update_by_name( + credential_name, + data={ + **credential_object_jsonified, + "display_name": merged_credential.display_name, + "updated_by": user_api_key_dict.user_id, + }, + ) + # Sync in-memory credential_list (skip if not in memory - e.g., proxy restarted) + _sync_in_memory_credential(normalized_patch, credential_name) from litellm.llms.github_copilot.github_host import normalized_github_host diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index b3c1f48bf99..3c8f7d6d64f 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -2392,12 +2392,12 @@ async def delete_team_models( if deleted_model_ids: await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable") - # Under MODEL_RECONCILE_LOCK, for the same reason as delete_model: the rows are + # Under DB_RECONCILE_LOCK, for the same reason as delete_model: the rows are # gone, but a reconcile holding a pre-delete snapshot would upsert these ids back # onto this pod. The lock orders the eviction after any in-flight reconcile. - from litellm.proxy.proxy_server import MODEL_RECONCILE_LOCK, proxy_config + from litellm.proxy.proxy_server import DB_RECONCILE_LOCK, proxy_config - async with MODEL_RECONCILE_LOCK: + async with DB_RECONCILE_LOCK: proxy_config.remove_auto_router_catalog_entries(frozenset(deleted_model_ids)) if llm_router is not None: for model_id in deleted_model_ids: @@ -2884,7 +2884,7 @@ async def delete_model( """ from litellm.proxy.proxy_server import ( - MODEL_RECONCILE_LOCK, + DB_RECONCILE_LOCK, general_settings_view, llm_router, premium_user, @@ -2945,13 +2945,13 @@ async def delete_model( ) ## DELETE FROM ROUTER ## - # Under MODEL_RECONCILE_LOCK. The db row is already gone, but a reconcile + # Under DB_RECONCILE_LOCK. The db row is already gone, but a reconcile # that snapshotted the db BEFORE that delete still lists this id as desired, # and its _add_deployment upserts the deployment straight back -- leaving # this pod serving a model the database no longer has, until the next # reconcile. Taking the lock orders this eviction after any such in-flight # reconcile's re-add, so the eviction is the last word. - async with MODEL_RECONCILE_LOCK: + async with DB_RECONCILE_LOCK: proxy_config.remove_auto_router_catalog_entries(frozenset({model_info.id})) if llm_router is not None: llm_router.apply(ChangeSet(delete=[model_info.id])) @@ -3892,7 +3892,7 @@ def reload_serving_verdict( Without it (no reconcile ran) every drop is reported, which is the safe direction. ``live_after`` is the router's serving state captured by the reload itself, while it - still held MODEL_RECONCILE_LOCK. Pass it whenever the caller has it: re-reading the + still held DB_RECONCILE_LOCK. Pass it whenever the caller has it: re-reading the router here instead means sampling it after the lock was released, where the NEXT reconcile's leading wipe (clear_cache un-serves every db model before reloading them) shows up as this reload having dropped them. Falling back to a fresh read is @@ -3970,13 +3970,13 @@ async def clear_cache() -> ReconcileOutcome: Returns what the reload saw (see ReconcileOutcome) so callers can pass it to raise_if_reload_degraded_serving. - Runs under MODEL_RECONCILE_LOCK for its whole extent, not just the reload at the + 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 add_deployment would re-acquire the same non-reentrant lock and deadlock. """ from litellm.proxy.proxy_server import ( - MODEL_RECONCILE_LOCK, + DB_RECONCILE_LOCK, llm_router, prisma_client, proxy_config, @@ -3988,7 +3988,7 @@ async def clear_cache() -> ReconcileOutcome: verbose_proxy_logger.debug("llm_router or prisma_client is None, skipping cache clear") return ReconcileOutcome(still_desired=None, live_after=None) - async with MODEL_RECONCILE_LOCK: + async with DB_RECONCILE_LOCK: try: # Only clear DB models, preserve config models verbose_proxy_logger.debug("Clearing only DB models, preserving config models") @@ -4060,7 +4060,7 @@ async def clear_cache() -> ReconcileOutcome: llm_router.registry().plugins("switchyard").pop(model_name, None) # Reload only DB models. _add_deployment_locked, not add_deployment: this - # coroutine already holds MODEL_RECONCILE_LOCK and asyncio.Lock is not + # coroutine already holds DB_RECONCILE_LOCK and asyncio.Lock is not # reentrant, so the public wrapper would deadlock against itself. outcome: Final = await proxy_config.add_deployment_locked( prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index fe633238e20..8c92a058f43 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3140,15 +3140,16 @@ experimental = False llm_router: Router | None = None llm_model_list: list | None = None heuristic_v1_tuning_baselines: Mapping[str, str] | None = None -# Serializes every model reconcile (ProxyConfig.add_deployment and clear_cache) so the -# read-modify-write of llm_router above is atomic. Without it, two concurrent model -# writes each reconcile the router against their OWN db snapshot, and the one holding -# the older snapshot evicts the deployment the newer one just added -- the db keeps the -# row, this pod stops serving it. Control-plane only (model create/update/delete and -# the config-sync tick), never on a completion path, so the serialization is free. -# Module-level rather than per-ProxyConfig because llm_router is a module global and a -# second ProxyConfig instance must not get its own independent lock over it. -MODEL_RECONCILE_LOCK: Final = asyncio.Lock() +# Serializes every db-to-memory reconcile (ProxyConfig.add_deployment and clear_cache) +# with every control-plane write to the state those reconciles rebuild: llm_router above +# and litellm.credential_list. Without it, two concurrent writes each reconcile against +# their OWN db snapshot, and the one holding the older snapshot evicts the deployment or +# credential the newer one just added -- the db keeps the row, this pod stops serving it. +# Control-plane only (model and credential create/update/delete and the config-sync +# tick), never on a completion path, so the serialization is free. Module-level rather +# than per-ProxyConfig because llm_router is a module global and a second ProxyConfig +# instance must not get its own independent lock over it. +DB_RECONCILE_LOCK: Final = asyncio.Lock() general_settings: dict = {} _GENERAL_SETTINGS_VIEW: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) @@ -8922,7 +8923,7 @@ class ProxyConfig: - Check if model id's in router already - If not, add to router - Serialized against every other model reconcile by MODEL_RECONCILE_LOCK, because + Serialized against every other model reconcile by DB_RECONCILE_LOCK, because the work below is a read-modify-write of the shared ``llm_router`` global: it reads the db into a snapshot and then makes the router match that snapshot. Two of those interleaving is not a lost update but an eviction -- the request whose @@ -8941,7 +8942,7 @@ class ProxyConfig: if not ui_settings_already_synced: await sync_ui_settings_to_general_settings(prisma_client) - async with MODEL_RECONCILE_LOCK: + async with DB_RECONCILE_LOCK: return await self.add_deployment_locked(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj) async def add_deployment_locked( @@ -8949,7 +8950,7 @@ class ProxyConfig: prisma_client: PrismaClient, proxy_logging_obj: ProxyLogging, ) -> ReconcileOutcome: - """add_deployment's body, minus the locking. MODEL_RECONCILE_LOCK MUST already + """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.""" @@ -9898,7 +9899,7 @@ class ProxyConfig: immediately, and by the periodic job in store_model_in_db-off deployments. Gated the same way as startup, so an admin who excluded search_tools from supported_db_objects opts out. - Serialized by MODEL_RECONCILE_LOCK for the reason add_deployment documents: the body is a + Serialized by DB_RECONCILE_LOCK for the reason add_deployment documents: the body is a read-modify-write of the shared ``llm_router`` global, so two of them interleaving lets the older snapshot's wholesale assignment land last and restore a tool the newer one deleted. The lock belongs here rather than in _init_search_tools_in_db, which _init_non_llm_objects_in_db @@ -9908,7 +9909,7 @@ class ProxyConfig: return if prisma_client is None: return - async with MODEL_RECONCILE_LOCK: + async with DB_RECONCILE_LOCK: await self._init_search_tools_in_db(prisma_client=prisma_client) @staticmethod diff --git a/tests/integration/_support/database_relay.py b/tests/integration/_support/database_relay.py index 2e9f2ec4a41..0ac1e1d2521 100644 --- a/tests/integration/_support/database_relay.py +++ b/tests/integration/_support/database_relay.py @@ -6,7 +6,6 @@ from contextlib import contextmanager from typing import Final from urllib.parse import parse_qsl, urlencode, urlsplit, urlunsplit - from tests.integration._support.async_server import LoopbackServer, close_writer from tests.integration._support.ports import allocate_port @@ -74,16 +73,37 @@ class DatabaseRelay: class HeldStatementRelay: - def __init__(self, upstream_host: str, upstream_port: int, trigger: bytes) -> None: + """Holds the first client statement carrying ``trigger`` until ``release``. With ``after``, only a + trigger statement that follows a statement carrying ``after`` is held, so a test can pin the hold to + one point of a longer sequence. Starts armed unless ``armed`` is false, then ``arm`` starts it.""" + + def __init__( + self, + upstream_host: str, + upstream_port: int, + trigger: bytes, + *, + after: bytes | None = None, + armed: bool = True, + ) -> None: self.port: Final = allocate_port() self._upstream_host: Final = upstream_host self._upstream_port: Final = upstream_port self._trigger: Final = trigger + self._after: Final = after + self._marker_length: Final = max(len(trigger), len(after or b"")) self._server: Final = LoopbackServer(self.port, self._serve) self._loop: Final = self._server.loop self._released: Final = asyncio.Event() + self._armed: Final = threading.Event() + if armed: + self._armed.set() + self._after_seen: Final = threading.Event() self.held: Final = threading.Event() + def arm(self) -> None: + self._armed.set() + def release(self) -> None: if not self._loop.is_closed(): self._loop.call_soon_threadsafe(self._released.set) @@ -96,7 +116,13 @@ class HeldStatementRelay: self._server.stop() def _holds(self, window: bytes) -> bool: - return not self.held.is_set() and self._trigger in window + if not self._armed.is_set() or self.held.is_set(): + return False + if self._after is not None and not self._after_seen.is_set(): + if self._after in window: + self._after_seen.set() + return False + return self._trigger in window async def _serve(self, client_reader: asyncio.StreamReader, client_writer: asyncio.StreamWriter) -> None: server_reader, server_writer = await asyncio.open_connection(self._upstream_host, self._upstream_port) @@ -109,7 +135,7 @@ class HeldStatementRelay: if inspect and self._holds(window): self.held.set() await self._released.wait() - tail = window[-(len(self._trigger) - 1) :] + tail = window[-(self._marker_length - 1) :] writer.write(chunk) await writer.drain() except (ConnectionError, asyncio.IncompleteReadError): @@ -295,10 +321,12 @@ def database_relay(database_url: str, trigger: bytes) -> Generator[tuple[Databas @contextmanager -def held_statement_relay(database_url: str, trigger: bytes) -> Generator[tuple[HeldStatementRelay, str]]: +def held_statement_relay( + database_url: str, trigger: bytes, *, after: bytes | None = None, armed: bool = True +) -> Generator[tuple[HeldStatementRelay, str]]: parts: Final = urlsplit(database_url) assert parts.hostname is not None and parts.port is not None, database_url - relay: Final = HeldStatementRelay(parts.hostname, parts.port, trigger) + relay: Final = HeldStatementRelay(parts.hostname, parts.port, trigger, after=after, armed=armed) try: relay.start() yield relay, _relayed_url(database_url, relay.port) diff --git a/tests/integration/management/test_credential_writes_survive_reload.py b/tests/integration/management/test_credential_writes_survive_reload.py new file mode 100644 index 00000000000..9fb7d7eb681 --- /dev/null +++ b/tests/integration/management/test_credential_writes_survive_reload.py @@ -0,0 +1,119 @@ +import time +import uuid +from concurrent.futures import ThreadPoolExecutor, wait +from pathlib import Path +from typing import Final, Literal + +import pytest +import yaml +from pydantic import JsonValue + +from tests.integration._support.client import JSON_OBJECT, Gateway +from tests.integration._support.database import read_rows, scratch_database +from tests.integration._support.database_relay import held_statement_relay, plaintext_database_url +from tests.integration._support.process import owned_proxy + +RELOAD_INTERVAL_SECONDS: Final = 6 +CREDENTIAL_SNAPSHOT: Final = b'FROM "public"."LiteLLM_CredentialsTable"' +CONFIG_ROW_READ: Final = b"environment_variables" +CredentialWrite = Literal["create", "update", "delete"] + + +def _reloading_config(directory: Path) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["general_settings"] = { + **config.get("general_settings", {}), + "store_model_in_db": True, + "proxy_config_reload_interval_seconds": RELOAD_INTERVAL_SECONDS, + } + config["litellm_settings"] = { + name: value + for name, value in config.get("litellm_settings", {}).items() + if name not in ("cache", "cache_params", "enable_redis_auth_cache") + } + path: Final = directory / "credential-reload.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _credential_body(name: str, display_name: str) -> dict[str, JsonValue]: + return { + "credential_name": name, + "display_name": display_name, + "credential_values": {"api_key": "synthetic-credential"}, + "credential_info": {"custom_llm_provider": "openai"}, + } + + +def _write_credential(proxy: Gateway, write: CredentialWrite, name: str) -> int: + match write: + case "create": + return proxy.request("POST", "/credentials", _credential_body(name, "created")).status_code + case "update": + return proxy.request( + "PATCH", f"/credentials/{name}", {"display_name": "updated", "credential_info": {}} + ).status_code + case "delete": + return proxy.request("DELETE", f"/credentials/{name}").status_code + + +def _served_display_name(proxy: Gateway, name: str) -> JsonValue: + response: Final = proxy.request("GET", f"/credentials/by_name/{name}") + if response.status_code == 404: + return None + assert response.status_code == 200, response.text + return JSON_OBJECT.validate_json(response.content)["display_name"] + + +def _stored_display_name(name: str, database_url: str) -> JsonValue: + rows: Final = read_rows( + 'SELECT display_name FROM "LiteLLM_CredentialsTable" WHERE credential_name = %s', + (name,), + database_url=database_url, + ) + return rows[0]["display_name"] if rows else None + + +def _served_display_names_until_the_next_reload(proxy: Gateway, name: str) -> frozenset[JsonValue]: + deadline: Final = time.monotonic() + RELOAD_INTERVAL_SECONDS / 3 + observed: Final = tuple(_served_display_name(proxy, name) for _ in iter(lambda: time.monotonic() < deadline, False)) + return frozenset(observed) + + +@pytest.mark.timeout(300) +@pytest.mark.parametrize("write", ("create", "update", "delete")) +def test_credential_write_is_not_undone_by_an_in_flight_config_reload( + gateway: Gateway, tmp_path: Path, write: CredentialWrite +) -> None: + name: Final = f"credential-{uuid.uuid4().hex}" + expected: Final = {"create": "created", "update": "updated", "delete": None}[write] + with ( + scratch_database() as database_url, + held_statement_relay(database_url, CONFIG_ROW_READ, after=CREDENTIAL_SNAPSHOT, armed=False) as ( + relay, + relayed_url, + ), + ): + overrides: Final = { + "DATABASE_URL": plaintext_database_url(relayed_url), + "LITELLM_CONFIG_PARAM_CACHE_TTL_SECONDS": "0", + } + with owned_proxy( + gateway, + tmp_path, + overrides, + config=_reloading_config(tmp_path), + remove_environment=("REDIS_HOST", "REDIS_PORT", "REDIS_URL"), + workers=1, + ) as proxy: + if write != "create": + proxy.post("/credentials", _credential_body(name, "original")) + relay.arm() + assert relay.held.wait(timeout=RELOAD_INTERVAL_SECONDS * 3), "no config reload reached the database" + with ThreadPoolExecutor(max_workers=1) as pool: + writing: Final = pool.submit(_write_credential, proxy, write, name) + _ = wait((writing,), timeout=2) + relay.release() + assert writing.result(timeout=30) == 200 + assert _stored_display_name(name, database_url) == expected + assert _served_display_names_until_the_next_reload(proxy, name) == frozenset({expected}) diff --git a/tests/unit/proxy/common_utils/test_registry_read_through.py b/tests/unit/proxy/common_utils/test_registry_read_through.py index da7b1f497c9..6c9104bb711 100644 --- a/tests/unit/proxy/common_utils/test_registry_read_through.py +++ b/tests/unit/proxy/common_utils/test_registry_read_through.py @@ -433,7 +433,7 @@ async def test_resync_model_deployments_mutates_router_under_model_reconcile_loc lock_states: list[bool] = [] def record_add_deployment(db_models) -> None: - lock_states.append(proxy_server.MODEL_RECONCILE_LOCK.locked()) + lock_states.append(proxy_server.DB_RECONCILE_LOCK.locked()) monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) monkeypatch.setattr(proxy_server, "store_model_in_db", True) @@ -443,7 +443,7 @@ async def test_resync_model_deployments_mutates_router_under_model_reconcile_loc assert await _resync_model_deployments("lock-scope-model") is True assert lock_states == [True] - assert not proxy_server.MODEL_RECONCILE_LOCK.locked() + assert not proxy_server.DB_RECONCILE_LOCK.locked() @pytest.mark.asyncio @@ -601,7 +601,7 @@ async def test_resync_model_deployments_by_ids_installs_only_the_published_rows( assert _served_ids(router.registry().deployments()) == {"primary-id", "m-1", "m-2"} assert _served_ids(proxy_server.llm_model_list) == {"primary-id", "m-1", "m-2"} - assert not proxy_server.MODEL_RECONCILE_LOCK.locked() + assert not proxy_server.DB_RECONCILE_LOCK.locked() @pytest.mark.asyncio @@ -619,7 +619,7 @@ async def test_resync_model_deployments_by_ids_leaves_the_router_alone_for_delet assert _served_ids(router.registry().deployments()) == {"primary-id"} assert proxy_server.llm_model_list is None - assert not proxy_server.MODEL_RECONCILE_LOCK.locked() + assert not proxy_server.DB_RECONCILE_LOCK.locked() @pytest.mark.asyncio @@ -1021,7 +1021,7 @@ async def test_unknown_model_name_read_through_reads_the_replica_without_waiting monkeypatch.setattr(proxy_server, "prisma_client", prisma_client) monkeypatch.setattr(proxy_server, "store_model_in_db", True) monkeypatch.setattr(proxy_server, "llm_router", MagicMock()) - monkeypatch.setattr(proxy_server, "MODEL_RECONCILE_LOCK", reload_in_progress) + monkeypatch.setattr(proxy_server, "DB_RECONCILE_LOCK", reload_in_progress) async with reload_in_progress: found: Final = await asyncio.wait_for(_resync_model_deployments("no-such-model"), timeout=2) @@ -1042,7 +1042,7 @@ async def test_resync_model_deployments_by_ids_installs_the_row_a_concurrent_rel router: Final = _router_serving("primary-id") _serve_db_models(monkeypatch, _Db(table), router) reload_in_progress: Final = asyncio.Lock() - monkeypatch.setattr(proxy_server, "MODEL_RECONCILE_LOCK", reload_in_progress) + monkeypatch.setattr(proxy_server, "DB_RECONCILE_LOCK", reload_in_progress) async with reload_in_progress: targeted_load: Final = asyncio.create_task(resync_model_deployments_by_ids(("m-1",))) @@ -1132,4 +1132,4 @@ async def test_resync_model_deployments_by_ids_installs_the_ordinary_rows_queued assert _served_ids(proxy_server.llm_model_list) == {"primary-id", "backup-id"} assert table.reads == 1 assert db.other_tables_read == [] - assert not proxy_server.MODEL_RECONCILE_LOCK.locked() + assert not proxy_server.DB_RECONCILE_LOCK.locked() diff --git a/tests/unit/proxy/credential_endpoints/test_endpoints.py b/tests/unit/proxy/credential_endpoints/test_endpoints.py index 10751def88d..f08fab25c57 100644 --- a/tests/unit/proxy/credential_endpoints/test_endpoints.py +++ b/tests/unit/proxy/credential_endpoints/test_endpoints.py @@ -1,6 +1,8 @@ """Tests for the credential management endpoints.""" +import asyncio import json +from collections.abc import Awaitable, Callable from contextlib import contextmanager from functools import reduce from types import SimpleNamespace @@ -19,11 +21,22 @@ import litellm from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.litellm_logging import get_masked_values from litellm.llms.github_copilot.per_user_auth import _SESSION_CACHE, _session_cache_key -from litellm.models.credentials import CredentialSource, UserConnectionFlowHandle +from litellm.models.credentials import ( + CreateCredentialItem, + CredentialSource, + UpdateCredentialItem, + UserConnectionFlowHandle, +) from litellm.proxy._types import LiteLLM_UserTable, LitellmUserRoles, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper, encrypt_value_helper -from litellm.proxy.credential_endpoints.endpoints import _decode_flow_handle, get_llm_router +from litellm.proxy.credential_endpoints.endpoints import ( + _decode_flow_handle, + create_credential, + delete_credential, + get_llm_router, + update_credential, +) from litellm.proxy.proxy_server import app from litellm.types.utils import CredentialItem @@ -3345,3 +3358,101 @@ def test_get_credential_by_model_masks_the_deployment_credentials(monkeypatch): with pytest.raises(ProxyException) as missing: asyncio.run(get_credential_by_model(MagicMock(), MagicMock(), "no-such-id", caller)) assert missing.value.code == "404" + + +class TestCredentialWritesHoldTheReconcileLock: + """A config reload snapshots the credential table, awaits the config rows, then prunes and + upserts ``litellm.credential_list`` from that snapshot. A credential write that lands inside + that gap is undone in memory until the next reload, so each write has to queue behind + DB_RECONCILE_LOCK, the lock the reload already holds.""" + + @staticmethod + async def _assert_write_waits_for_the_lock( + monkeypatch: pytest.MonkeyPatch, call_endpoint: Callable[[], Awaitable[dict[str, object]]], write: str + ) -> dict[str, object]: + lock: Final = asyncio.Lock() + monkeypatch.setattr("litellm.proxy.proxy_server.DB_RECONCILE_LOCK", lock) + async with lock: + task: Final = asyncio.create_task(call_endpoint()) + for _ in range(50): + await asyncio.sleep(0) + assert not task.done(), ( + f"{write} did not wait for DB_RECONCILE_LOCK, so an in-flight config reload can undo it in memory" + ) + return await asyncio.wait_for(task, timeout=5) + + async def test_create_credential_waits_for_an_in_flight_reload( + self, monkeypatch: pytest.MonkeyPatch, credential_store: Callable[..., None] + ) -> None: + credential_store(create=AsyncMock(return_value=None)) + + async def call_endpoint() -> dict[str, object]: + return await create_credential( + request=MagicMock(spec=Request), + fastapi_response=MagicMock(), + credential=CreateCredentialItem( + credential_name="fresh", + credential_values={"api_key": "sk-new"}, + credential_info={"custom_llm_provider": "openai"}, + ), + user_api_key_dict=_as_admin(), + llm_router=None, + ) + + response: Final = await self._assert_write_waits_for_the_lock(monkeypatch, call_endpoint, "create") + + assert response["success"] is True + assert [credential.credential_name for credential in litellm.credential_list] == ["fresh"] + + async def test_update_credential_waits_for_an_in_flight_reload( + self, monkeypatch: pytest.MonkeyPatch, credential_store: Callable[..., None] + ) -> None: + stored: Final = CredentialItem( + credential_name="existing", + display_name="before", + credential_values={"api_key": "sk-old"}, + credential_info={"custom_llm_provider": "openai"}, + ) + credential_store( + in_memory=(stored,), + find_by_name=AsyncMock(return_value=stored), + update_by_name=AsyncMock(return_value=None), + ) + + async def call_endpoint() -> dict[str, object]: + return await update_credential( + request=MagicMock(spec=Request), + fastapi_response=MagicMock(), + credential=UpdateCredentialItem(display_name="after", credential_info={}), + credential_name="existing", + user_api_key_dict=_as_admin(), + llm_router=None, + ) + + response: Final = await self._assert_write_waits_for_the_lock(monkeypatch, call_endpoint, "update") + + assert response["success"] is True + assert [credential.display_name for credential in litellm.credential_list] == ["after"] + + async def test_delete_credential_waits_for_an_in_flight_reload( + self, monkeypatch: pytest.MonkeyPatch, credential_store: Callable[..., None] + ) -> None: + stored: Final = CredentialItem( + credential_name="doomed", + credential_values={"api_key": "sk-old"}, + credential_info={"custom_llm_provider": "openai"}, + ) + credential_store(in_memory=(stored,), delete_by_name=AsyncMock(return_value=MagicMock())) + + async def call_endpoint() -> dict[str, object]: + return await delete_credential( + request=MagicMock(spec=Request), + fastapi_response=MagicMock(), + credential_name="doomed", + user_api_key_dict=_as_admin(), + ) + + response: Final = await self._assert_write_waits_for_the_lock(monkeypatch, call_endpoint, "delete") + + assert response["success"] is True + assert litellm.credential_list == [] 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 2fdb77af8f9..ee16b68b750 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -1571,7 +1571,7 @@ def deleted_auto_router_catalog(monkeypatch): config = proxy_server.ProxyConfig() config.auto_router_db_catalog = build_auto_router_catalog(rows) monkeypatch.setattr(proxy_server, "proxy_config", config) - monkeypatch.setattr(proxy_server, "MODEL_RECONCILE_LOCK", asyncio.Lock()) + monkeypatch.setattr(proxy_server, "DB_RECONCILE_LOCK", asyncio.Lock()) monkeypatch.setattr(proxy_server, "llm_router", Router(model_list=[])) monkeypatch.setattr(proxy_server, "_license_check", SimpleNamespace(auto_router_capability_limit=lambda: 1)) monkeypatch.setattr(proxy_server, "heuristic_v1_tuning_baselines", {}) @@ -5959,7 +5959,7 @@ class TestConcurrentModelWritesDoNotEvictEachOther: @pytest.mark.asyncio async def test_reconciles_serialize_so_no_stale_snapshot_can_evict(self, monkeypatch): - """MODEL_RECONCILE_LOCK admits one reconcile at a time. + """DB_RECONCILE_LOCK admits one reconcile at a time. The fake body awaits, which is the whole point: without the lock the gather below parks all five inside the critical section at that await and observed depth goes @@ -6088,7 +6088,7 @@ class TestConcurrentModelWritesDoNotEvictEachOther: class TestDeleteEvictionsHoldTheReconcileLock: """A delete evicts from ``llm_router`` directly instead of reconciling, so it must - take MODEL_RECONCILE_LOCK to do it. + take DB_RECONCILE_LOCK to do it. The db row is gone by then, but a reconcile that snapshotted the db BEFORE the row was deleted still lists that id as desired, and its ``_add_deployment`` upserts the @@ -6102,7 +6102,7 @@ class TestDeleteEvictionsHoldTheReconcileLock: async def _assert_evicts_under_lock(monkeypatch, call_endpoint, model_id: str, config) -> None: """Run ``call_endpoint`` with the lock already held and assert it blocks. - Holding MODEL_RECONCILE_LOCK stands in for a reconcile that is mid-flight. If + Holding DB_RECONCILE_LOCK stands in for a reconcile that is mid-flight. If the eviction takes the lock it cannot run until we release; if it does not, it runs straight through and the deployment is evicted while the "reconcile" is still in its critical section -- exactly the interleaving that resurrects it. @@ -6115,7 +6115,7 @@ class TestDeleteEvictionsHoldTheReconcileLock: own, exactly as here. """ lock = asyncio.Lock() - monkeypatch.setattr("litellm.proxy.proxy_server.MODEL_RECONCILE_LOCK", lock) + monkeypatch.setattr("litellm.proxy.proxy_server.DB_RECONCILE_LOCK", lock) stale_catalog = config.auto_router_db_catalog async with lock: @@ -6124,7 +6124,7 @@ class TestDeleteEvictionsHoldTheReconcileLock: for _ in range(50): await asyncio.sleep(0) assert not task.done(), ( - f"deleting {model_id} did not wait for MODEL_RECONCILE_LOCK -- an " + f"deleting {model_id} did not wait for DB_RECONCILE_LOCK -- an " f"in-flight reconcile can resurrect the deployment it just evicted" ) config.auto_router_db_catalog = stale_catalog