25b601e24a
The stale-readback bug fixed in f5917a0 was a class, not an instance. Two more:
- ram_optimizer's residency report cached for 15s and was never invalidated when
anything warmed a file, so warming a model and then looking at residency showed
the state from before the warm. warm_file_to_ram now invalidates it.
- _PID_KIND_CACHE was keyed on pid alone and never expired. Linux recycles PIDs, so
a stale entry could attribute a new process's VRAM to Ollama or ComfyUI -- inside
the very snapshot the yield barrier trusts to decide whether VRAM was released.
Now keyed by (pid, process start time) and bounded.
Unmanaged VRAM. Investigating a persistence-mode warning turned up a third GPU
consumer this service does not model: stt_relay.py, holding 842 MB for nearly three
days. It was bucketed as "system" alongside gnome-shell's 3.9 MB. That conflation
matters, because ComfyUI's memory can be reclaimed and a third party's cannot, and
the reclaim path assumed ComfyUI was always to blame for missing headroom.
Processes are now bucketed ollama | comfy | desktop | unmanaged. The breakdown
reports desktop_gb and unmanaged_gb separately and names the unmanaged processes;
when a reclaim-and-retry still fails, the error identifies them rather than
implying ComfyUI was at fault; and the dashboard shows the unreclaimable total, so
headroom the arbitrator can never give back is visible rather than inferred.
Checked and deliberately not changed: persistence mode reads Disabled, but
nvidia-persistenced is active and two clients hold the GPU open continuously, so
the driver never unloads. The nvidia-smi warning is legacy noise here and is not a
source of the profile drift.
Tests: 169 (was 164). The new ones cover the bucketing, and one existing test used
Xorg as its "unknown process" fixture -- correct before a display server had its own
bucket, wrong after.
Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
700 lines
28 KiB
Python
700 lines
28 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:
|
|
# Keep the shape identical to every other return path; callers index these.
|
|
return {"success": True, "filepath": filepath, "size_bytes": 0, "size_gb": 0.0,
|
|
"resident_pct": 0.0, "resident_bytes": 0, "measurable": True,
|
|
"method": "empty", "warm": True, "warm_confident": True}
|
|
|
|
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 invalidate_cache_report() -> None:
|
|
"""Drop the cached residency report.
|
|
|
|
Anything that changes what is resident must call this, or the report keeps serving
|
|
pre-change numbers for up to REPORT_TTL_S -- so warming a model and then looking at
|
|
residency showed the state from before the warm.
|
|
"""
|
|
_report_cache["report"] = None
|
|
|
|
|
|
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)
|
|
invalidate_cache_report()
|
|
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(),
|
|
}
|