File size: 6,946 Bytes
b72bf87 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 | """
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()}")
|