The repo is about to be pushed to a remote, so this covers what should never travel with it and what should not be baked into the source. .gitignore now covers credentials (.env, keys, tokens, .netrc), host-local config (*.local.json), the SQLite telemetry store and its WAL sidecars, logs, benchmark and sweep output, and the timestamped .bak files this project has accumulated before. Verified that no currently tracked file is caught by the new patterns. Also removed /home/drjones from tracked source, which an ignore file cannot help with. BASE_DIR now derives from the module's own location, the ComfyUI model directory falls back to ~/ComfyUI/models, and start_manager.sh resolves its interpreter through $HOME with a python3 fallback. All three still resolve to exactly the same paths on this machine; they just no longer hardcode one user's home directory into a published repository. Note for the record: the git history was scanned across all refs and contains no credentials. The password visible in `git remote -v` lives only in .git/config, which is never pushed. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
687 lines
27 KiB
Python
687 lines
27 KiB
Python
"""RAM Optimizer and Model Pre-warmer for High-Speed Switching.
|
|
|
|
Two things changed here versus the naive version:
|
|
|
|
1. Residency is *measured*, not assumed. mincore(2) tells us exactly what fraction of
|
|
each model file is resident in the Linux page cache, so "RAM Cache Hit" stops being
|
|
a guess based on how long a load took.
|
|
2. Warming is *budgeted*. This box has 64 GB of RAM and >33 GB of models; reading every
|
|
file top-to-bottom simply evicts whatever was warmed first. Files are now scored by
|
|
recency/frequency (from the telemetry store) and warmed until a byte budget is hit,
|
|
skipping anything already resident.
|
|
"""
|
|
import ctypes
|
|
import ctypes.util
|
|
import json
|
|
import logging
|
|
import os
|
|
import random
|
|
import time
|
|
from typing import Dict, List, Any, Optional, Tuple
|
|
|
|
import httpx
|
|
|
|
import telemetry_store
|
|
|
|
logger = logging.getLogger("ram_optimizer")
|
|
|
|
OLLAMA_API_BASE = "http://localhost:11434"
|
|
COMFY_API_BASE = "http://127.0.0.1:8188"
|
|
COMFY_MODELS_DIR = os.environ.get(
|
|
"HYPERSWAP_COMFY_MODELS", os.path.expanduser("~/ComfyUI/models"))
|
|
OLLAMA_MODEL_DIRS = [
|
|
"/usr/share/ollama/.ollama/models",
|
|
os.path.expanduser("~/.ollama/models"),
|
|
]
|
|
|
|
PAGE_SIZE = os.sysconf("SC_PAGE_SIZE")
|
|
# Files bigger than this are sampled rather than fully mapped for residency.
|
|
RESIDENCY_FULL_MAP_LIMIT = 2 * 1024 ** 3
|
|
RESIDENCY_SAMPLE_WINDOWS = 64
|
|
RESIDENCY_WINDOW_BYTES = 16 * 1024 * 1024
|
|
# A file at/above this residency is considered warm and is skipped by the warmer.
|
|
WARM_SKIP_THRESHOLD_PCT = 90.0
|
|
|
|
CATALOG_TTL_S = 30.0
|
|
|
|
|
|
def get_detailed_meminfo() -> Dict[str, Any]:
|
|
"""Parse /proc/meminfo for precise page cache and RAM stats."""
|
|
info = {}
|
|
try:
|
|
with open("/proc/meminfo", "r") as f:
|
|
for line in f:
|
|
parts = line.split(":")
|
|
if len(parts) == 2:
|
|
key = parts[0].strip()
|
|
val = parts[1].strip().split()[0]
|
|
info[key] = int(val) * 1024 # Convert kB to bytes
|
|
except Exception as e:
|
|
logger.error(f"Failed to read /proc/meminfo: {e}")
|
|
|
|
total = info.get("MemTotal", 0)
|
|
free = info.get("MemFree", 0)
|
|
available = info.get("MemAvailable", 0)
|
|
cached = info.get("Cached", 0) + info.get("Buffers", 0)
|
|
dirty = info.get("Dirty", 0)
|
|
used = total - free - cached
|
|
if used < 0:
|
|
used = total - available
|
|
|
|
return {
|
|
"total_bytes": total,
|
|
"total_gb": round(total / (1024**3), 2),
|
|
"used_bytes": used,
|
|
"used_gb": round(used / (1024**3), 2),
|
|
"cached_bytes": cached,
|
|
"cached_gb": round(cached / (1024**3), 2),
|
|
"free_bytes": free,
|
|
"free_gb": round(free / (1024**3), 2),
|
|
"available_bytes": available,
|
|
"available_gb": round(available / (1024**3), 2),
|
|
"dirty_bytes": dirty,
|
|
"dirty_mb": round(dirty / (1024**2), 2),
|
|
"cache_ratio_pct": round((cached / total * 100) if total > 0 else 0, 1),
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------- page residency
|
|
#
|
|
# Measuring page-cache residency turned out to be the subtle part.
|
|
#
|
|
# * cachestat(2) (Linux 6.5+) is the right tool: exact cached-page counts for an fd,
|
|
# no mmap, microseconds per call. But the kernel only permits it on files you own
|
|
# or can write -- the Ollama blobs are owned by uid `ollama`, so it returns EPERM.
|
|
# * mincore(2) does NOT fail closed for those files on this kernel: it reports every
|
|
# page as resident, which produced 128 GB of "resident" model weights on a box with
|
|
# 46 GB of page cache. It is therefore not used at all.
|
|
#
|
|
# So: cachestat where permitted, and an explicit read-throughput probe where it is not.
|
|
# Anything we cannot measure is reported as unmeasurable rather than guessed at.
|
|
|
|
_libc = None
|
|
_SYS_cachestat = 451 # x86_64
|
|
|
|
|
|
class _CachestatRange(ctypes.Structure):
|
|
_fields_ = [("off", ctypes.c_uint64), ("len", ctypes.c_uint64)]
|
|
|
|
|
|
class _Cachestat(ctypes.Structure):
|
|
_fields_ = [
|
|
("nr_cache", ctypes.c_uint64),
|
|
("nr_dirty", ctypes.c_uint64),
|
|
("nr_writeback", ctypes.c_uint64),
|
|
("nr_evicted", ctypes.c_uint64),
|
|
("nr_recently_evicted", ctypes.c_uint64),
|
|
]
|
|
|
|
|
|
def _get_libc():
|
|
global _libc
|
|
if _libc is None:
|
|
_libc = ctypes.CDLL(ctypes.util.find_library("c") or "libc.so.6", use_errno=True)
|
|
return _libc
|
|
|
|
|
|
def _cachestat(fd: int, offset: int, length: int) -> Optional[_Cachestat]:
|
|
"""Raw cachestat(2). Returns None if the kernel refuses (EPERM/ENOSYS)."""
|
|
libc = _get_libc()
|
|
rng = _CachestatRange(offset, length)
|
|
cs = _Cachestat()
|
|
ctypes.set_errno(0)
|
|
rc = libc.syscall(ctypes.c_long(_SYS_cachestat), ctypes.c_int(fd),
|
|
ctypes.byref(rng), ctypes.byref(cs), ctypes.c_uint(0))
|
|
if rc != 0:
|
|
return None
|
|
return cs
|
|
|
|
|
|
PROBE_WINDOWS = 12
|
|
PROBE_WINDOW_BYTES = 2 * 1024 * 1024
|
|
# Measured on this box: cold NVMe reads land around 0.35-0.5 GB/s, page-cache reads at
|
|
# 3.2-13 GB/s. 1.5 GB/s sits in the empty middle of that gap.
|
|
PROBE_CACHED_GBPS = 1.5
|
|
|
|
|
|
def _throughput_probe(fd: int, size: int, windows_override: Optional[int] = None) -> Dict[str, Any]:
|
|
"""Infer residency by timing reads of small windows spread across the file.
|
|
|
|
Used only where cachestat is not permitted (Ollama's blobs are owned by uid `ollama`).
|
|
|
|
Two details matter for correctness:
|
|
|
|
* Offsets are random per call. A fixed stride made the probe self-fulfilling: the
|
|
first pass faulted its 24 MB of sample windows into the page cache, and every pass
|
|
after that re-read exactly those windows and reported 100% resident for a file that
|
|
was almost entirely cold.
|
|
* Windows that read cold are handed straight back with FADV_DONTNEED. Those pages are
|
|
pollution the probe itself created, and leaving them behind would slowly warm the
|
|
cache with data nobody asked for.
|
|
"""
|
|
windows = min(windows_override or PROBE_WINDOWS, max(int(size // PROBE_WINDOW_BYTES), 1))
|
|
if windows <= 0:
|
|
return {"resident_pct": 0.0, "windows": 0}
|
|
|
|
max_off = max(size - PROBE_WINDOW_BYTES, 0)
|
|
offsets = sorted(random.randint(0, max_off) for _ in range(windows)) if max_off else [0]
|
|
buf = bytearray(PROBE_WINDOW_BYTES)
|
|
cached = 0
|
|
rates = []
|
|
for off in offsets:
|
|
length = min(PROBE_WINDOW_BYTES, size - off)
|
|
if length <= 0:
|
|
continue
|
|
view = memoryview(buf)[:length]
|
|
t0 = time.perf_counter()
|
|
os.preadv(fd, [view], off)
|
|
dt = time.perf_counter() - t0
|
|
gbps = (length / (1024 ** 3)) / dt if dt > 0 else 0.0
|
|
rates.append(gbps)
|
|
if gbps >= PROBE_CACHED_GBPS:
|
|
cached += 1
|
|
else:
|
|
# We just pulled this off disk; put it back the way we found it.
|
|
try:
|
|
os.posix_fadvise(fd, off, length, os.POSIX_FADV_DONTNEED)
|
|
except Exception:
|
|
pass
|
|
n = len(rates)
|
|
return {
|
|
"resident_pct": round((cached / n * 100) if n else 0.0, 1),
|
|
"windows": n,
|
|
"median_gbps": round(sorted(rates)[n // 2], 2) if n else 0.0,
|
|
"sampled_gb": round(n * PROBE_WINDOW_BYTES / (1024 ** 3), 3),
|
|
}
|
|
|
|
|
|
def page_residency(filepath: str, allow_probe: bool = True,
|
|
probe_windows: Optional[int] = None) -> Dict[str, Any]:
|
|
"""Measure what fraction of a file is resident in the Linux page cache."""
|
|
try:
|
|
size = os.path.getsize(filepath)
|
|
except OSError as e:
|
|
return {"success": False, "error": str(e), "resident_pct": 0.0, "measurable": False}
|
|
if size == 0:
|
|
return {"success": True, "resident_pct": 0.0, "size_bytes": 0, "measurable": True,
|
|
"method": "empty"}
|
|
|
|
try:
|
|
fd = os.open(filepath, os.O_RDONLY)
|
|
except OSError as e:
|
|
return {"success": False, "error": str(e), "resident_pct": 0.0, "measurable": False}
|
|
try:
|
|
cs = _cachestat(fd, 0, size)
|
|
if cs is not None:
|
|
total_pages = (size + PAGE_SIZE - 1) // PAGE_SIZE
|
|
pct = round((cs.nr_cache / total_pages * 100) if total_pages else 0.0, 1)
|
|
method, measurable = "cachestat", True
|
|
extra = {"dirty_pages": cs.nr_dirty, "evicted_pages": cs.nr_evicted}
|
|
elif allow_probe:
|
|
probe = _throughput_probe(fd, size, probe_windows)
|
|
pct = probe["resident_pct"]
|
|
method, measurable = "probe", True
|
|
extra = {"probe_windows": probe["windows"], "probe_median_gbps": probe.get("median_gbps")}
|
|
else:
|
|
return {"success": True, "filepath": filepath, "size_bytes": size,
|
|
"size_gb": round(size / (1024**3), 3), "resident_pct": None,
|
|
"measurable": False, "method": "unavailable", "warm": None,
|
|
"reason": "cachestat not permitted for this file (not owned by us)"}
|
|
|
|
return {
|
|
"success": True,
|
|
"filepath": filepath,
|
|
"size_bytes": size,
|
|
"size_gb": round(size / (1024**3), 3),
|
|
"resident_pct": pct,
|
|
"resident_bytes": int(size * pct / 100.0),
|
|
"method": method,
|
|
"measurable": measurable,
|
|
"warm": pct >= WARM_SKIP_THRESHOLD_PCT,
|
|
# Only an exact measurement is trustworthy enough to skip work on. A probe of a
|
|
# dozen 2 MB windows can clear 90% on a file that is mostly cold -- observed
|
|
# here as a 12.87 GB "already resident" blob that then loaded at 2.44 GB/s.
|
|
"warm_confident": (method == "cachestat" and pct >= WARM_SKIP_THRESHOLD_PCT)
|
|
or (method == "probe" and pct >= 100.0),
|
|
**extra,
|
|
}
|
|
except Exception as e:
|
|
return {"success": False, "error": str(e), "resident_pct": 0.0,
|
|
"size_bytes": size, "measurable": False}
|
|
finally:
|
|
os.close(fd)
|
|
|
|
|
|
def residency_capability() -> Dict[str, Any]:
|
|
"""Report whether exact residency is available, and how to enable it if not."""
|
|
catalog = get_model_catalog()
|
|
blocked = []
|
|
for f in catalog["ollama"]:
|
|
try:
|
|
fd = os.open(f["full_path"], os.O_RDONLY)
|
|
except OSError:
|
|
continue
|
|
try:
|
|
if _cachestat(fd, 0, 4096) is None:
|
|
blocked.append(f["full_path"])
|
|
finally:
|
|
os.close(fd)
|
|
break # one probe is enough; blobs share a directory and owner
|
|
if not blocked:
|
|
return {"exact_everywhere": True}
|
|
owner = ""
|
|
try:
|
|
import pwd
|
|
owner = pwd.getpwuid(os.stat(blocked[0]).st_uid).pw_name
|
|
except Exception:
|
|
owner = str(os.stat(blocked[0]).st_uid)
|
|
return {
|
|
"exact_everywhere": False,
|
|
"method_for_blocked": "probe",
|
|
"reason": f"cachestat(2) is only permitted on files you own or can write; "
|
|
f"Ollama blobs are owned by '{owner}'",
|
|
"hint": f"exact numbers for Ollama weights need read/write access, e.g. "
|
|
f"'sudo usermod -aG {owner} $USER' plus group-write on the blobs directory",
|
|
}
|
|
|
|
|
|
# ---------------------------------------------------------------- catalogs
|
|
|
|
_catalog_cache: Dict[str, Any] = {"ts": 0.0, "sig": None, "comfy": [], "ollama": []}
|
|
|
|
|
|
def _dir_signature(root: str) -> Tuple:
|
|
"""Cheap fingerprint of a model tree: (mtime, entry count) per subdirectory."""
|
|
sig = []
|
|
if not os.path.isdir(root):
|
|
return tuple(sig)
|
|
for dirpath, dirnames, filenames in os.walk(root):
|
|
try:
|
|
sig.append((dirpath, os.stat(dirpath).st_mtime_ns, len(filenames)))
|
|
except OSError:
|
|
continue
|
|
return tuple(sig)
|
|
|
|
|
|
def find_ollama_model_files() -> List[Dict[str, Any]]:
|
|
"""Map installed Ollama models to their on-disk GGUF blobs via the manifest tree.
|
|
|
|
Knowing the blob path is what lets us warm (or measure) a specific model's weights
|
|
without pulling them into VRAM.
|
|
"""
|
|
results: List[Dict[str, Any]] = []
|
|
seen = set()
|
|
for root in OLLAMA_MODEL_DIRS:
|
|
manifests = os.path.join(root, "manifests")
|
|
blobs = os.path.join(root, "blobs")
|
|
if not os.path.isdir(manifests):
|
|
continue
|
|
for dirpath, _, filenames in os.walk(manifests):
|
|
for tag in filenames:
|
|
manifest_path = os.path.join(dirpath, tag)
|
|
try:
|
|
with open(manifest_path) as f:
|
|
manifest = json.load(f)
|
|
except Exception:
|
|
continue
|
|
rel = os.path.relpath(dirpath, manifests)
|
|
parts = rel.split(os.sep)
|
|
# registry/namespace/name -> "name:tag", keeping non-library namespaces
|
|
name = parts[-1] if parts else rel
|
|
namespace = parts[-2] if len(parts) >= 2 else "library"
|
|
model_name = f"{name}:{tag}" if namespace == "library" else f"{namespace}/{name}:{tag}"
|
|
for layer in manifest.get("layers", []):
|
|
if layer.get("mediaType") != "application/vnd.ollama.image.model":
|
|
continue
|
|
digest = (layer.get("digest") or "").replace(":", "-")
|
|
blob_path = os.path.join(blobs, digest)
|
|
if not os.path.exists(blob_path):
|
|
continue
|
|
key = (model_name, blob_path)
|
|
if key in seen:
|
|
continue
|
|
seen.add(key)
|
|
size = layer.get("size") or os.path.getsize(blob_path)
|
|
results.append({
|
|
"model": model_name,
|
|
"filename": digest,
|
|
"full_path": blob_path,
|
|
"size_bytes": size,
|
|
"size_gb": round(size / (1024**3), 3),
|
|
"kind": "ollama",
|
|
})
|
|
return results
|
|
|
|
|
|
def find_comfy_model_files(force_refresh: bool = False) -> List[Dict[str, Any]]:
|
|
"""Discover all model files under ComfyUI models (cached).
|
|
|
|
This used to run inside the 1Hz telemetry snapshot, meaning a full recursive walk plus
|
|
a stat() of every checkpoint once per second per connected dashboard. It is now cached
|
|
behind a directory-mtime fingerprint.
|
|
"""
|
|
_refresh_catalog(force_refresh)
|
|
return _catalog_cache["comfy"]
|
|
|
|
|
|
def get_model_catalog(force_refresh: bool = False) -> Dict[str, Any]:
|
|
_refresh_catalog(force_refresh)
|
|
return {
|
|
"comfy": _catalog_cache["comfy"],
|
|
"ollama": _catalog_cache["ollama"],
|
|
"cached_at": _catalog_cache["ts"],
|
|
}
|
|
|
|
|
|
def _refresh_catalog(force: bool = False) -> None:
|
|
now = time.time()
|
|
if not force and (now - _catalog_cache["ts"]) < CATALOG_TTL_S:
|
|
return
|
|
sig = _dir_signature(COMFY_MODELS_DIR)
|
|
if not force and sig == _catalog_cache["sig"] and _catalog_cache["comfy"]:
|
|
_catalog_cache["ts"] = now
|
|
return
|
|
|
|
extensions = (".safetensors", ".ckpt", ".pt", ".bin", ".gguf", ".sft")
|
|
results = []
|
|
if os.path.exists(COMFY_MODELS_DIR):
|
|
for root, _, files in os.walk(COMFY_MODELS_DIR):
|
|
for file in files:
|
|
if not file.endswith(extensions):
|
|
continue
|
|
full_path = os.path.join(root, file)
|
|
try:
|
|
st = os.stat(full_path)
|
|
except OSError:
|
|
continue
|
|
results.append({
|
|
"filename": file,
|
|
"rel_path": os.path.relpath(full_path, COMFY_MODELS_DIR),
|
|
"full_path": full_path,
|
|
"category": os.path.relpath(root, COMFY_MODELS_DIR).split(os.sep)[0],
|
|
"size_bytes": st.st_size,
|
|
"size_mb": round(st.st_size / (1024**2), 2),
|
|
"size_gb": round(st.st_size / (1024**3), 3),
|
|
"mtime": st.st_mtime,
|
|
"kind": "comfy",
|
|
})
|
|
_catalog_cache.update({"ts": now, "sig": sig, "comfy": results,
|
|
"ollama": find_ollama_model_files()})
|
|
|
|
|
|
# ---------------------------------------------------------------- residency report
|
|
|
|
_report_cache: Dict[str, Any] = {"ts": 0.0, "report": None}
|
|
REPORT_TTL_S = 15.0
|
|
|
|
|
|
def get_cache_report(include_files: bool = True, force_refresh: bool = False) -> Dict[str, Any]:
|
|
"""Measured page-cache residency across the whole model catalog.
|
|
|
|
Deduplicated by blob path: several Ollama tags routinely point at the same GGUF, and
|
|
counting each tag separately produced more "resident" bytes than the box has RAM.
|
|
"""
|
|
now = time.time()
|
|
cached = _report_cache["report"]
|
|
if cached and not force_refresh and (now - _report_cache["ts"]) < REPORT_TTL_S:
|
|
return cached if include_files else {**cached, "files": []}
|
|
|
|
t0 = time.perf_counter()
|
|
catalog = get_model_catalog()
|
|
by_path: Dict[str, Dict[str, Any]] = {}
|
|
for f in list(catalog["ollama"]) + list(catalog["comfy"]):
|
|
path = f["full_path"]
|
|
name = f.get("model") or f.get("rel_path") or f.get("filename")
|
|
if path in by_path:
|
|
by_path[path]["aliases"].append(name)
|
|
continue
|
|
by_path[path] = {"entry": f, "name": name, "aliases": []}
|
|
|
|
entries = []
|
|
total_bytes = resident_bytes = 0
|
|
for path, meta in by_path.items():
|
|
f = meta["entry"]
|
|
res = page_residency(path)
|
|
size = f.get("size_bytes") or res.get("size_bytes") or 0
|
|
rb = res.get("resident_bytes", 0)
|
|
total_bytes += size
|
|
resident_bytes += rb
|
|
entries.append({
|
|
"name": meta["name"],
|
|
"aliases": meta["aliases"],
|
|
"kind": f.get("kind"),
|
|
"full_path": path,
|
|
"size_gb": round(size / (1024**3), 3),
|
|
"resident_pct": res.get("resident_pct", 0.0),
|
|
"resident_gb": round(rb / (1024**3), 3),
|
|
"warm": res.get("warm", False),
|
|
})
|
|
entries.sort(key=lambda e: e["resident_gb"], reverse=True)
|
|
|
|
report = {
|
|
"scan_ms": round((time.perf_counter() - t0) * 1000, 1),
|
|
"files_scanned": len(entries),
|
|
"unique_blobs": len(by_path),
|
|
"catalog_total_gb": round(total_bytes / (1024**3), 2),
|
|
"resident_total_gb": round(resident_bytes / (1024**3), 2),
|
|
"residency_pct": round((resident_bytes / total_bytes * 100) if total_bytes else 0, 1),
|
|
"warm_files": sum(1 for e in entries if e["warm"]),
|
|
"files": entries,
|
|
}
|
|
_report_cache.update({"ts": now, "report": report})
|
|
return report if include_files else {**report, "files": []}
|
|
|
|
|
|
# ---------------------------------------------------------------- warming
|
|
|
|
def warm_file_to_ram(filepath: str, chunk_size: int = 16 * 1024 * 1024,
|
|
skip_if_warm: bool = True, force: bool = False) -> Dict[str, Any]:
|
|
"""Pre-fault a file into the Linux page cache, skipping it only if confidently resident."""
|
|
if not os.path.exists(filepath):
|
|
return {"success": False, "error": f"File not found: {filepath}", "duration_ms": 0}
|
|
|
|
# Probe densely here: this decision skips real work, so it is worth 32 samples
|
|
# rather than 12.
|
|
before = page_residency(filepath, probe_windows=32)
|
|
if skip_if_warm and not force and before.get("warm_confident"):
|
|
return {
|
|
"success": True, "filepath": filepath, "skipped": True,
|
|
"reason": "already resident", "resident_pct": before.get("resident_pct"),
|
|
"method": before.get("method"),
|
|
"size_mb": round(before.get("size_bytes", 0) / (1024**2), 2),
|
|
"duration_ms": 0.0, "bytes_read": 0,
|
|
}
|
|
|
|
t0 = time.perf_counter()
|
|
file_size = os.path.getsize(filepath)
|
|
bytes_read = 0
|
|
try:
|
|
with open(filepath, "rb") as f:
|
|
try:
|
|
os.posix_fadvise(f.fileno(), 0, file_size, os.POSIX_FADV_WILLNEED)
|
|
except Exception:
|
|
pass
|
|
buf = bytearray(chunk_size)
|
|
while True:
|
|
n = f.readinto(buf)
|
|
if not n:
|
|
break
|
|
bytes_read += n
|
|
|
|
duration = time.perf_counter() - t0
|
|
after = page_residency(filepath, probe_windows=32)
|
|
return {
|
|
"success": True,
|
|
"filepath": filepath,
|
|
"skipped": False,
|
|
"size_bytes": file_size,
|
|
"size_mb": round(file_size / (1024**2), 2),
|
|
"bytes_read": bytes_read,
|
|
"duration_ms": round(duration * 1000, 2),
|
|
"speed_mb_s": round((bytes_read / (1024**2)) / duration if duration > 0 else 0, 2),
|
|
"resident_pct_before": before.get("resident_pct", 0.0),
|
|
"resident_pct_after": after.get("resident_pct", 0.0),
|
|
}
|
|
except Exception as e:
|
|
return {"success": False, "error": str(e),
|
|
"duration_ms": round((time.perf_counter() - t0) * 1000, 2)}
|
|
|
|
|
|
async def warm_ollama_model(model_name: str, keep_alive: str = "5m") -> Dict[str, Any]:
|
|
"""Warm an Ollama model into memory and measure time."""
|
|
t0 = time.perf_counter()
|
|
try:
|
|
async with httpx.AsyncClient(timeout=120.0) as client:
|
|
resp = await client.post(
|
|
f"{OLLAMA_API_BASE}/api/generate",
|
|
json={"model": model_name, "prompt": "", "keep_alive": keep_alive},
|
|
)
|
|
duration = time.perf_counter() - t0
|
|
if resp.status_code == 200:
|
|
data = resp.json()
|
|
res = {
|
|
"success": True,
|
|
"model": model_name,
|
|
"duration_ms": round(duration * 1000, 2),
|
|
"load_duration_ms": round(data.get("load_duration", 0) / 1e6, 2),
|
|
"total_duration_ms": round(data.get("total_duration", 0) / 1e6, 2),
|
|
}
|
|
telemetry_store.record_event({
|
|
"event_type": "Model Warm", "source": "warmer", "target": model_name,
|
|
"duration_ms": res["duration_ms"], "load_duration_ms": res["load_duration_ms"],
|
|
})
|
|
return res
|
|
return {
|
|
"success": False, "model": model_name,
|
|
"error": f"HTTP {resp.status_code}: {resp.text}",
|
|
"duration_ms": round(duration * 1000, 2),
|
|
}
|
|
except Exception as e:
|
|
return {"success": False, "model": model_name, "error": str(e),
|
|
"duration_ms": round((time.perf_counter() - t0) * 1000, 2)}
|
|
|
|
|
|
def warm_ollama_blob(model_name: str, force: bool = False) -> Dict[str, Any]:
|
|
"""Warm a specific Ollama model's GGUF into page cache without touching VRAM."""
|
|
for f in find_ollama_model_files():
|
|
if f["model"] == model_name:
|
|
res = warm_file_to_ram(f["full_path"], force=force)
|
|
res["model"] = model_name
|
|
return res
|
|
return {"success": False, "error": f"no blob found for model '{model_name}'"}
|
|
|
|
|
|
def _warm_priority(days: float = 30.0) -> Dict[str, float]:
|
|
"""Recency/frequency score per model name, from the persisted event log."""
|
|
try:
|
|
return {r["model"]: r["score"] for r in telemetry_store.model_usage_ranking(days)}
|
|
except Exception:
|
|
return {}
|
|
|
|
|
|
def build_warm_plan(budget_gb: Optional[float] = None) -> Dict[str, Any]:
|
|
"""Decide *what* to warm, in what order, within a byte budget.
|
|
|
|
Warming everything on a 64 GB box with 33+ GB of models just evicts the earliest
|
|
files, so we rank by usage (Ollama, from history) and recency (ComfyUI, by mtime),
|
|
then fill until the budget is spent. Already-resident files cost nothing.
|
|
"""
|
|
mem = get_detailed_meminfo()
|
|
if budget_gb is None:
|
|
# Leave headroom so warming never pushes the box into reclaim.
|
|
budget_gb = max((mem["available_bytes"] * 0.7) / (1024**3), 1.0)
|
|
budget_bytes = int(budget_gb * (1024**3))
|
|
|
|
catalog = get_model_catalog()
|
|
scores = _warm_priority()
|
|
now = time.time()
|
|
|
|
candidates = []
|
|
for f in catalog["ollama"]:
|
|
candidates.append({**f, "score": scores.get(f["model"], 0.0) + 0.5,
|
|
"name": f["model"]})
|
|
for f in catalog["comfy"]:
|
|
age_days = max((now - f.get("mtime", now)) / 86400.0, 0.01)
|
|
candidates.append({**f, "score": scores.get(f["rel_path"], 0.0) + 1.0 / (1.0 + age_days),
|
|
"name": f["rel_path"]})
|
|
|
|
candidates.sort(key=lambda c: c["score"], reverse=True)
|
|
|
|
plan, spent, skipped = [], 0, []
|
|
seen_paths = set()
|
|
for c in candidates:
|
|
if c["full_path"] in seen_paths:
|
|
continue
|
|
seen_paths.add(c["full_path"])
|
|
res = page_residency(c["full_path"], probe_windows=32)
|
|
entry = {
|
|
"name": c["name"], "kind": c["kind"], "full_path": c["full_path"],
|
|
"size_gb": c.get("size_gb", 0), "score": round(c["score"], 4),
|
|
"resident_pct": res.get("resident_pct", 0.0),
|
|
}
|
|
if res.get("warm_confident"):
|
|
entry["action"] = "already-warm"
|
|
skipped.append(entry)
|
|
continue
|
|
need = int(c.get("size_bytes", 0) * (1 - res.get("resident_pct", 0) / 100.0))
|
|
if spent + need > budget_bytes:
|
|
entry["action"] = "over-budget"
|
|
skipped.append(entry)
|
|
continue
|
|
spent += need
|
|
entry["action"] = "warm"
|
|
entry["bytes_to_read"] = need
|
|
plan.append(entry)
|
|
|
|
return {
|
|
"budget_gb": round(budget_gb, 2),
|
|
"planned_gb": round(spent / (1024**3), 2),
|
|
"warm_count": len(plan),
|
|
"skipped_count": len(skipped),
|
|
"plan": plan,
|
|
"skipped": skipped,
|
|
"meminfo": mem,
|
|
}
|
|
|
|
|
|
async def warm_all_models(budget_gb: Optional[float] = None,
|
|
include_vram_load: bool = False) -> Dict[str, Any]:
|
|
"""Warm the highest-value models into the page cache within a byte budget."""
|
|
t0 = time.perf_counter()
|
|
plan = build_warm_plan(budget_gb)
|
|
warmed = []
|
|
for entry in plan["plan"]:
|
|
res = warm_file_to_ram(entry["full_path"])
|
|
res["name"] = entry["name"]
|
|
res["kind"] = entry["kind"]
|
|
warmed.append(res)
|
|
# Budgets are computed up front, but the page cache is shared with the rest of
|
|
# the box; bail out if we start pushing the system into reclaim.
|
|
if get_detailed_meminfo()["available_gb"] < 4.0:
|
|
logger.warning("warm_all_models: stopping early, MemAvailable below 4 GB")
|
|
break
|
|
|
|
if include_vram_load:
|
|
try:
|
|
async with httpx.AsyncClient(timeout=10.0) as client:
|
|
tags = await client.get(f"{OLLAMA_API_BASE}/api/tags")
|
|
if tags.status_code == 200:
|
|
top = sorted(tags.json().get("models", []),
|
|
key=lambda m: _warm_priority().get(m.get("name"), 0),
|
|
reverse=True)[:1]
|
|
for m in top:
|
|
await warm_ollama_model(m.get("name"), keep_alive="1m")
|
|
except Exception as e:
|
|
logger.debug(f"optional VRAM preload skipped: {e}")
|
|
|
|
return {
|
|
"status": "completed",
|
|
"total_duration_ms": round((time.perf_counter() - t0) * 1000, 2),
|
|
"budget_gb": plan["budget_gb"],
|
|
"planned_gb": plan["planned_gb"],
|
|
"files_warmed": warmed,
|
|
"bytes_read": sum(w.get("bytes_read", 0) for w in warmed),
|
|
"skipped": plan["skipped"],
|
|
"meminfo_after": get_detailed_meminfo(),
|
|
}
|