Spaces:
Sleeping
Sleeping
| """Resource-aware local model manager (V6 phase_03). | |
| LOAD/UNLOAD/RELOAD/STATUS + memory pressure handling. The world keeps | |
| running while heavy models load/unload: expensive inference is synchronous | |
| per call, never on the physics thread. If a model does not fit, optional | |
| models unload first (priority order), then resolution/settings degrade, | |
| then the call fails explicitly — the simulation NEVER crashes for an | |
| optional model. | |
| """ | |
| import os | |
| import threading | |
| import time | |
| from typing import Any, Callable, Dict, Optional | |
| PRIORITY = ["EMBEDDING_MODEL", "TEXT_MODEL", "TTS_MODEL", "STT_MODEL", | |
| "VISION_MODEL", "IMAGE_MODEL"] # last = first to shed | |
| def _ram() -> Dict[str, float]: | |
| try: | |
| import psutil | |
| vm = psutil.virtual_memory() | |
| return {"percent": float(vm.percent), | |
| "avail_mb": float(vm.available / 1024 ** 2)} | |
| except Exception: | |
| return {"percent": -1.0, "avail_mb": -1.0} | |
| class ModelManager: | |
| def __init__(self, ram_limit_percent: float = 88.0, | |
| ram_reader=_ram): | |
| self.lock = threading.RLock() | |
| self.ram_limit = float(ram_limit_percent) | |
| self._ram_reader = ram_reader | |
| self.loaded: Dict[str, Dict[str, Any]] = {} # task -> {handle, loaded_at, uses} | |
| self._loaders: Dict[str, Callable[[], Any]] = {} | |
| self.events = [] | |
| def log(self, kind: str, detail: str) -> None: | |
| self.events.append({"ts": time.time(), "kind": kind, "detail": detail}) | |
| def register_loader(self, task: str, fn: Callable[[], Any]) -> None: | |
| with self.lock: | |
| self._loaders[task] = fn | |
| def status(self) -> Dict[str, Any]: | |
| with self.lock: | |
| return {"ram": self._ram_reader(), "ram_limit_percent": self.ram_limit, | |
| "loaded": {t: {"uses": v["uses"], "loaded_at": v["loaded_at"]} | |
| for t, v in self.loaded.items()}, | |
| "events": self.events[-20:]} | |
| def _pressure(self) -> bool: | |
| r = self._ram_reader() | |
| return r["percent"] >= self.ram_limit if r["percent"] >= 0 else False | |
| def _shed(self, keep: str) -> None: | |
| for task in reversed(PRIORITY): | |
| if task != keep and task in self.loaded: | |
| self.unload(task, reason="memory_pressure") | |
| if not self._pressure(): | |
| return | |
| def load(self, task: str) -> Any: | |
| with self.lock: | |
| if task in self.loaded: | |
| self.loaded[task]["uses"] += 1 | |
| return self.loaded[task]["handle"] | |
| if task not in self._loaders: | |
| raise RuntimeError(f"no loader registered for {task} (UNAVAILABLE)") | |
| if self._pressure(): | |
| self._shed(task) | |
| t0 = time.time() | |
| handle = self._loaders[task]() | |
| self.loaded[task] = {"handle": handle, "loaded_at": t0, "uses": 1} | |
| self.log("LOAD", f"{task} in {time.time() - t0:.1f}s") | |
| return handle | |
| def unload(self, task: str, reason: str = "manual") -> bool: | |
| with self.lock: | |
| h = self.loaded.pop(task, None) | |
| if h is None: | |
| return False | |
| try: | |
| handle = h["handle"] | |
| if hasattr(handle, "close"): | |
| handle.close() | |
| del handle | |
| except Exception: | |
| pass | |
| import gc | |
| gc.collect() | |
| self.log("UNLOAD", f"{task} ({reason})") | |
| return True | |
| def reload(self, task: str) -> Any: | |
| self.unload(task, reason="reload") | |
| return self.load(task) | |
| def ensure_unloaded_for_world(self) -> None: | |
| """Drop all optional models; world/physics unaffected (separate state).""" | |
| for task in list(self.loaded): | |
| self.unload(task, reason="world_priority") | |