diff --git a/packages/lmi/src/lmi/rate_limiter.py b/packages/lmi/src/lmi/rate_limiter.py index e7fadb53..974850f5 100644 --- a/packages/lmi/src/lmi/rate_limiter.py +++ b/packages/lmi/src/lmi/rate_limiter.py @@ -30,6 +30,8 @@ GLOBAL_RATE_LIMITER_TIMEOUT = float(os.environ.get("RATE_LIMITER_TIMEOUT", "60")) +RATE_LIMITER_REDIS_OP_TIMEOUT = 5.0 + MATCH_ALL = None MatchAllInputs: TypeAlias = Literal[None] # noqa: PYI061 MATCH_MACHINE_ID = "" @@ -170,7 +172,11 @@ def storage(self) -> RedisStorage | MemoryStorage: logger.info("Using in-memory rate limiter.") else: conn = f"{self._redis_scheme}://{self._redis_bare_url}" - self._storage = RedisStorage(conn) + self._storage = RedisStorage( + conn, + stream_timeout=RATE_LIMITER_REDIS_OP_TIMEOUT, + connect_timeout=RATE_LIMITER_REDIS_OP_TIMEOUT, + ) logger.info(f"Connected to redis instance for rate limiting: {conn}") return self._storage @@ -391,6 +397,10 @@ async def try_acquire( TimeoutError: if the acquire_timeout is exceeded. ValueError: if the weight exceeds the rate limit. Only raised if raise_impossible_limits was specified. + + Note: + If Redis is unreachable, the op raises a coredis RedisError (bounded by + RATE_LIMITER_REDIS_OP_TIMEOUT) that propagates instead of hanging. """ namespace, primary_key = await self.parse_namespace_and_primary_key( namespace_and_key, machine_id=machine_id diff --git a/packages/lmi/tests/test_rate_limiter.py b/packages/lmi/tests/test_rate_limiter.py index 32176fdd..be7a4948 100644 --- a/packages/lmi/tests/test_rate_limiter.py +++ b/packages/lmi/tests/test_rate_limiter.py @@ -14,7 +14,12 @@ from lmi.constants import CHARACTERS_PER_TOKEN_ASSUMPTION from lmi.embeddings import LiteLLMEmbeddingModel from lmi.llms import CommonLLMNames, LiteLLMModel -from lmi.rate_limiter import CROSSREF_BASE_URL, FALLBACK_RATE_LIMIT, GlobalRateLimiter +from lmi.rate_limiter import ( + CROSSREF_BASE_URL, + FALLBACK_RATE_LIMIT, + RATE_LIMITER_REDIS_OP_TIMEOUT, + GlobalRateLimiter, +) from lmi.types import LLMResult LLM_CONFIG_W_RATE_LIMITS = [ @@ -467,7 +472,12 @@ def test_storage_uses_expected_redis_scheme( with patch("lmi.rate_limiter.RedisStorage") as mock_redis_storage: limiter = GlobalRateLimiter(redis_url=redis_url) _ = limiter.storage - mock_redis_storage.assert_called_once_with(expected_storage_url) + # timeouts must reach RedisStorage so a stalled connection can't hang + mock_redis_storage.assert_called_once_with( + expected_storage_url, + stream_timeout=RATE_LIMITER_REDIS_OP_TIMEOUT, + connect_timeout=RATE_LIMITER_REDIS_OP_TIMEOUT, + ) @pytest.mark.asyncio async def test_parsing_namespace(self) -> None: