383 lines
14 KiB
Plaintext
383 lines
14 KiB
Plaintext
import os
|
||
import time
|
||
import hashlib
|
||
from urllib.parse import urlparse # <--- 新增导入
|
||
import litellm
|
||
from litellm.integrations.custom_logger import CustomLogger
|
||
from litellm.exceptions import ServiceUnavailableError, ContextWindowExceededError
|
||
import redis
|
||
import traceback
|
||
|
||
# ==========================================
|
||
# ⚠️ 入口说明
|
||
# ==========================================
|
||
# 当前 config.yaml 的 callbacks 使用的是 volc_limit_sync.proxy_handler_instance，
|
||
# 不是本文件。除非你把 callbacks 切回 daily_limit.py，否则这里的 MODEL_QUOTAS
|
||
# 修改不会影响当前 Docker 容器的监控与限额行为。
|
||
|
||
# ==========================================
|
||
# 🛠️ 调试配置
|
||
# ==========================================
|
||
# 日志级别配置
|
||
# 可选值: "DEBUG", "INFO", "WARNING", "ERROR"
|
||
# 修改此变量以控制日志详细程度，或使用环境变量 DAILY_LIMIT_LOG_LEVEL
|
||
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
|
||
|
||
|
||
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} [DailyLimit] {msg}", flush=True)
|
||
|
||
|
||
def dprint(msg):
|
||
log(msg, level="INFO")
|
||
|
||
|
||
dprint("Loading daily_limit.py module v20 (legacy entry; current config uses volc_limit_sync.py)...")
|
||
|
||
# ==========================================
|
||
# Redis 与 限额配置
|
||
# ==========================================
|
||
REDIS_HOST = os.getenv("REDIS_HOST", "litellm_redis")
|
||
REDIS_PORT = os.getenv("REDIS_PORT", "6379")
|
||
REDIS_PASSWORD = os.getenv("REDIS_PASSWORD", None)
|
||
|
||
MODEL_QUOTAS = {
|
||
"deepseek-v3-2-251201:ba0383d3": 1600000,
|
||
"doubao-seed-1-8-251228:ba0383d3": 1800000,
|
||
"glm-4-7-251222:ba0383d3": 1800000,
|
||
"doubao-seed-2-0-pro-260215:ba0383d3": 1900000,
|
||
"doubao-seed-2-0-lite-260215:ba0383d3": 1900000,
|
||
}
|
||
|
||
|
||
try:
|
||
redis_host = os.getenv("REDIS_HOST", "redis")
|
||
redis_port = int(os.getenv("REDIS_PORT", 6379))
|
||
redis_client = redis.Redis(host=redis_host, port=redis_port, db=0, decode_responses=True)
|
||
log(f"✅ Redis connected successfully to {redis_host}:{redis_port}", level="INFO")
|
||
except Exception as e:
|
||
log(f"❌ Redis connection failed: {str(e)}", level="ERROR")
|
||
redis_client = None
|
||
|
||
|
||
def clean_model_name(name):
|
||
if not name:
|
||
return ""
|
||
return name.replace("openai/", "").replace("azure/", "").replace("anthropic/", "")
|
||
|
||
|
||
def generate_api_base_hash(api_base):
|
||
"""
|
||
生成 api_base 的 Hash，强制忽略 URL 的路径部分。
|
||
例如:
|
||
'https://api.com/v1/' -> 'https://api.com'
|
||
'https://api.com/v3' -> 'https://api.com'
|
||
两者生成的 Hash 将完全一致。
|
||
"""
|
||
if not api_base:
|
||
return "default"
|
||
|
||
try:
|
||
# 1. 预处理：去空格
|
||
s = api_base.strip()
|
||
|
||
# 2. 解析 URL
|
||
parsed = urlparse(s)
|
||
|
||
# 3. 提取 Scheme 和 Netloc (域名+端口)
|
||
# 如果解析成功且包含 scheme 和 netloc，则重组
|
||
if parsed.scheme and parsed.netloc:
|
||
clean_base = f"{parsed.scheme}://{parsed.netloc}".lower()
|
||
else:
|
||
# 兜底：如果不是标准 URL，或者解析失败，尝试简单的去斜杠和小写化
|
||
# 这种情况比较少见，通常 litellm 会校验 URL 格式
|
||
clean_base = s.rstrip('/').lower()
|
||
|
||
except Exception:
|
||
# 极端情况下的兜底
|
||
clean_base = api_base.strip().rstrip('/').lower()
|
||
|
||
# 4. 计算 Hash
|
||
h = hashlib.md5(clean_base.encode()).hexdigest()[:8]
|
||
|
||
if DEBUG_MODE:
|
||
# Debugging hash mismatch
|
||
# 打印原始值和清洗后的值，方便确认
|
||
print(f"🔍 [Hash Debug] Base='{api_base}' -> Clean='{clean_base}' -> Hash='{h}'", flush=True)
|
||
|
||
return h
|
||
|
||
|
||
# ==========================================
|
||
# ⏰ 时间逻辑：北京时间 9:00 (UTC 01:00) 刷新
|
||
# ==========================================
|
||
def get_redis_key(clean_name, api_base=None):
|
||
now = time.time()
|
||
# UTC 01:00 视为新的一天，所以减去 1小时 (3600秒)
|
||
adjusted_time = now - 3600
|
||
date_str = time.strftime("%Y-%m-%d", time.gmtime(adjusted_time))
|
||
|
||
suffix = ""
|
||
if api_base:
|
||
suffix = f":{generate_api_base_hash(api_base)}"
|
||
|
||
return f"daily_usage:{date_str}:{clean_name}{suffix}"
|
||
|
||
|
||
# ==========================================
|
||
# 🎯 限额查找逻辑 (支持精确匹配)
|
||
# ==========================================
|
||
def get_daily_limit(clean_model, api_base):
|
||
# 1. 尝试精确匹配 (Model + Hash)
|
||
# 用户可以在 MODEL_QUOTAS 中配置 "model_name:hash": limit
|
||
current_hash = generate_api_base_hash(api_base)
|
||
specific_key = f"{clean_model}:{current_hash}"
|
||
|
||
limit = MODEL_QUOTAS.get(specific_key)
|
||
if limit:
|
||
return limit
|
||
|
||
# 2. 尝试精确模型名匹配
|
||
limit = MODEL_QUOTAS.get(clean_model)
|
||
if limit:
|
||
return limit
|
||
|
||
# 3. 尝试模糊模型名匹配 (向下兼容)
|
||
for quota_model, limit_val in MODEL_QUOTAS.items():
|
||
# 排除包含冒号的特定规则，防止误匹配
|
||
if ":" not in quota_model and quota_model in clean_model:
|
||
return limit_val
|
||
|
||
return None
|
||
|
||
|
||
# ==========================================
|
||
# 👻 核心逻辑：软更名踢出 & 自动恢复
|
||
# ==========================================
|
||
def soft_evict_from_router(target_upstream_model, target_api_base=None):
|
||
"""
|
||
不删除节点，而是修改节点的 model_name，让 Router 找不到它。
|
||
同时备份原始信息以便恢复。
|
||
"""
|
||
try:
|
||
from litellm.proxy.proxy_server import llm_router
|
||
|
||
if not llm_router or not llm_router.model_list:
|
||
return
|
||
|
||
evicted = False
|
||
target_upstream_clean = clean_model_name(target_upstream_model)
|
||
|
||
# Iterate over the actual list in memory
|
||
for i, model_entry in enumerate(llm_router.model_list):
|
||
params = model_entry.get("litellm_params", {})
|
||
upstream_model = params.get("model", "")
|
||
current_api_base = params.get("api_base", "")
|
||
|
||
# 匹配上游模型名称
|
||
if target_upstream_clean not in clean_model_name(upstream_model):
|
||
continue
|
||
|
||
# Strict api_base matching if provided
|
||
# ⚠️ 注意：这里比较的是原始字符串，为了确保精准匹配当前运行的实例
|
||
if target_api_base and current_api_base != target_api_base:
|
||
continue
|
||
|
||
current_alias = model_entry.get("model_name")
|
||
|
||
# 如果它还没有被标记为 blocked
|
||
if current_alias and "blocked" not in current_alias:
|
||
log(
|
||
f"👻 [ROUTER HIDE] Renaming alias '{current_alias}' -> '{current_alias}-blocked' for {target_upstream_model}",
|
||
level="INFO",
|
||
)
|
||
|
||
# Backup original state
|
||
model_entry["_daily_limit_backup"] = {"model_name": current_alias, "api_base": current_api_base}
|
||
|
||
# 1. 修改别名 - Modify IN PLACE
|
||
llm_router.model_list[i]["model_name"] = f"{current_alias}-blocked"
|
||
|
||
# 2. 破坏连接
|
||
llm_router.model_list[i]["litellm_params"]["api_base"] = "http://0.0.0.0"
|
||
|
||
evicted = True
|
||
|
||
if evicted:
|
||
# Re-initialize router to reflect changes immediately if possible
|
||
# But changing model_list in place is usually enough for next lookup
|
||
dprint("✅ Node effectively removed from routing pool.")
|
||
|
||
except Exception as e:
|
||
print(f"❌ Router eviction failed: {e}", flush=True)
|
||
|
||
|
||
last_restore_check = 0
|
||
RESTORE_CHECK_INTERVAL = 60 # Check every 60 seconds
|
||
|
||
|
||
def restore_blocked_models():
|
||
"""
|
||
检查是否有 blocked 的模型配额已重置，如果是则恢复。
|
||
"""
|
||
global last_restore_check
|
||
now = time.time()
|
||
if now - last_restore_check < RESTORE_CHECK_INTERVAL:
|
||
return
|
||
|
||
last_restore_check = now
|
||
|
||
try:
|
||
from litellm.proxy.proxy_server import llm_router
|
||
|
||
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")
|
||
|
||
# Check quota
|
||
daily_limit = get_daily_limit(clean_model, original_api_base)
|
||
|
||
should_restore = False
|
||
if not daily_limit:
|
||
should_restore = True
|
||
else:
|
||
key = get_redis_key(clean_model, original_api_base)
|
||
current_usage = redis_client.get(key)
|
||
current_val = float(current_usage) if current_usage else 0.0
|
||
|
||
if current_val < float(daily_limit):
|
||
should_restore = True
|
||
|
||
if should_restore:
|
||
dprint(f"♻️ [RESTORE] Restoring {backup['model_name']} (Quota OK)")
|
||
|
||
# Restore IN PLACE
|
||
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 to routing pool.")
|
||
|
||
except Exception as e:
|
||
log(f"❌ Restore check failed: {e}", level="ERROR")
|
||
|
||
|
||
# ==========================================
|
||
# Input Callback
|
||
# ==========================================
|
||
def quota_interceptor(kwargs, user_api_key_dict=None):
|
||
"""
|
||
Standard LiteLLM input callback signature:
|
||
def input_callback(kwargs, user_api_key_dict):
|
||
"""
|
||
try:
|
||
restore_blocked_models()
|
||
# check_and_hide(kwargs, source="input_callback") # Removed to avoid rigid blocking
|
||
except (ServiceUnavailableError, ContextWindowExceededError):
|
||
raise
|
||
except Exception:
|
||
# Don't let other errors block the request
|
||
pass
|
||
|
||
|
||
# ==========================================
|
||
# Handler & Logging
|
||
# ==========================================
|
||
class DailyLimitHandler(CustomLogger):
|
||
def __init__(self):
|
||
# Register the callback properly
|
||
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_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")
|
||
|
||
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_redis_key(clean_model, api_base)
|
||
current_total = redis_client.incrby(key, int(total_tokens))
|
||
redis_client.expire(key, 172800)
|
||
|
||
# Get Limit for display
|
||
daily_limit = get_daily_limit(clean_model, api_base)
|
||
display_limit = f"/{int(daily_limit):,}" if daily_limit else ""
|
||
|
||
# Show usage if DEBUG level is enabled
|
||
key_display = key.split(":", 2)[-1]
|
||
log(
|
||
f"📊 [Usage] {key_display}: +{total_tokens} (Today: {current_total:,}{display_limit})",
|
||
level="INFO",
|
||
)
|
||
|
||
# Enforcement Logic
|
||
if daily_limit and current_total >= float(daily_limit):
|
||
log(f"🚫 [BLOCK] {clean_model} Limit: {current_total}/{daily_limit}", level="WARNING")
|
||
soft_evict_from_router(raw_model, api_base)
|
||
else:
|
||
if DEBUG_MODE:
|
||
if not clean_model:
|
||
log("⚠️ [Usage Skip] No clean model name found", level="DEBUG")
|
||
if total_tokens <= 0:
|
||
log(f"⚠️ [Usage Skip] Total tokens is {total_tokens}", level="DEBUG")
|
||
|
||
except Exception as e:
|
||
log(f"❌ _record_usage failed: {e}\n{traceback.format_exc()}", level="ERROR")
|
||
|
||
|
||
proxy_handler_instance = DailyLimitHandler()
|