""" Memory Monitor — Tracks RAM usage and auto-unloads idle model shards. This ensures the system never runs out of RAM: 1. Monitors current RAM usage 2. Tracks which shards are loaded and when they were last used 3. When RAM is tight, unloads the least-recently-used shards 4. When a shard is needed again, reloads it from disk (via mmap) """ import os import time import psutil import threading from dataclasses import dataclass, field from typing import Dict, Optional from collections import OrderedDict @dataclass class ShardInfo: """Info about a loaded model shard.""" shard_id: int size_mb: float # Size of this shard in MB loaded_at: float # When it was loaded (timestamp) last_used: float # When it was last accessed access_count: int = 0 # How many times it was used class MemoryMonitor: """ Monitors RAM usage and manages shard lifecycle. Usage: monitor = MemoryMonitor(max_ram_gb=8.0) # Before loading a shard: if monitor.can_load(shard_size_mb=500): monitor.on_load(shard_id=3, size_mb=500) # After using a shard: monitor.on_access(shard_id=3) # When RAM is tight: monitor.evict_if_needed() # Unloads LRU shards """ def __init__(self, max_ram_gb: float = 8.0, eviction_threshold: float = 0.85): """ Args: max_ram_gb: Maximum RAM to use for model shards (default 8GB) eviction_threshold: Start evicting when RAM usage exceeds this (85%) """ self.max_ram_bytes = int(max_ram_gb * 1024 * 1024 * 1024) self.eviction_threshold = eviction_threshold self.loaded_shards: OrderedDict[int, ShardInfo] = OrderedDict() self._lock = threading.Lock() print(f"MemoryMonitor initialized: max {max_ram_gb}GB, " f"evict at {eviction_threshold:.0%}") def get_system_ram_usage(self) -> float: """Get current system RAM usage as a fraction (0-1).""" return psutil.virtual_memory().percent / 100.0 def get_available_ram_mb(self) -> float: """Get available RAM in MB.""" return psutil.virtual_memory().available / (1024 * 1024) def get_loaded_shards_size_mb(self) -> float: """Total size of all loaded shards in MB.""" return sum(s.size_mb for s in self.loaded_shards.values()) def can_load(self, shard_size_mb: float) -> bool: """Check if we can load a shard of the given size.""" with self._lock: available = self.get_available_ram_mb() # Need 20% headroom for KV cache + OS needed = shard_size_mb + (self.max_ram_bytes / (1024*1024)) * 0.2 if available > needed: return True # Try evicting first self._evict_until_available(needed) return self.get_available_ram_mb() > shard_size_mb def on_load(self, shard_id: int, size_mb: float): """Called when a shard is loaded into RAM.""" with self._lock: self.loaded_shards[shard_id] = ShardInfo( shard_id=shard_id, size_mb=size_mb, loaded_at=time.time(), last_used=time.time(), access_count=1, ) # Move to end (most recently used) self.loaded_shards.move_to_end(shard_id) print(f" [RAM] Loaded shard {shard_id} ({size_mb:.0f}MB). " f"Total: {self.get_loaded_shards_size_mb():.0f}MB, " f"Available: {self.get_available_ram_mb():.0f}MB") def on_access(self, shard_id: int): """Called when a shard is accessed (keeps it in RAM).""" with self._lock: if shard_id in self.loaded_shards: self.loaded_shards[shard_id].last_used = time.time() self.loaded_shards[shard_id].access_count += 1 self.loaded_shards.move_to_end(shard_id) # LRU def on_unload(self, shard_id: int): """Called when a shard is unloaded from RAM.""" with self._lock: if shard_id in self.loaded_shards: info = self.loaded_shards.pop(shard_id) print(f" [RAM] Unloaded shard {shard_id} ({info.size_mb:.0f}MB). " f"Was used {info.access_count}x. " f"Total: {self.get_loaded_shards_size_mb():.0f}MB") def evict_if_needed(self): """Evict least-recently-used shards if RAM is tight.""" with self._lock: usage = self.get_system_ram_usage() if usage > self.eviction_threshold: self._evict_until_available( self.max_ram_bytes / (1024 * 1024) * (1 - self.eviction_threshold) ) def _evict_until_available(self, needed_mb: float): """Evict LRU shards until we have enough available RAM.""" while self.loaded_shards and self.get_available_ram_mb() < needed_mb: # Pop the least recently used (first item in OrderedDict) shard_id, info = self.loaded_shards.popitem(last=False) print(f" [RAM] Evicting shard {shard_id} (LRU, " f"last used {time.time() - info.last_used:.0f}s ago, " f"{info.size_mb:.0f}MB)") # In a real implementation, this would call llama.cpp to unload the shard # For now, we just track it — the mmap system handles the actual eviction def get_status(self) -> dict: """Get current memory status.""" with self._lock: return { "system_ram_usage": f"{self.get_system_ram_usage():.0%}", "available_ram_mb": f"{self.get_available_ram_mb():.0f}MB", "loaded_shards": len(self.loaded_shards), "loaded_shards_size": f"{self.get_loaded_shards_size_mb():.0f}MB", "max_ram_gb": f"{self.max_ram_bytes / (1024**3):.1f}GB", "shards": [ { "id": s.shard_id, "size_mb": f"{s.size_mb:.0f}MB", "accesses": s.access_count, "last_used_ago": f"{time.time() - s.last_used:.0f}s ago", } for s in self.loaded_shards.values() ], } # Test the memory monitor if __name__ == "__main__": monitor = MemoryMonitor(max_ram_gb=8.0) print("\n" + "=" * 70) print("MEMORY MONITOR TEST") print("=" * 70) # Simulate loading shards for i in range(6): monitor.on_load(shard_id=i, size_mb=500) monitor.on_access(shard_id=i) print(f"\nStatus: {monitor.get_status()}") # Simulate eviction print("\nSimulating RAM pressure...") for i in range(10): monitor.on_load(shard_id=10 + i, size_mb=500) print(f"\nFinal status: {monitor.get_status()}")