mirror of
https://github.com/BerriAI/litellm.git
synced 2026-11-12 08:42:18 +00:00
fix(proxy): hold DB_RECONCILE_LOCK around credential create, update and delete (#50551)
* fix(proxy): hold MODEL_RECONCILE_LOCK around credential create, update and delete The config reload snapshots the credential table, awaits config rows, then prunes and upserts litellm.credential_list from that snapshot. A credential write landing inside that window was undone in memory until the next reload. The three endpoints now take the same lock the reload holds. Fixes #50273 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): take MODEL_RECONCILE_LOCK before the credential update reads its row Two PATCHes queued behind a reload read the same old row and the second overwrote the first. The relay now keeps enough bytes to spot a split after marker, and the new unit tests are fully typed. * test(credentials): annotate the reconcile lock test locals with Final * refactor(proxy): rename MODEL_RECONCILE_LOCK to DB_RECONCILE_LOCK The lock gates every db-to-memory reconcile and now the credential writes too, so the model-only name undersold what it protects. --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
Kerry Lu
GitHub
parent
b83a1f0608
commit
51b49973c9
9 changed files with 383 additions and 118 deletions
|
|
@ -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)
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
|
|
@ -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)
|
||||
|
|
|
|||
|
|
@ -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})
|
||||
|
|
@ -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()
|
||||
|
|
|
|||
|
|
@ -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 == []
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue