Initial commit: HyperSwap GPU Program Swapper with REST API, MCP 2.0, Real-Time Dashboard and Memory Orchestrator
This commit is contained in:
drjones
174
server.py
Normal file
174
server.py
Normal file
@@ -0,0 +1,174 @@
|
||||
"""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
|
||||
|
||||
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=["*"],
|
||||
)
|
||||
|
||||
# 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"])
|
||||
|
||||
|
||||
# ==========================================
|
||||
# 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()
|
||||
|
||||
return {
|
||||
"timestamp": asyncio.get_event_loop().time(),
|
||||
"gpu": gpu_stats,
|
||||
"ram": mem_stats,
|
||||
"ollama": ollama_state,
|
||||
"comfyui": comfy_state,
|
||||
"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)
|
||||
|
||||
# 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")
|
||||
Reference in New Issue
Block a user