feat: implement startup audit to block over-limit models on restart
- Add startup_audit() to check Redis usage against limits at startup - Call startup_audit() in active_monitor_loop - Add AGENTS.md for agent guidelines
This commit is contained in:
@@ -470,8 +470,73 @@ def restore_blocked_models():
|
||||
log(f"❌ Restore check failed: {e}", level="ERROR")
|
||||
|
||||
|
||||
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
|
||||
|
||||
# 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 not effective_limit:
|
||||
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 quota_interceptor(kwargs, user_api_key_dict=None):
|
||||
try:
|
||||
# 1. 尝试恢复(定期)
|
||||
restore_blocked_models()
|
||||
except (ServiceUnavailableError, ContextWindowExceededError):
|
||||
raise
|
||||
@@ -832,7 +897,15 @@ def active_monitor_loop():
|
||||
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
|
||||
|
||||
Reference in New Issue
Block a user