Initial commit: HyperSwap GPU Program Swapper with REST API, MCP 2.0, Real-Time Dashboard and Memory Orchestrator

mcp_server.py

@@ -0,0 +1,153 @@
+"""Model Context Protocol (MCP) Server for GPU Program Swapper and Memory Orchestrator."""
+import os
+import sys
+import json
+import asyncio
+import logging
+from typing import Dict, List, Any, Optional
+
+from mcp.server import MCPServer
+import ram_optimizer
+import vram_arbitrator
+
+logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(name)s: %(message)s")
+logger = logging.getLogger("gpu_swapper_mcp")
+
+mcp = MCPServer(
+    name="gpu-program-swapper",
+    version="1.0.0",
+    description="Orchestrates high-speed GPU VRAM hot-swaps between Ollama LLMs and ComfyUI with 64GB RAM cache telemetry."
+)
+
+# ==========================================
+# MCP TOOLS
+# ==========================================
+
+@mcp.tool()
+def get_gpu_status() -> str:
+    """Get live NVIDIA GPU hardware telemetry, VRAM allocation (Ollama vs ComfyUI vs System), temperature, power, and active compute PIDs."""
+    stats = vram_arbitrator.get_gpu_hardware_stats()
+    return json.dumps(stats, indent=2)
+
+@mcp.tool()
+def get_host_memory_status() -> str:
+    """Get host system RAM breakdown (Total, Apps Used, OS Page Cache containing models in RAM, Free memory) and cache hit percentage."""
+    mem = ram_optimizer.get_detailed_meminfo()
+    return json.dumps(mem, indent=2)
+
+@mcp.tool()
+async def switch_ollama_model(model_name: str, keep_alive: str = "30m") -> str:
+    """Hot-swap active Ollama LLM in VRAM. Measures swap latency (ms), eval speed (tokens/sec), and whether it was an instant RAM cache hit."""
+    res = await vram_arbitrator.switch_ollama_model(model_name, keep_alive=keep_alive)
+    return json.dumps(res, indent=2)
+
+@mcp.tool()
+async def soft_yield_ollama_vram(model_name: Optional[str] = None) -> str:
+    """Soft-yield Ollama VRAM down to 0 MB in ~15 milliseconds while keeping all model weights resident in 64GB host RAM cache."""
+    res = await vram_arbitrator.instant_free_ollama_vram(model_name)
+    return json.dumps(res, indent=2)
+
+@mcp.tool()
+async def purge_comfyui_vram() -> str:
+    """Purge loaded diffusion models and VRAM cache from the ComfyUI pipeline."""
+    res = await vram_arbitrator.instant_free_comfyui_vram()
+    return json.dumps(res, indent=2)
+
+@mcp.tool()
+async def prewarm_all_models_to_ram() -> str:
+    """Pre-read and fault all installed Ollama GGUF models and ComfyUI Safetensors checkpoints into the Linux OS Page Cache for PCIe-speed hot swapping."""
+    res = await ram_optimizer.warm_all_models()
+    return json.dumps(res, indent=2)
+
+@mcp.tool()
+async def prewarm_single_model(model_name: Optional[str] = None, filepath: Optional[str] = None) -> str:
+    """Pre-warm a specific Ollama model or individual model file path into Linux RAM cache."""
+    if model_name:
+        res = await ram_optimizer.warm_ollama_model(model_name, keep_alive="1m")
+    elif filepath:
+        res = ram_optimizer.warm_file_to_ram(filepath)
+    else:
+        res = {"error": "Either model_name or filepath must be provided"}
+    return json.dumps(res, indent=2)
+
+@mcp.tool()
+async def list_available_models() -> str:
+    """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()
+    result = {
+        "ollama_models": ollama_state.get("installed_models", []),
+        "comfy_models": comfy_models,
+    }
+    return json.dumps(result, indent=2)
+
+@mcp.tool()
+def get_switch_history(limit: int = 20) -> str:
+    """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 json.dumps(history[:limit], indent=2)
+
+@mcp.tool()
+async def run_model_switch_benchmark(iterations: int = 2) -> str:
+    """Run an automated benchmark swapping between available models to measure round-trip latency and RAM cache effectiveness."""
+    ollama_state = await vram_arbitrator.get_ollama_live_state()
+    installed = [m.get("name") for m in ollama_state.get("installed_models", []) if m.get("name")]
+    if len(installed) < 2:
+        return json.dumps({"error": "Need at least 2 installed models to benchmark", "installed": installed})
+    
+    m1, m2 = installed[0], installed[1]
+    results = []
+    
+    for i in range(iterations):
+        # Swap to m1
+        r1 = await vram_arbitrator.switch_ollama_model(m1, keep_alive="5m")
+        results.append({"iteration": i+1, "direction": f"-> {m1}", "latency_ms": r1.get("total_duration_ms", 0), "load_ms": r1.get("load_duration_ms", 0), "is_ram_hit": r1.get("is_ram_hit", False)})
+        await asyncio.sleep(0.5)
+        # Swap to m2
+        r2 = await vram_arbitrator.switch_ollama_model(m2, keep_alive="5m")
+        results.append({"iteration": i+1, "direction": f"-> {m2}", "latency_ms": r2.get("total_duration_ms", 0), "load_ms": r2.get("load_duration_ms", 0), "is_ram_hit": r2.get("is_ram_hit", False)})
+        await asyncio.sleep(0.5)
+        
+    avg_latency = round(sum(r["latency_ms"] for r in results) / len(results), 2) if results else 0
+    return json.dumps({
+        "status": "completed",
+        "tested_models": [m1, m2],
+        "iterations": iterations,
+        "average_swap_latency_ms": avg_latency,
+        "rounds": results,
+    }, indent=2)
+
+# ==========================================
+# MCP RESOURCES
+# ==========================================
+
+@mcp.resource("gpu://metrics/live")
+def get_live_metrics_resource() -> str:
+    """Live snapshot of GPU hardware sensors, VRAM, and RAM cache."""
+    gpu = vram_arbitrator.get_gpu_hardware_stats()
+    ram = ram_optimizer.get_detailed_meminfo()
+    return json.dumps({"gpu": gpu, "ram": ram}, indent=2)
+
+@mcp.resource("gpu://models/catalog")
+async def get_models_catalog_resource() -> str:
+    """Catalog of all local Ollama and ComfyUI models."""
+    return await list_available_models()
+
+@mcp.resource("gpu://history/switches")
+def get_switch_history_resource() -> str:
+    """Recent model switch events and latencies."""
+    return json.dumps(vram_arbitrator.get_switch_history(), indent=2)
+
+
+if __name__ == "__main__":
+    import argparse
+    parser = argparse.ArgumentParser(description="HyperSwap GPU Program Swapper MCP Server")
+    parser.add_argument("--stdio", action="store_true", help="Run in stdio mode (default)")
+    parser.add_argument("--sse", action="store_true", help="Run with SSE transport")
+    parser.add_argument("--port", type=int, default=8001, help="Port for SSE transport")
+    args = parser.parse_args()
+
+    if args.sse:
+        mcp.run(transport="sse", host="0.0.0.0", port=args.port)
+    else:
+        mcp.run(transport="stdio")