788 lines
28 KiB
Python
788 lines
28 KiB
Python
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()
|