243 lines
11 KiB
Python
243 lines
11 KiB
Python
"""FastAPI Backend Server with SSE Real-Time Telemetry and Model Orchestration API."""
|
|
import asyncio
|
|
import json
|
|
import logging
|
|
from typing import Dict, Any, Optional, List
|
|
from fastapi import FastAPI, Request, HTTPException, Query
|
|
from fastapi.responses import HTMLResponse, StreamingResponse, JSONResponse
|
|
from fastapi.staticfiles import StaticFiles
|
|
from fastapi.middleware.cors import CORSMiddleware
|
|
from pydantic import BaseModel, Field
|
|
|
|
import ram_optimizer
|
|
import vram_arbitrator
|
|
import overclock_manager
|
|
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
|
|
logger = logging.getLogger("model_manager_server")
|
|
|
|
app = FastAPI(
|
|
title="HyperSwap // GPU Program Swapper & Telemetry API",
|
|
version="1.0.0",
|
|
description="High-performance VRAM arbitration and 64GB RAM cache orchestrator for simultaneous Ollama and ComfyUI workloads on Linux.",
|
|
docs_url="/docs",
|
|
redoc_url="/redoc",
|
|
)
|
|
|
|
app.add_middleware(
|
|
CORSMiddleware,
|
|
allow_origins=["*"],
|
|
allow_credentials=True,
|
|
allow_methods=["*"],
|
|
allow_headers=["*"],
|
|
)
|
|
|
|
@app.on_event("startup")
|
|
async def on_startup():
|
|
await vram_arbitrator.arbitrator.start()
|
|
|
|
@app.on_event("shutdown")
|
|
async def on_shutdown():
|
|
await vram_arbitrator.arbitrator.stop()
|
|
|
|
# Pydantic Request Models
|
|
class SwitchRequest(BaseModel):
|
|
model: str = Field(..., description="Name of the Ollama model to hot-swap to in VRAM", example="qwen3.8fast:latest")
|
|
keep_alive: Optional[str] = Field("30m", description="Keep-alive duration in VRAM (e.g. 5m, 30m, 0)", example="30m")
|
|
|
|
class WarmRequest(BaseModel):
|
|
model_name: Optional[str] = Field(None, description="Ollama model name to warm into OS page cache", example="gemma4:26b")
|
|
filepath: Optional[str] = Field(None, description="Absolute file path of Safetensors/GGUF to warm into RAM", example="/home/drjones/ComfyUI/models/checkpoints/v1-5-pruned-emaonly-fp16.safetensors")
|
|
|
|
class BenchmarkRequest(BaseModel):
|
|
iterations: Optional[int] = Field(2, description="Number of back-and-forth switch iterations to measure", example=2)
|
|
models: Optional[List[str]] = Field(None, description="Optional pair of models to benchmark between", example=["qwen3.8fast:latest", "smtek/Qwen3.8-27B:Q2_K_XL"])
|
|
|
|
class OverclockApplyRequest(BaseModel):
|
|
profile: str = Field(..., description="Profile name: ollama | comfy | balanced", example="ollama")
|
|
|
|
class OverclockProfileUpdate(BaseModel):
|
|
config: Dict[str, Any] = Field(..., description="Profile settings dict", example={"power_limit_w": 370, "core_offset_mhz": 100})
|
|
|
|
class FanRequest(BaseModel):
|
|
mode: str = Field("auto", description="'auto' or 'manual'", example="manual")
|
|
percent: Optional[int] = Field(None, description="Fan speed 30-100 when mode=manual", example=70)
|
|
speed_pct: Optional[int] = Field(None, description="Alias for percent (30-100)", example=70)
|
|
|
|
|
|
# ==========================================
|
|
# REST API ENDPOINTS
|
|
# ==========================================
|
|
|
|
@app.get("/api/stats", summary="Full System Snapshot", tags=["Telemetry"])
|
|
async def get_all_stats() -> Dict[str, Any]:
|
|
"""Gather complete live snapshot of GPU hardware, host RAM, Ollama, ComfyUI, and switch history."""
|
|
gpu_stats = vram_arbitrator.get_gpu_hardware_stats()
|
|
mem_stats = ram_optimizer.get_detailed_meminfo()
|
|
ollama_state = await vram_arbitrator.get_ollama_live_state()
|
|
comfy_state = await vram_arbitrator.get_comfyui_live_state()
|
|
history = vram_arbitrator.get_switch_history()
|
|
comfy_models = ram_optimizer.find_comfy_model_files()
|
|
arbitrator_status = vram_arbitrator.arbitrator.get_status()
|
|
|
|
return {
|
|
"timestamp": asyncio.get_event_loop().time(),
|
|
"gpu": gpu_stats,
|
|
"ram": mem_stats,
|
|
"ollama": ollama_state,
|
|
"comfyui": comfy_state,
|
|
"arbitrator": arbitrator_status,
|
|
"history": history,
|
|
"comfy_models_count": len(comfy_models),
|
|
}
|
|
|
|
|
|
@app.get("/api/gpu", summary="GPU Sensors and VRAM Breakdown", tags=["Telemetry"])
|
|
async def get_gpu_metrics() -> Dict[str, Any]:
|
|
"""Retrieve detailed NVML sensors (utilization %, temp, power, fan, clocks, and per-process VRAM allocation)."""
|
|
return vram_arbitrator.get_gpu_hardware_stats()
|
|
|
|
@app.get("/api/memory", summary="Host RAM and Page Cache Breakdown", tags=["Telemetry"])
|
|
async def get_ram_metrics() -> Dict[str, Any]:
|
|
"""Retrieve precise host 64GB DDR5 RAM breakdown, active cache size, and cache hit ratios."""
|
|
return ram_optimizer.get_detailed_meminfo()
|
|
|
|
@app.get("/api/stream", summary="Real-Time SSE Telemetry Stream", tags=["Telemetry"])
|
|
async def sse_telemetry_stream(request: Request):
|
|
"""Server-Sent Events (SSE) streaming real-time statistics at 1Hz for dynamic dashboards."""
|
|
async def event_generator():
|
|
while True:
|
|
if await request.is_disconnected():
|
|
break
|
|
try:
|
|
stats = await get_all_stats()
|
|
yield f"data: {json.dumps(stats)}\n\n"
|
|
except Exception as e:
|
|
logger.error(f"SSE stream error: {e}")
|
|
yield f"data: {json.dumps({'error': str(e)})}\n\n"
|
|
await asyncio.sleep(1.0)
|
|
|
|
return StreamingResponse(
|
|
event_generator(),
|
|
media_type="text/event-stream",
|
|
headers={
|
|
"Cache-Control": "no-cache",
|
|
"Connection": "keep-alive",
|
|
"X-Accel-Buffering": "no",
|
|
}
|
|
)
|
|
|
|
@app.post("/api/switch-model", summary="Hot-Swap Ollama LLM in VRAM", tags=["Orchestration"])
|
|
async def api_switch_model(req: SwitchRequest):
|
|
"""Hot-swap the active Ollama model in VRAM and measure exact load duration and token evaluation speed."""
|
|
res = await vram_arbitrator.switch_ollama_model(req.model, keep_alive=req.keep_alive or "30m")
|
|
if not res.get("success"):
|
|
raise HTTPException(status_code=500, detail=res.get("error"))
|
|
return res
|
|
|
|
@app.post("/api/free-vram", summary="Soft-Yield Ollama VRAM", tags=["Orchestration"])
|
|
async def api_free_vram():
|
|
"""Instruct Ollama to instantly yield VRAM to 0MB in ~15ms while preserving model weights in the 64GB host RAM page cache."""
|
|
return await vram_arbitrator.instant_free_ollama_vram()
|
|
|
|
@app.post("/api/comfy-free", summary="Purge ComfyUI VRAM Cache", tags=["Orchestration"])
|
|
async def api_comfy_free():
|
|
"""Purge loaded diffusion models and VRAM cache from the ComfyUI pipeline."""
|
|
return await vram_arbitrator.instant_free_comfyui_vram()
|
|
|
|
@app.post("/api/warm-all", summary="Pre-warm All Models into RAM Cache", tags=["Memory Optimization"])
|
|
async def api_warm_all():
|
|
"""Pre-fault and read all installed Ollama GGUF models and ComfyUI Safetensors checkpoints into the Linux OS Page Cache."""
|
|
return await ram_optimizer.warm_all_models()
|
|
|
|
@app.post("/api/warm-model", summary="Pre-warm Single Model or File", tags=["Memory Optimization"])
|
|
async def api_warm_model(req: WarmRequest):
|
|
"""Pre-warm a specific Ollama model or individual file path into Linux RAM cache."""
|
|
if req.model_name:
|
|
return await ram_optimizer.warm_ollama_model(req.model_name, keep_alive="1m")
|
|
elif req.filepath:
|
|
return ram_optimizer.warm_file_to_ram(req.filepath)
|
|
else:
|
|
raise HTTPException(status_code=400, detail="model_name or filepath required")
|
|
|
|
@app.get("/api/models", summary="List All Installed Models", tags=["Catalog"])
|
|
async def api_get_models():
|
|
"""List all installed Ollama models and discovered ComfyUI model checkpoints/safetensors on disk with sizes and quantization levels."""
|
|
ollama_state = await vram_arbitrator.get_ollama_live_state()
|
|
comfy_models = ram_optimizer.find_comfy_model_files()
|
|
return {
|
|
"ollama_models": ollama_state.get("installed_models", []),
|
|
"comfy_models": comfy_models,
|
|
}
|
|
|
|
@app.get("/api/history", summary="Model Switch History Log", tags=["Analytics"])
|
|
async def api_get_history(limit: int = Query(20, description="Max history items to return")):
|
|
"""Get the recent history log of model switch events, swap durations (in ms), and RAM cache hit status."""
|
|
history = vram_arbitrator.get_switch_history()
|
|
return history[:limit]
|
|
|
|
@app.post("/api/benchmark", summary="Run Latency Benchmark", tags=["Analytics"])
|
|
async def api_run_benchmark(req: BenchmarkRequest):
|
|
"""Run an automated benchmark swapping between available models to measure round-trip latency and RAM cache effectiveness."""
|
|
from mcp_server import run_model_switch_benchmark
|
|
res_str = await run_model_switch_benchmark(iterations=req.iterations or 2)
|
|
return json.loads(res_str)
|
|
|
|
# ==========================================
|
|
# OVERCLOCK MANAGEMENT
|
|
# ==========================================
|
|
|
|
@app.get("/api/overclock", summary="Overclock Status & Profiles", tags=["Overclock"])
|
|
async def api_overclock_status():
|
|
"""Get live GPU overclock state, active profile, and all per-app profiles."""
|
|
return overclock_manager.get_status()
|
|
|
|
@app.post("/api/overclock/apply", summary="Apply Overclock Profile", tags=["Overclock"])
|
|
async def api_overclock_apply(req: OverclockApplyRequest):
|
|
"""Apply a named overclock profile (ollama | comfy | balanced) to the GPU immediately."""
|
|
res = overclock_manager.apply_profile(req.profile)
|
|
if not res.get("success"):
|
|
raise HTTPException(status_code=400, detail=res.get("error"))
|
|
return res
|
|
|
|
@app.get("/api/overclock/profiles", summary="List Overclock Profiles", tags=["Overclock"])
|
|
async def api_overclock_profiles():
|
|
"""List all overclock profiles with their current settings."""
|
|
return overclock_manager.get_profiles()
|
|
|
|
@app.post("/api/overclock/profiles/{name}", summary="Update Overclock Profile", tags=["Overclock"])
|
|
async def api_overclock_update_profile(name: str, req: OverclockProfileUpdate):
|
|
"""Update a profile's settings (persisted to disk)."""
|
|
res = overclock_manager.set_profile(name, req.config)
|
|
if not res.get("success"):
|
|
raise HTTPException(status_code=400, detail=res.get("error"))
|
|
return {"success": True, "profile": name, "profiles": res.get("profiles")}
|
|
|
|
@app.get("/api/overclock/fan", summary="Get GPU Fan Status", tags=["Overclock"])
|
|
@app.get("/api/gpu/fan", summary="Get GPU Fan Status", tags=["Overclock"])
|
|
async def api_get_fan_status():
|
|
"""Get current GPU fan control mode and speed."""
|
|
return overclock_manager.get_fan_status()
|
|
|
|
@app.post("/api/overclock/fan", summary="Set GPU Fan Speed", tags=["Overclock"])
|
|
@app.post("/api/gpu/fan", summary="Set GPU Fan Speed", tags=["Overclock"])
|
|
async def api_set_fan(req: FanRequest):
|
|
"""Set the GPU fan to manual speed (30-100%) or back to automatic control."""
|
|
pct = req.percent if req.percent is not None else req.speed_pct
|
|
if req.mode == "manual" and pct is not None:
|
|
return overclock_manager.set_fan_speed(pct)
|
|
return overclock_manager.set_fan_auto()
|
|
|
|
# Mount static web UI files
|
|
app.mount("/static", StaticFiles(directory="/home/drjones/unified-model-manager/static"), name="static")
|
|
|
|
@app.get("/", summary="Dashboard Web UI", tags=["UI"])
|
|
async def root_index():
|
|
with open("/home/drjones/unified-model-manager/static/index.html", "r") as f:
|
|
content = f.read()
|
|
return HTMLResponse(content=content)
|
|
|
|
if __name__ == "__main__":
|
|
import uvicorn
|
|
uvicorn.run("server:app", host="0.0.0.0", port=9090, reload=False, log_level="info")
|