src/bigru_t/utils/memory_monitor.py

"""
memory_monitor — Detecção avançada de uso de memória.

Implementa (item 1):
- RSS tracking (CPU)
- GPU memory tracking (allocated/reserved/peak)
- Leak detection (crescimento monotônico)
-OOM prevention (pre-allocation check)
- Memory budgeting (per-module)
- Snapshot diffing (antes/depois de operação)
- Alerting (warning/critical thresholds)

Detecção de erros:
- MemoryLeakError: crescimento monotônico em N amostras consecutivas
- MemoryBudgetExceededError: módulo excede quota
- OOMPredictedError: projeção indica OOM em K passos
- SnapshotMismatchError: snapshot diff inesperado
"""

from __future__ import annotations

import gc
import os
import time
import threading
import warnings
from collections import defaultdict, deque
from contextlib import contextmanager
from dataclasses import dataclass, field
from typing import Any, Callable, Dict, List, Optional, Tuple

import torch
import torch.nn as nn


# ============================================================================
# Exceções específicas
# ============================================================================

class MemoryLeakError(RuntimeError):
    """Crescimento monotônico de memória detectado (possível leak)."""


class MemoryBudgetExceededError(RuntimeError):
    """Módulo excedeu sua quota de memória."""


class OOMPredictedError(RuntimeError):
    """Projeção indica OOM em K passos."""


class SnapshotMismatchError(RuntimeError):
    """Snapshot diff inesperado (alocação não liberada)."""


# ============================================================================
# 1. MemorySnapshot — captura estado de memória em um instante
# ============================================================================

@dataclass
class MemorySnapshot:
    """Snapshot completo de uso de memória."""
    timestamp: float
    rss_bytes: int = 0           # CPU RSS (via psutil)
    cpu_percent: float = 0.0     # % do total de RAM
    gpu_allocated: int = 0       # GPU allocated (torch)
    gpu_reserved: int = 0        # GPU reserved (torch)
    gpu_peak: int = 0            # GPU peak (torch)
    gpu_total: int = 0           # GPU total (device property)
    label: str = ""

    def diff(self, other: "MemorySnapshot") -> "MemorySnapshotDiff":
        """Calcula diff (self - other): positivo = crescimento."""
        return MemorySnapshotDiff(
            rss_delta=self.rss_bytes - other.rss_bytes,
            gpu_allocated_delta=self.gpu_allocated - other.gpu_allocated,
            gpu_reserved_delta=self.gpu_reserved - other.gpu_reserved,
            time_delta=self.timestamp - other.timestamp,
            from_label=other.label,
            to_label=self.label,
        )


@dataclass
class MemorySnapshotDiff:
    """Diferença entre dois MemorySnapshots."""
    rss_delta: int
    gpu_allocated_delta: int
    gpu_reserved_delta: int
    time_delta: float
    from_label: str
    to_label: str

    def __repr__(self) -> str:
        return (
            f"MemoryDiff({self.from_label}→{self.to_label}: "
            f"rss={self.rss_delta:+d}B, "
            f"gpu_alloc={self.gpu_allocated_delta:+d}B, "
            f"dt={self.time_delta:.3f}s)"
        )


# ============================================================================
# 2. MemoryMonitor — monitora memória em runtime
# ============================================================================

class MemoryMonitor:
    """Monitor de uso de memória com leak detection.

    Recursos:
      * Snapshot a qualquer momento (capture_snapshot)
      * Diff entre snapshots (snapshot.diff)
      * Histórico circular (history_size amostras)
      * Leak detection: se rss cresce em N amostras consecutivas → alerta
      * OOM prediction: projeta tendência linear, previne OOM
      * Budgeting: cada módulo tem quota; exceder → MemoryBudgetExceededError
      * Auto-cleanup: gc.collect() + torch.cuda.empty_cache() quando crítico
    """

    def __init__(
        self,
        history_size: int = 100,
        leak_threshold_consecutive: int = 10,
        leak_threshold_bytes: int = 1024 * 1024,  # 1 MB
        warning_pct: float = 80.0,
        critical_pct: float = 95.0,
        auto_cleanup: bool = True,
    ):
        self.history_size = history_size
        self.leak_threshold_consecutive = leak_threshold_consecutive
        self.leak_threshold_bytes = leak_threshold_bytes
        self.warning_pct = warning_pct
        self.critical_pct = critical_pct
        self.auto_cleanup = auto_cleanup

        self._history: deque = deque(maxlen=history_size)
        self._budgets: Dict[str, int] = {}  # module_name -> max_bytes
        self._lock = threading.Lock()

        # Tenta importar psutil (opcional).
        try:
            import psutil
            self._psutil = psutil
        except ImportError:
            self._psutil = None

    def capture_snapshot(self, label: str = "") -> MemorySnapshot:
        """Captura snapshot atual de memória."""
        snap = MemorySnapshot(timestamp=time.time(), label=label)
        # CPU.
        if self._psutil is not None:
            try:
                process = self._psutil.Process()
                mem = process.memory_info()
                snap.rss_bytes = mem.rss
                vm = self._psutil.virtual_memory()
                snap.cpu_percent = snap.rss_bytes / vm.total * 100.0 if vm.total > 0 else 0.0
            except Exception:
                pass
        # GPU.
        if torch.cuda.is_available():
            try:
                snap.gpu_allocated = torch.cuda.memory_allocated()
                snap.gpu_reserved = torch.cuda.memory_reserved()
                snap.gpu_peak = torch.cuda.max_memory_allocated()
                snap.gpu_total = torch.cuda.get_device_properties(0).total_memory
            except Exception:
                pass
        with self._lock:
            self._history.append(snap)
        return snap

    def get_history(self) -> List[MemorySnapshot]:
        with self._lock:
            return list(self._history)

    def check_leak(self) -> Optional[MemoryLeakError]:
        """Verifica leak: crescimento monotônico em N amostras consecutivas.

        Returns: None se OK, MemoryLeakError se leak detectado.
        """
        with self._lock:
            history = list(self._history)
        if len(history) < self.leak_threshold_consecutive + 1:
            return None
        # Pega últimas N amostras.
        recent = history[-(self.leak_threshold_consecutive + 1):]
        # Verifica crescimento monotônico em rss ou gpu_allocated.
        rss_growing = all(
            recent[i+1].rss_bytes - recent[i].rss_bytes > self.leak_threshold_bytes
            for i in range(len(recent) - 1)
        ) and recent[-1].rss_bytes > recent[0].rss_bytes
        gpu_growing = all(
            recent[i+1].gpu_allocated - recent[i].gpu_allocated > self.leak_threshold_bytes
            for i in range(len(recent) - 1)
        ) and recent[-1].gpu_allocated > recent[0].gpu_allocated
        if rss_growing:
            total_growth = recent[-1].rss_bytes - recent[0].rss_bytes
            return MemoryLeakError(
                f"Possível leak de RSS: crescimento monotônico de "
                f"{total_growth / 1e6:.1f} MB em {len(recent)} amostras"
            )
        if gpu_growing:
            total_growth = recent[-1].gpu_allocated - recent[0].gpu_allocated
            return MemoryLeakError(
                f"Possível leak de GPU: crescimento monotônico de "
                f"{total_growth / 1e6:.1f} MB em {len(recent)} amostras"
            )
        return None

    def predict_oom(self, k_steps: int = 10) -> Optional[OOMPredictedError]:
        """Projeta tendência linear; se projeta OOM em k_steps, alerta.

        Usa regressão linear simples nos últimos min(20, len) snapshots.
        """
        with self._lock:
            history = list(self._history)
        if len(history) < 5:
            return None
        recent = history[-min(20, len(history)):]
        n = len(recent)
        # Regressão linear: y = a + b*t, onde y = rss_bytes, t = índice.
        ts = list(range(n))
        ys_rss = [s.rss_bytes for s in recent]
        ys_gpu = [s.gpu_allocated for s in recent]
        # Calcula b (slope) via least squares.
        mean_t = sum(ts) / n
        mean_y_rss = sum(ys_rss) / n
        mean_y_gpu = sum(ys_gpu) / n
        num_rss = sum((t - mean_t) * (y - mean_y_rss) for t, y in zip(ts, ys_rss))
        num_gpu = sum((t - mean_t) * (y - mean_y_gpu) for t, y in zip(ts, ys_gpu))
        den = sum((t - mean_t) ** 2 for t in ts) or 1
        slope_rss = num_rss / den
        slope_gpu = num_gpu / den
        # Projeta.
        proj_rss = ys_rss[-1] + slope_rss * k_steps
        proj_gpu = ys_gpu[-1] + slope_gpu * k_steps
        # Limites.
        total_ram = 0
        if self._psutil is not None:
            try:
                total_ram = self._psutil.virtual_memory().total
            except Exception:
                pass
        total_gpu = recent[-1].gpu_total if recent[-1].gpu_total > 0 else 0
        warnings_list = []
        if total_ram > 0 and proj_rss > self.critical_pct / 100.0 * total_ram:
            warnings_list.append(
                f"RSS projetado {proj_rss / 1e9:.2f} GB em {k_steps} passos "
                f"(>{self.critical_pct}% de {total_ram / 1e9:.2f} GB)"
            )
        if total_gpu > 0 and proj_gpu > self.critical_pct / 100.0 * total_gpu:
            warnings_list.append(
                f"GPU projetado {proj_gpu / 1e9:.2f} GB em {k_steps} passos "
                f"(>{self.critical_pct}% de {total_gpu / 1e9:.2f} GB)"
            )
        if warnings_list:
            return OOMPredictedError("; ".join(warnings_list))
        return None

    def set_budget(self, module_name: str, max_bytes: int) -> None:
        """Define quota de memória para um módulo."""
        self._budgets[module_name] = max_bytes

    def check_budget(self, module_name: str, current_bytes: int) -> Optional[MemoryBudgetExceededError]:
        """Verifica se módulo excedeu quota."""
        if module_name not in self._budgets:
            return None
        max_b = self._budgets[module_name]
        if current_bytes > max_b:
            return MemoryBudgetExceededError(
                f"Módulo '{module_name}' excedeu quota: "
                f"{current_bytes / 1e6:.1f} MB > {max_b / 1e6:.1f} MB"
            )
        return None

    def check_thresholds(self) -> Optional[str]:
        """Verifica se uso atual excede warning/critical thresholds."""
        snap = self.capture_snapshot()
        if snap.cpu_percent > self.critical_pct:
            if self.auto_cleanup:
                self.cleanup()
            return f"CRITICAL: CPU RSS {snap.cpu_percent:.1f}%"
        if snap.gpu_total > 0:
            gpu_pct = snap.gpu_allocated / snap.gpu_total * 100
            if gpu_pct > self.critical_pct:
                if self.auto_cleanup:
                    self.cleanup()
                return f"CRITICAL: GPU {gpu_pct:.1f}%"
            if gpu_pct > self.warning_pct:
                return f"WARNING: GPU {gpu_pct:.1f}%"
        if snap.cpu_percent > self.warning_pct:
            return f"WARNING: CPU RSS {snap.cpu_percent:.1f}%"
        return None

    def cleanup(self) -> None:
        """Limpa caches (gc + CUDA empty_cache)."""
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()

    def summary(self) -> str:
        """Resumo do estado atual."""
        snap = self.capture_snapshot()
        lines = [
            f"MemoryMonitor:",
            f"  RSS: {snap.rss_bytes / 1e6:.1f} MB ({snap.cpu_percent:.1f}%)",
        ]
        if snap.gpu_total > 0:
            gpu_pct = snap.gpu_allocated / snap.gpu_total * 100
            lines.append(f"  GPU allocated: {snap.gpu_allocated / 1e6:.1f} MB ({gpu_pct:.1f}%)")
            lines.append(f"  GPU reserved:  {snap.gpu_reserved / 1e6:.1f} MB")
            lines.append(f"  GPU peak:      {snap.gpu_peak / 1e6:.1f} MB")
        lines.append(f"  History: {len(self._history)} samples")
        lines.append(f"  Budgets: {len(self._budgets)} modules")
        return "\n".join(lines)


# ============================================================================
# 3. Context manager para snapshot diff
# ============================================================================

@contextmanager
def memory_diff(monitor: MemoryMonitor, label: str = "operation"):
    """Context manager que captura diff de memória antes/depois.

    Uso:
        >>> with memory_diff(monitor, "forward_pass") as diff:
        ...     y = model(X)
        >>> print(diff)  # MemorySnapshotDiff
    """
    before = monitor.capture_snapshot(label=f"{label}_before")
    diff_holder: Dict[str, MemorySnapshotDiff] = {}
    try:
        yield diff_holder
    finally:
        after = monitor.capture_snapshot(label=f"{label}_after")
        diff_holder["diff"] = after.diff(before)


# ============================================================================
# 4. ModuleMemoryProfiler — profile de memória por módulo
# ============================================================================

class ModuleMemoryProfiler:
    """Profile de memória por módulo nn.Module.

    Registra pico de memória durante forward de cada módulo.
    Útil para identificar módulos que consomem muita memória.
    """

    def __init__(self, monitor: Optional[MemoryMonitor] = None):
        self.monitor = monitor or MemoryMonitor()
        self.profiles: Dict[str, MemorySnapshotDiff] = {}

    def profile_forward(
        self,
        module: nn.Module,
        X: torch.Tensor,
        module_name: str = "",
    ) -> MemorySnapshotDiff:
        """Executa forward e captura diff de memória."""
        name = module_name or module.__class__.__name__
        before = self.monitor.capture_snapshot(label=f"{name}_before")
        try:
            with torch.no_grad():
                _ = module(X)
        except Exception as e:
            warnings.warn(f"Forward falhou para {name}: {e}")
        after = self.monitor.capture_snapshot(label=f"{name}_after")
        diff = after.diff(before)
        self.profiles[name] = diff
        return diff

    def get_profile(self, module_name: str) -> Optional[MemorySnapshotDiff]:
        return self.profiles.get(module_name)

    def summary(self) -> str:
        if not self.profiles:
            return "ModuleMemoryProfiler: sem profiles"
        lines = ["ModuleMemoryProfiler:"]
        # Ordena por RSS delta decrescente.
        sorted_profiles = sorted(
            self.profiles.items(),
            key=lambda x: x[1].rss_delta,
            reverse=True,
        )
        for name, diff in sorted_profiles:
            lines.append(f"  {name}: {diff}")
        return "\n".join(lines)