From eae3e7ea4700917c5299a16f4542b989a5a5d7b6 Mon Sep 17 00:00:00 2001 From: bittergreen Date: Wed, 9 Sep 2026 19:32:11 +0800 Subject: [PATCH 1/2] feat(llm): add configurable per-model Redis GCRA rate limiting --- docker/.env.example-full | 14 +- .../help/llm_qps_rate_limit.md | 83 ++++ .../help/llm_qps_rate_limit.md | 83 ++++ src/memos/configs/llm.py | 12 + src/memos/configs/llm_rate_limit.py | 124 +++++ src/memos/exceptions.py | 13 + src/memos/llms/openai.py | 58 ++- src/memos/llms/rate_limit.py | 275 +++++++++++ tests/configs/test_llm.py | 2 + tests/configs/test_llm_rate_limit.py | 176 +++++++ tests/llms/test_qps_rate_limit.py | 436 ++++++++++++++++++ tests/llms/test_qps_rate_limit_redis.py | 184 ++++++++ 12 files changed, 1434 insertions(+), 26 deletions(-) create mode 100644 docs/cn/open_source/open_source_api/help/llm_qps_rate_limit.md create mode 100644 docs/en/open_source/open_source_api/help/llm_qps_rate_limit.md create mode 100644 src/memos/configs/llm_rate_limit.py create mode 100644 src/memos/llms/rate_limit.py create mode 100644 tests/configs/test_llm_rate_limit.py create mode 100644 tests/llms/test_qps_rate_limit.py create mode 100644 tests/llms/test_qps_rate_limit_redis.py diff --git a/docker/.env.example-full b/docker/.env.example-full index d3d417181..78b29b614 100644 --- a/docker/.env.example-full +++ b/docker/.env.example-full @@ -53,6 +53,16 @@ DOCUMENT_PARSER_MODEL= # falls back to MEMREADER_GENERAL_MOD IMAGE_PARSER_MODEL= # falls back to MEMREADER_GENERAL_MODEL when omitted QWEN_MODEL=qwen-flash # optional qwen_llm slot when QWEN_API_KEY is set +## Optional per-model LLM QPS rate limiting (Redis GCRA) +# Disabled by default. Enable explicitly and tune rules to your provider quota. +MEMOS_LLM_RATE_LIMIT_ENABLED=false +MEMOS_LLM_RATE_LIMIT_RULES='{"gpt-4o-mini":{"qps":5,"burst":2,"max_wait_seconds":30,"queue_capacity":16,"retry_attempts":1}}' +# Rule keys must match actual model names; unlisted models are not limited. +# QPS/burst are shared across workers using the same Redis/DB and model key. +# queue_capacity is per process/model; max_wait_seconds is the permit-wait budget. +# Reuses MEMSCHEDULER_REDIS_HOST/PORT/DB/USERNAME/PASSWORD/SSL; host is required when enabled. +# See docs/cn/open_source/open_source_api/help/llm_qps_rate_limit.md. + ## Embedding & rerank # embedding dim EMBEDDING_DIMENSION=1024 @@ -111,9 +121,9 @@ ENABLE_INTERNET=false # Internet search backend (bocha | tavily) INTERNET_SEARCH_BACKEND=bocha # API key for BOCHA Search -BOCHA_API_KEY= # required if ENABLE_INTERNET=true and backend=bocha +BOCHA_API_KEY= your-bocha-api-key and backend=bocha # API key for Tavily Search -TAVILY_API_KEY= # required if ENABLE_INTERNET=true and backend=tavily +TAVILY_API_KEY= your-bocha-api-key and backend=tavily # default search mode SEARCH_MODE=fast # fast | fine | mixture # Slow retrieval strategy configuration, rewrite is the rewrite strategy diff --git a/docs/cn/open_source/open_source_api/help/llm_qps_rate_limit.md b/docs/cn/open_source/open_source_api/help/llm_qps_rate_limit.md new file mode 100644 index 000000000..d81019ccc --- /dev/null +++ b/docs/cn/open_source/open_source_api/help/llm_qps_rate_limit.md @@ -0,0 +1,83 @@ +# LLM GCRA 限流 + +## 环境变量 + +限流只暴露两个环境变量，默认关闭。当前仅限制明确选中的模型。 + +```dotenv +MEMOS_LLM_RATE_LIMIT_ENABLED=false +MEMOS_LLM_RATE_LIMIT_RULES='{"gpt-4o-mini":{"qps":5,"burst":2,"max_wait_seconds":30,"queue_capacity":16,"retry_attempts":1}}' +``` + +`RULES` 是 JSON 对象，键为实际请求的模型名。每条规则只支持以下五个参数，省略时使用代码默认值： + +| 参数 | 默认值 | 含义 | +|---|---:|---| +| qps | 5 | 全局持续放行速率；所有共享配额的 worker 合计 | +| burst | 2 | 空闲后最多立即放行的请求总数，不是 qps 加 burst | +| max_wait_seconds | 30 | 单次主/备用调用及其受控重试的累计许可等待预算，单位秒 | +| queue_capacity | 16 | 每进程、每模型的等待队列上限，包含正在申请的队首 | +| retry_attempts | 1 | 首次模型请求失败后最多重试次数；0 表示不重试 | + +上述示例与代码默认值及 `docker/.env.example-full` 一致，默认不启用。需要限流时显式设置 `MEMOS_LLM_RATE_LIMIT_ENABLED=true`。参数是开源部署的保守起点，不代表供应商保证的配额；应按实际配额、共享环境、worker 数、排队等待及限流错误调整。QPS 限流不等于 token 吞吐或模型在途并发限制。 + +- 未配置 `RULES` 时，默认只选择 `gpt-4o-mini`，使用上述默认值。 +- 显式配置 `RULES` 会替换整个模型规则集合；未列出的模型不受限流影响，不隐式追加默认模型。 +- `RULES={}` 不限制任何模型；删除某个模型的条目即可取消该模型限流。 +- `ENABLED=false` 关闭整个功能。 +- 模型名不支持通配符；qps 必须为有限正数，burst 和容量为正整数，重试次数为非负整数。 +- Shell/`.env` 示例的外层单引号用于保护 JSON；在部署平台直接填写环境变量值时，不包含外层单引号。 + +多个模型分别配置示例： + +```dotenv +MEMOS_LLM_RATE_LIMIT_RULES='{"gpt-4o-mini":{"qps":5,"burst":2,"max_wait_seconds":30,"queue_capacity":16,"retry_attempts":1},"qwen-flash":{"qps":10,"burst":2,"max_wait_seconds":3,"queue_capacity":8,"retry_attempts":0}}' +``` + +第二个模型仅为示例，不默认启用。 + +## Redis 与加载 + +Redis 连接复用现有 `MEMSCHEDULER_REDIS_HOST/PORT/DB/USERNAME/PASSWORD/SSL`。缺少 host 时，第一次受控调用报配置错误；默认关闭时不创建 Redis 客户端。密码通过部署 Secret 注入。 + +Redis key 自动按模型生成，无需 scope： + +```text +memos:llm:gcra:gpt-4o-mini +memos:llm:gcra:qwen-flash +``` + +同一 Redis/DB、相同前缀下，同名模型跨 worker/环境共享一个 TAT，不因 endpoint 或 API Key 不同而拆分。不同模型的 TAT 和本地队列独立；模型别名视为不同模型。各环境必须使用一致的速率和 burst。 + +配置在创建 LLM 配置对象时加载，不逐请求读环境变量，也不自行加载 `.env`。更新部署配置后需协调重启 worker。Python 显式 `rate_limit` 配置仍可覆盖环境值，配置对象保留内部运行参数用于程序化构造和测试；这些参数不再提供环境变量入口。 + +旧的 `MEMOS_LLM_RATE_LIMIT_*` 配置中，除 `ENABLED`、`RULES` 外均需移除，例如 MODELS、QPS、BURST、SCOPE、CONFIG_FILE、REDIS_* 和 WAIT_JITTER_SECONDS。加载时会对不支持的变量报错，避免旧配置被静默忽略。旧的模型规则中也应移除 enabled、scope、抖动及退避参数。不再支持通过环境变量指定独立 JSON 配置文件。 + +若从旧哈希 key 或旧前缀升级，请协调所有实例切换，避免新旧 key 同时放行；新 key 初始化时会恢复一个 burst。不要在运行中随意切换前缀。 + +## 内部行为 + +- Lua 使用 Redis TIME，原子读取、判断和更新 TAT，拒绝不推进 TAT；Python 使用 register_script，无需本机安装 Lua。 +- 每进程、每模型只有队首申请 Redis，其他线程通过 Condition 等待。获准后立即离队发起模型调用，不等模型返回。 +- Redis 建议等待时间后附加 0～10ms 抖动。无有效 Retry-After 时使用指数退避和抖动，退避基数 1s、上限 8s。这些是内部默认值，不需要部署配置。 +- 受控调用关闭 SDK 隐藏重试。连接/超时错误及 HTTP 408、409、429、5xx 可有限重试，每次重新申请许可；Retry-After 超过内部等待上限时不提前重试。 +- 流式请求只重试建立阶段，流开始后的错误不重放。调用方提前结束时应关闭生成器。 +- 队列满、等待超时、Redis 不可用分别抛出 LLMRateLimitQueueFullError、LLMRateLimitTimeoutError、LLMRateLimitUnavailableError，不通过备用模型绕过。 +- 默认 Redis 连接和读取超时 0.5s，故障策略为 closed，即停止受控调用。网络响应迟到时不发送模型请求，也不退还已消耗或状态不确定的许可。 +- max_wait_seconds 不包含模型网络耗时和失败退避，不是整个业务请求的总超时；外层仍需业务 deadline。同步 Redis I/O 最迟要等 socket 超时才能退出。 +- 本地队列不是持久任务队列；满队列、超时和进程退出不会自动延期任务。Redis 故障切换或淘汰 TAT 也可能重置额度。 + +## 范围与验证 + +当前接入 OpenAILLM 及其 Qwen、DeepSeek、MiniMax 子类的 Chat Completions，包括普通调用、流式建立和备用模型。Azure、Responses API、Ollama、VLLM 等独立实现暂未接入。 + +该版本仅控制 QPS，不控制 Token 用量、Token 增速或在途并发，不能保证解决供应商所有 429。 + +INFO 的 `[LLM_RATE_LIMIT] sending` 记录模型、尝试序号及许可等待时间；WARNING 记录重试、队列满、等待超时和 Redis 故障。新增日志不记录请求正文或凭据。 + +```sh +poetry run pytest tests/configs/ tests/llms/ -q +MEMOS_TEST_LOCAL_REDIS=1 poetry run pytest tests/llms/test_qps_rate_limit_redis.py -q +``` + +第二条启动隔离本地 Redis，仅 Unix socket、无 TCP、无持久化，不读取生产 Redis 配置。日志位于 pytest 管理的 `redis-gcra*` 临时目录；短路径临时 socket 退出时清理。 diff --git a/docs/en/open_source/open_source_api/help/llm_qps_rate_limit.md b/docs/en/open_source/open_source_api/help/llm_qps_rate_limit.md new file mode 100644 index 000000000..febabca9a --- /dev/null +++ b/docs/en/open_source/open_source_api/help/llm_qps_rate_limit.md @@ -0,0 +1,83 @@ +# LLM GCRA Rate Limiting + +## Environment Variables + +Rate limiting exposes only two environment variables and is disabled by default. Only explicitly selected models are limited. + +```dotenv +MEMOS_LLM_RATE_LIMIT_ENABLED=false +MEMOS_LLM_RATE_LIMIT_RULES='{"gpt-4o-mini":{"qps":5,"burst":2,"max_wait_seconds":30,"queue_capacity":16,"retry_attempts":1}}' +``` + +`RULES` is a JSON object keyed by the actual model name used in requests. Each rule accepts only the following five parameters. Omitted parameters use the code defaults: + +| Parameter | Default | Description | +|---|---:|---| +| qps | 5 | Global sustained admission rate, aggregated across all workers sharing the quota | +| burst | 2 | Total number of requests that can be admitted immediately after an idle period; not qps plus burst | +| max_wait_seconds | 30 | Cumulative permit-wait budget in seconds for a single primary/backup invocation and its managed retries | +| queue_capacity | 16 | Waiting queue capacity per process and model, including the head currently requesting a permit | +| retry_attempts | 1 | Maximum retries after the initial model request fails; 0 disables retries | + +The example matches the code defaults and `docker/.env.example-full`, with rate limiting disabled. Set `MEMOS_LLM_RATE_LIMIT_ENABLED=true` explicitly to enable it. These values are a conservative starting point for open-source deployments, not provider-guaranteed quotas. Adjust them based on your actual quota, shared environments, worker count, queue waits, and rate-limit errors. QPS limiting is not a token-throughput or in-flight concurrency limit. + +- When `RULES` is not set, only `gpt-4o-mini` is selected, using the defaults above. +- Explicit `RULES` replace the entire model rule set. Unlisted models are unaffected; the default model is not implicitly added. +- `RULES={}` limits no models. Remove a model entry to disable limiting for that model. +- `ENABLED=false` disables the entire feature. +- Model names do not support wildcards. qps must be finite and positive; burst and queue capacity must be positive integers; retry attempts must be a nonnegative integer. +- The outer single quotes in shell/`.env` examples protect the JSON. Omit them when entering the environment variable value directly in a deployment platform. + +Example with separate rules for multiple models: + +```dotenv +MEMOS_LLM_RATE_LIMIT_RULES='{"gpt-4o-mini":{"qps":5,"burst":2,"max_wait_seconds":30,"queue_capacity":16,"retry_attempts":1},"qwen-flash":{"qps":10,"burst":2,"max_wait_seconds":3,"queue_capacity":8,"retry_attempts":0}}' +``` + +The second model is illustrative and is not selected by default. + +## Redis and Configuration Loading + +The limiter reuses the existing `MEMSCHEDULER_REDIS_HOST/PORT/DB/USERNAME/PASSWORD/SSL` connection settings. A missing host causes a configuration error on the first limited invocation. No Redis client is created while the feature is disabled. Inject passwords through deployment secrets. + +Redis keys are generated automatically per model; no scope configuration is needed: + +```text +memos:llm:gcra:gpt-4o-mini +memos:llm:gcra:qwen-flash +``` + +Workers and environments using the same Redis instance, database, prefix, and model name share one theoretical arrival time (TAT). Different endpoints or API keys do not create separate quotas. Different models have independent TAT values and local queues; model aliases are treated as distinct models. All environments sharing a quota must use consistent qps and burst settings. + +Configuration is loaded when the LLM configuration object is created, not on every request. The limiter does not load `.env` itself. Coordinate worker restarts after changing deployment settings. Explicit Python `rate_limit` configuration can still override environment values. Configuration objects retain internal runtime parameters for programmatic construction and testing, but those parameters have no environment-variable interface. + +Remove all legacy `MEMOS_LLM_RATE_LIMIT_*` variables other than `ENABLED` and `RULES`, including MODELS, QPS, BURST, SCOPE, CONFIG_FILE, REDIS_*, and WAIT_JITTER_SECONDS. Unsupported variables cause a configuration-loading error rather than being silently ignored. Also remove enabled, scope, jitter, and backoff parameters from old model rules. Selecting a separate JSON configuration file through an environment variable is no longer supported. + +When migrating from hashed keys or an older prefix, coordinate the switch across all instances to avoid simultaneous admission through both old and new keys. A new key starts with a full burst allowance. Do not change prefixes arbitrarily while the system is running. + +## Internal Behavior + +- Lua uses Redis TIME to atomically read, check, and update TAT. Rejection does not advance TAT. Python uses register_script; a local Lua installation is not required. +- Only the queue head requests a Redis permit for each process and model. Other threads wait on a Condition. Once admitted, a request leaves the queue immediately and starts the model call without waiting for earlier model calls to finish. +- A random jitter of 0 to 10 ms is added to Redis's suggested wait. Without a valid Retry-After, retries use exponential backoff with jitter, a 1-second base, and an 8-second cap. These are internal defaults and need no deployment configuration. +- Managed invocations disable hidden SDK retries. Connection/timeout errors and HTTP 408, 409, 429, and 5xx responses can be retried within the configured retry limit. Every retry must acquire a new permit. If Retry-After exceeds the internal wait limit, the request is not retried earlier than requested. +- Streaming requests are retried only during establishment. Errors after streaming starts do not replay the stream. Callers that stop consuming early should close the generator. +- A full queue, an expired wait budget, and Redis unavailability raise LLMRateLimitQueueFullError, LLMRateLimitTimeoutError, and LLMRateLimitUnavailableError respectively. These errors do not bypass the limiter through a backup model. +- Redis connection and read timeouts default to 0.5 seconds. The default failure policy is closed, which stops limited invocations. A late Redis response does not cause a model request to be sent, and permits already consumed or with uncertain status are not refunded. +- max_wait_seconds excludes model network time and failure backoff. It is not a total business-request timeout; callers still need an outer deadline. Synchronous Redis I/O may need to wait until the socket timeout before exiting. +- The local queue is not a durable task queue. Queue overflow, wait timeouts, and process exits do not automatically defer tasks. Redis failover or eviction of TAT keys may also reset the allowance. + +## Scope and Verification + +The integration currently covers Chat Completions in OpenAILLM and its Qwen, DeepSeek, and MiniMax subclasses, including regular calls, streaming establishment, and backup models. Independent implementations such as Azure, the Responses API, Ollama, and VLLM are not integrated yet. + +This version controls QPS only, not token usage, token traffic growth, or in-flight concurrency. It cannot guarantee prevention of every provider-side 429 response. + +INFO-level `[LLM_RATE_LIMIT] sending` logs record the model, attempt number, and permit-wait duration. WARNING logs record retries, queue overflow, wait timeouts, and Redis failures. The new logs do not include request bodies or credentials. + +```sh +poetry run pytest tests/configs/ tests/llms/ -q +MEMOS_TEST_LOCAL_REDIS=1 poetry run pytest tests/llms/test_qps_rate_limit_redis.py -q +``` + +The second command starts an isolated local Redis instance using only a Unix socket, with no TCP listener or persistence. It does not read production Redis settings. Logs are stored in a pytest-managed `redis-gcra*` temporary directory; the short-path temporary socket is cleaned up on exit. diff --git a/src/memos/configs/llm.py b/src/memos/configs/llm.py index ef441b37f..9068d254a 100644 --- a/src/memos/configs/llm.py +++ b/src/memos/configs/llm.py @@ -3,6 +3,7 @@ from pydantic import Field, field_validator, model_validator from memos.configs.base import BaseConfig +from memos.configs.llm_rate_limit import LLMRateLimitConfig class BaseLLMConfig(BaseConfig): @@ -23,6 +24,17 @@ class BaseLLMConfig(BaseConfig): class OpenAILLMConfig(BaseLLMConfig): + rate_limit: LLMRateLimitConfig = Field(default_factory=LLMRateLimitConfig.load) + + @field_validator("rate_limit", mode="before") + @classmethod + def load_rate_limit(cls, value: Any) -> LLMRateLimitConfig: + if isinstance(value, LLMRateLimitConfig): + return value + if not isinstance(value, dict): + raise ValueError("rate_limit must be a configuration object") + return LLMRateLimitConfig.load(value) + api_key: str = Field(..., description="API key for OpenAI") api_base: str = Field( default="https://api.openai.com/v1", description="Base URL for OpenAI API" diff --git a/src/memos/configs/llm_rate_limit.py b/src/memos/configs/llm_rate_limit.py new file mode 100644 index 000000000..28df81bbb --- /dev/null +++ b/src/memos/configs/llm_rate_limit.py @@ -0,0 +1,124 @@ +"""Startup-loaded, opt-in QPS policies for OpenAI-compatible LLM calls.""" + +import math +import os + +from typing import Any, Literal + +from pydantic import ConfigDict, Field, TypeAdapter, model_validator + +from memos.configs.base import BaseConfig +from memos.exceptions import ConfigurationError + + +_RULE_FIELDS = frozenset({"qps", "burst", "max_wait_seconds", "queue_capacity", "retry_attempts"}) + + +class QPSLimitRule(BaseConfig): + """A quota pool's pacing, waiting and retry policy; burst is total capacity.""" + + enabled: bool = True + qps: float = Field(default=5.0, ge=0.001, le=1_000_000, allow_inf_nan=False) + burst: int = Field(default=2, ge=1, le=1_000_000) + max_wait_seconds: float = Field(default=30.0, gt=0, le=3600, allow_inf_nan=False) + queue_capacity: int = Field(default=16, ge=1, le=100_000) + wait_jitter_seconds: float = Field(default=0.01, ge=0, le=1, allow_inf_nan=False) + retry_attempts: int = Field( + default=1, ge=0, le=10, description="Retries after the first attempt" + ) + retry_initial_delay: float = Field(default=1.0, gt=0, le=300, allow_inf_nan=False) + retry_max_delay: float = Field(default=8.0, gt=0, le=300, allow_inf_nan=False) + + @model_validator(mode="after") + def validate_policy(self) -> "QPSLimitRule": + if self.retry_initial_delay > self.retry_max_delay: + raise ValueError("retry_initial_delay must not exceed retry_max_delay") + if math.ceil(1_000_000 / self.qps) * self.burst > 1_000_000_000_000: + raise ValueError("The burst recovery horizon must not exceed 1000000 seconds") + return self + + +class LLMRateLimitConfig(QPSLimitRule): + """Load enabled/rules from env and reuse scheduler Redis connection settings.""" + + model_config = ConfigDict(extra="forbid", strict=True, hide_input_in_errors=True) + + enabled: bool = False + rules: dict[str, dict[str, Any]] = Field( + default_factory=lambda: {"gpt-4o-mini": {}}, + description="Selected models and their QPS, burst, wait, queue and retry settings", + ) + redis_host: str | None = Field(default=None, min_length=1) + redis_port: int = Field(default=6379, ge=1, le=65535) + redis_db: int = Field(default=0, ge=0) + redis_username: str | None = None + redis_password: str | None = Field(default=None, repr=False) + redis_ssl: bool = False + redis_socket_timeout: float = Field(default=0.5, gt=0, le=30, allow_inf_nan=False) + key_prefix: str = Field(default="memos:llm:gcra", min_length=1, max_length=128) + failure_mode: Literal["closed", "open"] = "closed" + + def _base_rule(self) -> dict[str, Any]: + return self.model_dump(include=set(QPSLimitRule.model_fields) - {"model_schema"}) + + @model_validator(mode="after") + def validate_rules(self) -> "LLMRateLimitConfig": + for model, overrides in self.rules.items(): + if not model.strip() or model == "*": + raise ValueError("rules must use explicit nonblank model names") + if set(overrides) - _RULE_FIELDS: + raise ValueError( + "rules only support qps, burst, max_wait_seconds, queue_capacity, retry_attempts" + ) + QPSLimitRule.model_validate({**self._base_rule(), **overrides}) + if any(char in self.key_prefix for char in "{}") or not self.key_prefix.strip(): + raise ValueError("key_prefix must be nonblank and must not contain Redis hash tags") + return self + + def rule_for(self, model: str) -> QPSLimitRule | None: + """Resolve the actual wire model without applying policies to unlisted models.""" + if not self.enabled or model not in self.rules: + return None + return QPSLimitRule.model_validate({**self._base_rule(), **self.rules[model]}) + + @classmethod + def load(cls, overrides: dict[str, Any] | None = None) -> "LLMRateLimitConfig": + """Read configuration at construction time; never load or modify .env here.""" + supported = {"MEMOS_LLM_RATE_LIMIT_ENABLED", "MEMOS_LLM_RATE_LIMIT_RULES"} + 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) + ) + values: dict[str, Any] = {} + env_values: dict[str, str] = {} + for name in ("host", "port", "db", "username", "password", "ssl"): + raw = os.getenv(f"MEMSCHEDULER_REDIS_{name.upper()}") + if raw: + env_values[f"redis_{name}"] = raw + for name in ("enabled", "rules"): + raw = os.getenv(f"MEMOS_LLM_RATE_LIMIT_{name.upper()}") + if raw is not None: + env_values[name] = raw + explicit = overrides or {} + for name, raw in env_values.items(): + if name in explicit: + continue + annotation = cls.model_fields[name].annotation + adapter = TypeAdapter(annotation) + try: + values[name] = ( + adapter.validate_json(raw, strict=False) + if name == "rules" + else adapter.validate_python(raw, strict=False) + ) + except ValueError: + raise ConfigurationError( + f"Invalid LLM rate limit environment setting: {name}" + ) from None + return cls.model_validate({**values, **explicit}) diff --git a/src/memos/exceptions.py b/src/memos/exceptions.py index 28b83c3df..511e68d8f 100644 --- a/src/memos/exceptions.py +++ b/src/memos/exceptions.py @@ -24,6 +24,19 @@ class VectorDBError(MemOSError): ... class LLMError(MemOSError): ... +class LLMRateLimitError(LLMError): + """Local admission failed; do not bypass it through model retries or fallback.""" + + +class LLMRateLimitTimeoutError(LLMRateLimitError): ... + + +class LLMRateLimitQueueFullError(LLMRateLimitError): ... + + +class LLMRateLimitUnavailableError(LLMRateLimitError): ... + + class EmbedderError(MemOSError): ... diff --git a/src/memos/llms/openai.py b/src/memos/llms/openai.py index d01c56726..b61489b5f 100644 --- a/src/memos/llms/openai.py +++ b/src/memos/llms/openai.py @@ -10,6 +10,8 @@ from openai.types.chat.chat_completion_message_tool_call import ChatCompletionMessageToolCall from memos.configs.llm import AzureLLMConfig, OpenAILLMConfig +from memos.exceptions import ConfigurationError, LLMRateLimitError +from memos.llms import rate_limit from memos.llms.base import BaseLLM from memos.llms.utils import remove_thinking_tags from memos.log import get_logger @@ -99,13 +101,17 @@ def generate(self, messages: MessageList, **kwargs) -> str: logger.info(f"OpenAI LLM Request body: {request_body}") try: - response = self.client.chat.completions.create(**request_body) + response = rate_limit.create_completion( + self.client, request_body, self.config.rate_limit + ) cost_time = time.perf_counter() - start_time logger.info( f"Request body: {request_body}, Response from OpenAI: " f"{response.model_dump_json()}, Cost time: {cost_time}" ) return self._parse_response(response) + except (LLMRateLimitError, ConfigurationError): + raise except Exception as e: if not self.use_backup_client: raise @@ -117,7 +123,9 @@ def generate(self, messages: MessageList, **kwargs) -> str: **request_body, "model": self.config.backup_model_name_or_path or request_body["model"], } - backup_response = self.backup_client.chat.completions.create(**backup_body) + backup_response = rate_limit.create_completion( + self.backup_client, backup_body, self.config.rate_limit + ) cost_time = time.perf_counter() - start_time logger.info( f"Backup LLM request succeeded, Response: " @@ -141,30 +149,32 @@ def generate_stream(self, messages: MessageList, **kwargs) -> Generator[str, Non request_body["stream"] = True logger.info(f"OpenAI LLM Stream Request body: {request_body}") - response = self.client.chat.completions.create(**request_body) + response = rate_limit.create_completion(self.client, request_body, self.config.rate_limit) reasoning_started = False - - for chunk in response: - if not chunk.choices: - continue - delta = chunk.choices[0].delta - - # Support for custom 'reasoning_content' (if present in OpenAI-compatible models like Qwen, DeepSeek) - if hasattr(delta, "reasoning_content") and delta.reasoning_content: - if not reasoning_started and not self.config.remove_think_prefix: - yield "" - reasoning_started = True - yield delta.reasoning_content - elif hasattr(delta, "content") and delta.content: - if reasoning_started and not self.config.remove_think_prefix: - yield "" - reasoning_started = False - yield delta.content - - # Ensure we close the block if not already done - if reasoning_started and not self.config.remove_think_prefix: - yield "" + try: + for chunk in response: + if not chunk.choices: + continue + delta = chunk.choices[0].delta + + if hasattr(delta, "reasoning_content") and delta.reasoning_content: + if not reasoning_started and not self.config.remove_think_prefix: + yield "" + reasoning_started = True + yield delta.reasoning_content + elif hasattr(delta, "content") and delta.content: + if reasoning_started and not self.config.remove_think_prefix: + yield "" + reasoning_started = False + yield delta.content + + if reasoning_started and not self.config.remove_think_prefix: + yield "" + finally: + close = getattr(response, "close", None) + if callable(close): + close() def tool_call_parser(self, tool_calls: list[ChatCompletionMessageToolCall]) -> list[dict]: """Parse tool calls from OpenAI response.""" diff --git a/src/memos/llms/rate_limit.py b/src/memos/llms/rate_limit.py new file mode 100644 index 000000000..faa72cef3 --- /dev/null +++ b/src/memos/llms/rate_limit.py @@ -0,0 +1,275 @@ +"""Opt-in Redis GCRA admission for OpenAI-compatible chat completions.""" + +import math +import os +import random +import threading +import time + +from collections import deque +from datetime import datetime, timezone +from email.utils import parsedate_to_datetime +from typing import Any + +import openai + +from memos.configs.llm_rate_limit import LLMRateLimitConfig, QPSLimitRule +from memos.exceptions import ( + ConfigurationError, + LLMRateLimitQueueFullError, + LLMRateLimitTimeoutError, + LLMRateLimitUnavailableError, +) +from memos.log import get_logger + + +logger = get_logger(__name__) + +GCRA_LUA = """ +local interval = tonumber(ARGV[1]) +local burst = tonumber(ARGV[2]) +local clock = redis.call('TIME') +local now = tonumber(clock[1]) * 1000000 + tonumber(clock[2]) +local raw = redis.call('GET', KEYS[1]) +local tat = tonumber(raw) +if raw and not tat then + return redis.error_reply('Invalid GCRA state') +end +tat = tat or now +local wait = tat - (burst - 1) * interval - now +if wait > 0 then + return {0, math.ceil(wait)} +end +local next_tat = math.max(tat, now) + interval +redis.call('SET', KEYS[1], string.format('%.0f', next_tat), + 'PX', math.max(1, math.ceil((next_tat - now) / 1000))) +return {1, 0} +""" + + +def _create_redis_client(config: LLMRateLimitConfig) -> Any: + if not config.redis_host: + raise ConfigurationError("LLM QPS limiting requires a Redis host") + try: + import redis + + from redis.backoff import NoBackoff + from redis.retry import Retry + except ImportError as exc: + raise ConfigurationError( + "Install the mem-scheduler extras to enable LLM QPS limiting" + ) from exc + return redis.Redis( + host=config.redis_host, + port=config.redis_port, + db=config.redis_db, + username=config.redis_username, + password=config.redis_password, + ssl=config.redis_ssl, + socket_timeout=config.redis_socket_timeout, + socket_connect_timeout=config.redis_socket_timeout, + max_connections=2, + decode_responses=True, + retry=Retry(NoBackoff(), 0), + ) + + +class RedisGCRALimiter: + """One bounded queue per process/quota; only its head accesses Redis.""" + + def __init__(self, config: LLMRateLimitConfig, rule: QPSLimitRule, key: str) -> None: + self.config = config.model_copy(deep=True) + self.rule = rule.model_copy(deep=True) + self.redis_key = key + self._client = _create_redis_client(config) + self._script = self._client.register_script(GCRA_LUA) + self._condition = threading.Condition() + self._queue: deque[object] = deque() + + @property + def pending_count(self) -> int: + """Number of local waiters including the active queue head.""" + with self._condition: + return len(self._queue) + + def acquire(self, timeout_seconds: float | None = None) -> None: + """Wait for admission, bounded by a cumulative monotonic deadline.""" + from redis.exceptions import RedisError + + started = time.monotonic() + timeout = self.rule.max_wait_seconds if timeout_seconds is None else timeout_seconds + deadline = started + timeout + waiter = object() + with self._condition: + if len(self._queue) >= self.rule.queue_capacity: + logger.warning("[LLM_RATE_LIMIT] queue_full key=%s", self.redis_key) + raise LLMRateLimitQueueFullError("LLM permit queue is full") + self._queue.append(waiter) + checks = 0 + next_check = started + try: + while True: + with self._condition: + now = time.monotonic() + if now >= deadline: + raise LLMRateLimitTimeoutError("LLM permit waiting deadline exceeded") + head = self._queue[0] is waiter + if not head or now < next_check: + delay = min(deadline - now, next_check - now) if head else deadline - now + self._condition.wait(delay) + continue + try: + checks += 1 + reply = self._script( + keys=[self.redis_key], + args=[math.ceil(1_000_000 / self.rule.qps), self.rule.burst], + ) + except RedisError: + logger.warning( + "[LLM_RATE_LIMIT] redis_unavailable key=%s failure_mode=%s", + self.redis_key, + self.config.failure_mode, + ) + if self.config.failure_mode == "open" and time.monotonic() < deadline: + return + raise LLMRateLimitUnavailableError("Redis LLM limiter is unavailable") from None + if time.monotonic() >= deadline: + # A late successful reply consumes a permit; never refund uncertain sends. + raise LLMRateLimitTimeoutError("LLM permit waiting deadline exceeded") + if not isinstance(reply, list | tuple) or len(reply) != 2 or reply[0] not in (0, 1): + raise LLMRateLimitUnavailableError("Invalid Redis LLM limiter response") + if reply[0]: + logger.debug( + "[LLM_RATE_LIMIT] granted key=%s wait_ms=%.2f checks=%d", + self.redis_key, + (time.monotonic() - started) * 1000, + checks, + ) + return + wait_us = reply[1] + if ( + not isinstance(wait_us, int | float) + or not math.isfinite(wait_us) + or wait_us <= 0 + ): + raise LLMRateLimitUnavailableError("Invalid Redis LLM limiter wait time") + next_check = ( + time.monotonic() + + wait_us / 1_000_000 + + random.uniform(0, self.rule.wait_jitter_seconds) + ) + except LLMRateLimitTimeoutError: + logger.warning("[LLM_RATE_LIMIT] wait_timeout key=%s checks=%d", self.redis_key, checks) + raise + finally: + with self._condition: + self._queue.remove(waiter) + self._condition.notify_all() + + +_registry: dict[tuple, RedisGCRALimiter] = {} +_registry_lock = threading.Lock() + + +def _after_fork() -> None: + global _registry, _registry_lock + _registry = {} + _registry_lock = threading.Lock() + + +if hasattr(os, "register_at_fork"): + os.register_at_fork(after_in_child=_after_fork) + + +def get_limiter(config: LLMRateLimitConfig, rule: QPSLimitRule, model: str) -> RedisGCRALimiter: + """Share admission for the same actual model, independently of its endpoint.""" + key = f"{config.key_prefix}:{model}" + identity = ( + config.redis_host, + config.redis_port, + config.redis_db, + config.redis_username, + config.redis_password, + config.redis_ssl, + key, + ) + with _registry_lock: + current = _registry.get(identity) + if current is not None: + if ( + current.rule.qps != rule.qps + or current.rule.burst != rule.burst + or current.rule.queue_capacity != rule.queue_capacity + or current.rule.wait_jitter_seconds != rule.wait_jitter_seconds + or current.config.failure_mode != config.failure_mode + or current.config.redis_socket_timeout != config.redis_socket_timeout + ): + raise ConfigurationError("Conflicting LLM limiter settings for the same quota pool") + return current + limiter = RedisGCRALimiter(config, rule, key) + _registry[identity] = limiter + return limiter + + +def _retry_delay(error: openai.APIError, rule: QPSLimitRule, attempt: int) -> float | None: + if isinstance(error, openai.APIStatusError): + if error.status_code not in (408, 409, 429) and error.status_code < 500: + return None + headers = error.response.headers + requested_delay = None + try: + if "retry-after-ms" in headers: + requested_delay = float(headers["retry-after-ms"]) / 1000 + elif "retry-after" in headers: + try: + requested_delay = float(headers["retry-after"]) + except ValueError: + date = parsedate_to_datetime(headers["retry-after"]) + if date.tzinfo is None: + date = date.replace(tzinfo=timezone.utc) + requested_delay = (date - datetime.now(timezone.utc)).total_seconds() + except (ValueError, TypeError, OverflowError): + requested_delay = None + if requested_delay is not None and math.isfinite(requested_delay) and requested_delay >= 0: + return requested_delay if requested_delay <= rule.retry_max_delay else None + elif not isinstance(error, openai.APIConnectionError): + return None + cap = min(rule.retry_max_delay, rule.retry_initial_delay * 2**attempt) + return random.uniform(cap / 2, cap) + + +def create_completion(client: Any, body: dict, config: LLMRateLimitConfig) -> Any: + """Gate each SDK wire attempt; leave legacy behavior unchanged when unselected.""" + model = body["model"] + rule = config.rule_for(model) + if rule is None: + return client.chat.completions.create(**body) + limiter = get_limiter(config, rule, model) + managed_client = client.with_options(max_retries=0) + remaining_wait = rule.max_wait_seconds + for attempt in range(rule.retry_attempts + 1): + started = time.monotonic() + limiter.acquire(timeout_seconds=remaining_wait) + remaining_wait -= time.monotonic() - started + logger.info( + "[LLM_RATE_LIMIT] sending model=%s attempt=%d permit_wait_ms=%.2f", + model, + attempt + 1, + (time.monotonic() - started) * 1000, + ) + try: + return managed_client.chat.completions.create(**body) + except openai.APIError as error: + if attempt == rule.retry_attempts: + raise + delay = _retry_delay(error, rule, attempt) + if delay is None: + raise + logger.warning( + "[LLM_RATE_LIMIT] retry model=%s attempt=%d error_type=%s delay_s=%.3f", + model, + attempt + 1, + type(error).__name__, + delay, + ) + time.sleep(delay) diff --git a/tests/configs/test_llm.py b/tests/configs/test_llm.py index bc036fdb8..d0efed61e 100644 --- a/tests/configs/test_llm.py +++ b/tests/configs/test_llm.py @@ -47,6 +47,7 @@ def test_base_llm_config(): def test_openai_llm_config(): check_config_base_class( OpenAILLMConfig, + factory_fields=["rate_limit"], required_fields=["model_name_or_path", "api_key"], optional_fields=[ "temperature", @@ -150,6 +151,7 @@ def test_hf_llm_config(): def test_minimax_llm_config(): check_config_base_class( MinimaxLLMConfig, + factory_fields=["rate_limit"], required_fields=["model_name_or_path", "api_key"], optional_fields=[ "temperature", diff --git a/tests/configs/test_llm_rate_limit.py b/tests/configs/test_llm_rate_limit.py new file mode 100644 index 000000000..e9c44b080 --- /dev/null +++ b/tests/configs/test_llm_rate_limit.py @@ -0,0 +1,176 @@ +"""Startup policy loading without production environment or external services.""" + +import os + +from pathlib import Path + +import pytest + +from dotenv import dotenv_values + +from memos.configs.llm import LLMConfigFactory, OpenAILLMConfig +from memos.configs.llm_rate_limit import LLMRateLimitConfig +from memos.exceptions import ConfigurationError + + +@pytest.fixture(autouse=True) +def isolated_environment(monkeypatch): + for key in list(os.environ): + if key.startswith(("MEMOS_LLM_RATE_LIMIT_", "MEMSCHEDULER_REDIS_")): + monkeypatch.delenv(key) + + +def test_rules_are_the_only_model_selection(monkeypatch): + monkeypatch.setenv("MEMOS_LLM_RATE_LIMIT_ENABLED", "true") + monkeypatch.setenv("MEMOS_LLM_RATE_LIMIT_RULES", '{"extra": {"qps": 2, "burst": 1}}') + config = LLMRateLimitConfig.load() + assert config.rule_for("extra").qps == 2 + assert config.rule_for("gpt-4o-mini") is None + assert config.rule_for("unlisted") is None + assert "models" not in config.model_dump() + + +@pytest.mark.parametrize("rules", [None, '{"gpt-4o-mini":{}}']) +def test_default_rule_values(monkeypatch, rules): + assert LLMRateLimitConfig.load().enabled is False + monkeypatch.setenv("MEMOS_LLM_RATE_LIMIT_ENABLED", "true") + if rules is not None: + monkeypatch.setenv("MEMOS_LLM_RATE_LIMIT_RULES", rules) + rule = LLMRateLimitConfig.load().rule_for("gpt-4o-mini") + assert ( + rule.qps, + rule.burst, + rule.max_wait_seconds, + rule.queue_capacity, + rule.retry_attempts, + ) == (5, 2, 30, 16, 1) + + +def test_docker_full_example_is_disabled_and_matches_defaults(monkeypatch): + example = Path(__file__).resolve().parents[2] / "docker" / ".env.example-full" + values = dotenv_values(example, interpolate=False) + assert values["MEMOS_LLM_RATE_LIMIT_ENABLED"] == "false" + for name in ("MEMOS_LLM_RATE_LIMIT_ENABLED", "MEMOS_LLM_RATE_LIMIT_RULES"): + monkeypatch.setenv(name, values[name]) + config = LLMRateLimitConfig.load() + assert config.enabled is False + assert set(config.rules) == {values["MOS_CHAT_MODEL"]} + assert config.rule_for("gpt-4o-mini") is None + monkeypatch.setenv("MEMOS_LLM_RATE_LIMIT_ENABLED", "true") + rule = LLMRateLimitConfig.load().rule_for("gpt-4o-mini") + defaults = LLMRateLimitConfig(enabled=True).rule_for("gpt-4o-mini") + assert rule is not None + assert rule == defaults + + +def test_multiple_models_need_only_qps_and_burst(monkeypatch): + monkeypatch.setenv("MEMOS_LLM_RATE_LIMIT_ENABLED", "true") + monkeypatch.setenv( + "MEMOS_LLM_RATE_LIMIT_RULES", + '{"gpt-4o-mini":{"qps":5,"burst":2},"other":{"qps":3,"burst":1}}', + ) + config = LLMRateLimitConfig.load() + assert (config.rule_for("gpt-4o-mini").qps, config.rule_for("gpt-4o-mini").burst) == (5, 2) + assert (config.rule_for("other").qps, config.rule_for("other").burst) == (3, 1) + assert config.rule_for("unlisted") is None + assert "scope" not in config.model_dump() + + +@pytest.mark.parametrize( + "suffix", + ["SCOPE", "MODELS", "QPS", "BURST", "WAIT_JITTER_SECONDS", "REDIS_HOST", "CONFIG_FILE"], +) +def test_removed_environment_settings_are_rejected(monkeypatch, suffix): + monkeypatch.setenv(f"MEMOS_LLM_RATE_LIMIT_{suffix}", "obsolete") + with pytest.raises(ConfigurationError, match=suffix): + LLMRateLimitConfig.load() + + +def test_explicit_override_wins_over_invalid_environment(monkeypatch): + monkeypatch.setenv("MEMOS_LLM_RATE_LIMIT_RULES", "invalid") + assert LLMRateLimitConfig.load({"rules": {}}).rules == {} + with pytest.raises(ConfigurationError, match="rules"): + LLMRateLimitConfig.load() + + +def test_all_supported_rule_fields_and_internal_defaults(monkeypatch): + monkeypatch.setenv("MEMOS_LLM_RATE_LIMIT_ENABLED", "true") + monkeypatch.setenv( + "MEMOS_LLM_RATE_LIMIT_RULES", + '{"gpt-4o-mini":{"qps":3,"burst":1,"max_wait_seconds":10,' + '"queue_capacity":8,"retry_attempts":2}}', + ) + rule = LLMRateLimitConfig.load().rule_for("gpt-4o-mini") + assert ( + rule.qps, + rule.burst, + rule.max_wait_seconds, + rule.queue_capacity, + rule.retry_attempts, + ) == (3, 1, 10, 8, 2) + assert (rule.wait_jitter_seconds, rule.retry_initial_delay, rule.retry_max_delay) == ( + 0.01, + 1, + 8, + ) + + +@pytest.mark.parametrize( + "field", ["wait_jitter_seconds", "retry_initial_delay", "retry_max_delay", "enabled", "scope"] +) +def test_rules_do_not_expose_internal_tuning(monkeypatch, field): + monkeypatch.setenv("MEMOS_LLM_RATE_LIMIT_RULES", '{"extra":{"' + field + '":1}}') + with pytest.raises(ValueError): + LLMRateLimitConfig.load() + + +def test_empty_rules_disable_all_models(monkeypatch): + monkeypatch.setenv("MEMOS_LLM_RATE_LIMIT_ENABLED", "true") + monkeypatch.setenv("MEMOS_LLM_RATE_LIMIT_RULES", "{}") + assert LLMRateLimitConfig.load().rule_for("gpt-4o-mini") is None + + +def test_settings_are_snapshotted_not_read_for_every_request(monkeypatch): + monkeypatch.setenv("MEMOS_LLM_RATE_LIMIT_ENABLED", "true") + config = OpenAILLMConfig(api_key="test", model_name_or_path="gpt-4o-mini") + monkeypatch.setenv("MEMOS_LLM_RATE_LIMIT_ENABLED", "false") + assert config.rate_limit.rule_for("gpt-4o-mini") is not None + assert ( + OpenAILLMConfig(api_key="test", model_name_or_path="gpt-4o-mini").rate_limit.enabled + is False + ) + + +def test_factory_accepts_nested_policy_and_keeps_it_out_of_llm_body(): + config = LLMConfigFactory.model_validate( + { + "backend": "openai", + "config": { + "api_key": "test", + "model_name_or_path": "gpt-4o-mini", + "rate_limit": {"enabled": True, "qps": 4, "redis_host": "test.invalid"}, + }, + } + ) + assert config.config.rate_limit.qps == 4 + from memos.llms.openai import OpenAILLM + + llm = OpenAILLM.__new__(OpenAILLM) + llm.config = config.config + assert "rate_limit" not in llm._build_request_body([]) + + +def test_bad_rules_do_not_expose_contents(monkeypatch): + monkeypatch.setenv("MEMOS_LLM_RATE_LIMIT_RULES", "private-test-content-not-json") + with pytest.raises(ConfigurationError) as error: + LLMRateLimitConfig.load() + assert "private-test-content" not in str(error.value) + + +def test_scheduler_redis_reuse_and_explicit_override(monkeypatch): + monkeypatch.setenv("MEMSCHEDULER_REDIS_HOST", "fallback.invalid") + monkeypatch.setenv("MEMSCHEDULER_REDIS_DB", "3") + monkeypatch.setenv("MEMSCHEDULER_REDIS_SSL", "true") + config = LLMRateLimitConfig.load() + assert (config.redis_host, config.redis_db, config.redis_ssl) == ("fallback.invalid", 3, True) + assert LLMRateLimitConfig.load({"redis_db": 7}).redis_db == 7 diff --git a/tests/llms/test_qps_rate_limit.py b/tests/llms/test_qps_rate_limit.py new file mode 100644 index 000000000..167adb5ea --- /dev/null +++ b/tests/llms/test_qps_rate_limit.py @@ -0,0 +1,436 @@ +"""QPS limiter configuration, local queue and SDK integration tests.""" + +import json +import socket +import threading +import time + +from concurrent.futures import ThreadPoolExecutor +from unittest.mock import MagicMock + +import httpx +import openai +import pytest + +from pydantic import ValidationError + +from memos.configs.llm import OpenAILLMConfig +from memos.configs.llm_rate_limit import LLMRateLimitConfig, QPSLimitRule +from memos.exceptions import LLMRateLimitError, LLMRateLimitQueueFullError, LLMRateLimitTimeoutError +from memos.llms import rate_limit +from memos.llms.openai import OpenAILLM + + +@pytest.fixture(autouse=True) +def isolated_settings(monkeypatch): + import os + + for key in list(os.environ): + if key.startswith(("MEMOS_LLM_RATE_LIMIT_", "MEMSCHEDULER_REDIS_")): + monkeypatch.delenv(key) + monkeypatch.setattr(rate_limit, "_registry", {}) + + def reject_network(*args, **kwargs): + pytest.fail("Unit tests must mock external connections") + + monkeypatch.setattr(socket.socket, "connect", reject_network) + monkeypatch.setattr(socket.socket, "connect_ex", reject_network) + + +def settings(**kwargs): + return LLMRateLimitConfig(enabled=True, redis_host="test.invalid", **kwargs) + + +def test_env_and_explicit_precedence(monkeypatch): + monkeypatch.setenv("MEMOS_LLM_RATE_LIMIT_ENABLED", "true") + monkeypatch.setenv("MEMOS_LLM_RATE_LIMIT_RULES", '{"gpt-4o-mini":{"qps":8}}') + monkeypatch.setenv("MEMSCHEDULER_REDIS_HOST", "scheduler.invalid") + monkeypatch.setenv("MEMSCHEDULER_REDIS_DB", "3") + config = OpenAILLMConfig( + api_key="test", model_name_or_path="gpt-4o-mini", rate_limit={"redis_db": 4} + ).rate_limit + rule = config.rule_for("gpt-4o-mini") + assert (config.enabled, rule.qps, rule.burst) == (True, 8, 2) + assert (config.redis_host, config.redis_db) == ("scheduler.invalid", 4) + disabled = LLMRateLimitConfig.load({"enabled": False}) + assert disabled.rule_for("gpt-4o-mini") is None + + +def test_model_rules_inherit_defaults_and_omit_unselected_models(): + config = settings(qps=8, rules={"special": {"burst": 3}}) + assert config.rule_for("special").qps == 8 + assert config.rule_for("special").burst == 3 + assert config.rule_for("gpt-4o-mini") is None + + +@pytest.mark.parametrize( + "bad", + [ + {"qps": 0}, + {"qps": float("nan")}, + {"qps": float("inf")}, + {"burst": 0}, + {"queue_capacity": 0}, + {"max_wait_seconds": -1}, + {"retry_attempts": -1}, + {"retry_initial_delay": 3, "retry_max_delay": 1}, + {"burst": 1.5}, + {"rules": {"special": {"qps": -1}}}, + {"rules": {"special": {"typo": 1}}}, + {"scope": "shared"}, + ], +) +def test_invalid_parameters_fail_validation(bad): + with pytest.raises(ValidationError): + settings(**bad) + + +def make_limiter(monkeypatch, script, **kwargs): + config = settings(**kwargs) + fake = MagicMock() + fake.register_script.return_value = script + monkeypatch.setattr(rate_limit, "_create_redis_client", lambda _: fake) + return rate_limit.get_limiter(config, config.rule_for("gpt-4o-mini"), "gpt-4o-mini") + + +def test_wait_uses_returned_delay_and_passes_one_key(monkeypatch): + script = MagicMock(side_effect=[[0, 1000], [1, 0]]) + limiter = make_limiter(monkeypatch, script, wait_jitter_seconds=0) + limiter.acquire(timeout_seconds=1) + assert script.call_count == 2 + assert len(script.call_args.kwargs["keys"]) == 1 + assert script.call_args.kwargs["args"] == [200000, 2] + + +def test_wait_timeout_leaves_queue_usable(monkeypatch): + script = MagicMock(return_value=[0, 1000000]) + limiter = make_limiter(monkeypatch, script) + with pytest.raises(LLMRateLimitTimeoutError): + limiter.acquire(timeout_seconds=0.02) + script.return_value = [1, 0] + limiter.acquire(timeout_seconds=1) + + +def test_only_queue_head_checks_redis_and_queue_is_bounded(monkeypatch): + entered, release = threading.Event(), threading.Event() + + def script(**_): + entered.set() + assert release.wait(2) + return [1, 0] + + mocked = MagicMock(side_effect=script) + limiter = make_limiter(monkeypatch, mocked, queue_capacity=2) + with ThreadPoolExecutor(max_workers=2) as pool: + first = pool.submit(limiter.acquire, 1) + assert entered.wait(1) + second = pool.submit(limiter.acquire, 1) + deadline = time.monotonic() + 1 + while limiter.pending_count != 2 and time.monotonic() < deadline: + time.sleep(0.001) + try: + assert limiter.pending_count == 2 + assert mocked.call_count == 1 + with pytest.raises(LLMRateLimitQueueFullError): + limiter.acquire(1) + finally: + release.set() + first.result() + second.result() + assert limiter.pending_count == 0 + + +@pytest.mark.parametrize("mode,raises", [("closed", True), ("open", False)]) +def test_redis_outage_policy_and_cleanup(monkeypatch, mode, raises): + redis = pytest.importorskip("redis") + script = MagicMock(side_effect=redis.ConnectionError("unavailable")) + limiter = make_limiter(monkeypatch, script, failure_mode=mode) + if raises: + with pytest.raises(LLMRateLimitError): + limiter.acquire(1) + else: + limiter.acquire(1) + assert limiter.pending_count == 0 + + +def test_registry_rejects_conflicting_policy(monkeypatch): + from memos.exceptions import ConfigurationError + + make_limiter(monkeypatch, MagicMock(return_value=[1, 0])) + with pytest.raises(ConfigurationError): + make_limiter(monkeypatch, MagicMock(return_value=[1, 0]), qps=8) + + +def response_body(): + return { + "id": "test", + "object": "chat.completion", + "created": 0, + "model": "gpt-4o-mini", + "choices": [ + {"index": 0, "finish_reason": "stop", "message": {"role": "assistant", "content": "ok"}} + ], + } + + +def make_llm(monkeypatch, handler, **overrides): + config = OpenAILLMConfig( + api_key="test", + model_name_or_path="gpt-4o-mini", + api_base="https://api.test/v1", + rate_limit=settings(retry_initial_delay=0.001, retry_max_delay=0.01), + **overrides, + ) + llm = OpenAILLM(config) + llm.client.close() + llm.client = openai.Client( + api_key="test", + base_url="https://api.test/v1", + http_client=httpx.Client(transport=httpx.MockTransport(handler)), + ) + limiter = MagicMock() + monkeypatch.setattr(rate_limit, "get_limiter", MagicMock(return_value=limiter)) + return llm, limiter + + +def test_every_wire_retry_acquires_and_sdk_retries_do_not_bypass(monkeypatch): + wire_calls = [] + + def handler(_): + wire_calls.append(1) + return httpx.Response(429, json={"error": {"message": "limited", "type": "rate_limit"}}) + + llm, limiter = make_llm(monkeypatch, handler) + with pytest.raises(openai.RateLimitError): + llm.generate([{"role": "user", "content": "test"}]) + assert len(wire_calls) == limiter.acquire.call_count == 2 + llm.client.close() + + +def test_disabled_keeps_original_client_path(monkeypatch): + llm, _ = make_llm(monkeypatch, lambda _: httpx.Response(200, json=response_body())) + llm.config.rate_limit.enabled = False + assert llm.generate([]) == "ok" + rate_limit.get_limiter.assert_not_called() + llm.client.close() + + +def test_actual_model_override_selects_policy(monkeypatch): + llm, limiter = make_llm(monkeypatch, lambda _: httpx.Response(200, json=response_body())) + llm.config.rate_limit.rules["override"] = {"qps": 3} + assert llm.generate([], model_name_or_path="override") == "ok" + assert rate_limit.get_limiter.call_args.args[-1] == "override" + limiter.acquire.assert_called_once() + llm.client.close() + + +def test_local_timeout_does_not_trigger_backup(monkeypatch): + llm, limiter = make_llm(monkeypatch, lambda _: httpx.Response(200, json=response_body())) + llm.use_backup_client = True + llm.backup_client = MagicMock() + limiter.acquire.side_effect = LLMRateLimitTimeoutError("deadline") + with pytest.raises(LLMRateLimitTimeoutError): + llm.generate([]) + llm.backup_client.chat.completions.create.assert_not_called() + llm.client.close() + + +def test_stream_creation_is_limited_and_midstream_errors_are_not_replayed(monkeypatch): + class BrokenStream(httpx.SyncByteStream): + def __iter__(self): + chunk = { + "id": "test", + "object": "chat.completion.chunk", + "created": 0, + "model": "gpt-4o-mini", + "choices": [{"index": 0, "delta": {"content": "hello"}}], + } + yield ("data: " + json.dumps(chunk) + "\n\n").encode() + raise httpx.ReadError("stream interrupted") + + llm, limiter = make_llm( + monkeypatch, + lambda _: httpx.Response( + 200, headers={"content-type": "text/event-stream"}, stream=BrokenStream() + ), + ) + stream = llm.generate_stream([]) + assert next(stream) == "hello" + with pytest.raises(httpx.ReadError): + next(stream) + limiter.acquire.assert_called_once() + llm.client.close() + + +def test_closing_generator_closes_provider_stream(monkeypatch): + from types import SimpleNamespace + + llm, _ = make_llm(monkeypatch, lambda _: httpx.Response(200, json=response_body())) + provider = MagicMock() + provider.__iter__.return_value = iter( + [SimpleNamespace(choices=[SimpleNamespace(delta=SimpleNamespace(content="hello"))])] + ) + monkeypatch.setattr(rate_limit, "create_completion", lambda *_: provider) + stream = llm.generate_stream([]) + assert next(stream) == "hello" + stream.close() + provider.close.assert_called_once() + llm.client.close() + + +def test_backup_attempt_is_also_limited(monkeypatch): + llm, limiter = make_llm( + monkeypatch, + lambda _: httpx.Response(429, json={"error": {"message": "limited", "type": "rate_limit"}}), + ) + llm.config.rate_limit.retry_attempts = 0 + llm.config.rate_limit.rules["backup"] = {"qps": 3, "burst": 1} + llm.config.backup_model_name_or_path = "backup" + llm.use_backup_client = True + llm.backup_client = openai.Client( + api_key="test-backup", + base_url="https://backup.test/v1", + http_client=httpx.Client( + transport=httpx.MockTransport(lambda _: httpx.Response(200, json=response_body())) + ), + ) + assert llm.generate([]) == "ok" + assert limiter.acquire.call_count == 2 + assert [call.args[-1] for call in rate_limit.get_limiter.call_args_list] == [ + "gpt-4o-mini", + "backup", + ] + llm.client.close() + llm.backup_client.close() + + +def test_nonretryable_status_and_retry_after_budget(monkeypatch): + for status, headers in [(400, {}), (401, {}), (429, {"retry-after": "60"})]: + llm, limiter = make_llm( + monkeypatch, + lambda _, status=status, headers=headers: httpx.Response( + status, headers=headers, json={"error": {"message": "test", "type": "test"}} + ), + ) + with pytest.raises(openai.APIStatusError): + llm.generate([]) + limiter.acquire.assert_called_once() + llm.client.close() + + +def test_retry_after_is_honored(): + error = openai.RateLimitError( + "test", + response=httpx.Response( + 429, + headers={"retry-after-ms": "1500"}, + request=httpx.Request("POST", "https://api.test"), + ), + body=None, + ) + assert rate_limit._retry_delay(error, QPSLimitRule(), 0) == 1.5 + + +def test_permit_budget_is_shared_between_retries(monkeypatch): + replies = [ + httpx.Response(429, json={"error": {"message": "test"}}), + httpx.Response(200, json=response_body()), + ] + llm, limiter = make_llm(monkeypatch, lambda _: replies.pop(0)) + limiter.acquire.side_effect = lambda **_: time.sleep(0.002) + assert llm.generate([]) == "ok" + budgets = [call.kwargs["timeout_seconds"] for call in limiter.acquire.call_args_list] + assert 0 < budgets[1] < budgets[0] + llm.client.close() + + +def test_expired_waiter_does_not_send_when_redis_returns_late(monkeypatch): + def late(**_): + time.sleep(0.02) + return [1, 0] + + limiter = make_limiter(monkeypatch, late) + with pytest.raises(LLMRateLimitTimeoutError): + limiter.acquire(0.005) + assert limiter.pending_count == 0 + + +def test_interruption_removes_queue_head(monkeypatch): + script = MagicMock(side_effect=[KeyboardInterrupt(), [1, 0]]) + limiter = make_limiter(monkeypatch, script) + with pytest.raises(KeyboardInterrupt): + limiter.acquire(1) + limiter.acquire(1) + assert limiter.pending_count == 0 + + +@pytest.mark.parametrize("prefix_override", [{}, {"key_prefix": "test:custom:gcra"}]) +def test_model_keys_share_instances_but_isolate_models(monkeypatch, prefix_override): + monkeypatch.setattr(rate_limit, "_create_redis_client", lambda _: MagicMock()) + config = settings( + redis_password="test-password", + rules={"gpt-4o-mini": {}, "other": {"qps": 3, "burst": 1}}, + **prefix_override, + ) + rule = config.rule_for("gpt-4o-mini") + first = rate_limit.get_limiter(config, rule, "gpt-4o-mini") + second = rate_limit.get_limiter(config.model_copy(deep=True), rule, "gpt-4o-mini") + other = rate_limit.get_limiter(config, config.rule_for("other"), "other") + assert first is second + assert first is not other + assert first.redis_key == f"{config.key_prefix}:gpt-4o-mini" + assert other.redis_key == f"{config.key_prefix}:other" + assert (first.rule.qps, first.rule.burst) == (5, 2) + assert (other.rule.qps, other.rule.burst) == (3, 1) + first._queue.append(object()) + assert other.pending_count == 0 + assert "test-password" not in first.redis_key + assert "test-password" not in repr(config) + + +def test_fork_reset_drops_local_queue_registry(): + rate_limit._registry["test"] = object() + old_lock = rate_limit._registry_lock + rate_limit._after_fork() + assert rate_limit._registry == {} + assert old_lock is not rate_limit._registry_lock + + +@pytest.mark.parametrize("reply", [None, [0, 0], [0, "invalid"], [2, 0], [1, 0, 0]]) +def test_malformed_redis_reply_never_fails_open(monkeypatch, reply): + limiter = make_limiter(monkeypatch, MagicMock(return_value=reply), failure_mode="open") + with pytest.raises(LLMRateLimitError): + limiter.acquire(1) + assert limiter.pending_count == 0 + + +def test_qwen_wait_header_survives_managed_sdk_copy(monkeypatch): + from memos.configs.llm import QwenLLMConfig + from memos.llms.qwen import QwenLLM + + captured = [] + + def handler(request): + captured.append(request.headers.get("X-DashScope-Wait-Timeout")) + return httpx.Response(200, json=response_body()) + + llm = QwenLLM( + QwenLLMConfig( + api_key="test", + api_base="https://api.test/v1", + model_name_or_path="qwen-flash", + rate_limit=settings(rules={"qwen-flash": {}}), + ) + ) + headers = llm.config.default_headers + llm.client.close() + llm.client = openai.Client( + api_key="test", + base_url="https://api.test/v1", + default_headers=headers, + http_client=httpx.Client(transport=httpx.MockTransport(handler)), + ) + monkeypatch.setattr(rate_limit, "get_limiter", MagicMock(return_value=MagicMock())) + assert llm.generate([]) == "ok" + assert captured == [headers["X-DashScope-Wait-Timeout"]] + llm.client.close() diff --git a/tests/llms/test_qps_rate_limit_redis.py b/tests/llms/test_qps_rate_limit_redis.py new file mode 100644 index 000000000..edb39faaf --- /dev/null +++ b/tests/llms/test_qps_rate_limit_redis.py @@ -0,0 +1,184 @@ +"""Opt-in integration tests against an isolated local Redis, never .env Redis.""" + +import multiprocessing +import os +import shutil +import subprocess +import tempfile +import time + +import pytest + +from memos.configs.llm_rate_limit import LLMRateLimitConfig +from memos.exceptions import LLMRateLimitTimeoutError +from memos.llms import rate_limit + + +redis = pytest.importorskip("redis") +pytestmark = pytest.mark.skipif( + os.getenv("MEMOS_TEST_LOCAL_REDIS") != "1", + reason="Set MEMOS_TEST_LOCAL_REDIS=1 to start an isolated local Redis", +) + + +@pytest.fixture(scope="module") +def local_redis(tmp_path_factory): + executable = shutil.which("redis-server") + if not executable: + pytest.skip("Local redis-server is not installed") + root = tmp_path_factory.mktemp("redis-gcra") + # Unix socket paths must remain short, even when pytest's temp root is long. + with tempfile.TemporaryDirectory(prefix="memos-qps-", dir="/tmp") as socket_dir: + socket_path = os.path.join(socket_dir, "redis.sock") + with (root / "redis.log").open("w") as log: + process = subprocess.Popen( + [ + executable, + "--port", + "0", + "--unixsocket", + socket_path, + "--unixsocketperm", + "700", + "--save", + "", + "--appendonly", + "no", + "--dir", + str(root), + ], + stdout=log, + stderr=subprocess.STDOUT, + ) + client = redis.Redis( + unix_socket_path=socket_path, decode_responses=True, socket_timeout=2 + ) + try: + for _ in range(100): + assert process.poll() is None, "Local Redis startup failed" + try: + if client.ping(): + break + except redis.ConnectionError: + time.sleep(0.05) + else: + pytest.fail("Local Redis startup timed out") + yield client, socket_path + finally: + client.close() + if process.poll() is None: + process.terminate() + try: + process.wait(timeout=5) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + + +def test_real_script_total_burst_and_deny_does_not_advance_tat(local_redis): + client, _ = local_redis + script = client.register_script(rate_limit.GCRA_LUA) + key = "test:burst" + assert script(keys=[key], args=[200_000, 2]) == [1, 0] + assert script(keys=[key], args=[200_000, 2]) == [1, 0] + before = client.get(key) + denied = script(keys=[key], args=[200_000, 2]) + assert denied[0] == 0 and 0 < denied[1] <= 200_000 + assert client.get(key) == before + assert 0 < client.pttl(key) <= 400 + time.sleep(denied[1] / 1_000_000 + 0.01) + assert script(keys=[key], args=[200_000, 2]) == [1, 0] + + +def test_real_queue_refill_and_script_cache_recovery(local_redis, monkeypatch): + client, _ = local_redis + monkeypatch.setattr(rate_limit, "_create_redis_client", lambda _: client) + config = LLMRateLimitConfig( + enabled=True, redis_host="unused.invalid", qps=5, burst=2, wait_jitter_seconds=0 + ) + limiter = rate_limit.RedisGCRALimiter(config, config.rule_for("gpt-4o-mini"), "test:queue") + limiter.acquire() + limiter.acquire() + start = time.monotonic() + limiter.acquire() + assert time.monotonic() - start >= 0.15 + client.script_flush() + limiter.acquire() + assert limiter.pending_count == 0 + + +def _process_attempts(socket_path, gate, output): + client = redis.Redis(unix_socket_path=socket_path, decode_responses=True, socket_timeout=2) + try: + script = client.register_script(rate_limit.GCRA_LUA) + client.ping() + gate.wait(timeout=40) + # A 100-second refill interval isolates atomic initial-burst behavior from timing noise. + output.put( + sum(script(keys=["test:processes"], args=[100_000_000, 2])[0] for _ in range(20)) + ) + finally: + client.close() + + +def test_four_processes_share_one_atomic_burst(local_redis): + _, socket_path = local_redis + context = multiprocessing.get_context("spawn") + gate = context.Barrier(4) + output = context.Queue() + processes = [ + context.Process(target=_process_attempts, args=(socket_path, gate, output)) + for _ in range(4) + ] + try: + for process in processes: + process.start() + assert sum(output.get(timeout=60) for _ in processes) == 2 + for process in processes: + process.join(timeout=5) + assert process.exitcode == 0 + finally: + for process in processes: + if process.pid is not None: + if process.is_alive(): + process.terminate() + process.join(timeout=5) + output.close() + output.join_thread() + + +def test_corrupt_tat_is_not_silently_reset(local_redis): + client, _ = local_redis + client.set("test:corrupt", "invalid") + script = client.register_script(rate_limit.GCRA_LUA) + with pytest.raises(redis.ResponseError, match="Invalid GCRA state"): + script(keys=["test:corrupt"], args=[200_000, 2]) + assert client.get("test:corrupt") == "invalid" + + +def test_model_quota_exhaustion_does_not_block_another_model(local_redis, monkeypatch): + client, _ = local_redis + monkeypatch.setattr(rate_limit, "_create_redis_client", lambda _: client) + monkeypatch.setattr(rate_limit, "_registry", {}) + config = LLMRateLimitConfig( + enabled=True, + max_wait_seconds=0.01, + rules={ + "gpt-4o-mini": {"qps": 0.001, "burst": 1}, + "other": {"qps": 0.002, "burst": 2}, + }, + ) + primary = rate_limit.get_limiter(config, config.rule_for("gpt-4o-mini"), "gpt-4o-mini") + other = rate_limit.get_limiter(config, config.rule_for("other"), "other") + primary_key = f"{config.key_prefix}:gpt-4o-mini" + other_key = f"{config.key_prefix}:other" + primary.acquire() + before = client.get(primary_key) + with pytest.raises(LLMRateLimitTimeoutError): + primary.acquire() + other.acquire() + other.acquire() + with pytest.raises(LLMRateLimitTimeoutError): + other.acquire() + assert client.get(primary_key) == before + assert client.exists(other_key) == 1 From ea3f2626fa4aa7b1f4d5ba3882ea188b05a5354e Mon Sep 17 00:00:00 2001 From: bittergreen Date: Wed, 9 Sep 2026 20:13:56 +0800 Subject: [PATCH 2/2] test(llm): make retry permit budget test deterministic --- tests/llms/test_qps_rate_limit.py | 32 +++++++++++++++++++++++++------ 1 file changed, 26 insertions(+), 6 deletions(-) diff --git a/tests/llms/test_qps_rate_limit.py b/tests/llms/test_qps_rate_limit.py index 167adb5ea..3b73bb685 100644 --- a/tests/llms/test_qps_rate_limit.py +++ b/tests/llms/test_qps_rate_limit.py @@ -6,6 +6,7 @@ import time from concurrent.futures import ThreadPoolExecutor +from types import SimpleNamespace from unittest.mock import MagicMock import httpx @@ -332,16 +333,35 @@ def test_retry_after_is_honored(): def test_permit_budget_is_shared_between_retries(monkeypatch): + now = 100.0 + + def advance(seconds): + nonlocal now + now += seconds + + # Keep the clock local to the limiter so SDK/logging clocks are unaffected. + backoff = MagicMock(side_effect=advance) + monkeypatch.setattr(rate_limit, "time", SimpleNamespace(monotonic=lambda: now, sleep=backoff)) replies = [ httpx.Response(429, json={"error": {"message": "test"}}), httpx.Response(200, json=response_body()), ] - llm, limiter = make_llm(monkeypatch, lambda _: replies.pop(0)) - limiter.acquire.side_effect = lambda **_: time.sleep(0.002) - assert llm.generate([]) == "ok" - budgets = [call.kwargs["timeout_seconds"] for call in limiter.acquire.call_args_list] - assert 0 < budgets[1] < budgets[0] - llm.client.close() + + def respond(_): + advance(10.0) + return replies.pop(0) + + llm, limiter = make_llm(monkeypatch, respond) + limiter.acquire.side_effect = lambda **_: advance(0.25) + try: + assert llm.generate([]) == "ok" + budgets = [call.kwargs["timeout_seconds"] for call in limiter.acquire.call_args_list] + initial_budget = llm.config.rate_limit.rule_for("gpt-4o-mini").max_wait_seconds + # Deduct permit waiting only, excluding model I/O and retry backoff. + assert budgets == pytest.approx([initial_budget, initial_budget - 0.25]) + backoff.assert_called_once() + finally: + llm.client.close() def test_expired_waiter_does_not_send_when_redis_returns_late(monkeypatch):