diff --git a/volc_limit_sync.py b/volc_limit_sync.py index 7a7ae89..0e2a79c 100644 --- a/volc_limit_sync.py +++ b/volc_limit_sync.py @@ -109,6 +109,35 @@ PREBLOCK_RESERVE_TOKENS = int(os.getenv("PREBLOCK_RESERVE_TOKENS", "50000")) # max(声明值, 本下限) 兜住思考开销;2M 限额下 4096 偏差仅 0.2%。 THINKING_RESERVE_FLOOR = int(os.getenv("THINKING_RESERVE_FLOOR", "4096")) +# 在途预留(inflight)TTL:pre-hook 放行时预扣的 reserve 记入独立 +# daily_usage_inflight key,正常请求在结算回调中配对释放;若请求 +# 失败/进程崩溃/回调丢失,预留量由本 TTL 自动归零,防止额度被永久 +# 卡死。1 小时远大于正常结算延迟(秒级),兜底释放足够及时。 +PREBLOCK_INFLIGHT_TTL = int(os.getenv("PREBLOCK_INFLIGHT_TTL", "3600")) + +# pre-hook 写入请求 metadata 的预留量键名,结算回调据此配对释放 inflight。 +# 以 "_" 前缀避免与用户/代理自定义 metadata 冲突。 +PREBLOCK_RESERVE_META_KEY = "_volc_preblock_reserve" + +# Lua 原子闸门:把"读 usage + 读 inflight + 判 limit + 预扣"合并为 +# 一个 Redis 原子操作,堵住并发请求在同一 usage 快照上集体撞线的 +# TOCTOU 窗口(见 docs/volc-quota-notes.md 并发击穿分析)。 +# KEYS[1]=usage_key, KEYS[2]=inflight_key +# ARGV[1]=effective_limit, ARGV[2]=reserve, ARGV[3]=inflight TTL +# 返回: 1=放行(已预扣), 0=拦截(未预扣) +GATE_SCRIPT = """ +local usage = tonumber(redis.call('GET', KEYS[1]) or '0') +local inflight = tonumber(redis.call('GET', KEYS[2]) or '0') +local limit = tonumber(ARGV[1]) +local reserve = tonumber(ARGV[2]) +if usage + inflight + reserve >= limit then + return 0 +end +redis.call('INCRBY', KEYS[2], reserve) +redis.call('EXPIRE', KEYS[2], tonumber(ARGV[3])) +return 1 +""" + # stable_limit 缺失时(首次同步前/新增模型/TTL 过期/Redis 数据丢失)的保守兜底限额。 # 免费池常态 ~2M(首开赠送/昨日采多少返多少),若 fallback 裸奔到采集上限 2M, # 余额不足或已耗尽的模型会被烧穿扣费。900k 对 0 余额模型仍是 9 倍优惠, @@ -375,6 +404,16 @@ def get_usage_key(clean_model, api_base=None): return f"daily_usage:{date_str}:{clean_model}{suffix}" +def get_inflight_key(clean_model, api_base=None): + """在途预留 key:与 usage key 同日期轮换,独立账户,结算配对释放。 + + pre-hook 放行时 INCRBY reserve(Lua 原子闸门内完成),请求成功结算时 + 按该请求的 reserve DECRBY 释放,usage 只计实际消耗;失败/丢失的预留 + 由 PREBLOCK_INFLIGHT_TTL 自动归零。 + """ + return get_usage_key(clean_model, api_base) + ":inflight" + + def get_static_limit_for_hash(clean_model, current_hash=None): """Check static configuration in HARD_DAILY_LIMITS using an optional endpoint hash.""" if current_hash: @@ -862,13 +901,60 @@ def startup_audit(): log(f"❌ Startup audit failed: {e}", level="ERROR") +def _atomic_reserve(clean_model, api_base, effective_limit, reserve): + """Lua 原子闸门:usage + inflight + reserve >= limit 则拦截,否则预扣并放行。 + + 返回 (granted, usage_total)。granted=False 时 usage_total = usage + inflight + (供日志展示);granted=True 时预扣已入 inflight key。 + Redis 不可用/脚本失败时抛异常由调用方 fail-closed 处理。 + """ + usage_key = get_usage_key(clean_model, api_base) + inflight_key = get_inflight_key(clean_model, api_base) + granted = redis_client.eval( + GATE_SCRIPT, + 2, + usage_key, + inflight_key, + int(float(effective_limit)), + int(reserve), + int(PREBLOCK_INFLIGHT_TTL), + ) + usage = float(redis_client.get(usage_key) or 0) + inflight = float(redis_client.get(inflight_key) or 0) + return bool(granted), usage + inflight + + +def _release_inflight(clean_model, api_base, reserve): + """结算配对释放:请求完成后把 pre-hook 预扣的 reserve 从 inflight 账户扣除。 + + reserve<=0 或 inflight 已过期(TTL 兜底释放)时无操作。失败/丢失的预留 + 由 TTL 自动归零,这里只做正常结算路径的精确回收。 + """ + if not redis_client or not reserve: + return + try: + inflight_key = get_inflight_key(clean_model, api_base) + cur = float(redis_client.get(inflight_key) or 0) + if cur <= 0: + return + # DECRBY 到负数说明 inflight 已部分过期,直接 DEL 归零即可 + new_val = redis_client.decrby(inflight_key, int(reserve)) + if int(new_val) <= 0: + redis_client.delete(inflight_key) + except Exception as e: + log(f"❌ _release_inflight failed: {e}", level="ERROR") + + def _check_request_block(model_name, reserve=0.0): """ 请求前闸门:对外模型名 → 判定是否超限。 只拦截 MONITORED_MODEL_PAIRS 内的火山监控条目;newapi 等非监控渠道天然放行。 - reserve: 本次请求的预估消耗(max_tokens 或常量兜底),usage + reserve >= limit 即拦截, - 避免并发请求在闸门检查后集体撞线(限额是低频快照,用量是实时累计)。 - 返回 (clean_model, usage, limit) 表示超限,None 表示放行。 + reserve: 本次请求的预估消耗(max_tokens 或常量兜底)。 + 判定用 Lua 原子闸门:usage + inflight(在途预留) + reserve >= limit 即拦截, + inflight 账户在 Redis 内原子预扣,并发请求串行看到彼此的在途量,堵死 + "同一 usage 快照集体撞线"的 TOCTOU 击穿。放行时 reserve 已预扣入 inflight, + 由结算回调配对释放(见 _release_inflight)。 + 返回 (clean_model, usage, limit) 表示超限,None 表示放行(已预扣)。 fail-closed:Redis 不可用 / 限额获取失败 / 读取异常时,对监控条目返回 (clean_model, -1, -1) 标记拒绝(宁可错杀,不可漏放);非监控条目不受影响。 手动软踢:Redis key manual_evict:{config model 字段全名} 存在时返回 @@ -975,16 +1061,24 @@ def _check_request_block(model_name, reserve=0.0): return (clean, -1, -1) usage_key = get_usage_key(clean, api_base) - val = redis_client.get(usage_key) - usage = float(val or 0) # 已被软踢的条目恒 429(无论 reserve),避免踢后灰色区间(如 1.95M~2M) # 小请求被放行 → 路由原名不存在 → 400 路由错误 if entry_name == f"{model_name}-limit-reached" and "_daily_limit_backup" in entry: + usage = float(redis_client.get(usage_key) or 0) return (clean, usage, effective_limit) - if usage + float(reserve) >= float(effective_limit): - return (clean, usage, effective_limit) + # Lua 原子闸门:usage + inflight + reserve >= limit 即拒绝(不预扣), + # 否则把 reserve 预扣进 inflight 账户并放行。判定与预扣在 Redis 内 + # 原子完成,并发请求串行看到彼此的在途预留,堵死"同一 usage 快照上 + # 集体撞线"的 TOCTOU 击穿(08-18 13:14:41-50 根因,见诊断)。 + granted, _inflight_total = _atomic_reserve( + clean, api_base, effective_limit, reserve + ) + if not granted: + return (clean, _inflight_total, effective_limit) + # 放行:预扣已入 inflight,结算时配对释放 + return None return None except Exception as e: @@ -1081,7 +1175,19 @@ class DailyLimitHandler(CustomLogger): blocked = _check_request_block(data["model"], reserve=reserve) if not blocked: - return None + # 放行:Lua 原子闸门已把 reserve 预扣进 inflight 账户。 + # 把本次预留量记入请求 metadata,随 litellm 流转到 log_success_event + # 的 kwargs["litellm_params"]["metadata"],结算时配对释放(否则 + # inflight 只会增不减,额度被已放行请求永久占死)。 + try: + meta = data.get("metadata") or {} + if not isinstance(meta, dict): + meta = {} + meta[PREBLOCK_RESERVE_META_KEY] = int(reserve) + data["metadata"] = meta + except Exception: + pass # 写 metadata 失败仅影响结算释放,TTL 兜底仍会回收 + return data clean, usage, limit = blocked if usage < 0 or limit < 0: @@ -1157,6 +1263,20 @@ class DailyLimitHandler(CustomLogger): if clean_model and total_tokens > 0: key = get_usage_key(clean_model, api_base) + + # 结算配对释放:读回 pre-hook 预扣的 reserve(经 metadata 透传), + # 从 inflight 账户扣除。实际消耗 total_tokens 照常记入 usage; + # 预扣量 = 预估值,与实耗差额天然由 usage 侧的"只计实际"修正, + # inflight 只代表"已放行未结算"的在途量,结算即清零。 + released = 0 + try: + meta = litellm_params.get("metadata") or {} + if isinstance(meta, dict): + released = int(meta.get(PREBLOCK_RESERVE_META_KEY, 0) or 0) + except (TypeError, ValueError): + released = 0 + _release_inflight(clean_model, api_base, released) + current_total = redis_client.incrby(key, int(total_tokens)) redis_client.expire(key, 172800)