import os import sys import time import json import hmac 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 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 # ========================================== # ๐Ÿ› ๏ธ 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 BUFFER = 300000 # User requested 30w buffer SYNC_INTERVAL = 60 CHECK_INTERVAL = 10 THRESHOLD = 50000 # 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 = [ "_data_collaboration_resource_pack", "ๅไฝœๅฅ–ๅŠฑ่ฎกๅˆ’่ต„ๆบๅŒ…", "ๅ…่ดนๅœจ็บฟๆŽจ็†่ต„ๆบๅŒ…", "ๅ…่ดนๅœจ็บฟๆŽจ็†ๅŒ…", "่ต„ๆบๅŒ…", ] # Volcengine Name -> LiteLLM Model Name Map VOLC_MODEL_MAP = { "doubao-seed-1-6": ["doubao-seed-1-6-250615", "doubao-seed-1-6-251015"], "doubao-seed-1-8": "doubao-seed-1-8-251228", "doubao-seed-code": "doubao-seed-code-preview-251028", "doubao-seed-translation": "doubao-seed-translation-250915", "deepseek-v3": "deepseek-v3-250324", "deepseek-v3-1": "deepseek-v3-1-terminus", "deepseek-v3-2": "deepseek-v3-2-251201", "deepseek-r1": "deepseek-r1-250528", "kimi-k2": "kimi-k2-thinking-251104", "glm-4_7": "glm-4-7-251222", "glm-4-7": "glm-4-7-251222", } # Static Quotas (Fallback/Strict Limits) MODEL_QUOTAS = { # "deepseek-v3-2-251201:ba0383d3": 1600000, # "doubao-seed-1-8-251228:ba0383d3": 1800000, # "glm-4-7-251222:ba0383d3": 1800000, # "doubao-seed-1-8-251228:d57edef4": 1500000, } # ========================================== # ๐Ÿ”Œ 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} [VolcSync] {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/", "") def normalize_volc_name(raw_name): """Volcengine API response cleaning.""" 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 VOLC_MODEL_MAP.get(clean_name, 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): now = time.time() adjusted_time = now - 3600 # UTC+1 split date_str = time.strftime("%Y-%m-%d", time.gmtime(adjusted_time)) 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(clean_model, api_base): """Check static configuration in MODEL_QUOTAS.""" current_hash = generate_api_base_hash(api_base) # 1. Exact Match with Hash specific_key = f"{clean_model}:{current_hash}" if specific_key in MODEL_QUOTAS: return MODEL_QUOTAS[specific_key] # 2. Exact Model Name if clean_model in MODEL_QUOTAS: return MODEL_QUOTAS[clean_model] # 3. Fuzzy Match 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 get_effective_limit(clean_model, api_base, current_usage=None): """ Calculate effective limit combining Static Quotas and Dynamic Redis Quotas. Dynamic Limit = Current Usage + Redis Remaining Balance """ if not redis_client: return get_static_limit(clean_model, api_base) # 1. Get Static Limit static_limit = get_static_limit(clean_model, api_base) # 2. Get Dynamic Remaining Balance # Key set by daemon loop: limit:dynamic:volcengine:remaining:{clean_model} dynamic_key = f"limit:dynamic:volcengine:remaining:{clean_model}" dynamic_remaining = redis_client.get(dynamic_key) effective_limit = None # If we have dynamic data if dynamic_remaining is not None: try: # If usage not provided, fetch it if current_usage is None: usage_key = get_usage_key(clean_model, api_base) current_usage = float(redis_client.get(usage_key) or 0) # Dynamic Limit = The point where usage would consume all remaining balance # limit = usage + remaining dynamic_limit_val = float(current_usage) + float(dynamic_remaining) # If static limit exists, take the stricter one (minimum) # UNLESS static is much higher? No, min implies stricter. if static_limit: effective_limit = min(float(static_limit), dynamic_limit_val) else: effective_limit = dynamic_limit_val except (ValueError, TypeError): effective_limit = static_limit else: effective_limit = static_limit return effective_limit # ========================================== # ๐Ÿ‘ป Handler Logic (DailyLimitHandler) # ========================================== def soft_evict_from_router(target_upstream_model, target_api_base=None): 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) 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 if target_api_base and current_api_base != target_api_base: continue current_alias = model_entry.get("model_name") if current_alias and "blocked" not in current_alias: log( f"๐Ÿ‘ป [ROUTER HIDE] Blocking '{current_alias}' for {target_upstream_model}", level="INFO", ) model_entry["_daily_limit_backup"] = { "model_name": current_alias, "api_base": current_api_base, } llm_router.model_list[i]["model_name"] = f"{current_alias}-blocked" llm_router.model_list[i]["litellm_params"]["api_base"] = ( "http://0.0.0.0" ) evicted = True if evicted: 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 def restore_blocked_models(): 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 # 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 not effective_limit: 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 quota_interceptor(kwargs, user_api_key_dict=None): try: 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_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_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 ) display_limit = f"/{int(effective_limit):,}" if effective_limit else "" # Get Dynamic Remaining for Log dynamic_key = f"limit:dynamic:volcengine:remaining:{clean_model}" rem_val = redis_client.get(dynamic_key) rem_display = "" if rem_val is not None: try: # Display Real Balance (Remaining + Buffer) rem_int = int(float(rem_val)) bal_est = rem_int + BUFFER rem_display = f", Bal: {bal_est:,}" 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: Dynamic key not found: {dynamic_key}") key_display = key.split(":", 2)[-1] log( f"๐Ÿ“Š [Usage] {key_display}: +{total_tokens} (Today: {current_total:,}{display_limit}{rem_display})", level="INFO", ) if effective_limit 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 = "" while True: try: 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() items = data.get("Result", {}).get("List", []) if items: all_packages.extend(items) next_token = data.get("Result", {}).get("NextToken", "") if not next_token or next_token == current_token: break 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 = {} print(f"\n--- Raw Package Processing ---", flush=True) for pkg in packages: try: raw_code = pkg.get("ConfigurationCode", "") raw_name = pkg.get("ConfigurationName", "Unknown") target_raw = raw_code if raw_code else raw_name model_name = normalize_volc_name(target_raw) available = pkg.get("AvailableAmount", 0) unit = pkg.get("Unit", "token") # Debug log to trace mapping if DEBUG_MODE or "glm" in model_name or "doubao" in model_name: print( f"DEBUG: Pkg Raw='{target_raw}' -> Norm='{model_name}', Avail={available}", flush=True, ) if unit.lower() == "token": model_counts[model_name] = model_counts.get(model_name, 0) + int( available ) except Exception as e: print(f"Error processing package: {e}") continue return model_counts def active_monitor_loop(): # Config Discovery 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") print("Auto-discovering Volcengine models from config.yaml...") # Robust Parsing Logic current_entry = {"is_volc": False, "upstream_model": None} for line in lines: sline = line.strip() # 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" ): INTERESTED_MODELS.add(current_entry["upstream_model"]) # Reset for new entry current_entry = {"is_volc": False, "upstream_model": None} continue # Check for Volcengine Base URL if "api_base:" in sline and "HUOSHAN_BASE_URL" in sline: current_entry["is_volc"] = True # 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/", "") ) current_entry["upstream_model"] = clean except IndexError: pass # Commit the last entry if current_entry.get("is_volc") and current_entry.get("upstream_model"): INTERESTED_MODELS.add(current_entry["upstream_model"]) print( f"Auto-discovered {len(INTERESTED_MODELS)} Active Volcengine models (Upstream Names): {INTERESTED_MODELS}" ) except Exception as e: print(f"Config Auto-discovery failed: {e}") if not redis_client: print("โŒ Redis not connected. Exiting monitor loop.") sys.exit(1) print(f"โœ… Redis connected at {REDIS_HOST}:{REDIS_PORT}") 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)...") 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") # Idle Check usage_key_global = f"daily_usage:{today_str}:global_total" current_usage_global = int(redis_client.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 if delta <= 0: 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 # Logic: Adjust intervals based on distance from BUFFER # User specified thresholds (Distance < X): # < 10w -> 2min (120s) # < 30w -> 5min (300s) # < 50w -> 10min (600s) # < 80w -> 30min (1800s) # < 100w -> 60min (3600s) # >= 100w -> 60min (3600s) if global_min_remaining < 100_000: base_interval = 120 elif global_min_remaining < 300_000: base_interval = 300 elif global_min_remaining < 500_000: base_interval = 600 elif global_min_remaining < 800_000: base_interval = 1800 else: base_interval = 3600 if is_idle: current_sync_interval = min(base_interval * 4, 3600) if current_sync_interval < 300: current_sync_interval = 300 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" ) print( f"{'๐Ÿ’ค' if is_idle else 'โšก'} Sync Triggered ({trigger_reason})... Delta: {delta}, Interval: {current_sync_interval}s" ) 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(): remaining = balance - BUFFER print( f"DEBUG: Checking {model_name} (Bal: {balance}, Rem: {remaining})", flush=True, ) # Target Resolution targets = [] if model_name in VOLC_MODEL_MAP: val = VOLC_MODEL_MAP[model_name] targets = val if isinstance(val, list) else [val] if not targets: targets = [model_name] # Empty string case else: targets = [model_name] for target in targets: # Only set keys for models we are actually using (found in config) # BUT we must set ALL keys to ensure we don't miss any aliases. # Optimization: Check if target is in INTERESTED_MODELS if target in INTERESTED_MODELS: configured_models_found = True # Use normalized key for storage # Format: limit:dynamic:volcengine:remaining:{model_name} dynamic_key = f"limit:dynamic:volcengine:remaining:{target}" # Set with expiry (slightly longer than max interval) redis_client.setex(dynamic_key, 7200, int(remaining)) print( f"DEBUG: Set Redis {dynamic_key} = {int(remaining)}", flush=True, ) if remaining < temp_min_remaining: temp_min_remaining = remaining else: if DEBUG_MODE: print( f"DEBUG: Skipping {target} (Not in INTERESTED_MODELS)", flush=True, ) global_min_remaining = temp_min_remaining if model_balances else 0 last_usage_map["global"] = current_usage_global last_sync_time = now 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}") traceback.print_exc() 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()