augment17/claude-code-backend
0
1# -*- coding: utf-8 -*-2"""3survival_watchdog.py — HF Space Survival & Resource Monitor4━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━5Two threats to a Hugging Face Free Space:6 1. OOM crash (RAM > 16 GB → container killed without warning)7 2. Idle sleep (HF puts a Space to sleep after ~48h of inactivity;8 the wake-up latency is 30-60 seconds, breaking loops)9 10This module runs as a background asyncio task and:11 - Monitors psutil every 60 s12 - If RAM > RAM_KILL_THRESHOLD: kills the heaviest non-essential process13 - If CPU < CPU_IDLE_THRESHOLD for > IDLE_GRACE_MINUTES: fires a14 compute spike (1M-iteration sum) to reset HF's idle timer15 - Exposes live metrics as a JSON-serialisable dict for /api/metrics16"""17 18import os19import asyncio20import logging21import time22from typing import Dict, Any23 24import psutil25 26logger = logging.getLogger("survival_watchdog")27 28RAM_KILL_THRESHOLD = float(os.environ.get("WD_RAM_KILL_PCT", "88")) # %29CPU_IDLE_THRESHOLD = float(os.environ.get("WD_CPU_IDLE_PCT", "4")) # %30IDLE_GRACE_MINUTES = int(os.environ.get("WD_IDLE_GRACE_MIN", "8")) # minutes31CHECK_INTERVAL_SECS = int(os.environ.get("WD_CHECK_SECS", "60")) # seconds32 33# Processes that should NEVER be killed (guarded by name prefix)34PROTECTED_PROCESS_NAMES = {35 "python", "uvicorn", "node", "npm", "git", "bash", "sh",36}37 38_metrics_snapshot: Dict[str, Any] = {}39 40 41def get_metrics() -> Dict[str, Any]:42 """Return the latest resource snapshot (called by /api/metrics endpoint)."""43 return _metrics_snapshot44 45 46class SurvivalWatchdog:47 def __init__(self, own_pid: int = None):48 self.own_pid = own_pid or os.getpid()49 self._idle_ticks = 0 # consecutive checks where CPU < threshold50 51 # ── Core Monitor Loop ─────────────────────────────────────────────────────52 53 async def run(self):54 """Main loop — runs forever as an asyncio background task."""55 logger.info("[Watchdog] Started. RAM kill threshold: %.0f%% | CPU idle threshold: %.0f%%",56 RAM_KILL_THRESHOLD, CPU_IDLE_THRESHOLD)57 while True:58 try:59 await self._check()60 except Exception as e:61 logger.error(f"[Watchdog] Error in check loop: {e}")62 await asyncio.sleep(CHECK_INTERVAL_SECS)63 64 async def _check(self):65 global _metrics_snapshot66 67 # ── Collect ───────────────────────────────────────────────────────────68 vm = psutil.virtual_memory()69 cpu = psutil.cpu_percent(interval=1)70 disk = psutil.disk_usage("/tmp")71 net = psutil.net_io_counters()72 73 ram_used_gb = vm.used / (1024 ** 3)74 ram_total_gb = vm.total / (1024 ** 3)75 ram_pct = vm.percent76 77 _metrics_snapshot = {78 "cpu_percent": round(cpu, 1),79 "ram_used_gb": round(ram_used_gb, 2),80 "ram_total_gb": round(ram_total_gb, 2),81 "ram_percent": round(ram_pct, 1),82 "ram_free_gb": round((vm.total - vm.used) / (1024 ** 3), 2),83 "disk_used_gb": round(disk.used / (1024 ** 3), 2),84 "disk_free_gb": round(disk.free / (1024 ** 3), 2),85 "net_sent_mb": round(net.bytes_sent / (1024 ** 2), 1),86 "net_recv_mb": round(net.bytes_recv / (1024 ** 2), 1),87 "timestamp": time.time(),88 "idle_ticks": self._idle_ticks,89 "status": "ok",90 }91 92 # ── OOM Defence ───────────────────────────────────────────────────────93 if ram_pct >= RAM_KILL_THRESHOLD:94 logger.warning(95 "[Watchdog] ⚠ RAM at %.1f%% (%.2f/%.2f GB) — initiating OOM defence.",96 ram_pct, ram_used_gb, ram_total_gb97 )98 killed = self._kill_heaviest_safe()99 _metrics_snapshot["oom_kill"] = killed100 _metrics_snapshot["status"] = "oom_defence"101 102 # ── Idle Sleep Defence ────────────────────────────────────────────────103 idle_grace_ticks = (IDLE_GRACE_MINUTES * 60) // CHECK_INTERVAL_SECS104 105 if cpu < CPU_IDLE_THRESHOLD:106 self._idle_ticks += 1107 else:108 self._idle_ticks = 0109 110 if self._idle_ticks >= idle_grace_ticks:111 logger.info(112 "[Watchdog] Space has been idle for ~%d min — firing CPU wake-up spike.",113 IDLE_GRACE_MINUTES114 )115 await self._cpu_spike()116 self._idle_ticks = 0117 _metrics_snapshot["status"] = "wake_spike_fired"118 119 logger.debug(120 "[Watchdog] CPU=%.1f%% | RAM=%.1f%% (%.2fGB free) | Idle ticks=%d",121 cpu, ram_pct, _metrics_snapshot["ram_free_gb"], self._idle_ticks122 )123 124 # ── OOM Kill ─────────────────────────────────────────────────────────────125 126 def _kill_heaviest_safe(self) -> str:127 """128 Finds the non-protected process using the most RAM and kills it.129 Returns a string description of what was killed (or 'none').130 """131 candidates = []132 for proc in psutil.process_iter(["pid", "name", "memory_percent"]):133 try:134 info = proc.info135 pid = info["pid"]136 name = (info["name"] or "").lower()137 mem = info["memory_percent"] or 0.0138 139 if pid == self.own_pid:140 continue141 if any(name.startswith(p) for p in PROTECTED_PROCESS_NAMES):142 continue143 candidates.append((mem, pid, name))144 except (psutil.NoSuchProcess, psutil.AccessDenied):145 pass146 147 if not candidates:148 logger.warning("[Watchdog] No safe kill candidates found — RAM is used by protected processes.")149 return "none"150 151 candidates.sort(reverse=True)152 mem_pct, pid, name = candidates[0]153 try:154 psutil.Process(pid).kill()155 logger.warning("[Watchdog] Killed '%s' (PID %d, %.1f%% RAM) for OOM defence.", name, pid, mem_pct)156 return f"{name}:{pid}"157 except Exception as e:158 logger.error(f"[Watchdog] Failed to kill PID {pid}: {e}")159 return "kill_failed"160 161 # ── CPU Spike (Idle Prevention) ───────────────────────────────────────────162 163 async def _cpu_spike(self):164 """165 Runs a CPU-bound task in an executor so it doesn't block the event loop.166 Uses a 1-million-iteration sum — takes ~50 ms on 2 vCPUs.167 Just enough to reset HF's idle detector without burning quota.168 """169 loop = asyncio.get_event_loop()170 await loop.run_in_executor(None, lambda: sum(i * i for i in range(1_000_000)))171 logger.debug("[Watchdog] CPU spike complete.")172