diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index bc64a4c7ef0..c76b9be7c33 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -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": [ { diff --git a/litellm/proxy/common_utils/config_sync_pubsub.py b/litellm/proxy/common_utils/config_sync_pubsub.py index da74851b2d8..e9aaa0e81e4 100644 --- a/litellm/proxy/common_utils/config_sync_pubsub.py +++ b/litellm/proxy/common_utils/config_sync_pubsub.py @@ -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: diff --git a/litellm/proxy/common_utils/registry_read_through.py b/litellm/proxy/common_utils/registry_read_through.py index 22df11429c0..487c6feb68d 100644 --- a/litellm/proxy/common_utils/registry_read_through.py +++ b/litellm/proxy/common_utils/registry_read_through.py @@ -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: diff --git a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py index 2d88882a2ac..3a2de12d108 100644 --- a/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_access_group_management_endpoints.py @@ -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, diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index cdfe3c90182..74fa8e3b3c1 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -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" diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index b0ccf74c63b..3e004f1ced7 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -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() diff --git a/tests/integration/management/test_added_backup_deployment_serves_peer_fallbacks.py b/tests/integration/management/test_added_backup_deployment_serves_peer_fallbacks.py new file mode 100644 index 00000000000..942bed6c165 --- /dev/null +++ b/tests/integration/management/test_added_backup_deployment_serves_peer_fallbacks.py @@ -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 diff --git a/tests/unit/proxy/common_utils/test_config_sync_pubsub.py b/tests/unit/proxy/common_utils/test_config_sync_pubsub.py index 92cfb6a71e5..cbb37c964c6 100644 --- a/tests/unit/proxy/common_utils/test_config_sync_pubsub.py +++ b/tests/unit/proxy/common_utils/test_config_sync_pubsub.py @@ -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 diff --git a/tests/unit/proxy/common_utils/test_registry_read_through.py b/tests/unit/proxy/common_utils/test_registry_read_through.py index 04729167707..da7b1f497c9 100644 --- a/tests/unit/proxy/common_utils/test_registry_read_through.py +++ b/tests/unit/proxy/common_utils/test_registry_read_through.py @@ -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() diff --git a/tests/unit/proxy/management_endpoints/test_access_group_management.py b/tests/unit/proxy/management_endpoints/test_access_group_management.py index c36576beaa9..d79af0af519 100644 --- a/tests/unit/proxy/management_endpoints/test_access_group_management.py +++ b/tests/unit/proxy/management_endpoints/test_access_group_management.py @@ -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 == [] diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 14d93f368c1..872d1607351 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -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,) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 13a064517a4..7d9e9839a5f 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -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;