import os
import time
from typing import List, AsyncGenerator, Optional
import httpx
from fastapi import FastAPI, HTTPException
from fastapi.responses import StreamingResponse, ORJSONResponse, HTMLResponse
from fastapi.middleware.cors import CORSMiddleware
from pydantic import BaseModel, Field
# ===== FastAPI app =====
app = FastAPI(
title="Qwen3 Main Router",
description="Main router / load balancer for Qwen3 mini servers",
version="1.0.0",
default_response_class=ORJSONResponse,
)
app.add_middleware(
CORSMiddleware,
allow_origins=["*"],
allow_credentials=True,
allow_methods=["*"],
allow_headers=["*"],
)
# Comma-separated list of mini server base URLs
# Example:
# MINI_SERVERS="https://username-mini1.hf.space,https://username-mini2.hf.space"
MINI_SERVERS = [
"https://antaram-server1.hf.space",
"https://antaram-server2.hf.space",
]
http_client: Optional[httpx.AsyncClient] = None
# Usage stats per mini for /gui
MINI_USAGE = {} # { base_url: {"total_requests": int, "last_used": float or None} }
@app.on_event("startup")
async def startup():
global http_client, MINI_USAGE
http_client = httpx.AsyncClient(
timeout=httpx.Timeout(300.0, connect=10.0),
limits=httpx.Limits(max_keepalive_connections=50, max_connections=100),
http2=True,
)
MINI_USAGE = {base_url: {"total_requests": 0, "last_used": None} for base_url in MINI_SERVERS}
@app.on_event("shutdown")
async def shutdown():
global http_client
if http_client:
await http_client.aclose()
# ===== Shared models (same as mini server) =====
class Message(BaseModel):
role: str
content: str
class Config:
extra = "ignore"
class ChatRequest(BaseModel):
messages: List[Message]
temperature: float = Field(default=0.6, ge=0.0, le=2.0)
top_p: float = Field(default=0.95, ge=0.0, le=1.0)
max_tokens: int = Field(default=4096, ge=1, le=32768)
stream: bool = Field(default=True)
class Config:
extra = "ignore"
class SimpleChatRequest(BaseModel):
prompt: str
temperature: float = Field(default=0.6, ge=0.0, le=2.0)
top_p: float = Field(default=0.95, ge=0.0, le=1.0)
max_tokens: int = Field(default=4096, ge=1, le=32768)
stream: bool = Field(default=True)
class Config:
extra = "ignore"
# ===== Mini server coordination helpers =====
async def reserve_on_mini(base_url: str) -> bool:
"""Try to reserve a slot on the given mini.
Returns True if reserved, False if busy/unreachable.
"""
try:
resp = await http_client.post(f"{base_url}/reserve", timeout=5.0)
if resp.status_code == 200:
return True
return False # 429 or anything else
except Exception:
return False
async def release_on_mini(base_url: str) -> None:
"""Best-effort release; ignore errors."""
try:
await http_client.post(f"{base_url}/release", timeout=5.0)
except Exception:
pass
async def choose_mini() -> str:
"""Iterate minis and grab the first one that grants a /reserve.
Encodes the logic:
- if mini1 is working/processing (full), try mini2, etc.
"""
if not MINI_SERVERS:
raise HTTPException(status_code=503, detail="No mini servers configured")
for base_url in MINI_SERVERS:
if await reserve_on_mini(base_url):
usage = MINI_USAGE.setdefault(base_url, {"total_requests": 0, "last_used": None})
usage["total_requests"] += 1
usage["last_used"] = time.time()
return base_url
raise HTTPException(status_code=503, detail="All mini servers are busy")
# ===== Proxy helpers =====
async def proxy_sse_to_mini(path: str, payload: dict) -> AsyncGenerator[bytes, None]:
"""Streaming proxy:
frontend -> main -> mini (SSE) -> main -> frontend
"""
mini_url = await choose_mini()
full_url = f"{mini_url}{path}"
try:
async with http_client.stream(
"POST",
full_url,
json=payload,
headers={"Accept": "text/event-stream"},
) as resp:
if resp.status_code != 200:
body = await resp.aread()
raise HTTPException(
status_code=resp.status_code,
detail=f"Mini error: {body.decode(errors='ignore')}",
)
async for chunk in resp.aiter_raw():
# pass SSE bytes straight through
yield chunk
finally:
await release_on_mini(mini_url)
async def proxy_json_to_mini(path: str, payload: dict) -> ORJSONResponse:
mini_url = await choose_mini()
full_url = f"{mini_url}{path}"
try:
resp = await http_client.post(full_url, json=payload)
data = resp.json()
return ORJSONResponse(content=data, status_code=resp.status_code)
finally:
await release_on_mini(mini_url)
# ===== Public endpoints for frontend =====
@app.get("/")
async def root():
return {
"status": "ok",
"message": "Main Qwen3 router is running",
"mini_servers": MINI_SERVERS,
}
@app.get("/v1/models")
async def list_models():
# You can keep this static; it's the same model ID in all minis
return {
"object": "list",
"data": [{
"id": "qwen3-0.6b",
"object": "model",
"created": int(time.time()),
"owned_by": "cluster",
}],
}
@app.get("/health")
async def health():
"""Aggregated health + status across minis."""
results = []
for base_url in MINI_SERVERS:
mini_health = None
mini_status = None
try:
# Low-level LLM backend health from mini
resp_h = await http_client.get(f"{base_url}/health", timeout=5.0)
mini_health = resp_h.json()
except Exception as e:
mini_health = {"status": "unreachable", "error": str(e)}
try:
# Load status from mini
resp_s = await http_client.get(f"{base_url}/status", timeout=5.0)
mini_status = resp_s.json()
except Exception as e:
mini_status = {"status": "unknown", "error": str(e)}
usage = MINI_USAGE.get(base_url, {"total_requests": 0, "last_used": None})
results.append(
{
"mini": base_url,
"health": mini_health,
"status": mini_status,
"usage": usage,
}
)
return {"mini_servers": results}
@app.post("/v1/chat/completions")
async def chat_completions(request: ChatRequest):
payload = request.dict()
if request.stream:
return StreamingResponse(
proxy_sse_to_mini("/v1/chat/completions", payload),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
"Transfer-Encoding": "chunked",
},
)
return await proxy_json_to_mini("/v1/chat/completions", payload)
@app.post("/chat")
async def simple_chat(request: SimpleChatRequest):
payload = {
"prompt": request.prompt,
"temperature": request.temperature,
"top_p": request.top_p,
"max_tokens": request.max_tokens,
"stream": request.stream,
}
if request.stream:
return StreamingResponse(
proxy_sse_to_mini("/chat", payload),
media_type="text/event-stream",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
},
)
return await proxy_json_to_mini("/chat", payload)
@app.post("/chat/raw")
async def raw_chat(request: SimpleChatRequest):
payload = {
"prompt": request.prompt,
"temperature": request.temperature,
"top_p": request.top_p,
"max_tokens": request.max_tokens,
"stream": True,
}
return StreamingResponse(
proxy_sse_to_mini("/chat/raw", payload),
media_type="text/plain",
headers={
"Cache-Control": "no-cache",
"Connection": "keep-alive",
"X-Accel-Buffering": "no",
},
)
@app.post("/fast")
async def fast_chat(prompt: str = "", max_tokens: int = 512):
payload = {
"prompt": prompt,
"max_tokens": max_tokens,
"stream": False,
"temperature": 0.6,
"top_p": 0.95,
}
return await proxy_json_to_mini("/fast", payload)
# ===== Simple GUI at /gui =====
@app.get("/gui", response_class=HTMLResponse)
async def gui():
"""Simple HTML dashboard showing mini servers with lights and stats.
Frontend: GET mainserver.space/gui
"""
return """
Qwen3 Cluster GUI
Qwen3 Cluster
Main router view of all mini servers. Lights = idle, busy, or offline.
Loading...
"""
if __name__ == "__main__":
import uvicorn
uvicorn.run(
app,
host="0.0.0.0",
port=8000,
loop="uvloop",
http="httptools",
access_log=False,
workers=1,
)