drjones/gpu-program-swapper
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add barrier-confirmed yielding, measured residency, persistence and closed-loop tuning

Browse Source 
Nine changes, in rough order of how much they affect real behaviour:

1. VRAM yield is now a barrier. Posting keep_alive:0 only asks Ollama to unload;
   measured here, the HTTP call returns in 63ms while the driver takes a further
   77ms to release 14.9GB. Returning inside that window is how ComfyUI ends up
   allocating into VRAM that is still occupied. instant_free_ollama_vram() polls
   NVML until the allocation is actually gone and reports request/confirm split.

2. ComfyUI VRAM is no longer purged 1.5s after every prompt, which forced a full
   checkpoint reload on each workflow iteration. It is held for 30s of genuinely
   empty queue, with an immediate purge when Ollama actually asks for the memory.

3. Cache-hit classification uses achieved bandwidth (size / load duration) rather
   than a fixed `load_duration < 2500ms`. That constant called a 12.9GB model read
   at 2.9GB/s a cold load, and a 0.5GB model read from NVMe a cache hit.

4. Page-cache residency is measured, not assumed. mincore(2) reported 128GB
   resident on a box with 46GB of page cache: the kernel only permits page-cache
   introspection on files you own, and the Ollama blobs are owned by uid ollama,
   for which mincore answers "all resident" instead of failing. Uses cachestat(2)
   where permitted and a randomised read-rate probe elsewhere, labelling which was
   used. Fixed-offset probing was self-fulfilling, so windows are random and cold
   ones are returned with FADV_DONTNEED.

5. Warming is budgeted and ranked by recency/frequency instead of reading every
   file top-to-bottom, which on 64GB of RAM just evicts whatever was warmed first.

6. Telemetry and events persist to SQLite (~0.38 MB/hour) instead of living in a
   50-entry in-memory deque, so /api/analytics/profiles can finally answer whether
   an overclock profile actually delivers more tok/s.

7. Thermal governor walks the overclock back on sustained heat or hardware
   throttling, with hysteresis, fed from the existing sampler.

8. Autotune sweeps a clock offset, benchmarks decode at each step, watches for Xid
   errors and degenerate output, and restores the profile in a finally block.

9. Stock clocks/power/fans are restored on shutdown and via systemd ExecStopPost.
   Nothing previously undid a locked clock or a manually pinned fan.

Also: one shared 1Hz telemetry sampler fanned out to SSE subscribers rather than
every client re-running the whole snapshot; wall-clock timestamps in place of the
event loop's monotonic clock; cached nvidia-smi shell-outs; quieter httpx logging.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
drjones
2026-08-28 08:57:35 -07:00
parent 6acf235bd7
commit 5431144b2e
12 changed files with 2779 additions and 362 deletions
Download Patch File Download Diff File Expand all files Collapse all files

658
ram_optimizer.py
View File

@@ -1,16 +1,48 @@
"""RAM Optimizer and Model Pre-warmer for High-Speed Switching."""
import os
import glob
import time
import httpx
"""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
from typing import Dict, List, Any
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 = "/home/drjones/ComfyUI/models"
COMFY_MODELS_DIR = os.environ.get("HYPERSWAP_COMFY_MODELS", "/home/drjones/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."""
@@ -25,7 +57,7 @@ def get_detailed_meminfo() -> Dict[str, Any]:
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)
@@ -51,43 +83,439 @@ def get_detailed_meminfo() -> Dict[str, Any]:
"cache_ratio_pct": round((cached / total * 100) if total > 0 else 0, 1),
}
def warm_file_to_ram(filepath: str, chunk_size: int = 16 * 1024 * 1024) -> Dict[str, Any]:
"""Pre-fault/read file into Linux OS Page Cache at maximum disk read speed."""
# ---------------------------------------------------------------- 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) -> 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(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) -> 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)
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,
**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) -> Dict[str, Any]:
"""Pre-fault a file into the Linux page cache, skipping it if already resident."""
if not os.path.exists(filepath):
return {"success": False, "error": f"File not found: {filepath}", "duration_ms": 0}
before = page_residency(filepath)
if skip_if_warm and before.get("warm"):
return {
"success": True, "filepath": filepath, "skipped": True,
"reason": "already resident", "resident_pct": before.get("resident_pct"),
"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:
# Hint kernel that we will read this sequentially
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
duration_ms = round(duration * 1000, 2)
speed_mb_s = round((bytes_read / (1024**2)) / duration if duration > 0 else 0, 2)
after = page_residency(filepath)
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": duration_ms,
"speed_mb_s": speed_mb_s,
"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)}
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."""
@@ -101,76 +529,148 @@ async def warm_ollama_model(model_name: str, keep_alive: str = "5m") -> Dict[str
duration = time.perf_counter() - t0
if resp.status_code == 200:
data = resp.json()
return {
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),
}
else:
return {
"success": False,
"model": model_name,
"error": f"HTTP {resp.status_code}: {resp.text}",
"duration_ms": round(duration * 1000, 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)}
return {"success": False, "model": model_name, "error": str(e),
"duration_ms": round((time.perf_counter() - t0) * 1000, 2)}
def find_comfy_model_files() -> List[Dict[str, Any]]:
"""Discover all model files under ComfyUI models."""
results = []
extensions = ("*.safetensors", "*.ckpt", "*.pt", "*.bin")
if os.path.exists(COMFY_MODELS_DIR):
for root, _, files in os.walk(COMFY_MODELS_DIR):
for file in files:
if any(file.endswith(ext.replace("*", "")) for ext in extensions):
full_path = os.path.join(root, file)
rel_path = os.path.relpath(full_path, COMFY_MODELS_DIR)
size = os.path.getsize(full_path)
results.append({
"filename": file,
"rel_path": rel_path,
"full_path": full_path,
"size_bytes": size,
"size_mb": round(size / (1024**2), 2),
"size_gb": round(size / (1024**3), 3),
})
return results
async def warm_all_models() -> Dict[str, Any]:
"""Warm all available Ollama and ComfyUI models into Linux RAM Cache."""
t0 = time.perf_counter()
warmed_ollama = []
warmed_comfy = []
# 1. Ollama models
def warm_ollama_blob(model_name: str) -> 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"])
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:
async with httpx.AsyncClient(timeout=10.0) as client:
tags_resp = await client.get(f"{OLLAMA_API_BASE}/api/tags")
if tags_resp.status_code == 200:
models = tags_resp.json().get("models", [])
for m in models:
name = m.get("name")
res = await warm_ollama_model(name, keep_alive="1m")
warmed_ollama.append(res)
except Exception as e:
logger.error(f"Error discovering Ollama models: {e}")
# 2. ComfyUI models
comfy_files = find_comfy_model_files()
for f in comfy_files:
res = warm_file_to_ram(f["full_path"])
warmed_comfy.append(res)
total_duration_ms = round((time.perf_counter() - t0) * 1000, 2)
meminfo = get_detailed_meminfo()
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"])
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"):
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": total_duration_ms,
"ollama_models_warmed": warmed_ollama,
"comfy_files_warmed": warmed_comfy,
"meminfo_after": meminfo,
"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(),
}