memory_monitor.py

"""
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()}")