170 lines
7.2 KiB
Python
170 lines
7.2 KiB
Python
"""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
|
|
import overclock_manager
|
|
|
|
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.tool()
|
|
def get_gpu_fan_status() -> str:
|
|
"""Get current GPU fan control mode (auto vs manual) and target speed."""
|
|
status = overclock_manager.get_fan_status()
|
|
return json.dumps(status, indent=2)
|
|
|
|
@mcp.tool()
|
|
def set_gpu_fan_speed(mode: str = "auto", percent: Optional[int] = None) -> str:
|
|
"""Set GPU fan speed mode ('auto' or 'manual') with target percent (30-100%)."""
|
|
if mode.lower() == "manual" and percent is not None:
|
|
res = overclock_manager.set_fan_speed(percent)
|
|
else:
|
|
res = overclock_manager.set_fan_auto()
|
|
return json.dumps(res, 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")
|