Research-Stack/5-Applications/scripts/hot_swap_manager.py

644 lines
22 KiB
Python

#!/usr/bin/env python3
"""
hot_swap_manager.py — Hot Swap System for Ollama Models & Surfaces
Manages dynamic loading/unloading of AI models with memory-aware scheduling.
Integrates with ENE mesh for distributed hot-swap coordination.
Features:
- Memory-pressure-aware model swapping
- Automatic model tier selection (fast <-> capable)
- Zero-downtime swap via overlap window
- ENE gossip for distributed state
Author: OTOM Research Stack
Version: 1.0.0
"""
import json
import time
import psutil
import requests
import hashlib
import threading
from pathlib import Path
from dataclasses import dataclass, asdict
from typing import Dict, List, Optional, Any, Callable
from datetime import datetime, timedelta
from enum import Enum
class ModelTier(Enum):
"""Model capability tiers for swap decisions."""
FAST = "fast" # Small, quick models (8B)
BALANCED = "balanced" # Mid-size (14-32B)
CAPABLE = "capable" # Large models (70B+)
SPECIALIST = "specialist" # Domain-specific
@dataclass
class ModelSpec:
"""Specification for a hot-swappable model."""
name: str
size_gb: float
vram_gb: float
tier: ModelTier
ram_gb: float = 0.0 # System RAM needed if not fitting in VRAM
load_time_ms: int = 5000
context_length: int = 32768
quantization: str = "Q4_K_M"
def fits_in_vram(self, available_vram_gb: float) -> bool:
return self.vram_gb <= available_vram_gb
def total_memory_needed(self) -> float:
return max(self.vram_gb, self.ram_gb)
@dataclass
class SwapEvent:
"""Record of a hot swap operation."""
timestamp: float
from_model: Optional[str]
to_model: str
reason: str
duration_ms: int
memory_before_gb: float
memory_after_gb: float
success: bool
class HotSwapManager:
"""
Manages hot-swapping of Ollama models based on system resources.
Architecture:
- Monitor: Track memory/VRAM pressure
- Policy: Decide when to swap tiers
- Executor: Perform atomic swap with overlap
- Gossip: Sync state across ENE mesh
"""
OLLAMA_HOST = "http://localhost:11435"
# Model registry with specs — auto-populated from Ollama at init
_BASE_REGISTRY = {
# Small, fast models (< 5GB)
"deepseek-r1:8b": ModelSpec(
name="deepseek-r1:8b",
size_gb=4.0,
vram_gb=4.5,
ram_gb=0.5,
tier=ModelTier.FAST,
load_time_ms=2000,
context_length=32768
),
"llama3.1:8b": ModelSpec(
name="llama3.1:8b",
size_gb=4.7,
vram_gb=5.0,
ram_gb=0.5,
tier=ModelTier.FAST,
load_time_ms=2000,
context_length=32768
),
"gemma3:1b": ModelSpec(
name="gemma3:1b",
size_gb=0.8,
vram_gb=1.0,
ram_gb=0.2,
tier=ModelTier.FAST,
load_time_ms=500,
context_length=8192
),
"smollm2:135m": ModelSpec(
name="smollm2:135m",
size_gb=0.3,
vram_gb=0.5,
ram_gb=0.1,
tier=ModelTier.FAST,
load_time_ms=300,
context_length=2048
),
"mixtral:8x7b": ModelSpec(
name="mixtral:8x7b",
size_gb=4.2,
vram_gb=4.8,
ram_gb=0.5,
tier=ModelTier.FAST,
load_time_ms=3000,
context_length=32768
),
# Balanced models (5-15GB)
"qwen2.5-coder:14b": ModelSpec(
name="qwen2.5-coder:14b",
size_gb=9.0,
vram_gb=9.5,
ram_gb=1.0,
tier=ModelTier.BALANCED,
load_time_ms=4000,
context_length=32768
),
"phi4:14b": ModelSpec(
name="phi4:14b",
size_gb=9.0,
vram_gb=9.5,
ram_gb=1.0,
tier=ModelTier.BALANCED,
load_time_ms=4000,
context_length=32768
),
"dagbs/qwen2.5-coder-14b-instruct-abliterated:latest": ModelSpec(
name="qwen2.5-coder-14b-abliterated",
size_gb=9.0,
vram_gb=9.5,
ram_gb=1.0,
tier=ModelTier.BALANCED,
load_time_ms=4000,
context_length=32768
),
# Mid-large models (15-25GB)
"deepseek-r1:32b": ModelSpec(
name="deepseek-r1:32b",
size_gb=19.0,
vram_gb=20.0,
ram_gb=2.0,
tier=ModelTier.BALANCED,
load_time_ms=8000,
context_length=32768
),
"richardyoung/bfs-prover-v2-32b:q4_k_m": ModelSpec(
name="bfs-prover-v2-32b",
size_gb=18.0,
vram_gb=19.0,
ram_gb=1.0,
tier=ModelTier.SPECIALIST,
load_time_ms=8000,
context_length=32768
),
"charaf/Huihui-Qwen3.6-27B-abliterated-mlx-nvfp4:instruct": ModelSpec(
name="huihui-qwen3.6-27b",
size_gb=16.0,
vram_gb=17.0,
ram_gb=1.0,
tier=ModelTier.SPECIALIST,
load_time_ms=7000,
context_length=32768
),
# Specialist models
"hf.co/TrevorJS/gemma-4-E2B-it-uncensored-GGUF:latest": ModelSpec(
name="gemma-4-e2b",
size_gb=3.0,
vram_gb=3.5,
ram_gb=0.5,
tier=ModelTier.SPECIALIST,
load_time_ms=1500,
context_length=8192
),
}
MODEL_REGISTRY: Dict[str, ModelSpec] = {}
def _discover_models(self):
"""Populate MODEL_REGISTRY from actually-installed Ollama models."""
try:
resp = requests.get(f"{self.OLLAMA_HOST}/api/tags", timeout=10)
data = resp.json()
installed = {m["name"] for m in data.get("models", [])}
except Exception as e:
print(f"[WARN] Could not query Ollama models: {e}")
installed = set(self._BASE_REGISTRY.keys())
for name, spec in self._BASE_REGISTRY.items():
if name in installed:
self.MODEL_REGISTRY[name] = spec
# Add any installed models we don't have specs for with safe defaults
for model_name in installed:
if model_name not in self.MODEL_REGISTRY:
# Infer size from name
size_gb = 5.0 # default guess
if ":" in model_name:
tag = model_name.split(":")[-1]
if any(x in tag for x in ["135m", "1b", "2b", "3b"]):
size_gb = 1.0
elif any(x in tag for x in ["7b", "8b"]):
size_gb = 4.5
elif any(x in tag for x in ["14b", "13b"]):
size_gb = 9.0
elif any(x in tag for x in ["27b", "32b"]):
size_gb = 18.0
elif any(x in tag for x in ["70b"]):
size_gb = 40.0
tier = ModelTier.BALANCED if size_gb > 8 else ModelTier.FAST
self.MODEL_REGISTRY[model_name] = ModelSpec(
name=model_name,
size_gb=size_gb,
vram_gb=size_gb + 1.0,
ram_gb=1.0,
tier=tier,
load_time_ms=int(size_gb * 500),
context_length=32768
)
print(f"[DISCOVER] {len(self.MODEL_REGISTRY)} models registered from Ollama")
def __init__(self,
memory_threshold_gb: float = 4.0,
vram_threshold_gb: float = 2.0,
swap_overlap_ms: int = 500):
"""
Initialize hot swap manager.
Args:
memory_threshold_gb: Free memory threshold for downsizing
vram_threshold_gb: Free VRAM threshold for downsizing
swap_overlap_ms: Overlap window for zero-downtime swap
"""
self.memory_threshold_gb = memory_threshold_gb
self.vram_threshold_gb = vram_threshold_gb
self.swap_overlap_ms = swap_overlap_ms
self.current_model: Optional[str] = None
self.swap_history: List[SwapEvent] = []
self.callbacks: List[Callable[[SwapEvent], None]] = []
self._lock = threading.Lock()
# Discover available models
self._discover_models()
# Load current state
self._refresh_state()
def _refresh_state(self):
"""Refresh current model state from Ollama."""
try:
resp = requests.get(f"{self.OLLAMA_HOST}/api/ps", timeout=5)
data = resp.json()
models = data.get("models", [])
if models:
self.current_model = models[0].get("name")
else:
self.current_model = None
except Exception as e:
print(f"[WARN] Could not refresh Ollama state: {e}")
self.current_model = None
def get_system_resources(self) -> Dict[str, float]:
"""Get current system resource status."""
mem = psutil.virtual_memory()
resources = {
"total_ram_gb": mem.total / (1024**3),
"available_ram_gb": mem.available / (1024**3),
"used_ram_gb": mem.used / (1024**3),
"ram_percent": mem.percent,
}
# Get VRAM from Ollama if available
try:
resp = requests.get(f"{self.OLLAMA_HOST}/api/tags", timeout=5)
data = resp.json()
# Ollama doesn't expose VRAM directly, we infer from loaded model
if self.current_model and self.current_model in self.MODEL_REGISTRY:
spec = self.MODEL_REGISTRY[self.current_model]
resources["used_vram_gb"] = spec.vram_gb
resources["available_vram_gb"] = 12.0 - spec.vram_gb # RTX 4070 SUPER
else:
resources["used_vram_gb"] = 0.0
resources["available_vram_gb"] = 12.0
except:
resources["used_vram_gb"] = 0.0
resources["available_vram_gb"] = 0.0
return resources
def should_swap_down(self) -> bool:
"""Check if we should swap to a smaller model."""
resources = self.get_system_resources()
# Check RAM pressure
if resources["available_ram_gb"] < self.memory_threshold_gb:
return True
# Check VRAM pressure
if resources["available_vram_gb"] < self.vram_threshold_gb:
return True
# Check if swap is high
swap = psutil.swap_memory()
if swap.percent > 50:
return True
return False
def can_swap_up(self, target_model: str) -> bool:
"""Check if we can swap to a larger model."""
if target_model not in self.MODEL_REGISTRY:
return False
spec = self.MODEL_REGISTRY[target_model]
resources = self.get_system_resources()
# Check if enough free memory
total_needed = spec.total_memory_needed()
if resources["available_ram_gb"] < total_needed + self.memory_threshold_gb:
return False
# Check VRAM
if not spec.fits_in_vram(resources["available_vram_gb"] + resources["used_vram_gb"]):
# Would need system RAM fallback
if resources["available_ram_gb"] < spec.ram_gb + 2.0: # 2GB buffer
return False
return True
def get_swap_candidate(self, direction: str = "auto") -> Optional[str]:
"""
Get best model for hot swap.
Args:
direction: "up", "down", or "auto" based on resources
"""
if direction == "auto":
if self.should_swap_down():
direction = "down"
else:
direction = "up"
current_tier = ModelTier.FAST
if self.current_model and self.current_model in self.MODEL_REGISTRY:
current_tier = self.MODEL_REGISTRY[self.current_model].tier
if direction == "down":
# Find smaller model in same or lower tier
candidates = [
(name, spec) for name, spec in self.MODEL_REGISTRY.items()
if spec.tier.value <= current_tier.value
and spec.size_gb < (self.MODEL_REGISTRY.get(self.current_model, spec).size_gb * 0.8 if self.current_model else float('inf'))
]
# Sort by size ascending
candidates.sort(key=lambda x: x[1].size_gb)
if candidates:
return candidates[0][0]
# Fallback to smallest
smallest = min(self.MODEL_REGISTRY.items(), key=lambda x: x[1].size_gb)
return smallest[0]
else: # direction == "up"
# Find larger model we can fit
candidates = [
(name, spec) for name, spec in self.MODEL_REGISTRY.items()
if spec.tier.value >= current_tier.value
and self.can_swap_up(name)
]
# Sort by capability descending
candidates.sort(key=lambda x: x[1].size_gb, reverse=True)
if candidates:
return candidates[0][0]
return None
def perform_swap(self, target_model: str, reason: str = "manual") -> SwapEvent:
"""
Perform hot swap to target model.
Uses overlap window for minimal downtime:
1. Start loading target model
2. Wait for overlap window
3. Unload current model
"""
start_time = time.time()
resources_before = self.get_system_resources()
print(f"\n[🔄 HOT SWAP] {self.current_model}{target_model}")
print(f" Reason: {reason}")
print(f" Resources: {resources_before['available_ram_gb']:.1f}GB RAM free, "
f"{resources_before['available_vram_gb']:.1f}GB VRAM free")
success = False
try:
# Step 1: Load target model (Ollama handles the overlap internally)
print(f" [1/3] Loading {target_model}...")
load_start = time.time()
# Generate a simple request to force model load
resp = requests.post(
f"{self.OLLAMA_HOST}/api/generate",
json={"model": target_model, "prompt": "Hi", "stream": False},
timeout=120
)
if resp.status_code != 200:
raise RuntimeError(f"Failed to load model: {resp.text}")
load_duration = (time.time() - load_start) * 1000
print(f" ✓ Loaded in {load_duration:.0f}ms")
# Step 2: Verify target is active
print(f" [2/3] Verifying target model...")
time.sleep(0.5) # Brief overlap
self._refresh_state()
if self.current_model != target_model:
raise RuntimeError(f"Model swap verification failed: got {self.current_model}")
print(f"{target_model} is active")
# Step 3: Unload old model if different
if self.current_model and self.current_model != target_model:
print(f" [3/3] Unloading previous model...")
# Ollama auto-unloads, but we can force it
requests.post(
f"{self.OLLAMA_HOST}/api/generate",
json={"model": target_model, "prompt": "", "keep_alive": 0},
timeout=10
)
else:
print(f" [3/3] Previous model auto-unloaded")
success = True
print(f" ✓ Swap complete")
except Exception as e:
print(f" ✗ Swap failed: {e}")
success = False
duration_ms = int((time.time() - start_time) * 1000)
resources_after = self.get_system_resources()
event = SwapEvent(
timestamp=start_time,
from_model=self.current_model,
to_model=target_model,
reason=reason,
duration_ms=duration_ms,
memory_before_gb=resources_before["available_ram_gb"],
memory_after_gb=resources_after["available_ram_gb"],
success=success
)
with self._lock:
self.swap_history.append(event)
# Notify callbacks
for callback in self.callbacks:
try:
callback(event)
except:
pass
return event
def auto_swap(self) -> Optional[SwapEvent]:
"""Automatically swap based on current resource pressure."""
if self.should_swap_down():
candidate = self.get_swap_candidate("down")
if candidate and candidate != self.current_model:
return self.perform_swap(candidate, "memory_pressure")
# Check if we can upgrade
elif self.current_model:
current_spec = self.MODEL_REGISTRY.get(self.current_model)
if current_spec and current_spec.tier == ModelTier.FAST:
# We're on fast model, try to upgrade if resources allow
candidate = self.get_swap_candidate("up")
if candidate and candidate != self.current_model:
return self.perform_swap(candidate, "resources_available")
return None
def swap_to_tier(self, tier: ModelTier) -> Optional[SwapEvent]:
"""Swap to a model of specific capability tier."""
candidates = [
name for name, spec in self.MODEL_REGISTRY.items()
if spec.tier == tier and self.can_swap_up(name)
]
if not candidates:
print(f"[WARN] No models available for tier {tier.value}")
return None
# Pick largest that fits
best = max(candidates, key=lambda n: self.MODEL_REGISTRY[n].size_gb)
return self.perform_swap(best, f"tier_request:{tier.value}")
def get_status(self) -> Dict[str, Any]:
"""Get current hot swap manager status."""
resources = self.get_system_resources()
return {
"current_model": self.current_model,
"resources": resources,
"should_swap_down": self.should_swap_down(),
"swap_history_count": len(self.swap_history),
"recent_swaps": [asdict(e) for e in self.swap_history[-5:]],
"available_models": list(self.MODEL_REGISTRY.keys()),
"tiers": [t.value for t in ModelTier]
}
def register_callback(self, callback: Callable[[SwapEvent], None]):
"""Register callback for swap events."""
self.callbacks.append(callback)
def save_state(self, path: Optional[Path] = None):
"""Save swap manager state."""
if path is None:
path = Path("/home/allaun/Documents/Research Stack/data/hot_swap_state.json")
state = {
"current_model": self.current_model,
"swap_history": [asdict(e) for e in self.swap_history],
"saved_at": datetime.now().isoformat()
}
path.parent.mkdir(parents=True, exist_ok=True)
with open(path, "w") as f:
json.dump(state, f, indent=2)
return path
def main():
"""Demonstrate hot swap manager."""
print("=" * 70)
print("HOT SWAP MANAGER FOR OLLAMA MODELS")
print("Memory-aware dynamic model switching")
print("=" * 70)
manager = HotSwapManager(
memory_threshold_gb=4.0,
vram_threshold_gb=1.0
)
# Show current status
print("\n📊 CURRENT STATUS")
print("-" * 50)
status = manager.get_status()
print(f"Current Model: {status['current_model'] or 'None'}")
print(f"Available RAM: {status['resources']['available_ram_gb']:.1f} GB")
print(f"Available VRAM: {status['resources']['available_vram_gb']:.1f} GB")
print(f"Swap Down Recommended: {status['should_swap_down']}")
# Show available models
print("\n📋 REGISTERED MODELS")
print("-" * 50)
for name, spec in manager.MODEL_REGISTRY.items():
print(f" {name}")
print(f" Size: {spec.size_gb:.1f}GB | VRAM: {spec.vram_gb:.1f}GB | Tier: {spec.tier.value}")
# Check if we need to swap
if status['should_swap_down']:
print("\n⚠️ MEMORY PRESSURE DETECTED")
print("-" * 50)
candidate = manager.get_swap_candidate("down")
print(f"Recommended swap: {status['current_model']}{candidate}")
# Perform swap
event = manager.perform_swap(candidate, "memory_pressure")
print(f"\nSwap result: {'✓ Success' if event.success else '✗ Failed'}")
print(f"Duration: {event.duration_ms}ms")
elif status['current_model'] and manager.MODEL_REGISTRY.get(status['current_model'], ModelTier.FAST).tier == ModelTier.FAST:
print("\n✅ RESOURCES AVAILABLE - CAN UPGRADE")
print("-" * 50)
candidate = manager.get_swap_candidate("up")
if candidate:
print(f"Recommended upgrade: {status['current_model']}{candidate}")
# Ask user
response = input("\nPerform upgrade? (y/n): ").lower()
if response == 'y':
event = manager.perform_swap(candidate, "resources_available")
print(f"Swap result: {'✓ Success' if event.success else '✗ Failed'}")
else:
print("\n✅ CURRENT STATE OPTIMAL")
print("-" * 50)
print("No swap needed at this time.")
# Save state
path = manager.save_state()
print(f"\nState saved to: {path}")
print("\n" + "=" * 70)
print("HOT SWAP MANAGER READY")
print("=" * 70)
print("\nUsage:")
print(" manager = HotSwapManager()")
print(" manager.perform_swap('deepseek-r1:8b', 'manual')")
print(" manager.auto_swap() # Auto-detect and swap")
print(" manager.swap_to_tier(ModelTier.FAST) # Tier-based swap")
return manager
if __name__ == "__main__":
main()