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