mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-29 06:33:28 +00:00
fix(proxy): keep first-match semantics in the credential upsert and filter deletes in one pass
Index the existing list by first occurrence so a name that appears twice (a config duplicate) still has the entry find_credential reads replaced, as before. Replace the index/sort/pop delete with a single in-place filter. Move the accessor tests under tests/unit so a CI shard owns them, hoist imports, drop narrative comments, assert whole values.
This commit is contained in:
parent
3d33258b43
commit
27c192e1ae
5 changed files with 75 additions and 91 deletions
|
|
@ -31,19 +31,16 @@ class CredentialAccessor:
|
|||
|
||||
@staticmethod
|
||||
def upsert_credentials(credentials: list[CredentialItem]):
|
||||
"""Add credentials to the list, replacing any existing entry with the same name in place.
|
||||
|
||||
Runs on every config reload with the full DB credential table, so it indexes the
|
||||
current list once instead of rescanning it per credential.
|
||||
"""
|
||||
index_by_name: dict[str, int] = {cred.credential_name: i for i, cred in enumerate(litellm.credential_list)}
|
||||
"""Add credentials to the list, replacing the first existing entry with the same name in place."""
|
||||
index_by_name: Final[dict[str, int]] = {}
|
||||
for i, cred in enumerate(litellm.credential_list):
|
||||
index_by_name.setdefault(cred.credential_name, i)
|
||||
|
||||
for credential in credentials:
|
||||
existing_index = index_by_name.get(credential.credential_name)
|
||||
if existing_index is not None:
|
||||
litellm.credential_list[existing_index] = credential
|
||||
else:
|
||||
index_by_name[credential.credential_name] = len(litellm.credential_list)
|
||||
litellm.credential_list.append(credential)
|
||||
|
||||
|
||||
|
|
|
|||
|
|
@ -9656,18 +9656,11 @@ class ProxyConfig:
|
|||
config: Final = await self.get_config(config_file_path=user_config_file_path)
|
||||
credential_list: Final = self.load_credential_list(config=config)
|
||||
|
||||
## COMBINED LIST ##
|
||||
combined_list: Final = db_credentials + credential_list
|
||||
|
||||
## DELETE ##
|
||||
names_to_keep: Final = {cred.credential_name for cred in combined_list}
|
||||
idx_to_delete: Final = [
|
||||
idx
|
||||
for idx, credential in enumerate(litellm.credential_list)
|
||||
if credential.credential_name not in names_to_keep
|
||||
]
|
||||
for idx in sorted(idx_to_delete, reverse=True):
|
||||
litellm.credential_list.pop(idx)
|
||||
names_to_keep: Final = {cred.credential_name for cred in db_credentials} | {
|
||||
cred.credential_name for cred in credential_list
|
||||
}
|
||||
litellm.credential_list[:] = [cred for cred in litellm.credential_list if cred.credential_name in names_to_keep]
|
||||
|
||||
async def get_credentials(self, prisma_client: PrismaClient):
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -1,42 +0,0 @@
|
|||
"""CredentialAccessor.upsert_credentials: the per-reload sync of DB credentials into litellm.credential_list."""
|
||||
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from litellm.types.utils import CredentialItem
|
||||
|
||||
|
||||
def _cred(name: str, api_key: str) -> CredentialItem:
|
||||
return CredentialItem(credential_name=name, credential_values={"api_key": api_key}, credential_info={})
|
||||
|
||||
|
||||
def test_upsert_credentials_replaces_by_name_in_place_and_appends_new(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(litellm, "credential_list", [_cred("a", "old-a"), _cred("b", "old-b"), _cred("c", "old-c")])
|
||||
|
||||
CredentialAccessor.upsert_credentials([_cred("b", "new-b"), _cred("d", "new-d"), _cred("a", "new-a")])
|
||||
|
||||
assert [(c.credential_name, c.credential_values["api_key"]) for c in litellm.credential_list] == [
|
||||
("a", "new-a"),
|
||||
("b", "new-b"),
|
||||
("c", "old-c"),
|
||||
("d", "new-d"),
|
||||
]
|
||||
|
||||
|
||||
def test_upsert_credentials_scales_linearly_with_credential_count(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
# The proxy re-upserts every DB credential on each config reload (default every 30s) on the serving
|
||||
# event loop. 20k credentials took ~40s with the previous nested scan; linear takes milliseconds.
|
||||
n = 20_000
|
||||
monkeypatch.setattr(litellm, "credential_list", [_cred(f"cred-{i}", f"k{i}") for i in range(n)])
|
||||
refreshed = [_cred(f"cred-{i}", f"k{i}-v2") for i in range(n)]
|
||||
|
||||
started = time.perf_counter()
|
||||
CredentialAccessor.upsert_credentials(refreshed)
|
||||
elapsed = time.perf_counter() - started
|
||||
|
||||
assert len(litellm.credential_list) == n
|
||||
assert litellm.credential_list[n - 1].credential_values["api_key"] == f"k{n - 1}-v2"
|
||||
assert elapsed < 2.0, f"upsert of {n} credentials took {elapsed:.1f}s; expected linear-time behaviour"
|
||||
51
tests/unit/litellm_core_utils/test_credential_accessor.py
Normal file
51
tests/unit/litellm_core_utils/test_credential_accessor.py
Normal file
|
|
@ -0,0 +1,51 @@
|
|||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from litellm.types.utils import CredentialItem
|
||||
|
||||
|
||||
def _credential(name: str, api_key: str) -> CredentialItem:
|
||||
return CredentialItem(credential_name=name, credential_values={"api_key": api_key}, credential_info={})
|
||||
|
||||
|
||||
def _names_and_keys() -> list[tuple[str, str]]:
|
||||
return [(cred.credential_name, cred.credential_values["api_key"]) for cred in litellm.credential_list]
|
||||
|
||||
|
||||
def test_upsert_credentials_replaces_by_name_in_place_and_appends_new(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setattr(
|
||||
litellm, "credential_list", [_credential("a", "old-a"), _credential("b", "old-b"), _credential("c", "old-c")]
|
||||
)
|
||||
|
||||
CredentialAccessor.upsert_credentials(
|
||||
[_credential("b", "new-b"), _credential("d", "new-d"), _credential("a", "new-a")]
|
||||
)
|
||||
|
||||
assert _names_and_keys() == [("a", "new-a"), ("b", "new-b"), ("c", "old-c"), ("d", "new-d")]
|
||||
|
||||
|
||||
def test_upsert_credentials_replaces_the_first_duplicate_that_find_credential_reads(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
monkeypatch.setattr(litellm, "credential_list", [_credential("a", "first"), _credential("a", "second")])
|
||||
|
||||
CredentialAccessor.upsert_credentials([_credential("a", "from-db")])
|
||||
|
||||
assert _names_and_keys() == [("a", "from-db"), ("a", "second")]
|
||||
assert CredentialAccessor.get_credential_values("a") == {"api_key": "from-db"}
|
||||
|
||||
|
||||
def test_upsert_credentials_scales_linearly_with_credential_count(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
n = 20_000
|
||||
monkeypatch.setattr(litellm, "credential_list", [_credential(f"cred-{i}", f"k{i}") for i in range(n)])
|
||||
refreshed = [_credential(f"cred-{i}", f"k{i}-v2") for i in range(n)]
|
||||
|
||||
started = time.perf_counter()
|
||||
CredentialAccessor.upsert_credentials(refreshed)
|
||||
elapsed = time.perf_counter() - started
|
||||
|
||||
assert litellm.credential_list == refreshed
|
||||
assert elapsed < 2.0, f"upsert of {n} credentials took {elapsed:.1f}s; the per-reload sync must stay linear"
|
||||
|
|
@ -13,6 +13,7 @@ import json
|
|||
import logging
|
||||
import os
|
||||
import re
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Mapping
|
||||
from contextlib import nullcontext
|
||||
from dataclasses import dataclass
|
||||
|
|
@ -29,6 +30,7 @@ from pydantic import JsonValue, TypeAdapter, ValidationError
|
|||
|
||||
import litellm
|
||||
from litellm.proxy._types import CommonProxyErrors, ConfigGeneralSettings
|
||||
from litellm.types.utils import CredentialItem
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
|
||||
from litellm.types.router import Deployment, LiteLLM_Params
|
||||
from litellm.proxy.proxy_server import (
|
||||
|
|
@ -2064,9 +2066,7 @@ def test_ProxyConfig_load_credential_list_invalid_entry_raises():
|
|||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _credential_item(name: str) -> "CredentialItem":
|
||||
from litellm.types.utils import CredentialItem
|
||||
|
||||
def _credential(name: str) -> CredentialItem:
|
||||
return CredentialItem(credential_name=name, credential_values={"api_key": f"key-{name}"}, credential_info={})
|
||||
|
||||
|
||||
|
|
@ -2076,50 +2076,35 @@ async def test_ProxyConfig_delete_credentials_drops_only_names_missing_from_db_a
|
|||
) -> None:
|
||||
pc = ProxyConfig()
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"credential_list",
|
||||
[_credential_item("db-kept"), _credential_item("stale"), _credential_item("from-config")],
|
||||
litellm, "credential_list", [_credential("db-kept"), _credential("stale"), _credential("from-config")]
|
||||
)
|
||||
monkeypatch.setattr(
|
||||
pc, "get_config", AsyncMock(return_value={"credential_list": [_credential("from-config").model_dump()]})
|
||||
)
|
||||
|
||||
async def fake_get_config(*args: object, **kwargs: object) -> dict[str, object]:
|
||||
return {
|
||||
"credential_list": [
|
||||
{"credential_name": "from-config", "credential_values": {"api_key": "cfg"}, "credential_info": {}}
|
||||
]
|
||||
}
|
||||
await pc.delete_credentials([_credential("db-kept")])
|
||||
|
||||
monkeypatch.setattr(pc, "get_config", fake_get_config)
|
||||
|
||||
await pc.delete_credentials([_credential_item("db-kept")])
|
||||
|
||||
assert [c.credential_name for c in litellm.credential_list] == ["db-kept", "from-config"]
|
||||
assert [cred.credential_name for cred in litellm.credential_list] == ["db-kept", "from-config"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_ProxyConfig_delete_credentials_scales_linearly_with_credential_count(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
# Runs on the serving event loop every config reload; 20k credentials took ~40s with the previous
|
||||
# nested scan and blocked every request and health probe for that long.
|
||||
import time
|
||||
|
||||
n = 20_000
|
||||
pc = ProxyConfig()
|
||||
monkeypatch.setattr(litellm, "credential_list", [_credential_item(f"cred-{i}") for i in range(n)])
|
||||
|
||||
async def fake_get_config(*args: object, **kwargs: object) -> dict[str, object]:
|
||||
return {}
|
||||
|
||||
monkeypatch.setattr(pc, "get_config", fake_get_config)
|
||||
db_credentials = [_credential_item(f"cred-{i}") for i in range(n) if i != 7]
|
||||
monkeypatch.setattr(litellm, "credential_list", [_credential(f"cred-{i}") for i in range(n)])
|
||||
monkeypatch.setattr(pc, "get_config", AsyncMock(return_value={}))
|
||||
db_credentials = [cred for cred in litellm.credential_list if cred.credential_name != "cred-7"]
|
||||
|
||||
started = time.perf_counter()
|
||||
await pc.delete_credentials(db_credentials)
|
||||
elapsed = time.perf_counter() - started
|
||||
|
||||
assert len(litellm.credential_list) == n - 1
|
||||
assert all(c.credential_name != "cred-7" for c in litellm.credential_list)
|
||||
assert elapsed < 2.0, f"delete_credentials over {n} credentials took {elapsed:.1f}s; expected linear-time behaviour"
|
||||
assert litellm.credential_list == db_credentials
|
||||
assert elapsed < 2.0, (
|
||||
f"delete_credentials over {n} credentials took {elapsed:.1f}s; the per-reload sync must stay linear"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue