mirror of
https://github.com/BerriAI/litellm.git
synced 2026-09-13 23:11:40 +00:00
fix(proxy): converge sibling pods on model CRUD via redis pubsub
Model reads are served from each pod's in-memory router, so a create/update/delete handled by one pod stayed invisible to its siblings until their next config reload (up to PROXY_CONFIG_RELOAD_INTERVAL_SECONDS). Behind a load balancer that made a deleted model keep reappearing in the Admin UI on refresh. The pod that handled the write now publishes a notification on the coordination Redis; every pod subscribes at startup and re-runs the same DB reconcile the reload timer performs. The notification is only a trigger, never a payload, so the DB stays the single source of truth and the periodic reload remains the backstop. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
tin
Devin AI
parent
6f1625d23b
commit
0e8fa4e081
8 changed files with 645 additions and 2 deletions
|
|
@ -2,7 +2,11 @@ import os
|
|||
import sys
|
||||
from typing import List, Literal, Optional
|
||||
|
||||
from litellm.litellm_core_utils.env_utils import get_env_int, get_env_int_or_none
|
||||
from litellm.litellm_core_utils.env_utils import (
|
||||
get_env_float,
|
||||
get_env_int,
|
||||
get_env_int_or_none,
|
||||
)
|
||||
|
||||
DEFAULT_HEALTH_CHECK_PROMPT = str(os.getenv("DEFAULT_HEALTH_CHECK_PROMPT", "test from litellm"))
|
||||
AZURE_DEFAULT_RESPONSES_API_VERSION = str(os.getenv("AZURE_DEFAULT_RESPONSES_API_VERSION", "preview"))
|
||||
|
|
@ -1478,6 +1482,11 @@ PROXY_BUDGET_RESCHEDULER_MAX_TIME = int(os.getenv("PROXY_BUDGET_RESCHEDULER_MAX_
|
|||
PROXY_BATCH_WRITE_AT = int(os.getenv("PROXY_BATCH_WRITE_AT", 10)) # in seconds, increased from 10
|
||||
PROXY_CONFIG_RELOAD_INTERVAL_SECONDS = get_env_int("PROXY_CONFIG_RELOAD_INTERVAL_SECONDS", 30)
|
||||
|
||||
MODEL_CHANGE_PUBSUB_ENABLED = os.getenv("MODEL_CHANGE_PUBSUB_ENABLED", "true").lower() in ["true", "1"]
|
||||
MODEL_CHANGE_PUBSUB_CHANNEL = os.getenv("MODEL_CHANGE_PUBSUB_CHANNEL", "litellm:model_changes")
|
||||
MODEL_CHANGE_PUBSUB_POLL_TIMEOUT_SECONDS = get_env_float("MODEL_CHANGE_PUBSUB_POLL_TIMEOUT_SECONDS", 0.5)
|
||||
MODEL_CHANGE_PUBSUB_RECONNECT_SECONDS = get_env_float("MODEL_CHANGE_PUBSUB_RECONNECT_SECONDS", 5.0)
|
||||
|
||||
# APScheduler Configuration - MEMORY LEAK FIX
|
||||
# These settings prevent memory leaks in APScheduler's normalize() and _apply_jitter() functions
|
||||
APSCHEDULER_COALESCE = os.getenv("APSCHEDULER_COALESCE", "True").lower() in [
|
||||
|
|
|
|||
|
|
@ -21,6 +21,21 @@ def get_env_int(env_var: str, default: int) -> int:
|
|||
return default
|
||||
|
||||
|
||||
def get_env_float(env_var: str, default: float) -> float:
|
||||
"""Parse an environment variable as a float, falling back to default on invalid values.
|
||||
|
||||
Same forgiving behaviour as `get_env_int`, so a typo in a tuning knob cannot crash
|
||||
the process at import time.
|
||||
"""
|
||||
raw = os.getenv(env_var)
|
||||
if raw is None:
|
||||
return default
|
||||
try:
|
||||
return float(raw.strip())
|
||||
except (ValueError, TypeError):
|
||||
return default
|
||||
|
||||
|
||||
def get_env_int_or_none(env_var: str) -> int | None:
|
||||
"""Parse an environment variable as an integer, returning None when it is unset or unusable.
|
||||
|
||||
|
|
|
|||
|
|
@ -22,6 +22,7 @@ from litellm.proxy.management_endpoints.model_management_endpoints import (
|
|||
reload_serving_verdict,
|
||||
clear_cache,
|
||||
)
|
||||
from litellm.proxy.model_change_broadcast import broadcast_model_change
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.repositories.model_repository import ModelRepository
|
||||
from litellm.types.proxy.management_endpoints.model_management_endpoints import (
|
||||
|
|
@ -419,6 +420,7 @@ async def create_model_group(
|
|||
live_before_reload = live_model_ids_snapshot()
|
||||
|
||||
await clear_cache()
|
||||
await broadcast_model_change(operation="updated")
|
||||
_raise_http_if_reload_degraded_serving(
|
||||
before=live_before_reload,
|
||||
written_models=updated_pairs,
|
||||
|
|
@ -679,6 +681,7 @@ async def update_access_group(
|
|||
# Clear cache and reload models to pick up the access group changes
|
||||
live_before_reload = live_model_ids_snapshot()
|
||||
await clear_cache()
|
||||
await broadcast_model_change(operation="updated")
|
||||
_raise_http_if_reload_degraded_serving(
|
||||
before=live_before_reload,
|
||||
written_models=list({**dict(stripped_pairs), **dict(updated_pairs)}.items()),
|
||||
|
|
@ -781,6 +784,7 @@ async def delete_access_group(
|
|||
# Clear cache and reload models to pick up the access group changes
|
||||
live_before_reload = live_model_ids_snapshot()
|
||||
await clear_cache()
|
||||
await broadcast_model_change(operation="updated")
|
||||
_raise_http_if_reload_degraded_serving(
|
||||
before=live_before_reload,
|
||||
written_models=removed_pairs,
|
||||
|
|
|
|||
|
|
@ -49,6 +49,7 @@ from litellm.proxy.management_endpoints.team_endpoints import (
|
|||
update_team as _legacy_update_team,
|
||||
)
|
||||
from litellm.proxy.management_helpers.audit_logs import create_object_audit_log
|
||||
from litellm.proxy.model_change_broadcast import broadcast_model_change
|
||||
from litellm.proxy.utils import PrismaClient
|
||||
from litellm.repositories.model_repository import ModelRepository
|
||||
from litellm.repositories.table_repositories import ModelTableRepository
|
||||
|
|
@ -325,6 +326,7 @@ async def patch_model(
|
|||
# Clear cache and reload models (uses config setting or defaults to preserving config models for DB updates)
|
||||
live_before_reload = live_model_ids_snapshot()
|
||||
await clear_cache()
|
||||
await broadcast_model_change(operation="updated", model_id=model_id)
|
||||
|
||||
## CREATE AUDIT LOG ##
|
||||
asyncio.create_task(
|
||||
|
|
@ -430,6 +432,7 @@ async def _set_model_blocked_status(
|
|||
|
||||
live_before_reload = live_model_ids_snapshot()
|
||||
await clear_cache()
|
||||
await broadcast_model_change(operation="updated", model_id=data.model_id)
|
||||
|
||||
asyncio.create_task(
|
||||
create_object_audit_log(
|
||||
|
|
@ -1186,6 +1189,8 @@ async def delete_model(
|
|||
if llm_router is not None:
|
||||
llm_router.delete_deployment(id=model_info.id)
|
||||
|
||||
await broadcast_model_change(operation="deleted", model_id=model_info.id)
|
||||
|
||||
# Runs after the row delete so the sibling check sees post-delete state.
|
||||
if model_params.model_info.team_id is not None:
|
||||
await _remove_unbacked_team_models(
|
||||
|
|
@ -1370,6 +1375,10 @@ async def add_new_model(
|
|||
prisma_client=prisma_client,
|
||||
)
|
||||
await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj)
|
||||
await broadcast_model_change(
|
||||
operation="created",
|
||||
model_id=model_response.model_id if model_response is not None else None,
|
||||
)
|
||||
# don't let failed slack alert block the /model/new response
|
||||
_alerting = general_settings.get("alerting", []) or []
|
||||
if "slack" in _alerting:
|
||||
|
|
@ -1543,6 +1552,7 @@ async def update_model(
|
|||
# Clear cache and reload models (uses config setting or defaults to preserving config models for DB updates)
|
||||
live_before_reload = live_model_ids_snapshot()
|
||||
await clear_cache()
|
||||
await broadcast_model_change(operation="updated", model_id=_model_id)
|
||||
## CREATE AUDIT LOG ##
|
||||
asyncio.create_task(
|
||||
create_object_audit_log(
|
||||
|
|
|
|||
230
litellm/proxy/model_change_broadcast.py
Normal file
230
litellm/proxy/model_change_broadcast.py
Normal file
|
|
@ -0,0 +1,230 @@
|
|||
"""
|
||||
Cross-pod invalidation for model CRUD.
|
||||
|
||||
Each pod serves reads (`/model/info`, request routing) from its own in-memory
|
||||
`llm_router`, reconciled with the DB every `PROXY_CONFIG_RELOAD_INTERVAL_SECONDS`.
|
||||
A write handled by one pod is therefore invisible to its siblings for up to a full
|
||||
interval, so a model deleted in the Admin UI keeps showing up on refreshes the load
|
||||
balancer sends elsewhere.
|
||||
|
||||
The pod that handled the write publishes a notification on the coordination Redis;
|
||||
subscribers react by re-running the same DB reconcile the timer already performs.
|
||||
The notification is a trigger, never a payload: siblings always derive state from
|
||||
the DB, so a dropped, duplicated or reordered message can only cost latency, and the
|
||||
periodic reconcile remains the backstop.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import contextlib
|
||||
from collections.abc import Mapping
|
||||
from typing import Awaitable, Callable, Final, Literal, Protocol
|
||||
|
||||
from pydantic import BaseModel, ConfigDict, ValidationError
|
||||
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
from litellm._uuid import uuid
|
||||
from litellm.constants import (
|
||||
MODEL_CHANGE_PUBSUB_CHANNEL,
|
||||
MODEL_CHANGE_PUBSUB_ENABLED,
|
||||
MODEL_CHANGE_PUBSUB_POLL_TIMEOUT_SECONDS,
|
||||
MODEL_CHANGE_PUBSUB_RECONNECT_SECONDS,
|
||||
)
|
||||
|
||||
ModelChangeOperation = Literal["created", "updated", "deleted"]
|
||||
|
||||
PROCESS_ID: Final[str] = str(uuid.uuid4())
|
||||
|
||||
|
||||
class ModelChangeNotification(BaseModel):
|
||||
model_config = ConfigDict(frozen=True)
|
||||
|
||||
operation: ModelChangeOperation
|
||||
model_id: str | None = None
|
||||
origin: str
|
||||
|
||||
|
||||
class RedisPubSubConnection(Protocol):
|
||||
async def subscribe(self, *channels: str) -> None: ...
|
||||
|
||||
async def get_message(
|
||||
self,
|
||||
ignore_subscribe_messages: bool = ...,
|
||||
timeout: float | None = ...,
|
||||
) -> Mapping[str, object] | None: ...
|
||||
|
||||
async def aclose(self) -> None: ...
|
||||
|
||||
|
||||
class RedisPubSubClient(Protocol):
|
||||
async def publish(self, channel: str, message: str) -> int: ...
|
||||
|
||||
def pubsub(self) -> RedisPubSubConnection: ...
|
||||
|
||||
|
||||
class RedisPubSubBackend(Protocol):
|
||||
"""The slice of `RedisCache` this module needs."""
|
||||
|
||||
def check_and_fix_namespace(self, key: str) -> str: ...
|
||||
|
||||
def init_async_client(self) -> RedisPubSubClient: ...
|
||||
|
||||
|
||||
def _coordination_redis() -> RedisPubSubBackend | None:
|
||||
from litellm.proxy.proxy_server import redis_usage_cache
|
||||
|
||||
return redis_usage_cache
|
||||
|
||||
|
||||
async def broadcast_model_change(
|
||||
operation: ModelChangeOperation,
|
||||
model_id: str | None = None,
|
||||
redis_cache: RedisPubSubBackend | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Tell sibling pods that the model table changed. Never raises: the write it
|
||||
follows has already succeeded, and the periodic reconcile still converges.
|
||||
"""
|
||||
if not MODEL_CHANGE_PUBSUB_ENABLED:
|
||||
return
|
||||
|
||||
backend = redis_cache if redis_cache is not None else _coordination_redis()
|
||||
if backend is None:
|
||||
return
|
||||
|
||||
notification = ModelChangeNotification(operation=operation, model_id=model_id, origin=PROCESS_ID)
|
||||
try:
|
||||
client = backend.init_async_client()
|
||||
await client.publish(
|
||||
backend.check_and_fix_namespace(MODEL_CHANGE_PUBSUB_CHANNEL),
|
||||
notification.model_dump_json(),
|
||||
)
|
||||
except Exception as e: # noqa: BLE001 # no redis error may fail a write that already succeeded
|
||||
verbose_proxy_logger.warning(
|
||||
"Could not broadcast model change (%s, model_id=%s) to other pods: %s. "
|
||||
"They will pick it up on their next config reload.",
|
||||
operation,
|
||||
model_id,
|
||||
str(e),
|
||||
)
|
||||
|
||||
|
||||
def _parse_notification(payload: object) -> ModelChangeNotification | None:
|
||||
raw = payload.decode() if isinstance(payload, bytes) else payload
|
||||
if not isinstance(raw, str):
|
||||
return None
|
||||
try:
|
||||
return ModelChangeNotification.model_validate_json(raw)
|
||||
except ValidationError:
|
||||
return None
|
||||
|
||||
|
||||
class ModelChangeSubscriber:
|
||||
"""
|
||||
Listens for model-change notifications and re-runs `reconcile` for each burst.
|
||||
|
||||
`reconcile` is injected (the proxy passes `ProxyConfig.add_deployment`), so the
|
||||
subscriber owns no knowledge of how the router is rebuilt.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
redis_cache: RedisPubSubBackend,
|
||||
reconcile: Callable[[], Awaitable[None]],
|
||||
origin: str = PROCESS_ID,
|
||||
poll_timeout_seconds: float = MODEL_CHANGE_PUBSUB_POLL_TIMEOUT_SECONDS,
|
||||
reconnect_seconds: float = MODEL_CHANGE_PUBSUB_RECONNECT_SECONDS,
|
||||
) -> None:
|
||||
self._redis_cache = redis_cache
|
||||
self._reconcile = reconcile
|
||||
self._origin = origin
|
||||
self._poll_timeout_seconds = poll_timeout_seconds
|
||||
self._reconnect_seconds = reconnect_seconds
|
||||
|
||||
@property
|
||||
def channel(self) -> str:
|
||||
return self._redis_cache.check_and_fix_namespace(MODEL_CHANGE_PUBSUB_CHANNEL)
|
||||
|
||||
async def listen_forever(self) -> None:
|
||||
while True:
|
||||
try:
|
||||
await self.listen_once()
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as e: # noqa: BLE001 # a subscriber that dies leaves this pod on the reload interval
|
||||
verbose_proxy_logger.warning(
|
||||
"Model-change subscriber dropped off %s: %s. Reconnecting in %ss; the "
|
||||
"periodic config reload keeps this pod converging meanwhile.",
|
||||
self.channel,
|
||||
str(e),
|
||||
self._reconnect_seconds,
|
||||
)
|
||||
await asyncio.sleep(self._reconnect_seconds)
|
||||
|
||||
async def listen_once(self) -> None:
|
||||
"""
|
||||
One connection's lifetime. Notifications are coalesced: a burst of writes
|
||||
(deleting several models in a row) triggers a single reconcile once the channel
|
||||
goes quiet, instead of one full DB reload per message.
|
||||
"""
|
||||
connection = self._redis_cache.init_async_client().pubsub()
|
||||
await connection.subscribe(self.channel)
|
||||
verbose_proxy_logger.info("Subscribed to model changes from other pods on %s", self.channel)
|
||||
try:
|
||||
pending = False
|
||||
while True:
|
||||
message = await connection.get_message(
|
||||
ignore_subscribe_messages=True,
|
||||
timeout=self._poll_timeout_seconds,
|
||||
)
|
||||
if message is None:
|
||||
if pending:
|
||||
pending = False
|
||||
await self._reconcile()
|
||||
continue
|
||||
pending = pending or self._is_foreign_change(message)
|
||||
finally:
|
||||
with contextlib.suppress(Exception):
|
||||
await connection.aclose()
|
||||
|
||||
def _is_foreign_change(self, message: Mapping[str, object]) -> bool:
|
||||
if message.get("type") != "message":
|
||||
return False
|
||||
notification = _parse_notification(message.get("data"))
|
||||
if notification is None:
|
||||
verbose_proxy_logger.debug("Ignoring unparseable model-change notification: %s", message.get("data"))
|
||||
return False
|
||||
return notification.origin != self._origin
|
||||
|
||||
|
||||
class ModelChangeSubscriberHandle:
|
||||
"""
|
||||
Owns the subscriber task for the lifetime of the process. Holding the reference
|
||||
matters: an unreferenced asyncio task can be garbage collected mid-flight.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self._task: asyncio.Task[None] | None = None
|
||||
|
||||
def start(
|
||||
self,
|
||||
redis_cache: RedisPubSubBackend | None,
|
||||
reconcile: Callable[[], Awaitable[None]],
|
||||
) -> None:
|
||||
if not MODEL_CHANGE_PUBSUB_ENABLED or redis_cache is None:
|
||||
return
|
||||
self.stop()
|
||||
subscriber = ModelChangeSubscriber(redis_cache=redis_cache, reconcile=reconcile)
|
||||
self._task = asyncio.create_task(subscriber.listen_forever())
|
||||
|
||||
def stop(self) -> None:
|
||||
if self._task is None:
|
||||
return
|
||||
self._task.cancel()
|
||||
self._task = None
|
||||
|
||||
@property
|
||||
def is_running(self) -> bool:
|
||||
return self._task is not None and not self._task.done()
|
||||
|
||||
|
||||
model_change_subscriber = ModelChangeSubscriberHandle()
|
||||
|
|
@ -213,7 +213,7 @@ def generate_feedback_box():
|
|||
import contextlib
|
||||
from collections import defaultdict
|
||||
from contextlib import asynccontextmanager
|
||||
from functools import lru_cache
|
||||
from functools import lru_cache, partial
|
||||
|
||||
import litellm
|
||||
import litellm._redis
|
||||
|
|
@ -468,6 +468,7 @@ from litellm.proxy.middleware.billable_request_metrics_middleware import (
|
|||
BillableRequestMetricsMiddleware,
|
||||
BillingRecorder,
|
||||
)
|
||||
from litellm.proxy.model_change_broadcast import model_change_subscriber
|
||||
from litellm.proxy.plugin_routes import (
|
||||
register_plugins_from_config,
|
||||
)
|
||||
|
|
@ -797,6 +798,9 @@ def cleanup_router_config_variables():
|
|||
async def proxy_shutdown_event():
|
||||
global prisma_client, master_key, user_custom_auth, user_custom_key_generate, user_custom_key_update
|
||||
verbose_proxy_logger.info("Shutting down LiteLLM Proxy Server")
|
||||
|
||||
model_change_subscriber.stop()
|
||||
|
||||
if prisma_client:
|
||||
verbose_proxy_logger.debug("Disconnecting from Prisma")
|
||||
await prisma_client.disconnect()
|
||||
|
|
@ -8122,6 +8126,16 @@ class ProxyStartupEvent:
|
|||
# this will load all existing models on proxy startup
|
||||
await proxy_config.add_deployment(prisma_client=prisma_client, proxy_logging_obj=proxy_logging_obj)
|
||||
|
||||
### REACT TO MODEL CHANGES MADE ON OTHER PODS ###
|
||||
model_change_subscriber.start(
|
||||
redis_cache=redis_usage_cache,
|
||||
reconcile=partial(
|
||||
proxy_config.add_deployment,
|
||||
prisma_client=prisma_client,
|
||||
proxy_logging_obj=proxy_logging_obj,
|
||||
),
|
||||
)
|
||||
|
||||
### GET STORED CREDENTIALS ###
|
||||
scheduler.add_job(
|
||||
proxy_config.get_credentials,
|
||||
|
|
|
|||
|
|
@ -1,3 +1,4 @@
|
|||
import asyncio
|
||||
import json
|
||||
import os
|
||||
import sys
|
||||
|
|
@ -1938,6 +1939,100 @@ class TestAddAndDeleteModelLifecycle:
|
|||
assert str(exc_info.value.code) == "400"
|
||||
|
||||
|
||||
class TestModelCrudNotifiesOtherPods:
|
||||
"""Regression: a model deleted (or added) on one pod stayed in every sibling pod's
|
||||
in-memory router until its next config reload, so the Admin UI kept showing the
|
||||
deleted model on refreshes the load balancer sent to another pod."""
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_add_and_delete_publish_model_change_notifications(self):
|
||||
import fakeredis.aioredis
|
||||
|
||||
from litellm.constants import MODEL_CHANGE_PUBSUB_CHANNEL
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
ModelInfoDelete,
|
||||
add_new_model,
|
||||
)
|
||||
from litellm.proxy.management_endpoints.model_management_endpoints import (
|
||||
delete_model as delete_model_endpoint,
|
||||
)
|
||||
from litellm.proxy.model_change_broadcast import ModelChangeNotification
|
||||
|
||||
class FakeRedisCache:
|
||||
def __init__(self, client: fakeredis.aioredis.FakeRedis) -> None:
|
||||
self._client = client
|
||||
|
||||
def check_and_fix_namespace(self, key: str) -> str:
|
||||
return key
|
||||
|
||||
def init_async_client(self) -> fakeredis.aioredis.FakeRedis:
|
||||
return self._client
|
||||
|
||||
model_id = "notify-siblings-model-1"
|
||||
admin_user = UserAPIKeyAuth(user_id="test-admin", user_role=LitellmUserRoles.PROXY_ADMIN)
|
||||
db_row = LiteLLM_ProxyModelTable(
|
||||
model_id=model_id,
|
||||
model_name="notify-model",
|
||||
litellm_params={"model": "openai/gpt-4.1-nano"},
|
||||
model_info={"id": model_id},
|
||||
created_by="test-admin",
|
||||
updated_by="test-admin",
|
||||
)
|
||||
|
||||
mock_prisma = MagicMock()
|
||||
mock_prisma.db = MagicMock()
|
||||
mock_prisma.db.litellm_proxymodeltable = AsyncMock()
|
||||
mock_prisma.db.litellm_proxymodeltable.create = AsyncMock(return_value=db_row)
|
||||
mock_prisma.db.litellm_proxymodeltable.find_unique = AsyncMock(return_value=db_row)
|
||||
mock_prisma.db.litellm_proxymodeltable.delete = AsyncMock(return_value=db_row)
|
||||
|
||||
mock_proxy_config = MagicMock()
|
||||
mock_proxy_config.add_deployment = AsyncMock()
|
||||
|
||||
mock_router = MagicMock()
|
||||
mock_router.get_model_ids.return_value = [model_id]
|
||||
|
||||
redis_cache = FakeRedisCache(fakeredis.aioredis.FakeRedis(server=fakeredis.FakeServer()))
|
||||
sibling_pod = redis_cache.init_async_client().pubsub()
|
||||
await sibling_pod.subscribe(MODEL_CHANGE_PUBSUB_CHANNEL)
|
||||
|
||||
_PS = "litellm.proxy.proxy_server"
|
||||
_ENCRYPT = "litellm.proxy.management_endpoints.model_management_endpoints.encrypt_value_helper"
|
||||
with (
|
||||
patch(f"{_PS}.prisma_client", mock_prisma),
|
||||
patch(f"{_PS}.store_model_in_db", True),
|
||||
patch(f"{_PS}.proxy_config", mock_proxy_config),
|
||||
patch(f"{_PS}.proxy_logging_obj", MagicMock()),
|
||||
patch(f"{_PS}.general_settings", {}),
|
||||
patch(f"{_PS}.premium_user", True),
|
||||
patch(f"{_PS}.llm_router", mock_router),
|
||||
patch(f"{_PS}.redis_usage_cache", redis_cache),
|
||||
patch(_ENCRYPT, side_effect=lambda value, **kwargs: value),
|
||||
):
|
||||
await add_new_model(
|
||||
model_params=Deployment(
|
||||
model_name="notify-model",
|
||||
litellm_params=LiteLLM_Params(model="openai/gpt-4.1-nano", api_key="fake-key"),
|
||||
model_info={"id": model_id},
|
||||
),
|
||||
user_api_key_dict=admin_user,
|
||||
)
|
||||
await delete_model_endpoint(
|
||||
model_info=ModelInfoDelete(id=model_id),
|
||||
user_api_key_dict=admin_user,
|
||||
)
|
||||
|
||||
notifications = []
|
||||
deadline = asyncio.get_event_loop().time() + 3.0
|
||||
while len(notifications) < 2 and asyncio.get_event_loop().time() < deadline:
|
||||
message = await sibling_pod.get_message(ignore_subscribe_messages=True, timeout=0.05)
|
||||
if message is not None:
|
||||
notifications.append(ModelChangeNotification.model_validate_json(message["data"]))
|
||||
|
||||
assert [notification.operation for notification in notifications] == ["created", "deleted"]
|
||||
assert {notification.model_id for notification in notifications} == {model_id}
|
||||
|
||||
|
||||
class TestDeleteTeamBYOKModelGhost:
|
||||
"""Regression for issue #22594.
|
||||
|
||||
|
|
|
|||
266
tests/test_litellm/proxy/test_model_change_broadcast.py
Normal file
266
tests/test_litellm/proxy/test_model_change_broadcast.py
Normal file
|
|
@ -0,0 +1,266 @@
|
|||
import asyncio
|
||||
import json
|
||||
|
||||
import fakeredis.aioredis
|
||||
import pytest
|
||||
|
||||
from litellm.constants import MODEL_CHANGE_PUBSUB_CHANNEL
|
||||
from litellm.proxy.model_change_broadcast import (
|
||||
ModelChangeNotification,
|
||||
ModelChangeSubscriber,
|
||||
ModelChangeSubscriberHandle,
|
||||
broadcast_model_change,
|
||||
)
|
||||
|
||||
|
||||
class FakeRedisBackend:
|
||||
"""Stands in for `RedisCache`, exposing only what the broadcast module uses."""
|
||||
|
||||
def __init__(self, client: fakeredis.aioredis.FakeRedis, namespace: str | None = None) -> None:
|
||||
self._client = client
|
||||
self._namespace = namespace
|
||||
|
||||
def check_and_fix_namespace(self, key: str) -> str:
|
||||
if self._namespace is None:
|
||||
return key
|
||||
return f"{self._namespace}:{key}"
|
||||
|
||||
def init_async_client(self) -> fakeredis.aioredis.FakeRedis:
|
||||
return self._client
|
||||
|
||||
|
||||
class ExplodingRedisBackend:
|
||||
def check_and_fix_namespace(self, key: str) -> str:
|
||||
return key
|
||||
|
||||
def init_async_client(self) -> fakeredis.aioredis.FakeRedis:
|
||||
raise ConnectionError("redis is down")
|
||||
|
||||
|
||||
def _backend(namespace: str | None = None) -> FakeRedisBackend:
|
||||
server = fakeredis.FakeServer()
|
||||
return FakeRedisBackend(fakeredis.aioredis.FakeRedis(server=server), namespace=namespace)
|
||||
|
||||
|
||||
async def _drain(pubsub, expected: int, timeout: float = 2.0) -> list[dict]:
|
||||
deadline = asyncio.get_event_loop().time() + timeout
|
||||
messages: list[dict] = []
|
||||
while len(messages) < expected and asyncio.get_event_loop().time() < deadline:
|
||||
message = await pubsub.get_message(ignore_subscribe_messages=True, timeout=0.05)
|
||||
if message is not None:
|
||||
messages.append(message)
|
||||
return messages
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_broadcast_publishes_notification_on_channel():
|
||||
backend = _backend()
|
||||
pubsub = backend.init_async_client().pubsub()
|
||||
await pubsub.subscribe(MODEL_CHANGE_PUBSUB_CHANNEL)
|
||||
|
||||
await broadcast_model_change(operation="deleted", model_id="model-1", redis_cache=backend)
|
||||
|
||||
messages = await _drain(pubsub, expected=1)
|
||||
assert len(messages) == 1
|
||||
notification = ModelChangeNotification.model_validate_json(messages[0]["data"])
|
||||
assert notification.operation == "deleted"
|
||||
assert notification.model_id == "model-1"
|
||||
assert notification.origin != ""
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_broadcast_respects_redis_namespace():
|
||||
backend = _backend(namespace="tenant-a")
|
||||
pubsub = backend.init_async_client().pubsub()
|
||||
await pubsub.subscribe(f"tenant-a:{MODEL_CHANGE_PUBSUB_CHANNEL}")
|
||||
|
||||
await broadcast_model_change(operation="created", model_id="model-1", redis_cache=backend)
|
||||
|
||||
assert len(await _drain(pubsub, expected=1)) == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_broadcast_never_raises_when_redis_is_unavailable():
|
||||
await broadcast_model_change(operation="deleted", model_id="model-1", redis_cache=ExplodingRedisBackend())
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_broadcast_is_a_noop_without_a_coordination_redis():
|
||||
await broadcast_model_change(operation="deleted", model_id="model-1", redis_cache=None)
|
||||
|
||||
|
||||
class ReconcileSpy:
|
||||
def __init__(self) -> None:
|
||||
self.calls = 0
|
||||
|
||||
async def __call__(self) -> None:
|
||||
self.calls += 1
|
||||
|
||||
|
||||
async def _run_subscriber(subscriber: ModelChangeSubscriber) -> "asyncio.Task[None]":
|
||||
task = asyncio.create_task(subscriber.listen_once())
|
||||
await asyncio.sleep(0.2)
|
||||
return task
|
||||
|
||||
|
||||
async def _stop(task: "asyncio.Task[None]") -> None:
|
||||
task.cancel()
|
||||
with pytest.raises(asyncio.CancelledError):
|
||||
await task
|
||||
|
||||
|
||||
async def _wait_for_reconciles(spy: ReconcileSpy, expected: int, timeout: float = 3.0) -> None:
|
||||
deadline = asyncio.get_event_loop().time() + timeout
|
||||
while spy.calls < expected and asyncio.get_event_loop().time() < deadline:
|
||||
await asyncio.sleep(0.05)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subscriber_reconciles_on_a_change_from_another_pod():
|
||||
"""The regression: a model deleted on pod A must reach pod B's router without
|
||||
waiting for the periodic config reload."""
|
||||
backend = _backend()
|
||||
spy = ReconcileSpy()
|
||||
subscriber = ModelChangeSubscriber(
|
||||
redis_cache=backend,
|
||||
reconcile=spy,
|
||||
origin="pod-b",
|
||||
poll_timeout_seconds=0.05,
|
||||
)
|
||||
task = await _run_subscriber(subscriber)
|
||||
|
||||
await broadcast_model_change(operation="deleted", model_id="model-1", redis_cache=backend)
|
||||
await _wait_for_reconciles(spy, expected=1)
|
||||
await _stop(task)
|
||||
|
||||
assert spy.calls == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subscriber_ignores_its_own_changes():
|
||||
backend = _backend()
|
||||
spy = ReconcileSpy()
|
||||
origin = "pod-a"
|
||||
subscriber = ModelChangeSubscriber(
|
||||
redis_cache=backend,
|
||||
reconcile=spy,
|
||||
origin=origin,
|
||||
poll_timeout_seconds=0.05,
|
||||
)
|
||||
task = await _run_subscriber(subscriber)
|
||||
|
||||
own_notification = ModelChangeNotification(operation="deleted", model_id="model-1", origin=origin)
|
||||
await backend.init_async_client().publish(subscriber.channel, own_notification.model_dump_json())
|
||||
await asyncio.sleep(0.5)
|
||||
await _stop(task)
|
||||
|
||||
assert spy.calls == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subscriber_coalesces_a_burst_into_one_reconcile():
|
||||
backend = _backend()
|
||||
spy = ReconcileSpy()
|
||||
subscriber = ModelChangeSubscriber(
|
||||
redis_cache=backend,
|
||||
reconcile=spy,
|
||||
origin="pod-b",
|
||||
poll_timeout_seconds=0.05,
|
||||
)
|
||||
task = await _run_subscriber(subscriber)
|
||||
|
||||
for model_id in ("model-1", "model-2", "model-3"):
|
||||
await broadcast_model_change(operation="deleted", model_id=model_id, redis_cache=backend)
|
||||
await _wait_for_reconciles(spy, expected=1)
|
||||
await asyncio.sleep(0.3)
|
||||
await _stop(task)
|
||||
|
||||
assert spy.calls == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_reconciles_while_running_and_stops_cleanly():
|
||||
backend = _backend()
|
||||
spy = ReconcileSpy()
|
||||
handle = ModelChangeSubscriberHandle()
|
||||
|
||||
async def publish_from_another_pod(model_id: str) -> None:
|
||||
notification = ModelChangeNotification(operation="deleted", model_id=model_id, origin="pod-a")
|
||||
await backend.init_async_client().publish(
|
||||
backend.check_and_fix_namespace(MODEL_CHANGE_PUBSUB_CHANNEL),
|
||||
notification.model_dump_json(),
|
||||
)
|
||||
|
||||
handle.start(redis_cache=backend, reconcile=spy)
|
||||
assert handle.is_running
|
||||
await asyncio.sleep(0.2)
|
||||
await publish_from_another_pod("model-1")
|
||||
await _wait_for_reconciles(spy, expected=1)
|
||||
|
||||
handle.stop()
|
||||
await asyncio.sleep(0.1)
|
||||
assert not handle.is_running
|
||||
assert spy.calls == 1
|
||||
|
||||
await publish_from_another_pod("model-2")
|
||||
await asyncio.sleep(0.5)
|
||||
assert spy.calls == 1
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_handle_is_a_noop_without_a_coordination_redis():
|
||||
handle = ModelChangeSubscriberHandle()
|
||||
handle.start(redis_cache=None, reconcile=ReconcileSpy())
|
||||
assert not handle.is_running
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subscriber_ignores_unparseable_payloads():
|
||||
backend = _backend()
|
||||
spy = ReconcileSpy()
|
||||
subscriber = ModelChangeSubscriber(
|
||||
redis_cache=backend,
|
||||
reconcile=spy,
|
||||
origin="pod-b",
|
||||
poll_timeout_seconds=0.05,
|
||||
)
|
||||
task = await _run_subscriber(subscriber)
|
||||
client = backend.init_async_client()
|
||||
|
||||
await client.publish(subscriber.channel, "not json")
|
||||
await client.publish(subscriber.channel, json.dumps({"operation": "exploded"}))
|
||||
await asyncio.sleep(0.5)
|
||||
await _stop(task)
|
||||
|
||||
assert spy.calls == 0
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_subscriber_resubscribes_after_a_failed_reconcile():
|
||||
"""A reconcile that blows up (DB hiccup) must not take the subscriber down for
|
||||
the rest of the pod's life."""
|
||||
backend = _backend()
|
||||
calls: list[int] = []
|
||||
|
||||
async def flaky_reconcile() -> None:
|
||||
calls.append(len(calls))
|
||||
if len(calls) == 1:
|
||||
raise RuntimeError("db unavailable")
|
||||
|
||||
subscriber = ModelChangeSubscriber(
|
||||
redis_cache=backend,
|
||||
reconcile=flaky_reconcile,
|
||||
origin="pod-b",
|
||||
poll_timeout_seconds=0.05,
|
||||
reconnect_seconds=0.05,
|
||||
)
|
||||
task = asyncio.create_task(subscriber.listen_forever())
|
||||
await asyncio.sleep(0.2)
|
||||
|
||||
await broadcast_model_change(operation="deleted", model_id="model-1", redis_cache=backend)
|
||||
await asyncio.sleep(0.5)
|
||||
await broadcast_model_change(operation="deleted", model_id="model-2", redis_cache=backend)
|
||||
await asyncio.sleep(0.5)
|
||||
await _stop(task)
|
||||
|
||||
assert len(calls) == 2
|
||||
Loading…
Add table
Reference in a new issue