Merge ace963b7a8 into f285229b51

litellm/proxy/hooks/dynamic_rate_limiter.py

@@ -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,

tests/test_litellm/proxy/hooks/test_dynamic_rate_limiter.py

@@ -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])