Initial commit: HyperSwap GPU Program Swapper with REST API, MCP 2.0, Real-Time Dashboard and Memory Orchestrator
This commit is contained in:
drjones
176
ram_optimizer.py
Normal file
176
ram_optimizer.py
Normal file
@@ -0,0 +1,176 @@
|
||||
"""RAM Optimizer and Model Pre-warmer for High-Speed Switching."""
|
||||
import os
|
||||
import glob
|
||||
import time
|
||||
import httpx
|
||||
import logging
|
||||
from typing import Dict, List, Any
|
||||
|
||||
logger = logging.getLogger("ram_optimizer")
|
||||
|
||||
OLLAMA_API_BASE = "http://localhost:11434"
|
||||
COMFY_API_BASE = "http://127.0.0.1:8188"
|
||||
COMFY_MODELS_DIR = "/home/drjones/ComfyUI/models"
|
||||
|
||||
def get_detailed_meminfo() -> Dict[str, Any]:
|
||||
"""Parse /proc/meminfo for precise page cache and RAM stats."""
|
||||
info = {}
|
||||
try:
|
||||
with open("/proc/meminfo", "r") as f:
|
||||
for line in f:
|
||||
parts = line.split(":")
|
||||
if len(parts) == 2:
|
||||
key = parts[0].strip()
|
||||
val = parts[1].strip().split()[0]
|
||||
info[key] = int(val) * 1024 # Convert kB to bytes
|
||||
except Exception as e:
|
||||
logger.error(f"Failed to read /proc/meminfo: {e}")
|
||||
|
||||
total = info.get("MemTotal", 0)
|
||||
free = info.get("MemFree", 0)
|
||||
available = info.get("MemAvailable", 0)
|
||||
cached = info.get("Cached", 0) + info.get("Buffers", 0)
|
||||
dirty = info.get("Dirty", 0)
|
||||
used = total - free - cached
|
||||
if used < 0:
|
||||
used = total - available
|
||||
|
||||
return {
|
||||
"total_bytes": total,
|
||||
"total_gb": round(total / (1024**3), 2),
|
||||
"used_bytes": used,
|
||||
"used_gb": round(used / (1024**3), 2),
|
||||
"cached_bytes": cached,
|
||||
"cached_gb": round(cached / (1024**3), 2),
|
||||
"free_bytes": free,
|
||||
"free_gb": round(free / (1024**3), 2),
|
||||
"available_bytes": available,
|
||||
"available_gb": round(available / (1024**3), 2),
|
||||
"dirty_bytes": dirty,
|
||||
"dirty_mb": round(dirty / (1024**2), 2),
|
||||
"cache_ratio_pct": round((cached / total * 100) if total > 0 else 0, 1),
|
||||
}
|
||||
|
||||
def warm_file_to_ram(filepath: str, chunk_size: int = 16 * 1024 * 1024) -> Dict[str, Any]:
|
||||
"""Pre-fault/read file into Linux OS Page Cache at maximum disk read speed."""
|
||||
if not os.path.exists(filepath):
|
||||
return {"success": False, "error": f"File not found: {filepath}", "duration_ms": 0}
|
||||
|
||||
t0 = time.perf_counter()
|
||||
file_size = os.path.getsize(filepath)
|
||||
bytes_read = 0
|
||||
try:
|
||||
with open(filepath, "rb") as f:
|
||||
# Hint kernel that we will read this sequentially
|
||||
try:
|
||||
os.posix_fadvise(f.fileno(), 0, file_size, os.POSIX_FADV_WILLNEED)
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
buf = bytearray(chunk_size)
|
||||
while True:
|
||||
n = f.readinto(buf)
|
||||
if not n:
|
||||
break
|
||||
bytes_read += n
|
||||
|
||||
duration = time.perf_counter() - t0
|
||||
duration_ms = round(duration * 1000, 2)
|
||||
speed_mb_s = round((bytes_read / (1024**2)) / duration if duration > 0 else 0, 2)
|
||||
return {
|
||||
"success": True,
|
||||
"filepath": filepath,
|
||||
"size_bytes": file_size,
|
||||
"size_mb": round(file_size / (1024**2), 2),
|
||||
"bytes_read": bytes_read,
|
||||
"duration_ms": duration_ms,
|
||||
"speed_mb_s": speed_mb_s,
|
||||
}
|
||||
except Exception as e:
|
||||
return {"success": False, "error": str(e), "duration_ms": round((time.perf_counter() - t0) * 1000, 2)}
|
||||
|
||||
async def warm_ollama_model(model_name: str, keep_alive: str = "5m") -> Dict[str, Any]:
|
||||
"""Warm an Ollama model into memory and measure time."""
|
||||
t0 = time.perf_counter()
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=120.0) as client:
|
||||
resp = await client.post(
|
||||
f"{OLLAMA_API_BASE}/api/generate",
|
||||
json={"model": model_name, "prompt": "", "keep_alive": keep_alive},
|
||||
)
|
||||
duration = time.perf_counter() - t0
|
||||
if resp.status_code == 200:
|
||||
data = resp.json()
|
||||
return {
|
||||
"success": True,
|
||||
"model": model_name,
|
||||
"duration_ms": round(duration * 1000, 2),
|
||||
"load_duration_ms": round(data.get("load_duration", 0) / 1e6, 2),
|
||||
"total_duration_ms": round(data.get("total_duration", 0) / 1e6, 2),
|
||||
}
|
||||
else:
|
||||
return {
|
||||
"success": False,
|
||||
"model": model_name,
|
||||
"error": f"HTTP {resp.status_code}: {resp.text}",
|
||||
"duration_ms": round(duration * 1000, 2),
|
||||
}
|
||||
except Exception as e:
|
||||
return {"success": False, "model": model_name, "error": str(e), "duration_ms": round((time.perf_counter() - t0) * 1000, 2)}
|
||||
|
||||
def find_comfy_model_files() -> List[Dict[str, Any]]:
|
||||
"""Discover all model files under ComfyUI models."""
|
||||
results = []
|
||||
extensions = ("*.safetensors", "*.ckpt", "*.pt", "*.bin")
|
||||
if os.path.exists(COMFY_MODELS_DIR):
|
||||
for root, _, files in os.walk(COMFY_MODELS_DIR):
|
||||
for file in files:
|
||||
if any(file.endswith(ext.replace("*", "")) for ext in extensions):
|
||||
full_path = os.path.join(root, file)
|
||||
rel_path = os.path.relpath(full_path, COMFY_MODELS_DIR)
|
||||
size = os.path.getsize(full_path)
|
||||
results.append({
|
||||
"filename": file,
|
||||
"rel_path": rel_path,
|
||||
"full_path": full_path,
|
||||
"size_bytes": size,
|
||||
"size_mb": round(size / (1024**2), 2),
|
||||
"size_gb": round(size / (1024**3), 3),
|
||||
})
|
||||
return results
|
||||
|
||||
async def warm_all_models() -> Dict[str, Any]:
|
||||
"""Warm all available Ollama and ComfyUI models into Linux RAM Cache."""
|
||||
t0 = time.perf_counter()
|
||||
warmed_ollama = []
|
||||
warmed_comfy = []
|
||||
|
||||
# 1. Ollama models
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=10.0) as client:
|
||||
tags_resp = await client.get(f"{OLLAMA_API_BASE}/api/tags")
|
||||
if tags_resp.status_code == 200:
|
||||
models = tags_resp.json().get("models", [])
|
||||
for m in models:
|
||||
name = m.get("name")
|
||||
res = await warm_ollama_model(name, keep_alive="1m")
|
||||
warmed_ollama.append(res)
|
||||
except Exception as e:
|
||||
logger.error(f"Error discovering Ollama models: {e}")
|
||||
|
||||
# 2. ComfyUI models
|
||||
comfy_files = find_comfy_model_files()
|
||||
for f in comfy_files:
|
||||
res = warm_file_to_ram(f["full_path"])
|
||||
warmed_comfy.append(res)
|
||||
|
||||
total_duration_ms = round((time.perf_counter() - t0) * 1000, 2)
|
||||
meminfo = get_detailed_meminfo()
|
||||
|
||||
return {
|
||||
"status": "completed",
|
||||
"total_duration_ms": total_duration_ms,
|
||||
"ollama_models_warmed": warmed_ollama,
|
||||
"comfy_files_warmed": warmed_comfy,
|
||||
"meminfo_after": meminfo,
|
||||
}
|
||||
Reference in New Issue
Block a user