src/bigru_t/utils/oom_guard.py

"""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_<timestamp>.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"]