Skip to content

fix(lmi): bound Redis I/O in GlobalRateLimiter [CLD-320] #464

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ypicard merged 1 commit into from
Jul 1, 2026
Merged

fix(lmi): bound Redis I/O in GlobalRateLimiter [CLD-320] #464

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 11 additions & 1 deletion packages/lmi/src/lmi/rate_limiter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = "<machine_id>"
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
14 changes: 12 additions & 2 deletions packages/lmi/tests/test_rate_limiter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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 = [
Expand Down Expand Up @@ -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:
Expand Down
Loading