HAKO upload: hako/memory_manager.py

hako/memory_manager.py

@@ -0,0 +1,121 @@
+"""Aggressive memory & storage management for HAKO.
+
+Proposition P-MEM (bounded-resource operation). If every phase satisfies
+    (i)  RSS(t) <= ram_soft_cap for all t (enforced by watchdog degrade),
+    (ii) disk blobs not referenced by decomposed artifacts are deleted,
+    (iii) telemetry/checkpoints are size-capped and compressed,
+then the pipeline operates indefinitely within a fixed resource envelope.
+Proof sketch: (i) bounds the state space of the resident set by construction
+(degrade-or-abort policy with monotone batch reduction); (ii) makes total disk
+usage a sum of compressed artifacts (monotone, bounded by artifact budget);
+(iii) bounds every log by its cap. Induction over phases completes the bound.
+"""
+from __future__ import annotations
+
+import gc
+import logging
+import os
+import shutil
+import time
+from pathlib import Path
+from typing import Iterable
+
+import psutil
+
+log = logging.getLogger("hako.mem")
+
+
+class MemoryManager:
+    def __init__(self, cfg) -> None:
+        self.cfg = cfg
+        vm = psutil.virtual_memory()
+        self.ram_soft_cap = int(cfg.ram_soft_cap_ratio * vm.total)
+        self._last_report = 0.0
+
+    # ------------------------------------------------------------------ ram
+    def rss(self) -> int:
+        return psutil.Process(os.getpid()).memory_info().rss
+
+    def pressure(self) -> float:
+        """RSS pressure in [0, 1+] relative to the soft cap."""
+        return self.rss() / max(1, self.ram_soft_cap)
+
+    def check(self, degrade_hooks: Iterable | None = None) -> str:
+        """Watchdog: returns 'ok' | 'degraded'. Calls hooks to shrink state."""
+        p = self.pressure()
+        if p < 0.85:
+            return "ok"
+        gc.collect()
+        if p < 1.0:
+            return "ok"
+        action = "degraded"
+        for hook in degrade_hooks or []:
+            try:
+                hook()
+            except Exception as exc:  # pragma: no cover - defensive
+                log.warning("degrade hook failed: %s", exc)
+        gc.collect()
+        if self.pressure() > 1.05:
+            raise MemoryError(
+                f"RSS {self.rss()/2**30:.2f}GiB exceeded soft cap "
+                f"{self.ram_soft_cap/2**30:.2f}GiB after degradation")
+        return action
+
+    def report(self, tag: str = "") -> None:
+        now = time.time()
+        if now - self._last_report < 15:
+            return
+        self._last_report = now
+        vm = psutil.virtual_memory()
+        log.info("mem[%s] rss=%.2fGiB cap=%.2fGiB sys_avail=%.2fGiB",
+                 tag, self.rss() / 2**30, self.ram_soft_cap / 2**30,
+                 vm.available / 2**30)
+
+    # ----------------------------------------------------------------- disk
+    def disk_free(self) -> int:
+        return psutil.disk_usage("/").free
+
+    def ensure_disk(self, need_bytes: int) -> None:
+        if self.disk_free() >= max(need_bytes, self.cfg.disk_min_free_bytes):
+            return
+        freed = self.sweep_caches()
+        log.info("disk sweep freed %.2f MiB", freed / 2**20)
+        if self.disk_free() < max(need_bytes, self.cfg.disk_min_free_bytes):
+            raise OSError("disk exhausted even after cache sweep")
+
+    def sweep_caches(self) -> int:
+        """Delete HF hub blobs and work leftovers that have decomposed twins."""
+        freed = 0
+        hf_blob_dirs = [
+            Path(self.cfg.hf_cache) / "hub",
+        ]
+        for base in hf_blob_dirs:
+            if not base.exists():
+                continue
+            for model_dir in base.iterdir():
+                marker = Path(self.cfg.artifacts_dir) / (
+                    model_dir.name.replace("models--", "") + ".decomposed.json")
+                if marker.exists():
+                    # decomposed copy exists -> raw weights are expendable
+                    size = sum(f.stat().st_size for f in model_dir.rglob("*")
+                               if f.is_file())
+                    shutil.rmtree(model_dir, ignore_errors=True)
+                    freed += size
+        gc.collect()
+        return freed
+
+    @staticmethod
+    def cap_file_size(path: Path, max_bytes: int, trim_head: bool = True) -> None:
+        """Keep a log under max_bytes by trimming the oldest half (JSONL-safe)."""
+        p = Path(path)
+        if not p.exists() or p.stat().st_size <= max_bytes:
+            return
+        lines = p.read_text(encoding="utf-8", errors="ignore").splitlines(keepends=True)
+        keep = lines[len(lines) // 2:] if trim_head else lines[:len(lines) // 2]
+        tmp = p.with_suffix(p.suffix + ".tmp")
+        tmp.write_text("".join(keep), encoding="utf-8")
+        tmp.replace(p)
+
+    @staticmethod
+    def rm_tree_quiet(path: Path) -> None:
+        shutil.rmtree(path, ignore_errors=True)