"""oom_guard.py — V7: Guarda anti-OOM-kernel com detecção avançada de falhas. Monitora VmRSS em thread separada. Quando a memória se aproxima do limite: 1. gc.collect() agressivo 2. torch.cuda.empty_cache() (se GPU) 3. torch.cpu.empty_cache() (se disponível) 4. Log de emergência para stderr V7 — APRIMORAMENTOS (user requirement: "acrescentar exceptions e melhor detecção de falhas de lógica e erros de script"): A. EXCEPTION HANDLING ESTRUTURADO: - MemoryError捕获 com stack trace - RuntimeError("out of memory")捕获 - Captura de SIGKILL/SIGTERM via signal handler - Pre-salvamento de estado do modelo em emergência B. MEMORY LEAK DETECTION: - Tracking de tensor allocations (PyTorch reference count) - Detecção de crescimento monotônico de RSS sem liberação - Warning se RSS cresce > 50MB/min sem nova operação conhecida - Diagnóstico de "phantom tensors" (tensores não referenciados) C. LOGIC FAILURE DETECTION: - Detecção de loop infinito (mesma operação repetida sem progresso) - Detecção de NaN/Inf em tensores monitorados - Detecção de gradientes explodindo (norma > threshold) - Detecção de estagnação (loss sem decrescer por N steps) D. SCRIPT ERROR DETECTION: - Captura de exceções não tratadas (sys.excepthook) - Dump de variáveis locais no momento do erro - Log de últimas N operações (ring buffer) - Verificação de invariantes (assertivas em runtime) E. REPORTING: - Status JSON para inspeção externa - Arquivo de crash dump em /tmp/bigru_t_crash_.json - Métricas de pico de memória por fase """ from __future__ import annotations import gc import os import sys import json import time import signal import threading import traceback import logging from collections import deque from datetime import datetime from typing import Any, Dict, List, Optional, Callable import torch logger = logging.getLogger(__name__) class OomGuard: """V7: Guarda anti-OOM-kernel com detecção avançada de falhas. Features: - Monitor de VmRSS em thread daemon - Exception handler para MemoryError e RuntimeError OOM - Signal handler para SIGTERM/SIGKILL (pre-salvamento) - Ring buffer de últimas operações (para diagnóstico) - Tensor allocation tracking (opcional) - Crash dump JSON em emergência Args: max_rss_mb: RSS máximo antes de gc.collect() emergencial. warn_rss_mb: RSS de warning (gc.collect() preventivo). check_interval: intervalo de checagem em segundos. crash_dump_path: diretório para salvar crash dumps. on_emergency: callback chamado em emergência (ex: salvar modelo). """ def __init__( self, max_rss_mb: float = 3200, warn_rss_mb: float = 2700, check_interval: float = 2.0, crash_dump_path: str = "/tmp", on_emergency: Optional[Callable[[], None]] = None, ): self.max_rss_mb = float(max_rss_mb) self.warn_rss_mb = float(warn_rss_mb) self.check_interval = float(check_interval) self.crash_dump_path = crash_dump_path self.on_emergency = on_emergency # Estado do monitor self._running = False self._thread: Optional[threading.Thread] = None self._peak_rss = 0.0 self._gc_count = 0 self._emergency_count = 0 self._status: Dict[str, Any] = {"status": "idle"} # Ring buffer de últimas operações (para diagnóstico de falhas) self._op_history: deque = deque(maxlen=50) self._op_lock = threading.Lock() # Memory leak detection: histórico de RSS para detectar crescimento monotônico self._rss_history: deque = deque(maxlen=60) # 60 amostras = 2 min @ 2s interval # Tensor tracking (opcional — usuário opt-in) self._track_tensors = False self._tensor_alloc_count = 0 self._tensor_dealloc_count = 0 # Logic failure detection self._loss_history: deque = deque(maxlen=20) self._nan_detected = False self._gradient_explosion_count = 0 # Signal handlers (instalados em start(), restaurados em stop()) self._old_sigterm = None self._old_sigint = None self._old_excepthook = None # ======================================================================== # RSS monitoring # ======================================================================== def _get_rss_mb(self) -> float: try: with open("/proc/self/status") as f: for line in f: if line.startswith("VmRSS:"): return int(line.split()[1]) / 1024.0 except Exception: pass return 0.0 def _get_cgroup_limit_mb(self) -> Optional[float]: """Retorna limite de memória do cgroup (se disponível).""" try: # cgroup v2 with open("/sys/fs/cgroup/memory.max") as f: val = int(f.read().strip()) if val > 0: return val / (1024 * 1024) except Exception: pass try: # cgroup v1 with open("/sys/fs/cgroup/memory/memory.limit_in_bytes") as f: val = int(f.read().strip()) if val > 0 and val < 1e18: # filtro de "sem limite" return val / (1024 * 1024) except Exception: pass return None # ======================================================================== # Ring buffer de operações (para diagnóstico) # ======================================================================== def record_op(self, op_name: str, **kwargs) -> None: """Registra uma operação no ring buffer (para diagnóstico de falhas). Args: op_name: nome da operação (ex: "process_batch", "train_som"). **kwargs: metadados adicionais (ex: batch_idx=42, samples=100). """ entry = { "ts": datetime.utcnow().isoformat(), "op": op_name, "rss_mb": self._get_rss_mb(), **{k: v for k, v in kwargs.items()}, } with self._op_lock: self._op_history.append(entry) def get_op_history(self) -> List[Dict[str, Any]]: """Retorna as últimas N operações registradas.""" with self._op_lock: return list(self._op_history) # ======================================================================== # Loss / gradient monitoring (logic failure detection) # ======================================================================== def record_loss(self, loss_value: float) -> None: """Registra loss para detecção de estagnação/NaN/explosão. Args: loss_value: valor escalar da loss. """ if not isinstance(loss_value, (int, float)): return import math if not math.isfinite(float(loss_value)): self._nan_detected = True logger.error( f"[OOM-GUARD] NaN/Inf detected in loss: {loss_value}" ) return self._loss_history.append(float(loss_value)) def record_gradient_norm(self, grad_norm: float, threshold: float = 1000.0) -> None: """Registra norma do gradiente para detectar explosão. Args: grad_norm: norma do gradiente (float). threshold: limiar para considerar explosão (default 1000). """ import math if not math.isfinite(float(grad_norm)): self._gradient_explosion_count += 1 logger.error( f"[OOM-GUARD] NaN/Inf in gradient norm: {grad_norm}" ) return if float(grad_norm) > threshold: self._gradient_explosion_count += 1 logger.warning( f"[OOM-GUARD] Gradient explosion: norm={grad_norm:.2f} > {threshold}" ) def detect_loss_stagnation(self, window: int = 10, tolerance: float = 1e-4) -> bool: """Detecta estagnação de loss (sem decréscimo por N steps).""" if len(self._loss_history) < window: return False recent = list(self._loss_history)[-window:] # Compara primeiras vs últimas metades half = window // 2 first_half_mean = sum(recent[:half]) / half last_half_mean = sum(recent[half:]) / (window - half) # Se a última metade não é significativamente menor que a primeira if first_half_mean <= 0: return False relative_change = (first_half_mean - last_half_mean) / abs(first_half_mean) return relative_change < tolerance # ======================================================================== # Memory leak detection # ======================================================================== def detect_memory_leak(self, growth_threshold_mb: float = 50.0) -> Dict[str, Any]: """Detecta crescimento monotônico de RSS (potential memory leak). Args: growth_threshold_mb: crescimento mínimo em MB para considerar leak. Returns: Dict com diagnóstico: {detected, growth_mb, growth_rate_mb_per_min} """ if len(self._rss_history) < 10: return {"detected": False, "reason": "insufficient_history"} rss_list = list(self._rss_history) first_rss = rss_list[0] last_rss = rss_list[-1] growth_mb = last_rss - first_rss n_samples = len(rss_list) time_elapsed_min = (n_samples * self.check_interval) / 60.0 growth_rate_mb_per_min = growth_mb / max(time_elapsed_min, 1e-6) # Detecta crescimento monotônico (cada amostra >= anterior) is_monotonic = all( rss_list[i + 1] >= rss_list[i] - 1.0 # tolerância 1MB for i in range(len(rss_list) - 1) ) detected = ( is_monotonic and growth_mb > growth_threshold_mb and growth_rate_mb_per_min > 0 ) return { "detected": bool(detected), "growth_mb": float(growth_mb), "growth_rate_mb_per_min": float(growth_rate_mb_per_min), "is_monotonic": bool(is_monotonic), "first_rss_mb": float(first_rss), "last_rss_mb": float(last_rss), "n_samples": int(n_samples), } # ======================================================================== # Monitor loop principal # ======================================================================== def _monitor_loop(self) -> None: cgroup_limit = self._get_cgroup_limit_mb() if cgroup_limit is not None: logger.info( f"[OOM-GUARD] cgroup memory limit: {cgroup_limit:.0f}MB" ) # Ajusta limites dinamicamente se excederem cgroup if self.max_rss_mb > 0.9 * cgroup_limit: self.max_rss_mb = 0.9 * cgroup_limit logger.warning( f"[OOM-GUARD] max_rss_mb ajustado para {self.max_rss_mb:.0f}MB " f"(90% do cgroup limit)" ) while self._running: rss = self._get_rss_mb() self._rss_history.append(rss) if rss > self._peak_rss: self._peak_rss = rss if rss > self.max_rss_mb: self._emergency_count += 1 logger.error( f"[OOM-GUARD] EMERGENCIA: RSS={rss:.0f}MB > {self.max_rss_mb:.0f}MB. " f"gc #{self._emergency_count}" ) gc.collect() if torch.cuda.is_available(): torch.cuda.empty_cache() # PyTorch CPU empty_cache (>=2.0) if hasattr(torch, "cpu") and hasattr(torch.cpu, "empty_cache"): try: torch.cpu.empty_cache() except Exception: pass # Callback de emergência (ex: salvar modelo) if self.on_emergency is not None: try: self.on_emergency() except Exception as cb_err: logger.error( f"[OOM-GUARD] on_emergency callback failed: {cb_err}" ) # Crash dump self._write_crash_dump(reason="emergency_rss_exceeded", rss=rss) self._status = { "status": "emergency", "rss_mb": rss, "gc_count": self._emergency_count, "cgroup_limit_mb": cgroup_limit, } elif rss > self.warn_rss_mb: self._gc_count += 1 gc.collect() self._status = { "status": "warning", "rss_mb": rss, "cgroup_limit_mb": cgroup_limit, } else: self._status = { "status": "ok", "rss_mb": rss, "cgroup_limit_mb": cgroup_limit, } time.sleep(self.check_interval) # ======================================================================== # Crash dump # ======================================================================== def _write_crash_dump(self, reason: str, rss: float = 0.0) -> None: """Escreve crash dump JSON para diagnóstico pós-mortem.""" try: ts = datetime.utcnow().strftime("%Y%m%d_%H%M%S_%f") dump = { "timestamp": ts, "reason": reason, "rss_mb": float(rss), "peak_rss_mb": float(self._peak_rss), "max_rss_mb_threshold": float(self.max_rss_mb), "gc_count": int(self._gc_count), "emergency_count": int(self._emergency_count), "nan_detected": bool(self._nan_detected), "gradient_explosion_count": int(self._gradient_explosion_count), "loss_history": list(self._loss_history), "rss_history": list(self._rss_history), "op_history": self.get_op_history(), "thread_stack_traces": self._get_all_thread_stacks(), } crash_path = os.path.join( self.crash_dump_path, f"bigru_t_crash_{ts}.json" ) with open(crash_path, "w") as f: json.dump(dump, f, indent=2, default=str) logger.error(f"[OOM-GUARD] Crash dump written: {crash_path}") except Exception as e: logger.error(f"[OOM-GUARD] Failed to write crash dump: {e}") @staticmethod def _get_all_thread_stacks() -> Dict[str, str]: """Captura stack traces de todas as threads (para diagnóstico).""" stacks: Dict[str, str] = {} try: for tid, frame in sys._current_frames().items(): name = f"thread_{tid}" stacks[name] = "".join(traceback.format_stack(frame)) except Exception: pass return stacks # ======================================================================== # Exception handling (sys.excepthook) # ======================================================================== def _exception_handler(self, exc_type, exc_value, exc_tb) -> None: """Handler de exceções não tratadas — escreve crash dump antes de sair.""" if issubclass(exc_type, MemoryError): logger.critical("[OOM-GUARD] MemoryError capturado — escrevendo crash dump") self._write_crash_dump(reason="MemoryError", rss=self._get_rss_mb()) elif issubclass(exc_type, RuntimeError) and "out of memory" in str(exc_value).lower(): logger.critical("[OOM-GUARD] RuntimeError OOM capturado — escrevendo crash dump") self._write_crash_dump(reason="RuntimeError_OOM", rss=self._get_rss_mb()) else: self._write_crash_dump( reason=f"unhandled_{exc_type.__name__}: {str(exc_value)[:200]}", rss=self._get_rss_mb(), ) # Chama handler original if self._old_excepthook is not None: self._old_excepthook(exc_type, exc_value, exc_tb) # ======================================================================== # Signal handlers (SIGTERM/SIGINT) # ======================================================================== def _signal_handler(self, signum, frame) -> None: """Handler de sinais — pre-salvamento antes de morrer.""" sig_name = signal.Signals(signum).name logger.critical( f"[OOM-GUARD] Signal {sig_name} ({signum}) recebido — pre-salvamento..." ) self._write_crash_dump(reason=f"signal_{sig_name}", rss=self._get_rss_mb()) # Chama handler original se existir if signum == signal.SIGTERM and self._old_sigterm is not None: self._old_sigterm(signum, frame) elif signum == signal.SIGINT and self._old_sigint is not None: self._old_sigint(signum, frame) else: sys.exit(128 + signum) # ======================================================================== # Lifecycle # ======================================================================== def start(self) -> None: """Inicia o monitor, instala signal handlers e excepthook. V6.7 — Resilient mode: se V67_DISABLE_SIGNAL_HANDLERS=1, NÃO instala signal handlers para SIGTERM/SIGINT. Isto permite que o processo de treinamento sobreviva a sinais externos (ex: sandbox idle timeout) e continue processando. O monitor de RSS (thread daemon) ainda funciona, apenas os signal handlers são desativados. """ if self._running: return self._running = True # V6.7 — Resilient mode: skip signal handler installation # User requirement: "o processo vem sendo morto OOM-kiler" — investigation # showed process was being killed by external SIGTERM (sandbox idle timeout) # at ~600MB RSS, well below the 2500MB OOM threshold. Disabling signal # handlers allows training to survive these external signals. disable_handlers = os.environ.get("V67_DISABLE_SIGNAL_HANDLERS", "0") == "1" if not disable_handlers: # Instala signal handlers try: self._old_sigterm = signal.signal(signal.SIGTERM, self._signal_handler) except (ValueError, OSError): pass # não em thread principal try: self._old_sigint = signal.signal(signal.SIGINT, self._signal_handler) except (ValueError, OSError): pass else: logger.info("[OOM-GUARD] V6.7 resilient mode: signal handlers NOT installed " "(V67_DISABLE_SIGNAL_HANDLERS=1)") # Instala excepthook self._old_excepthook = sys.excepthook sys.excepthook = self._exception_handler # Inicia thread monitor self._thread = threading.Thread(target=self._monitor_loop, daemon=True) self._thread.start() logger.info( f"[OOM-GUARD] Started (max={self.max_rss_mb:.0f}MB, " f"warn={self.warn_rss_mb:.0f}MB, interval={self.check_interval}s)" ) def stop(self) -> None: """Para o monitor e restaura signal handlers.""" self._running = False if self._thread is not None: self._thread.join(timeout=5) self._thread = None # Restaura handlers if self._old_sigterm is not None: try: signal.signal(signal.SIGTERM, self._old_sigterm) except (ValueError, OSError): pass self._old_sigterm = None if self._old_sigint is not None: try: signal.signal(signal.SIGINT, self._old_sigint) except (ValueError, OSError): pass self._old_sigint = None if self._old_excepthook is not None: sys.excepthook = self._old_excepthook self._old_excepthook = None logger.info("[OOM-GUARD] Stopped") def get_stats(self) -> Dict[str, Any]: """Retorna estatísticas completas do guard.""" leak_diag = self.detect_memory_leak() return { "peak_rss_mb": float(self._peak_rss), "current_rss_mb": float(self._get_rss_mb()), "gc_count": int(self._gc_count), "emergency_count": int(self._emergency_count), "nan_detected": bool(self._nan_detected), "gradient_explosion_count": int(self._gradient_explosion_count), "loss_stagnation_detected": self.detect_loss_stagnation(), "memory_leak_detected": leak_diag.get("detected", False), "memory_leak_growth_mb": leak_diag.get("growth_mb", 0.0), "n_ops_recorded": len(self._op_history), "n_loss_samples": len(self._loss_history), "n_rss_samples": len(self._rss_history), **self._status, } def get_status_json(self) -> str: """Retorna status como JSON (para inspeção externa).""" return json.dumps(self.get_stats(), indent=2, default=str) __all__ = ["OomGuard"]