mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-31 06:56:59 +00:00
fix(proxy): apply new model rows on every worker as soon as they are written (#49034)
* fix(proxy): apply published model rows on every worker before the debounced resync Config sync messages now carry the ids of model rows a create, update, or upsert wrote. Every worker that receives one loads those rows by id right away, including while its debounced full resync is still waiting, so a group added one deployment at a time is complete on every worker as soon as the write returns. Deletes, bulk writes, and messages from older pods carry no ids and keep the full resync path * fix(proxy): load published model rows from the primary under the reconcile lock Read the by-id rows through the writer so a lagging read replica cannot put an old row back, and read them inside MODEL_RECONCILE_LOCK so a concurrent delete or reload stays the last word. Look the ids up in bounded chunks, reload DB credentials only when a row names one this worker lacks, and leave auto-router rows to the debounced full resync. Add an integration test that boots two proxies and needs the peer to fall back to a just-added backup within 4 s. * test(proxy): cover a published model row that names no credential * fix(proxy): keep the unknown model name read-through on the replica and outside the lock Only the by-id load that follows a published write reads the primary inside the reconcile lock. A request for a model name a worker does not know reads through the replica routing before taking the lock again, as on main, so unknown names put no load on the primary and do not wait behind a reload. * fix(proxy): keep access group edits from dropping staying deployments and batch published config ids * fix(proxy): read access group rows from the primary so quick edits are not lost to a lagging replica * fix(proxy): install ordinary rows queued with an auto router and check access group edits before writing A batch of published model ids that held one auto router skipped every row in it, so ordinary rows written at the same moment waited for the full reload. The by-id load now leaves only the auto router rows to that reload. An access group update now reads every deployment from the primary, answers 404 for an unknown group and 400 for an unknown model id or name before it writes anything, then removes the group from the deployments that leave before it adds it to the ones that join, so a failed write leaves the group narrower than asked, never wider. * fix(proxy): write an access group update in one transaction after a fresh read The update now checks model names first, reads every deployment from the primary, and commits every join and leave in one batch right after that read, so a failed write changes nothing and a field another admin set during validation is kept. Other workers hear about the rows only after the commit. * fix(proxy): lock the model rows while an access group update reads and rewrites them * test(proxy): run the model resync tests on the real installer and a real router The by-id resync tests now install rows through ProxyConfig's real installer and credential loader on a real Router, faking only the database, and the config sync test row model is frozen. * fix(proxy): lock only the rows an access group update touches and answer 404 before 400 An access group update locked every model row, so a lock on any unrelated row stalled it for 30 s and then failed it with a 500, and unrelated model edits on other proxies waited behind it. It now takes a per-group advisory lock, locks only the rows that carry the group or that it selects, and reads them again once locked. The group existence check runs before the model name check again, so a missing group answers 404 as on main. * fix(proxy): never wait on a locked row in an access group update The update now takes its row locks with NOWAIT. When another edit holds one of the group's rows, it rolls back, holding nothing, and tries again for about 10 s before answering 409. Holding rows while waiting deadlocked against team deletes and let a concurrent model edit undo the update's change. The per-group lock also moves to the two-key advisory space, apart from team ids. * fix(proxy): retry access group row locks for 30 s with SKIP LOCKED and lock only the rows model_ids select * fix(proxy): take the access group lock with pg_try_advisory_xact_lock and retry while another edit holds it --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com>
This commit is contained in:
tessellark
GitHub
parent
d1758cda91
commit
080e643edc
12 changed files with 2049 additions and 222 deletions
|
|
@ -1648,7 +1648,7 @@
|
|||
},
|
||||
"/access_group/{access_group}/update": {
|
||||
"put": {
|
||||
"description": "Update an access group's model names.\n\nThis will:\n1. Remove the access group from all current deployments\n2. Add the access group to all deployments for the new model_names list\n\nExample:\n```bash\ncurl -X PUT 'http://localhost:4000/access_group/production-models/update' \\\n -H \"Authorization: Bearer $LITELLM_MASTER_KEY\" \\\n -H 'Content-Type: application/json' \\\n -d '{\n \"model_names\": [\"gpt-4\", \"claude-3-sonnet\"]\n }'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n- model_names: List[str] - New list of model groups to include\n\nReturns:\n- NewModelGroupResponse with the updated access group details\n\nRaises:\n- HTTPException 400: If any model names don't exist\n- HTTPException 404: If access group not found",
|
||||
"description": "Update an access group's model names.\n\nThis will:\n1. Add the access group to all deployments for the new model_names list\n2. Remove the access group from every other deployment\n\nExample:\n```bash\ncurl -X PUT 'http://localhost:4000/access_group/production-models/update' \\\n -H \"Authorization: Bearer $LITELLM_MASTER_KEY\" \\\n -H 'Content-Type: application/json' \\\n -d '{\n \"model_names\": [\"gpt-4\", \"claude-3-sonnet\"]\n }'\n```\n\nParameters:\n- access_group: str - The access group name (URL path parameter)\n- model_names: List[str] - New list of model groups to include\n\nReturns:\n- NewModelGroupResponse with the updated access group details\n\nRaises:\n- HTTPException 400: If any model names don't exist\n- HTTPException 404: If access group not found\n- HTTPException 409: If another edit holds the group or one of its models for 30 s",
|
||||
"operationId": "update_access_group_access_group__access_group__update_put",
|
||||
"parameters": [
|
||||
{
|
||||
|
|
|
|||
|
|
@ -2,9 +2,12 @@ import asyncio
|
|||
import json
|
||||
import random
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable
|
||||
from collections.abc import AsyncIterator, Awaitable, Callable
|
||||
from dataclasses import asdict, dataclass
|
||||
from typing import TYPE_CHECKING, Final, Protocol, cast # noqa: TID251 # untyped prisma/redis boundary needs cast
|
||||
from itertools import chain
|
||||
from typing import TYPE_CHECKING, Final, Protocol, TypeAlias, cast # noqa: TID251 # untyped prisma/redis boundary
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.repositories.prisma_protocols import RowT_co, TableActions
|
||||
|
|
@ -106,24 +109,69 @@ _pubsub_capable_client: Final = pubsub_capable_client
|
|||
@dataclass(frozen=True, slots=True)
|
||||
class _ConfigChangeMessage:
|
||||
object_type: str
|
||||
model_ids: tuple[str, ...] = ()
|
||||
|
||||
|
||||
def _config_change_message_json(object_type: str) -> str:
|
||||
return json.dumps(asdict(_ConfigChangeMessage(object_type=object_type)))
|
||||
@dataclass(frozen=True, slots=True)
|
||||
class _PubSubEnvelope:
|
||||
data: str | bytes
|
||||
|
||||
|
||||
async def publish_config_change(redis_cache: "RedisCache | None", object_type: str) -> None:
|
||||
class _WrittenModelRow(BaseModel):
|
||||
model_config = ConfigDict(from_attributes=True, frozen=True)
|
||||
|
||||
model_id: str
|
||||
|
||||
|
||||
_CONFIG_CHANGE_MESSAGE: Final = TypeAdapter(_ConfigChangeMessage)
|
||||
_PUBSUB_ENVELOPE: Final = TypeAdapter(_PubSubEnvelope)
|
||||
_MODEL_TABLE_NAME: Final = "litellm_proxymodeltable"
|
||||
_TARGETED_WRITE_ACTION_NAMES: Final[frozenset[str]] = frozenset({"create", "update", "upsert"})
|
||||
|
||||
PublishConfigChange: TypeAlias = Callable[[str, tuple[str, ...]], Awaitable[None]]
|
||||
|
||||
|
||||
def _config_change_message_json(object_type: str, model_ids: tuple[str, ...]) -> str:
|
||||
return json.dumps(asdict(_ConfigChangeMessage(object_type=object_type, model_ids=model_ids)))
|
||||
|
||||
|
||||
def parse_config_change(message: object) -> _ConfigChangeMessage | None:
|
||||
try:
|
||||
payload: Final = (
|
||||
message if isinstance(message, (str, bytes)) else _PUBSUB_ENVELOPE.validate_python(message).data
|
||||
)
|
||||
return _CONFIG_CHANGE_MESSAGE.validate_json(payload)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
def written_model_ids(object_type: str, action_name: str, result: object) -> tuple[str, ...]:
|
||||
if object_type != _MODEL_TABLE_NAME or action_name not in _TARGETED_WRITE_ACTION_NAMES:
|
||||
return ()
|
||||
try:
|
||||
return (_WrittenModelRow.model_validate(result).model_id,)
|
||||
except ValidationError:
|
||||
return ()
|
||||
|
||||
|
||||
async def publish_config_change(
|
||||
redis_cache: "RedisCache | None", object_type: str, model_ids: tuple[str, ...] = ()
|
||||
) -> None:
|
||||
if redis_cache is None:
|
||||
return
|
||||
try:
|
||||
client: Final = pubsub_capable_client(redis_cache)
|
||||
await client.publish(config_sync_channel(redis_cache), _config_change_message_json(object_type))
|
||||
await client.publish(config_sync_channel(redis_cache), _config_change_message_json(object_type, model_ids))
|
||||
except Exception as e: # noqa: BLE001 # best-effort publish; writes must never fail on redis errors
|
||||
verbose_proxy_logger.warning("config sync publish for %s failed: %s", object_type, e)
|
||||
|
||||
|
||||
async def publish_config_change_for_object_type(object_type: str) -> None:
|
||||
await publish_config_change(redis_cache=coordination_redis_cache(), object_type=object_type)
|
||||
async def publish_config_change_for_object_type(object_type: str, model_ids: tuple[str, ...] = ()) -> None:
|
||||
await publish_config_change(redis_cache=coordination_redis_cache(), object_type=object_type, model_ids=model_ids)
|
||||
|
||||
|
||||
async def publish_model_rows_written(model_ids: tuple[str, ...]) -> None:
|
||||
await publish_config_change_for_object_type(_MODEL_TABLE_NAME, model_ids)
|
||||
|
||||
|
||||
async def publish_config_param_change(param_name: str) -> None:
|
||||
|
|
@ -139,7 +187,7 @@ async def publish_config_param_change(param_name: str) -> None:
|
|||
class _PublishOnWriteActions:
|
||||
__slots__ = ("_actions", "_object_type", "_publish")
|
||||
|
||||
def __init__(self, actions: object, object_type: str, publish: Callable[[str], Awaitable[None]]) -> None:
|
||||
def __init__(self, actions: object, object_type: str, publish: PublishConfigChange) -> None:
|
||||
self._actions = actions
|
||||
self._object_type = object_type
|
||||
self._publish = publish
|
||||
|
|
@ -157,7 +205,7 @@ class _PublishOnWriteActions:
|
|||
**kwargs: object, # kwargs-ok: transparent passthrough to untyped prisma action
|
||||
) -> object:
|
||||
result: Final = await write_action(*args, **kwargs)
|
||||
await publish(object_type)
|
||||
await publish(object_type, written_model_ids(object_type, name, result))
|
||||
return result
|
||||
|
||||
return _write_then_publish
|
||||
|
|
@ -166,7 +214,7 @@ class _PublishOnWriteActions:
|
|||
def wrap_table_actions_for_config_sync(
|
||||
actions: "TableActions[RowT_co]",
|
||||
table_name: str,
|
||||
publish: Callable[[str], Awaitable[None]] = publish_config_change_for_object_type,
|
||||
publish: PublishConfigChange = publish_config_change_for_object_type,
|
||||
) -> "TableActions[RowT_co]":
|
||||
if table_name not in _CONFIG_SYNCED_TABLE_NAMES:
|
||||
return actions
|
||||
|
|
@ -174,8 +222,14 @@ def wrap_table_actions_for_config_sync(
|
|||
return cast("TableActions[RowT_co]", wrapped) # cast-ok: dynamic write-through proxy keeps the wrapped row type
|
||||
|
||||
|
||||
def _published_model_ids(messages: tuple[object, ...]) -> tuple[str, ...]:
|
||||
changes: Final = (parse_config_change(message) for message in messages)
|
||||
return tuple(dict.fromkeys(chain.from_iterable(change.model_ids for change in changes if change is not None)))
|
||||
|
||||
|
||||
class ConfigSyncSubscriber:
|
||||
__slots__ = (
|
||||
"_apply_model_ids",
|
||||
"_backoff_initial_seconds",
|
||||
"_backoff_max_seconds",
|
||||
"_debounce_seconds",
|
||||
|
|
@ -195,6 +249,7 @@ class ConfigSyncSubscriber:
|
|||
self,
|
||||
redis_cache: "RedisCache",
|
||||
resync_callbacks: tuple[Callable[[], Awaitable[None]], ...],
|
||||
apply_model_ids: Callable[[tuple[str, ...]], Awaitable[None]] | None = None,
|
||||
debounce_seconds: float = CONFIG_SYNC_DEBOUNCE_SECONDS,
|
||||
jitter_max_seconds: float = CONFIG_SYNC_JITTER_MAX_SECONDS,
|
||||
min_resync_interval_seconds: float = CONFIG_SYNC_MIN_RESYNC_INTERVAL_SECONDS,
|
||||
|
|
@ -206,6 +261,7 @@ class ConfigSyncSubscriber:
|
|||
) -> None:
|
||||
self._redis_cache = redis_cache
|
||||
self._resync_callbacks = resync_callbacks
|
||||
self._apply_model_ids = apply_model_ids
|
||||
self._debounce_seconds = debounce_seconds
|
||||
self._jitter_max_seconds = jitter_max_seconds
|
||||
self._min_resync_interval_seconds = min_resync_interval_seconds
|
||||
|
|
@ -266,32 +322,64 @@ class ConfigSyncSubscriber:
|
|||
message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=_POLL_TIMEOUT_SECONDS)
|
||||
if message is None:
|
||||
continue
|
||||
await self._apply(pubsub, message)
|
||||
await self._debounce_and_resync(pubsub)
|
||||
|
||||
async def _debounce_and_resync(self, pubsub: _ConfigSyncPubSub) -> None:
|
||||
await self._sleep(self._debounce_seconds + self._rng.uniform(0.0, self._jitter_max_seconds))
|
||||
await self._wait_for_min_resync_interval()
|
||||
async def _debounce_and_resync(self, pubsub: ConfigSyncPubSub) -> None:
|
||||
await self._apply_until(
|
||||
pubsub, self._monotonic() + self._debounce_seconds + self._rng.uniform(0.0, self._jitter_max_seconds)
|
||||
)
|
||||
await self._apply_until(pubsub, self._next_resync_at())
|
||||
await self._drain_pending(pubsub)
|
||||
await self._run_resync_callbacks()
|
||||
self._last_resync_at = self._monotonic()
|
||||
|
||||
async def _wait_for_min_resync_interval(self) -> None:
|
||||
def _next_resync_at(self) -> float:
|
||||
if self._last_resync_at is None:
|
||||
return
|
||||
seconds_until_next_resync = self._min_resync_interval_seconds - (self._monotonic() - self._last_resync_at)
|
||||
if seconds_until_next_resync <= 0:
|
||||
return
|
||||
verbose_proxy_logger.debug(
|
||||
"config sync resync throttled for %.1fs to cap fleet-wide reload rate",
|
||||
seconds_until_next_resync,
|
||||
)
|
||||
await self._sleep(seconds_until_next_resync)
|
||||
return self._monotonic()
|
||||
next_resync_at: Final = self._last_resync_at + self._min_resync_interval_seconds
|
||||
throttled_for: Final = next_resync_at - self._monotonic()
|
||||
if throttled_for > 0:
|
||||
verbose_proxy_logger.debug(
|
||||
"config sync resync throttled for %.1fs to cap fleet-wide reload rate",
|
||||
throttled_for,
|
||||
)
|
||||
return next_resync_at
|
||||
|
||||
async def _apply_until(self, pubsub: ConfigSyncPubSub, deadline: float) -> None:
|
||||
while True:
|
||||
remaining = deadline - self._monotonic()
|
||||
if remaining <= 0:
|
||||
return
|
||||
message = await pubsub.get_message(
|
||||
ignore_subscribe_messages=True, timeout=min(remaining, _POLL_TIMEOUT_SECONDS)
|
||||
)
|
||||
if message is not None:
|
||||
await self._apply(pubsub, message)
|
||||
|
||||
@staticmethod
|
||||
async def _drain_pending(pubsub: ConfigSyncPubSub) -> None:
|
||||
while await pubsub.get_message(ignore_subscribe_messages=True, timeout=0) is not None:
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
async def _ready_messages(pubsub: ConfigSyncPubSub) -> AsyncIterator[object]:
|
||||
while (message := await pubsub.get_message(ignore_subscribe_messages=True, timeout=0)) is not None:
|
||||
yield message
|
||||
|
||||
async def _apply(self, pubsub: ConfigSyncPubSub, message: object) -> None:
|
||||
apply: Final = self._apply_model_ids
|
||||
if apply is None:
|
||||
return
|
||||
ready: Final = (message, *[ready_message async for ready_message in self._ready_messages(pubsub)])
|
||||
model_ids: Final = _published_model_ids(ready)
|
||||
if not model_ids:
|
||||
return
|
||||
try:
|
||||
await apply(model_ids)
|
||||
except Exception as e: # noqa: BLE001 # the debounced full resync retries what a targeted apply could not
|
||||
verbose_proxy_logger.warning("config sync targeted apply of %s failed: %s", model_ids, e)
|
||||
|
||||
async def _run_resync_callbacks(self) -> None:
|
||||
for callback in self._resync_callbacks:
|
||||
try:
|
||||
|
|
|
|||
|
|
@ -15,12 +15,13 @@ or the first routed request on this replica would pick a tier it has not loaded.
|
|||
import asyncio
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Final
|
||||
from typing import TYPE_CHECKING, Final, Literal, TypeAlias
|
||||
|
||||
from pydantic import TypeAdapter, ValidationError
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.caching.in_memory_cache import InMemoryCache
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
|
@ -33,8 +34,14 @@ if TYPE_CHECKING:
|
|||
)
|
||||
|
||||
from litellm.integrations.custom_guardrail import CustomGuardrail
|
||||
from litellm.repositories.prisma_protocols import TableActions
|
||||
from litellm.types.agents import AgentResponse
|
||||
|
||||
_ModelTable: TypeAlias = TableActions[prisma_models.LiteLLM_ProxyModelTable]
|
||||
_ModelRows: TypeAlias = Sequence[prisma_models.LiteLLM_ProxyModelTable]
|
||||
|
||||
ModelLoadOutcome: TypeAlias = Literal["loaded", "missing", "no_router", "auto_router"]
|
||||
|
||||
READ_THROUGH_MISS_TTL_SECONDS: Final = 2.0
|
||||
READ_THROUGH_RESYNC_WINDOW_SECONDS: Final = 5.0
|
||||
READ_THROUGH_MAX_RESYNCS_PER_WINDOW: Final = 20
|
||||
|
|
@ -123,47 +130,114 @@ async def _resync_model_deployments(model_name: str) -> bool:
|
|||
return False
|
||||
prisma_client: Final = proxy_server.prisma_client
|
||||
assert prisma_client is not None
|
||||
table: Final = ModelRepository(prisma_client).table
|
||||
name_filter: Final[LiteLLM_ProxyModelTableWhereInput] = {"model_name": model_name}
|
||||
id_filter: Final[LiteLLM_ProxyModelTableWhereInput] = {"model_id": model_name}
|
||||
rows: Final = await table.find_many(where=name_filter) or await table.find_many(where=id_filter)
|
||||
rows: Final = await _find_rows_by_name_or_id(model_name, ModelRepository(prisma_client).table)
|
||||
if not rows:
|
||||
return False
|
||||
router: Final = proxy_server.llm_router
|
||||
if router is None:
|
||||
await proxy_server.proxy_config.add_deployment(
|
||||
prisma_client=prisma_client, proxy_logging_obj=proxy_server.proxy_logging_obj
|
||||
)
|
||||
return proxy_server.llm_router is not None
|
||||
if _holds_an_auto_router(rows):
|
||||
await proxy_server.proxy_config.add_deployment(
|
||||
prisma_client=prisma_client, proxy_logging_obj=proxy_server.proxy_logging_obj
|
||||
)
|
||||
return _model_is_loaded(model_name)
|
||||
async with proxy_server.MODEL_RECONCILE_LOCK:
|
||||
outcome: Final = await _install_model_rows(rows, _always)
|
||||
if outcome in ("no_router", "auto_router"):
|
||||
await _reload_every_model()
|
||||
return outcome == "loaded" or _model_is_loaded(model_name)
|
||||
|
||||
|
||||
async def resync_model_deployments_by_ids(model_ids: tuple[str, ...]) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient
|
||||
from litellm.repositories.model_repository import ModelRepository
|
||||
|
||||
if not _db_backed_registries_enabled("models"):
|
||||
return
|
||||
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:
|
||||
rows: Final = await _find_rows_by_ids(model_ids, writer_table)
|
||||
_ = await _install_model_rows(_without_auto_routers(rows), _names_an_unloaded_credential)
|
||||
|
||||
|
||||
async def _find_rows_by_name_or_id(model_name: str, table: "_ModelTable") -> "_ModelRows":
|
||||
name_filter: Final[LiteLLM_ProxyModelTableWhereInput] = {"model_name": model_name}
|
||||
id_filter: Final[LiteLLM_ProxyModelTableWhereInput] = {"model_id": model_name}
|
||||
return await table.find_many(where=name_filter) or await table.find_many(where=id_filter)
|
||||
|
||||
|
||||
async def _find_rows_by_ids(model_ids: tuple[str, ...], table: "_ModelTable") -> "_ModelRows":
|
||||
from litellm.repositories.chunked_in import find_many_in
|
||||
|
||||
return list(await find_many_in(table, "model_id", model_ids))
|
||||
|
||||
|
||||
async def _install_model_rows(
|
||||
rows: "_ModelRows",
|
||||
needs_credentials: Callable[["_ModelRows"], bool],
|
||||
) -> ModelLoadOutcome:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
router: Final = proxy_server.llm_router
|
||||
if not rows:
|
||||
return "missing"
|
||||
if router is None:
|
||||
return "no_router"
|
||||
if _holds_an_auto_router(rows):
|
||||
return "auto_router"
|
||||
if needs_credentials(rows):
|
||||
prisma_client: Final = proxy_server.prisma_client
|
||||
assert prisma_client is not None
|
||||
await proxy_server.proxy_config.get_credentials(prisma_client=prisma_client)
|
||||
proxy_server.proxy_config._add_deployment(db_models=rows)
|
||||
proxy_server.llm_model_list = router.registry().deployments()
|
||||
proxy_server.proxy_config._add_deployment(db_models=rows)
|
||||
proxy_server.llm_model_list = router.registry().deployments()
|
||||
return "loaded"
|
||||
|
||||
|
||||
async def _reload_every_model() -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
|
||||
prisma_client: Final = proxy_server.prisma_client
|
||||
assert prisma_client is not None
|
||||
await proxy_server.proxy_config.add_deployment(
|
||||
prisma_client=prisma_client, proxy_logging_obj=proxy_server.proxy_logging_obj
|
||||
)
|
||||
|
||||
|
||||
def _always(rows: "_ModelRows") -> bool:
|
||||
return True
|
||||
|
||||
|
||||
def _holds_an_auto_router(rows: Sequence["prisma_models.LiteLLM_ProxyModelTable"]) -> bool:
|
||||
def _holds_an_auto_router(rows: "_ModelRows") -> bool:
|
||||
return any(_names_an_auto_router(row.litellm_params) for row in rows)
|
||||
|
||||
|
||||
def _without_auto_routers(rows: "_ModelRows") -> "_ModelRows":
|
||||
return list(row for row in rows if not _names_an_auto_router(row.litellm_params))
|
||||
|
||||
|
||||
def _names_an_unloaded_credential(rows: "_ModelRows") -> bool:
|
||||
return any(_names_an_unloaded_credential_param(row.litellm_params) for row in rows)
|
||||
|
||||
|
||||
_STORED_LITELLM_PARAMS: Final = TypeAdapter(Mapping[str, object])
|
||||
|
||||
|
||||
def _names_an_auto_router(stored_litellm_params: object) -> bool:
|
||||
def _decrypted_param(stored_litellm_params: object, key: str) -> str | None:
|
||||
try:
|
||||
params: Final = _STORED_LITELLM_PARAMS.validate_python(stored_litellm_params)
|
||||
except ValidationError:
|
||||
return False
|
||||
stored_model: Final = params.get("model")
|
||||
if not isinstance(stored_model, str):
|
||||
return False
|
||||
model: Final = decrypt_value_helper(value=stored_model, key="model", return_original_value=True)
|
||||
return isinstance(model, str) and model.startswith("auto_router/")
|
||||
return None
|
||||
stored_value: Final = params.get(key)
|
||||
if not isinstance(stored_value, str):
|
||||
return None
|
||||
value: Final = decrypt_value_helper(value=stored_value, key=key, return_original_value=True)
|
||||
return value if isinstance(value, str) else None
|
||||
|
||||
|
||||
def _names_an_auto_router(stored_litellm_params: object) -> bool:
|
||||
model: Final = _decrypted_param(stored_litellm_params, "model")
|
||||
return model is not None and model.startswith("auto_router/")
|
||||
|
||||
|
||||
def _names_an_unloaded_credential_param(stored_litellm_params: object) -> bool:
|
||||
credential_name: Final = _decrypted_param(stored_litellm_params, "litellm_credential_name")
|
||||
return credential_name is not None and CredentialAccessor.find_credential(credential_name) is None
|
||||
|
||||
|
||||
async def _resync_guardrails(guardrail_name: str) -> bool:
|
||||
|
|
|
|||
|
|
@ -12,23 +12,30 @@ Endpoints here:
|
|||
- DELETE /access_group/{access_group}/budget - Clear an access group's shared budget
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
import itertools
|
||||
import json
|
||||
from collections.abc import Mapping, Sequence
|
||||
from datetime import datetime
|
||||
import time
|
||||
from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence
|
||||
from datetime import datetime, timedelta
|
||||
from types import MappingProxyType
|
||||
from typing import TYPE_CHECKING, Annotated, Any, Final, Protocol
|
||||
|
||||
from fastapi import APIRouter, Depends, HTTPException
|
||||
from pydantic import BaseModel, ConfigDict, TypeAdapter
|
||||
from typing_extensions import ReadOnly, TypedDict
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.auth.user_api_key_auth import user_api_key_auth
|
||||
from litellm.proxy.common_utils.config_sync_pubsub import publish_model_rows_written
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
UserApiKeyCache,
|
||||
model_access_group_cache_key,
|
||||
model_access_group_registry_cache_key,
|
||||
)
|
||||
from litellm.proxy.db.routing_prisma_wrapper import WriterPinnedClient, writer_wrapper
|
||||
from litellm.proxy.management_endpoints.common_utils import validate_budget_duration
|
||||
|
||||
# Clear cache and reload models to pick up the access group changes
|
||||
|
|
@ -116,6 +123,28 @@ class _ModelTableClient(Protocol):
|
|||
async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ...
|
||||
|
||||
|
||||
class _ModelTableBatchActions(Protocol):
|
||||
def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> None: ...
|
||||
|
||||
|
||||
class _ModelTableBatch(Protocol):
|
||||
@property
|
||||
def litellm_proxymodeltable(self) -> _ModelTableBatchActions: ...
|
||||
|
||||
|
||||
class _ModelTableTxActions(Protocol):
|
||||
async def find_many(self) -> Sequence[_DeploymentRow]: ...
|
||||
|
||||
|
||||
class _ModelTableTx(Protocol):
|
||||
@property
|
||||
def litellm_proxymodeltable(self) -> _ModelTableTxActions: ...
|
||||
|
||||
async def query_raw(self, query: str, *args: object) -> object: ...
|
||||
|
||||
def batch_(self) -> contextlib.AbstractAsyncContextManager[_ModelTableBatch]: ...
|
||||
|
||||
|
||||
class _BudgetRow(Protocol):
|
||||
@property
|
||||
def budget_id(self) -> str: ...
|
||||
|
|
@ -168,7 +197,42 @@ class _ModelAccessGroupBudgetTableClient(Protocol):
|
|||
|
||||
|
||||
def _model_table(prisma_client: PrismaClient) -> _ModelTableClient:
|
||||
return ModelRepository(prisma_client).table
|
||||
return ModelRepository(WriterPinnedClient(prisma_client.db)).table
|
||||
|
||||
|
||||
_TRY_LOCK_ACCESS_GROUP: Final = (
|
||||
"SELECT pg_try_advisory_xact_lock(hashtext('model_access_group'), hashtext($1)) AS locked"
|
||||
)
|
||||
_LOCK_MODEL_ROWS: Final = (
|
||||
'SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = ANY($1::text[]) '
|
||||
"ORDER BY model_id FOR UPDATE SKIP LOCKED"
|
||||
)
|
||||
_MEMBERSHIP_TX_TIMEOUT: Final = timedelta(seconds=30)
|
||||
_LOCK_RETRY_BUDGET_SECONDS: Final = _MEMBERSHIP_TX_TIMEOUT.total_seconds()
|
||||
_LOCK_RETRY_BACKOFF_SECONDS: Final = (0.05, 0.1, 0.2, 0.4, 0.8)
|
||||
_LOCK_RETRY_LONGEST_DELAY_SECONDS: Final = 1.0
|
||||
|
||||
|
||||
class _LockedModelRow(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
model_id: str
|
||||
|
||||
|
||||
_LOCKED_MODEL_ROWS: Final = TypeAdapter(tuple[_LockedModelRow, ...])
|
||||
|
||||
|
||||
class _AccessGroupLock(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
locked: bool
|
||||
|
||||
|
||||
_ACCESS_GROUP_LOCK: Final = TypeAdapter(tuple[_AccessGroupLock])
|
||||
|
||||
|
||||
def _model_table_tx(prisma_client: PrismaClient) -> contextlib.AbstractAsyncContextManager[_ModelTableTx]:
|
||||
return writer_wrapper(prisma_client.db).tx(timeout=_MEMBERSHIP_TX_TIMEOUT)
|
||||
|
||||
|
||||
def _model_access_group_budget_table(prisma_client: PrismaClient) -> _ModelAccessGroupBudgetTableClient:
|
||||
|
|
@ -379,11 +443,10 @@ async def _tag_deployment_with_access_group(
|
|||
"""Write `access_group` into one deployment's model_info; returns the
|
||||
(model_id, updated model_info) pair when a write happened, None when the
|
||||
deployment already carried the group."""
|
||||
updated_model_info, was_modified = add_access_group_to_deployment(
|
||||
model_info=dict(_readable_model_info_or_raise(model_id=model_id, model_info=model_info)),
|
||||
access_group=access_group,
|
||||
updated_model_info: Final = _model_info_with_access_group(
|
||||
model_id=model_id, model_info=model_info, access_group=access_group
|
||||
)
|
||||
if not was_modified:
|
||||
if updated_model_info is None:
|
||||
return None
|
||||
await _model_table(prisma_client).update(
|
||||
where={"model_id": model_id},
|
||||
|
|
@ -393,6 +456,141 @@ async def _tag_deployment_with_access_group(
|
|||
return (model_id, updated_model_info)
|
||||
|
||||
|
||||
def _model_info_with_access_group(model_id: str, model_info: object, access_group: str) -> Mapping[str, object] | None:
|
||||
updated_model_info, was_modified = add_access_group_to_deployment(
|
||||
model_info=dict(_readable_model_info_or_raise(model_id=model_id, model_info=model_info)),
|
||||
access_group=access_group,
|
||||
)
|
||||
return updated_model_info if was_modified else None
|
||||
|
||||
|
||||
def _model_info_without_access_group(
|
||||
model_id: str, model_info: object, access_group: str
|
||||
) -> Mapping[str, object] | None:
|
||||
updated_model_info, was_modified = remove_access_group_from_deployment(
|
||||
model_info=dict(_readable_model_info_or_raise(model_id=model_id, model_info=model_info)),
|
||||
access_group=access_group,
|
||||
)
|
||||
return updated_model_info if was_modified else None
|
||||
|
||||
|
||||
def _membership_rewrite(deployment: _DeploymentRow, selected: bool, access_group: str) -> Mapping[str, object] | None:
|
||||
if selected:
|
||||
return _model_info_with_access_group(
|
||||
model_id=deployment.model_id, model_info=deployment.model_info, access_group=access_group
|
||||
)
|
||||
return _model_info_without_access_group(
|
||||
model_id=deployment.model_id, model_info=deployment.model_info, access_group=access_group
|
||||
)
|
||||
|
||||
|
||||
def _membership_rewrites(
|
||||
deployments: Sequence[_DeploymentRow], selected_model_ids: frozenset[str], access_group: str
|
||||
) -> tuple[tuple[str, Mapping[str, object]], ...]:
|
||||
rewrites: Final = (
|
||||
(
|
||||
deployment.model_id,
|
||||
_membership_rewrite(
|
||||
deployment=deployment,
|
||||
selected=deployment.model_id in selected_model_ids,
|
||||
access_group=access_group,
|
||||
),
|
||||
)
|
||||
for deployment in deployments
|
||||
)
|
||||
return tuple((model_id, model_info) for model_id, model_info in rewrites if model_info is not None)
|
||||
|
||||
|
||||
def _carries_access_group(model_info: object, access_group: str) -> bool:
|
||||
access_groups: Final = (model_info_as_mapping(model_info) or {}).get("access_groups")
|
||||
return isinstance(access_groups, list) and access_group in access_groups
|
||||
|
||||
|
||||
def _model_ids_to_lock(
|
||||
deployments: Sequence[_DeploymentRow], access_group: str, data: UpdateModelGroupRequest
|
||||
) -> frozenset[str]:
|
||||
selected_model_ids: Final = frozenset(data.model_ids or ())
|
||||
selected_model_names: Final = frozenset(() if data.model_ids else data.model_names or ())
|
||||
return frozenset(
|
||||
deployment.model_id
|
||||
for deployment in deployments
|
||||
if deployment.model_id in selected_model_ids
|
||||
or deployment.model_name in selected_model_names
|
||||
or _carries_access_group(deployment.model_info, access_group)
|
||||
)
|
||||
|
||||
|
||||
async def _raise_404_unless_the_primary_carries(access_group: str, prisma_client: PrismaClient) -> None:
|
||||
if access_group not in _access_groups_of_deployments(await _model_table(prisma_client).find_many()):
|
||||
raise _http_error(404, f"Access group '{access_group}' not found")
|
||||
|
||||
|
||||
async def _rewrite_membership_unless_locked(
|
||||
prisma_client: PrismaClient, access_group: str, data: UpdateModelGroupRequest
|
||||
) -> tuple[tuple[_DeploymentRow, ...], tuple[tuple[str, Mapping[str, object]], ...]] | None:
|
||||
async with _model_table_tx(prisma_client) as tx:
|
||||
(group_lock,) = _ACCESS_GROUP_LOCK.validate_python(await tx.query_raw(_TRY_LOCK_ACCESS_GROUP, access_group))
|
||||
if not group_lock.locked:
|
||||
return None
|
||||
model_ids_to_lock: Final = _model_ids_to_lock(
|
||||
deployments=await tx.litellm_proxymodeltable.find_many(), access_group=access_group, data=data
|
||||
)
|
||||
locked_rows: Final = _LOCKED_MODEL_ROWS.validate_python(
|
||||
await tx.query_raw(_LOCK_MODEL_ROWS, sorted(model_ids_to_lock))
|
||||
)
|
||||
if frozenset(row.model_id for row in locked_rows) != model_ids_to_lock:
|
||||
return None
|
||||
deployments: Final = tuple(
|
||||
deployment
|
||||
for deployment in await tx.litellm_proxymodeltable.find_many()
|
||||
if deployment.model_id in model_ids_to_lock
|
||||
)
|
||||
if access_group not in _access_groups_of_deployments(deployments):
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail={"error": f"Access group '{access_group}' not found"},
|
||||
)
|
||||
selected_deployments: Final = _selected_deployments_or_400(deployments=deployments, data=data)
|
||||
rewrites: Final = _membership_rewrites(
|
||||
deployments=deployments,
|
||||
selected_model_ids=frozenset(deployment.model_id for deployment in selected_deployments),
|
||||
access_group=access_group,
|
||||
)
|
||||
if rewrites:
|
||||
async with tx.batch_() as batch:
|
||||
for model_id, model_info in rewrites:
|
||||
batch.litellm_proxymodeltable.update(
|
||||
where={"model_id": model_id}, data={"model_info": json.dumps(model_info)}
|
||||
)
|
||||
return selected_deployments, rewrites
|
||||
|
||||
|
||||
def _lock_retry_delays() -> Iterator[float]:
|
||||
return itertools.chain(_LOCK_RETRY_BACKOFF_SECONDS, itertools.repeat(_LOCK_RETRY_LONGEST_DELAY_SECONDS))
|
||||
|
||||
|
||||
async def _rewrite_membership_once_unlocked(
|
||||
prisma_client: PrismaClient,
|
||||
access_group: str,
|
||||
data: UpdateModelGroupRequest,
|
||||
retry_budget_seconds: float = _LOCK_RETRY_BUDGET_SECONDS,
|
||||
sleep: Callable[[float], Awaitable[None]] = asyncio.sleep,
|
||||
clock: Callable[[], float] = time.monotonic,
|
||||
) -> tuple[tuple[_DeploymentRow, ...], tuple[tuple[str, Mapping[str, object]], ...]]:
|
||||
deadline: Final = clock() + retry_budget_seconds
|
||||
for delay in _lock_retry_delays():
|
||||
if (
|
||||
rewritten := await _rewrite_membership_unless_locked(
|
||||
prisma_client=prisma_client, access_group=access_group, data=data
|
||||
)
|
||||
) is not None:
|
||||
return rewritten
|
||||
if clock() + delay > deadline:
|
||||
break
|
||||
await sleep(delay)
|
||||
raise _http_error(409, f"Another edit holds access group '{access_group}' or one of its models, try again")
|
||||
|
||||
|
||||
def _readable_model_info_or_raise(model_id: str, model_info: object) -> Mapping[str, object]:
|
||||
"""These helpers rewrite the model_info column wholesale, so a present-but-unreadable
|
||||
value must refuse loudly rather than be silently replaced with a fresh object; an
|
||||
|
|
@ -412,11 +610,10 @@ async def _strip_access_group_from_deployment(
|
|||
"""Remove `access_group` from one deployment's model_info; returns the
|
||||
(model_id, updated model_info) pair when a write happened, None when the
|
||||
deployment did not carry the group."""
|
||||
updated_model_info, was_modified = remove_access_group_from_deployment(
|
||||
model_info=dict(_readable_model_info_or_raise(model_id=model_id, model_info=model_info)),
|
||||
access_group=access_group,
|
||||
updated_model_info: Final = _model_info_without_access_group(
|
||||
model_id=model_id, model_info=model_info, access_group=access_group
|
||||
)
|
||||
if not was_modified:
|
||||
if updated_model_info is None:
|
||||
return None
|
||||
await _model_table(prisma_client).update(
|
||||
where={"model_id": model_id},
|
||||
|
|
@ -493,6 +690,35 @@ async def update_specific_deployments_with_access_group(
|
|||
return tuple(pair for pair in tagged if pair is not None)
|
||||
|
||||
|
||||
def _selected_deployments_or_400(
|
||||
deployments: Sequence[_DeploymentRow], data: UpdateModelGroupRequest
|
||||
) -> tuple[_DeploymentRow, ...]:
|
||||
if data.model_ids:
|
||||
known_model_ids: Final = frozenset(deployment.model_id for deployment in deployments)
|
||||
unknown_model_id: Final = next(
|
||||
(model_id for model_id in data.model_ids if model_id not in known_model_ids), None
|
||||
)
|
||||
if unknown_model_id is not None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={"error": f"Deployment with model_id '{unknown_model_id}' not found in Database."},
|
||||
)
|
||||
selected_model_ids: Final = frozenset(data.model_ids)
|
||||
return tuple(deployment for deployment in deployments if deployment.model_id in selected_model_ids)
|
||||
model_names: Final = data.model_names or ()
|
||||
known_model_names: Final = frozenset(deployment.model_name for deployment in deployments)
|
||||
unknown_model_name: Final = next((name for name in model_names if name not in known_model_names), None)
|
||||
if unknown_model_name is not None:
|
||||
raise HTTPException(
|
||||
status_code=400,
|
||||
detail={
|
||||
"error": f"Can't find model '{unknown_model_name}' in Database. Access group management is only supported for database models."
|
||||
},
|
||||
)
|
||||
selected_model_names: Final = frozenset(model_names)
|
||||
return tuple(deployment for deployment in deployments if deployment.model_name in selected_model_names)
|
||||
|
||||
|
||||
async def _find_deployment_or_400(model_id: str, prisma_client: PrismaClient) -> object:
|
||||
deployment: Final = await _model_table(prisma_client).find_unique(where={"model_id": model_id})
|
||||
if deployment is None:
|
||||
|
|
@ -536,10 +762,10 @@ async def get_all_access_groups_from_db(
|
|||
Returns:
|
||||
Dict[str, AccessGroupInfo]: Dictionary mapping access_group name to info
|
||||
"""
|
||||
# Get all deployments
|
||||
deployments: Final = await ModelRepository(prisma_client).table.find_many()
|
||||
return dict(_access_groups_of_deployments(await ModelRepository(prisma_client).table.find_many()))
|
||||
|
||||
# Build access group map
|
||||
|
||||
def _access_groups_of_deployments(deployments: Sequence[_DeploymentRow]) -> Mapping[str, AccessGroupInfo]:
|
||||
model_names_by_group: Final[dict[str, list[str]]] = {}
|
||||
|
||||
for deployment in deployments:
|
||||
|
|
@ -835,8 +1061,8 @@ async def update_access_group(
|
|||
Update an access group's model names.
|
||||
|
||||
This will:
|
||||
1. Remove the access group from all current deployments
|
||||
2. Add the access group to all deployments for the new model_names list
|
||||
1. Add the access group to all deployments for the new model_names list
|
||||
2. Remove the access group from every other deployment
|
||||
|
||||
Example:
|
||||
```bash
|
||||
|
|
@ -858,6 +1084,7 @@ async def update_access_group(
|
|||
Raises:
|
||||
- HTTPException 400: If any model names don't exist
|
||||
- HTTPException 404: If access group not found
|
||||
- HTTPException 409: If another edit holds the group or one of its models for 30 s
|
||||
"""
|
||||
from litellm.proxy.proxy_server import llm_router, prisma_client
|
||||
|
||||
|
|
@ -879,26 +1106,14 @@ async def update_access_group(
|
|||
detail={"error": "Either model_names or model_ids must be provided and non-empty"},
|
||||
)
|
||||
|
||||
use_model_ids: Final = has_model_ids
|
||||
|
||||
# Validation: Check if access group exists
|
||||
try:
|
||||
access_groups_map: Final = await get_all_access_groups_from_db(prisma_client=prisma_client)
|
||||
if access_group not in access_groups_map:
|
||||
raise HTTPException(
|
||||
status_code=404,
|
||||
detail={"error": f"Access group '{access_group}' not found"},
|
||||
)
|
||||
await _raise_404_unless_the_primary_carries(access_group=access_group, prisma_client=prisma_client)
|
||||
except HTTPException:
|
||||
raise
|
||||
except Exception as e:
|
||||
raise HTTPException(
|
||||
status_code=500,
|
||||
detail={"error": f"Failed to check access group existence: {e}"},
|
||||
)
|
||||
raise _http_error(500, f"Failed to check access group existence: {e}")
|
||||
|
||||
# Validation: Check if all new models exist (only if using model_names path)
|
||||
if not use_model_ids and has_model_names:
|
||||
if not has_model_ids and has_model_names:
|
||||
assert data.model_names is not None
|
||||
missing_models: Final = await _missing_models_after_read_through(
|
||||
model_names=data.model_names,
|
||||
|
|
@ -912,43 +1127,19 @@ async def update_access_group(
|
|||
)
|
||||
|
||||
try:
|
||||
# Step 1: Remove access group from ALL DB deployments (skip config models)
|
||||
all_deployments: Final = await _model_table(prisma_client).find_many()
|
||||
|
||||
stripped: Final = [
|
||||
await _strip_access_group_from_deployment(
|
||||
model_id=deployment.model_id,
|
||||
model_info=deployment.model_info,
|
||||
access_group=access_group,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
for deployment in all_deployments
|
||||
]
|
||||
stripped_pairs: Final = tuple(pair for pair in stripped if pair is not None)
|
||||
|
||||
# Step 2: Add access group using the appropriate method
|
||||
if use_model_ids:
|
||||
assert data.model_ids is not None
|
||||
updated_pairs = await update_specific_deployments_with_access_group(
|
||||
model_ids=data.model_ids,
|
||||
access_group=access_group,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
else:
|
||||
assert data.model_names is not None
|
||||
updated_pairs = await update_deployments_with_access_group(
|
||||
model_names=data.model_names,
|
||||
access_group=access_group,
|
||||
prisma_client=prisma_client,
|
||||
)
|
||||
models_updated: Final = len(updated_pairs)
|
||||
selected_deployments, rewrites = await _rewrite_membership_once_unlocked(
|
||||
prisma_client=prisma_client, access_group=access_group, data=data
|
||||
)
|
||||
if rewrites:
|
||||
await publish_model_rows_written(tuple(model_id for model_id, _ in rewrites))
|
||||
models_updated: Final = len(selected_deployments)
|
||||
|
||||
# Clear cache and reload models to pick up the access group changes
|
||||
live_before_reload: Final = live_model_ids_snapshot()
|
||||
reload_outcome: Final = await clear_cache()
|
||||
_raise_http_if_reload_degraded_serving(
|
||||
before=live_before_reload,
|
||||
written_models=list({**dict(stripped_pairs), **dict(updated_pairs)}.items()),
|
||||
written_models=rewrites,
|
||||
access_group=access_group,
|
||||
still_desired=reload_outcome.still_desired,
|
||||
live_after=reload_outcome.live_after,
|
||||
|
|
|
|||
|
|
@ -779,7 +779,11 @@ async def _auto_router_capability_slot(
|
|||
limit=limit,
|
||||
)
|
||||
yield tables.litellm_proxymodeltable
|
||||
await publish_config_change(redis_cache=coordination_redis_cache(), object_type="litellm_proxymodeltable")
|
||||
await publish_config_change(
|
||||
redis_cache=coordination_redis_cache(),
|
||||
object_type="litellm_proxymodeltable",
|
||||
model_ids=() if model_id is None else (model_id,),
|
||||
)
|
||||
|
||||
|
||||
ENFORCE_RPM_TPM_ON_MODEL_ADD_SETTING: Final = "enforce_rpm_tpm_on_model_add"
|
||||
|
|
|
|||
|
|
@ -549,6 +549,7 @@ from litellm.proxy.common_utils.periodic_reload_schedule import (
|
|||
write_reload_interval,
|
||||
)
|
||||
from litellm.proxy.common_utils.proxy_state import ProxyState
|
||||
from litellm.proxy.common_utils.registry_read_through import resync_model_deployments_by_ids
|
||||
from litellm.proxy.common_utils.reset_budget_job import ResetBudgetJob
|
||||
from litellm.proxy.common_utils.responses_stream_errors import ResponsesStreamErrorState
|
||||
from litellm.proxy.common_utils.scheduled_job_stagger import (
|
||||
|
|
@ -8750,6 +8751,7 @@ class ProxyConfig:
|
|||
subscriber: Final = ConfigSyncSubscriber(
|
||||
redis_cache=redis_cache,
|
||||
resync_callbacks=(_resync_config_from_db,),
|
||||
apply_model_ids=resync_model_deployments_by_ids,
|
||||
)
|
||||
self.config_sync_subscriber = subscriber
|
||||
subscriber.start()
|
||||
|
|
|
|||
|
|
@ -0,0 +1,109 @@
|
|||
import uuid
|
||||
from collections.abc import Iterator
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
from pydantic import TypeAdapter
|
||||
from redis import Redis
|
||||
|
||||
from tests.integration._support.client import Gateway, eventually, list_value, object_value, string_value
|
||||
from tests.integration._support.database import scratch_database
|
||||
from tests.integration._support.process import graceful_stop_seconds, owned_proxy
|
||||
from tests.integration._support.redis_process import owned_redis
|
||||
|
||||
pytestmark: Final = [pytest.mark.timeout(int(4 * graceful_stop_seconds() + 120))]
|
||||
|
||||
COORDINATION_CONFIG: Final = Path(__file__).resolve().parents[1] / "coordination_redis_proxy_config.yaml"
|
||||
CONFIG_SYNC_CHANNEL: Final = "litellm_proxy.config_change"
|
||||
SUBSCRIBER_COUNTS: Final = TypeAdapter(tuple[tuple[str, int], ...])
|
||||
PERIODIC_RELOAD_SECONDS: Final = 3600
|
||||
PEER_PICKUP_SECONDS: Final = 4
|
||||
PROVIDER_KEY: Final = "integration-provider-key"
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _proxy_pair(gateway: Gateway, directory: Path) -> Iterator[tuple[Gateway, Gateway]]:
|
||||
with scratch_database("integration_model_add_sync") as database, owned_redis(directory) as coordination:
|
||||
shared_stores: Final = {
|
||||
"DATABASE_URL": database,
|
||||
"REDIS_HOST": coordination.host,
|
||||
"REDIS_PORT": str(coordination.port),
|
||||
"PROXY_CONFIG_RELOAD_INTERVAL_SECONDS": str(PERIODIC_RELOAD_SECONDS),
|
||||
}
|
||||
with (
|
||||
owned_proxy(gateway, directory / "writer", shared_stores, config=COORDINATION_CONFIG) as writer,
|
||||
owned_proxy(gateway, directory / "peer", shared_stores, config=COORDINATION_CONFIG) as peer,
|
||||
Redis(host=coordination.host, port=coordination.port, socket_timeout=2) as redis,
|
||||
):
|
||||
eventually(
|
||||
lambda: SUBSCRIBER_COUNTS.validate_python(redis.pubsub_numsub(CONFIG_SYNC_CHANNEL))[0][1],
|
||||
lambda count: count >= 2,
|
||||
seconds=30,
|
||||
)
|
||||
yield writer, peer
|
||||
|
||||
|
||||
@contextmanager
|
||||
def _failing_upstream_model(gateway: Gateway, upstream_model: str) -> Iterator[None]:
|
||||
script: Final = f"{gateway.upstream_url}/__scripts/{upstream_model}"
|
||||
httpx.post(script, json={"statuses": [500]}, timeout=5, trust_env=False).raise_for_status()
|
||||
try:
|
||||
yield
|
||||
finally:
|
||||
httpx.delete(script, timeout=5, trust_env=False).raise_for_status()
|
||||
|
||||
|
||||
def _add_deployment(writer: Gateway, group: str, upstream_model: str, order: int) -> None:
|
||||
writer.post(
|
||||
"/model/new",
|
||||
{
|
||||
"model_name": group,
|
||||
"litellm_params": {
|
||||
"model": f"openai/{upstream_model}",
|
||||
"api_key": PROVIDER_KEY,
|
||||
"api_base": f"{writer.upstream_url}/v1",
|
||||
"order": order,
|
||||
"num_retries": 0,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
|
||||
def _listed_models(proxy: Gateway) -> frozenset[str]:
|
||||
return frozenset(string_value(object_value(entry)["id"]) for entry in list_value(proxy.get("/v1/models")["data"]))
|
||||
|
||||
|
||||
def _chat(proxy: Gateway, group: str) -> httpx.Response:
|
||||
return proxy.request(
|
||||
"POST",
|
||||
"/v1/chat/completions",
|
||||
{"model": group, "messages": [{"role": "user", "content": f"order fallback {uuid.uuid4().hex}"}]},
|
||||
)
|
||||
|
||||
|
||||
def test_a_backup_added_on_one_proxy_catches_the_peers_failing_primary_right_away(
|
||||
gateway: Gateway, tmp_path: Path
|
||||
) -> None:
|
||||
group: Final = f"integration-{uuid.uuid4().hex}"
|
||||
primary: Final = f"integration-primary-{uuid.uuid4().hex}"
|
||||
backup: Final = f"integration-backup-{uuid.uuid4().hex}"
|
||||
with _proxy_pair(gateway, tmp_path) as (writer, peer), _failing_upstream_model(gateway, primary):
|
||||
_add_deployment(writer, group, primary, order=1)
|
||||
eventually(lambda: _listed_models(peer), lambda listed: group in listed, seconds=30)
|
||||
|
||||
_add_deployment(writer, group, backup, order=2)
|
||||
on_writer: Final = _chat(writer, group)
|
||||
on_peer: Final = eventually(
|
||||
lambda: _chat(peer, group),
|
||||
lambda response: response.status_code == 200,
|
||||
seconds=PEER_PICKUP_SECONDS,
|
||||
return_last_on_timeout=True,
|
||||
)
|
||||
|
||||
assert on_writer.status_code == 200, on_writer.text
|
||||
assert object_value(on_writer.json())["model"] == backup, on_writer.text
|
||||
assert on_peer.status_code == 200, on_peer.text
|
||||
assert object_value(on_peer.json())["model"] == backup, on_peer.text
|
||||
|
|
@ -5,6 +5,7 @@ from typing import Callable, Coroutine, Final, Iterable, List, Optional, Tuple
|
|||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from pydantic import BaseModel, ConfigDict
|
||||
from redis.asyncio import Redis
|
||||
|
||||
import litellm
|
||||
|
|
@ -102,8 +103,8 @@ class _ScriptedPubSubClient:
|
|||
|
||||
|
||||
class _QueuePubSub:
|
||||
def __init__(self, initial_messages: Iterable[str] = ()) -> None:
|
||||
self.queue: "asyncio.Queue[str]" = asyncio.Queue()
|
||||
def __init__(self, initial_messages: Iterable[object] = ()) -> None:
|
||||
self.queue: "asyncio.Queue[object]" = asyncio.Queue()
|
||||
for message in initial_messages:
|
||||
self.queue.put_nowait(message)
|
||||
self.subscribed_channels: List[str] = []
|
||||
|
|
@ -112,7 +113,7 @@ class _QueuePubSub:
|
|||
async def subscribe(self, *channels: str) -> None:
|
||||
self.subscribed_channels.extend(channels)
|
||||
|
||||
async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[str]:
|
||||
async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[object]:
|
||||
if timeout == 0:
|
||||
try:
|
||||
return self.queue.get_nowait()
|
||||
|
|
@ -128,7 +129,7 @@ class _QueuePubSub:
|
|||
|
||||
|
||||
class _BrokenPubSub(_QueuePubSub):
|
||||
async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[str]:
|
||||
async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[object]:
|
||||
raise ConnectionError("connection lost")
|
||||
|
||||
|
||||
|
|
@ -138,11 +139,11 @@ class _CloseFailingBrokenPubSub(_BrokenPubSub):
|
|||
|
||||
|
||||
class _EmptyPollsThenMessagePubSub(_QueuePubSub):
|
||||
def __init__(self, empty_polls: int, initial_messages: Iterable[str] = ()) -> None:
|
||||
def __init__(self, empty_polls: int, initial_messages: Iterable[object] = ()) -> None:
|
||||
super().__init__(initial_messages=initial_messages)
|
||||
self.remaining_empty_polls = empty_polls
|
||||
|
||||
async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[str]:
|
||||
async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[object]:
|
||||
if timeout != 0 and self.remaining_empty_polls > 0:
|
||||
self.remaining_empty_polls -= 1
|
||||
return None
|
||||
|
|
@ -157,6 +158,38 @@ class _FakeClock:
|
|||
return self.now
|
||||
|
||||
|
||||
class _VirtualTimePubSub(_QueuePubSub):
|
||||
"""Timed polls move the injected clock forward instead of waiting on the wall clock.
|
||||
|
||||
A delivery scheduled at a clock time is returned by the first poll whose window
|
||||
reaches it, with the clock set to that time, so a test asserts when each resync
|
||||
and targeted apply ran in subscriber time.
|
||||
"""
|
||||
|
||||
def __init__(self, clock: _FakeClock, deliveries: Iterable[Tuple[float, object]] = ()) -> None:
|
||||
super().__init__()
|
||||
self._clock = clock
|
||||
self.pending: List[Tuple[float, object]] = sorted(deliveries, key=lambda delivery: delivery[0])
|
||||
|
||||
async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[object]:
|
||||
await asyncio.sleep(0)
|
||||
due_by = self._clock.now + timeout
|
||||
if self.pending and self.pending[0][0] <= due_by:
|
||||
delivered_at, message = self.pending.pop(0)
|
||||
self._clock.now = max(self._clock.now, delivered_at)
|
||||
return message
|
||||
self._clock.now = due_by
|
||||
return None
|
||||
|
||||
|
||||
def _change(*model_ids: str) -> str:
|
||||
return json.dumps({"object_type": "litellm_proxymodeltable", "model_ids": list(model_ids)})
|
||||
|
||||
|
||||
def _redis_delivery(message: str) -> dict:
|
||||
return {"type": "message", "pattern": None, "channel": b"litellm_proxy.config_change", "data": message.encode()}
|
||||
|
||||
|
||||
class _ScriptedPubSubRedisClient(Redis):
|
||||
def __init__(self, pubsubs: Iterable[_QueuePubSub]) -> None:
|
||||
self._scripted_pubsubs = iter(pubsubs)
|
||||
|
|
@ -204,7 +237,16 @@ async def test_publish_sends_object_type_json_on_channel() -> None:
|
|||
assert len(client.published) == 1
|
||||
channel, message = client.published[0]
|
||||
assert channel == "litellm_proxy.config_change"
|
||||
assert json.loads(message) == {"object_type": "litellm_proxymodeltable"}
|
||||
assert json.loads(message) == {"object_type": "litellm_proxymodeltable", "model_ids": []}
|
||||
|
||||
|
||||
async def test_publish_carries_the_written_model_ids() -> None:
|
||||
client = _RecordingRedisClient()
|
||||
cache = _FakeRedisCache(client)
|
||||
|
||||
await publish_config_change(redis_cache=cache, object_type="litellm_proxymodeltable", model_ids=("m-1", "m-2"))
|
||||
|
||||
assert json.loads(client.published[0][1]) == {"object_type": "litellm_proxymodeltable", "model_ids": ["m-1", "m-2"]}
|
||||
|
||||
|
||||
async def test_publish_uses_namespaced_channel() -> None:
|
||||
|
|
@ -235,7 +277,7 @@ async def test_publish_reaches_cluster_derived_pubsub_clients() -> None:
|
|||
await publish_config_change(redis_cache=cache, object_type="litellm_proxymodeltable")
|
||||
|
||||
assert client.published == [
|
||||
(CONFIG_SYNC_CHANNEL, json.dumps({"object_type": "litellm_proxymodeltable"}))
|
||||
(CONFIG_SYNC_CHANNEL, json.dumps({"object_type": "litellm_proxymodeltable", "model_ids": []}))
|
||||
]
|
||||
|
||||
|
||||
|
|
@ -312,29 +354,210 @@ class _MaxJitterRandom(random.Random):
|
|||
return b
|
||||
|
||||
|
||||
async def test_debounce_sleep_adds_jitter_from_injected_rng() -> None:
|
||||
pubsub = _QueuePubSub(initial_messages=[json.dumps({"object_type": "litellm_proxymodeltable"})])
|
||||
cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub]))
|
||||
sleeps: List[float] = []
|
||||
fired = asyncio.Event()
|
||||
def _virtual_time_subscriber(
|
||||
pubsub: _VirtualTimePubSub,
|
||||
clock: _FakeClock,
|
||||
events: List[str],
|
||||
resynced: asyncio.Event,
|
||||
*,
|
||||
debounce_seconds: float = 0.0,
|
||||
jitter_max_seconds: float = 0.0,
|
||||
min_resync_interval_seconds: float = 10.0,
|
||||
rng: Optional[random.Random] = None,
|
||||
expected_resyncs: int = 1,
|
||||
apply_error: Optional[Exception] = None,
|
||||
) -> ConfigSyncSubscriber:
|
||||
async def resync() -> None:
|
||||
events.append(f"resync@{clock.now}")
|
||||
if sum(event.startswith("resync@") for event in events) >= expected_resyncs:
|
||||
resynced.set()
|
||||
|
||||
async def recording_sleep(seconds: float) -> None:
|
||||
sleeps.append(seconds)
|
||||
async def apply(model_ids: Tuple[str, ...]) -> None:
|
||||
events.append(f"apply:{','.join(model_ids)}@{clock.now}")
|
||||
if apply_error is not None:
|
||||
raise apply_error
|
||||
|
||||
subscriber = ConfigSyncSubscriber(
|
||||
redis_cache=cache,
|
||||
resync_callbacks=(_recording_callback([], "resync", fired),),
|
||||
debounce_seconds=1.0,
|
||||
jitter_max_seconds=4.0,
|
||||
rng=_MaxJitterRandom(),
|
||||
sleep=recording_sleep,
|
||||
return ConfigSyncSubscriber(
|
||||
redis_cache=_FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])),
|
||||
resync_callbacks=(resync,),
|
||||
apply_model_ids=apply,
|
||||
debounce_seconds=debounce_seconds,
|
||||
jitter_max_seconds=jitter_max_seconds,
|
||||
min_resync_interval_seconds=min_resync_interval_seconds,
|
||||
rng=rng,
|
||||
monotonic=clock,
|
||||
)
|
||||
|
||||
subscriber.start()
|
||||
await asyncio.wait_for(fired.wait(), timeout=5)
|
||||
await subscriber.stop()
|
||||
|
||||
assert sleeps == [5.0]
|
||||
async def _run_until_resynced(subscriber: ConfigSyncSubscriber, resynced: asyncio.Event) -> None:
|
||||
subscriber.start()
|
||||
try:
|
||||
await asyncio.wait_for(resynced.wait(), timeout=5)
|
||||
finally:
|
||||
await subscriber.stop()
|
||||
|
||||
|
||||
async def test_debounce_waits_out_the_jitter_from_injected_rng() -> None:
|
||||
clock = _FakeClock()
|
||||
pubsub = _VirtualTimePubSub(clock, deliveries=[(1000.0, _change())])
|
||||
events: List[str] = []
|
||||
resynced = asyncio.Event()
|
||||
subscriber = _virtual_time_subscriber(
|
||||
pubsub, clock, events, resynced, debounce_seconds=1.0, jitter_max_seconds=4.0, rng=_MaxJitterRandom()
|
||||
)
|
||||
|
||||
await _run_until_resynced(subscriber, resynced)
|
||||
|
||||
assert events == ["resync@1005.0"]
|
||||
|
||||
|
||||
async def test_published_model_ids_apply_on_receipt_while_the_full_resync_still_debounces() -> None:
|
||||
clock = _FakeClock()
|
||||
pubsub = _VirtualTimePubSub(clock, deliveries=[(1000.0, _change("m-1"))])
|
||||
events: List[str] = []
|
||||
resynced = asyncio.Event()
|
||||
subscriber = _virtual_time_subscriber(
|
||||
pubsub, clock, events, resynced, debounce_seconds=1.0, jitter_max_seconds=4.0, rng=_MaxJitterRandom()
|
||||
)
|
||||
|
||||
await _run_until_resynced(subscriber, resynced)
|
||||
|
||||
assert events == ["apply:m-1@1000.0", "resync@1005.0"]
|
||||
|
||||
|
||||
async def test_model_ids_published_inside_the_throttle_wait_apply_without_waiting() -> None:
|
||||
clock = _FakeClock()
|
||||
pubsub = _VirtualTimePubSub(
|
||||
clock, deliveries=[(1000.0, _change()), (1004.0, _change("m-2")), (1006.0, _change("m-3", "m-4"))]
|
||||
)
|
||||
events: List[str] = []
|
||||
resynced = asyncio.Event()
|
||||
subscriber = _virtual_time_subscriber(pubsub, clock, events, resynced, expected_resyncs=2)
|
||||
|
||||
await _run_until_resynced(subscriber, resynced)
|
||||
|
||||
assert events == ["resync@1000.0", "apply:m-2@1004.0", "apply:m-3,m-4@1006.0", "resync@1010.0"]
|
||||
|
||||
|
||||
async def test_model_ids_already_waiting_apply_together_in_one_call() -> None:
|
||||
clock = _FakeClock()
|
||||
pubsub = _VirtualTimePubSub(
|
||||
clock, deliveries=[(1000.0, _change("m-1")), (1000.0, _change("m-2", "m-1")), (1000.0, _change("m-3"))]
|
||||
)
|
||||
events: List[str] = []
|
||||
resynced = asyncio.Event()
|
||||
subscriber = _virtual_time_subscriber(pubsub, clock, events, resynced)
|
||||
|
||||
await _run_until_resynced(subscriber, resynced)
|
||||
|
||||
assert events == ["apply:m-1,m-2,m-3@1000.0", "resync@1000.0"]
|
||||
|
||||
|
||||
async def test_redis_shaped_message_applies_its_model_ids() -> None:
|
||||
clock = _FakeClock()
|
||||
pubsub = _VirtualTimePubSub(clock, deliveries=[(1000.0, _redis_delivery(_change("m-1")))])
|
||||
events: List[str] = []
|
||||
resynced = asyncio.Event()
|
||||
subscriber = _virtual_time_subscriber(pubsub, clock, events, resynced)
|
||||
|
||||
await _run_until_resynced(subscriber, resynced)
|
||||
|
||||
assert events == ["apply:m-1@1000.0", "resync@1000.0"]
|
||||
|
||||
|
||||
async def test_messages_without_usable_model_ids_only_schedule_the_full_resync() -> None:
|
||||
clock = _FakeClock()
|
||||
pubsub = _VirtualTimePubSub(
|
||||
clock,
|
||||
deliveries=[
|
||||
(1000.0, "change"),
|
||||
(1004.0, json.dumps({"object_type": "litellm_proxymodeltable"})),
|
||||
(1005.0, json.dumps({"object_type": "litellm_proxymodeltable", "model_ids": "m-1"})),
|
||||
(1006.0, {"type": "message", "pattern": None, "channel": b"c", "data": 7}),
|
||||
(1007.0, None),
|
||||
],
|
||||
)
|
||||
events: List[str] = []
|
||||
resynced = asyncio.Event()
|
||||
subscriber = _virtual_time_subscriber(pubsub, clock, events, resynced, expected_resyncs=2)
|
||||
|
||||
await _run_until_resynced(subscriber, resynced)
|
||||
|
||||
assert events == ["resync@1000.0", "resync@1010.0"]
|
||||
|
||||
|
||||
async def test_failing_targeted_apply_still_runs_the_full_resync() -> None:
|
||||
clock = _FakeClock()
|
||||
pubsub = _VirtualTimePubSub(clock, deliveries=[(1000.0, _change("m-1"))])
|
||||
events: List[str] = []
|
||||
resynced = asyncio.Event()
|
||||
subscriber = _virtual_time_subscriber(pubsub, clock, events, resynced, apply_error=RuntimeError("db down"))
|
||||
|
||||
await _run_until_resynced(subscriber, resynced)
|
||||
|
||||
assert events == ["apply:m-1@1000.0", "resync@1000.0"]
|
||||
|
||||
|
||||
async def test_subscriber_without_an_apply_callback_only_resyncs() -> None:
|
||||
clock = _FakeClock()
|
||||
pubsub = _VirtualTimePubSub(clock, deliveries=[(1000.0, _change("m-1"))])
|
||||
resyncs: List[str] = []
|
||||
fired = asyncio.Event()
|
||||
subscriber = ConfigSyncSubscriber(
|
||||
redis_cache=_FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])),
|
||||
resync_callbacks=(_recording_callback(resyncs, "resync", fired),),
|
||||
debounce_seconds=0.0,
|
||||
jitter_max_seconds=0.0,
|
||||
monotonic=clock,
|
||||
)
|
||||
|
||||
await _run_until_resynced(subscriber, fired)
|
||||
|
||||
assert resyncs == ["resync"]
|
||||
|
||||
|
||||
async def test_peer_holds_every_published_row_before_its_next_full_resync() -> None:
|
||||
publisher = _RecordingRedisClient()
|
||||
cache = _FakeRedisCache(publisher)
|
||||
await publish_config_change(redis_cache=cache, object_type="litellm_proxymodeltable", model_ids=("m-order-1",))
|
||||
await publish_config_change(redis_cache=cache, object_type="litellm_proxymodeltable", model_ids=("m-order-2",))
|
||||
on_the_wire = [_redis_delivery(message) for _, message in publisher.published]
|
||||
clock = _FakeClock()
|
||||
pubsub = _VirtualTimePubSub(clock, deliveries=[(1000.0, on_the_wire[0]), (1007.0, on_the_wire[1])])
|
||||
db_rows_written_at = {"m-order-1": 1000.0, "m-order-2": 1007.0}
|
||||
registry: set = set()
|
||||
snapshots: List[Tuple[str, float, frozenset]] = []
|
||||
resynced = asyncio.Event()
|
||||
|
||||
async def apply(model_ids: Tuple[str, ...]) -> None:
|
||||
registry.update(model_ids)
|
||||
snapshots.append(("apply", clock.now, frozenset(registry)))
|
||||
|
||||
async def resync() -> None:
|
||||
registry.update(row for row, written_at in db_rows_written_at.items() if written_at <= clock.now)
|
||||
snapshots.append(("resync", clock.now, frozenset(registry)))
|
||||
if sum(kind == "resync" for kind, _, _ in snapshots) == 2:
|
||||
resynced.set()
|
||||
|
||||
subscriber = ConfigSyncSubscriber(
|
||||
redis_cache=_FakeRedisCache(_ScriptedPubSubRedisClient([pubsub])),
|
||||
resync_callbacks=(resync,),
|
||||
apply_model_ids=apply,
|
||||
debounce_seconds=1.0,
|
||||
jitter_max_seconds=5.0,
|
||||
min_resync_interval_seconds=10.0,
|
||||
rng=_MaxJitterRandom(),
|
||||
monotonic=clock,
|
||||
)
|
||||
|
||||
await _run_until_resynced(subscriber, resynced)
|
||||
|
||||
assert snapshots == [
|
||||
("apply", 1000.0, frozenset({"m-order-1"})),
|
||||
("resync", 1006.0, frozenset({"m-order-1"})),
|
||||
("apply", 1007.0, frozenset({"m-order-1", "m-order-2"})),
|
||||
("resync", 1016.0, frozenset({"m-order-1", "m-order-2"})),
|
||||
]
|
||||
|
||||
|
||||
def test_default_jitter_window_is_nonzero() -> None:
|
||||
|
|
@ -345,92 +568,43 @@ def test_default_min_resync_interval_caps_reload_rate() -> None:
|
|||
assert CONFIG_SYNC_MIN_RESYNC_INTERVAL_SECONDS > CONFIG_SYNC_JITTER_MAX_SECONDS
|
||||
|
||||
|
||||
def _throttled_subscriber(
|
||||
cache: object,
|
||||
events: List[str],
|
||||
fired: asyncio.Event,
|
||||
clock: _FakeClock,
|
||||
min_resync_interval_seconds: float = 10.0,
|
||||
) -> ConfigSyncSubscriber:
|
||||
async def recording_sleep(seconds: float) -> None:
|
||||
events.append(f"sleep:{seconds}")
|
||||
await asyncio.sleep(0)
|
||||
|
||||
async def resync() -> None:
|
||||
events.append("resync")
|
||||
fired.set()
|
||||
|
||||
return ConfigSyncSubscriber(
|
||||
redis_cache=cache,
|
||||
resync_callbacks=(resync,),
|
||||
debounce_seconds=0.0,
|
||||
jitter_max_seconds=0.0,
|
||||
min_resync_interval_seconds=min_resync_interval_seconds,
|
||||
sleep=recording_sleep,
|
||||
monotonic=clock,
|
||||
)
|
||||
|
||||
|
||||
async def test_resync_arriving_inside_min_interval_waits_out_the_remainder() -> None:
|
||||
pubsub = _QueuePubSub()
|
||||
cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub]))
|
||||
events: List[str] = []
|
||||
fired = asyncio.Event()
|
||||
clock = _FakeClock()
|
||||
subscriber = _throttled_subscriber(cache=cache, events=events, fired=fired, clock=clock)
|
||||
pubsub = _VirtualTimePubSub(clock, deliveries=[(1000.0, "change"), (1004.0, "change")])
|
||||
events: List[str] = []
|
||||
resynced = asyncio.Event()
|
||||
subscriber = _virtual_time_subscriber(pubsub, clock, events, resynced, expected_resyncs=2)
|
||||
|
||||
subscriber.start()
|
||||
pubsub.queue.put_nowait("change")
|
||||
await asyncio.wait_for(fired.wait(), timeout=5)
|
||||
fired.clear()
|
||||
clock.now += 4.0
|
||||
pubsub.queue.put_nowait("change")
|
||||
await asyncio.wait_for(fired.wait(), timeout=5)
|
||||
await subscriber.stop()
|
||||
await _run_until_resynced(subscriber, resynced)
|
||||
|
||||
assert events == ["sleep:0.0", "resync", "sleep:0.0", "sleep:6.0", "resync"]
|
||||
assert events == ["resync@1000.0", "resync@1010.0"]
|
||||
|
||||
|
||||
async def test_resync_after_min_interval_elapsed_is_not_throttled() -> None:
|
||||
pubsub = _QueuePubSub()
|
||||
cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub]))
|
||||
events: List[str] = []
|
||||
fired = asyncio.Event()
|
||||
clock = _FakeClock()
|
||||
subscriber = _throttled_subscriber(cache=cache, events=events, fired=fired, clock=clock)
|
||||
pubsub = _VirtualTimePubSub(clock, deliveries=[(1000.0, "change"), (1030.0, "change")])
|
||||
events: List[str] = []
|
||||
resynced = asyncio.Event()
|
||||
subscriber = _virtual_time_subscriber(pubsub, clock, events, resynced, expected_resyncs=2)
|
||||
|
||||
subscriber.start()
|
||||
pubsub.queue.put_nowait("change")
|
||||
await asyncio.wait_for(fired.wait(), timeout=5)
|
||||
fired.clear()
|
||||
clock.now += 30.0
|
||||
pubsub.queue.put_nowait("change")
|
||||
await asyncio.wait_for(fired.wait(), timeout=5)
|
||||
await subscriber.stop()
|
||||
await _run_until_resynced(subscriber, resynced)
|
||||
|
||||
assert events == ["sleep:0.0", "resync", "sleep:0.0", "resync"]
|
||||
assert events == ["resync@1000.0", "resync@1030.0"]
|
||||
|
||||
|
||||
async def test_writes_during_the_throttle_wait_collapse_into_the_next_resync() -> None:
|
||||
pubsub = _QueuePubSub()
|
||||
cache = _FakeRedisCache(_ScriptedPubSubRedisClient([pubsub]))
|
||||
events: List[str] = []
|
||||
fired = asyncio.Event()
|
||||
clock = _FakeClock()
|
||||
subscriber = _throttled_subscriber(cache=cache, events=events, fired=fired, clock=clock)
|
||||
pubsub = _VirtualTimePubSub(
|
||||
clock, deliveries=[(1000.0, "change")] + [(1004.0 + offset, "change") for offset in range(5)]
|
||||
)
|
||||
events: List[str] = []
|
||||
resynced = asyncio.Event()
|
||||
subscriber = _virtual_time_subscriber(pubsub, clock, events, resynced, expected_resyncs=2)
|
||||
|
||||
subscriber.start()
|
||||
pubsub.queue.put_nowait("change")
|
||||
await asyncio.wait_for(fired.wait(), timeout=5)
|
||||
fired.clear()
|
||||
for _ in range(5):
|
||||
pubsub.queue.put_nowait("change")
|
||||
await asyncio.wait_for(fired.wait(), timeout=5)
|
||||
await asyncio.sleep(0.1)
|
||||
await subscriber.stop()
|
||||
await _run_until_resynced(subscriber, resynced)
|
||||
|
||||
assert events.count("resync") == 2
|
||||
assert pubsub.queue.empty()
|
||||
assert events == ["resync@1000.0", "resync@1010.0"]
|
||||
assert pubsub.pending == []
|
||||
|
||||
|
||||
async def test_polls_without_messages_do_not_trigger_resyncs() -> None:
|
||||
|
|
@ -646,13 +820,39 @@ class _AllWritesTableActions:
|
|||
return action
|
||||
|
||||
|
||||
def _recording_publish(calls: List[Tuple[str, str]]) -> Callable[[str], Coroutine[None, None, None]]:
|
||||
async def publish(object_type: str) -> None:
|
||||
def _recording_publish(calls: List[Tuple[str, str]]) -> Callable[[str, Tuple[str, ...]], Coroutine[None, None, None]]:
|
||||
async def publish(object_type: str, model_ids: Tuple[str, ...]) -> None:
|
||||
calls.append(("publish", object_type))
|
||||
|
||||
return publish
|
||||
|
||||
|
||||
def _recording_publish_with_ids(
|
||||
calls: List[Tuple[str, Tuple[str, ...]]],
|
||||
) -> Callable[[str, Tuple[str, ...]], Coroutine[None, None, None]]:
|
||||
async def publish(object_type: str, model_ids: Tuple[str, ...]) -> None:
|
||||
calls.append((object_type, model_ids))
|
||||
|
||||
return publish
|
||||
|
||||
|
||||
class _ModelRow(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
model_id: str
|
||||
model_name: str
|
||||
|
||||
|
||||
class _ModelRowTableActions:
|
||||
def __getattr__(self, name: str) -> Callable[..., Coroutine[None, None, object]]:
|
||||
async def action(*args: object, **kwargs: object) -> object:
|
||||
if name.endswith("_many"):
|
||||
return 2
|
||||
return _ModelRow(model_id="m-1", model_name="gpt-5.2")
|
||||
|
||||
return action
|
||||
|
||||
|
||||
def test_wrapper_passes_through_unsynced_tables() -> None:
|
||||
actions = object()
|
||||
|
||||
|
|
@ -722,6 +922,46 @@ async def test_wrapper_publishes_for_every_write_action(action_name: str) -> Non
|
|||
assert publish_calls == [("publish", "litellm_guardrailstable")]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"action_name, expected_ids",
|
||||
[
|
||||
("create", ("m-1",)),
|
||||
("update", ("m-1",)),
|
||||
("upsert", ("m-1",)),
|
||||
("delete", ()),
|
||||
("create_many", ()),
|
||||
("update_many", ()),
|
||||
("delete_many", ()),
|
||||
],
|
||||
)
|
||||
async def test_wrapper_publishes_the_written_model_id_for_row_writes(
|
||||
action_name: str, expected_ids: Tuple[str, ...]
|
||||
) -> None:
|
||||
publish_calls: List[Tuple[str, Tuple[str, ...]]] = []
|
||||
wrapped = wrap_table_actions_for_config_sync(
|
||||
actions=_ModelRowTableActions(),
|
||||
table_name="litellm_proxymodeltable",
|
||||
publish=_recording_publish_with_ids(publish_calls),
|
||||
)
|
||||
|
||||
await getattr(wrapped, action_name)(data={})
|
||||
|
||||
assert publish_calls == [("litellm_proxymodeltable", expected_ids)]
|
||||
|
||||
|
||||
async def test_wrapper_publishes_no_model_ids_for_other_tables() -> None:
|
||||
publish_calls: List[Tuple[str, Tuple[str, ...]]] = []
|
||||
wrapped = wrap_table_actions_for_config_sync(
|
||||
actions=_ModelRowTableActions(),
|
||||
table_name="litellm_guardrailstable",
|
||||
publish=_recording_publish_with_ids(publish_calls),
|
||||
)
|
||||
|
||||
await wrapped.create(data={})
|
||||
|
||||
assert publish_calls == [("litellm_guardrailstable", ())]
|
||||
|
||||
|
||||
async def test_model_repository_write_publishes_via_live_coordination_cache() -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.proxy_server import _set_redis_usage_cache
|
||||
|
|
@ -747,7 +987,7 @@ async def test_model_repository_write_publishes_via_live_coordination_cache() ->
|
|||
assert len(client.published) == 1
|
||||
channel, message = client.published[0]
|
||||
assert channel == CONFIG_SYNC_CHANNEL
|
||||
assert json.loads(message) == {"object_type": "litellm_proxymodeltable"}
|
||||
assert json.loads(message) == {"object_type": "litellm_proxymodeltable", "model_ids": ["m-1"]}
|
||||
|
||||
|
||||
async def test_ui_settings_write_publishes_via_live_coordination_cache() -> None:
|
||||
|
|
@ -774,7 +1014,7 @@ async def test_ui_settings_write_publishes_via_live_coordination_cache() -> None
|
|||
assert len(client.published) == 1
|
||||
channel, message = client.published[0]
|
||||
assert channel == CONFIG_SYNC_CHANNEL
|
||||
assert json.loads(message) == {"object_type": "litellm_uisettings"}
|
||||
assert json.loads(message) == {"object_type": "litellm_uisettings", "model_ids": []}
|
||||
|
||||
|
||||
async def _publish_calls_for_invalidated_param(param_name: str) -> List[Tuple[str, str]]:
|
||||
|
|
@ -799,7 +1039,7 @@ async def test_invalidate_config_param_publishes_params_a_resync_applies(param_n
|
|||
assert len(published) == 1
|
||||
channel, message = published[0]
|
||||
assert channel == CONFIG_SYNC_CHANNEL
|
||||
assert json.loads(message) == {"object_type": param_name}
|
||||
assert json.loads(message) == {"object_type": param_name, "model_ids": []}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("param_name", _STARTUP_ONLY_CONFIG_PARAM_NAMES)
|
||||
|
|
@ -932,6 +1172,46 @@ async def test_proxy_config_subscriber_resyncs_deployments_only() -> None:
|
|||
assert subscriber._task is None
|
||||
|
||||
|
||||
async def test_proxy_config_subscriber_applies_published_model_ids_from_the_db(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from litellm.proxy import proxy_server
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
rows = [MagicMock()]
|
||||
prisma_client = MagicMock()
|
||||
prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=rows)
|
||||
router = MagicMock()
|
||||
router.registry.return_value.deployments.return_value = ["deployment"]
|
||||
installed: List[object] = []
|
||||
|
||||
def install(db_models: object) -> None:
|
||||
installed.append(db_models)
|
||||
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
|
||||
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
|
||||
monkeypatch.setattr(proxy_server, "llm_router", router)
|
||||
monkeypatch.setattr(proxy_server, "llm_model_list", None)
|
||||
monkeypatch.setattr(proxy_server.proxy_config, "get_credentials", AsyncMock())
|
||||
monkeypatch.setattr(proxy_server.proxy_config, "_add_deployment", install)
|
||||
cache = _FakeRedisCache(_ScriptedPubSubRedisClient([_QueuePubSub()]))
|
||||
config = ProxyConfig()
|
||||
|
||||
config.start_config_sync_subscriber(prisma_client=prisma_client, proxy_logging_obj=MagicMock(), redis_cache=cache)
|
||||
subscriber = config.config_sync_subscriber
|
||||
assert subscriber is not None
|
||||
apply = subscriber._apply_model_ids
|
||||
assert apply is not None
|
||||
await apply(("m-1", "m-2"))
|
||||
await config.stop_config_sync_subscriber()
|
||||
|
||||
prisma_client.db.litellm_proxymodeltable.find_many.assert_awaited_once_with(
|
||||
where={"model_id": {"in": ["m-1", "m-2"]}}
|
||||
)
|
||||
assert installed == [rows]
|
||||
assert proxy_server.llm_model_list == ["deployment"]
|
||||
|
||||
|
||||
async def test_proxy_config_does_not_start_subscriber_without_coordination_redis() -> None:
|
||||
from litellm.proxy.proxy_server import ProxyConfig
|
||||
|
||||
|
|
|
|||
|
|
@ -1,4 +1,5 @@
|
|||
import asyncio
|
||||
from collections.abc import Mapping, Sequence
|
||||
from typing import TYPE_CHECKING, Final
|
||||
|
||||
import pytest
|
||||
|
|
@ -6,7 +7,10 @@ import pytest
|
|||
from litellm.proxy.common_utils.registry_read_through import RegistryReadThrough
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from litellm import Router
|
||||
from litellm.proxy._types import LiteLLM_ProxyModelTable
|
||||
from litellm.proxy.agent_endpoints.agent_registry import AgentRegistry
|
||||
from litellm.types.utils import CredentialItem
|
||||
|
||||
|
||||
def nothing_loaded(_key: str) -> bool:
|
||||
|
|
@ -487,6 +491,178 @@ async def test_resync_model_deployments_loads_db_credentials_before_reconciling_
|
|||
installed.assert_called_once_with(db_models=rows, credential={"api_key": "sk-from-db"})
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def stored_params_salt(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
monkeypatch.setenv("LITELLM_SALT_KEY", "sk-salt-stored-model-rows")
|
||||
|
||||
|
||||
def _stored_model_row(model_id: str, stored_params: Mapping[str, str]) -> "LiteLLM_ProxyModelTable":
|
||||
from litellm.proxy._types import LiteLLM_ProxyModelTable
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
|
||||
|
||||
return LiteLLM_ProxyModelTable(
|
||||
model_id=model_id,
|
||||
model_name="gpt-failover-group",
|
||||
litellm_params={key: encrypt_value_helper(value) for key, value in stored_params.items()},
|
||||
model_info={"id": model_id},
|
||||
)
|
||||
|
||||
|
||||
def _backup_row(model_id: str, model: str = "openai/gpt-backup") -> "LiteLLM_ProxyModelTable":
|
||||
return _stored_model_row(model_id, {"model": model, "api_key": "sk-backup"})
|
||||
|
||||
|
||||
def _router_serving(*model_ids: str) -> "Router":
|
||||
from litellm import Router
|
||||
|
||||
return Router(
|
||||
model_list=[
|
||||
{
|
||||
"model_name": "gpt-failover-group",
|
||||
"litellm_params": {"model": "openai/gpt-primary", "api_key": "sk-primary"},
|
||||
"model_info": {"id": model_id},
|
||||
}
|
||||
for model_id in model_ids
|
||||
]
|
||||
)
|
||||
|
||||
|
||||
def _served_ids(deployments: "Sequence[Mapping[str, Mapping[str, object]]] | None") -> frozenset[object]:
|
||||
return frozenset(deployment["model_info"]["id"] for deployment in deployments or ())
|
||||
|
||||
|
||||
class _ProxyModelTable:
|
||||
def __init__(self, rows: "tuple[LiteLLM_ProxyModelTable, ...]") -> None:
|
||||
self.rows = rows
|
||||
self.reads = 0
|
||||
|
||||
async def find_many(
|
||||
self, where: Mapping[str, Mapping[str, Sequence[str]]]
|
||||
) -> "tuple[LiteLLM_ProxyModelTable, ...]":
|
||||
self.reads += 1
|
||||
return tuple(row for row in self.rows if row.model_id in where["model_id"]["in"])
|
||||
|
||||
async def create(self, data: Mapping[str, object]) -> None:
|
||||
raise AssertionError(f"a model load wrote a row: {data}")
|
||||
|
||||
|
||||
class _CredentialsTable:
|
||||
def __init__(self, rows: "tuple[CredentialItem, ...]") -> None:
|
||||
self.rows = rows
|
||||
|
||||
async def find_many(self) -> "tuple[CredentialItem, ...]":
|
||||
return self.rows
|
||||
|
||||
|
||||
class _Db:
|
||||
def __init__(
|
||||
self, litellm_proxymodeltable: _ProxyModelTable, litellm_credentialstable: _CredentialsTable | None = None
|
||||
) -> None:
|
||||
self.litellm_proxymodeltable = litellm_proxymodeltable
|
||||
self.litellm_credentialstable = litellm_credentialstable or _CredentialsTable(())
|
||||
self.other_tables_read: list[str] = []
|
||||
|
||||
def __getattr__(self, table_name: str) -> object:
|
||||
self.other_tables_read.append(table_name)
|
||||
raise AttributeError(table_name)
|
||||
|
||||
|
||||
class _PrismaClientWithoutConfigRows:
|
||||
def __init__(self, db: object) -> None:
|
||||
self.db = db
|
||||
|
||||
async def get_generic_data(self, key: str, value: str, table_name: str) -> None:
|
||||
return None
|
||||
|
||||
|
||||
def _serve_db_models(monkeypatch: pytest.MonkeyPatch, db: object, router: "Router | None") -> None:
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", _PrismaClientWithoutConfigRows(db))
|
||||
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
|
||||
monkeypatch.setattr(proxy_server, "llm_router", router)
|
||||
monkeypatch.setattr(proxy_server, "llm_model_list", None)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("stored_params_salt")
|
||||
async def test_resync_model_deployments_by_ids_installs_only_the_published_rows(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy.common_utils.registry_read_through import resync_model_deployments_by_ids
|
||||
|
||||
router: Final = _router_serving("primary-id")
|
||||
_serve_db_models(
|
||||
monkeypatch, _Db(_ProxyModelTable((_backup_row("m-1"), _backup_row("m-2"), _backup_row("m-3")))), router
|
||||
)
|
||||
|
||||
await resync_model_deployments_by_ids(("m-1", "m-2"))
|
||||
|
||||
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()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("stored_params_salt")
|
||||
async def test_resync_model_deployments_by_ids_leaves_the_router_alone_for_deleted_rows(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy.common_utils.registry_read_through import resync_model_deployments_by_ids
|
||||
|
||||
router: Final = _router_serving("primary-id")
|
||||
_serve_db_models(monkeypatch, _Db(_ProxyModelTable(())), router)
|
||||
|
||||
await resync_model_deployments_by_ids(("deleted-id",))
|
||||
|
||||
assert _served_ids(router.registry().deployments()) == {"primary-id"}
|
||||
assert proxy_server.llm_model_list is None
|
||||
assert not proxy_server.MODEL_RECONCILE_LOCK.locked()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("stored_params_salt")
|
||||
async def test_resync_model_deployments_by_ids_without_a_router_leaves_the_reload_to_the_debounced_resync(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy.common_utils.registry_read_through import resync_model_deployments_by_ids
|
||||
|
||||
table: Final = _ProxyModelTable((_backup_row("m-1"),))
|
||||
db: Final = _Db(table)
|
||||
_serve_db_models(monkeypatch, db, None)
|
||||
|
||||
await resync_model_deployments_by_ids(("m-1",))
|
||||
|
||||
assert proxy_server.llm_router is None
|
||||
assert table.reads == 1
|
||||
assert db.other_tables_read == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resync_model_deployments_by_ids_respects_supported_db_objects(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy.common_utils.registry_read_through import resync_model_deployments_by_ids
|
||||
|
||||
prisma_client: Final = MagicMock()
|
||||
prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(
|
||||
side_effect=AssertionError("db hit for an object type this replica does not load")
|
||||
)
|
||||
monkeypatch.setattr(proxy_server, "prisma_client", prisma_client)
|
||||
monkeypatch.setattr(proxy_server, "store_model_in_db", True)
|
||||
monkeypatch.setattr(proxy_server, "general_settings", {"supported_db_objects": ["guardrails"]})
|
||||
|
||||
await resync_model_deployments_by_ids(("gated-out-id",))
|
||||
|
||||
prisma_client.db.litellm_proxymodeltable.find_many.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_resync_model_deployments_respects_supported_db_objects(monkeypatch):
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
|
@ -786,3 +962,174 @@ async def test_resync_model_deployments_reconciles_every_db_row_when_the_missed_
|
|||
|
||||
assert await _resync_model_deployments(router_name) is router_holds_the_name_after_reconcile
|
||||
reconciled.assert_called_once_with(prisma_client=prisma_client, proxy_logging_obj=proxy_server.proxy_logging_obj)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("stored_params_salt")
|
||||
async def test_resync_model_deployments_by_ids_reads_the_writer_not_a_lagging_replica(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from litellm.proxy.common_utils.registry_read_through import resync_model_deployments_by_ids
|
||||
from litellm.proxy.db.prisma_client import PrismaWrapper
|
||||
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
|
||||
|
||||
writer: Final = MagicMock(name="writer_prisma")
|
||||
writer.litellm_proxymodeltable = _ProxyModelTable((_backup_row("backup-id"),))
|
||||
replica_table: Final = _ProxyModelTable(())
|
||||
replica: Final = MagicMock(name="lagging_replica_prisma")
|
||||
replica.litellm_proxymodeltable = replica_table
|
||||
router: Final = _router_serving("primary-id")
|
||||
_serve_db_models(
|
||||
monkeypatch,
|
||||
RoutingPrismaWrapper(
|
||||
writer=PrismaWrapper(original_prisma=writer, iam_token_db_auth=False),
|
||||
reader=PrismaWrapper(original_prisma=replica, iam_token_db_auth=False),
|
||||
),
|
||||
router,
|
||||
)
|
||||
|
||||
await resync_model_deployments_by_ids(("backup-id",))
|
||||
|
||||
assert router.registry().has_deployment("backup-id")
|
||||
assert replica_table.reads == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unknown_model_name_read_through_reads_the_replica_without_waiting_for_a_reload(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import AsyncMock, MagicMock
|
||||
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy.common_utils.registry_read_through import _resync_model_deployments
|
||||
from litellm.proxy.db.prisma_client import PrismaWrapper
|
||||
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
|
||||
|
||||
writer: Final = MagicMock(name="writer_prisma")
|
||||
writer.litellm_proxymodeltable = SimpleNamespace(find_many=AsyncMock(return_value=[]), create=AsyncMock())
|
||||
replica: Final = MagicMock(name="replica_prisma")
|
||||
replica.litellm_proxymodeltable = SimpleNamespace(find_many=AsyncMock(return_value=[]), create=AsyncMock())
|
||||
prisma_client: Final = MagicMock()
|
||||
prisma_client.db = RoutingPrismaWrapper(
|
||||
writer=PrismaWrapper(original_prisma=writer, iam_token_db_auth=False),
|
||||
reader=PrismaWrapper(original_prisma=replica, iam_token_db_auth=False),
|
||||
)
|
||||
reload_in_progress: Final = asyncio.Lock()
|
||||
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)
|
||||
|
||||
async with reload_in_progress:
|
||||
found: Final = await asyncio.wait_for(_resync_model_deployments("no-such-model"), timeout=2)
|
||||
|
||||
assert found is False
|
||||
writer.litellm_proxymodeltable.find_many.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("stored_params_salt")
|
||||
async def test_resync_model_deployments_by_ids_installs_the_row_a_concurrent_reload_left_current(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy.common_utils.registry_read_through import resync_model_deployments_by_ids
|
||||
|
||||
table: Final = _ProxyModelTable((_backup_row("m-1", model="openai/gpt-before-update"),))
|
||||
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)
|
||||
|
||||
async with reload_in_progress:
|
||||
targeted_load: Final = asyncio.create_task(resync_model_deployments_by_ids(("m-1",)))
|
||||
for _ in range(5):
|
||||
await asyncio.sleep(0)
|
||||
assert not router.registry().has_deployment("m-1")
|
||||
table.rows = (_backup_row("m-1", model="openai/gpt-after-update"),)
|
||||
await targeted_load
|
||||
|
||||
installed: Final = router.registry().deployment("m-1")
|
||||
assert installed is not None
|
||||
assert installed.litellm_params.model == "openai/gpt-after-update"
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("stored_params_salt")
|
||||
@pytest.mark.parametrize(
|
||||
("stored_credential_name", "credentials_on_this_worker", "api_key_this_worker_uses"),
|
||||
[
|
||||
("shared-openai", (), "sk-from-db"),
|
||||
("shared-openai", ("another-openai",), "sk-from-db"),
|
||||
("shared-openai", ("shared-openai",), "sk-on-this-worker"),
|
||||
(None, (), None),
|
||||
],
|
||||
)
|
||||
async def test_resync_model_deployments_by_ids_reloads_credentials_only_for_one_this_worker_lacks(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
stored_credential_name: str | None,
|
||||
credentials_on_this_worker: tuple[str, ...],
|
||||
api_key_this_worker_uses: str | None,
|
||||
) -> None:
|
||||
import litellm
|
||||
from litellm.litellm_core_utils.credential_accessor import CredentialAccessor
|
||||
from litellm.proxy.common_utils.encrypt_decrypt_utils import encrypt_value_helper
|
||||
from litellm.proxy.common_utils.registry_read_through import resync_model_deployments_by_ids
|
||||
from litellm.types.utils import CredentialItem
|
||||
|
||||
model_params: Final = {"model": "openai/gpt-backup"}
|
||||
row: Final = _stored_model_row(
|
||||
"backup-id",
|
||||
model_params if stored_credential_name is None else model_params | {"litellm_credential_name": stored_credential_name},
|
||||
)
|
||||
stored_credential: Final = CredentialItem(
|
||||
credential_name="shared-openai",
|
||||
credential_values={"api_key": encrypt_value_helper("sk-from-db")},
|
||||
credential_info={},
|
||||
)
|
||||
router: Final = _router_serving("primary-id")
|
||||
monkeypatch.setattr(
|
||||
litellm,
|
||||
"credential_list",
|
||||
[
|
||||
CredentialItem(credential_name=name, credential_values={"api_key": "sk-on-this-worker"}, credential_info={})
|
||||
for name in credentials_on_this_worker
|
||||
],
|
||||
)
|
||||
_serve_db_models(monkeypatch, _Db(_ProxyModelTable((row,)), _CredentialsTable((stored_credential,))), router)
|
||||
|
||||
await resync_model_deployments_by_ids(("backup-id",))
|
||||
|
||||
credential: Final = CredentialAccessor.find_credential("shared-openai")
|
||||
assert (credential.credential_values.get("api_key") if credential else None) == api_key_this_worker_uses
|
||||
assert router.registry().has_deployment("backup-id")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.usefixtures("stored_params_salt")
|
||||
async def test_resync_model_deployments_by_ids_installs_the_ordinary_rows_queued_with_an_auto_router(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
) -> None:
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy.common_utils.registry_read_through import resync_model_deployments_by_ids
|
||||
|
||||
table: Final = _ProxyModelTable(
|
||||
(
|
||||
_stored_model_row("auto-router-id", {"model": "auto_router/complexity-router"}),
|
||||
_backup_row("backup-id"),
|
||||
)
|
||||
)
|
||||
db: Final = _Db(table)
|
||||
router: Final = _router_serving("primary-id")
|
||||
_serve_db_models(monkeypatch, db, router)
|
||||
|
||||
await resync_model_deployments_by_ids(("auto-router-id", "backup-id"))
|
||||
|
||||
assert _served_ids(router.registry().deployments()) == {"primary-id", "backup-id"}
|
||||
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()
|
||||
|
|
|
|||
|
|
@ -687,33 +687,152 @@ class _FakeAccessGroupBudgetTable:
|
|||
return self.rows.pop(access_group_name, None)
|
||||
|
||||
|
||||
class _FakeClock:
|
||||
def __init__(self) -> None:
|
||||
self.now = 0.0
|
||||
self.slept: list[float] = []
|
||||
|
||||
def __call__(self) -> float:
|
||||
return self.now
|
||||
|
||||
async def sleep(self, seconds: float) -> None:
|
||||
self.slept.append(seconds)
|
||||
self.now += seconds
|
||||
|
||||
|
||||
class _FakeModelTable:
|
||||
def __init__(self, journal: list[str], deployments) -> None:
|
||||
def __init__(self, journal: list[str], deployments, failing_model_ids=frozenset()) -> None:
|
||||
self.journal = journal
|
||||
self.deployments = list(deployments)
|
||||
self.failing_model_ids = failing_model_ids
|
||||
self.updates: list[tuple[dict, dict]] = []
|
||||
self.rows_locked_by_another_session: frozenset[str] = frozenset()
|
||||
self.group_locked_by_another_session = False
|
||||
self.edits_the_other_session_commits: dict[str, dict] = {}
|
||||
self.other_session_commits_after_refusals: int | None = None
|
||||
self.lock_refusals = 0
|
||||
|
||||
async def find_many(self, where=None, **kwargs):
|
||||
return list(self.deployments)
|
||||
model_names = (where or {}).get("model_name", {}).get("in")
|
||||
return [d for d in self.deployments if model_names is None or d.model_name in model_names]
|
||||
|
||||
def lock_group(self, wait):
|
||||
if not self.group_locked_by_another_session:
|
||||
return True
|
||||
if wait:
|
||||
raise TimeoutError("lock_timeout expired waiting for the access group lock")
|
||||
self._refuse_lock()
|
||||
return False
|
||||
|
||||
def lock_rows(self, model_ids, skip_locked):
|
||||
held = self.rows_locked_by_another_session.intersection(model_ids)
|
||||
if not held:
|
||||
return list(model_ids)
|
||||
if not skip_locked:
|
||||
raise TimeoutError(f"timed out waiting for the row lock on {sorted(held)}")
|
||||
self._refuse_lock()
|
||||
return [model_id for model_id in model_ids if model_id not in held]
|
||||
|
||||
def _refuse_lock(self):
|
||||
self.lock_refusals += 1
|
||||
if self.lock_refusals == self.other_session_commits_after_refusals:
|
||||
self._other_session_commits()
|
||||
|
||||
def _other_session_commits(self):
|
||||
self.deployments = [
|
||||
_FakeDeployment(
|
||||
model_id=d.model_id,
|
||||
model_name=d.model_name,
|
||||
model_info=self.edits_the_other_session_commits.get(d.model_id, d.model_info),
|
||||
)
|
||||
for d in self.deployments
|
||||
]
|
||||
self.rows_locked_by_another_session = frozenset()
|
||||
self.group_locked_by_another_session = False
|
||||
|
||||
async def find_unique(self, where, include=None):
|
||||
return next((d for d in self.deployments if d.model_id == where["model_id"]), None)
|
||||
|
||||
async def update(self, where, data, include=None):
|
||||
if where["model_id"] in self.failing_model_ids:
|
||||
raise RuntimeError(f"database write to {where['model_id']} failed")
|
||||
self.journal.append(f"model_table.update:{where['model_id']}")
|
||||
self.updates.append((dict(where), dict(data)))
|
||||
return None
|
||||
|
||||
async def create(self, data, **kwargs):
|
||||
raise AssertionError("access group endpoints never create deployments")
|
||||
|
||||
|
||||
class _FakeModelTableBatch:
|
||||
"""Queues model-table updates and applies them all or none when the batch closes, like a Prisma batch."""
|
||||
|
||||
def __init__(self, model_table: _FakeModelTable) -> None:
|
||||
self.model_table = model_table
|
||||
self.queued: list[tuple[dict, dict]] = []
|
||||
self.litellm_proxymodeltable = SimpleNamespace(update=self._queue_update)
|
||||
|
||||
def _queue_update(self, where, data):
|
||||
self.queued.append((dict(where), dict(data)))
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
if exc_type is not None:
|
||||
return False
|
||||
failing = [where["model_id"] for where, _ in self.queued if where["model_id"] in self.model_table.failing_model_ids]
|
||||
if failing:
|
||||
raise RuntimeError(f"batch write to {failing} failed, nothing applied")
|
||||
for where, data in self.queued:
|
||||
await self.model_table.update(where=where, data=data)
|
||||
return False
|
||||
|
||||
|
||||
class _FakeModelTableTx:
|
||||
"""An interactive transaction on the primary that journals the row lock and the read it guards."""
|
||||
|
||||
def __init__(self, journal: list[str], model_table: _FakeModelTable) -> None:
|
||||
self.journal = journal
|
||||
self.model_table = model_table
|
||||
self.litellm_proxymodeltable = SimpleNamespace(find_many=self._find_many)
|
||||
|
||||
async def query_raw(self, query, *args):
|
||||
if "advisory_xact_lock" in query:
|
||||
self.journal.append(f"access_group.lock:{args[0]}")
|
||||
return [{"locked": self.model_table.lock_group(wait="pg_try_" not in query)}]
|
||||
if "FOR UPDATE" in query:
|
||||
model_ids = sorted(args[0] if args else (d.model_id for d in self.model_table.deployments))
|
||||
self.journal.append(f"model_table.lock:{','.join(model_ids)}")
|
||||
locked = self.model_table.lock_rows(model_ids, skip_locked="SKIP LOCKED" in query)
|
||||
return [{"model_id": model_id} for model_id in locked]
|
||||
self.journal.append(f"model_table.query:{query}")
|
||||
return []
|
||||
|
||||
async def _find_many(self):
|
||||
self.journal.append("model_table.read")
|
||||
return await self.model_table.find_many()
|
||||
|
||||
def batch_(self):
|
||||
return _FakeModelTableBatch(self.model_table)
|
||||
|
||||
async def __aenter__(self):
|
||||
return self
|
||||
|
||||
async def __aexit__(self, exc_type, exc, tb):
|
||||
return False
|
||||
|
||||
|
||||
class _FakePrismaClient:
|
||||
def __init__(self, journal: list[str], deployments=()) -> None:
|
||||
def __init__(self, journal: list[str], deployments=(), failing_model_ids=frozenset()) -> None:
|
||||
self.budget_table = _FakeBudgetTable(journal)
|
||||
self.access_group_budget_table = _FakeAccessGroupBudgetTable(journal, self.budget_table)
|
||||
self.model_table = _FakeModelTable(journal, deployments)
|
||||
self.model_table = _FakeModelTable(journal, deployments, failing_model_ids)
|
||||
self.db = SimpleNamespace(
|
||||
litellm_budgettable=self.budget_table,
|
||||
litellm_modelaccessgroupbudgettable=self.access_group_budget_table,
|
||||
litellm_proxymodeltable=self.model_table,
|
||||
tx=lambda timeout=None: _FakeModelTableTx(journal, self.model_table),
|
||||
)
|
||||
|
||||
def jsonify_object(self, data):
|
||||
|
|
@ -784,6 +903,12 @@ def _proxy_with_stubbed_reload(prisma):
|
|||
yield
|
||||
|
||||
|
||||
def _serve_model_groups(router, model_groups):
|
||||
router.registry.return_value.names.side_effect = lambda model_name=None, *, kind="groups", **_: (
|
||||
list(model_groups) if kind == "groups" else []
|
||||
)
|
||||
|
||||
|
||||
def _eviction_journal(access_group):
|
||||
"""Both auth cache keys, in the order a write path has to evict them."""
|
||||
from litellm.proxy.common_utils.user_api_key_cache import (
|
||||
|
|
@ -1290,3 +1415,605 @@ async def test_deleting_an_access_group_that_never_had_a_budget_still_evicts():
|
|||
assert response.models_updated == 1
|
||||
assert prisma.access_group_budget_table.rows == {}
|
||||
_assert_evicted_after_write(journal, "prod-models", "access_group_budget.delete:prod-models")
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"selection",
|
||||
[{"model_ids": ["stays", "joins"]}, {"model_names": ["gpt-stays", "gpt-joins"]}],
|
||||
)
|
||||
async def test_updating_an_access_group_writes_only_the_deployments_that_join_or_leave_it(selection):
|
||||
import json
|
||||
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
|
||||
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
|
||||
update_access_group,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
UpdateModelGroupRequest,
|
||||
)
|
||||
|
||||
prisma = _FakePrismaClient(
|
||||
[],
|
||||
deployments=[
|
||||
_deployment(model_id="stays", model_name="gpt-stays"),
|
||||
_deployment(model_id="leaves", model_name="gpt-leaves"),
|
||||
_deployment(model_id="joins", model_name="gpt-joins", access_groups=()),
|
||||
],
|
||||
)
|
||||
|
||||
with _proxy_with_stubbed_reload(prisma):
|
||||
_serve_model_groups(proxy_server.llm_router, ["gpt-stays", "gpt-leaves", "gpt-joins"])
|
||||
response = await update_access_group(
|
||||
access_group="prod-models",
|
||||
data=UpdateModelGroupRequest(**selection),
|
||||
user_api_key_dict=_admin(),
|
||||
)
|
||||
|
||||
written_access_groups = {
|
||||
where["model_id"]: json.loads(data["model_info"])["access_groups"]
|
||||
for where, data in prisma.model_table.updates
|
||||
}
|
||||
assert written_access_groups == {"joins": ["prod-models"], "leaves": []}
|
||||
assert response.models_updated == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"selection",
|
||||
[{"model_ids": ["stays", "joins"]}, {"model_names": ["gpt-stays", "gpt-joins"]}],
|
||||
)
|
||||
async def test_updating_an_access_group_reads_the_primary_not_a_lagging_replica(selection):
|
||||
import json
|
||||
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy.db.prisma_client import PrismaWrapper
|
||||
from litellm.proxy.db.routing_prisma_wrapper import RoutingPrismaWrapper
|
||||
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
|
||||
update_access_group,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
UpdateModelGroupRequest,
|
||||
)
|
||||
|
||||
primary = _FakePrismaClient(
|
||||
[],
|
||||
deployments=[
|
||||
_deployment(model_id="stays", model_name="gpt-stays"),
|
||||
_deployment(model_id="leaves", model_name="gpt-leaves"),
|
||||
_deployment(model_id="joins", model_name="gpt-joins", access_groups=()),
|
||||
],
|
||||
)
|
||||
replica_before_the_last_edit = _FakePrismaClient(
|
||||
[],
|
||||
deployments=[
|
||||
_deployment(model_id="stays", model_name="gpt-stays"),
|
||||
_deployment(model_id="leaves", model_name="gpt-leaves", access_groups=()),
|
||||
_deployment(model_id="joins", model_name="gpt-joins"),
|
||||
],
|
||||
)
|
||||
primary.db = RoutingPrismaWrapper(
|
||||
writer=PrismaWrapper(original_prisma=primary.db, iam_token_db_auth=False),
|
||||
reader=PrismaWrapper(original_prisma=replica_before_the_last_edit.db, iam_token_db_auth=False),
|
||||
)
|
||||
|
||||
with _proxy_with_stubbed_reload(primary):
|
||||
_serve_model_groups(proxy_server.llm_router, ["gpt-stays", "gpt-leaves", "gpt-joins"])
|
||||
await update_access_group(
|
||||
access_group="prod-models",
|
||||
data=UpdateModelGroupRequest(**selection),
|
||||
user_api_key_dict=_admin(),
|
||||
)
|
||||
|
||||
written_access_groups = {
|
||||
where["model_id"]: json.loads(data["model_info"])["access_groups"]
|
||||
for where, data in primary.model_table.updates
|
||||
}
|
||||
assert written_access_groups == {"joins": ["prod-models"], "leaves": []}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("failing_model_id", ["leaves", "joins"])
|
||||
async def test_a_failed_access_group_update_leaves_every_deployment_unchanged(failing_model_id):
|
||||
from fastapi import HTTPException
|
||||
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
|
||||
update_access_group,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
UpdateModelGroupRequest,
|
||||
)
|
||||
|
||||
deployments = [
|
||||
_deployment(model_id="stays", model_name="gpt-stays"),
|
||||
_deployment(model_id="leaves", model_name="gpt-leaves"),
|
||||
_deployment(model_id="joins", model_name="gpt-joins", access_groups=()),
|
||||
]
|
||||
prisma = _FakePrismaClient([], deployments=deployments, failing_model_ids=frozenset({failing_model_id}))
|
||||
|
||||
with _proxy_with_stubbed_reload(prisma):
|
||||
_serve_model_groups(proxy_server.llm_router, ["gpt-stays", "gpt-leaves", "gpt-joins"])
|
||||
with pytest.raises(HTTPException) as raised:
|
||||
await update_access_group(
|
||||
access_group="prod-models",
|
||||
data=UpdateModelGroupRequest(model_ids=["stays", "joins"]),
|
||||
user_api_key_dict=_admin(),
|
||||
)
|
||||
|
||||
assert raised.value.status_code == 500
|
||||
assert prisma.model_table.updates == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_access_group_update_keeps_a_model_info_field_another_admin_set_while_names_were_validated():
|
||||
import json
|
||||
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
|
||||
update_access_group,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
UpdateModelGroupRequest,
|
||||
)
|
||||
|
||||
prisma = _FakePrismaClient(
|
||||
[],
|
||||
deployments=[
|
||||
_deployment(model_id="stays", model_name="gpt-stays"),
|
||||
_deployment(model_id="joins", model_name="gpt-joins", access_groups=()),
|
||||
],
|
||||
)
|
||||
|
||||
def another_admin_edits_joins_while_names_are_checked(model_name=None, *, kind="groups", **_):
|
||||
if kind != "groups":
|
||||
return []
|
||||
prisma.model_table.deployments = [
|
||||
deployment
|
||||
if deployment.model_id != "joins"
|
||||
else _FakeDeployment(
|
||||
model_id="joins",
|
||||
model_name="gpt-joins",
|
||||
model_info={"access_groups": [], "description": "set by another admin"},
|
||||
)
|
||||
for deployment in prisma.model_table.deployments
|
||||
]
|
||||
return ["gpt-stays", "gpt-joins"]
|
||||
|
||||
with _proxy_with_stubbed_reload(prisma):
|
||||
proxy_server.llm_router.registry.return_value.names.side_effect = (
|
||||
another_admin_edits_joins_while_names_are_checked
|
||||
)
|
||||
await update_access_group(
|
||||
access_group="prod-models",
|
||||
data=UpdateModelGroupRequest(model_names=["gpt-stays", "gpt-joins"]),
|
||||
user_api_key_dict=_admin(),
|
||||
)
|
||||
|
||||
written = {where["model_id"]: json.loads(data["model_info"]) for where, data in prisma.model_table.updates}
|
||||
assert written["joins"]["description"] == "set by another admin"
|
||||
assert written["joins"]["access_groups"] == ["prod-models"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_access_group_update_locks_the_rows_it_reads_and_tells_other_workers_after_the_commit():
|
||||
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
|
||||
update_access_group,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
UpdateModelGroupRequest,
|
||||
)
|
||||
|
||||
journal: list[str] = []
|
||||
prisma = _FakePrismaClient(
|
||||
journal,
|
||||
deployments=[
|
||||
_deployment(model_id="stays", model_name="gpt-stays"),
|
||||
_deployment(model_id="leaves", model_name="gpt-leaves"),
|
||||
_deployment(model_id="joins", model_name="gpt-joins", access_groups=()),
|
||||
],
|
||||
)
|
||||
|
||||
async def record_publish(model_ids):
|
||||
journal.append(f"publish:{','.join(sorted(model_ids))}")
|
||||
|
||||
with (
|
||||
_proxy_with_stubbed_reload(prisma),
|
||||
patch( # test-quality-ok: the endpoint publishes through its module-level import; Redis is the real seam
|
||||
"litellm.proxy.management_endpoints.model_access_group_management_endpoints.publish_model_rows_written",
|
||||
new=record_publish,
|
||||
),
|
||||
):
|
||||
await update_access_group(
|
||||
access_group="prod-models",
|
||||
data=UpdateModelGroupRequest(model_ids=["stays", "joins"]),
|
||||
user_api_key_dict=_admin(),
|
||||
)
|
||||
|
||||
assert journal == [
|
||||
"access_group.lock:prod-models",
|
||||
"model_table.read",
|
||||
"model_table.lock:joins,leaves,stays",
|
||||
"model_table.read",
|
||||
"model_table.update:leaves",
|
||||
"model_table.update:joins",
|
||||
"publish:joins,leaves",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_access_group_update_that_changes_no_membership_writes_and_announces_nothing():
|
||||
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
|
||||
update_access_group,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
UpdateModelGroupRequest,
|
||||
)
|
||||
|
||||
journal: list[str] = []
|
||||
prisma = _FakePrismaClient(
|
||||
journal,
|
||||
deployments=[
|
||||
_deployment(model_id="stays", model_name="gpt-stays"),
|
||||
_deployment(model_id="also-stays", model_name="gpt-also-stays"),
|
||||
_deployment(model_id="outside", model_name="gpt-outside", access_groups=()),
|
||||
],
|
||||
)
|
||||
|
||||
async def record_publish(model_ids):
|
||||
journal.append(f"publish:{','.join(sorted(model_ids))}")
|
||||
|
||||
with (
|
||||
_proxy_with_stubbed_reload(prisma),
|
||||
patch( # test-quality-ok: the endpoint publishes through its module-level import; Redis is the real seam
|
||||
"litellm.proxy.management_endpoints.model_access_group_management_endpoints.publish_model_rows_written",
|
||||
new=record_publish,
|
||||
),
|
||||
):
|
||||
response = await update_access_group(
|
||||
access_group="prod-models",
|
||||
data=UpdateModelGroupRequest(model_ids=["stays", "also-stays"]),
|
||||
user_api_key_dict=_admin(),
|
||||
)
|
||||
|
||||
assert response.models_updated == 2
|
||||
assert journal == [
|
||||
"access_group.lock:prod-models",
|
||||
"model_table.read",
|
||||
"model_table.lock:also-stays,stays",
|
||||
"model_table.read",
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("ignored_model_names", [None, ["gpt-unrelated"]])
|
||||
async def test_an_access_group_update_goes_through_while_another_session_holds_a_model_row_outside_the_group(
|
||||
ignored_model_names,
|
||||
):
|
||||
import json
|
||||
|
||||
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
|
||||
update_access_group,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
UpdateModelGroupRequest,
|
||||
)
|
||||
|
||||
prisma = _FakePrismaClient(
|
||||
[],
|
||||
deployments=[
|
||||
_deployment(model_id="stays", model_name="gpt-stays"),
|
||||
_deployment(model_id="leaves", model_name="gpt-leaves"),
|
||||
_deployment(model_id="joins", model_name="gpt-joins", access_groups=()),
|
||||
_deployment(model_id="unrelated", model_name="gpt-unrelated", access_groups=("other-group",)),
|
||||
],
|
||||
)
|
||||
prisma.model_table.rows_locked_by_another_session = frozenset({"unrelated"})
|
||||
|
||||
with _proxy_with_stubbed_reload(prisma):
|
||||
response = await update_access_group(
|
||||
access_group="prod-models",
|
||||
data=UpdateModelGroupRequest(model_ids=["stays", "joins"], model_names=ignored_model_names),
|
||||
user_api_key_dict=_admin(),
|
||||
)
|
||||
|
||||
written_access_groups = {
|
||||
where["model_id"]: json.loads(data["model_info"])["access_groups"]
|
||||
for where, data in prisma.model_table.updates
|
||||
}
|
||||
assert written_access_groups == {"joins": ["prod-models"], "leaves": []}
|
||||
assert response.models_updated == 2
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_access_group_update_retries_a_locked_row_and_keeps_the_edit_committed_there():
|
||||
import json
|
||||
|
||||
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
|
||||
update_access_group,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
UpdateModelGroupRequest,
|
||||
)
|
||||
|
||||
prisma = _FakePrismaClient(
|
||||
[],
|
||||
deployments=[
|
||||
_deployment(model_id="stays", model_name="gpt-stays"),
|
||||
_deployment(model_id="joins", model_name="gpt-joins", access_groups=()),
|
||||
],
|
||||
)
|
||||
prisma.model_table.rows_locked_by_another_session = frozenset({"joins"})
|
||||
prisma.model_table.edits_the_other_session_commits = {
|
||||
"joins": {"access_groups": [], "description": "saved while the update waited"}
|
||||
}
|
||||
prisma.model_table.other_session_commits_after_refusals = 1
|
||||
|
||||
with _proxy_with_stubbed_reload(prisma):
|
||||
await update_access_group(
|
||||
access_group="prod-models",
|
||||
data=UpdateModelGroupRequest(model_ids=["stays", "joins"]),
|
||||
user_api_key_dict=_admin(),
|
||||
)
|
||||
|
||||
written = {where["model_id"]: json.loads(data["model_info"]) for where, data in prisma.model_table.updates}
|
||||
assert written == {"joins": {"access_groups": ["prod-models"], "description": "saved while the update waited"}}
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_access_group_emptied_while_the_update_retried_a_locked_row_answers_404_and_writes_nothing():
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
|
||||
update_access_group,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
UpdateModelGroupRequest,
|
||||
)
|
||||
|
||||
prisma = _FakePrismaClient(
|
||||
[],
|
||||
deployments=[
|
||||
_deployment(model_id="leaves", model_name="gpt-leaves"),
|
||||
_deployment(model_id="joins", model_name="gpt-joins", access_groups=()),
|
||||
],
|
||||
)
|
||||
prisma.model_table.rows_locked_by_another_session = frozenset({"leaves"})
|
||||
prisma.model_table.edits_the_other_session_commits = {"leaves": {"access_groups": []}}
|
||||
prisma.model_table.other_session_commits_after_refusals = 1
|
||||
|
||||
with _proxy_with_stubbed_reload(prisma), pytest.raises(HTTPException) as raised:
|
||||
await update_access_group(
|
||||
access_group="prod-models",
|
||||
data=UpdateModelGroupRequest(model_ids=["joins"]),
|
||||
user_api_key_dict=_admin(),
|
||||
)
|
||||
|
||||
assert raised.value.status_code == 404
|
||||
assert prisma.model_table.updates == []
|
||||
|
||||
|
||||
def _another_session_holds_a_member_row(model_table):
|
||||
model_table.rows_locked_by_another_session = frozenset({"leaves"})
|
||||
|
||||
|
||||
def _another_edit_holds_the_group(model_table):
|
||||
model_table.group_locked_by_another_session = True
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"hold", [_another_session_holds_a_member_row, _another_edit_holds_the_group], ids=["member_row", "group"]
|
||||
)
|
||||
async def test_an_access_group_update_answers_409_and_writes_nothing_while_another_session_keeps_a_lock_it_needs(
|
||||
hold,
|
||||
):
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
|
||||
_rewrite_membership_once_unlocked,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
UpdateModelGroupRequest,
|
||||
)
|
||||
|
||||
prisma = _FakePrismaClient(
|
||||
[],
|
||||
deployments=[
|
||||
_deployment(model_id="stays", model_name="gpt-stays"),
|
||||
_deployment(model_id="leaves", model_name="gpt-leaves"),
|
||||
],
|
||||
)
|
||||
hold(prisma.model_table)
|
||||
clock = _FakeClock()
|
||||
|
||||
with pytest.raises(HTTPException) as raised:
|
||||
await _rewrite_membership_once_unlocked(
|
||||
prisma_client=prisma,
|
||||
access_group="prod-models",
|
||||
data=UpdateModelGroupRequest(model_ids=["stays"]),
|
||||
retry_budget_seconds=0.5,
|
||||
sleep=clock.sleep,
|
||||
clock=clock,
|
||||
)
|
||||
|
||||
assert raised.value.status_code == 409
|
||||
assert 0 < clock.now <= 0.5
|
||||
assert prisma.model_table.lock_refusals == len(clock.slept) + 1
|
||||
assert prisma.model_table.updates == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
"hold", [_another_session_holds_a_member_row, _another_edit_holds_the_group], ids=["member_row", "group"]
|
||||
)
|
||||
async def test_an_access_group_update_waits_out_a_lock_held_for_over_ten_seconds_and_writes_the_edit(hold):
|
||||
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
|
||||
_rewrite_membership_once_unlocked,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
UpdateModelGroupRequest,
|
||||
)
|
||||
|
||||
prisma = _FakePrismaClient(
|
||||
[],
|
||||
deployments=[
|
||||
_deployment(model_id="stays", model_name="gpt-stays"),
|
||||
_deployment(model_id="leaves", model_name="gpt-leaves"),
|
||||
],
|
||||
)
|
||||
hold(prisma.model_table)
|
||||
prisma.model_table.other_session_commits_after_refusals = 20
|
||||
clock = _FakeClock()
|
||||
|
||||
_, rewrites = await _rewrite_membership_once_unlocked(
|
||||
prisma_client=prisma,
|
||||
access_group="prod-models",
|
||||
data=UpdateModelGroupRequest(model_ids=["stays"]),
|
||||
sleep=clock.sleep,
|
||||
clock=clock,
|
||||
)
|
||||
|
||||
assert clock.now > 10
|
||||
assert [model_id for model_id, _ in rewrites] == ["leaves"]
|
||||
assert [where["model_id"] for where, _ in prisma.model_table.updates] == ["leaves"]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_updating_a_missing_access_group_answers_404_before_checking_the_model_names():
|
||||
from fastapi import HTTPException
|
||||
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
|
||||
update_access_group,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
UpdateModelGroupRequest,
|
||||
)
|
||||
|
||||
prisma = _FakePrismaClient(
|
||||
[], deployments=[_deployment(model_id="other", model_name="gpt-other", access_groups=("other-group",))]
|
||||
)
|
||||
|
||||
with _proxy_with_stubbed_reload(prisma):
|
||||
_serve_model_groups(proxy_server.llm_router, ["gpt-other"])
|
||||
with pytest.raises(HTTPException) as raised:
|
||||
await update_access_group(
|
||||
access_group="prod-models",
|
||||
data=UpdateModelGroupRequest(model_names=["no-such-model"]),
|
||||
user_api_key_dict=_admin(),
|
||||
)
|
||||
|
||||
assert raised.value.status_code == 404
|
||||
assert prisma.model_table.updates == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_an_access_group_update_answers_500_when_the_deployments_cannot_be_read():
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
|
||||
update_access_group,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
UpdateModelGroupRequest,
|
||||
)
|
||||
|
||||
prisma = _FakePrismaClient([], deployments=[_deployment(model_id="stays")])
|
||||
prisma.model_table.find_many = AsyncMock(side_effect=RuntimeError("database unreachable"))
|
||||
|
||||
with _proxy_with_stubbed_reload(prisma), pytest.raises(HTTPException) as raised:
|
||||
await update_access_group(
|
||||
access_group="prod-models",
|
||||
data=UpdateModelGroupRequest(model_ids=["stays"]),
|
||||
user_api_key_dict=_admin(),
|
||||
)
|
||||
|
||||
assert raised.value.status_code == 500
|
||||
assert prisma.model_table.updates == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("model_ids", [["no-such-deployment"], ["joins", "no-such-deployment"]])
|
||||
async def test_updating_an_access_group_with_an_unknown_model_id_writes_no_deployment(model_ids):
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
|
||||
update_access_group,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
UpdateModelGroupRequest,
|
||||
)
|
||||
|
||||
prisma = _FakePrismaClient(
|
||||
[],
|
||||
deployments=[_deployment(model_id="stays"), _deployment(model_id="joins", access_groups=())],
|
||||
)
|
||||
|
||||
with _proxy_with_stubbed_reload(prisma), pytest.raises(HTTPException) as raised:
|
||||
await update_access_group(
|
||||
access_group="prod-models",
|
||||
data=UpdateModelGroupRequest(model_ids=model_ids),
|
||||
user_api_key_dict=_admin(),
|
||||
)
|
||||
|
||||
assert raised.value.status_code == 400
|
||||
assert prisma.model_table.updates == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_updating_an_access_group_no_deployment_carries_answers_404_and_writes_nothing():
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
|
||||
update_access_group,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
UpdateModelGroupRequest,
|
||||
)
|
||||
|
||||
prisma = _FakePrismaClient([], deployments=[_deployment(model_id="other", access_groups=("other-group",))])
|
||||
|
||||
with _proxy_with_stubbed_reload(prisma), pytest.raises(HTTPException) as raised:
|
||||
await update_access_group(
|
||||
access_group="prod-models",
|
||||
data=UpdateModelGroupRequest(model_ids=["other"]),
|
||||
user_api_key_dict=_admin(),
|
||||
)
|
||||
|
||||
assert raised.value.status_code == 404
|
||||
assert prisma.model_table.updates == []
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_updating_an_access_group_with_a_model_name_only_the_config_serves_writes_no_deployment():
|
||||
from fastapi import HTTPException
|
||||
|
||||
import litellm.proxy.proxy_server as proxy_server
|
||||
from litellm.proxy.management_endpoints.model_access_group_management_endpoints import (
|
||||
update_access_group,
|
||||
)
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
UpdateModelGroupRequest,
|
||||
)
|
||||
|
||||
prisma = _FakePrismaClient(
|
||||
[],
|
||||
deployments=[
|
||||
_deployment(model_id="stays", model_name="gpt-stays"),
|
||||
_deployment(model_id="joins", model_name="gpt-joins", access_groups=()),
|
||||
],
|
||||
)
|
||||
|
||||
with _proxy_with_stubbed_reload(prisma):
|
||||
_serve_model_groups(proxy_server.llm_router, ["gpt-stays", "gpt-joins", "config-only-model"])
|
||||
with pytest.raises(HTTPException) as raised:
|
||||
await update_access_group(
|
||||
access_group="prod-models",
|
||||
data=UpdateModelGroupRequest(model_names=["gpt-joins", "config-only-model"]),
|
||||
user_api_key_dict=_admin(),
|
||||
)
|
||||
|
||||
assert raised.value.status_code == 400
|
||||
assert prisma.model_table.updates == []
|
||||
|
|
|
|||
|
|
@ -6508,7 +6508,11 @@ class TestStrategyRouterWriteValidation:
|
|||
assert fake.tx_obj.raw_calls == []
|
||||
return
|
||||
assert handle is fake.tx_obj.litellm_proxymodeltable
|
||||
published.assert_awaited_once_with(redis_cache=None, object_type="litellm_proxymodeltable")
|
||||
published.assert_awaited_once_with(
|
||||
redis_cache=None,
|
||||
object_type="litellm_proxymodeltable",
|
||||
model_ids=(model_id,) if model_id else (),
|
||||
)
|
||||
(lock_sql, lock_params), (count_sql, count_params) = fake.tx_obj.raw_calls
|
||||
assert "pg_advisory_xact_lock($1)" in lock_sql and "count" not in lock_sql
|
||||
assert lock_params == (AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY,)
|
||||
|
|
|
|||
5
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
5
ui/litellm-dashboard/src/lib/http/schema.d.ts
generated
vendored
|
|
@ -699,8 +699,8 @@ export interface paths {
|
|||
* @description Update an access group's model names.
|
||||
*
|
||||
* This will:
|
||||
* 1. Remove the access group from all current deployments
|
||||
* 2. Add the access group to all deployments for the new model_names list
|
||||
* 1. Add the access group to all deployments for the new model_names list
|
||||
* 2. Remove the access group from every other deployment
|
||||
*
|
||||
* Example:
|
||||
* ```bash
|
||||
|
|
@ -722,6 +722,7 @@ export interface paths {
|
|||
* Raises:
|
||||
* - HTTPException 400: If any model names don't exist
|
||||
* - HTTPException 404: If access group not found
|
||||
* - HTTPException 409: If another edit holds the group or one of its models for 30 s
|
||||
*/
|
||||
put: operations["update_access_group_access_group__access_group__update_put"];
|
||||
post?: never;
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue