Spaces:
Sleeping
Sleeping
File size: 7,564 Bytes
12ab90a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 | # -*- 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.")
|