scriptoriumadmin/litellm
1
0
Fork 0
mirror of https://github.com/BerriAI/litellm.git synced 2026-09-24 00:52:24 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions 91

fix(proxy): publish auth cache invalidations in the background so a wedged coordination Redis cannot stall user updates (#42534)

Browse source 
* fix(proxy): bound auth cache invalidation publish so a wedged coordination Redis cannot stall user updates

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): resolve publish callable at call time in evict_and_broadcast

The keyword-only default bound publish_auth_cache_invalidation at
function-definition time, so tests patching the module attribute observed
zero calls. Default to None, resolve the real publisher inside the body,
and keep the keyword-shaped cache_key call the existing contract asserts

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): publish auth cache invalidations in the background so a wedged coordination Redis costs handlers nothing

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

* fix(proxy): cap in-flight auth cache invalidation publishes so a wedge cannot drain the redis pool

Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>

---------

Co-authored-by: yucheng <yucheng@berri.ai>
Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com>
This commit is contained in:
devin-ai-integration[bot] 2026-09-22 23:33:32 -07:00 • committed by GitHub
parent 320ad73f56
commit 40ec84caa2
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
5 changed files with 514 additions and 21 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

59
litellm/proxy/common_utils/auth_cache_invalidation_pubsub.py
View file

@ -18,6 +18,10 @@ if TYPE_CHECKING:
AUTH_CACHE_INVALIDATION_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation"
_POLL_TIMEOUT_SECONDS: Final = 1.0
_MAX_PENDING_PUBLISHES: Final = 1024
_MAX_IN_FLIGHT_PUBLISHES: Final = 16
_pending_publishes: Final[set[asyncio.Task[None]]] = set() # mutable-ok: strong refs keep background publishes alive
_in_flight_publishes: Final = asyncio.Semaphore(_MAX_IN_FLIGHT_PUBLISHES)
_BACKOFF_INITIAL_SECONDS: Final = 5.0
_BACKOFF_MAX_SECONDS: Final = 60.0
@ -67,6 +71,21 @@ def _message_from_data(data: object) -> _CacheInvalidationMessage | None:
)
async def _publish_to_redis(redis_cache: "RedisCache", cache_key: str, message: str) -> None:
try:
client: Final = _pubsub_capable_client(redis_cache)
if client is None:
verbose_proxy_logger.debug(
"auth cache invalidation publish for %s skipped: cluster redis client has no pub/sub support",
cache_key,
)
return
async with _in_flight_publishes:
await client.publish(auth_cache_invalidation_channel(redis_cache), message)
except Exception as e: # noqa: BLE001 # best-effort publish; mutations must never fail on redis errors
verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, e)
async def publish_auth_cache_invalidation(
cache_key: str, new_value: float | None = None, ttl: float | None = None
) -> None:
@ -80,24 +99,34 @@ async def publish_auth_cache_invalidation(
writes the value into its additional in-memory caches rather than deleting
the key. A spend reset uses this so the handler's self-delivered message
cannot erase the freshly-written post-reset counter or floor marker.
The Redis round trip runs as a background task: this call returns once the
publish has been handed to the event loop, so a Redis that accepts
connections but never replies costs the caller nothing. The DB write has
already committed and the local eviction already happened, so the caller
has nothing to do with the publish result. At most 16 publishes hold a
Redis connection at once; the rest wait in the task set, so a wedge cannot
drain the shared connection pool.
"""
redis_cache: Final = coordination_redis_cache()
if redis_cache is None:
return
try:
client: Final = _pubsub_capable_client(redis_cache)
if client is None:
verbose_proxy_logger.debug(
"auth cache invalidation publish for %s skipped: cluster redis client has no pub/sub support",
cache_key,
)
return
await client.publish(
auth_cache_invalidation_channel(redis_cache),
_cache_invalidation_message_json(cache_key, new_value=new_value, ttl=ttl),
_pending_publishes.difference_update({task for task in _pending_publishes if task.done()})
if len(_pending_publishes) >= _MAX_PENDING_PUBLISHES:
verbose_proxy_logger.warning(
"auth cache invalidation publish for %s dropped: %d publishes already waiting on redis; "
"other workers keep their cached copy until its TTL expires",
cache_key,
len(_pending_publishes),
)
except Exception as e: # noqa: BLE001 # best-effort publish; mutations must never fail on redis errors
verbose_proxy_logger.warning("auth cache invalidation publish for %s failed: %s", cache_key, e)
return
task: Final = asyncio.create_task(
_publish_to_redis(
redis_cache, cache_key, _cache_invalidation_message_json(cache_key, new_value=new_value, ttl=ttl)
)
)
_pending_publishes.add(task)
await asyncio.sleep(0)
async def evict_and_broadcast(cache_keys: Sequence[str], user_api_key_cache: "UserApiKeyCache") -> None:
@ -106,8 +135,8 @@ async def evict_and_broadcast(cache_keys: Sequence[str], user_api_key_cache: "Us
Every endpoint that mutates a cached object must call this: auth serves those objects
cache-first with no freshness check, so a mutation that leaves the entry in place keeps the
stale object enforced until its TTL expires (LIT-3803). Best-effort on both steps: the DB write
has already committed, so a cache backend error must not fail the endpoint.
stale object enforced until its TTL expires (LIT-3803). Best-effort: the DB write has already
committed, so a cache backend error must not fail the endpoint.
"""
for cache_key in cache_keys:
try:

8
tests/integration/contracts.json
View file

@ -2019,6 +2019,14 @@
],
"tests/integration/observability/test_callback_delivery.py::test_response_survives_raising_success_deployment_hook[videos]": [
"other.observability.callbacks.raising_success_deployment_hook_keeps_response"
],
"tests/integration/management/test_user_updates_wedged_coordination_redis.py::test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged": [
"mgmt.user.update.budget_change_returns_promptly_with_wedged_coordination_redis",
"mgmt.user.bulk_update.budget_change_returns_promptly_with_wedged_coordination_redis",
"mgmt.customer.update.budget_change_returns_promptly_with_wedged_coordination_redis",
"mgmt.key.reset_spend.returns_promptly_with_wedged_coordination_redis",
"mgmt.auth_cache_invalidation.publish_parked_by_short_redis_wedge_lands_after_recovery",
"mgmt.auth_cache_invalidation.burst_with_worker_kill_keeps_serving_while_redis_wedged"
]
},
"browser": {

12
tests/integration/coordination_redis_proxy_config.yaml Normal file
View file

@ -0,0 +1,12 @@
model_list: []
general_settings:
master_key: os.environ/LITELLM_MASTER_KEY
database_url: os.environ/DATABASE_URL
store_model_in_db: true
disable_spend_logs: false
proxy_batch_write_at: 1
coordination_redis:
host: os.environ/REDIS_HOST
port: os.environ/REDIS_PORT
router_settings:
disable_cooldowns: true

379
tests/integration/management/test_user_updates_wedged_coordination_redis.py Normal file
View file

@ -0,0 +1,379 @@
import os
import signal
import time
import uuid
from collections.abc import Callable, Mapping
from concurrent.futures import ThreadPoolExecutor
from pathlib import Path
from typing import Final
from urllib.parse import urlsplit, urlunsplit
import psutil
import psycopg
import pytest
from psycopg import sql
from pydantic import JsonValue
from redis import Redis
from redis.client import PubSub
from tests.integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value
from tests.integration._support.process import owned_proxy
from tests.integration._support.redis_process import owned_redis
_USERS: Final = 60
_BURST: Final = 30
_HANDLER_BUDGET_SECONDS: Final = 0.75
_BULK_BUDGET_SECONDS: Final = 2.0
_CHANNEL: Final = "litellm_proxy.auth_cache_invalidation"
def _timed_post(candidate: Gateway, path: str, body: Mapping[str, JsonValue], timeout: float = 15) -> float:
started: Final = time.monotonic()
response: Final = candidate.client.request(
"POST",
path,
json=body,
headers={"Authorization": f"Bearer {candidate.key}"},
timeout=timeout,
)
elapsed: Final = time.monotonic() - started
assert response.status_code == 200, f"POST {path}: {response.status_code} {response.text} after {elapsed:.3f}s"
return elapsed
def _received(pubsub: PubSub) -> tuple[dict[str, JsonValue], ...]:
messages: list[dict[str, JsonValue]] = []
while True:
message = pubsub.get_message(ignore_subscribe_messages=True, timeout=0)
if message is None:
return tuple(messages)
data = message.get("data")
if isinstance(data, (bytes, str)):
messages.append(JSON_OBJECT.validate_json(data))
def _worker_pid(port: int) -> int:
for process in psutil.process_iter():
parent = process.parent()
if parent is None:
continue
try:
cmdline = parent.cmdline()
own_cmdline = process.cmdline()
except (psutil.NoSuchProcess, psutil.AccessDenied):
continue
if (
"integration._support.proxy" in cmdline
and "--port" in cmdline
and str(port) in cmdline
and not any("prisma" in part for part in own_cmdline)
):
return process.pid
raise AssertionError(f"no uvicorn worker found under the owned proxy on port {port}")
def _burst_call(
index: int, users: tuple[str, ...], key: str, team_id: str, customer_id: str
) -> tuple[str, dict[str, JsonValue]]:
match index % 5:
case 0:
return "/user/update", {"user_id": users[index], "max_budget": 200.0 + index}
case 1:
return "/user/update", {"user_id": users[index], "tpm_limit": 1000 + index}
case 2:
return "/key/update", {"key": key, "max_budget": 7.0 + index}
case 3:
return "/team/update", {"team_id": team_id, "max_budget": 7.0 + index}
case _:
return "/customer/update", {"user_id": customer_id, "max_budget": 7.0 + index}
@pytest.mark.timeout(240)
@pytest.mark.covers(
"mgmt.user.update.budget_change_returns_promptly_with_wedged_coordination_redis",
"mgmt.user.bulk_update.budget_change_returns_promptly_with_wedged_coordination_redis",
"mgmt.customer.update.budget_change_returns_promptly_with_wedged_coordination_redis",
"mgmt.key.reset_spend.returns_promptly_with_wedged_coordination_redis",
"mgmt.auth_cache_invalidation.publish_parked_by_short_redis_wedge_lands_after_recovery",
"mgmt.auth_cache_invalidation.burst_with_worker_kill_keeps_serving_while_redis_wedged",
)
def test_user_budget_updates_return_promptly_while_coordination_redis_is_wedged(
gateway: Gateway, tmp_path: Path, record_property: Callable[[str, object], None]
) -> None:
original: Final = os.environ["DATABASE_URL"]
identity: Final = "integration_wedged_redis_" + uuid.uuid4().hex
parsed: Final = urlsplit(original)
database_url: Final = urlunsplit((parsed.scheme, parsed.netloc, "/" + identity, "", ""))
timings: dict[str, float] = {}
with psycopg.connect(original, autocommit=True) as admin:
admin.execute(sql.SQL("CREATE DATABASE {}").format(sql.Identifier(identity)))
try:
results_dir: Final = Path(os.environ.get("INTEGRATION_RESULTS_DIR", str(tmp_path)))
prior_logs: Final = frozenset(results_dir.glob("owned-proxy-*.log"))
with (
owned_redis(tmp_path) as coordination,
owned_proxy(
gateway,
tmp_path,
{
"DATABASE_URL": database_url,
"REDIS_HOST": coordination.host,
"REDIS_PORT": str(coordination.port),
},
config=Path("tests/integration/coordination_redis_proxy_config.yaml"),
workers=2,
) as candidate,
Redis(host=coordination.host, port=coordination.port, socket_timeout=1) as subscriber_client,
):
pubsub: Final = subscriber_client.pubsub()
pubsub.subscribe(_CHANNEL)
received: list[dict[str, JsonValue]] = []
def drained() -> tuple[dict[str, JsonValue], ...]:
received.extend(_received(pubsub))
return tuple(received)
eventually(
lambda: subscriber_client.pubsub_numsub(_CHANNEL)[0][1],
lambda count: count >= 3,
seconds=15,
)
users: Final = tuple(f"{identity}_u{index}" for index in range(_USERS))
for user_id in users:
candidate.post("/user/new", {"user_id": user_id, "auto_create_key": False, "max_budget": 10.0})
key: Final = string_value(
candidate.post("/key/generate", {"user_id": users[0], "max_budget": 5.0})["key"]
)
team_id: Final = string_value(
candidate.post("/team/new", {"team_alias": identity, "max_budget": 5.0})["team_id"]
)
customer_id: Final = identity + "_cust"
candidate.post("/customer/new", {"user_id": customer_id, "max_budget": 5.0})
drained()
timings["h1_healthy"] = _timed_post(
candidate, "/user/update", {"user_id": users[0], "max_budget": 11.0}
)
assert timings["h1_healthy"] < _HANDLER_BUDGET_SECONDS, (
f"healthy /user/update took {timings['h1_healthy']:.3f}s"
)
eventually(
drained,
lambda messages: any(message.get("cache_key") == users[0] for message in messages),
seconds=10,
)
timings["h2_healthy_control"] = _timed_post(
candidate, "/user/update", {"user_id": users[0], "tpm_limit": 1000}
)
assert timings["h2_healthy_control"] < _HANDLER_BUDGET_SECONDS, (
f"healthy control update took {timings['h2_healthy_control']:.3f}s"
)
coordination.signal(signal.SIGSTOP)
try:
timings["s2_wedged_control"] = _timed_post(
candidate, "/user/update", {"user_id": users[0], "tpm_limit": 1000}
)
assert timings["s2_wedged_control"] < _HANDLER_BUDGET_SECONDS, (
f"control update without a cache-relevant field took {timings['s2_wedged_control']:.3f}s"
)
timings["s1_user_update"] = _timed_post(
candidate, "/user/update", {"user_id": users[0], "max_budget": 98.0}
)
assert timings["s1_user_update"] < _HANDLER_BUDGET_SECONDS, (
f"/user/update with max_budget took {timings['s1_user_update']:.3f}s "
"with a wedged coordination Redis"
)
timings["s3_bulk_update"] = _timed_post(
candidate, "/user/bulk_update", {"all_users": True, "user_updates": {"max_budget": 79.0}}
)
assert timings["s3_bulk_update"] < _BULK_BUDGET_SECONDS, (
f"/user/bulk_update over {_USERS} users took {timings['s3_bulk_update']:.3f}s "
"with a wedged coordination Redis"
)
timings["s4_key_update"] = _timed_post(
candidate, "/key/update", {"key": key, "max_budget": 6.0}, timeout=60
)
assert timings["s4_key_update"] < 30, (
f"/key/update hung for {timings['s4_key_update']:.3f}s with a wedged coordination Redis"
)
timings["s5_team_update"] = _timed_post(
candidate, "/team/update", {"team_id": team_id, "max_budget": 6.0}, timeout=60
)
assert timings["s5_team_update"] < 30, (
f"/team/update hung for {timings['s5_team_update']:.3f}s with a wedged coordination Redis"
)
timings["s6_customer_update"] = _timed_post(
candidate, "/customer/update", {"user_id": customer_id, "max_budget": 6.0}
)
assert timings["s6_customer_update"] < _HANDLER_BUDGET_SECONDS, (
f"/customer/update took {timings['s6_customer_update']:.3f}s with a wedged coordination Redis"
)
timings["s7_reset_spend"] = _timed_post(candidate, f"/key/{key}/reset_spend", {"reset_to": 0})
assert timings["s7_reset_spend"] < _HANDLER_BUDGET_SECONDS, (
f"/key/<key>/reset_spend took {timings['s7_reset_spend']:.3f}s with a wedged coordination Redis"
)
missing_started: Final = time.monotonic()
missing: Final = candidate.request(
"POST", "/user/update", {"user_id": users[0], "max_budget": "not-a-number"}
)
timings["s8_invalid_body"] = time.monotonic() - missing_started
assert missing.status_code // 100 == 4, (
f"/user/update with an invalid body returned {missing.status_code} "
f"in {timings['s8_invalid_body']:.3f}s"
)
assert timings["s8_invalid_body"] < _HANDLER_BUDGET_SECONDS, (
f"/user/update with an invalid body took {timings['s8_invalid_body']:.3f}s"
)
def burst_request(path: str, body: Mapping[str, JsonValue]) -> tuple[object, float]:
started: Final = time.monotonic()
try:
response: Final = candidate.client.request(
"POST",
path,
json=body,
headers={"Authorization": f"Bearer {candidate.key}"},
timeout=60,
)
return response.status_code, time.monotonic() - started
except Exception as error: # noqa: BLE001 # the killed worker drops in-flight requests
return error, time.monotonic() - started
port: Final = candidate.client.base_url.port
assert port is not None, f"owned proxy client has no port: {candidate.client.base_url}"
with ThreadPoolExecutor(_BURST) as pool:
futures: Final = [
pool.submit(
burst_request,
*_burst_call(i, users, key, team_id, customer_id),
)
for i in range(_BURST)
]
os.kill(_worker_pid(port), signal.SIGKILL)
results: Final = [future.result() for future in futures]
responses: Final = [(status, elapsed) for status, elapsed in results if isinstance(status, int)]
failures: Final = [status for status, _elapsed in responses if status != 200]
assert not failures, f"burst responses that were not 200: {failures}"
transport_errors: Final = [status for status, _elapsed in results if not isinstance(status, int)]
assert len(transport_errors) <= 3, (
f"{len(transport_errors)} requests raised transport errors: {transport_errors!r}"
)
elapsed_sorted: Final = sorted(
elapsed for i, (status, elapsed) in enumerate(results) if i % 5 in (0, 1, 4) and status == 200
)
timings["c1_burst_p95"] = elapsed_sorted[int(len(elapsed_sorted) * 0.95) - 1]
assert timings["c1_burst_p95"] < _HANDLER_BUDGET_SECONDS, (
f"burst p95 {timings['c1_burst_p95']:.3f}s"
)
eventually(
lambda: candidate.request("GET", "/health/liveliness").status_code,
lambda status: status == 200,
seconds=15,
)
timings["c1_survivor"] = _timed_post(
candidate, "/user/update", {"user_id": users[0], "tpm_limit": 2000}
)
assert timings["c1_survivor"] < _HANDLER_BUDGET_SECONDS, (
f"control update on the surviving worker took {timings['c1_survivor']:.3f}s"
)
finally:
coordination.signal(signal.SIGCONT)
wedged_keys: Final = {users[i] for i in range(_BURST) if i % 5 == 0 and i != 0} | {f"team_id:{team_id}"}
def proxy_log() -> str:
return "".join(
path.read_text() for path in results_dir.glob("owned-proxy-*.log") if path not in prior_logs
)
team_wedged_key: Final = f"team_id:{team_id}"
eventually(
proxy_log,
lambda text: (
all(
f"publish for {wedged_key} failed" in text
for wedged_key in wedged_keys
if wedged_key != team_wedged_key
)
and (
f"publish for {team_wedged_key} failed" in text
or f"internal usage cache entry {team_wedged_key}" in text
)
),
seconds=45,
)
marker: Final = len(received)
drained()
recovered_keys: Final = {str(message.get("cache_key")) for message in received[marker:]}
assert recovered_keys.isdisjoint(wedged_keys), (
f"wedged publishes unexpectedly landed after recovery: {sorted(recovered_keys & wedged_keys)}"
)
coordination.signal(signal.SIGSTOP)
try:
timings["r1b_short_wedge_a"] = _timed_post(
candidate, "/user/update", {"user_id": users[4], "max_budget": 15.0}
)
assert timings["r1b_short_wedge_a"] < _HANDLER_BUDGET_SECONDS, (
f"/user/update inside a short wedge took {timings['r1b_short_wedge_a']:.3f}s"
)
timings["r1b_short_wedge_b"] = _timed_post(
candidate, "/user/update", {"user_id": users[5], "max_budget": 16.0}
)
assert timings["r1b_short_wedge_b"] < _HANDLER_BUDGET_SECONDS, (
f"/user/update inside a short wedge took {timings['r1b_short_wedge_b']:.3f}s"
)
finally:
coordination.signal(signal.SIGCONT)
eventually(
drained,
lambda messages: {str(message.get("cache_key")) for message in messages} >= {users[4], users[5]},
seconds=10,
)
timings["r2_resumed"] = _timed_post(
candidate, "/user/update", {"user_id": users[1], "max_budget": 12.0}
)
assert timings["r2_resumed"] < _HANDLER_BUDGET_SECONDS, (
f"post-recovery /user/update took {timings['r2_resumed']:.3f}s"
)
eventually(
drained,
lambda messages: any(message.get("cache_key") == users[1] for message in messages),
seconds=10,
)
coordination.stop()
timings["f1_refused"] = _timed_post(
candidate, "/user/update", {"user_id": users[2], "max_budget": 13.0}
)
assert timings["f1_refused"] < _HANDLER_BUDGET_SECONDS, (
f"/user/update with refused coordination Redis took {timings['f1_refused']:.3f}s"
)
coordination.start()
restarted_pubsub: Final = subscriber_client.pubsub()
restarted_pubsub.subscribe(_CHANNEL)
restarted_received: list[dict[str, JsonValue]] = []
def drained_after_restart() -> tuple[dict[str, JsonValue], ...]:
restarted_received.extend(_received(restarted_pubsub))
return tuple(restarted_received)
eventually(
lambda: subscriber_client.pubsub_numsub(_CHANNEL)[0][1],
lambda count: count >= 3,
seconds=30,
)
timings["f2_restarted"] = _timed_post(
candidate, "/user/update", {"user_id": users[3], "max_budget": 14.0}
)
assert timings["f2_restarted"] < _HANDLER_BUDGET_SECONDS, (
f"/user/update after Redis restart took {timings['f2_restarted']:.3f}s"
)
eventually(
drained_after_restart,
lambda messages: any(message.get("cache_key") == users[3] for message in messages),
seconds=10,
)
info_last: Final = object_value(candidate.get("/user/info", {"user_id": users[-1]})["user_info"])
assert info_last["max_budget"] == 79.0, info_last
info_user3: Final = object_value(candidate.get("/user/info", {"user_id": users[3]})["user_info"])
assert info_user3["max_budget"] == 14.0, info_user3
finally:
admin.execute(sql.SQL("DROP DATABASE {} WITH (FORCE)").format(sql.Identifier(identity)))
record_property("cell_elapsed_seconds", timings)

77
tests/test_litellm/proxy/common_utils/test_auth_cache_invalidation_pubsub.py
View file

@ -1,17 +1,20 @@
import asyncio
import hashlib
import json
from typing import Iterable, List, Optional, Tuple
import time
from collections.abc import Iterable
from unittest.mock import patch
import pytest
from redis.asyncio import Redis
import litellm.proxy.common_utils.auth_cache_invalidation_pubsub as pubsub_module
from litellm.caching.in_memory_cache import InMemoryCache
from litellm.proxy._types import UserAPIKeyAuth
from litellm.proxy.common_utils.auth_cache_invalidation_pubsub import (
AUTH_CACHE_INVALIDATION_CHANNEL,
AuthCacheInvalidationSubscriber,
evict_and_broadcast,
publish_auth_cache_invalidation,
)
from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
@ -19,13 +22,29 @@ from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache
class _RecordingRedisClient(Redis):
def __init__(self) -> None:
self.published: List[Tuple[str, str]] = []
self.published: list[tuple[str, str]] = []
async def publish(self, channel: str, message: str) -> int:
self.published.append((channel, message))
return 1
class _WedgedPublishRedisClient(Redis):
def __init__(self) -> None:
self.attempted: list[str] = []
self.in_flight = 0
self.max_in_flight = 0
self.release = asyncio.Event()
async def publish(self, channel: str, message: str) -> int:
self.in_flight += 1
self.max_in_flight = max(self.max_in_flight, self.in_flight)
self.attempted.append(message)
await self.release.wait()
self.in_flight -= 1
return 1
class _FailingPublishRedisClient(Redis):
def __init__(self) -> None:
pass
@ -36,16 +55,16 @@ class _FailingPublishRedisClient(Redis):
class _QueuePubSub:
def __init__(self, initial_messages: Iterable[object] = ()) -> None:
self.queue: "asyncio.Queue[object]" = asyncio.Queue()
self.queue: asyncio.Queue[object] = asyncio.Queue()
for message in initial_messages:
self.queue.put_nowait(message)
self.subscribed_channels: List[str] = []
self.subscribed_channels: list[str] = []
self.closed = False
async def subscribe(self, *channels: str) -> None:
self.subscribed_channels.extend(channels)
async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> Optional[object]:
async def get_message(self, *, ignore_subscribe_messages: bool, timeout: float) -> object | None:
try:
return await asyncio.wait_for(self.queue.get(), timeout)
except asyncio.TimeoutError:
@ -64,7 +83,7 @@ class _ScriptedPubSubRedisClient(Redis):
class _FakeRedisCache:
def __init__(self, client: object, namespace: Optional[str] = None) -> None:
def __init__(self, client: object, namespace: str | None = None) -> None:
self._client = client
self.namespace = namespace
@ -222,3 +241,49 @@ async def test_subscriber_ignores_malformed_messages() -> None:
subscriber._apply_message(None)
assert cache.in_memory_cache.get_cache("project_id:p-1") is not None
@pytest.mark.asyncio
async def test_evict_and_broadcast_evicts_locally_and_returns_while_redis_publish_never_answers() -> None:
cache = UserApiKeyCache()
cache.set_cache("user-wedged", UserAPIKeyAuth(user_id="user-wedged"), model_type=UserAPIKeyAuth)
client = _WedgedPublishRedisClient()
with patch(
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
return_value=_FakeRedisCache(client=client),
):
started = time.monotonic()
await evict_and_broadcast(cache_keys=("user-wedged",), user_api_key_cache=cache)
elapsed = time.monotonic() - started
assert elapsed < 0.1, f"handler waited {elapsed:.3f}s on a publish that never answers"
assert cache.get_cache("user-wedged", model_type=UserAPIKeyAuth) is None
assert client.attempted == [json.dumps({"cache_key": "user-wedged"})], "publish was not handed to redis"
client.release.set()
await asyncio.sleep(0)
@pytest.mark.asyncio
async def test_publish_holds_at_most_sixteen_redis_connections_while_redis_is_wedged(
monkeypatch: pytest.MonkeyPatch,
) -> None:
monkeypatch.setattr(pubsub_module, "_in_flight_publishes", asyncio.Semaphore(16))
monkeypatch.setattr(pubsub_module, "_pending_publishes", set())
client = _WedgedPublishRedisClient()
with patch(
"litellm.proxy.common_utils.auth_cache_invalidation_pubsub.coordination_redis_cache",
return_value=_FakeRedisCache(client=client),
):
for i in range(64):
await publish_auth_cache_invalidation(cache_key=f"user-{i}")
await asyncio.sleep(0)
await asyncio.sleep(0)
assert client.max_in_flight == 16, f"publish tasks held {client.max_in_flight} redis connections at once"
assert len(client.attempted) == 16, "waiters called publish before a semaphore slot freed"
client.release.set()
await asyncio.gather(*pubsub_module._pending_publishes) # pyright: ignore[reportPrivateUsage] # drain module-level tasks
assert len(client.attempted) == 64