feat(proxy): T-X12 采纳 enterprise-ai-gateway 双桶限流——TPM token 桶 + 语义缓存命中率透出
- auth.RateLimiter.allow_tokens:TPM 令牌桶(容量=tpm_cap,连续回流 tpm/60 每秒, 与 RPM 桶共用 per-key 锁、判定相互独立;tpm_cap<=0 不限) - ProxyConfig.limits.tpm_per_key(默认 60000,0=不限);chat_completions 在并发槽 之后做 TPM 预扣判定(est=字符/3 与计费同口径),拒绝时释放并发槽并回 429 (文案区分并发/RPM/TPM 三种原因) - semcache.stats 增 misses 与 hit_rate;/admin/stats 新增 semcache 段透出 (entries/hits_exact/hits_semantic/misses/hit_rate,与账本 h_g 互补) - 新增 tests/test_tpm_limiter.py 6 项(容量/回流/禁用/双桶独立/key 隔离/命中率)
This commit is contained in:
@@ -89,6 +89,7 @@ class RateLimiter:
|
|||||||
|
|
||||||
def __init__(self, concurrent_per_key: int = 2):
|
def __init__(self, concurrent_per_key: int = 2):
|
||||||
self._tokens: Dict[int, tuple[float, float]] = {} # key_id -> (tokens, last_ts)
|
self._tokens: Dict[int, tuple[float, float]] = {} # key_id -> (tokens, last_ts)
|
||||||
|
self._tpm: Dict[int, tuple[float, float]] = {} # key_id -> (tokens, last_ts)
|
||||||
self._locks: Dict[int, threading.Lock] = {}
|
self._locks: Dict[int, threading.Lock] = {}
|
||||||
self._sems: Dict[int, threading.Semaphore] = {}
|
self._sems: Dict[int, threading.Semaphore] = {}
|
||||||
self._global = threading.Lock()
|
self._global = threading.Lock()
|
||||||
@@ -108,6 +109,27 @@ class RateLimiter:
|
|||||||
self._tokens[key_id] = (tokens - 1.0, now)
|
self._tokens[key_id] = (tokens - 1.0, now)
|
||||||
return True
|
return True
|
||||||
|
|
||||||
|
def allow_tokens(self, key_id: int, est_tokens: int, tpm_cap: int) -> bool:
|
||||||
|
"""TPM 令牌桶放行判定(T-X12,采纳 enterprise-ai-gateway 双桶设计)。
|
||||||
|
|
||||||
|
容量 = tpm_cap,回流速率 = tpm/60 每秒(连续时间回流补充);
|
||||||
|
est_tokens 为本次预扣估算(字符/3,与预扣计费同口径)。
|
||||||
|
tpm_cap <= 0 视为不限流;与 RPM 桶共用 per-key 锁,判定相互独立。
|
||||||
|
"""
|
||||||
|
if tpm_cap <= 0:
|
||||||
|
return True
|
||||||
|
now = time.monotonic()
|
||||||
|
with self._global:
|
||||||
|
lock = self._locks.setdefault(key_id, threading.Lock())
|
||||||
|
with lock:
|
||||||
|
tokens, last = self._tpm.get(key_id, (float(tpm_cap), now))
|
||||||
|
tokens = min(float(tpm_cap), tokens + (now - last) * (tpm_cap / 60.0))
|
||||||
|
if tokens < est_tokens:
|
||||||
|
self._tpm[key_id] = (tokens, now)
|
||||||
|
return False
|
||||||
|
self._tpm[key_id] = (tokens - est_tokens, now)
|
||||||
|
return True
|
||||||
|
|
||||||
def acquire_slot(self, key_id: int) -> bool:
|
def acquire_slot(self, key_id: int) -> bool:
|
||||||
"""并发槽(非阻塞);返回 False = 超并发上限(429)。"""
|
"""并发槽(非阻塞);返回 False = 超并发上限(429)。"""
|
||||||
with self._global:
|
with self._global:
|
||||||
|
|||||||
@@ -52,6 +52,7 @@ class ProxyConfig:
|
|||||||
rpm_per_key: int = 10
|
rpm_per_key: int = 10
|
||||||
day_req_cap: int = 200
|
day_req_cap: int = 200
|
||||||
concurrent_per_key: int = 2
|
concurrent_per_key: int = 2
|
||||||
|
tpm_per_key: int = 60000 # T-X12:每分钟 token 桶(0 = 不限)
|
||||||
max_body_chars: int = 60000
|
max_body_chars: int = 60000
|
||||||
semcache_enabled: bool = True
|
semcache_enabled: bool = True
|
||||||
sim_threshold: float = 0.92
|
sim_threshold: float = 0.92
|
||||||
@@ -159,6 +160,7 @@ def build_proxy_config(settings_dict: Dict[str, Any]) -> ProxyConfig:
|
|||||||
rpm_per_key=int(limits.get("rpm_per_key", 10) or 10),
|
rpm_per_key=int(limits.get("rpm_per_key", 10) or 10),
|
||||||
day_req_cap=int(limits.get("day_req_cap", 200) or 200),
|
day_req_cap=int(limits.get("day_req_cap", 200) or 200),
|
||||||
concurrent_per_key=int(limits.get("concurrent_per_key", 2) or 2),
|
concurrent_per_key=int(limits.get("concurrent_per_key", 2) or 2),
|
||||||
|
tpm_per_key=int(limits.get("tpm_per_key", 60000) or 0),
|
||||||
max_body_chars=int(limits.get("max_body_chars", 60000) or 60000),
|
max_body_chars=int(limits.get("max_body_chars", 60000) or 60000),
|
||||||
semcache_enabled=bool(sem.get("enabled", True)),
|
semcache_enabled=bool(sem.get("enabled", True)),
|
||||||
sim_threshold=float(sem.get("sim_threshold", 0.92) or 0.92),
|
sim_threshold=float(sem.get("sim_threshold", 0.92) or 0.92),
|
||||||
|
|||||||
@@ -86,6 +86,14 @@ def build_proxy_router(cfg: ProxyConfig, pool, settings_provider=None,
|
|||||||
_limits(), time.time())
|
_limits(), time.time())
|
||||||
if not _limits().acquire_slot(ctx["key_id"]):
|
if not _limits().acquire_slot(ctx["key_id"]):
|
||||||
raise QuotaError("并发请求已达该 key 上限")
|
raise QuotaError("并发请求已达该 key 上限")
|
||||||
|
# T-X12(采纳 enterprise-ai-gateway 双桶限流):TPM 桶——RPM 之外加
|
||||||
|
# token 维度(预扣估算 = 字符/3,与计费同口径;cfg.tpm_per_key=0 不限)
|
||||||
|
tpm_cap = int(getattr(cfg, "tpm_per_key", 0) or 0)
|
||||||
|
if tpm_cap > 0:
|
||||||
|
est_tokens = max(1, len(raw_body) // 3)
|
||||||
|
if not _limits().allow_tokens(ctx["key_id"], est_tokens, tpm_cap):
|
||||||
|
_limits().release_slot(ctx["key_id"])
|
||||||
|
raise QuotaError("该 key 的每分钟 token 额度(TPM)已用尽")
|
||||||
|
|
||||||
# 语义分析器 live 分流(D-G7:mode=live 才启用;任何异常不影响代理可用性)。
|
# 语义分析器 live 分流(D-G7:mode=live 才启用;任何异常不影响代理可用性)。
|
||||||
# T-X4 修复:settings_provider 由 api.py 注入(原实现引用未导入的
|
# T-X4 修复:settings_provider 由 api.py 注入(原实现引用未导入的
|
||||||
@@ -217,10 +225,18 @@ def build_proxy_router(cfg: ProxyConfig, pool, settings_provider=None,
|
|||||||
in_hit = s["in_hit_tok"]
|
in_hit = s["in_hit_tok"]
|
||||||
in_miss = s["in_miss_tok"]
|
in_miss = s["in_miss_tok"]
|
||||||
h_p = in_hit / (in_hit + in_miss) if (in_hit + in_miss) else 0.0
|
h_p = in_hit / (in_hit + in_miss) if (in_hit + in_miss) else 0.0
|
||||||
|
# T-X12:进程内语义缓存计数透出(entries/hits/misses/hit_rate)
|
||||||
|
sem_stats = None
|
||||||
|
if cfg.semcache_enabled:
|
||||||
|
try:
|
||||||
|
sem_stats = _get_semcache(cfg, ledger).stats()
|
||||||
|
except Exception: # noqa: BLE001
|
||||||
|
sem_stats = None
|
||||||
return {"requests": n, "h_g": round(h_g, 4), "h_p": round(h_p, 4),
|
return {"requests": n, "h_g": round(h_g, 4), "h_p": round(h_p, 4),
|
||||||
"revenue_milli": s["revenue_milli"], "cost_milli": s["cost_milli"],
|
"revenue_milli": s["revenue_milli"], "cost_milli": s["cost_milli"],
|
||||||
"margin_milli": s["revenue_milli"] - s["cost_milli"],
|
"margin_milli": s["revenue_milli"] - s["cost_milli"],
|
||||||
"by_bucket": s["by_bucket"],
|
"by_bucket": s["by_bucket"],
|
||||||
|
"semcache": sem_stats,
|
||||||
"today": today,
|
"today": today,
|
||||||
"upstream_failover": upstream_failover_stats()}
|
"upstream_failover": upstream_failover_stats()}
|
||||||
|
|
||||||
|
|||||||
@@ -104,6 +104,7 @@ class SemanticCache:
|
|||||||
self._clock = now # 假时钟注入(None = time.time)
|
self._clock = now # 假时钟注入(None = time.time)
|
||||||
self.hits_exact = 0
|
self.hits_exact = 0
|
||||||
self.hits_semantic = 0
|
self.hits_semantic = 0
|
||||||
|
self.misses = 0
|
||||||
self._rebuild()
|
self._rebuild()
|
||||||
|
|
||||||
# ---------- 时钟 ----------
|
# ---------- 时钟 ----------
|
||||||
@@ -166,6 +167,7 @@ class SemanticCache:
|
|||||||
# L2
|
# L2
|
||||||
g = grams(norm_text)
|
g = grams(norm_text)
|
||||||
if len(g) < 3:
|
if len(g) < 3:
|
||||||
|
self.misses += 1
|
||||||
return None
|
return None
|
||||||
w_q = _weight(g)
|
w_q = _weight(g)
|
||||||
candidates: Dict[str, int] = {}
|
candidates: Dict[str, int] = {}
|
||||||
@@ -195,6 +197,7 @@ class SemanticCache:
|
|||||||
best_score = score
|
best_score = score
|
||||||
best_key = key
|
best_key = key
|
||||||
if best_key is None or best_score < self.sim_threshold:
|
if best_key is None or best_score < self.sim_threshold:
|
||||||
|
self.misses += 1
|
||||||
return None
|
return None
|
||||||
cand = self._l1[best_key]
|
cand = self._l1[best_key]
|
||||||
cand.hits += 1
|
cand.hits += 1
|
||||||
@@ -244,8 +247,11 @@ class SemanticCache:
|
|||||||
pass # 持久化失败不影响内存缓存(可重建)
|
pass # 持久化失败不影响内存缓存(可重建)
|
||||||
|
|
||||||
def stats(self) -> Dict[str, Any]:
|
def stats(self) -> Dict[str, Any]:
|
||||||
|
total = self.hits_exact + self.hits_semantic + self.misses
|
||||||
return {"entries": len(self._l1), "hits_exact": self.hits_exact,
|
return {"entries": len(self._l1), "hits_exact": self.hits_exact,
|
||||||
"hits_semantic": self.hits_semantic}
|
"hits_semantic": self.hits_semantic, "misses": self.misses,
|
||||||
|
"hit_rate": round((self.hits_exact + self.hits_semantic) / total, 4)
|
||||||
|
if total else 0.0}
|
||||||
|
|
||||||
|
|
||||||
class SingleFlight:
|
class SingleFlight:
|
||||||
|
|||||||
@@ -0,0 +1,70 @@
|
|||||||
|
"""TPM 双桶限流与缓存命中率透出(T-X12,采纳 enterprise-ai-gateway 设计)。"""
|
||||||
|
import time
|
||||||
|
|
||||||
|
from gateway.proxy.auth import RateLimiter
|
||||||
|
from gateway.proxy.semcache import SemanticCache
|
||||||
|
|
||||||
|
|
||||||
|
def test_tpm_bucket_capacity_and_consume():
|
||||||
|
"""容量内放行并扣减;超过余量拒绝。"""
|
||||||
|
rl = RateLimiter()
|
||||||
|
assert rl.allow_tokens(1, 100, 1000) is True
|
||||||
|
assert rl.allow_tokens(1, 500, 1000) is True
|
||||||
|
# 已扣 600,剩 400 + 回流(瞬时近似 0):700 > 400+eps -> 拒
|
||||||
|
assert rl.allow_tokens(1, 700, 1000) is False
|
||||||
|
|
||||||
|
|
||||||
|
def test_tpm_bucket_continuous_refill(monkeypatch):
|
||||||
|
"""连续时间回流:elapsed * tpm/60 补充(容量封顶)。"""
|
||||||
|
rl = RateLimiter()
|
||||||
|
clock = {"t": 100.0}
|
||||||
|
monkeypatch.setattr(time, "monotonic", lambda: clock["t"])
|
||||||
|
assert rl.allow_tokens(2, 1000, 2000) is True # 剩 1000
|
||||||
|
clock["t"] = 130.0 # 30s 回流 1000 -> 满 2000
|
||||||
|
assert rl.allow_tokens(2, 1500, 2000) is True # 扣 1500 剩 500
|
||||||
|
clock["t"] = 145.0 # 15s 回流 500 -> 1000
|
||||||
|
assert rl.allow_tokens(2, 1000, 2000) is True # 精确扣空
|
||||||
|
assert rl.allow_tokens(2, 1, 2000) is False # 空桶拒绝
|
||||||
|
|
||||||
|
|
||||||
|
def test_tpm_disabled_when_cap_zero():
|
||||||
|
"""tpm_cap <= 0 = 不限流(恒放行)。"""
|
||||||
|
rl = RateLimiter()
|
||||||
|
assert rl.allow_tokens(3, 10**9, 0) is True
|
||||||
|
|
||||||
|
|
||||||
|
def test_tpm_independent_from_rpm_bucket():
|
||||||
|
"""双桶相互独立:打满 TPM 不影响 RPM 放行,RPM 打空不影响 TPM 记账。"""
|
||||||
|
rl = RateLimiter()
|
||||||
|
assert rl.allow_tokens(4, 5000, 5000) is True # TPM 打空
|
||||||
|
assert rl.allow(4, 10) is True # RPM 仍放行
|
||||||
|
for _ in range(9): # 共 10 次 = 打满 RPM
|
||||||
|
assert rl.allow(4, 10) is True
|
||||||
|
assert rl.allow(4, 10) is False # RPM 空桶
|
||||||
|
assert rl.allow_tokens(4, 1, 5000) is False # TPM 仍空(未被 RPM 动过)
|
||||||
|
|
||||||
|
|
||||||
|
def test_tpm_per_key_isolation():
|
||||||
|
"""按 key 隔离:一个 key 打满不影响另一个。"""
|
||||||
|
rl = RateLimiter()
|
||||||
|
assert rl.allow_tokens(5, 5000, 5000) is True
|
||||||
|
assert rl.allow_tokens(6, 5000, 5000) is True
|
||||||
|
assert rl.allow_tokens(5, 5000, 5000) is False
|
||||||
|
|
||||||
|
|
||||||
|
def test_semcache_stats_hit_rate():
|
||||||
|
"""stats 透出 misses 与 hit_rate(命中率口径 = hits / (hits + misses))。"""
|
||||||
|
class _S:
|
||||||
|
def semcache_rows(self, limit=0):
|
||||||
|
return []
|
||||||
|
|
||||||
|
def put_semcache(self, *a, **k):
|
||||||
|
pass
|
||||||
|
|
||||||
|
c = SemanticCache(_S(), max_entries=10, sim_threshold=0.5, promote_frequency=5)
|
||||||
|
c.put("k1", "请解释一下递归函数的概念", "答案", "m", doc_version=1)
|
||||||
|
assert c.lookup("k1", "请解释一下递归函数的概念", doc_version=1) is not None
|
||||||
|
assert c.lookup("k9", "完全无关的问题内容呢", doc_version=1) is None
|
||||||
|
s = c.stats()
|
||||||
|
assert s["hits_exact"] == 1 and s["misses"] == 1
|
||||||
|
assert s["hit_rate"] == 0.5
|
||||||
Reference in New Issue
Block a user