Random load balance

app.py

@@ -26,38 +26,29 @@ def log(msg):
 @app.get("/")
 def root():
     healthy = [k for k, v in key_status.items() if v["healthy"]]
-    return {
-        "status": "ok",
-        "keys_loaded": len(OLLAMA_KEYS),
-        "healthy_keys": len(healthy),
-        "keys_status": {v["prefix"]: {"success": v["success"], "healthy": v["healthy"]} for v in key_status.values()}
-    }
+    return {"status": "ok", "keys_loaded": len(OLLAMA_KEYS), "healthy": len(healthy),
+            "keys_status": {v["prefix"]: {"s": v["success"], "h": v["healthy"]} for v in key_status.values()}}
 
 def get_healthy_keys():
-    """Get all healthy keys"""
     return [k for k, v in key_status.items() if v["healthy"]]
 
 @app.post("/v1/chat/completions")
 async def chat(req: Request):
     auth_key = req.headers.get("Authorization", "").replace("Bearer ", "")
     if auth_key != MASTER_API_KEY:
-        log(f"AUTH FAIL")
         return JSONResponse({"error": "Unauthorized"}, status_code=401)
     
     body = await req.json()
     model = body.get("model", "?")
     
-    # Get all healthy keys
     healthy = get_healthy_keys()
     if not healthy:
-        # Reset all
-        for v in key_status.values():
-            v["healthy"] = True
+        for v in key_status.values(): v["healthy"] = True
         healthy = OLLAMA_KEYS[:2]
     
-    # RANDOMLY pick up to 2 keys
+    # RANDOMLY pick 2 keys
     selected = random.sample(healthy, min(2, len(healthy)))
-    log(f"REQ: model='{model}' | Using: {[key_status[k]['prefix'] for k in selected]}")
+    log(f"REQ: {model} -> {[key_status[k]['prefix'] for k in selected]}")
     
     for key in selected:
         ki = key_status[key]
@@ -65,32 +56,25 @@ async def chat(req: Request):
             start = time.time()
             async with httpx.AsyncClient(timeout=30.0) as client:
                 resp = await client.post(f"{BASE_URL}/v1/chat/completions", json=body, headers={"Authorization": f"Bearer {key}"})
-            elapsed = time.time() - start
             
             if resp.status_code == 200:
                 ki["success"] += 1
-                log(f"OK: key#{ki['index']} ({ki['prefix']}) in {elapsed:.2f}s")
+                log(f"OK: key#{ki['index']} in {time.time()-start:.2f}s")
                 return Response(resp.content, status_code=200)
-            
             elif resp.status_code == 429:
                 ki["healthy"] = False
-                log(f"RATE LIMIT: key#{ki['index']} - marking unhealthy")
                 continue
-                
         except Exception as e:
             ki["healthy"] = False
-            log(f"FAIL: key#{ki['index']} - {str(e)[:30]}")
             continue
     
-    log(f"ALL FAILED")
-    return JSONResponse({"error": "all keys failed"}, status_code=500)
+    return JSONResponse({"error": "all failed"}, status_code=500)
 
 @app.get("/v1/models")
 def models(req: Request):
     auth_key = req.headers.get("Authorization", "").replace("Bearer ", "")
     if auth_key != MASTER_API_KEY:
         return JSONResponse({"error": "Unauthorized"}, status_code=401)
-    
     for key in get_healthy_keys()[:2]:
         try:
             resp = httpx.get(f"{BASE_URL}/v1/models", headers={"Authorization": f"Bearer {key}"}, timeout=10)