survival_watchdog.py

# -*- coding: utf-8 -*-
"""
survival_watchdog.py — HF Space Survival & Resource Monitor
━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━
Two threats to a Hugging Face Free Space:
  1. OOM crash  (RAM > 16 GB → container killed without warning)
  2. Idle sleep (HF puts a Space to sleep after ~48h of inactivity;
                 the wake-up latency is 30-60 seconds, breaking loops)

This module runs as a background asyncio task and:
  - Monitors psutil every 60 s
  - If RAM > RAM_KILL_THRESHOLD: kills the heaviest non-essential process
  - If CPU < CPU_IDLE_THRESHOLD for > IDLE_GRACE_MINUTES: fires a
    compute spike (1M-iteration sum) to reset HF's idle timer
  - Exposes live metrics as a JSON-serialisable dict for /api/metrics
"""

import os
import asyncio
import logging
import time
from typing import Dict, Any

import psutil

logger = logging.getLogger("survival_watchdog")

RAM_KILL_THRESHOLD  = float(os.environ.get("WD_RAM_KILL_PCT",   "88"))  # %
CPU_IDLE_THRESHOLD  = float(os.environ.get("WD_CPU_IDLE_PCT",   "4"))   # %
IDLE_GRACE_MINUTES  = int(os.environ.get("WD_IDLE_GRACE_MIN",   "8"))   # minutes
CHECK_INTERVAL_SECS = int(os.environ.get("WD_CHECK_SECS",       "60"))  # seconds

# Processes that should NEVER be killed (guarded by name prefix)
PROTECTED_PROCESS_NAMES = {
    "python", "uvicorn", "node", "npm", "git", "bash", "sh",
}

_metrics_snapshot: Dict[str, Any] = {}


def get_metrics() -> Dict[str, Any]:
    """Return the latest resource snapshot (called by /api/metrics endpoint)."""
    return _metrics_snapshot


class SurvivalWatchdog:
    def __init__(self, own_pid: int = None):
        self.own_pid = own_pid or os.getpid()
        self._idle_ticks = 0  # consecutive checks where CPU < threshold

    # ── Core Monitor Loop ─────────────────────────────────────────────────────

    async def run(self):
        """Main loop — runs forever as an asyncio background task."""
        logger.info("[Watchdog] Started. RAM kill threshold: %.0f%% | CPU idle threshold: %.0f%%",
                    RAM_KILL_THRESHOLD, CPU_IDLE_THRESHOLD)
        while True:
            try:
                await self._check()
            except Exception as e:
                logger.error(f"[Watchdog] Error in check loop: {e}")
            await asyncio.sleep(CHECK_INTERVAL_SECS)

    async def _check(self):
        global _metrics_snapshot

        # ── Collect ───────────────────────────────────────────────────────────
        vm  = psutil.virtual_memory()
        cpu = psutil.cpu_percent(interval=1)
        disk = psutil.disk_usage("/tmp")
        net  = psutil.net_io_counters()

        ram_used_gb  = vm.used / (1024 ** 3)
        ram_total_gb = vm.total / (1024 ** 3)
        ram_pct      = vm.percent

        _metrics_snapshot = {
            "cpu_percent":      round(cpu, 1),
            "ram_used_gb":      round(ram_used_gb, 2),
            "ram_total_gb":     round(ram_total_gb, 2),
            "ram_percent":      round(ram_pct, 1),
            "ram_free_gb":      round((vm.total - vm.used) / (1024 ** 3), 2),
            "disk_used_gb":     round(disk.used / (1024 ** 3), 2),
            "disk_free_gb":     round(disk.free / (1024 ** 3), 2),
            "net_sent_mb":      round(net.bytes_sent / (1024 ** 2), 1),
            "net_recv_mb":      round(net.bytes_recv / (1024 ** 2), 1),
            "timestamp":        time.time(),
            "idle_ticks":       self._idle_ticks,
            "status":           "ok",
        }

        # ── OOM Defence ───────────────────────────────────────────────────────
        if ram_pct >= RAM_KILL_THRESHOLD:
            logger.warning(
                "[Watchdog] ⚠ RAM at %.1f%% (%.2f/%.2f GB) — initiating OOM defence.",
                ram_pct, ram_used_gb, ram_total_gb
            )
            killed = self._kill_heaviest_safe()
            _metrics_snapshot["oom_kill"] = killed
            _metrics_snapshot["status"] = "oom_defence"

        # ── Idle Sleep Defence ────────────────────────────────────────────────
        idle_grace_ticks = (IDLE_GRACE_MINUTES * 60) // CHECK_INTERVAL_SECS

        if cpu < CPU_IDLE_THRESHOLD:
            self._idle_ticks += 1
        else:
            self._idle_ticks = 0

        if self._idle_ticks >= idle_grace_ticks:
            logger.info(
                "[Watchdog] Space has been idle for ~%d min — firing CPU wake-up spike.",
                IDLE_GRACE_MINUTES
            )
            await self._cpu_spike()
            self._idle_ticks = 0
            _metrics_snapshot["status"] = "wake_spike_fired"

        logger.debug(
            "[Watchdog] CPU=%.1f%% | RAM=%.1f%% (%.2fGB free) | Idle ticks=%d",
            cpu, ram_pct, _metrics_snapshot["ram_free_gb"], self._idle_ticks
        )

    # ── OOM Kill ─────────────────────────────────────────────────────────────

    def _kill_heaviest_safe(self) -> str:
        """
        Finds the non-protected process using the most RAM and kills it.
        Returns a string description of what was killed (or 'none').
        """
        candidates = []
        for proc in psutil.process_iter(["pid", "name", "memory_percent"]):
            try:
                info = proc.info
                pid  = info["pid"]
                name = (info["name"] or "").lower()
                mem  = info["memory_percent"] or 0.0

                if pid == self.own_pid:
                    continue
                if any(name.startswith(p) for p in PROTECTED_PROCESS_NAMES):
                    continue
                candidates.append((mem, pid, name))
            except (psutil.NoSuchProcess, psutil.AccessDenied):
                pass

        if not candidates:
            logger.warning("[Watchdog] No safe kill candidates found — RAM is used by protected processes.")
            return "none"

        candidates.sort(reverse=True)
        mem_pct, pid, name = candidates[0]
        try:
            psutil.Process(pid).kill()
            logger.warning("[Watchdog] Killed '%s' (PID %d, %.1f%% RAM) for OOM defence.", name, pid, mem_pct)
            return f"{name}:{pid}"
        except Exception as e:
            logger.error(f"[Watchdog] Failed to kill PID {pid}: {e}")
            return "kill_failed"

    # ── CPU Spike (Idle Prevention) ───────────────────────────────────────────

    async def _cpu_spike(self):
        """
        Runs a CPU-bound task in an executor so it doesn't block the event loop.
        Uses a 1-million-iteration sum — takes ~50 ms on 2 vCPUs.
        Just enough to reset HF's idle detector without burning quota.
        """
        loop = asyncio.get_event_loop()
        await loop.run_in_executor(None, lambda: sum(i * i for i in range(1_000_000)))
        logger.debug("[Watchdog] CPU spike complete.")