vibe-code
23778a8f7c
核心改动:不再用 BUFFER 猜滞后窗口消耗,而是用同步时点快照 直接计算累计限额。 - BUFFER=300000 → SMALL_BUFFER=50000(仅覆盖非消耗性跳降) - get_effective_limit: limit = B0 - SMALL_BUFFER,不依赖 current_usage - 同步循环:写 balance_at_sync + usage_at_sync 替代 remaining/stable_limit - 突发检测/低额度判断:读 balance_at_sync 替代 remaining - volc_consume.py: vls.BUFFER → vls.SMALL_BUFFER - debug/debug_volc_sync.py: 同步更新 - docs/volc-quota-notes.md: 新增 §7.6 v6 设计文档,更新各节口径
1551 lines
59 KiB
Python
1551 lines
59 KiB
Python
import os
|
||
import sys
|
||
import time
|
||
import json
|
||
import hmac
|
||
import math
|
||
import random
|
||
import hashlib
|
||
import datetime
|
||
import requests
|
||
import redis
|
||
import traceback
|
||
import threading
|
||
from urllib.parse import urlparse
|
||
|
||
# ==========================================
|
||
# 📦 LiteLLM Imports (Safe Handling)
|
||
# ==========================================
|
||
try:
|
||
import litellm
|
||
from litellm.integrations.custom_logger import CustomLogger
|
||
from litellm.exceptions import (
|
||
ServiceUnavailableError,
|
||
ContextWindowExceededError,
|
||
RateLimitError,
|
||
)
|
||
|
||
LITELLM_AVAILABLE = True
|
||
except ImportError:
|
||
LITELLM_AVAILABLE = False
|
||
|
||
# Define dummy classes to prevent errors if run standalone without litellm
|
||
class CustomLogger:
|
||
pass
|
||
|
||
class ServiceUnavailableError(Exception):
|
||
pass
|
||
|
||
class ContextWindowExceededError(Exception):
|
||
pass
|
||
|
||
class RateLimitError(Exception):
|
||
pass
|
||
|
||
|
||
# ==========================================
|
||
# 🏷️ Unified Log Tag (统一日志标识，便于过滤)
|
||
# ==========================================
|
||
# 本脚本所有输出统一带 [VolcSync] 前缀，方便日志筛选：
|
||
# docker logs litellm 2>&1 | grep VolcSync
|
||
# 实现方式：模块级重绑定 print（`print` 在模块命名空间中按调用时查找），
|
||
# 一处覆盖 log()/dprint 与全部裸 print 调用；对模块外部无任何影响。
|
||
LOG_TAG = "[VolcSync]"
|
||
|
||
_original_print = print
|
||
|
||
|
||
def _tagged_print(*args, **kwargs):
|
||
# 统一北京时间戳（显式 UTC+8，与业务口径一致，不依赖容器 TZ）。
|
||
# 性能：µs 级/次，相对打印本身（含 flush）可忽略。
|
||
ts = time.strftime("%m-%d %H:%M:%S", time.gmtime(time.time() + 8 * 3600))
|
||
if args and isinstance(args[0], str):
|
||
args = (args[0].lstrip("\n"), *args[1:])
|
||
_original_print(f"[{ts}] {LOG_TAG}", *args, **kwargs)
|
||
|
||
|
||
print = _tagged_print
|
||
|
||
|
||
# ==========================================
|
||
# 🛠️ Configuration & Constants
|
||
# ==========================================
|
||
# 日志级别配置
|
||
LOG_LEVEL = os.getenv("DAILY_LIMIT_LOG_LEVEL", "INFO").upper()
|
||
|
||
|
||
def get_log_level_value(level_name):
|
||
levels = {"DEBUG": 10, "INFO": 20, "WARNING": 30, "ERROR": 40}
|
||
return levels.get(level_name, 20)
|
||
|
||
|
||
CURRENT_LEVEL_VAL = get_log_level_value(LOG_LEVEL)
|
||
DEBUG_MODE = CURRENT_LEVEL_VAL <= 10
|
||
|
||
# Volcengine API Config
|
||
ACCESS_KEY = os.getenv(
|
||
"VOLC_ACCESS_KEY", "AKLTNGYxZmUyNzNmZjk3NGMyZmEzNjVjZmY4NGM5ZjcxMjU"
|
||
)
|
||
SECRET_KEY = os.getenv(
|
||
"VOLC_SECRET_KEY", "Tm1VNE5XSmlORGs1WXpGbU5EVmtabUUyT1RZeVltTTBNalkxTWpjeE1tTQ=="
|
||
)
|
||
HOST = "open.volcengineapi.com"
|
||
REGION = "cn-beijing"
|
||
|
||
# Daemon Config
|
||
SMALL_BUFFER = 50000 # 覆盖资源包过期/API 未结算等非消耗性余额跳降
|
||
CHECK_INTERVAL = 10
|
||
|
||
# 前置闸门预留量：请求开始判定 usage + reserve >= limit 即 429。
|
||
# reserve 优先取本次请求自带 max_tokens（真实上限，零采集损失），
|
||
# 未设置时 fallback 本常量。堵住"并发请求穿过闸门后集体撞线"的窗口。
|
||
PREBLOCK_RESERVE_TOKENS = int(os.getenv("PREBLOCK_RESERVE_TOKENS", "50000"))
|
||
|
||
# stable_limit 缺失时（首次同步前/新增模型/TTL 过期）的保守兜底限额。
|
||
# 免费模型余额常态 < 2M（昨日采多少返多少），若 fallback 裸奔到采集上限 2M，
|
||
# 余额不足的模型会被烧穿扣费。宁可短时收紧，不可裸奔。
|
||
FALLBACK_LIMIT = int(os.getenv("FALLBACK_LIMIT", "900000"))
|
||
|
||
# Redis Config
|
||
REDIS_HOST = os.getenv(
|
||
"REDIS_HOST", "redis"
|
||
) # Default to 'redis' (docker service name)
|
||
REDIS_PORT = int(os.getenv("REDIS_PORT", 6379))
|
||
REDIS_PASSWORD = os.getenv("REDIS_PASSWORD", None)
|
||
|
||
# Suffixes to clean from Volcengine response
|
||
SUFFIXES_TO_REMOVE = [
|
||
"_pack_data_collaboration",
|
||
"pack_data_collaboration",
|
||
"_data_collaboration_resource_pack",
|
||
"_data_collaboration",
|
||
"data_collaboration",
|
||
"协作奖励计划资源包",
|
||
"免费在线推理资源包",
|
||
"免费在线推理包",
|
||
"资源包",
|
||
"_free_infer_res_pack",
|
||
"_pack_free_infer",
|
||
"_free_inference_resource_pack",
|
||
"free_inference_resource_pack",
|
||
"_infer_res_pack",
|
||
"-free-infer-res-pack", # Matches: DeepSeek-V3-free-infer-res-pack
|
||
]
|
||
|
||
# Manual package-name aliases for monitored models.
|
||
# 常规版本号模型（如 xxx-260215）会根据 config.yaml 自动推导成 xxx，
|
||
# 这里只保留特殊别名或一对多映射，减少新增/删除监控模型时需要手改的地方。
|
||
MANUAL_MONITORED_MODEL_ALIASES = {
|
||
"doubao-seed-1-6": ["doubao-seed-1-6-250615", "doubao-seed-1-6-251015"],
|
||
"doubao-seed-code": "doubao-seed-code-preview-251028",
|
||
"doubao-seed-translation": "doubao-seed-translation-250915",
|
||
"doubao-seed-2-0-code": "doubao-seed-2-0-code-preview-260215",
|
||
# 火山资源包名用 MMDD（deepseek-v4-flash-0731），上游模型名用 -ga-YYMMDD（-ga-260731），
|
||
# 两者不一致，必须手工映射，否则资源包会被 "Not in monitored models" 跳过。
|
||
"deepseek-v4-flash-0731": "deepseek-v4-flash-ga-260731",
|
||
}
|
||
|
||
# Auto-discovered monitored upstream models and endpoint pairs from config.yaml
|
||
MONITORED_MODEL_PAIRS = set()
|
||
MONITORED_MODEL_ALIASES = {}
|
||
|
||
# ==========================================
|
||
# 🛑 Hard Daily Limits (Local Override / 本地每日硬限额)
|
||
# ==========================================
|
||
# Configured directly here to bypass API latency.
|
||
# These limits are applied LOCALLY and update immediately.
|
||
# 这里的限额是本地强制执行的，不受API延迟影响，用于防止超支。
|
||
#
|
||
# Format: "clean-model-name": limit_in_tokens
|
||
# Or with hash: "clean-model-name:hash": limit_in_tokens
|
||
HARD_DAILY_LIMITS = {
|
||
# 火山免费额度规则（2026-08）：每模型每日采集上限 2M，次日按采集量发放资源包。
|
||
# 请求前闸门（async_pre_call_hook）已不滞后 → 无需打折（旧 1.8M/480 万为 post-hoc 时代遗产）。
|
||
# 统一 2M，余额相关的动态限额由 balance_at_sync - (usage_now - usage_at_sync) 实时。
|
||
"__DEFAULT__": 2_000_000,
|
||
}
|
||
|
||
# ==========================================
|
||
# 🔌 Redis Connection
|
||
# ==========================================
|
||
try:
|
||
redis_client = redis.Redis(
|
||
host=REDIS_HOST, port=REDIS_PORT, password=REDIS_PASSWORD, decode_responses=True
|
||
)
|
||
# Check connection only if running as script, to avoid noise on import
|
||
# redis_client.ping()
|
||
except Exception as e:
|
||
print(f"❌ Redis connection failed: {str(e)}", flush=True)
|
||
redis_client = None
|
||
|
||
|
||
# ==========================================
|
||
# 📝 Logging Helpers
|
||
# ==========================================
|
||
def log(msg, level="INFO"):
|
||
msg_val = get_log_level_value(level)
|
||
if msg_val >= CURRENT_LEVEL_VAL:
|
||
icon = {"DEBUG": "🔍", "INFO": "ℹ️", "WARNING": "⚠️", "ERROR": "❌"}.get(
|
||
level, "📝"
|
||
)
|
||
print(f"{icon} {msg}", flush=True)
|
||
|
||
|
||
def dprint(msg):
|
||
log(msg, level="INFO")
|
||
|
||
|
||
# ==========================================
|
||
# 🧮 Utility Functions (Hash, Clean, Sign)
|
||
# ==========================================
|
||
def generate_api_base_hash(api_base):
|
||
"""Generate robust hash from api_base ignoring path variations."""
|
||
if not api_base:
|
||
return "default"
|
||
try:
|
||
s = api_base.strip()
|
||
parsed = urlparse(s)
|
||
if parsed.scheme and parsed.netloc:
|
||
clean_base = f"{parsed.scheme}://{parsed.netloc}".lower()
|
||
else:
|
||
clean_base = s.rstrip("/").lower()
|
||
except Exception:
|
||
clean_base = api_base.strip().rstrip("/").lower()
|
||
|
||
h = hashlib.md5(clean_base.encode()).hexdigest()[:8]
|
||
return h
|
||
|
||
|
||
def clean_model_name(name):
|
||
"""LiteLLM style cleaning: remove provider prefixes."""
|
||
if not name:
|
||
return ""
|
||
return (
|
||
name.replace("openai/", "")
|
||
.replace("azure/", "")
|
||
.replace("anthropic/", "")
|
||
.replace("volcengine/", "")
|
||
)
|
||
|
||
|
||
def strip_model_version_suffix(model_name):
|
||
"""Convert foo-bar-260215 -> foo-bar for common versioned upstream names."""
|
||
if not model_name:
|
||
return ""
|
||
|
||
parts = model_name.rsplit("-", 1)
|
||
if len(parts) == 2 and len(parts[1]) == 6 and parts[1].isdigit():
|
||
return parts[0]
|
||
|
||
return model_name
|
||
|
||
|
||
def build_monitored_model_aliases(monitored_models):
|
||
"""
|
||
Build alias -> monitored upstream model(s) mapping.
|
||
Common versioned names are derived automatically from config.yaml.
|
||
"""
|
||
alias_map = {}
|
||
|
||
def add_alias(alias_name, target_model):
|
||
clean_alias = clean_model_name(alias_name)
|
||
clean_target = clean_model_name(target_model)
|
||
if not clean_alias or not clean_target:
|
||
return
|
||
alias_map.setdefault(clean_alias, [])
|
||
if clean_target not in alias_map[clean_alias]:
|
||
alias_map[clean_alias].append(clean_target)
|
||
|
||
for target_model in sorted(monitored_models):
|
||
add_alias(target_model, target_model)
|
||
|
||
base_alias = strip_model_version_suffix(target_model)
|
||
if base_alias != target_model:
|
||
add_alias(base_alias, target_model)
|
||
|
||
for alias_name, targets in MANUAL_MONITORED_MODEL_ALIASES.items():
|
||
if isinstance(targets, list):
|
||
for target_model in targets:
|
||
add_alias(alias_name, target_model)
|
||
else:
|
||
add_alias(alias_name, targets)
|
||
|
||
return alias_map
|
||
|
||
|
||
def normalize_volc_name(raw_name):
|
||
"""Normalize Volcengine package names before monitored-model resolution."""
|
||
if not raw_name:
|
||
return "unknown"
|
||
clean_name = raw_name
|
||
for suffix in sorted(SUFFIXES_TO_REMOVE, key=len, reverse=True):
|
||
if clean_name.endswith(suffix):
|
||
clean_name = clean_name.replace(suffix, "")
|
||
break
|
||
clean_name = clean_name.replace("_", "-").replace(".", "-")
|
||
clean_name = clean_name.lower().strip()
|
||
return clean_name
|
||
|
||
|
||
def sign(key, msg):
|
||
return hmac.new(key, msg.encode("utf-8"), hashlib.sha256).digest()
|
||
|
||
|
||
def get_signature_key(key, date_stamp, region_name, service_name):
|
||
k_date = sign(key.encode("utf-8"), date_stamp)
|
||
k_region = sign(k_date, region_name)
|
||
k_service = sign(k_region, service_name)
|
||
k_signing = sign(k_service, "request")
|
||
return k_signing
|
||
|
||
|
||
def get_signed_headers(action, version, body_dict, service="billing"):
|
||
t = datetime.datetime.utcnow()
|
||
amz_date = t.strftime("%Y%m%dT%H%M%SZ")
|
||
date_stamp = t.strftime("%Y%m%d")
|
||
request_body = json.dumps(body_dict)
|
||
payload_hash = hashlib.sha256(request_body.encode("utf-8")).hexdigest()
|
||
|
||
canonical_headers = f"content-type:application/json\nhost:{HOST}\nx-content-sha256:{payload_hash}\nx-date:{amz_date}\n"
|
||
signed_headers = "content-type;host;x-content-sha256;x-date"
|
||
canonical_request = f"POST\n/\nAction={action}&Version={version}\n{canonical_headers}\n{signed_headers}\n{payload_hash}"
|
||
|
||
credential_scope = f"{date_stamp}/{REGION}/{service}/request"
|
||
string_to_sign = f"HMAC-SHA256\n{amz_date}\n{credential_scope}\n{hashlib.sha256(canonical_request.encode('utf-8')).hexdigest()}"
|
||
|
||
signing_key = get_signature_key(SECRET_KEY, date_stamp, REGION, service)
|
||
signature = hmac.new(
|
||
signing_key, string_to_sign.encode("utf-8"), hashlib.sha256
|
||
).hexdigest()
|
||
|
||
authorization_header = f"HMAC-SHA256 Credential={ACCESS_KEY}/{credential_scope}, SignedHeaders={signed_headers}, Signature={signature}"
|
||
|
||
return {
|
||
"Content-Type": "application/json",
|
||
"Host": HOST,
|
||
"X-Date": amz_date,
|
||
"X-Content-Sha256": payload_hash,
|
||
"Authorization": authorization_header,
|
||
}, request_body
|
||
|
||
|
||
# ==========================================
|
||
# 🎯 Limit Logic (Unified)
|
||
# ==========================================
|
||
def get_usage_key(clean_model, api_base=None):
|
||
# Enforce Beijing Time (UTC+8) for daily key rotation at 00:00 CST
|
||
# Beijing 00:00 = UTC 16:00 (Previous Day) -> simple offset: UTC_timestamp + 8 hours
|
||
now = time.time()
|
||
beijing_time_ts = now + (8 * 3600)
|
||
date_str = time.strftime("%Y-%m-%d", time.gmtime(beijing_time_ts))
|
||
|
||
suffix = f":{generate_api_base_hash(api_base)}" if api_base else ""
|
||
return f"daily_usage:{date_str}:{clean_model}{suffix}"
|
||
|
||
|
||
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:
|
||
specific_key = f"{clean_model}:{current_hash}"
|
||
if specific_key in HARD_DAILY_LIMITS:
|
||
return HARD_DAILY_LIMITS[specific_key]
|
||
|
||
# 2. Exact Model Name
|
||
if clean_model in HARD_DAILY_LIMITS:
|
||
return HARD_DAILY_LIMITS[clean_model]
|
||
|
||
# 3. Fuzzy Match
|
||
for quota_model, limit_val in HARD_DAILY_LIMITS.items():
|
||
if quota_model == "__DEFAULT__":
|
||
continue
|
||
if ":" not in quota_model and quota_model in clean_model:
|
||
return limit_val
|
||
|
||
# 4. Universal Default Fallback
|
||
if "__DEFAULT__" in HARD_DAILY_LIMITS:
|
||
return HARD_DAILY_LIMITS["__DEFAULT__"]
|
||
|
||
return None
|
||
|
||
|
||
def get_static_limit(clean_model, api_base):
|
||
"""Check static configuration in HARD_DAILY_LIMITS."""
|
||
current_hash = generate_api_base_hash(api_base) if api_base else None
|
||
return get_static_limit_for_hash(clean_model, current_hash)
|
||
|
||
|
||
def get_monitored_static_limit(clean_model, model_hashes=None):
|
||
"""
|
||
Resolve the static limit for monitored-model debug output.
|
||
If multiple endpoint hashes exist, show the strictest applicable static cap.
|
||
"""
|
||
specific_limits = []
|
||
|
||
for current_hash in model_hashes or []:
|
||
limit_val = get_static_limit_for_hash(clean_model, current_hash)
|
||
if limit_val is not None:
|
||
specific_limits.append(float(limit_val))
|
||
|
||
if specific_limits:
|
||
return min(specific_limits)
|
||
|
||
return get_static_limit_for_hash(clean_model, None)
|
||
|
||
|
||
def calculate_effective_limit(static_limit, dynamic_limit):
|
||
"""Return the effective cap after combining static and dynamic limits."""
|
||
try:
|
||
if dynamic_limit is None:
|
||
return float(static_limit) if static_limit is not None else None
|
||
|
||
dynamic_limit = float(dynamic_limit)
|
||
if static_limit is None:
|
||
return dynamic_limit
|
||
|
||
return min(float(static_limit), dynamic_limit)
|
||
except (ValueError, TypeError):
|
||
return float(static_limit) if static_limit is not None else None
|
||
|
||
|
||
def format_limit_debug_fields(static_limit=None, balance_at_sync=None, usage_at_sync=None, effective_limit=None):
|
||
"""Render explicit limit labels for debug logs."""
|
||
fields = []
|
||
if static_limit is not None:
|
||
fields.append(f"StaticLimit={int(float(static_limit)):,}")
|
||
if balance_at_sync is not None:
|
||
fields.append(f"BalAtSync={int(float(balance_at_sync)):,}")
|
||
if usage_at_sync is not None:
|
||
fields.append(f"UsageAtSync={int(float(usage_at_sync)):,}")
|
||
if effective_limit is not None:
|
||
fields.append(f"EffectiveLimit={int(float(effective_limit)):,}")
|
||
return " | ".join(fields) if fields else "EffectiveLimit=unknown"
|
||
|
||
|
||
def get_effective_limit(clean_model, api_base, current_usage=None):
|
||
"""
|
||
Calculate effective limit using the Snapshot-Offset strategy.
|
||
limit = balance_at_sync + usage_at_sync - SMALL_BUFFER (= B0 - SMALL_BUFFER)
|
||
_check_request_block 用 current_usage + reserve >= limit 撞线,语义正确。
|
||
current_usage 参数仅用于无同步数据的 fallback 路径,非公式必需。
|
||
"""
|
||
if not redis_client:
|
||
return get_static_limit(clean_model, api_base)
|
||
|
||
# 1. Get Static Limit (Hard Config)
|
||
static_limit = get_static_limit(clean_model, api_base)
|
||
|
||
# 2. Get Dynamic Limit from Snapshot
|
||
# balance_at_sync + usage_at_sync = 当日初始免费池 B0(常数)
|
||
# 累计限额 = B0 - SMALL_BUFFER,两次同步间不变
|
||
balance_key = f"limit:dynamic:volcengine:balance_at_sync:{clean_model}"
|
||
usage_key = f"limit:dynamic:volcengine:usage_at_sync:{clean_model}"
|
||
|
||
balance_val = redis_client.get(balance_key)
|
||
usage_val = redis_client.get(usage_key)
|
||
|
||
effective_limit = None
|
||
|
||
if balance_val is not None and usage_val is not None:
|
||
try:
|
||
balance_at_sync = float(balance_val)
|
||
usage_at_sync = float(usage_val)
|
||
# 累计限额 = B0 - SMALL_BUFFER,不依赖 current_usage
|
||
dynamic_limit = balance_at_sync + usage_at_sync - SMALL_BUFFER
|
||
dynamic_limit = max(0, dynamic_limit)
|
||
|
||
effective_limit = calculate_effective_limit(static_limit, dynamic_limit)
|
||
except (ValueError, TypeError):
|
||
effective_limit = calculate_effective_limit(static_limit, None)
|
||
else:
|
||
# Fallback: 无同步数据时用保守兜底
|
||
fallback = float(static_limit) if static_limit is not None else float(FALLBACK_LIMIT)
|
||
fallback = min(fallback, float(FALLBACK_LIMIT))
|
||
effective_limit = calculate_effective_limit(fallback, None)
|
||
|
||
return effective_limit
|
||
|
||
|
||
# ==========================================
|
||
# 👻 Handler Logic (DailyLimitHandler)
|
||
# ==========================================
|
||
def soft_evict_from_router(target_model, target_api_base):
|
||
"""
|
||
Temporarily disable a model in the router by renaming it.
|
||
"""
|
||
if not LITELLM_AVAILABLE:
|
||
return
|
||
|
||
try:
|
||
from litellm.proxy.proxy_server import llm_router
|
||
|
||
if not llm_router:
|
||
return
|
||
|
||
# Prepare target identifiers
|
||
target_clean = clean_model_name(target_model)
|
||
target_hash = generate_api_base_hash(target_api_base)
|
||
|
||
for i, model_entry in enumerate(llm_router.model_list):
|
||
# Skip if already evicted/backed up
|
||
if "_daily_limit_backup" in model_entry:
|
||
continue
|
||
|
||
params = model_entry.get("litellm_params", {})
|
||
m_model = params.get("model", "")
|
||
m_base = params.get("api_base", "")
|
||
|
||
# Check Match
|
||
m_clean = clean_model_name(m_model)
|
||
m_hash = generate_api_base_hash(m_base)
|
||
|
||
# Match on Name AND Hash (Strict)
|
||
if m_clean == target_clean and m_hash == target_hash:
|
||
dprint(
|
||
f"🚫 Soft-evicting {model_entry.get('model_name')} (Limit Reached)"
|
||
)
|
||
|
||
# Create Backup
|
||
model_entry["_daily_limit_backup"] = {
|
||
"model_name": model_entry.get("model_name"),
|
||
"api_base": m_base,
|
||
}
|
||
|
||
# Disable by renaming (Router won't find it)
|
||
model_entry["model_name"] = (
|
||
f"{model_entry.get('model_name')}-limit-reached"
|
||
)
|
||
|
||
# Optional: break after finding first match?
|
||
# Or continue to evict all instances (e.g. load balancing same model)?
|
||
# Let's evict ALL matches to be safe.
|
||
# return
|
||
|
||
except Exception as e:
|
||
log(f"❌ Soft evict failed: {e}", level="ERROR")
|
||
|
||
|
||
last_restore_check = 0
|
||
|
||
|
||
def restore_blocked_models():
|
||
"""
|
||
Periodically check if blocked models have recovered quota.
|
||
Optimized for Volcengine:
|
||
- Quota usually resets in Beijing Time mornings (uncertain time).
|
||
- Logic:
|
||
- During Beijing Morning (00:00 - 14:00): Check every 30 mins.
|
||
- Other times: Check every 4 hours (low probability of reset).
|
||
"""
|
||
global last_restore_check
|
||
now = time.time()
|
||
|
||
# Calculate Dynamic Interval based on Beijing Time
|
||
# UTC+8 Logic
|
||
utc_now_struct = time.gmtime(now)
|
||
beijing_hour = (utc_now_struct.tm_hour + 8) % 24
|
||
|
||
# Morning Window: 07:00 - 13:00
|
||
if 7 <= beijing_hour < 13:
|
||
interval = 1800 # 30 minutes
|
||
else:
|
||
interval = 14400 # 4 hours
|
||
|
||
if now - last_restore_check < interval:
|
||
return
|
||
|
||
last_restore_check = now
|
||
|
||
try:
|
||
from litellm.proxy.proxy_server import llm_router
|
||
# ... (rest of the function) ...
|
||
|
||
if not llm_router or not llm_router.model_list:
|
||
return
|
||
|
||
restored_count = 0
|
||
for i, model_entry in enumerate(llm_router.model_list):
|
||
backup = model_entry.get("_daily_limit_backup")
|
||
if not backup:
|
||
continue
|
||
|
||
params = model_entry.get("litellm_params", {})
|
||
upstream_model = params.get("model", "")
|
||
clean_model = clean_model_name(upstream_model)
|
||
original_api_base = backup.get("api_base")
|
||
|
||
# 只恢复监控内的火山条目（clean + 端点 hash），与 soft_evict/请求闸门口径一致，
|
||
# 避免按 clean 名误恢复同名不同端点（如 newapi 渠道）被踢的条目
|
||
if MONITORED_MODEL_PAIRS and (
|
||
clean_model,
|
||
generate_api_base_hash(params.get("api_base", "")),
|
||
) not in MONITORED_MODEL_PAIRS:
|
||
continue
|
||
|
||
# Check Quota
|
||
# Need current usage to calculate dynamic limit correctly?
|
||
# get_effective_limit handles fetching usage if None passed.
|
||
effective_limit = get_effective_limit(clean_model, original_api_base)
|
||
|
||
should_restore = False
|
||
if effective_limit is None:
|
||
should_restore = True
|
||
else:
|
||
usage_key = get_usage_key(clean_model, original_api_base)
|
||
current_val = float(redis_client.get(usage_key) or 0)
|
||
if current_val < float(effective_limit):
|
||
should_restore = True
|
||
|
||
if should_restore:
|
||
dprint(f"♻️ [RESTORE] Restoring {backup['model_name']} (Quota OK)")
|
||
llm_router.model_list[i]["model_name"] = backup["model_name"]
|
||
llm_router.model_list[i]["litellm_params"]["api_base"] = (
|
||
original_api_base
|
||
)
|
||
del llm_router.model_list[i]["_daily_limit_backup"]
|
||
restored_count += 1
|
||
|
||
if restored_count > 0:
|
||
dprint(f"✅ Restored {restored_count} models.")
|
||
|
||
except Exception as e:
|
||
log(f"❌ Restore check failed: {e}", level="ERROR")
|
||
|
||
|
||
def restore_model_if_healthy(clean_model, allowed_hashes=None):
|
||
"""
|
||
若指定模型已被软踢（_daily_limit_backup）且当前限额已恢复（usage < effective limit），
|
||
立即恢复其路由条目。供同步守护在写健康限额后调用，不等 restore_blocked_models 的
|
||
30min/4h 闸门；否则额度到账后模型会带额度空转被禁（有额度但 no healthy deployments）。
|
||
allowed_hashes: 限定匹配的 api_base hash 集合（None = 不限，按 clean 名匹配）。
|
||
返回恢复的条目数；任何异常 fail-safe 返回 0。
|
||
"""
|
||
if not LITELLM_AVAILABLE or not redis_client:
|
||
return 0
|
||
try:
|
||
from litellm.proxy.proxy_server import llm_router
|
||
|
||
if not llm_router or not llm_router.model_list:
|
||
return 0
|
||
|
||
restored = 0
|
||
for i, model_entry in enumerate(llm_router.model_list):
|
||
backup = model_entry.get("_daily_limit_backup")
|
||
if not backup:
|
||
continue
|
||
|
||
params = model_entry.get("litellm_params", {})
|
||
m_model = params.get("model", "")
|
||
if clean_model_name(m_model) != clean_model:
|
||
continue
|
||
|
||
m_base = params.get("api_base", "")
|
||
if allowed_hashes is not None and generate_api_base_hash(
|
||
m_base
|
||
) not in allowed_hashes:
|
||
continue
|
||
|
||
effective_limit = get_effective_limit(clean_model, m_base)
|
||
should_restore = False
|
||
if effective_limit is None:
|
||
should_restore = True
|
||
else:
|
||
usage_key = get_usage_key(clean_model, m_base)
|
||
current_val = float(redis_client.get(usage_key) or 0)
|
||
if current_val < float(effective_limit):
|
||
should_restore = True
|
||
|
||
if should_restore:
|
||
dprint(f"♻️ [RESTORE] Restoring {backup['model_name']} (Quota OK)")
|
||
llm_router.model_list[i]["model_name"] = backup["model_name"]
|
||
llm_router.model_list[i]["litellm_params"]["api_base"] = m_base
|
||
del llm_router.model_list[i]["_daily_limit_backup"]
|
||
restored += 1
|
||
|
||
if restored > 0:
|
||
log(f"✅ Restored {restored} entry/entries of {clean_model}.", level="INFO")
|
||
return restored
|
||
|
||
except Exception as e:
|
||
log(f"❌ Restore failed for {clean_model}: {e}", level="ERROR")
|
||
return 0
|
||
|
||
|
||
def startup_audit():
|
||
"""
|
||
🚀 Startup Audit:
|
||
System initialization check. Iterate through all models in the router.
|
||
If a model is found to be already over the limit in Redis, kick it out immediately.
|
||
系统启动自检：遍历 Router 中所有模型，如果发现 Redis 中记录已超额，直接踢出。
|
||
"""
|
||
if not redis_client or not LITELLM_AVAILABLE:
|
||
return
|
||
|
||
try:
|
||
from litellm.proxy.proxy_server import llm_router
|
||
|
||
if not llm_router or not llm_router.model_list:
|
||
dprint("⚠️ Startup audit skipped: llm_router not ready.")
|
||
return
|
||
|
||
dprint("🚀 Starting Startup Audit (Pre-emptive Block)...")
|
||
blocked_count = 0
|
||
|
||
# Create a copy to iterate safely
|
||
# We only need to read configs to find targets
|
||
for model_entry in llm_router.model_list:
|
||
# Skip if already blocked (shouldn't happen on fresh start, but good practice)
|
||
if "_daily_limit_backup" in model_entry:
|
||
continue
|
||
|
||
params = model_entry.get("litellm_params", {})
|
||
raw_model = params.get("model", "")
|
||
api_base = params.get("api_base", "")
|
||
|
||
# 1. Clean Name
|
||
clean_model = clean_model_name(raw_model)
|
||
if not clean_model:
|
||
continue
|
||
|
||
# 只审计监控内的火山模型（clean + 端点 hash），与请求闸门口径一致，
|
||
# 避免误判同名非监控渠道（如 newapi 别名模型）超限而误踢
|
||
if MONITORED_MODEL_PAIRS and (
|
||
clean_model,
|
||
generate_api_base_hash(api_base),
|
||
) not in MONITORED_MODEL_PAIRS:
|
||
continue
|
||
|
||
# 2. Check Limit & Usage
|
||
# We use effective limit logic which checks both static and dynamic
|
||
effective_limit = get_effective_limit(clean_model, api_base)
|
||
|
||
if effective_limit is None:
|
||
continue
|
||
|
||
usage_key = get_usage_key(clean_model, api_base)
|
||
current_val = redis_client.get(usage_key)
|
||
|
||
if current_val and float(current_val) >= float(effective_limit):
|
||
dprint(
|
||
f"🚫 [STARTUP BLOCK] Found {raw_model} over limit ({int(float(current_val))}/{int(effective_limit)}). Evicting now."
|
||
)
|
||
soft_evict_from_router(raw_model, api_base)
|
||
blocked_count += 1
|
||
|
||
if blocked_count > 0:
|
||
dprint(f"✅ Startup Audit Complete: Blocked {blocked_count} models.")
|
||
else:
|
||
dprint("✅ Startup Audit Complete: All models healthy.")
|
||
|
||
except ImportError:
|
||
dprint("⚠️ Startup audit skipped: litellm not found.")
|
||
except Exception as e:
|
||
log(f"❌ Startup audit 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 表示放行。任何异常 fail-open。
|
||
"""
|
||
if not redis_client or not LITELLM_AVAILABLE:
|
||
return None
|
||
try:
|
||
from litellm.proxy.proxy_server import llm_router
|
||
|
||
if not llm_router or not llm_router.model_list:
|
||
return None
|
||
|
||
for entry in llm_router.model_list:
|
||
entry_name = entry.get("model_name") or ""
|
||
# 匹配请求名；被限额软踢的条目（原名被重命名为 {name}-limit-reached）
|
||
# 也参与判定，保证超限模型无论是否已踢都返回 429，而非 router 400
|
||
if entry_name != model_name:
|
||
if not (
|
||
"_daily_limit_backup" in entry
|
||
and entry_name == f"{model_name}-limit-reached"
|
||
):
|
||
continue
|
||
|
||
params = entry.get("litellm_params", {})
|
||
raw_model = params.get("model", "")
|
||
api_base = params.get("api_base", "")
|
||
|
||
clean = clean_model_name(raw_model)
|
||
h = generate_api_base_hash(api_base)
|
||
|
||
# 只拦已纳入监控的火山模型（模型 + 端点粒度）
|
||
if MONITORED_MODEL_PAIRS and (clean, h) not in MONITORED_MODEL_PAIRS:
|
||
continue
|
||
|
||
effective_limit = get_effective_limit(clean, api_base)
|
||
if effective_limit is None:
|
||
continue
|
||
|
||
usage_key = get_usage_key(clean, api_base)
|
||
val = redis_client.get(usage_key)
|
||
usage = float(val or 0)
|
||
|
||
if usage + float(reserve) >= float(effective_limit):
|
||
return (clean, usage, effective_limit)
|
||
|
||
return None
|
||
except Exception as e:
|
||
log(f"❌ Pre-block check failed: {e}", level="ERROR")
|
||
return None
|
||
|
||
|
||
def quota_interceptor(kwargs, user_api_key_dict=None):
|
||
try:
|
||
# 1. 尝试恢复（定期）
|
||
restore_blocked_models()
|
||
except (ServiceUnavailableError, ContextWindowExceededError):
|
||
raise
|
||
except Exception:
|
||
pass
|
||
|
||
|
||
class DailyLimitHandler(CustomLogger):
|
||
def __init__(self):
|
||
if LITELLM_AVAILABLE:
|
||
if quota_interceptor not in litellm.input_callback:
|
||
litellm.input_callback.append(quota_interceptor)
|
||
dprint("✅ quota_interceptor registered")
|
||
if self not in litellm.callbacks:
|
||
litellm.callbacks.append(self)
|
||
dprint("✅ DailyLimitHandler registered")
|
||
|
||
async def async_pre_call_hook(self, user_api_key_dict, cache, data, call_type):
|
||
"""
|
||
请求前限额闸门：在 route_request 之前拒绝超限请求（429），
|
||
不进入 router → 不触发 num_retries / fallback 放大。
|
||
"""
|
||
if call_type not in ("completion", "acompletion", "text_completion"):
|
||
return None
|
||
if not data or not data.get("model"):
|
||
return None
|
||
|
||
# 预留本次请求的预估消耗：优先真实 max_tokens（零采集损失），
|
||
# 未设置/非法时 fallback 常量，避免并发请求集体撞线超限。
|
||
try:
|
||
reserve = float(
|
||
data.get("max_tokens")
|
||
or data.get("max_completion_tokens")
|
||
or PREBLOCK_RESERVE_TOKENS
|
||
)
|
||
except (TypeError, ValueError):
|
||
reserve = float(PREBLOCK_RESERVE_TOKENS)
|
||
|
||
blocked = _check_request_block(data["model"], reserve=reserve)
|
||
if not blocked:
|
||
return None
|
||
|
||
clean, usage, limit = blocked
|
||
log(
|
||
f"🚫 [PRE-BLOCK] {data['model']} daily limit reached ({int(usage):,}+{int(reserve):,}/{int(limit):,})",
|
||
level="WARNING",
|
||
)
|
||
raise RateLimitError(
|
||
message=f"Daily limit reached for model '{data['model']}' ({int(usage):,}/{int(limit):,})",
|
||
llm_provider="volcengine",
|
||
model=data["model"],
|
||
)
|
||
|
||
async def async_log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||
self._record_usage(kwargs, response_obj)
|
||
|
||
def log_success_event(self, kwargs, response_obj, start_time, end_time):
|
||
self._record_usage(kwargs, response_obj)
|
||
|
||
def _record_usage(self, kwargs, response_obj):
|
||
if not redis_client:
|
||
return
|
||
try:
|
||
litellm_params = kwargs.get("litellm_params") or {}
|
||
raw_model = litellm_params.get("model") or kwargs.get("model")
|
||
clean_model = clean_model_name(raw_model)
|
||
api_base = litellm_params.get("api_base")
|
||
|
||
# FILTER: Only track models explicitly discovered as VolcEngine in config.yaml
|
||
# This prevents tracking NewAPI aliases that point to other providers
|
||
current_hash = generate_api_base_hash(api_base)
|
||
if (
|
||
MONITORED_MODEL_PAIRS
|
||
and (clean_model, current_hash) not in MONITORED_MODEL_PAIRS
|
||
):
|
||
return
|
||
|
||
usage = None
|
||
if hasattr(response_obj, "usage"):
|
||
usage = response_obj.usage
|
||
elif isinstance(response_obj, dict):
|
||
usage = response_obj.get("usage")
|
||
if hasattr(usage, "model_dump"):
|
||
usage = usage.model_dump()
|
||
elif hasattr(usage, "dict"):
|
||
usage = usage.dict()
|
||
|
||
if not usage:
|
||
return
|
||
total_tokens = usage.get("total_tokens", 0)
|
||
|
||
if clean_model and total_tokens > 0:
|
||
key = get_usage_key(clean_model, api_base)
|
||
current_total = redis_client.incrby(key, int(total_tokens))
|
||
redis_client.expire(key, 172800)
|
||
|
||
# Get Limit (Static + Dynamic)
|
||
effective_limit = get_effective_limit(
|
||
clean_model, api_base, current_usage=current_total
|
||
)
|
||
limit_str = (
|
||
f" | Limit: {int(effective_limit):,}" if effective_limit is not None else ""
|
||
)
|
||
|
||
# Get Balance At Sync for Log
|
||
balance_key = f"limit:dynamic:volcengine:balance_at_sync:{clean_model}"
|
||
bal_val = redis_client.get(balance_key)
|
||
bal_display = ""
|
||
if bal_val is not None:
|
||
try:
|
||
bal_int = int(float(bal_val))
|
||
bal_display = f" | BalAtSync: {bal_int:,}"
|
||
except Exception as e:
|
||
print(
|
||
f"DEBUG: Error calculating balance for {clean_model}: {e}"
|
||
)
|
||
pass
|
||
else:
|
||
# Debug why key is missing
|
||
if DEBUG_MODE:
|
||
print(f"DEBUG: Balance key not found: {balance_key}")
|
||
|
||
key_display = key.split(":", 2)[-1]
|
||
log(
|
||
f"📊 [Usage] {key_display}: +{total_tokens} (Today: {current_total:,}{bal_display}{limit_str})",
|
||
level="INFO",
|
||
)
|
||
|
||
if effective_limit is not None and float(current_total) >= float(effective_limit):
|
||
log(
|
||
f"🚫 [BLOCK] {clean_model} Limit Hit: {current_total}/{effective_limit}",
|
||
level="WARNING",
|
||
)
|
||
soft_evict_from_router(raw_model, api_base)
|
||
|
||
except Exception as e:
|
||
log(
|
||
f"❌ _record_usage failed: {e}\n{traceback.format_exc()}", level="ERROR"
|
||
)
|
||
|
||
|
||
# Module Level Instance
|
||
proxy_handler_instance = DailyLimitHandler()
|
||
|
||
|
||
# ==========================================
|
||
# 🔄 Daemon Logic (Active Monitor Loop)
|
||
# ==========================================
|
||
def fetch_all_packages():
|
||
action = "ListResourcePackages"
|
||
version = "2022-01-01"
|
||
all_packages = []
|
||
current_token = ""
|
||
seen_tokens = set()
|
||
|
||
while True:
|
||
try:
|
||
payload = {
|
||
"ResourceType": "Package",
|
||
"MaxResults": "20",
|
||
# 服务端只返回生效中的数据包，跳过 UsedUp/Expired/Refunded 等
|
||
# 余额为 0 的无用包，显著减少翻页次数与 QPS 占用。
|
||
# 周期型包（月度重置）状态仍为 Effective，不会被误过滤。
|
||
"Status": "Effective",
|
||
"NextToken": current_token,
|
||
}
|
||
headers, body = get_signed_headers(action, version, payload)
|
||
url = f"https://{HOST}/?Action={action}&Version={version}"
|
||
response = requests.post(url, data=body, headers=headers, timeout=60)
|
||
response.raise_for_status()
|
||
data = response.json()
|
||
items = data.get("Result", {}).get("List", [])
|
||
if items:
|
||
all_packages.extend(items)
|
||
next_token = data.get("Result", {}).get("NextToken", "")
|
||
|
||
# Infinite Loop Protection
|
||
# 火山 API 翻页到头后 NextToken 会循环回起始 token（实测数据只有
|
||
# ~8 页时第 9 页绕回第一个 token），必须用 seen_tokens 拦下载断。
|
||
if not next_token:
|
||
break
|
||
if next_token == current_token:
|
||
break
|
||
if next_token in seen_tokens:
|
||
log(
|
||
f"⚠️ VolcEngine pagination cycle detected (token {next_token} repeated). "
|
||
f"Stopping after {len(all_packages)} packages.",
|
||
level="WARNING",
|
||
)
|
||
break
|
||
|
||
seen_tokens.add(next_token)
|
||
current_token = next_token
|
||
time.sleep(0.2)
|
||
except Exception as e:
|
||
print(f"Error fetching packages: {e}")
|
||
break
|
||
return all_packages
|
||
|
||
|
||
def process_packages_breakdown(packages):
|
||
model_counts = {}
|
||
seen_package_hashes = set()
|
||
|
||
print(f"\n--- Raw Package Processing ---", flush=True)
|
||
|
||
for pkg in packages:
|
||
try:
|
||
# 0. Deduplication (Exact Match)
|
||
# VolcEngine pagination bugs might return duplicate items or pages.
|
||
# We assume identical JSON objects are duplicates.
|
||
pkg_hash = hashlib.md5(json.dumps(pkg, sort_keys=True).encode()).hexdigest()
|
||
if pkg_hash in seen_package_hashes:
|
||
if DEBUG_MODE:
|
||
print(f"DEBUG: Skipping duplicate package (Hash: {pkg_hash})")
|
||
continue
|
||
seen_package_hashes.add(pkg_hash)
|
||
|
||
raw_code = pkg.get("ConfigurationCode", "")
|
||
raw_name = pkg.get("ConfigurationName", "Unknown")
|
||
target_raw = raw_code if raw_code else raw_name
|
||
|
||
# 1. Normalize name (strips suffixes, handles case/hyphens)
|
||
norm_result = normalize_volc_name(target_raw)
|
||
|
||
status = pkg.get("Status", "")
|
||
if status == "Effective":
|
||
available = pkg.get("AvailableAmount", 0)
|
||
else:
|
||
available = 0
|
||
|
||
unit = pkg.get("Unit", "token")
|
||
|
||
# 2. Resolve to monitored upstream models.
|
||
# 普通版本号模型会自动从 config.yaml 推导；特殊别名走手工覆盖。
|
||
target_models = MONITORED_MODEL_ALIASES.get(norm_result, [norm_result])
|
||
|
||
# Debug log
|
||
if (
|
||
DEBUG_MODE
|
||
or "doubao" in str(target_models)
|
||
or "glm" in str(target_models)
|
||
):
|
||
print(
|
||
f"DEBUG: Pkg Raw='{target_raw}' -> Norm='{norm_result}' -> Targets={target_models}, Avail={int(float(available)):,}",
|
||
flush=True,
|
||
)
|
||
|
||
if unit.lower() == "token":
|
||
for model_name in target_models:
|
||
if model_name == "unknown":
|
||
continue
|
||
# Aggregate balances (Handle float string '123.00' -> float -> int)
|
||
try:
|
||
amount = int(float(available))
|
||
except ValueError:
|
||
print(
|
||
f"Error converting available amount '{available}' to int for {model_name}"
|
||
)
|
||
amount = 0
|
||
|
||
model_counts[model_name] = model_counts.get(model_name, 0) + amount
|
||
|
||
except Exception as e:
|
||
print(f"Error processing package: {e}")
|
||
continue
|
||
|
||
return model_counts
|
||
|
||
|
||
def discover_monitored_models():
|
||
"""
|
||
🔎 从 config.yaml 自动发现纳入监控的火山模型，填充 MONITORED_MODEL_PAIRS /
|
||
MONITORED_MODEL_ALIASES。返回监控模型名 set。
|
||
独立函数化：volc_consume.py 手动运行（无守护线程）时也可复用。
|
||
"""
|
||
global MONITORED_MODEL_PAIRS, MONITORED_MODEL_ALIASES
|
||
MONITORED_MODEL_PAIRS.clear()
|
||
MONITORED_MODEL_ALIASES.clear()
|
||
MONITORED_MODELS = set()
|
||
try:
|
||
if os.path.exists("config.yaml"):
|
||
with open("config.yaml", "r", encoding="utf-8") as f:
|
||
content = f.read()
|
||
lines = content.split("\n")
|
||
|
||
print(
|
||
"Auto-discovering monitored Volcengine models from config.yaml / 正在从 config.yaml 自动发现纳入监控的火山模型...",
|
||
flush=True,
|
||
)
|
||
|
||
# Robust Parsing Logic
|
||
current_entry = {
|
||
"is_volc": False,
|
||
"upstream_model": None,
|
||
"api_base_val": None,
|
||
}
|
||
|
||
for line in lines:
|
||
sline = line.strip()
|
||
|
||
# Skip commented lines (YAML comments). Otherwise commented-out
|
||
# api_base lines (e.g. a disabled model block) leak into the
|
||
# currently-open entry and can misclassify a NewAPI entry as
|
||
# VolcEngine.
|
||
if sline.startswith("#"):
|
||
continue
|
||
|
||
# Start of a new model definition
|
||
if sline.startswith("- model_name:"):
|
||
# Commit previous entry if valid
|
||
if (
|
||
current_entry.get("is_volc")
|
||
and current_entry.get("upstream_model")
|
||
and current_entry.get("api_base_val")
|
||
):
|
||
MONITORED_MODELS.add(current_entry["upstream_model"])
|
||
# Compute Hash
|
||
h = generate_api_base_hash(current_entry["api_base_val"])
|
||
MONITORED_MODEL_PAIRS.add((current_entry["upstream_model"], h))
|
||
print(
|
||
f"✅ [Config] Added monitored pair: {current_entry['upstream_model']} | Hash: {h} | Base: {str(current_entry['api_base_val'])[:30]}...",
|
||
flush=True,
|
||
)
|
||
|
||
# Reset for new entry
|
||
current_entry = {
|
||
"is_volc": False,
|
||
"upstream_model": None,
|
||
"api_base_val": None,
|
||
}
|
||
continue
|
||
|
||
# Capture API Base & Check for Volcengine
|
||
if "api_base:" in sline:
|
||
try:
|
||
raw_val = sline.split(":", 1)[1].strip()
|
||
# Resolve Env Var
|
||
if raw_val.startswith("os.environ/"):
|
||
env_key = raw_val.replace("os.environ/", "")
|
||
resolved_val = os.getenv(env_key, "")
|
||
if not resolved_val:
|
||
print(
|
||
f"⚠️ [Config] Warning: Env var {env_key} resolved to EMPTY!",
|
||
flush=True,
|
||
)
|
||
else:
|
||
resolved_val = raw_val
|
||
|
||
current_entry["api_base_val"] = resolved_val
|
||
|
||
# Identify as Volcengine if using HUOSHAN env or specific host
|
||
if (
|
||
"HUOSHAN_BASE_URL" in sline
|
||
or "volcengineapi" in resolved_val
|
||
):
|
||
current_entry["is_volc"] = True
|
||
except Exception:
|
||
pass
|
||
|
||
# Capture Upstream Model Name
|
||
# Matches: "model: openai/deepseek..."
|
||
if sline.startswith("model:"):
|
||
try:
|
||
raw_val = sline.split(":", 1)[1].strip()
|
||
# Clean prefixes
|
||
clean = (
|
||
raw_val.replace("openai/", "")
|
||
.replace("azure/", "")
|
||
.replace("anthropic/", "")
|
||
.replace("volcengine/", "")
|
||
)
|
||
current_entry["upstream_model"] = clean
|
||
except IndexError:
|
||
pass
|
||
|
||
# Commit the last entry
|
||
if (
|
||
current_entry.get("is_volc")
|
||
and current_entry.get("upstream_model")
|
||
and current_entry.get("api_base_val")
|
||
):
|
||
MONITORED_MODELS.add(current_entry["upstream_model"])
|
||
h = generate_api_base_hash(current_entry["api_base_val"])
|
||
MONITORED_MODEL_PAIRS.add((current_entry["upstream_model"], h))
|
||
if DEBUG_MODE:
|
||
print(
|
||
f"DEBUG: Added monitored pair: {current_entry['upstream_model']} @ {h}"
|
||
)
|
||
|
||
MONITORED_MODEL_ALIASES = build_monitored_model_aliases(MONITORED_MODELS)
|
||
print(
|
||
f"Auto-discovered / 已自动发现 {len(MONITORED_MODELS)} 个纳入监控的火山上游模型: {MONITORED_MODELS}",
|
||
flush=True,
|
||
)
|
||
print(
|
||
f"Monitored model pair hashes / 监控模型端点哈希: {[p[1] for p in MONITORED_MODEL_PAIRS]}",
|
||
flush=True,
|
||
)
|
||
if DEBUG_MODE:
|
||
print(
|
||
f"DEBUG: Monitored model aliases / 监控模型别名: {MONITORED_MODEL_ALIASES}",
|
||
flush=True,
|
||
)
|
||
|
||
except Exception as e:
|
||
print(f"Config Auto-discovery failed / 配置自动发现失败: {e}", flush=True)
|
||
|
||
return MONITORED_MODELS
|
||
|
||
|
||
def _maybe_trigger_daily_consume():
|
||
"""
|
||
🕰️ 每日北京时间 23:00 触发免费额度消耗任务（洗过期资源包）。
|
||
- consume:last_run:{北京日期} 去重（SETNX，当天只触发一次）
|
||
- daemon 线程执行，不阻塞主循环；线程内 0-10 分钟随机抖动，避免固定 23:00 被风控识别
|
||
- 延迟 import volc_consume，避免模块级循环 import
|
||
"""
|
||
if not redis_client:
|
||
return
|
||
try:
|
||
now = time.time()
|
||
beijing_ts = now + (8 * 3600)
|
||
beijing_hour = time.gmtime(beijing_ts).tm_hour
|
||
if beijing_hour != 23:
|
||
return
|
||
date_str = time.strftime("%Y-%m-%d", time.gmtime(beijing_ts))
|
||
lock_key = f"consume:last_run:{date_str}"
|
||
if redis_client.set(lock_key, "1", nx=True, ex=172800) is None:
|
||
return # 今天已触发过
|
||
|
||
log("🚀 触发每日消耗任务（洗过期资源包防退费）...", level="INFO")
|
||
|
||
def _run():
|
||
try:
|
||
time.sleep(random.randint(0, 600))
|
||
import volc_consume # 延迟导入，避免模块级循环 import
|
||
|
||
volc_consume.run_consume_all()
|
||
except Exception as e:
|
||
log(f"❌ 消耗任务异常: {e}\n{traceback.format_exc()}", level="ERROR")
|
||
|
||
t = threading.Thread(target=_run, daemon=True)
|
||
t.start()
|
||
except Exception as e:
|
||
log(f"❌ 消耗任务调度失败: {e}", level="ERROR")
|
||
|
||
|
||
def compute_sync_interval(min_balance, is_idle):
|
||
"""按最小余额 + 空闲状态计算同步间隔(秒)。
|
||
min_balance 是 balance_at_sync 中的最小值(实际余额,非扣 SMALL_BUFFER 后)。"""
|
||
if min_balance < 100_000:
|
||
base_interval = 300 # 紧急 5 minutes
|
||
elif min_balance < 400_000:
|
||
base_interval = 600 # 低余额 10 minutes
|
||
else:
|
||
base_interval = 1800 # 常规 30 minutes
|
||
|
||
if is_idle:
|
||
interval = min(base_interval * 4, 3600)
|
||
if interval < 300:
|
||
interval = 300
|
||
else:
|
||
interval = base_interval
|
||
return interval
|
||
|
||
|
||
def active_monitor_loop():
|
||
# Config Discovery
|
||
MONITORED_MODELS = discover_monitored_models()
|
||
|
||
if not redis_client:
|
||
print(
|
||
"❌ Redis not connected. Exiting monitor loop / Redis 未连接，退出监控循环.",
|
||
flush=True,
|
||
)
|
||
sys.exit(1)
|
||
|
||
print(
|
||
f"✅ Redis connected at / Redis 已连接于 {REDIS_HOST}:{REDIS_PORT}", flush=True
|
||
)
|
||
|
||
# ==========================================
|
||
# 🚀 Run Startup Audit Immediately
|
||
# ==========================================
|
||
# Give LiteLLM a moment to initialize router if running concurrently
|
||
time.sleep(2)
|
||
startup_audit()
|
||
|
||
last_usage_map = {}
|
||
|
||
last_sync_time = 0
|
||
current_sync_interval = 0
|
||
global_min_remaining = 0
|
||
IDLE_THRESHOLD = 300
|
||
idle_start_time = 0
|
||
|
||
print(
|
||
"Starting Active Monitor Loop (Safety-Margin + Idle Awareness) / 正在启动主动监控循环 (安全余量 + 空闲感知)...",
|
||
flush=True,
|
||
)
|
||
print(f"SmallBuffer / 缓冲: {SMALL_BUFFER}", flush=True)
|
||
|
||
while True:
|
||
try:
|
||
time.sleep(CHECK_INTERVAL)
|
||
_maybe_trigger_daily_consume()
|
||
now = time.time()
|
||
|
||
# Enforce Beijing Time for daily operations
|
||
beijing_time_ts = now + (8 * 3600)
|
||
today_str = time.strftime("%Y-%m-%d", time.gmtime(beijing_time_ts))
|
||
|
||
# Idle Check - Only exit idle if a "low quota" model has usage
|
||
LOW_QUOTA_THRESHOLD = 800000 # 判定为“低额度”的阈值
|
||
model_usages = {}
|
||
any_low_quota_usage = False
|
||
|
||
for m in MONITORED_MODELS:
|
||
m_usage = 0
|
||
target_hashes = [p[1] for p in MONITORED_MODEL_PAIRS if p[0] == m]
|
||
for h in target_hashes:
|
||
val = redis_client.get(f"daily_usage:{today_str}:{m}:{h}")
|
||
if val: m_usage += float(val)
|
||
val_no_hash = redis_client.get(f"daily_usage:{today_str}:{m}")
|
||
if val_no_hash: m_usage += float(val_no_hash)
|
||
model_usages[m] = m_usage
|
||
|
||
# Calculate delta for this specific model
|
||
m_delta = m_usage - last_usage_map.get(m, 0.0)
|
||
|
||
# Check if this model is considered "low quota"
|
||
bal_key = f"limit:dynamic:volcengine:balance_at_sync:{m}"
|
||
bal_val = redis_client.get(bal_key)
|
||
# 未知额度视为充足,避免首次启动时所有模型被误判为"低额度"
|
||
remaining = float(bal_val) if bal_val is not None else float(LOW_QUOTA_THRESHOLD)
|
||
|
||
if m_delta > 0 and remaining < LOW_QUOTA_THRESHOLD:
|
||
any_low_quota_usage = True
|
||
|
||
time_diff = now - last_sync_time
|
||
|
||
if not any_low_quota_usage:
|
||
if idle_start_time == 0:
|
||
idle_start_time = now
|
||
else:
|
||
idle_start_time = 0
|
||
|
||
is_idle = (idle_start_time > 0) and (
|
||
(now - idle_start_time) > IDLE_THRESHOLD
|
||
)
|
||
|
||
# Dynamic Interval: 按最小正剩余余额分档（见 compute_sync_interval）
|
||
current_sync_interval = compute_sync_interval(
|
||
global_min_remaining, is_idle
|
||
)
|
||
|
||
# =====================================================
|
||
# 🚀 Per-Model Burst Detection (精细化突发流量检测)
|
||
# =====================================================
|
||
burst_triggered = False
|
||
burst_reason = ""
|
||
|
||
for model in MONITORED_MODELS:
|
||
# 1. Get Current Usage (calculated above)
|
||
curr_usage = model_usages.get(model, 0.0)
|
||
|
||
# 2. Get Last Synced Usage (Snapshot)
|
||
last_usage = last_usage_map.get(model, 0.0)
|
||
|
||
# 3. Calculate Delta since last sync
|
||
model_delta = curr_usage - last_usage
|
||
|
||
# 4. Get Current Balance for Threshold Calculation
|
||
bal_key = f"limit:dynamic:volcengine:balance_at_sync:{model}"
|
||
bal_val = redis_client.get(bal_key)
|
||
current_balance = float(bal_val) if bal_val else 0.0
|
||
|
||
# 5. Determine Dynamic Threshold based on Balance
|
||
# Rich models get loose limits, Poor models get strict limits.
|
||
if current_balance > 1000000:
|
||
model_threshold = 500000
|
||
elif current_balance > 500000:
|
||
model_threshold = 100000
|
||
elif current_balance > 100000:
|
||
model_threshold = 50000
|
||
else:
|
||
model_threshold = 10000 # Strict for poor models
|
||
|
||
# 6. Check Trigger
|
||
if model_delta > model_threshold:
|
||
burst_triggered = True
|
||
burst_reason = f"{model} (+{model_delta:,} > {model_threshold:,})"
|
||
break # One burst is enough to trigger sync
|
||
|
||
should_sync = (
|
||
burst_triggered
|
||
or (time_diff > current_sync_interval)
|
||
or (last_sync_time == 0)
|
||
)
|
||
|
||
if should_sync:
|
||
trigger_reason = (
|
||
f"Burst: {burst_reason}" if burst_triggered else "Timer Expired"
|
||
)
|
||
# For Timer Expired, we don't have a specific delta to show, or we can calculate global delta if needed.
|
||
# But for now, let's just show the reason.
|
||
print(
|
||
f"{'💤' if is_idle else '⚡'} Sync Triggered ({trigger_reason}) / 同步触发... Interval: {current_sync_interval}s",
|
||
flush=True,
|
||
)
|
||
|
||
packages = fetch_all_packages()
|
||
model_balances = process_packages_breakdown(packages)
|
||
|
||
print(
|
||
f"DEBUG: Processed models / 已处理模型: {list(model_balances.keys())}",
|
||
flush=True,
|
||
)
|
||
|
||
temp_min_remaining = float("inf")
|
||
configured_models_found = False
|
||
|
||
for model_name, balance in model_balances.items():
|
||
print(
|
||
f"DEBUG: Checking {model_name} (Bal: {int(balance):,})",
|
||
flush=True,
|
||
)
|
||
|
||
# Target Resolution
|
||
targets = [model_name]
|
||
|
||
for target in targets:
|
||
# STRICT FILTER: Only sync models explicitly configured with VolcEngine Base URL
|
||
if target not in MONITORED_MODELS:
|
||
if DEBUG_MODE:
|
||
print(
|
||
f"DEBUG: Skipping {target} (Not in monitored models / 未纳入监控模型)",
|
||
flush=True,
|
||
)
|
||
continue
|
||
|
||
# Fetch Current Usage from Redis
|
||
current_usage = 0
|
||
target_hashes = [
|
||
p[1] for p in MONITORED_MODEL_PAIRS if p[0] == target
|
||
]
|
||
|
||
# Sum usage across all configured endpoints for this model
|
||
for h in target_hashes:
|
||
now_ts = time.time()
|
||
beijing_time_ts = now_ts + (8 * 3600)
|
||
usage_date = time.strftime(
|
||
"%Y-%m-%d", time.gmtime(beijing_time_ts)
|
||
)
|
||
u_key = f"daily_usage:{usage_date}:{target}:{h}"
|
||
val = redis_client.get(u_key)
|
||
if val:
|
||
current_usage += float(val)
|
||
|
||
balance_key = f"limit:dynamic:volcengine:balance_at_sync:{target}"
|
||
usage_key = f"limit:dynamic:volcengine:usage_at_sync:{target}"
|
||
|
||
# Store snapshot: balance_at_sync, usage_at_sync
|
||
redis_client.setex(balance_key, 86400, int(balance))
|
||
redis_client.setex(usage_key, 86400, int(current_usage))
|
||
|
||
if DEBUG_MODE:
|
||
static_limit = get_monitored_static_limit(
|
||
target, target_hashes
|
||
)
|
||
print(
|
||
f"DEBUG: {target} Bal={int(balance):,} | Usage={int(current_usage):,} | BalAtSync={int(balance):,} UsageAtSync={int(current_usage):,}",
|
||
flush=True,
|
||
)
|
||
|
||
# 即时恢复:同步发现限额已恢复(usage < effective limit)时立刻解除软踢,
|
||
restore_model_if_healthy(target, target_hashes)
|
||
|
||
# Track min balance for frequency control
|
||
configured_models_found = True
|
||
|
||
if balance > 0 and balance < temp_min_remaining:
|
||
temp_min_remaining = balance
|
||
|
||
# If all models are dead (temp_min_remaining is still inf), set to 0 for logic but handle gracefully
|
||
# If we found models but all were <= 0, then temp_min_remaining is inf.
|
||
|
||
# Logic:
|
||
# 1. No models found at all -> Default to Safe (300k)
|
||
# 2. Models found, but all 0 -> "All Dead Mode"
|
||
# 3. Models found, some > 0 -> Use min of >0 models
|
||
|
||
if not configured_models_found:
|
||
global_min_remaining = 300000 # Default
|
||
elif temp_min_remaining == float("inf"):
|
||
# All active models are <= 0
|
||
# Check time of day to decide if we should panic
|
||
utc_now_struct = time.gmtime(now)
|
||
beijing_hour = (utc_now_struct.tm_hour + 8) % 24
|
||
if 7 <= beijing_hour < 13:
|
||
# Morning Reset Window: Check relatively frequently
|
||
global_min_remaining = (
|
||
200000 # Pretend we have some quota to get ~5min interval
|
||
)
|
||
else:
|
||
# 返还"不定时"（可能下午到账）：All Dead 不能睡 60 分钟，
|
||
# 否则额度到账后最多空转 60 分钟才解锁。收紧到 30 分钟。
|
||
global_min_remaining = (
|
||
300000 # 30-minute sync while waiting for reset
|
||
)
|
||
else:
|
||
global_min_remaining = temp_min_remaining
|
||
|
||
# Update snapshots for next iteration delta calculation
|
||
for m, u_val in model_usages.items():
|
||
last_usage_map[m] = u_val
|
||
last_sync_time = now
|
||
|
||
if global_min_remaining <= 100000:
|
||
mode_icon = "🆘 紧急模式"
|
||
elif global_min_remaining <= 300000:
|
||
mode_icon = "🚨 低余额模式"
|
||
else:
|
||
mode_icon = "ℹ️ 常规模式"
|
||
# 用本次同步后的新 Min Rem 重算，避免"旧 interval + 新 Min Rem"混打误导
|
||
next_interval = compute_sync_interval(global_min_remaining, is_idle)
|
||
print(
|
||
f"{mode_icon}: 下一次同步 {next_interval}s (Min Rem: {int(global_min_remaining):,}, {'空闲' if is_idle else '活跃'}) / 同步间隔 {next_interval}s",
|
||
flush=True,
|
||
)
|
||
|
||
except Exception as e:
|
||
print(f"Error in monitor loop / 监控循环出错: {e}", flush=True)
|
||
print(traceback.format_exc(), end="")
|
||
time.sleep(5)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
active_monitor_loop()
|
||
else:
|
||
# Auto-start when loaded as a module/callback
|
||
print(
|
||
"🚀 Initializing VolcSync Background Monitor... / 正在初始化火山引擎同步后台监控...",
|
||
flush=True,
|
||
)
|
||
t = threading.Thread(target=active_monitor_loop, daemon=True)
|
||
t.start()
|