diff --git a/router/router.py b/router/router.py index bfd8d7a..0b7449d 100644 --- a/router/router.py +++ b/router/router.py @@ -117,7 +117,9 @@ def check_gpu_health(model): def available_models(): return [m for m in GPU_URLS if check_gpu_health(m)["status"] in ("healthy","saturated")] -def estimate_tokens(msgs): return sum(len(str(m.get("content",""))) for m in msgs) // 4 +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 is_gpu_busy(model): """Check if GPU is at or near max concurrent capacity.""" @@ -150,6 +152,9 @@ def route(rd, tier): allowed = TIER_MODELS.get(tier, ["qwen3.5-9b-vlm"]) avail = [m for m in available_models() if m in allowed] if not avail: return {"model": allowed[0], "reason": "all_saturated", "saturated": True} + # Check if all available GPUs are at max capacity + if all(is_gpu_busy(m) for m in avail): + return {"model": avail[0], "reason": "all_saturated", "saturated": True} req = rd.get("model","auto") if req != "auto": @@ -190,8 +195,8 @@ def route(rd, tier): # TIER 3: Heavy reasoning — extremely large context or very long conversations if t > 50000 or turns > 25: - # Dense first (98K, purpose-built for reasoning), then MoE/VLM 131K - candidates = [m for m in ["qwen3.6-27B-code","qwen3.6-35B-A3B","qwen3.5-9b-vlm"] if m in avail] + # MoE first (131K context handles heavy sessions), then Dense (98K reasoning), then Light (131K fallback) + candidates = [m for m in ["qwen3.6-35B-A3B","qwen3.6-27B-code","qwen3.5-9b-vlm"] if m in avail] result = select_best_gpu(candidates, "heavy_reasoning") if result: return result @@ -265,6 +270,17 @@ def chat(): # 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 + session_id = request.headers.get("X-Session-Id", "") + session_tokens = 0 + if session_id and r: + try: + prev = int(r.get("session:" + session_id) or 0) + current = estimate_tokens(rd.get("messages",[])) + session_tokens = max(prev, current) # context only grows + r.set("session:" + session_id, session_tokens, ex=86400) # TTL 24h + except Exception: pass + d = route(rd, tier) queue_start = time.time() @@ -307,7 +323,7 @@ def chat(): for raw in resp.iter_content(chunk_size=None, decode_unicode=True): if raw: yield clean_unicode(raw) bcast() - ctx_remaining = GPU_CONTEXT.get(model, 65536) - estimate_tokens(rd.get("messages",[])) + ctx_remaining = GPU_CONTEXT.get(model, 65536) - max(session_tokens, estimate_tokens(rd.get("messages",[]))) sse_resp = Response(stream_with_context(gen()), mimetype="text/event-stream") sse_resp.headers["X-Context-Remaining"] = str(max(0, ctx_remaining)) sse_resp.headers["X-Context-Model"] = model @@ -317,7 +333,7 @@ def chat(): msg = c.get("message",{}) if not msg.get("content") and msg.get("reasoning_content"): msg["content"] = msg["reasoning_content"] - ctx_remaining = GPU_CONTEXT.get(model, 65536) - estimate_tokens(rd.get("messages",[])) + ctx_remaining = GPU_CONTEXT.get(model, 65536) - max(session_tokens, estimate_tokens(rd.get("messages",[]))) data["routing"] = {"model":model,"reason":reason,"gpu":url,"tier":tier,"agent":agent,"latency_ms":lat,"active_gpu":gpu_active_count(model),"context_remaining": max(0, ctx_remaining)} resp = jsonify(data) resp.headers["X-Context-Remaining"] = str(max(0, ctx_remaining))