932 lines
40 KiB
Python
932 lines
40 KiB
Python
import os, json, time, logging, traceback, threading, queue, statistics, math
|
|
import requests, redis
|
|
from flask import Flask, request, jsonify, Response, stream_with_context
|
|
|
|
|
|
|
|
# Phase 2: Atomic session token update Redis Lua script
|
|
SESSION_LUA_SCRIPT = """
|
|
local key = KEYS[1]
|
|
local new_val = tonumber(ARGV[1])
|
|
local current = tonumber(redis.call('GET', key) or 0)
|
|
local max_val = math.max(current, new_val)
|
|
redis.call('SET', key, max_val, 'EX', 86400)
|
|
return max_val
|
|
"""
|
|
|
|
|
|
REDIS_URL = os.environ.get("REDIS_URL", "redis://redis:6379")
|
|
GPU_MOE_URL = os.environ.get("GPU_MOE_URL", "http://192.168.68.15:8080/v1")
|
|
GPU_DENSE_URL = os.environ.get("GPU_DENSE_URL", "http://192.168.68.8:8080/v1")
|
|
GPU_LIGHT_URL = os.environ.get("GPU_LIGHT_URL", "http://192.168.68.110:8080/v1")
|
|
|
|
GPU_SIDECARS = {
|
|
"qwen3.6-35B-A3B": "http://192.168.68.15:8090",
|
|
"qwen3.6-27B-code": "http://192.168.68.8:8090",
|
|
"gemma-4-12b": "http://192.168.68.110:8090",
|
|
}
|
|
GPU_URLS = {
|
|
"qwen3.6-35B-A3B": GPU_MOE_URL,
|
|
"qwen3.6-27B-code": GPU_DENSE_URL,
|
|
"gemma-4-12b": GPU_LIGHT_URL,
|
|
}
|
|
# Max concurrent requests per GPU (based on llama.cpp --parallel)
|
|
GPU_MAX_CONCURRENT = {
|
|
"qwen3.6-35B-A3B": 2, # 2 slots (cross-agent spread prevents overheating)
|
|
"qwen3.6-27B-code": 2, # 2 slots (128K context frees VRAM)
|
|
"gemma-4-12b": 2, # 2 slots (7.1GB VRAM)
|
|
}
|
|
|
|
# Context window sizes (tokens) — used for compaction signals
|
|
GPU_CONTEXT = {
|
|
"qwen3.6-35B-A3B": 262144,
|
|
"qwen3.6-27B-code": 262144,
|
|
"gemma-4-12b": 262144,
|
|
}
|
|
|
|
TIER_MODELS = {
|
|
"starter": ["gemma-4-12b"],
|
|
"professional": ["qwen3.6-35B-A3B", "qwen3.6-27B-code", "gemma-4-12b"],
|
|
"enterprise": ["qwen3.6-35B-A3B", "qwen3.6-27B-code", "gemma-4-12b"],
|
|
}
|
|
# API keys loaded from API_KEYS env var (set in docker-compose.yml)
|
|
# Fallback is dev-only — production MUST set API_KEYS env var
|
|
API_KEYS = json.loads(os.environ.get("API_KEYS", json.dumps({
|
|
"sk-dev-local-only": {"tier": "enterprise", "agent": "dev"},
|
|
})))
|
|
# Rate limits: requests per minute per API key tier
|
|
RATE_LIMIT_RPM = {
|
|
"enterprise": 120,
|
|
"professional": 60,
|
|
"starter": 20,
|
|
}
|
|
|
|
def check_rate_limit(api_key, tier):
|
|
"""Token bucket rate limiter using Redis. Returns (allowed, retry_after_or_remaining, reset_seconds)."""
|
|
if not get_redis():
|
|
return True, 999, 60
|
|
limit = RATE_LIMIT_RPM.get(tier, 30)
|
|
key = f"ratelimit:{api_key}"
|
|
current = int(get_redis().get(key) or 0)
|
|
if current >= limit:
|
|
ttl = r.ttl(key)
|
|
retry = max(ttl, 1) if ttl and ttl > 0 else 60
|
|
return False, retry, 0
|
|
pipe = r.pipeline()
|
|
pipe.incr(key)
|
|
pipe.expire(key, 60) # 1-minute sliding window
|
|
pipe.execute()
|
|
remaining = limit - (current + 1)
|
|
reset_seconds = r.ttl(key) or 60
|
|
return True, remaining, reset_seconds
|
|
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s [ROUTER] %(levelname)s %(message)s")
|
|
log = logging.getLogger("router")
|
|
# Redis connection — initialized lazily, retries on first use
|
|
def get_redis():
|
|
global r
|
|
if r is not None:
|
|
try:
|
|
r.ping()
|
|
return r
|
|
except Exception:
|
|
r = None
|
|
try:
|
|
r = redis.from_url(REDIS_URL, decode_responses=True)
|
|
r.ping()
|
|
return r
|
|
except Exception:
|
|
return None
|
|
|
|
r = None
|
|
try: get_redis()
|
|
except Exception: pass
|
|
except Exception: r = None
|
|
|
|
|
|
def counter_audit_loop():
|
|
"""Every 30s, check GPU slots and reset counters if all slots idle."""
|
|
while True:
|
|
time.sleep(30)
|
|
if not get_redis(): continue
|
|
for model, url in GPU_URLS.items():
|
|
try:
|
|
resp = requests.get(url.replace("/v1","") + "/slots",
|
|
headers={"Authorization": "Bearer not-needed"}, timeout=5)
|
|
if resp.status_code == 200:
|
|
slots = resp.json()
|
|
all_idle = all(not s.get("is_processing", False) for s in slots)
|
|
if all_idle:
|
|
current = int(get_redis().get("active:" + model) or 0)
|
|
if current > 0:
|
|
get_redis().set("active:" + model, 0)
|
|
log.info("AUDIT: Reset stuck counter for %s (was %d)", model, current)
|
|
except Exception:
|
|
pass
|
|
|
|
threading.Thread(target=counter_audit_loop, daemon=True).start()
|
|
|
|
app = Flask(__name__)
|
|
sse_subscribers = []; sse_lock = threading.Lock()
|
|
|
|
def gpu_active_count(model):
|
|
"""Get number of in-flight requests for a GPU."""
|
|
if r:
|
|
return int(r.get("active:" + model) or 0)
|
|
return 0
|
|
|
|
def gpu_incr(model):
|
|
if get_redis(): get_redis().incr("active:" + model)
|
|
|
|
def gpu_decr(model):
|
|
rd = get_redis()
|
|
if rd:
|
|
v = rd.decr("active:" + model)
|
|
if v and int(v) < 0:
|
|
get_redis().set("active:" + model, 0) # never go negative
|
|
def check_gpu_health(model, sidecar_timeout=5, gpu_timeout=3):
|
|
url = GPU_SIDECARS.get(model)
|
|
if not url: return {"status": "unknown"}
|
|
try:
|
|
resp = requests.get(url, timeout=sidecar_timeout)
|
|
if resp.status_code == 200:
|
|
d = resp.json()
|
|
pct = (d.get("vram_used_mb",0) / max(d.get("vram_total_mb",1), 1)) * 100
|
|
status = "healthy" # VRAM usage != saturation; busy slots handled by is_gpu_busy()
|
|
vram_warning = pct >= 95
|
|
# Also check if llama.cpp endpoint is actually responding
|
|
gpu_url = GPU_URLS.get(model, "")
|
|
try:
|
|
hr = requests.get(gpu_url.replace("/v1","") + "/health", headers={"Authorization": "Bearer not-needed"}, timeout=gpu_timeout)
|
|
if hr.status_code != 200:
|
|
status = "down"
|
|
except Exception:
|
|
status = "down"
|
|
return {"status": status, "vram_warning": vram_warning, "vram_used_mb": d.get("vram_used_mb"), "vram_total_mb": d.get("vram_total_mb"), "vram_pct": round(pct,1), "temp_c": d.get("temp_c"), "gpu_util_pct": d.get("gpu_util_pct"), "gpu_name": d.get("gpu_name"), "power_w": d.get("power_w"), "power_limit_w": d.get("power_limit_w")}
|
|
except Exception: pass
|
|
return {"status": "down"}
|
|
|
|
def available_models(): return [m for m in GPU_URLS if check_gpu_health(m)["status"] in ("healthy","saturated")]
|
|
|
|
def estimate_tokens(msgs):
|
|
"""Estimate token count from messages. Uses JSON length / 3.5 (closer to real tokenizer ratios for dense text)."""
|
|
return len(json.dumps(msgs, default=str)) // 3.5
|
|
|
|
def store_perf_record(model, agent, tier, reason, queue_ms, inference_ms, prompt_tokens, completion_tokens, stream):
|
|
"""Store detailed performance record in Redis for analytics."""
|
|
if not get_redis(): return
|
|
try:
|
|
total_ms = queue_ms + inference_ms
|
|
tps = completion_tokens / (inference_ms / 1000) if inference_ms > 0 and completion_tokens > 0 else 0
|
|
rec = json.dumps({
|
|
"ts": time.time(),
|
|
"model": model, "agent": agent, "tier": tier, "reason": reason,
|
|
"queue_ms": round(queue_ms, 1),
|
|
"inference_ms": round(inference_ms, 1),
|
|
"total_ms": round(total_ms, 1),
|
|
"prompt_tokens": prompt_tokens,
|
|
"completion_tokens": completion_tokens,
|
|
"tokens_per_sec": round(tps, 1),
|
|
"stream": stream
|
|
})
|
|
# Global recent list (last 500)
|
|
r.lpush("perf:recent", rec)
|
|
r.ltrim("perf:recent", 0, 499)
|
|
# Per-model list (last 200)
|
|
r.lpush("perf:model:" + model, rec)
|
|
r.ltrim("perf:model:" + model, 0, 199)
|
|
# Per-reason list (last 200)
|
|
r.lpush("perf:reason:" + reason, rec)
|
|
r.ltrim("perf:reason:" + reason, 0, 199)
|
|
# Per-agent list (last 200)
|
|
r.lpush("perf:agent:" + agent, rec)
|
|
r.ltrim("perf:agent:" + agent, 0, 199)
|
|
except Exception:
|
|
pass
|
|
|
|
def is_gpu_busy(model):
|
|
"""Check if GPU is at or near max concurrent capacity."""
|
|
active = gpu_active_count(model)
|
|
max_c = GPU_MAX_CONCURRENT.get(model, 1)
|
|
return active >= max_c
|
|
|
|
def select_best_gpu(candidates, reason, agent=""):
|
|
"""Pick best GPU, spreading agents across GPUs to prevent hotspots."""
|
|
# Count how many distinct agents are on each GPU
|
|
gpu_agent_counts = {}
|
|
if r:
|
|
for m in GPU_URLS:
|
|
count = 0
|
|
for ak in API_KEYS.values():
|
|
if r.get("agent_gpu:" + ak["agent"] + ":" + m):
|
|
count += 1
|
|
gpu_agent_counts[m] = count
|
|
# First pass: prefer GPUs with 0 other agents (fresh GPU for this agent)
|
|
for m in candidates:
|
|
if not is_gpu_busy(m) and gpu_agent_counts.get(m, 0) == 0:
|
|
return {"model": m, "reason": reason}
|
|
# Second pass: prefer GPU this agent is NOT already on (skip own GPU)
|
|
if agent:
|
|
for m in candidates:
|
|
if not is_gpu_busy(m) and not r.get("agent_gpu:" + agent + ":" + m):
|
|
return {"model": m, "reason": reason}
|
|
# Third pass: any non-busy GPU
|
|
for m in candidates:
|
|
if not is_gpu_busy(m):
|
|
return {"model": m, "reason": reason}
|
|
# All busy — pick least loaded
|
|
best = None
|
|
best_load = 999
|
|
for m in candidates:
|
|
load = gpu_active_count(m)
|
|
if load < best_load:
|
|
best_load = load
|
|
best = m
|
|
if best:
|
|
return {"model": best, "reason": "load_balanced_" + reason}
|
|
return None
|
|
|
|
|
|
|
|
|
|
|
|
# Phase 1: Circuit Breaker for GPU Hosts (Approved by Abiba)
|
|
def is_circuit_tripped(model):
|
|
"""Check if a GPU host is currently blacklisted."""
|
|
if not get_redis():
|
|
return False
|
|
return r.exists("circuit:" + model + ":open")
|
|
|
|
def trip_circuit(model, duration=30):
|
|
"""Blacklist a GPU host for a specified duration."""
|
|
if not get_redis():
|
|
return
|
|
key = "circuit:" + model + ":open"
|
|
r.set(key, 1, ex=duration)
|
|
r.incr("circuit:" + model + ":count")
|
|
log.warning("CIRCUIT_TRIPPED: %s blacklisted for %ds", model, duration)
|
|
|
|
def half_open_probe(model):
|
|
"""Check if a GPU host can be un-blacklisted."""
|
|
if not get_redis():
|
|
return True
|
|
key = "circuit:" + model + ":open"
|
|
if not r.exists(key):
|
|
return True # no circuit
|
|
return False # still open
|
|
|
|
def route(rd, tier, agent=""):
|
|
msgs = rd.get("messages",[]); t = estimate_tokens(msgs)
|
|
sys = any(m.get("role")=="system" for m in msgs)
|
|
turns = len([m for m in msgs if m.get("role") in ("user","assistant")])
|
|
hints = rd.get("routing_hints",{})
|
|
allowed = TIER_MODELS.get(tier, ["gemma-4-12b"])
|
|
# Phase 1: Filter out models with tripped circuit breakers
|
|
avail = [m for m in available_models() if m in allowed and not is_circuit_tripped(m)]
|
|
if not avail: return {"model": allowed[0], "reason": "all_saturated", "saturated": True}
|
|
if all(is_gpu_busy(m) for m in avail):
|
|
return {"model": avail[0], "reason": "all_saturated", "saturated": True}
|
|
|
|
# GUARD: multimodal -> VLM only (sole vision model)
|
|
has_image = any(
|
|
isinstance(m.get("content"), list) and
|
|
any(p.get("type") == "image_url" for p in m["content"] if isinstance(p, dict))
|
|
for m in msgs
|
|
)
|
|
if has_image:
|
|
if "gemma-4-12b" in avail and not is_gpu_busy("gemma-4-12b"):
|
|
return {"model": "gemma-4-12b", "reason": "vision"}
|
|
elif "gemma-4-12b" in avail:
|
|
return {"model": "gemma-4-12b", "reason": "vision_saturated", "saturated": True}
|
|
else:
|
|
return {"model": allowed[0], "reason": "vision_unavailable"}
|
|
|
|
req = rd.get("model","auto")
|
|
if req != "auto":
|
|
target = req if req in avail else avail[0]
|
|
if is_gpu_busy(target) and req in allowed:
|
|
alts = [m for m in avail if m != target and m in allowed]
|
|
if alts:
|
|
alt = select_best_gpu(alts, "explicit", agent)
|
|
if alt: return alt
|
|
return {"model": target, "reason": "explicit"}
|
|
|
|
if hints:
|
|
if hints.get("priority")=="speed" and "gemma-4-12b" in avail:
|
|
return select_best_gpu(["gemma-4-12b"], "hint_speed", agent) or {"model":"gemma-4-12b","reason":"hint_speed"}
|
|
if hints.get("priority")=="quality" and "qwen3.6-35B-A3B" in avail:
|
|
return select_best_gpu(["qwen3.6-35B-A3B"], "hint_quality", agent) or {"model":"qwen3.6-35B-A3B","reason":"hint_quality"}
|
|
if hints.get("priority")=="code" and "qwen3.6-27B-code" in avail:
|
|
return select_best_gpu(["qwen3.6-27B-code"], "hint_code", agent) or {"model":"qwen3.6-27B-code","reason":"hint_code"}
|
|
|
|
first_msg = msgs[0].get("content","") if msgs else ""
|
|
words = len(first_msg.split()) if isinstance(first_msg, str) else 99
|
|
|
|
# TIER 1: Tiny - single-turn micro queries -> VLM (fastest)
|
|
if not sys and turns <= 1 and t <= 300 and words <= 100 and "gemma-4-12b" in avail:
|
|
if not is_gpu_busy("gemma-4-12b"):
|
|
return {"model":"gemma-4-12b","reason":"tiny"}
|
|
fallback = [m for m in ["qwen3.6-27B-code","qwen3.6-35B-A3B"] if m in avail]
|
|
result = select_best_gpu(fallback, "tiny_fallback", agent)
|
|
if result: return result
|
|
|
|
# TIER 2: Light - moderate chat -> VLM first (fastest), Dense fallback
|
|
if t <= 5000 and turns <= 4:
|
|
candidates = [m for m in ["gemma-4-12b","qwen3.6-27B-code","qwen3.6-35B-A3B"] if m in avail]
|
|
result = select_best_gpu(candidates, "light", agent)
|
|
if result: return result
|
|
|
|
# TIER 3: Medium - quality matters -> MoE primary (60%), Dense spillover (40%)
|
|
if t <= 30000:
|
|
candidates = moe_spillover(avail, ["qwen3.6-35B-A3B","qwen3.6-27B-code","gemma-4-12b"])
|
|
result = select_best_gpu(candidates, "medium", agent)
|
|
if result: return result
|
|
|
|
# TIER 4: Heavy - quality first -> Dense primary, MoE fallback
|
|
if t > 30000:
|
|
candidates = [m for m in ["qwen3.6-27B-code","qwen3.6-35B-A3B","gemma-4-12b"] if m in avail]
|
|
result = select_best_gpu(candidates, "heavy", agent)
|
|
if result: return result
|
|
|
|
# TIER 5: Default - MoE primary (60%), Dense spillover (40%)
|
|
candidates = moe_spillover(avail, ["qwen3.6-35B-A3B","qwen3.6-27B-code","gemma-4-12b"])
|
|
result = select_best_gpu(candidates, "default", agent)
|
|
if result: return result
|
|
return {"model":avail[0],"reason":"last_resort"}
|
|
|
|
|
|
def moe_spillover(avail, default_order):
|
|
"""Spill 40% of MoE-first traffic to Dense to prevent Strix Halo overheating.
|
|
Only applies when MoE is first candidate, available, and not busy."""
|
|
import random
|
|
if (default_order[0] == "qwen3.6-35B-A3B"
|
|
and "qwen3.6-35B-A3B" in avail
|
|
and not is_gpu_busy("qwen3.6-35B-A3B")
|
|
and "qwen3.6-27B-code" in avail
|
|
and not is_gpu_busy("qwen3.6-27B-code")
|
|
and random.random() < 0.4):
|
|
# Swap: Dense first, MoE second
|
|
return ["qwen3.6-27B-code","qwen3.6-35B-A3B"] + [m for m in default_order[2:] if m in avail and m not in ("qwen3.6-27B-code","qwen3.6-35B-A3B")]
|
|
return [m for m in default_order if m in avail]
|
|
def clean_unicode(text):
|
|
if not isinstance(text, str): return text
|
|
text = text.replace(chr(0x2014), "-"); text = text.replace(chr(0x2013), "-")
|
|
text = text.replace(chr(0x2018), "'"); text = text.replace(chr(0x2019), "'")
|
|
text = text.replace(chr(0x201C), '"'); text = text.replace(chr(0x201D), '"')
|
|
text = text.replace(chr(0x2026), "..."); text = text.replace(chr(0x00A0), " ")
|
|
return text.encode("ascii", "ignore").decode("ascii")
|
|
|
|
def clean_response(d):
|
|
if isinstance(d, dict): return {k: clean_response(v) for k,v in d.items()}
|
|
if isinstance(d, list): return [clean_response(v) for v in d]
|
|
if isinstance(d, str): return clean_unicode(d)
|
|
return d
|
|
|
|
def get_metrics():
|
|
d = {"gpus":[],"route_counts":{},"agent_counts":{},"tier_counts":{},"recent":[],"timestamp":time.time(),"active_requests":{}}
|
|
for m in GPU_URLS:
|
|
h = check_gpu_health(m)
|
|
d["gpus"].append({"id":m,"gpu_name":h.get("gpu_name",m),"status":h.get("status"),"vram_used_mb":h.get("vram_used_mb"),"vram_total_mb":h.get("vram_total_mb"),"vram_pct":h.get("vram_pct"),"temp_c":h.get("temp_c"),"gpu_util_pct":h.get("gpu_util_pct"),"power_w":h.get("power_w"),"power_limit_w":h.get("power_limit_w"),"active_requests":gpu_active_count(m), "max_concurrent": GPU_MAX_CONCURRENT.get(m, 1)})
|
|
d["active_requests"][m] = gpu_active_count(m)
|
|
if r:
|
|
try:
|
|
for m in GPU_URLS: d["route_counts"][m] = int(r.get("routes:"+m) or 0)
|
|
for k,v in API_KEYS.items():
|
|
c = int(r.get("routes:agent:"+v["agent"]) or 0)
|
|
if c>0: d["agent_counts"][v["agent"]] = c
|
|
for t in TIER_MODELS: d["tier_counts"][t] = int(r.get("routes:tier:"+t) or 0)
|
|
raw = r.lrange("routes:recent",0,49)
|
|
d["recent"] = [json.loads(x) for x in raw] if raw else []
|
|
except Exception: pass
|
|
return d
|
|
|
|
def bcast():
|
|
data = get_metrics(); payload = json.dumps(data)
|
|
with sse_lock:
|
|
dead = []
|
|
for q in sse_subscribers:
|
|
try: q.put(payload)
|
|
except Exception: dead.append(q)
|
|
for q in dead: sse_subscribers.remove(q)
|
|
|
|
QUEUE_TIMEOUT = int(os.environ.get("QUEUE_TIMEOUT", "30")) # max seconds to queue before 503
|
|
|
|
@app.route("/v1/chat/completions", methods=["POST"])
|
|
def chat():
|
|
try:
|
|
rd = request.get_json(force=True)
|
|
ak = request.headers.get("Authorization","").replace("Bearer ","")
|
|
if not ak or ak not in API_KEYS:
|
|
log.warning("AUTH_REJECTED: no/invalid API key from %s", request.remote_addr)
|
|
return jsonify({"error": "Unauthorized — valid API key required"}), 401
|
|
ki = API_KEYS[ak]
|
|
tier, agent = ki["tier"], ki["agent"]
|
|
|
|
# Phase 0: dual-key transition — log deprecated key usage
|
|
if ki.get("deprecated"):
|
|
new_key = next((k for k, v in API_KEYS.items()
|
|
if v.get("agent") == agent and not v.get("deprecated")), None)
|
|
log.warning("DEPRECATED_KEY: agent=%s using old key %s...%s — switch to %s...%s",
|
|
agent, ak[:12], ak[-8:],
|
|
new_key[:12] if new_key else "N/A",
|
|
new_key[-8:] if new_key else "N/A")
|
|
if r:
|
|
r.incr("deprecated_usage:" + agent)
|
|
|
|
# Rate limit check
|
|
allowed, rl_val, reset_sec = check_rate_limit(ak, tier)
|
|
if not allowed:
|
|
resp = jsonify({"error": "Rate limit exceeded", "retry_after_s": rl_val})
|
|
resp.headers["Retry-After"] = str(rl_val)
|
|
resp.headers["X-RateLimit-Limit"] = str(RATE_LIMIT_RPM.get(tier, 30))
|
|
resp.headers["X-RateLimit-Remaining"] = "0"
|
|
resp.headers["X-RateLimit-Reset"] = str(int(time.time() + rl_val))
|
|
log.warning("RATE_LIMIT: %s (%s) exceeded limit", agent, ak[-8:])
|
|
return resp, 429
|
|
|
|
# Allow agent to override queue timeout via header
|
|
q_timeout = int(request.headers.get("X-Queue-Timeout", str(QUEUE_TIMEOUT)))
|
|
|
|
# Cross-turn context tracking: accumulate tokens per session (Phase 2: atomic Lua)
|
|
session_id = request.headers.get("X-Session-Id", "")
|
|
session_tokens = 0
|
|
if session_id and r:
|
|
try:
|
|
current = estimate_tokens(rd.get("messages",[]))
|
|
# Atomic GET/MAX/SET via Lua script prevents race conditions
|
|
session_tokens = r.eval(SESSION_LUA_SCRIPT, 1, "session:" + session_id, current)
|
|
except Exception: pass
|
|
|
|
d = route(rd, tier, agent)
|
|
queue_start = time.time()
|
|
|
|
# Queue loop: wait for a GPU slot instead of immediate 503
|
|
while d.get("saturated"):
|
|
elapsed = time.time() - queue_start
|
|
if elapsed > q_timeout:
|
|
resp = jsonify({"error": "All GPUs saturated", "queued_s": round(elapsed,1), "retry_after_s": 5})
|
|
resp.headers["Retry-After"] = "5"
|
|
log.warning("QUEUE_TIMEOUT: %s waited %.1fs, all GPUs saturated", agent, elapsed)
|
|
return resp, 503
|
|
time.sleep(0.5) # poll every 500ms
|
|
d = route(rd, tier, agent)
|
|
|
|
queue_ms = (time.time() - queue_start) * 1000
|
|
if queue_ms > 500:
|
|
log.info("QUEUED: %s waited %.0fms before slot opened", agent, queue_ms)
|
|
model, reason, url = d["model"], d["reason"], GPU_URLS[d["model"]]
|
|
|
|
# Stash rate limit values for response headers
|
|
_rl_remaining = rl_val
|
|
_rl_limit = RATE_LIMIT_RPM.get(tier, 30)
|
|
_rl_reset = reset_sec
|
|
is_stream = rd.get("stream", False)
|
|
|
|
gpu_incr(model)
|
|
|
|
log.info("ROUTE: %s -> %s (%s) stream=%s active=%d/%d", agent, model, reason, is_stream, gpu_active_count(model), GPU_MAX_CONCURRENT.get(model,1))
|
|
# Track which GPU this agent is using (TTL 120s covers typical request)
|
|
if r and agent:
|
|
try: r.setex("agent_gpu:" + agent + ":" + model, 120, "1")
|
|
except: pass
|
|
if r:
|
|
try:
|
|
r.incr("routes:"+model); r.incr("routes:tier:"+tier); r.incr("routes:agent:"+agent)
|
|
r.incr("ts:"+model+":"+time.strftime("%Y%m%d%H"))
|
|
r.lpush("routes:recent", json.dumps({"ts":time.time(),"model":model,"reason":reason,"tier":tier,"agent":agent,"queue_ms": round(queue_ms,1)}))
|
|
r.ltrim("routes:recent",0,999)
|
|
except Exception: pass
|
|
start = time.time()
|
|
resp = requests.post(url+"/chat/completions", json=rd,
|
|
headers={"Content-Type":"application/json","Authorization":"Bearer not-needed"}, timeout=300, stream=is_stream)
|
|
lat = int((time.time()-start)*1000)
|
|
gpu_decr(model)
|
|
|
|
if resp.status_code != 200:
|
|
if resp.status_code in (502, 504):
|
|
trip_circuit(model, 30)
|
|
return jsonify({"error":"GPU error "+str(resp.status_code)}), 502
|
|
if is_stream:
|
|
# Buffer SSE chunks, handle split lines for large responses
|
|
chunks = []
|
|
stream_timings = {}
|
|
buf = "" # accumulate partial lines
|
|
for raw in resp.iter_content(chunk_size=None, decode_unicode=True):
|
|
if raw:
|
|
cleaned = clean_unicode(raw)
|
|
chunks.append(cleaned)
|
|
buf += cleaned
|
|
# Process complete lines from buffer
|
|
while "\n" in buf:
|
|
line, buf = buf.split("\n", 1)
|
|
line = line.strip()
|
|
if line.startswith("data: ") and not stream_timings:
|
|
js = line[6:].strip()
|
|
if js.startswith("{") and "timings" in js and "predicted_n" in js:
|
|
try:
|
|
tj = json.loads(js).get("timings", {})
|
|
if tj:
|
|
stream_timings = tj
|
|
except: pass
|
|
# Store perf record with real token counts from stream
|
|
if stream_timings:
|
|
pt = stream_timings.get("prompt_n", 0)
|
|
ct = stream_timings.get("predicted_n", 0)
|
|
tps = stream_timings.get("predicted_per_second", 0)
|
|
gen_ms = stream_timings.get("predicted_ms", lat)
|
|
store_perf_record(model, agent, tier, reason, queue_ms, gen_ms, pt, ct, True)
|
|
else:
|
|
store_perf_record(model, agent, tier, reason, queue_ms, lat, estimate_tokens(rd.get("messages",[])), 0, True)
|
|
# Yield all chunks to client
|
|
def gen():
|
|
for c in chunks: yield c
|
|
bcast()
|
|
ctx_remaining = GPU_CONTEXT.get(model, 65536) - max(session_tokens, estimate_tokens(rd.get("messages",[])))
|
|
ctx_pct = ctx_remaining / GPU_CONTEXT.get(model, 65536) * 100
|
|
ctx_warning = "compact_urgent" if ctx_pct < 5 else ("compact_recommended" if ctx_pct < 15 else ("compact_soon" if ctx_pct < 30 else "ok"))
|
|
sse_resp = Response(stream_with_context(gen()), mimetype="text/event-stream")
|
|
sse_resp.headers["X-RateLimit-Limit"] = str(_rl_limit)
|
|
sse_resp.headers["X-RateLimit-Remaining"] = str(max(0, _rl_remaining))
|
|
sse_resp.headers["X-RateLimit-Reset"] = str(int(time.time() + _rl_reset))
|
|
sse_resp.headers["X-Context-Remaining"] = str(max(0, ctx_remaining))
|
|
sse_resp.headers["X-Context-Warning"] = ctx_warning
|
|
sse_resp.headers["X-Context-Model"] = model
|
|
return sse_resp
|
|
data = clean_response(resp.json())
|
|
for c in data.get("choices",[]):
|
|
msg = c.get("message",{})
|
|
if not msg.get("content") and msg.get("reasoning_content"):
|
|
msg["content"] = msg["reasoning_content"]
|
|
# Extract performance data from llama.cpp response
|
|
usage = data.get("usage", {})
|
|
timings = data.get("timings", {})
|
|
prompt_tokens = usage.get("prompt_tokens", 0)
|
|
completion_tokens = usage.get("completion_tokens", 0)
|
|
inference_ms = lat # total GPU round-trip
|
|
store_perf_record(model, agent, tier, reason, queue_ms, inference_ms, prompt_tokens, completion_tokens, False)
|
|
ctx_remaining = GPU_CONTEXT.get(model, 65536) - max(session_tokens, estimate_tokens(rd.get("messages",[])))
|
|
ctx_pct = ctx_remaining / GPU_CONTEXT.get(model, 65536) * 100
|
|
ctx_warning = "compact_urgent" if ctx_pct < 5 else ("compact_recommended" if ctx_pct < 15 else ("compact_soon" if ctx_pct < 30 else "ok"))
|
|
data["routing"] = {"model":model,"reason":reason,"gpu":url,"tier":tier,"agent":agent,"latency_ms":lat,"queue_ms": round(queue_ms,1),"active_gpu":gpu_active_count(model),"context_remaining": max(0, ctx_remaining),"context_pct": round(ctx_pct,1),"context_warning": ctx_warning}
|
|
resp = jsonify(data)
|
|
resp.headers["X-RateLimit-Limit"] = str(_rl_limit)
|
|
resp.headers["X-RateLimit-Remaining"] = str(max(0, _rl_remaining))
|
|
resp.headers["X-RateLimit-Reset"] = str(int(time.time() + _rl_reset))
|
|
resp.headers["X-Context-Remaining"] = str(max(0, ctx_remaining))
|
|
resp.headers["X-Context-Warning"] = ctx_warning
|
|
resp.headers["X-Context-Model"] = model
|
|
bcast()
|
|
return resp
|
|
except requests.Timeout:
|
|
gpu_decr(model)
|
|
trip_circuit(model, 30)
|
|
log.error("TIMEOUT: %s -> %s (Circuit tripped)", agent, model)
|
|
return jsonify({"error":"timeout"}), 504
|
|
except Exception as e:
|
|
gpu_decr(model)
|
|
log.error("Error: %s\n%s", e, traceback.format_exc())
|
|
return jsonify({"error":str(e)}), 500
|
|
|
|
@app.route("/metrics/performance")
|
|
def performance():
|
|
"""Per-request performance analytics with percentiles per model/reason/agent."""
|
|
if not get_redis(): return jsonify({"error": "Redis unavailable"}), 503
|
|
try:
|
|
window_hours = int(request.args.get("window", "24").replace("h",""))
|
|
model_filter = request.args.get("model", "all")
|
|
|
|
# Load recent records
|
|
cutoff = time.time() - (window_hours * 3600)
|
|
raw = r.lrange("perf:recent", 0, -1)
|
|
records = []
|
|
for x in raw:
|
|
try:
|
|
rec = json.loads(x)
|
|
if rec["ts"] >= cutoff:
|
|
records.append(rec)
|
|
except: pass
|
|
|
|
# Filter by model if specified
|
|
if model_filter != "all":
|
|
records = [r for r in records if r["model"] == model_filter]
|
|
|
|
if not records:
|
|
return jsonify({"models": [], "reasons": [], "agents": [], "summary": {"total_requests": 0}})
|
|
|
|
def pct(values, p):
|
|
if len(values) < 2: return round(values[0], 1) if values else 0
|
|
return round(statistics.quantiles(sorted(values), n=100, method='inclusive')[min(p-1, 98)], 1)
|
|
|
|
# Per-model stats
|
|
model_groups = {}
|
|
for rec in records:
|
|
m = rec["model"]
|
|
if m not in model_groups: model_groups[m] = []
|
|
model_groups[m].append(rec)
|
|
|
|
models = []
|
|
for m, recs in sorted(model_groups.items()):
|
|
latencies = [r["total_ms"] for r in recs]
|
|
tps_vals = [r["tokens_per_sec"] for r in recs if r["tokens_per_sec"] > 0]
|
|
non_stream = [r for r in recs if not r["stream"]]
|
|
queue_times = [r["queue_ms"] for r in non_stream]
|
|
models.append({
|
|
"model": m,
|
|
"count": len(recs),
|
|
"stream_pct": round(len([r for r in recs if r["stream"]]) / len(recs) * 100, 1),
|
|
"latency": {
|
|
"avg": round(statistics.mean(latencies), 1),
|
|
"p50": pct(latencies, 50),
|
|
"p95": pct(latencies, 95),
|
|
"p99": pct(latencies, 99)
|
|
},
|
|
"throughput": {
|
|
"avg_tokens_per_sec": round(statistics.mean(tps_vals), 1) if tps_vals else 0,
|
|
"p50": pct(tps_vals, 50) if tps_vals else 0,
|
|
"p95": pct(tps_vals, 95) if tps_vals else 0,
|
|
},
|
|
"queue": {
|
|
"avg_ms": round(statistics.mean(queue_times), 1) if queue_times else 0,
|
|
"p95_ms": pct(queue_times, 95) if queue_times else 0,
|
|
} if queue_times else None
|
|
})
|
|
|
|
# Per-reason stats
|
|
reason_groups = {}
|
|
for rec in records:
|
|
rsn = rec["reason"]
|
|
if rsn not in reason_groups: reason_groups[rsn] = []
|
|
reason_groups[rsn].append(rec)
|
|
|
|
reasons = []
|
|
for rsn, recs in sorted(reason_groups.items(), key=lambda x: -len(x[1])):
|
|
latencies = [r["total_ms"] for r in recs]
|
|
reasons.append({
|
|
"reason": rsn,
|
|
"count": len(recs),
|
|
"avg_total_ms": round(statistics.mean(latencies), 1),
|
|
"p95_total_ms": pct(latencies, 95)
|
|
})
|
|
|
|
# Per-agent stats
|
|
agent_groups = {}
|
|
for rec in records:
|
|
ag = rec["agent"]
|
|
if ag not in agent_groups: agent_groups[ag] = []
|
|
agent_groups[ag].append(rec)
|
|
|
|
agents = []
|
|
for ag, recs in sorted(agent_groups.items(), key=lambda x: -len(x[1])):
|
|
latencies = [r["total_ms"] for r in recs]
|
|
tps_vals = [r["tokens_per_sec"] for r in recs if r["tokens_per_sec"] > 0]
|
|
agents.append({
|
|
"agent": ag,
|
|
"count": len(recs),
|
|
"avg_total_ms": round(statistics.mean(latencies), 1),
|
|
"avg_tokens_per_sec": round(statistics.mean(tps_vals), 1) if tps_vals else 0
|
|
})
|
|
|
|
all_lat = [r["total_ms"] for r in records]
|
|
all_tps = [r["tokens_per_sec"] for r in records if r["tokens_per_sec"] > 0]
|
|
summary = {
|
|
"total_requests": len(records),
|
|
"window_hours": window_hours,
|
|
"latency": {
|
|
"avg_ms": round(statistics.mean(all_lat), 1),
|
|
"p50_ms": pct(all_lat, 50),
|
|
"p95_ms": pct(all_lat, 95),
|
|
"p99_ms": pct(all_lat, 99)
|
|
},
|
|
"throughput_avg_tps": round(statistics.mean(all_tps), 1) if all_tps else 0
|
|
}
|
|
|
|
return jsonify({"models": models, "reasons": reasons, "agents": agents, "summary": summary})
|
|
except Exception as e:
|
|
return jsonify({"error": str(e)}), 500
|
|
|
|
@app.route("/metrics/scatter")
|
|
def scatter():
|
|
"""Return individual data points for scatter plots (prompt_tokens vs latency)."""
|
|
if not get_redis(): return jsonify({"error": "Redis unavailable"}), 503
|
|
try:
|
|
window_hours = int(request.args.get("window", "24").replace("h",""))
|
|
model_filter = request.args.get("model", "all")
|
|
cutoff = time.time() - (window_hours * 3600)
|
|
raw = r.lrange("perf:recent", 0, -1)
|
|
points = []
|
|
for x in raw:
|
|
try:
|
|
rec = json.loads(x)
|
|
if rec["ts"] >= cutoff:
|
|
if model_filter == "all" or rec["model"] == model_filter:
|
|
points.append({
|
|
"model": rec["model"],
|
|
"agent": rec["agent"],
|
|
"reason": rec["reason"],
|
|
"prompt_tokens": int(rec.get("prompt_tokens", 0)),
|
|
"completion_tokens": rec.get("completion_tokens", 0),
|
|
"inference_ms": round(rec["inference_ms"], 1),
|
|
"tokens_per_sec": rec.get("tokens_per_sec", 0),
|
|
"stream": rec.get("stream", False)
|
|
})
|
|
except: pass
|
|
return jsonify({"points": points, "count": len(points)})
|
|
except Exception as e:
|
|
return jsonify({"error": str(e)}), 500
|
|
|
|
@app.route("/v1/models")
|
|
def models():
|
|
def _h(m): return check_gpu_health(m, sidecar_timeout=1.5, gpu_timeout=1)
|
|
return jsonify({"object":"list","data":[{"id":m,"object":"model","owned_by":"syslog","status":_h(m).get("status"),"gpu":_h(m).get("gpu_name")} for m in GPU_URLS]})
|
|
|
|
@app.route("/health")
|
|
def health():
|
|
gpus = {}
|
|
for m in GPU_URLS:
|
|
h = check_gpu_health(m, sidecar_timeout=1.5, gpu_timeout=1)
|
|
h["active_requests"] = gpu_active_count(m)
|
|
h["max_concurrent"] = GPU_MAX_CONCURRENT.get(m, 1)
|
|
gpus[m] = h
|
|
return jsonify({"status":"healthy","redis":"connected" if r else "down","gpus":gpus,"available_models":available_models()})
|
|
|
|
@app.route("/metrics")
|
|
def metrics(): return jsonify(get_metrics())
|
|
|
|
@app.route("/metrics/timeseries")
|
|
def metrics_timeseries():
|
|
period = request.args.get("period", "day"); models_list = list(GPU_URLS.keys())
|
|
data = {"models": {}, "labels": []}
|
|
if period == "day":
|
|
buckets = [time.strftime("%Y%m%d%H", time.gmtime(time.time()-h*3600)) for h in range(23,-1,-1)]
|
|
data["labels"] = [time.strftime("%H:00", time.gmtime(time.time()-h*3600)) for h in range(23,-1,-1)]
|
|
elif period == "week":
|
|
buckets = [time.strftime("%Y%m%d", time.gmtime(time.time()-d*86400)) for d in range(6,-1,-1)]
|
|
data["labels"] = [time.strftime("%a", time.gmtime(time.time()-d*86400)) for d in range(6,-1,-1)]
|
|
else:
|
|
buckets = [time.strftime("%Y%m%d", time.gmtime(time.time()-d*86400)) for d in range(29,-1,-1)]
|
|
data["labels"] = [time.strftime("%m/%d", time.gmtime(time.time()-d*86400)) for d in range(29,-1,-1)]
|
|
if r:
|
|
for model in models_list:
|
|
counts = []
|
|
for bucket in buckets:
|
|
total = 0
|
|
if period in ("week","month"):
|
|
for hh in range(24): total += int(r.get("ts:"+model+":"+bucket+"{:02d}".format(hh)) or 0)
|
|
else: total = int(r.get("ts:"+model+":"+bucket) or 0)
|
|
counts.append(total)
|
|
data["models"][model] = counts
|
|
return jsonify(data)
|
|
|
|
|
|
|
|
@app.route("/metrics/circuit-breaker")
|
|
def metrics_circuit_breaker():
|
|
"""Expose circuit breaker status per model. Phase 1."""
|
|
result = {}
|
|
if r:
|
|
for model in GPU_URLS:
|
|
key = "circuit:" + model + ":open"
|
|
duration = r.ttl(key)
|
|
trip_count = int(r.get("circuit:" + model + ":count") or 0)
|
|
result[model] = {
|
|
"tripped": r.exists(key),
|
|
"remaining_ttl": duration,
|
|
"trip_count": trip_count
|
|
}
|
|
return jsonify(result)
|
|
|
|
@app.route("/stream")
|
|
def stream():
|
|
def ev():
|
|
q = queue.Queue()
|
|
with sse_lock: sse_subscribers.append(q)
|
|
try:
|
|
yield "data: "+json.dumps(get_metrics())+"\n\n"
|
|
while True:
|
|
try: yield "data: "+q.get(timeout=3)+"\n\n"
|
|
except queue.Empty: yield "data: "+json.dumps(get_metrics())+"\n\n"
|
|
except GeneratorExit: pass
|
|
finally:
|
|
with sse_lock:
|
|
if q in sse_subscribers: sse_subscribers.remove(q)
|
|
return Response(stream_with_context(ev()), mimetype="text/event-stream",
|
|
headers={"Cache-Control":"no-cache","X-Accel-Buffering":"no","Access-Control-Allow-Origin":"*"})
|
|
|
|
# ── Phase 0: Admin Key Management ──
|
|
ADMIN_KEY = os.environ.get("ADMIN_KEY", "")
|
|
|
|
def _admin_auth():
|
|
"""Require admin key for management endpoints."""
|
|
if not ADMIN_KEY:
|
|
return False, "ADMIN_KEY not configured on server"
|
|
ak = request.headers.get("Authorization","").replace("Bearer ","")
|
|
if ak != ADMIN_KEY:
|
|
return False, "Admin key required"
|
|
return True, None
|
|
|
|
@app.route("/admin/keys")
|
|
def admin_keys():
|
|
"""List all API keys (masked) with agent, tier, and deprecation status."""
|
|
ok, err = _admin_auth()
|
|
if not ok: return jsonify({"error": err}), 401
|
|
keys = []
|
|
for key, info in API_KEYS.items():
|
|
masked = key[:8] + "..." + key[-8:]
|
|
keys.append({
|
|
"masked": masked,
|
|
"prefix": key[:8],
|
|
"agent": info["agent"],
|
|
"tier": info["tier"],
|
|
"deprecated": info.get("deprecated", False),
|
|
"length": len(key)
|
|
})
|
|
return jsonify({
|
|
"total": len(keys),
|
|
"active": sum(1 for k in keys if not k["deprecated"]),
|
|
"deprecated": sum(1 for k in keys if k["deprecated"]),
|
|
"keys": sorted(keys, key=lambda k: (k["deprecated"], k["agent"]))
|
|
})
|
|
|
|
@app.route("/admin/keys/deprecation-summary")
|
|
def admin_deprecation_summary():
|
|
"""Summary of deprecated key usage (from Redis logs, if available)."""
|
|
ok, err = _admin_auth()
|
|
if not ok: return jsonify({"error": err}), 401
|
|
deprecated_agents = []
|
|
for key, info in API_KEYS.items():
|
|
if info.get("deprecated"):
|
|
# Check Redis for usage count
|
|
count = 0
|
|
if r:
|
|
count = int(r.get("deprecated_usage:" + info["agent"]) or 0)
|
|
deprecated_agents.append({
|
|
"agent": info["agent"],
|
|
"deprecated_uses": count,
|
|
"needs_migration": count > 0
|
|
})
|
|
return jsonify({
|
|
"deprecated_agents": sorted(deprecated_agents, key=lambda d: -d["deprecated_uses"]),
|
|
"recommendation": "Run POST /admin/keys/revoke to remove keys with 0 usage"
|
|
})
|
|
|
|
@app.route("/admin/keys/generate", methods=["POST"])
|
|
def admin_generate_key():
|
|
"""Generate a new API key for an agent. Body: {"agent": "Name", "tier": "enterprise"}"""
|
|
ok, err = _admin_auth()
|
|
if not ok: return jsonify({"error": err}), 401
|
|
body = request.get_json(force=True)
|
|
agent = body.get("agent", "").strip()
|
|
tier = body.get("tier", "enterprise")
|
|
if not agent:
|
|
return jsonify({"error": "agent field required"}), 400
|
|
if tier not in ("starter", "professional", "enterprise"):
|
|
return jsonify({"error": "tier must be starter/professional/enterprise"}), 400
|
|
# Generate secure key
|
|
import secrets, hashlib
|
|
prefix = hashlib.sha256(secrets.token_bytes(12)).hexdigest()[:8]
|
|
suffix = secrets.token_hex(20)
|
|
new_key = f"sk-{prefix}-{suffix}"
|
|
# Update in-memory dict (note: not persisted across restarts without env var update)
|
|
API_KEYS[new_key] = {"tier": tier, "agent": agent}
|
|
log.info("KEY_GENERATED: agent=%s tier=%s key=%s...%s", agent, tier, new_key[:8], new_key[-8:])
|
|
return jsonify({
|
|
"agent": agent,
|
|
"tier": tier,
|
|
"key": new_key,
|
|
"masked": new_key[:8] + "..." + new_key[-8:],
|
|
"warning": "This key exists in memory only. Update API_KEYS env var and redeploy to persist."
|
|
}), 201
|
|
|
|
@app.route("/admin/keys/revoke", methods=["POST"])
|
|
def admin_revoke_key():
|
|
"""Revoke a deprecated key. Body: {"agent": "Name"} or {"key_prefix": "sk-xxxx"}"""
|
|
ok, err = _admin_auth()
|
|
if not ok: return jsonify({"error": err}), 401
|
|
body = request.get_json(force=True)
|
|
agent = body.get("agent", "")
|
|
key_prefix = body.get("key_prefix", "")
|
|
revoked = []
|
|
keys_to_remove = []
|
|
for key, info in API_KEYS.items():
|
|
if not info.get("deprecated"):
|
|
continue
|
|
if agent and info["agent"] == agent:
|
|
keys_to_remove.append(key)
|
|
elif key_prefix and key.startswith(key_prefix):
|
|
keys_to_remove.append(key)
|
|
for key in keys_to_remove:
|
|
info = API_KEYS.pop(key)
|
|
revoked.append({"agent": info["agent"], "masked": key[:8] + "..." + key[-8:]})
|
|
log.warning("KEY_REVOKED: agent=%s key=%s...%s", info["agent"], key[:8], key[-8:])
|
|
return jsonify({
|
|
"revoked": len(revoked),
|
|
"keys": revoked,
|
|
"remaining_total": len(API_KEYS),
|
|
"warning": "Memory-only revoke. Update API_KEYS env var and redeploy to persist."
|
|
})
|
|
|
|
if __name__ == "__main__":
|
|
log.info("Router on :9000 (load-aware)")
|
|
app.run(host="0.0.0.0", port=9000, debug=False)
|