mirror of
https://github.com/BerriAI/litellm.git
synced 2026-10-06 02:48:13 +00:00
fix(proxy): reject dynamic limiter requests on usage lookup failure
This commit is contained in:
parent
90e4962c81
commit
ace963b7a8
2 changed files with 71 additions and 1 deletions
|
|
@ -8,6 +8,8 @@ from collections.abc import Callable
|
|||
from datetime import datetime
|
||||
from typing import Final
|
||||
|
||||
from fastapi import HTTPException
|
||||
|
||||
import litellm
|
||||
from litellm import ModelResponse, Router
|
||||
from litellm._logging import verbose_proxy_logger
|
||||
|
|
@ -177,7 +179,10 @@ class _PROXY_DynamicRateLimitHandler(CustomLogger):
|
|||
verbose_proxy_logger.exception(
|
||||
"litellm.proxy.hooks.dynamic_rate_limiter.py::check_available_usage: Exception occurred - %s", e
|
||||
)
|
||||
return None, None, None, None, None
|
||||
raise HTTPException(
|
||||
status_code=503,
|
||||
detail="Dynamic rate limiter usage is unavailable; request rejected.",
|
||||
) from e
|
||||
|
||||
async def async_pre_call_hook(
|
||||
self,
|
||||
|
|
|
|||
|
|
@ -1,8 +1,13 @@
|
|||
import asyncio
|
||||
from datetime import datetime, timezone
|
||||
from unittest.mock import AsyncMock, MagicMock, patch
|
||||
|
||||
import pytest
|
||||
from fastapi import HTTPException
|
||||
|
||||
from litellm.caching.caching import DualCache
|
||||
from litellm.proxy._types import UserAPIKeyAuth
|
||||
from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError
|
||||
from litellm.proxy.hooks.dynamic_rate_limiter import (
|
||||
DynamicRateLimiterCache,
|
||||
_PROXY_DynamicRateLimitHandler,
|
||||
|
|
@ -42,3 +47,63 @@ async def test_handler_threads_time_fn_to_internal_cache():
|
|||
)
|
||||
await handler.internal_usage_cache.async_set_cache_sadd(model="my-fake-model", value=["p1", "p2"])
|
||||
assert await handler.internal_usage_cache.async_get_cache(model="my-fake-model") == 2
|
||||
|
||||
|
||||
def fake_handler(*, cache_error=None, tpm=100, rpm=10):
|
||||
usage_cache = MagicMock()
|
||||
usage_cache.async_get_cache = AsyncMock(side_effect=cache_error, return_value=None)
|
||||
handler = _PROXY_DynamicRateLimitHandler(internal_usage_cache=usage_cache)
|
||||
router = MagicMock()
|
||||
router.get_model_group_info.return_value = MagicMock(tpm=tpm, rpm=rpm)
|
||||
router.get_model_group_usage = AsyncMock(return_value=(0, 0))
|
||||
handler.update_variables(router)
|
||||
return handler, usage_cache, router
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_cache_lookup_failure_rejects_before_router_call():
|
||||
handler, _, router = fake_handler(cache_error=ConnectionError("cache offline"))
|
||||
with pytest.raises(HTTPException) as exc:
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", metadata={}),
|
||||
cache=DualCache(),
|
||||
data={"model": "test-model"},
|
||||
call_type="completion",
|
||||
)
|
||||
assert exc.value.status_code == 503
|
||||
assert "cache offline" not in str(exc.value.detail)
|
||||
router.get_model_group_usage.assert_not_awaited()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_zero_quota_still_raises_rate_limit():
|
||||
handler, _, router = fake_handler(tpm=0)
|
||||
with pytest.raises(ProxyRateLimitError) as exc:
|
||||
await handler.async_pre_call_hook(
|
||||
user_api_key_dict=UserAPIKeyAuth(api_key="sk-test", metadata={}),
|
||||
cache=DualCache(),
|
||||
data={"model": "gpt-4o-mini"},
|
||||
call_type="completion",
|
||||
)
|
||||
assert exc.value.status_code == 429
|
||||
router.get_model_group_usage.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_unlimited_capacity_still_admits_request():
|
||||
handler, _, router = fake_handler(tpm=None, rpm=None)
|
||||
user = UserAPIKeyAuth(api_key="sk-test", metadata={})
|
||||
with patch("litellm.proxy.hooks.dynamic_rate_limiter.asyncio.create_task") as create_task:
|
||||
assert await handler.async_pre_call_hook(user, DualCache(), {"model": "test-model"}, "completion") is None
|
||||
create_task.assert_not_called()
|
||||
router.get_model_group_usage.assert_awaited_once()
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_successful_lookup_keeps_active_project_accounting():
|
||||
handler, _, _ = fake_handler()
|
||||
user = UserAPIKeyAuth(api_key="sk-test", token="token-123", metadata={})
|
||||
handler.internal_usage_cache.async_set_cache_sadd = AsyncMock() # type: ignore[method-assign]
|
||||
await handler.async_pre_call_hook(user, DualCache(), {"model": "test-model"}, "completion")
|
||||
await asyncio.sleep(0)
|
||||
handler.internal_usage_cache.async_set_cache_sadd.assert_awaited_once_with(model="test-model", value=[user.token])
|
||||
|
|
|
|||
Loading…
Add table
Reference in a new issue