feat: add dynamic rate limiting for Volcengine models
This commit introduces a dual-layer rate limiting system: 1. volc_limit_sync.py: A daemon that adaptively syncs Volcengine API balances to Redis based on traffic and safety margins. 2. daily_limit.py: Updated to enforce dynamic limits by checking Redis (Current Usage + Remaining Balance) and handling soft evictions.
This commit is contained in:
@@ -0,0 +1,456 @@
|
||||
import os
|
||||
import sys
|
||||
import time
|
||||
import json
|
||||
import hmac
|
||||
import hashlib
|
||||
import datetime
|
||||
import requests
|
||||
import redis
|
||||
|
||||
# ================= Configuration =================
|
||||
ACCESS_KEY = os.getenv(
|
||||
"VOLC_ACCESS_KEY", "AKLTNGYxZmUyNzNmZjk3NGMyZmEzNjVjZmY4NGM5ZjcxMjU"
|
||||
)
|
||||
SECRET_KEY = os.getenv(
|
||||
"VOLC_SECRET_KEY", "Tm1VNE5XSmlORGs1WXpGbU5EVmtabUUyT1RZeVltTTBNalkxTWpjeE1tTQ=="
|
||||
)
|
||||
HOST = "open.volcengineapi.com"
|
||||
REGION = "cn-beijing"
|
||||
|
||||
BUFFER = 300000 # User requested 30w buffer
|
||||
SYNC_INTERVAL = 60
|
||||
CHECK_INTERVAL = 10
|
||||
THRESHOLD = 50000
|
||||
|
||||
SUFFIXES_TO_REMOVE = [
|
||||
"_data_collaboration_resource_pack",
|
||||
"协作奖励计划资源包",
|
||||
"免费在线推理资源包",
|
||||
"免费在线推理包",
|
||||
"资源包",
|
||||
]
|
||||
|
||||
# 火山引擎清洗后名称 -> LiteLLM 模型名称的映射表
|
||||
# 用户 TODO: 请根据您的实际部署情况手动配置此映射
|
||||
#
|
||||
# Key 格式 (火山清洗后名称):
|
||||
# 1. 移除后缀 (如 "_data_collaboration_resource_pack")
|
||||
# 2. 替换下划线和点为中划线 ("_.", -> "-")
|
||||
# 3. 转小写
|
||||
# 示例: "DeepSeek_V3.2_..." -> "deepseek-v3-2"
|
||||
#
|
||||
# Value 格式 (LiteLLM 清洗后名称):
|
||||
# config.yaml 中定义的 model_name (不带 openai/ 等前缀)
|
||||
# 类型: 字符串 或 字符串列表
|
||||
# 示例 (一对一): "deepseek-v3-2-251201"
|
||||
# 示例 (一对多, 共享额度): ["doubao-seed-1-6-250615", "doubao-seed-1-6-251015"]
|
||||
VOLC_MODEL_MAP = {
|
||||
"doubao-seed-1-6": ["doubao-seed-1-6-250615", "doubao-seed-1-6-251015"], # 豆包 1.6
|
||||
# "doubao-seed-1-6-thinking": "", # 豆包 1.6 Thinking
|
||||
"doubao-seed-1-8": "doubao-seed-1-8-251228", # 豆包 1.8
|
||||
"doubao-seed-code": "doubao-seed-code-preview-251028", # 豆包 Code
|
||||
"doubao-seed-translation": "doubao-seed-translation-250915", # 豆包 Translation
|
||||
"deepseek-v3": "deepseek-v3-250324", # DeepSeek V3 (通用)
|
||||
"deepseek-v3-1": "deepseek-v3-1-terminus", # DeepSeek V3.1
|
||||
"deepseek-v3-2": "deepseek-v3-2-251201", # DeepSeek V3.2
|
||||
"deepseek-r1": "deepseek-r1-250528", # DeepSeek R1
|
||||
"kimi-k2": "kimi-k2-thinking-251104", # Kimi K2
|
||||
"glm-4_7": "glm-4-7-251222",
|
||||
}
|
||||
# Redis Configuration
|
||||
REDIS_HOST = os.getenv("REDIS_HOST", "redis")
|
||||
REDIS_PORT = int(os.getenv("REDIS_PORT", 6379))
|
||||
REDIS_PASSWORD = os.getenv("REDIS_PASSWORD", None)
|
||||
|
||||
|
||||
# ================= Signing Logic =================
|
||||
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"):
|
||||
# 1. Prepare Time
|
||||
t = datetime.datetime.utcnow()
|
||||
amz_date = t.strftime("%Y%m%dT%H%M%SZ")
|
||||
date_stamp = t.strftime("%Y%m%d")
|
||||
|
||||
# 2. Process Body
|
||||
request_body = json.dumps(body_dict)
|
||||
payload_hash = hashlib.sha256(request_body.encode("utf-8")).hexdigest()
|
||||
|
||||
# 3. Canonical Request
|
||||
method = "POST"
|
||||
canonical_uri = "/"
|
||||
canonical_querystring = f"Action={action}&Version={version}"
|
||||
content_type = "application/json"
|
||||
|
||||
# Headers must be sorted
|
||||
canonical_headers = f"content-type:{content_type}\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"{method}\n{canonical_uri}\n{canonical_querystring}\n{canonical_headers}\n{signed_headers}\n{payload_hash}"
|
||||
|
||||
# 4. String to Sign
|
||||
algorithm = "HMAC-SHA256"
|
||||
credential_scope = f"{date_stamp}/{REGION}/{service}/request"
|
||||
string_to_sign = f"{algorithm}\n{amz_date}\n{credential_scope}\n{hashlib.sha256(canonical_request.encode('utf-8')).hexdigest()}"
|
||||
|
||||
# 5. Calculate Signature
|
||||
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()
|
||||
|
||||
# 6. Authorization Header
|
||||
authorization_header = f"{algorithm} Credential={ACCESS_KEY}/{credential_scope}, SignedHeaders={signed_headers}, Signature={signature}"
|
||||
|
||||
return {
|
||||
"Content-Type": content_type,
|
||||
"Host": HOST,
|
||||
"X-Date": amz_date,
|
||||
"X-Content-Sha256": payload_hash,
|
||||
"Authorization": authorization_header,
|
||||
}, request_body
|
||||
|
||||
|
||||
# ================= Helper Functions =================
|
||||
def normalize_model_name(raw_name):
|
||||
"""
|
||||
Align Volcengine response names with LiteLLM upstream names.
|
||||
1. Remove suffixes.
|
||||
2. Replace _ with -.
|
||||
3. Replace . with -. (Crucial: V3.2 -> v3-2)
|
||||
4. Lowercase.
|
||||
5. Check VOLC_MODEL_MAP for explicit overrides.
|
||||
"""
|
||||
if not raw_name:
|
||||
return "unknown"
|
||||
|
||||
clean_name = raw_name
|
||||
|
||||
# Step 1: Remove suffixes
|
||||
for suffix in sorted(SUFFIXES_TO_REMOVE, key=len, reverse=True):
|
||||
if clean_name.endswith(suffix):
|
||||
clean_name = clean_name.replace(suffix, "")
|
||||
break
|
||||
|
||||
# Step 2 & 3: Replace _ and . with -
|
||||
clean_name = clean_name.replace("_", "-").replace(".", "-")
|
||||
|
||||
# Step 4: Lowercase everything
|
||||
clean_name = clean_name.lower().strip()
|
||||
|
||||
# Step 5: Check VOLC_MODEL_MAP for explicit overrides
|
||||
return VOLC_MODEL_MAP.get(clean_name, clean_name)
|
||||
|
||||
|
||||
def fetch_all_packages():
|
||||
"""Fetch all package pages from Volcengine API."""
|
||||
action = "ListResourcePackages"
|
||||
version = "2022-01-01"
|
||||
all_packages = []
|
||||
current_token = ""
|
||||
|
||||
print(f"[{datetime.datetime.now()}] Fetching packages from Volcengine...")
|
||||
|
||||
while True:
|
||||
try:
|
||||
# User doc: NextToken 首次查询传 ""
|
||||
payload = {
|
||||
"ResourceType": "Package",
|
||||
"MaxResults": "20",
|
||||
"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()
|
||||
|
||||
# Debug: print first item of first page to verify fields
|
||||
# if not current_token and DEBUG_MODE:
|
||||
# pass
|
||||
|
||||
result_block = data.get("Result", {})
|
||||
items = result_block.get("List", [])
|
||||
|
||||
if items:
|
||||
all_packages.extend(items)
|
||||
|
||||
next_token = result_block.get("NextToken", "")
|
||||
if not next_token or next_token == current_token:
|
||||
break
|
||||
|
||||
current_token = next_token
|
||||
time.sleep(0.2) # Rate limit protection
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error fetching packages: {e}")
|
||||
break
|
||||
|
||||
return all_packages
|
||||
|
||||
|
||||
def active_monitor_loop():
|
||||
# 4. In-Line Config Parser (Config Awareness)
|
||||
# Always run auto-discovery to populate INTERESTED_MODELS (The Whitelist)
|
||||
INTERESTED_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")
|
||||
current_model_name = None
|
||||
|
||||
print("Auto-discovering Volcengine models from config.yaml...")
|
||||
|
||||
for line in lines:
|
||||
sline = line.strip()
|
||||
if sline.startswith("- model_name:"):
|
||||
current_model_name = sline.split(":", 1)[1].strip()
|
||||
# Remove quotes if present
|
||||
if current_model_name.startswith(
|
||||
'"'
|
||||
) or current_model_name.startswith("'"):
|
||||
current_model_name = current_model_name[1:-1]
|
||||
|
||||
if current_model_name:
|
||||
# Heuristic: Check for Volcengine signals in the block
|
||||
if "api_base:" in sline and ("HUOSHAN" in sline or "volc" in sline):
|
||||
INTERESTED_MODELS.add(current_model_name)
|
||||
# Also check provider prefix in model name
|
||||
if "model:" in sline and "volcengine/" in sline:
|
||||
INTERESTED_MODELS.add(current_model_name)
|
||||
|
||||
print(
|
||||
f"Auto-discovered {len(INTERESTED_MODELS)} Active Volcengine models: {INTERESTED_MODELS}"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Config Auto-discovery failed: {e}")
|
||||
|
||||
# Connect to Redis
|
||||
try:
|
||||
r = redis.Redis(
|
||||
host=REDIS_HOST,
|
||||
port=REDIS_PORT,
|
||||
password=REDIS_PASSWORD,
|
||||
decode_responses=True,
|
||||
)
|
||||
r.ping()
|
||||
print(f"Connected to Redis at {REDIS_HOST}:{REDIS_PORT}")
|
||||
except Exception as e:
|
||||
print(f"Failed to connect to Redis: {e}")
|
||||
sys.exit(1)
|
||||
|
||||
last_usage_map = {} # Track usage per model
|
||||
last_sync_time = 0
|
||||
# Dynamic interval state (Start with aggressive check to initialize)
|
||||
current_sync_interval = 0
|
||||
global_min_remaining = 0
|
||||
IDLE_THRESHOLD = 300 # 5 minutes of no usage -> Start relaxing
|
||||
idle_start_time = 0
|
||||
|
||||
print("Starting Active Monitor Loop (Safety-Margin + Idle Awareness)...")
|
||||
print(f"Buffer: {BUFFER}, Threshold: {THRESHOLD}")
|
||||
|
||||
while True:
|
||||
try:
|
||||
time.sleep(CHECK_INTERVAL)
|
||||
|
||||
now = time.time()
|
||||
today_str = datetime.datetime.now().strftime("%Y-%m-%d")
|
||||
|
||||
# 1. Fetch current global usage
|
||||
usage_key_global = f"daily_usage:{today_str}:global_total"
|
||||
current_usage_global = int(r.get(usage_key_global) or 0)
|
||||
|
||||
last_usage_global = sum(last_usage_map.values()) if last_usage_map else 0
|
||||
delta = current_usage_global - last_usage_global
|
||||
time_diff = now - last_sync_time
|
||||
|
||||
# 2. Idle Detection Logic
|
||||
if delta <= 0: # No traffic
|
||||
if idle_start_time == 0:
|
||||
idle_start_time = now
|
||||
else:
|
||||
idle_start_time = 0 # Traffic detected, reset idle timer
|
||||
|
||||
is_idle = (idle_start_time > 0) and (
|
||||
(now - idle_start_time) > IDLE_THRESHOLD
|
||||
)
|
||||
|
||||
# 3. Determine Dynamic Interval
|
||||
# Logic: Base Interval (based on balance) + Idle Multiplier
|
||||
if global_min_remaining > 2_000_000:
|
||||
base_interval = 1800 # 30 min (>200w)
|
||||
elif global_min_remaining > 1_000_000:
|
||||
base_interval = 600 # 10 min (100w-200w)
|
||||
elif global_min_remaining > 600_000:
|
||||
base_interval = 300 # 5 min (60w-100w)
|
||||
elif global_min_remaining > 300_000:
|
||||
base_interval = 60 # 1 min (30w-60w)
|
||||
else:
|
||||
base_interval = 10 # Panic (<30w Remaining)
|
||||
|
||||
# Apply Idle Relaxation
|
||||
# If system is idle for >5 min, we can relax the check interval.
|
||||
# Max idle interval extended to 3600s (1 hour) as requested.
|
||||
if is_idle:
|
||||
# Exponential backoff or just aggressive multiplier?
|
||||
# Let's quadruple the base interval, capped at 3600s
|
||||
current_sync_interval = min(base_interval * 4, 3600)
|
||||
if current_sync_interval < 300:
|
||||
current_sync_interval = 300 # Minimum 5 min sleep if idle
|
||||
else:
|
||||
current_sync_interval = base_interval
|
||||
|
||||
should_sync = (
|
||||
(delta > THRESHOLD)
|
||||
or (time_diff > current_sync_interval)
|
||||
or (last_sync_time == 0)
|
||||
)
|
||||
|
||||
if should_sync:
|
||||
trigger_reason = (
|
||||
"Burst Traffic" if delta > THRESHOLD else "Timer Expired"
|
||||
)
|
||||
status_icon = "💤" if is_idle else "⚡"
|
||||
print(
|
||||
f"{status_icon} Sync Triggered ({trigger_reason})... Delta: {delta}, Interval: {current_sync_interval}s"
|
||||
)
|
||||
|
||||
# Fetch and Process
|
||||
packages = fetch_all_packages()
|
||||
model_balances = process_packages_breakdown(packages)
|
||||
|
||||
# Update Redis and calculate Global Min Remaining (for next interval)
|
||||
temp_min_remaining = float("inf")
|
||||
|
||||
# Logic Fix: Only active targets affect the panic timer.
|
||||
# VOLC_MODEL_MAP is for Translation Only.
|
||||
# INTERESTED_MODELS is the Source of Truth for Whitelist.
|
||||
|
||||
configured_models_found = False
|
||||
|
||||
for model_name, balance in model_balances.items():
|
||||
remaining = balance - BUFFER
|
||||
|
||||
# If VOLC_MODEL_MAP is not empty, we ignore anything NOT in it.
|
||||
is_interested = True
|
||||
|
||||
# 1. Resolve Targets (Translation)
|
||||
# Translate Volc Name -> LiteLLM Name(s) using MAP if available
|
||||
targets = []
|
||||
if model_name in VOLC_MODEL_MAP:
|
||||
val = VOLC_MODEL_MAP[model_name]
|
||||
if val:
|
||||
targets = val if isinstance(val, list) else [val]
|
||||
else:
|
||||
# Empty string = Use default name (Auto Match)
|
||||
targets = [model_name]
|
||||
else:
|
||||
targets = [model_name]
|
||||
|
||||
# 2. Filter Targets (Whitelist)
|
||||
# Only keep targets that are actually in INTERESTED_MODELS (config.yaml)
|
||||
|
||||
final_targets_to_update = []
|
||||
is_active_for_panic = False
|
||||
|
||||
for target in targets:
|
||||
# Check against Whitelist
|
||||
matched = False
|
||||
if not INTERESTED_MODELS:
|
||||
# Fallback: Monitor everything if config discovery failed
|
||||
matched = True
|
||||
else:
|
||||
for interested in INTERESTED_MODELS:
|
||||
# Exact match (Highest Priority)
|
||||
if target == interested:
|
||||
matched = True
|
||||
# Heuristic match (only if target is raw volc name)
|
||||
# If target was explicitly mapped, it implies exact intent.
|
||||
# If target is raw, we allow fuzzy.
|
||||
elif target == model_name and (
|
||||
target in interested or interested in target
|
||||
):
|
||||
matched = True
|
||||
|
||||
if matched:
|
||||
break
|
||||
|
||||
if matched:
|
||||
final_targets_to_update.append(target)
|
||||
is_active_for_panic = True
|
||||
|
||||
# 3. Safety Calculation
|
||||
# Only active models affect the system panic level
|
||||
if is_active_for_panic:
|
||||
configured_models_found = True
|
||||
if remaining < temp_min_remaining:
|
||||
temp_min_remaining = remaining
|
||||
|
||||
# Redis Update: We still update Redis for the active targets
|
||||
for target in final_targets_to_update:
|
||||
key = f"limit:dynamic:volcengine:remaining:{target}"
|
||||
r.setex(key, 7200, int(remaining))
|
||||
print(f"Updated {target}: Bal={balance:,} -> Rem={remaining:,}")
|
||||
|
||||
# Update State
|
||||
global_min_remaining = temp_min_remaining if model_balances else 0
|
||||
last_usage_map["global"] = current_usage_global
|
||||
last_sync_time = now
|
||||
|
||||
# Log next scheduled check
|
||||
if global_min_remaining > 10_000_000:
|
||||
print(
|
||||
f"🌊 High Tide Mode: Sleeping for 30m (Min Rem: {int(global_min_remaining):,})"
|
||||
)
|
||||
elif global_min_remaining > 2_000_000:
|
||||
print(
|
||||
f"⛵ Safe Sailing Mode: Sleeping for 5m (Min Rem: {int(global_min_remaining):,})"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error in monitor loop: {e}")
|
||||
time.sleep(5)
|
||||
|
||||
|
||||
def process_packages_breakdown(packages):
|
||||
"""Return dictionary of {model_name: total_available}"""
|
||||
model_counts = {}
|
||||
|
||||
for pkg in packages:
|
||||
raw_code = pkg.get("ConfigurationCode", "")
|
||||
raw_name = pkg.get("ConfigurationName", "Unknown")
|
||||
target_raw = raw_code if raw_code else raw_name
|
||||
|
||||
model_name = normalize_model_name(target_raw)
|
||||
|
||||
try:
|
||||
available = int(pkg.get("AvailableAmount", 0))
|
||||
unit = pkg.get("Unit", "token").lower()
|
||||
if unit == "token":
|
||||
model_counts[model_name] = model_counts.get(model_name, 0) + available
|
||||
except ValueError:
|
||||
continue
|
||||
|
||||
return model_counts
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
active_monitor_loop()
|
||||
Reference in New Issue
Block a user