Skip to content

feat(llm): add configurable per-model Redis GCRA rate limiting - #2353

Merged
bittergreen merged 2 commits into
MemTensor:dev-v2.0.34from
bittergreen:wq-dev-v2.0.34
Sep 9, 2026
Merged

feat(llm): add configurable per-model Redis GCRA rate limiting#2353
bittergreen merged 2 commits into
MemTensor:dev-v2.0.34from
bittergreen:wq-dev-v2.0.34

Conversation

@bittergreen

Copy link
Copy Markdown
Collaborator

Description

feat(llm): add configurable per-model Redis GCRA rate limiting

Type of change

  • New feature (non-breaking change which adds functionality)

How Has This Been Tested?

  • Unit Test

Checklist

  • I have performed a self-review of my own code | 我已自行检查了自己的代码
  • I have commented my code in hard-to-understand areas | 我已在难以理解的地方对代码进行了注释
  • I have added tests that prove my fix is effective or that my feature works | 我已添加测试以证明我的修复有效或功能正常
  • I have created related documentation issue/PR in MemOS-Docs (if applicable) | 我已在 MemOS-Docs 中创建了相关的文档 issue/PR(如果适用)
  • I have linked the issue to this PR (if applicable) | 我已将 issue 链接到此 PR(如果适用)
  • I have mentioned the person who will review this PR | 我已提及将审查此 PR 的人

Reviewer Checklist

  • closes #xxxx (Replace xxxx with the GitHub issue number)
  • Made sure Checks passed
  • Tests have been provided

@Memtensor-AI Memtensor-AI added area:core MOS 编排层 / 框架底座 / 跨模块问题 area:docs 文档、示例 area:model llm + embedder + reranker status:in-progress Someone or AI is working on it | 人工或 AI 正在处理 labels Sep 9, 2026
@Memtensor-AI

Memtensor-AI commented Sep 9, 2026

Copy link
Copy Markdown
Collaborator

🤖 Open Code Review

Target: PR #2353
Task: b05728eb9fc2e2d3
Base: dev-v2.0.34
Head: wq-dev-v2.0.34

🔍 OpenCodeReview found 10 issue(s) in this PR.


1. src/memos/llms/openai.py (L152-L156)

If rate_limit.create_completion(...) raises (e.g. LLMRateLimitError, ConfigurationError, or any other exception), response is never assigned. The finally block then attempts getattr(response, 'close', None), which triggers UnboundLocalError: local variable 'response' referenced before assignment, masking the original exception entirely.

Move the assignment inside the try block and guard the finally against the unbound case:

response = None
try:
    response = rate_limit.create_completion(self.client, request_body, self.config.rate_limit)
    for chunk in response:
        ...
finally:
    if response is not None:
        close = getattr(response, "close", None)
        if callable(close):
            close()

Alternatively, initialize response = None before the try and keep the existing finally guard as if response is not None.


2. src/memos/configs/llm_rate_limit.py (L88-L97)

Typos in MEMSCHEDULER_REDIS_* variable names (e.g. MEMSCHEDULER_REDIS_PASWORD) are silently ignored — the misspelled variable is simply not read, so the default (None) is used and the Redis connection runs unauthenticated or with wrong settings. The MEMOS_LLM_RATE_LIMIT_* namespace gets a strict unknown-variable check, but the MEMSCHEDULER_REDIS_* namespace has no equivalent guard. Consider adding the same validation for the Redis namespace so operators get an actionable error instead of a silent misconfiguration.

💡 Suggested Change

Before:

        unsupported = sorted(
            name
            for name in os.environ
            if name.startswith("MEMOS_LLM_RATE_LIMIT_") and name not in supported
        )
        if unsupported:
            raise ConfigurationError(
                "Only MEMOS_LLM_RATE_LIMIT_ENABLED and MEMOS_LLM_RATE_LIMIT_RULES are supported; "
                "remove: " + ", ".join(unsupported)
            )

After:

        supported_redis = {
            f"MEMSCHEDULER_REDIS_{n.upper()}"
            for n in ("host", "port", "db", "username", "password", "ssl", "socket_timeout")
        }
        unsupported_redis = sorted(
            name
            for name in os.environ
            if name.startswith("MEMSCHEDULER_REDIS_") and name not in supported_redis
        )
        if unsupported_redis:
            raise ConfigurationError(
                "Unrecognised MEMSCHEDULER_REDIS_* variable(s); remove: "
                + ", ".join(unsupported_redis)
            )
        unsupported = sorted(
            name
            for name in os.environ
            if name.startswith("MEMOS_LLM_RATE_LIMIT_") and name not in supported
        )
        if unsupported:
            raise ConfigurationError(
                "Only MEMOS_LLM_RATE_LIMIT_ENABLED and MEMOS_LLM_RATE_LIMIT_RULES are supported; "
                "remove: " + ", ".join(unsupported)
            )

3. src/memos/configs/llm_rate_limit.py (L120-L123)

The original ValueError from pydantic's TypeAdapter is suppressed with from None, which cuts the traceback chain. When an operator passes an invalid value (e.g. a non-numeric string for redis_port), they only see the field name — not which validation rule was violated or what the bad value was. Use from err to preserve the cause and make misconfiguration much easier to diagnose.

💡 Suggested Change

Before:

            except ValueError:
                raise ConfigurationError(
                    f"Invalid LLM rate limit environment setting: {name}"
                ) from None

After:

            except ValueError as err:
                raise ConfigurationError(
                    f"Invalid LLM rate limit environment setting: {name}"
                ) from err

4. src/memos/llms/rate_limit.py (L118-L120)

When next_check is in the past at the point of re-entry (i.e., next_check - now < 0), min(deadline - now, next_check - now) evaluates to a negative value. Python's Condition.wait(timeout) with a non-positive timeout returns immediately without sleeping, creating a tight busy-wait spin loop that hammers Redis with Lua script executions on every iteration until the deadline expires.

Fix: clamp next_check - now to at least a small positive floor (e.g., 1 ms) so a stale next_check still causes at least a minimal sleep:

delay = min(deadline - now, max(0.001, next_check - now)) if head else deadline - now
💡 Suggested Change

Before:

                    delay = min(deadline - now, next_check - now) if head else deadline - now
                        self._condition.wait(delay)
                        continue

After:

                    delay = min(deadline - now, max(0.001, next_check - now)) if head else deadline - now
                        self._condition.wait(delay)
                        continue

5. src/memos/llms/rate_limit.py (L174-L177)

After a fork, the child process inherits all of the parent's open file descriptors, including the TCP sockets held by each RedisGCRALimiter._client connection pool. Discarding _registry prevents new calls from routing to stale limiters, but does not close the inherited socket FDs. Both parent and child now share the same underlying socket to Redis; a response intended for the parent can be read by the child (and vice-versa), silently corrupting the GCRA state or causing unexpected ResponseErrors.

Fix: close each client's connection pool before replacing the registry:

def _after_fork() -> None:
    global _registry, _registry_lock
    old_registry = _registry
    _registry = {}
    _registry_lock = threading.Lock()
    for limiter in old_registry.values():
        try:
            limiter._client.close()
        except Exception:
            pass
💡 Suggested Change

Before:

def _after_fork() -> None:
    global _registry, _registry_lock
    _registry = {}
    _registry_lock = threading.Lock()

After:

def _after_fork() -> None:
    global _registry, _registry_lock
    old_registry = _registry
    _registry = {}
    _registry_lock = threading.Lock()
    for limiter in old_registry.values():
        try:
            limiter._client.close()
        except Exception:
            pass

6. src/memos/llms/rate_limit.py (L133-L134)

When failure_mode='open' causes an early return, the finally block correctly removes the waiter and calls notify_all(). This wakes every other thread sleeping in self._condition.wait(...). Each woken thread finds itself at the (new) queue head, immediately calls Redis — which is still down — hits another RedisError, and with failure_mode='open' returns as well, triggering another notify_all(). This creates an O(queue_capacity) cascade of rapid Redis calls fanning out in waves during an outage.

Consider adding a short sleep before returning on the open failure path to provide a natural backoff, or using a shared "circuit open" flag to suppress subsequent Redis calls while the outage is in progress.


7. tests/configs/test_llm_rate_limit.py (L50-L51)

If docker/.env.example-full is absent (e.g., a CI checkout that excludes the docker/ directory), dotenv_values silently returns an empty dict. The subsequent dict key accesses — values["MEMOS_LLM_RATE_LIMIT_ENABLED"], values["MEMOS_LLM_RATE_LIMIT_RULES"], and values["MOS_CHAT_MODEL"] — will all raise a KeyError with no indication that the file is missing, making the failure hard to diagnose.

Add an existence check before reading the file:

assert example.exists(), f"Required fixture file not found: {example}"
💡 Suggested Change

Before:

    example = Path(__file__).resolve().parents[2] / "docker" / ".env.example-full"
    values = dotenv_values(example, interpolate=False)

After:

    example = Path(__file__).resolve().parents[2] / "docker" / ".env.example-full"
    assert example.exists(), f"Required fixture file not found: {example}"
    values = dotenv_values(example, interpolate=False)

8. tests/llms/test_qps_rate_limit.py (L267-L269)

from types import SimpleNamespace is already imported at module level (line 9). This local re-import is redundant and can be removed.

💡 Suggested Change

Before:

    from types import SimpleNamespace

    llm, _ = make_llm(monkeypatch, lambda _: httpx.Response(200, json=response_body()))

After:

    llm, _ = make_llm(monkeypatch, lambda _: httpx.Response(200, json=response_body()))

9. tests/llms/test_qps_rate_limit.py (L207)

The expected count 2 is hardcoded rather than derived from the configured retry_attempts value. retry_attempts defaults to 1 (one retry after the first attempt), which gives 2 total wire calls — but if that default changes, this assertion will silently validate the wrong contract. Tie the expectation to the actual config to make the contract explicit and resilient to config changes.

💡 Suggested Change

Before:

    assert len(wire_calls) == limiter.acquire.call_count == 2

After:

    expected_calls = llm.config.rate_limit.rule_for("gpt-4o-mini").retry_attempts + 1
    assert len(wire_calls) == limiter.acquire.call_count == expected_calls

10. tests/llms/test_qps_rate_limit.py (L118-L121)

assert release.wait(2) runs inside a ThreadPoolExecutor worker. If the wait times out, AssertionError is raised in the worker thread, swallowed by the executor, and only re-surfaced when first.result() is called in the finally block — after release.set() has already been called, making the timeout condition impossible to reproduce and the failure difficult to diagnose. Raise an explicit exception instead so the cause is clear.

💡 Suggested Change

Before:

    def script(**_):
        entered.set()
        assert release.wait(2)
        return [1, 0]

After:

    def script(**_):
        entered.set()
        if not release.wait(2):
            raise RuntimeError("Timed out waiting for release in worker thread")
        return [1, 0]

Generated by cloud-assistant via Open Code Review.

@Memtensor-AI

Copy link
Copy Markdown
Collaborator

⚠️ Automated Test Results: ENV ISSUE

The test environment encountered an issue that requires manual attention.

Details: Executor error: Command failed: git fetch --deepen 200 base dev-v2.0.34
git@github.com: Permission denied (publickey).
fatal: Could not read from remote repository.

Please make sure you have the correct access rights
and the repository exists.
Branch: wq-dev-v2.0.34

@bittergreen bittergreen assigned bittergreen and unassigned wustzdy Sep 9, 2026
@bittergreen
bittergreen requested review from wustzdy and removed request for WeiminLee, endxxxx and shinetata September 9, 2026 12:45
@Memtensor-AI

Copy link
Copy Markdown
Collaborator

✅ Automated Test Results: PASSED

All tests passed (72/72 executed, 5 skipped). memos_github_open_source/smoke: 1/1, memos_python_core/changed-repo-python: 71 passed, 5 skipped. Duration: 8s [advisory, non-gating] AI-generated tests on branch test/auto-gen-b05728eb9fc2e2d3-20260909203616: 156/161 passed, 5 failed — these do NOT affect the PR verdict; review the branch manually.

Branch: wq-dev-v2.0.34

@Memtensor-AI Memtensor-AI added status:ready Ready for implementation; waiting for assignee or AI dispatch | 可进入实现,等待认领或派发 and removed status:in-progress Someone or AI is working on it | 人工或 AI 正在处理 labels Sep 9, 2026
@bittergreen
bittergreen merged commit 0f74774 into MemTensor:dev-v2.0.34 Sep 9, 2026
18 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:core MOS 编排层 / 框架底座 / 跨模块问题 area:docs 文档、示例 area:model llm + embedder + reranker status:ready Ready for implementation; waiting for assignee or AI dispatch | 可进入实现,等待认领或派发

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants