Compare commits
7 Commits
172d812820
...
b53b026ced
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
b53b026ced | ||
|
|
2948b0b440 | ||
|
|
fb104ac9c0 | ||
|
|
25b601e24a | ||
|
|
f5917a0464 | ||
|
|
868d82794d | ||
|
|
bacaf50713 |
60
README.md
60
README.md
@@ -20,6 +20,21 @@
|
||||
## 1. Feature Matrix
|
||||
|
||||
### ⚡ Bidirectional VRAM Hot-Swapping & Arbitration
|
||||
|
||||
* **Both directions are now automatic.** Yielding Ollama for ComfyUI always was; the
|
||||
reverse was not, despite "bidirectional" in this heading. Which way an LLM fails when
|
||||
it cannot fit depends on configuration: with `n_gpu_layers` left to Ollama it spills
|
||||
layers to the CPU and reports `size_vram < size` (roughly an order of magnitude slower,
|
||||
and silent). With `n_gpu_layers` pinned — 99 on this box — it refuses outright with
|
||||
`cudaMalloc failed: out of memory`. Both are handled: the spill triggers a reclaim from
|
||||
an idle ComfyUI, and the hard failure is caught by `switch_ollama_model`, which reclaims
|
||||
and retries once. Measured: a 12.87 GB model that returned HTTP 500 from Ollama directly
|
||||
now loads through HyperSwap after reclaiming 6.83 GB, at 3.85 GB/s.
|
||||
* **A busy LLM is not a failed yield.** A model mid-generation cannot unload; the
|
||||
`keep_alive: 0` request queues behind it and applies when it finishes. That is reported
|
||||
as `busy` (returning in ~610 ms) rather than blocking, with per-model backoff and a
|
||||
detached watcher that logs the eventual release. Only VRAM held while the GPU sits
|
||||
*idle* counts as a fault.
|
||||
* **Confirmed Soft-Yield (barrier, not fire-and-forget)**: Releases Ollama VRAM allocations (`keep_alive: 0`) down to 0 MB, then **waits on NVML until the driver has actually freed the allocation** before letting ComfyUI proceed. Posting `keep_alive: 0` only *asks* Ollama to unload; on this box the HTTP call returns in ~63 ms while the driver takes a further ~77 ms to release 14.9 GB. Returning during that window is how diffusion ends up allocating into VRAM that is still occupied.
|
||||
* **Idle-Aware ComfyUI Purge**: Diffusion checkpoints are held for `COMFY_IDLE_PURGE_S` (30 s) of genuinely empty queue rather than purged 1.5 s after every prompt — iterating on a workflow no longer pays a full checkpoint reload per run. An immediate purge still happens the moment Ollama actually asks for VRAM (`POST /api/request-vram`).
|
||||
* **Real-Time ComfyUI WebSocket & Watchdog Listener**: Subscribes directly to `ws://127.0.0.1:8188/ws`. The WebSocket is the primary signal; a connection-pooled watchdog polls `/queue` at 1 Hz purely as a fallback, backing off to 3 s while the socket is healthy.
|
||||
@@ -96,6 +111,28 @@ only if the card actually needs it.
|
||||
* **Honest gain reporting**: gain against the profile's *current* setting is reported separately from the spread across values tried. Conflating them turns a flat result into a headline "+102%".
|
||||
* **Safety**: refuses to start while ComfyUI is executing, suspends the arbitrator's automatic profile switching for the duration (otherwise a diffusion benchmark trips `trigger_comfy_priority`, which reapplies the whole profile and overwrites the clock being measured), and restores the original profile in a `finally` block — including on exception or cancellation.
|
||||
|
||||
### 🩺 Dependency Self-Check (`health.py`)
|
||||
* `GET /api/health` verifies **everything this service depends on**: NVML, passwordless
|
||||
sudo for `nvidia-smi`, fan control through the headless X server, overclock drift, the
|
||||
telemetry store, residency-measurement capability, model directories, the ComfyUI
|
||||
WebSocket, and both upstream HTTP services.
|
||||
* Each check reports **what is broken, what that breaks, and how to fix it** — not just a
|
||||
red light. Shown on the dashboard as a badge that expands only when something is wrong.
|
||||
* It exists because fan control once failed for an entire session, recoverably and
|
||||
silently: the unit started before the X server that owns the GPU was accepting
|
||||
connections, the assignment failed with `Error resolving target specification 'gpu:0'`,
|
||||
nothing retried, and nothing ever asked whether fans worked. That failure now shows up
|
||||
in three places — a retry, a drift check, and this endpoint.
|
||||
|
||||
### 🧮 Honest VRAM Accounting
|
||||
* Processes are bucketed **ollama / comfy / desktop / unmanaged** rather than into one
|
||||
catch-all. On this machine a long-running `stt_relay.py` held 842 MB for three days
|
||||
while the compositor held 3.9 MB; a single "system" number reported them as one figure.
|
||||
* That distinction matters because ComfyUI's memory **can** be reclaimed and a third
|
||||
party's **cannot**. `unmanaged_gb` is headroom the arbitrator can never give back, so it
|
||||
is reported explicitly, shown on the dashboard, and named in the error when a
|
||||
reclaim-and-retry still cannot fit a model.
|
||||
|
||||
### 🗄️ Persistent Telemetry Store (`telemetry_store.py`)
|
||||
* Swap history used to be an in-memory `deque(maxlen=50)` that evaporated on every restart. Telemetry and events now persist to SQLite (WAL, single writer thread, batched 1 Hz inserts, automatic retention pruning) at roughly **0.4 MB per hour**.
|
||||
* This is what makes the app's central question answerable: **`GET /api/analytics/profiles` compares decode throughput per overclock profile**, joined against the thermals recorded while that profile was active.
|
||||
@@ -104,7 +141,7 @@ only if the card actually needs it.
|
||||
* **Live Hardware Telemetry**: GPU utilization %, GPU temperature (°C), power draw (W), fan speeds (%), and graphics/memory clock frequencies (MHz).
|
||||
* **Live Dual-Axis Time-Series Chart**: Real-time graphical visualization of VRAM usage (GB) and Host RAM Cache (GB) with zero frontend polling overhead.
|
||||
* **Interactive Control Center**: Trigger model hot-swaps, soft-yields, cache pre-warms, fan adjustments, and benchmarks directly from the web interface.
|
||||
* **Server-Sent Events (SSE)**: A single background sampler produces one 1 Hz snapshot and fans it out to every subscriber via `GET /api/stream`. Previously each connected client independently re-ran the whole snapshot — NVML, `/proc/meminfo`, an HTTP round-trip each to Ollama and ComfyUI, and a recursive walk of the ComfyUI models tree with a `stat()` per checkpoint — once per second, so opening the dashboard in three tabs tripled the load on the thing it was measuring. Slow clients drop stale frames instead of stalling the sampler.
|
||||
* **Server-Sent Events (SSE)**: A single background sampler produces one 1 Hz snapshot and fans it out to every subscriber via `GET /api/stream`. Frames are trimmed: the installed-model catalog was 81% of a 13.1 KB payload and changes only when a model is pulled, so it is sent on a subscriber's first frame and whenever it changes. Steady-state frames dropped 14041 → 3664 bytes (**74% smaller**; 135 → 38 MB/hour across three tabs), while `/api/stats` still returns the complete snapshot. Previously each connected client independently re-ran the whole snapshot — NVML, `/proc/meminfo`, an HTTP round-trip each to Ollama and ComfyUI, and a recursive walk of the ComfyUI models tree with a `stat()` per checkpoint — once per second, so opening the dashboard in three tabs tripled the load on the thing it was measuring. Slow clients drop stale frames instead of stalling the sampler.
|
||||
|
||||
### 🤖 Model Context Protocol (MCP 2.0) Server
|
||||
* **23 Native Agentic Tools**: Allows AI agents (Antigravity CLI, Claude Desktop, Cursor) to manage GPU resources, trigger model hot-swaps, measure page-cache residency, read persisted performance analytics, drive the thermal governor, and run overclock sweeps.
|
||||
@@ -116,6 +153,26 @@ only if the card actually needs it.
|
||||
|
||||
---
|
||||
|
||||
## 1a. Tests
|
||||
|
||||
```bash
|
||||
/home/drjones/comfy-mcp-venv/bin/python -m pytest tests/ -q # 182 passed in ~3.7s
|
||||
```
|
||||
|
||||
Hermetic: no GPU, no network, no sleeps. An autouse fixture stubs `overclock_manager._sh`
|
||||
— the single choke point for every `nvidia-smi`/`nvidia-settings` write — so no test can
|
||||
mutate the card, and `HYPERSWAP_DB` is redirected before `telemetry_store` imports.
|
||||
|
||||
The suite deliberately **pins empirically measured constants**, so that a future edit
|
||||
which contradicts the hardware fails loudly rather than silently:
|
||||
|
||||
| Pinned fact | Measured value | Why it is pinned |
|
||||
| :--- | :--- | :--- |
|
||||
| Warm model load | 12.87 GB in 4901 ms = 2.63 GB/s | The cache-hit threshold must stay below this, or no load can ever qualify |
|
||||
| Cold model load | 12.87 GB in 34267 ms = 0.38 GB/s | Separates a genuine cold read from a partial hit |
|
||||
| Busy yield | VRAM held at ≥50% GPU utilisation | A mid-generation model is finishing, not failing |
|
||||
| Residency confidence | probe trusted only at 100% | A 12-window probe once cleared 90% on a mostly-cold file |
|
||||
|
||||
## 2. Architectural Overview
|
||||
|
||||
```mermaid
|
||||
@@ -207,6 +264,7 @@ The HyperSwap server runs on port `9090` by default. Interactive OpenAPI/Swagger
|
||||
| `/api/analytics/models` | `GET` | Recency/frequency model ranking used to prioritise the warm budget. |
|
||||
| `/api/history?durable=true` | `GET` | Swap history from the persistent store rather than the in-memory ring. |
|
||||
| `/api/db` | `GET` | Store location, row counts and how many hours of history are held. |
|
||||
| `/api/health` | `GET` | Dependency self-check: NVML, sudo, fan control, drift, store, upstreams — each with impact and remediation. |
|
||||
|
||||
### Governor & Autotune Endpoints
|
||||
|
||||
|
||||
@@ -111,6 +111,10 @@ def _subsample(values: List[int], max_steps: int) -> List[int]:
|
||||
"""
|
||||
if len(values) <= max_steps:
|
||||
return values
|
||||
if max_steps < 2:
|
||||
# One step cannot span a range; take the top, which is what a caller asking for
|
||||
# a single sample almost certainly wants. Guards a ZeroDivisionError below.
|
||||
return values[-1:]
|
||||
idx = [round(i * (len(values) - 1) / (max_steps - 1)) for i in range(max_steps)]
|
||||
return sorted({values[i] for i in idx})
|
||||
|
||||
|
||||
188
health.py
Normal file
188
health.py
Normal file
@@ -0,0 +1,188 @@
|
||||
"""Dependency self-check.
|
||||
|
||||
Motivation: fan control failed for an entire session because the service started before
|
||||
the headless X server that owns the GPU was accepting connections. The failure was real,
|
||||
recoverable and completely invisible -- it appeared once, inside one field of one log
|
||||
line, and nothing ever asked "is fan control actually working?"
|
||||
|
||||
Everything HyperSwap needs is checked here, each with a plain statement of what breaks
|
||||
when it is missing and how to fix it. A degraded dependency should be loud.
|
||||
"""
|
||||
import asyncio
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from typing import Any, Dict, List
|
||||
|
||||
import httpx
|
||||
|
||||
import overclock_manager
|
||||
import ram_optimizer
|
||||
import telemetry_store
|
||||
import vram_arbitrator
|
||||
|
||||
logger = logging.getLogger("health")
|
||||
|
||||
OK, DEGRADED, FAILED = "ok", "degraded", "failed"
|
||||
|
||||
|
||||
def _check(name: str, status: str, detail: str, impact: str = "",
|
||||
fix: str = "") -> Dict[str, Any]:
|
||||
return {"name": name, "status": status, "detail": detail,
|
||||
"impact": impact, "fix": fix}
|
||||
|
||||
|
||||
def _check_nvml() -> Dict[str, Any]:
|
||||
if not vram_arbitrator.NVML_AVAILABLE:
|
||||
return _check("nvml", FAILED, "pynvml did not initialise",
|
||||
"No GPU telemetry, and VRAM yields cannot be confirmed",
|
||||
"Check the NVIDIA driver and that pynvml is installed in the venv")
|
||||
stats = vram_arbitrator.get_gpu_hardware_stats()
|
||||
if not stats.get("available"):
|
||||
return _check("nvml", FAILED, stats.get("error", "unavailable"),
|
||||
"No GPU telemetry", "Check the NVIDIA driver")
|
||||
return _check("nvml", OK, f"{stats.get('device_name')}, "
|
||||
f"{stats.get('vram_total_gb')} GB")
|
||||
|
||||
|
||||
def _check_sudo_smi() -> Dict[str, Any]:
|
||||
r = overclock_manager._smi("--query-gpu=name", "--format=csv,noheader")
|
||||
if r["rc"] != 0:
|
||||
return _check("nvidia-smi (sudo)", FAILED, r.get("err") or "non-zero exit",
|
||||
"Power limits and clock locks cannot be applied",
|
||||
"Passwordless sudo for /usr/bin/nvidia-smi is required "
|
||||
"(see /etc/sudoers.d/)")
|
||||
return _check("nvidia-smi (sudo)", OK, "passwordless sudo works")
|
||||
|
||||
|
||||
def _check_fan_control() -> Dict[str, Any]:
|
||||
"""The check that would have caught the startup race."""
|
||||
if not overclock_manager.is_headless_x_running():
|
||||
return _check("fan control", FAILED,
|
||||
f"no X server found on {overclock_manager.HEADLESS_DISPLAY}",
|
||||
"Fan speed cannot be read or set; the thermal governor cannot "
|
||||
"raise the fan floor when the card gets hot",
|
||||
f"Start the headless X server on {overclock_manager.HEADLESS_DISPLAY}")
|
||||
status = overclock_manager.get_fan_status(force=True)
|
||||
if status.get("target_speed_pct") is None and not status.get("manual"):
|
||||
# Auto mode legitimately reports no target; probe the control attribute instead.
|
||||
probe = overclock_manager._nvidia_settings("-q", "[gpu:0]/GPUFanControlState")
|
||||
if probe.get("rc") != 0 or overclock_manager._fan_target_missing(probe):
|
||||
return _check("fan control", FAILED,
|
||||
probe.get("err") or "GPU target not resolvable",
|
||||
"Fan control unavailable; the governor cannot cool the card",
|
||||
"Check Coolbits and that X on "
|
||||
f"{overclock_manager.HEADLESS_DISPLAY} owns the GPU")
|
||||
return _check("fan control", OK, f"mode={status.get('mode')}")
|
||||
|
||||
|
||||
def _check_profile_drift() -> Dict[str, Any]:
|
||||
drift = overclock_manager.profile_drift()
|
||||
if not drift.get("applied_since_start"):
|
||||
return _check("overclock profile", DEGRADED,
|
||||
"no profile has been successfully applied since startup",
|
||||
"The card may not be running the settings this app reports",
|
||||
"Apply a profile, or check the nvidia-smi/fan checks above")
|
||||
if drift.get("drifted"):
|
||||
return _check("overclock profile", DEGRADED, drift.get("reason", "drifted"),
|
||||
"Reported settings do not match the hardware",
|
||||
"The sampler reconciles once a minute; POST /api/overclock/apply "
|
||||
"to force it now")
|
||||
return _check("overclock profile", OK,
|
||||
f"{drift['profile']} @ {drift['power_limit_actual_w']}W")
|
||||
|
||||
|
||||
async def _check_http(name: str, url: str, impact: str, fix: str) -> Dict[str, Any]:
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=3.0) as c:
|
||||
r = await c.get(url)
|
||||
if r.status_code == 200:
|
||||
return _check(name, OK, "reachable")
|
||||
return _check(name, DEGRADED, f"HTTP {r.status_code}", impact, fix)
|
||||
except Exception as e:
|
||||
return _check(name, FAILED, str(e)[:120], impact, fix)
|
||||
|
||||
|
||||
def _check_comfy_ws() -> Dict[str, Any]:
|
||||
arb = vram_arbitrator.arbitrator
|
||||
if not arb.running:
|
||||
return _check("arbitrator", FAILED, "background engine not running",
|
||||
"No automatic VRAM handoff between Ollama and ComfyUI",
|
||||
"Restart the service")
|
||||
if not arb.connected_ws:
|
||||
return _check("comfyui websocket", DEGRADED, "not connected",
|
||||
"Falling back to 1 Hz polling; handoffs react more slowly",
|
||||
"Check that ComfyUI is running and reachable on :8188")
|
||||
return _check("comfyui websocket", OK, "subscribed")
|
||||
|
||||
|
||||
def _check_store() -> Dict[str, Any]:
|
||||
info = telemetry_store.db_info()
|
||||
if not info.get("exists"):
|
||||
return _check("telemetry store", DEGRADED, "database not created yet",
|
||||
"No persisted history, so profile comparison cannot be computed",
|
||||
"It is created on first write; check the directory is writable")
|
||||
if not os.access(telemetry_store.DB_PATH, os.W_OK):
|
||||
return _check("telemetry store", FAILED, "database not writable",
|
||||
"Telemetry and swap events are being dropped",
|
||||
f"Check permissions on {telemetry_store.DB_PATH}")
|
||||
return _check("telemetry store", OK,
|
||||
f"{info.get('size_mb')} MB, {info.get('coverage_hours')} h of history")
|
||||
|
||||
|
||||
def _check_residency() -> Dict[str, Any]:
|
||||
cap = ram_optimizer.residency_capability()
|
||||
if cap.get("exact_everywhere"):
|
||||
return _check("residency measurement", OK, "cachestat available for all models")
|
||||
return _check("residency measurement", DEGRADED, cap.get("reason", ""),
|
||||
"Ollama weight residency is estimated by read-rate probe, not measured",
|
||||
cap.get("hint", ""))
|
||||
|
||||
|
||||
def _check_model_dirs() -> Dict[str, Any]:
|
||||
comfy_dir = ram_optimizer.COMFY_MODELS_DIR
|
||||
if not os.path.isdir(comfy_dir):
|
||||
return _check("model directories", DEGRADED,
|
||||
f"ComfyUI model directory not found: {comfy_dir}",
|
||||
"ComfyUI checkpoints cannot be catalogued or pre-warmed",
|
||||
"Set HYPERSWAP_COMFY_MODELS to the right path")
|
||||
catalog = ram_optimizer.get_model_catalog()
|
||||
return _check("model directories", OK,
|
||||
f"{len(catalog['ollama'])} Ollama blobs, {len(catalog['comfy'])} ComfyUI files")
|
||||
|
||||
|
||||
async def run_health_checks() -> Dict[str, Any]:
|
||||
"""Run every dependency check. Never raises."""
|
||||
t0 = time.perf_counter()
|
||||
loop = asyncio.get_running_loop()
|
||||
|
||||
sync_checks = [_check_nvml, _check_sudo_smi, _check_fan_control,
|
||||
_check_profile_drift, _check_store, _check_residency,
|
||||
_check_model_dirs, _check_comfy_ws]
|
||||
results: List[Dict[str, Any]] = []
|
||||
for fn in sync_checks:
|
||||
try:
|
||||
results.append(await loop.run_in_executor(None, fn))
|
||||
except Exception as e:
|
||||
results.append(_check(fn.__name__, FAILED, f"check raised: {e}"))
|
||||
|
||||
results.extend(await asyncio.gather(
|
||||
_check_http("ollama", f"{vram_arbitrator.OLLAMA_API_BASE}/api/tags",
|
||||
"No LLM orchestration", "Start the ollama service"),
|
||||
_check_http("comfyui", f"{vram_arbitrator.COMFY_API_BASE}/system_stats",
|
||||
"No diffusion arbitration", "Start ComfyUI on :8188"),
|
||||
))
|
||||
|
||||
failed = [r for r in results if r["status"] == FAILED]
|
||||
degraded = [r for r in results if r["status"] == DEGRADED]
|
||||
overall = FAILED if failed else (DEGRADED if degraded else OK)
|
||||
return {
|
||||
"status": overall,
|
||||
"checked_at": time.time(),
|
||||
"duration_ms": round((time.perf_counter() - t0) * 1000, 1),
|
||||
"summary": (f"{len(results) - len(failed) - len(degraded)} ok, "
|
||||
f"{len(degraded)} degraded, {len(failed)} failed"),
|
||||
"failed": [r["name"] for r in failed],
|
||||
"degraded": [r["name"] for r in degraded],
|
||||
"checks": results,
|
||||
}
|
||||
@@ -8,6 +8,7 @@ from typing import Dict, List, Any, Optional
|
||||
|
||||
from mcp.server import MCPServer
|
||||
import autotune
|
||||
import health
|
||||
import overclock_manager
|
||||
import ram_optimizer
|
||||
import telemetry_store
|
||||
@@ -136,6 +137,14 @@ def set_gpu_fan_speed(mode: str = "auto", percent: Optional[int] = None) -> str:
|
||||
res = overclock_manager.set_fan_auto()
|
||||
return json.dumps(res, indent=2)
|
||||
|
||||
@mcp.tool()
|
||||
async def check_system_health() -> str:
|
||||
"""Check every dependency HyperSwap needs (NVML, sudo nvidia-smi, fan control via the
|
||||
headless X server, Ollama, ComfyUI, the telemetry store, model directories) and report
|
||||
what is broken, what it breaks, and how to fix it."""
|
||||
return json.dumps(await health.run_health_checks(), indent=2, default=str)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def get_page_cache_residency(include_files: bool = True) -> str:
|
||||
"""Measure how much of each model on disk is genuinely resident in the Linux page cache.
|
||||
|
||||
@@ -18,6 +18,7 @@ import logging
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import time
|
||||
from typing import Dict, Any, Optional, List
|
||||
|
||||
logger = logging.getLogger("overclock_manager")
|
||||
@@ -278,12 +279,40 @@ def apply_profile(name: str, overrides: Optional[Dict[str, Any]] = None) -> Dict
|
||||
"detail": "skipped: this driver accepts clock offsets and ignores them"}),
|
||||
"fan": apply_fan_control(fan_mode, fan_speed),
|
||||
}
|
||||
result["gpu"] = get_gpu_state()
|
||||
result["fan_status"] = get_fan_status()
|
||||
result["overrides"] = overrides or {}
|
||||
|
||||
# Invalidate BEFORE reading back. These caches exist so the dashboard's polling does
|
||||
# not fork sudo every few seconds, but reading through them here reported the
|
||||
# pre-change value: a profile that had just moved the card 370W -> 320W came back
|
||||
# claiming 370W, so the API contradicted nvidia-smi.
|
||||
_STATE_CACHE["value"] = None
|
||||
_FAN_CACHE["value"] = None
|
||||
result["gpu"] = get_gpu_state(force=True)
|
||||
result["fan_status"] = get_fan_status(force=True)
|
||||
result["overrides"] = overrides or {}
|
||||
|
||||
# Say plainly whether the card ended up where the profile asked it to.
|
||||
intended_w = int(cfg.get("power_limit_w", 0))
|
||||
actual_w = result["gpu"].get("power_limit_w")
|
||||
result["verified"] = {
|
||||
"power_limit_intended_w": intended_w,
|
||||
"power_limit_actual_w": actual_w,
|
||||
"power_limit_ok": (actual_w is None or intended_w == 0
|
||||
or abs(float(actual_w) - intended_w) < 1.0),
|
||||
"fan_mode_intended": fan_mode,
|
||||
"fan_mode_actual": result["fan_status"].get("mode"),
|
||||
"fan_ok": result["fan"].get("applied", False),
|
||||
}
|
||||
if not result["verified"]["power_limit_ok"]:
|
||||
logger.warning(f"Profile '{name}' asked for {intended_w}W but the card reports "
|
||||
f"{actual_w}W")
|
||||
if not result["verified"]["fan_ok"]:
|
||||
logger.warning(f"Profile '{name}' could not set fans: "
|
||||
f"{result['fan'].get('detail')}")
|
||||
|
||||
global _APPLIED_ONCE, _FAN_AVAILABLE
|
||||
if result["verified"]["power_limit_ok"]:
|
||||
_APPLIED_ONCE = True
|
||||
if result["verified"]["fan_ok"]:
|
||||
_FAN_AVAILABLE = True
|
||||
ACTIVE_PROFILE = name
|
||||
_LAST_RESULT = result
|
||||
logger.info(f"Overclock profile applied: {name} -> {json.dumps(result, default=str)}")
|
||||
@@ -332,14 +361,39 @@ def is_headless_x_running() -> bool:
|
||||
return r["rc"] == 0
|
||||
|
||||
|
||||
def apply_fan_control(mode: str, speed_pct: int) -> Dict[str, Any]:
|
||||
FAN_RETRY_ATTEMPTS = 5
|
||||
FAN_RETRY_DELAY_S = 2.0
|
||||
|
||||
|
||||
def _fan_target_missing(result: Dict[str, Any]) -> bool:
|
||||
"""True when nvidia-settings could not see the GPU at all.
|
||||
|
||||
On boot this service can start before the headless X server on :8 that owns the GPU
|
||||
is accepting connections, and the fan assignment fails with 'Error resolving target
|
||||
specification'. Nothing retried, so the fans were simply never configured for the
|
||||
whole session and the failure was only visible deep in a log line.
|
||||
"""
|
||||
text = ((result.get("err") or "") + (result.get("out") or "")).lower()
|
||||
return ("error resolving target" in text or "no targets match" in text
|
||||
or "cannot open display" in text)
|
||||
|
||||
|
||||
def apply_fan_control(mode: str, speed_pct: int, _attempt: int = 0) -> Dict[str, Any]:
|
||||
global FAN_MANUAL
|
||||
if mode == "auto":
|
||||
r = _nvidia_settings("-a", "[gpu:0]/GPUFanControlState=0")
|
||||
ok = r["rc"] == 0
|
||||
if not ok and _fan_target_missing(r) and _attempt < FAN_RETRY_ATTEMPTS:
|
||||
logger.info(f"Fan control target not ready (attempt {_attempt + 1}/"
|
||||
f"{FAN_RETRY_ATTEMPTS}); X on {HEADLESS_DISPLAY} may still be "
|
||||
f"starting — retrying in {FAN_RETRY_DELAY_S}s")
|
||||
time.sleep(FAN_RETRY_DELAY_S)
|
||||
return apply_fan_control(mode, speed_pct, _attempt + 1)
|
||||
if ok:
|
||||
FAN_MANUAL = False
|
||||
return {"applied": ok, "mode": "auto", "detail": r.get("out") or r.get("err")}
|
||||
_FAN_CACHE["value"] = None
|
||||
return {"applied": ok, "mode": "auto", "detail": r.get("out") or r.get("err"),
|
||||
"attempts": _attempt + 1}
|
||||
|
||||
speed_pct = max(30, min(100, int(speed_pct)))
|
||||
r = _nvidia_settings(
|
||||
@@ -422,10 +476,76 @@ def restore_safe(reason: str = "shutdown") -> Dict[str, Any]:
|
||||
return result
|
||||
|
||||
|
||||
_APPLIED_ONCE = False
|
||||
# Set once fan control has worked at least once, so drift checks do not fire forever on
|
||||
# a machine that simply has no fan control available.
|
||||
_FAN_AVAILABLE = False
|
||||
|
||||
|
||||
def profile_drift() -> Dict[str, Any]:
|
||||
"""Compare what the active profile asks for against what the card actually reports.
|
||||
|
||||
ACTIVE_PROFILE defaults to "balanced" at import, which is indistinguishable from
|
||||
"balanced was successfully applied" -- so a startup apply that failed left the app
|
||||
confidently reporting a profile it had never put on the hardware. This makes the
|
||||
difference visible instead.
|
||||
"""
|
||||
profiles = load_profiles()
|
||||
cfg = profiles.get(ACTIVE_PROFILE, {})
|
||||
state = get_gpu_state()
|
||||
intended = int(cfg.get("power_limit_w", 0) or 0)
|
||||
actual = state.get("power_limit_w")
|
||||
power_drift = bool(intended and actual is not None
|
||||
and abs(float(actual) - intended) >= 1.0)
|
||||
|
||||
# Fan mode is checked too. The headless X server that owns the GPU can still be
|
||||
# starting when this unit does, and the fan assignment then fails; the in-process
|
||||
# retries cover a short delay, but if X arrives later nothing else would ever notice
|
||||
# that the profile's fan setting was never applied.
|
||||
fan_intended = cfg.get("fan_mode", "auto")
|
||||
fan_actual = None
|
||||
fan_drift = False
|
||||
if _FAN_AVAILABLE:
|
||||
fan_actual = get_fan_status().get("mode")
|
||||
fan_drift = bool(fan_actual and fan_actual != fan_intended)
|
||||
|
||||
drifted = power_drift or fan_drift or not _APPLIED_ONCE
|
||||
reasons = []
|
||||
if not _APPLIED_ONCE:
|
||||
reasons.append("no profile has been successfully applied since startup")
|
||||
if power_drift:
|
||||
reasons.append(f"card reports {actual}W, profile asks {intended}W")
|
||||
if fan_drift:
|
||||
reasons.append(f"fans are {fan_actual}, profile asks {fan_intended}")
|
||||
return {
|
||||
"profile": ACTIVE_PROFILE,
|
||||
"applied_since_start": _APPLIED_ONCE,
|
||||
"power_limit_intended_w": intended,
|
||||
"power_limit_actual_w": actual,
|
||||
"fan_mode_intended": fan_intended,
|
||||
"fan_mode_actual": fan_actual,
|
||||
"fan_available": _FAN_AVAILABLE,
|
||||
"drifted": drifted,
|
||||
"reason": "; ".join(reasons) or None,
|
||||
}
|
||||
|
||||
|
||||
def reconcile_profile() -> Dict[str, Any]:
|
||||
"""Re-apply the active profile if the hardware has drifted away from it."""
|
||||
drift = profile_drift()
|
||||
if not drift["drifted"]:
|
||||
return {"reconciled": False, "drift": drift}
|
||||
logger.warning(f"Overclock drift detected — re-applying '{ACTIVE_PROFILE}': "
|
||||
f"{drift['reason']}")
|
||||
res = apply_profile(ACTIVE_PROFILE)
|
||||
return {"reconciled": True, "drift": drift, "result": res.get("verified")}
|
||||
|
||||
|
||||
def get_status() -> Dict[str, Any]:
|
||||
"""Full overclock status for the dashboard."""
|
||||
return {
|
||||
"active_profile": ACTIVE_PROFILE,
|
||||
"drift": profile_drift(),
|
||||
"offsets_supported": offsets_supported(),
|
||||
"effective_levers": (["power_limit", "clock_lock", "mem_lock", "fan"]
|
||||
+ (["offsets"] if offsets_supported() else [])),
|
||||
|
||||
5
pytest.ini
Normal file
5
pytest.ini
Normal file
@@ -0,0 +1,5 @@
|
||||
[pytest]
|
||||
testpaths = tests
|
||||
addopts = -ra
|
||||
filterwarnings =
|
||||
ignore::DeprecationWarning
|
||||
@@ -203,8 +203,10 @@ def page_residency(filepath: str, allow_probe: bool = True,
|
||||
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"}
|
||||
# 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)
|
||||
@@ -415,6 +417,16 @@ _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.
|
||||
|
||||
@@ -510,6 +522,7 @@ def warm_file_to_ram(filepath: str, chunk_size: int = 16 * 1024 * 1024,
|
||||
|
||||
duration = time.perf_counter() - t0
|
||||
after = page_residency(filepath, probe_windows=32)
|
||||
invalidate_cache_report()
|
||||
return {
|
||||
"success": True,
|
||||
"filepath": filepath,
|
||||
|
||||
62
server.py
62
server.py
@@ -16,6 +16,7 @@ from fastapi.middleware.cors import CORSMiddleware
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
import autotune
|
||||
import health
|
||||
import overclock_manager
|
||||
import ram_optimizer
|
||||
import telemetry_store
|
||||
@@ -49,6 +50,8 @@ class TelemetryBroker:
|
||||
layer, so neither needs to poll the GPU on its own.
|
||||
"""
|
||||
|
||||
RECONCILE_EVERY_N = 60 # once a minute at 1 Hz
|
||||
|
||||
def __init__(self, interval_s: float = 1.0) -> None:
|
||||
self.interval_s = interval_s
|
||||
self.snapshot: Dict[str, Any] = {}
|
||||
@@ -57,6 +60,7 @@ class TelemetryBroker:
|
||||
self.running = False
|
||||
self.samples = 0
|
||||
self.last_sample_ms = 0.0
|
||||
self._models_fp = None
|
||||
# Set on shutdown so open SSE generators finish instead of holding the server up.
|
||||
self.closing = False
|
||||
|
||||
@@ -87,6 +91,26 @@ class TelemetryBroker:
|
||||
with contextlib.suppress(asyncio.CancelledError):
|
||||
await self.task
|
||||
|
||||
def _stream_frame(self, snap: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Trim the snapshot for streaming.
|
||||
|
||||
The installed-model catalog is 10.6 KB of a 13.1 KB payload -- 81% -- and it
|
||||
changes only when a model is pulled or removed, yet it was re-sent to every
|
||||
subscriber every second (135 MB/hour across three tabs). It is sent on the first
|
||||
frame and whenever it changes; otherwise the client keeps what it has.
|
||||
|
||||
/api/stats still returns the complete snapshot, so API consumers are unaffected.
|
||||
"""
|
||||
ollama = snap.get("ollama", {})
|
||||
models = ollama.get("installed_models") or []
|
||||
fp = hash(tuple(sorted(m.get("name", "") for m in models)))
|
||||
if fp == self._models_fp:
|
||||
trimmed_ollama = {k: v for k, v in ollama.items() if k != "installed_models"}
|
||||
trimmed_ollama["installed_models_unchanged"] = True
|
||||
return {**snap, "ollama": trimmed_ollama}
|
||||
self._models_fp = fp
|
||||
return snap
|
||||
|
||||
def subscribe(self) -> asyncio.Queue:
|
||||
q: asyncio.Queue = asyncio.Queue(maxsize=2)
|
||||
self.subscribers.add(q)
|
||||
@@ -107,19 +131,27 @@ class TelemetryBroker:
|
||||
# Feed the governor and the durable store from the sample we already have.
|
||||
thermal_governor.governor.observe(snap.get("gpu", {}),
|
||||
overclock_manager.ACTIVE_PROFILE)
|
||||
|
||||
# Cheap, infrequent check that the card still matches the active profile.
|
||||
# A startup apply can fail silently (the headless X server may not be up
|
||||
# yet), and an external tool can move the power limit underneath us.
|
||||
if self.samples % self.RECONCILE_EVERY_N == 0:
|
||||
await asyncio.get_running_loop().run_in_executor(
|
||||
None, overclock_manager.reconcile_profile)
|
||||
telemetry_store.record_telemetry(
|
||||
snap.get("gpu", {}), snap.get("ram", {}),
|
||||
profile=overclock_manager.ACTIVE_PROFILE,
|
||||
throttle_reasons=",".join(snap.get("gpu", {}).get("throttle_reasons") or []),
|
||||
)
|
||||
|
||||
frame = self._stream_frame(snap)
|
||||
for q in list(self.subscribers):
|
||||
if q.full():
|
||||
# Slow client: drop the stale frame rather than stalling the sampler.
|
||||
with contextlib.suppress(asyncio.QueueEmpty):
|
||||
q.get_nowait()
|
||||
with contextlib.suppress(asyncio.QueueFull):
|
||||
q.put_nowait(snap)
|
||||
q.put_nowait(frame)
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as e:
|
||||
@@ -279,6 +311,16 @@ async def get_all_stats() -> Dict[str, Any]:
|
||||
return await broker.get()
|
||||
|
||||
|
||||
@app.get("/api/health", summary="Dependency Self-Check", tags=["Telemetry"])
|
||||
async def api_health():
|
||||
"""Check everything HyperSwap depends on, with impact and remediation for each.
|
||||
|
||||
Returns overall `status` of ok | degraded | failed. Exists because fan control once
|
||||
failed for a whole session -- recoverably, and completely silently.
|
||||
"""
|
||||
return await health.run_health_checks()
|
||||
|
||||
|
||||
@app.get("/api/gpu", summary="GPU Sensors and VRAM Breakdown", tags=["Telemetry"])
|
||||
async def get_gpu_metrics() -> Dict[str, Any]:
|
||||
"""Detailed NVML sensors (utilization, temp, power, fan, clocks, throttle reasons, per-process VRAM)."""
|
||||
@@ -296,6 +338,7 @@ async def sse_telemetry_stream(request: Request):
|
||||
q = broker.subscribe()
|
||||
try:
|
||||
snap = await broker.get()
|
||||
# Full snapshot first: a new subscriber has no cached catalog yet.
|
||||
yield f"data: {json.dumps(snap)}\n\n"
|
||||
while not broker.closing:
|
||||
if await request.is_disconnected():
|
||||
@@ -336,9 +379,20 @@ async def api_switch_model(req: SwitchRequest):
|
||||
return res
|
||||
|
||||
@app.post("/api/free-vram", summary="Soft-Yield Ollama VRAM", tags=["Orchestration"])
|
||||
async def api_free_vram(confirm: bool = Query(True, description="Wait for the driver to actually release the allocation")):
|
||||
"""Yield Ollama's VRAM and wait for the release to be confirmed by NVML."""
|
||||
return await vram_arbitrator.instant_free_ollama_vram(confirm=confirm)
|
||||
async def api_free_vram(
|
||||
confirm: bool = Query(True, description="Wait for the driver to actually release the allocation"),
|
||||
wait_for_generation: bool = Query(False, description="Block until an in-flight generation finishes, instead of returning 'busy'"),
|
||||
):
|
||||
"""Yield Ollama's VRAM and wait for NVML to confirm the release.
|
||||
|
||||
Returns `outcome`: "released", "busy" (model is mid-generation; the unload is queued
|
||||
and will apply when it finishes) or "stuck" (VRAM held with an idle GPU).
|
||||
"""
|
||||
return await vram_arbitrator.instant_free_ollama_vram(
|
||||
confirm=confirm,
|
||||
timeout_s=(vram_arbitrator.YIELD_CONFIRM_TIMEOUT_BLOCKING_S
|
||||
if wait_for_generation else None),
|
||||
)
|
||||
|
||||
@app.post("/api/comfy-free", summary="Purge ComfyUI VRAM Cache", tags=["Orchestration"])
|
||||
async def api_comfy_free():
|
||||
|
||||
200
static/app.js
200
static/app.js
@@ -33,8 +33,9 @@ function initSSE() {
|
||||
function updateDashboard(data) {
|
||||
if (!data) return;
|
||||
|
||||
// Governor state rides along in the shared snapshot — no extra polling needed.
|
||||
// Governor and arbitration state ride along in the shared snapshot.
|
||||
if (data.governor) renderGovernor(data.governor);
|
||||
if (data.arbitrator) renderArbitrator(data.arbitrator, data.gpu);
|
||||
|
||||
// 1. GPU VRAM Stats
|
||||
const gpu = data.gpu || {};
|
||||
@@ -153,6 +154,29 @@ function updateDashboard(data) {
|
||||
|
||||
document.getElementById('legend-ram-used').textContent = `${ram.used_gb || 0} GB`;
|
||||
document.getElementById('legend-ram-cached').textContent = `${ram.cached_gb || 0} GB`;
|
||||
|
||||
// Split the old catch-all: desktop compositor vs VRAM we cannot reclaim.
|
||||
const vbd = gpu.breakdown || {};
|
||||
const unmanagedEl = document.getElementById('legend-vram-unmanaged');
|
||||
if (unmanagedEl) unmanagedEl.textContent = `${vbd.unmanaged_gb ?? 0} GB`;
|
||||
const sysEl = document.getElementById('legend-system');
|
||||
if (sysEl) sysEl.textContent = `${vbd.desktop_gb ?? 0} GB`;
|
||||
|
||||
// Measured ComfyUI facts, replacing the panel's previously hardcoded ones.
|
||||
const held = document.getElementById('comfy-vram-held');
|
||||
if (held) held.textContent = `${vbd.comfyui_gb ?? 0} GB`;
|
||||
const pcie = document.getElementById('comfy-pcie');
|
||||
if (pcie) {
|
||||
const tx = (gpu.pcie_tx_kbps || 0) / 1024, rx = (gpu.pcie_rx_kbps || 0) / 1024;
|
||||
pcie.textContent = `${tx.toFixed(1)} / ${rx.toFixed(1)} MB/s`;
|
||||
}
|
||||
const purge = document.getElementById('comfy-purge-state');
|
||||
if (purge && data.arbitrator) {
|
||||
const a = data.arbitrator;
|
||||
purge.textContent = a.pending_purge && a.comfy_idle_s != null
|
||||
? `in ${Math.max((a.idle_purge_after_s || 0) - a.comfy_idle_s, 0).toFixed(0)}s`
|
||||
: (a.comfy_active ? 'generating' : 'idle');
|
||||
}
|
||||
document.getElementById('legend-ram-free').textContent = `${ram.free_gb || 0} GB`;
|
||||
}
|
||||
|
||||
@@ -176,9 +200,13 @@ function updateDashboard(data) {
|
||||
document.getElementById('ollama-context').textContent = 'Idle';
|
||||
}
|
||||
|
||||
// The catalog is omitted from stream frames when unchanged, so keep the last one.
|
||||
if (ollama.installed_models && ollama.installed_models.length > 0) {
|
||||
document.getElementById('ollama-total-models').textContent = ollama.installed_models.length;
|
||||
updateModelSelect(ollama.installed_models, ollama.active_model_name);
|
||||
currentInstalledModels = ollama.installed_models;
|
||||
}
|
||||
if (currentInstalledModels && currentInstalledModels.length > 0) {
|
||||
document.getElementById('ollama-total-models').textContent = currentInstalledModels.length;
|
||||
updateModelSelect(currentInstalledModels, ollama.active_model_name);
|
||||
}
|
||||
} else {
|
||||
document.getElementById('ollama-status-text').textContent = 'OFFLINE';
|
||||
@@ -220,9 +248,14 @@ function updateDashboard(data) {
|
||||
const history = data.history || [];
|
||||
const logContainer = document.getElementById('switch-log-container');
|
||||
if (history.length > 0) {
|
||||
const latest = history[0];
|
||||
document.getElementById('ollama-last-swap').textContent = `${latest.duration_ms} ms`;
|
||||
document.getElementById('ollama-cache-hit').textContent = latest.cache_status || 'OK';
|
||||
// These tiles are labelled "last swap time" and "RAM hit status", so they must come
|
||||
// from an actual model swap. Reading history[0] meant a ComfyUI purge showed its
|
||||
// 1.65 ms duration as the swap time and its "Cleaned" status as a cache result.
|
||||
const lastSwap = history.find(h => h.event_type === 'LLM Model Switch');
|
||||
document.getElementById('ollama-last-swap').textContent =
|
||||
lastSwap ? `${lastSwap.duration_ms} ms` : '\u2014';
|
||||
document.getElementById('ollama-cache-hit').textContent =
|
||||
lastSwap ? (lastSwap.cache_status || 'OK') : 'no swaps yet';
|
||||
|
||||
logContainer.innerHTML = history.slice(0, 10).map(item => {
|
||||
const isHit = (item.cache_status || '').includes('RAM Cache Hit') || (item.cache_status || '').includes('RAM-Cached');
|
||||
@@ -881,3 +914,158 @@ document.addEventListener('DOMContentLoaded', () => {
|
||||
if (d.last_result) renderSweep(d.last_result);
|
||||
}).catch(() => {});
|
||||
});
|
||||
|
||||
|
||||
// ---------------------------------------------------------------- arbitration
|
||||
|
||||
function renderArbitrator(arb, gpu) {
|
||||
const el = (id) => document.getElementById(id);
|
||||
if (!el('arb-action')) return;
|
||||
const c = arb.counters || {};
|
||||
|
||||
// VRAM held by processes HyperSwap cannot reclaim. Worth showing: it is headroom the
|
||||
// arbitrator can never give back, no matter how much it purges.
|
||||
const bd = (gpu && gpu.breakdown) || {};
|
||||
const un = el('arb-unmanaged');
|
||||
if (un) {
|
||||
const procs = bd.unmanaged || [];
|
||||
un.innerHTML = procs.length
|
||||
? `<span class="text-amber-400">${bd.unmanaged_gb} GB unreclaimable</span> — ` +
|
||||
procs.map(p => `${p.name} (${p.vram_mb} MB)`).join(', ')
|
||||
: '';
|
||||
}
|
||||
|
||||
el('arb-action').textContent = arb.last_action || 'Idle';
|
||||
el('arb-yields').textContent = c.yields ?? 0;
|
||||
el('arb-busy').textContent = c.yield_deferred_busy ?? 0;
|
||||
el('arb-later').textContent = c.deferred_releases ?? 0;
|
||||
el('arb-stalled').textContent = c.yield_stalled ?? 0;
|
||||
el('arb-purges').textContent = c.purges ?? 0;
|
||||
el('arb-defpurge').textContent = c.deferred_purges ?? 0;
|
||||
|
||||
const ws = el('arb-ws');
|
||||
ws.textContent = arb.connected_ws ? 'ComfyUI WS live' : 'WS down — polling';
|
||||
ws.className = 'text-xs font-mono ' + (arb.connected_ws ? 'text-emerald-400' : 'text-amber-400');
|
||||
|
||||
// Show why we are holding off, and the idle countdown before ComfyUI is purged.
|
||||
const parts = [];
|
||||
const backoff = arb.yield_backoff || {};
|
||||
for (const [model, secs] of Object.entries(backoff)) {
|
||||
parts.push(`waiting ${secs}s before asking '${model}' again`);
|
||||
}
|
||||
if (arb.pending_purge && arb.comfy_idle_s != null) {
|
||||
const left = Math.max((arb.idle_purge_after_s || 0) - arb.comfy_idle_s, 0).toFixed(0);
|
||||
parts.push(`ComfyUI idle ${arb.comfy_idle_s}s — holding its checkpoints ${left}s longer`);
|
||||
}
|
||||
el('arb-backoff').textContent = parts.join(' · ');
|
||||
}
|
||||
|
||||
|
||||
// ---------------------------------------------------------------- health
|
||||
|
||||
async function fetchHealth(verbose = false) {
|
||||
const badge = document.getElementById('health-badge');
|
||||
const body = document.getElementById('health-body');
|
||||
if (!badge) return;
|
||||
try {
|
||||
const d = await (await fetch('/api/health')).json();
|
||||
const style = {
|
||||
ok: 'bg-emerald-950/70 border border-emerald-800 text-emerald-300',
|
||||
degraded: 'bg-amber-950/70 border border-amber-800 text-amber-300',
|
||||
failed: 'bg-rose-950/70 border border-rose-800 text-rose-300',
|
||||
}[d.status];
|
||||
badge.className = `px-2 py-1 text-xs font-bold rounded-lg ${style}`;
|
||||
badge.textContent = `${d.status.toUpperCase()} · ${d.summary}`;
|
||||
|
||||
// Healthy checks collapse to one line; anything wrong is shown in full with the
|
||||
// impact and the fix, because that is the information you actually need.
|
||||
const problems = d.checks.filter(c => c.status !== 'ok');
|
||||
const shown = (verbose || problems.length) ? (verbose ? d.checks : problems) : [];
|
||||
body.innerHTML = shown.map(c => {
|
||||
const colour = c.status === 'ok' ? 'text-emerald-400'
|
||||
: c.status === 'degraded' ? 'text-amber-400' : 'text-rose-400';
|
||||
const mark = c.status === 'ok' ? '✓' : c.status === 'degraded' ? '!' : '✗';
|
||||
let html = `<div><span class="${colour}">${mark} ${c.name}</span>` +
|
||||
`<span class="text-slate-400"> — ${c.detail}</span></div>`;
|
||||
if (c.status !== 'ok') {
|
||||
if (c.impact) html += `<div class="text-slate-500 pl-4">impact: ${c.impact}</div>`;
|
||||
if (c.fix) html += `<div class="text-slate-500 pl-4">fix: ${c.fix}</div>`;
|
||||
}
|
||||
return html;
|
||||
}).join('') || '<div class="text-emerald-500">All dependencies healthy.</div>';
|
||||
} catch (e) {
|
||||
badge.className = 'px-2 py-1 text-xs font-bold rounded-lg bg-rose-950/70 border border-rose-800 text-rose-300';
|
||||
badge.textContent = 'UNREACHABLE';
|
||||
body.innerHTML = `<div class="text-rose-400">${e}</div>`;
|
||||
}
|
||||
}
|
||||
|
||||
document.addEventListener('DOMContentLoaded', () => {
|
||||
fetchHealth(false);
|
||||
setInterval(() => fetchHealth(false), 60000);
|
||||
});
|
||||
|
||||
|
||||
// ---------------------------------------------------------------- honest labels
|
||||
|
||||
async function fetchMeasuredResidency() {
|
||||
const el = document.getElementById('legend-models-resident');
|
||||
if (!el) return;
|
||||
try {
|
||||
const d = await (await fetch('/api/cache/report?files=false')).json();
|
||||
// Deliberately distinct from the page-cache figure beside it: that is every file
|
||||
// the kernel has cached; this is model weights, measured.
|
||||
el.textContent = `${d.resident_total_gb} GB of ${d.catalog_total_gb} GB`;
|
||||
const cr = document.getElementById('comfy-resident');
|
||||
if (cr) cr.textContent = `${d.residency_pct}%`;
|
||||
} catch (e) {
|
||||
el.textContent = 'unavailable';
|
||||
}
|
||||
}
|
||||
|
||||
async function fetchVersion() {
|
||||
const el = document.getElementById('app-version');
|
||||
if (!el) return;
|
||||
try {
|
||||
const d = await (await fetch('/openapi.json')).json();
|
||||
el.textContent = `v${d.info.version}`;
|
||||
} catch (e) {
|
||||
el.textContent = 'v?';
|
||||
}
|
||||
}
|
||||
|
||||
document.addEventListener('DOMContentLoaded', () => {
|
||||
fetchVersion();
|
||||
fetchMeasuredResidency();
|
||||
setInterval(fetchMeasuredResidency, 30000);
|
||||
});
|
||||
|
||||
|
||||
// ---------------------------------------------------------------- durable swap stats
|
||||
|
||||
let lastSwapFromStore = null;
|
||||
|
||||
async function fetchLastSwapFromStore() {
|
||||
// The in-memory ring empties on restart, which blanked the swap tiles even though
|
||||
// every swap is on disk. Fall back to the persisted log so a restart does not erase
|
||||
// what the dashboard knows.
|
||||
try {
|
||||
const rows = await (await fetch('/api/history?durable=true&limit=50')).json();
|
||||
const swap = (rows || []).find(r => r.event_type === 'LLM Model Switch');
|
||||
if (!swap) return;
|
||||
lastSwapFromStore = swap;
|
||||
const t = document.getElementById('ollama-last-swap');
|
||||
const c = document.getElementById('ollama-cache-hit');
|
||||
if (t && (!t.textContent || t.textContent.includes('\u2014') || t.textContent.includes('--'))) {
|
||||
t.textContent = `${Math.round(swap.duration_ms)} ms`;
|
||||
}
|
||||
if (c && (!c.textContent || c.textContent.includes('\u2014') || c.textContent.includes('--')
|
||||
|| c.textContent === 'no swaps yet')) {
|
||||
c.textContent = swap.cache_status || 'OK';
|
||||
}
|
||||
} catch (e) { /* the live stream is the primary source; this is only a fallback */ }
|
||||
}
|
||||
|
||||
document.addEventListener('DOMContentLoaded', () => {
|
||||
setTimeout(fetchLastSwapFromStore, 1500);
|
||||
});
|
||||
|
||||
@@ -24,7 +24,7 @@
|
||||
HYPERSWAP
|
||||
</h1>
|
||||
<span class="text-xs uppercase tracking-widest px-2 py-0.5 rounded bg-cyan-950/80 text-cyan-400 border border-cyan-800 font-mono">
|
||||
v1.0-DEPLOY
|
||||
<span id="app-version">v...</span>
|
||||
</span>
|
||||
</div>
|
||||
<p class="text-xs text-slate-400 font-mono">NVIDIA RTX 4080 SUPER 16GB // 64GB DDR5 RAM // Ubuntu Linux</p>
|
||||
@@ -115,9 +115,14 @@
|
||||
</div>
|
||||
<div class="flex items-center space-x-1.5">
|
||||
<span class="w-2.5 h-2.5 rounded-full bg-slate-500"></span>
|
||||
<span class="text-slate-400">System:</span>
|
||||
<span class="text-slate-400" title="Desktop compositor and display server">Desktop:</span>
|
||||
<span id="legend-system" class="text-slate-200 font-bold">0 GB</span>
|
||||
</div>
|
||||
<div class="flex items-center space-x-1.5">
|
||||
<span class="w-2.5 h-2.5 rounded-full bg-slate-800 border border-slate-700"></span>
|
||||
<span class="text-slate-400" title="Held by processes HyperSwap cannot reclaim">Unmanaged:</span>
|
||||
<span id="legend-vram-unmanaged" class="text-amber-400 font-bold">0 GB</span>
|
||||
</div>
|
||||
<div class="flex items-center space-x-1.5">
|
||||
<span class="w-2.5 h-2.5 rounded-full bg-slate-800 border border-slate-700"></span>
|
||||
<span class="text-slate-400">Free:</span>
|
||||
@@ -159,7 +164,7 @@
|
||||
</div>
|
||||
<div class="flex items-center space-x-1.5">
|
||||
<span class="w-2.5 h-2.5 rounded-full bg-amber-400"></span>
|
||||
<span class="text-slate-400">Models in RAM:</span>
|
||||
<span class="text-slate-400">Page Cache (all files):</span>
|
||||
<span id="legend-ram-cached" class="text-amber-400 font-bold">0 GB</span>
|
||||
</div>
|
||||
<div class="flex items-center space-x-1.5">
|
||||
@@ -167,6 +172,11 @@
|
||||
<span class="text-slate-400">Free RAM:</span>
|
||||
<span id="legend-ram-free" class="text-emerald-400 font-bold">0 GB</span>
|
||||
</div>
|
||||
<div class="flex items-center space-x-1.5" title="Measured with cachestat(2) where the kernel permits it, and a read-rate probe where it does not. This is model weights only, not the whole page cache.">
|
||||
<span class="w-2.5 h-2.5 rounded-full bg-emerald-500"></span>
|
||||
<span class="text-slate-400">Models resident (measured):</span>
|
||||
<span id="legend-models-resident" class="text-emerald-400 font-bold">—</span>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
@@ -286,26 +296,22 @@
|
||||
<!-- Comfy Feature Checklist -->
|
||||
<div class="mt-4 space-y-2 text-xs">
|
||||
<div class="p-2.5 rounded-lg bg-slate-950 border border-slate-800 space-y-1.5">
|
||||
<div class="flex items-center justify-between text-slate-300">
|
||||
<span class="flex items-center space-x-1.5">
|
||||
<i class="fa-solid fa-check text-emerald-400 text-xs"></i>
|
||||
<span>Host Pinned Memory:</span>
|
||||
</span>
|
||||
<span class="font-mono font-bold text-emerald-400">53.6 GB Staging Buffer</span>
|
||||
<div class="space-y-2 text-xs">
|
||||
<div class="flex justify-between">
|
||||
<span class="text-slate-400">VRAM held by ComfyUI:</span>
|
||||
<span id="comfy-vram-held" class="font-mono font-bold text-emerald-400">—</span>
|
||||
</div>
|
||||
<div class="flex items-center justify-between text-slate-300">
|
||||
<span class="flex items-center space-x-1.5">
|
||||
<i class="fa-solid fa-check text-emerald-400 text-xs"></i>
|
||||
<span>Async PCIe Offloading:</span>
|
||||
</span>
|
||||
<span class="font-mono font-bold text-cyan-400">Enabled (2 Streams)</span>
|
||||
<div class="flex justify-between">
|
||||
<span class="text-slate-400">PCIe throughput (TX/RX):</span>
|
||||
<span id="comfy-pcie" class="font-mono font-bold text-cyan-400">—</span>
|
||||
</div>
|
||||
<div class="flex items-center justify-between text-slate-300">
|
||||
<span class="flex items-center space-x-1.5">
|
||||
<i class="fa-solid fa-check text-emerald-400 text-xs"></i>
|
||||
<span>Fast Disk RAM Mmap:</span>
|
||||
</span>
|
||||
<span class="font-mono font-bold text-amber-400">Active</span>
|
||||
<div class="flex justify-between">
|
||||
<span class="text-slate-400">Checkpoints resident in RAM:</span>
|
||||
<span id="comfy-resident" class="font-mono font-bold text-amber-400">—</span>
|
||||
</div>
|
||||
<div class="flex justify-between">
|
||||
<span class="text-slate-400">Idle purge:</span>
|
||||
<span id="comfy-purge-state" class="font-mono font-bold text-slate-300">—</span>
|
||||
</div>
|
||||
</div>
|
||||
</div>
|
||||
@@ -668,6 +674,79 @@
|
||||
<!-- ============ NEXT-LEVEL PANELS: governor / residency / analytics / autotune ============ -->
|
||||
<div class="grid grid-cols-1 xl:grid-cols-2 gap-5 mt-5">
|
||||
|
||||
|
||||
<!-- System health: makes a silently-broken dependency loud -->
|
||||
<div class="bg-slate-900/80 border border-slate-800 rounded-2xl p-5 xl:col-span-2">
|
||||
<div class="flex items-center justify-between pb-3 border-b border-slate-800">
|
||||
<div class="flex items-center space-x-2">
|
||||
<div class="p-2 rounded-lg bg-sky-950/80 border border-sky-800 text-sky-400">
|
||||
<i class="fa-solid fa-stethoscope text-sm"></i>
|
||||
</div>
|
||||
<div>
|
||||
<h3 class="font-bold text-slate-100 text-sm">System Health</h3>
|
||||
<p class="text-xs text-slate-400">Every dependency, with impact and how to fix it</p>
|
||||
</div>
|
||||
</div>
|
||||
<div class="flex items-center space-x-2">
|
||||
<span id="health-badge" class="px-2 py-1 text-xs font-bold rounded-lg bg-slate-800 text-slate-400">checking…</span>
|
||||
<button onclick="fetchHealth(true)" class="px-2.5 py-1 text-xs font-semibold rounded-lg bg-slate-800 border border-slate-700 text-slate-300 hover:bg-slate-700 transition">Recheck</button>
|
||||
</div>
|
||||
</div>
|
||||
<div id="health-body" class="mt-4 space-y-1.5 text-[11px] font-mono"></div>
|
||||
</div>
|
||||
|
||||
<!-- VRAM Arbitration — the core handoff, previously invisible -->
|
||||
<div class="bg-slate-900/80 border border-slate-800 rounded-2xl p-5 xl:col-span-2">
|
||||
<div class="flex items-center justify-between pb-3 border-b border-slate-800">
|
||||
<div class="flex items-center space-x-2">
|
||||
<div class="p-2 rounded-lg bg-emerald-950/80 border border-emerald-800 text-emerald-400">
|
||||
<i class="fa-solid fa-right-left text-sm"></i>
|
||||
</div>
|
||||
<div>
|
||||
<h3 class="font-bold text-slate-100 text-sm">VRAM Arbitration</h3>
|
||||
<p class="text-xs text-slate-400">Who holds the GPU, and how handoffs are going</p>
|
||||
</div>
|
||||
</div>
|
||||
<span id="arb-ws" class="text-xs font-mono text-slate-500">—</span>
|
||||
</div>
|
||||
<div class="mt-4">
|
||||
<div id="arb-action" class="text-sm text-slate-200 bg-slate-950/60 border border-slate-800 rounded-lg px-3 py-2 mb-3 font-mono">Idle</div>
|
||||
<div id="arb-backoff" class="text-[11px] font-mono text-amber-400 mb-3"></div>
|
||||
<div id="arb-unmanaged" class="text-[11px] font-mono text-slate-400 mb-3"></div>
|
||||
<div class="grid grid-cols-3 sm:grid-cols-6 gap-2 text-center">
|
||||
<div class="bg-slate-950/60 rounded-lg p-2 border border-slate-800">
|
||||
<div id="arb-yields" class="text-lg font-bold text-emerald-400">0</div>
|
||||
<div class="text-[10px] text-slate-500 uppercase leading-tight">Released</div>
|
||||
</div>
|
||||
<div class="bg-slate-950/60 rounded-lg p-2 border border-slate-800">
|
||||
<div id="arb-busy" class="text-lg font-bold text-cyan-400">0</div>
|
||||
<div class="text-[10px] text-slate-500 uppercase leading-tight">Deferred<br>(busy)</div>
|
||||
</div>
|
||||
<div class="bg-slate-950/60 rounded-lg p-2 border border-slate-800">
|
||||
<div id="arb-later" class="text-lg font-bold text-cyan-400">0</div>
|
||||
<div class="text-[10px] text-slate-500 uppercase leading-tight">Landed<br>later</div>
|
||||
</div>
|
||||
<div class="bg-slate-950/60 rounded-lg p-2 border border-slate-800">
|
||||
<div id="arb-stalled" class="text-lg font-bold text-rose-400">0</div>
|
||||
<div class="text-[10px] text-slate-500 uppercase leading-tight">Stalled</div>
|
||||
</div>
|
||||
<div class="bg-slate-950/60 rounded-lg p-2 border border-slate-800">
|
||||
<div id="arb-purges" class="text-lg font-bold text-fuchsia-400">0</div>
|
||||
<div class="text-[10px] text-slate-500 uppercase leading-tight">Comfy<br>purges</div>
|
||||
</div>
|
||||
<div class="bg-slate-950/60 rounded-lg p-2 border border-slate-800">
|
||||
<div id="arb-defpurge" class="text-lg font-bold text-slate-300">0</div>
|
||||
<div class="text-[10px] text-slate-500 uppercase leading-tight">Purges<br>deferred</div>
|
||||
</div>
|
||||
</div>
|
||||
<p class="text-[11px] text-slate-500 mt-3">
|
||||
<span class="text-cyan-400">Deferred</span> is healthy — an LLM mid-generation cannot unload, so the
|
||||
request queues and applies the moment it finishes. Only <span class="text-rose-400">stalled</span>
|
||||
(VRAM held while the GPU sits idle) indicates a real problem.
|
||||
</p>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
<!-- Thermal Governor -->
|
||||
<div class="bg-slate-900/80 border border-slate-800 rounded-2xl p-5">
|
||||
<div class="flex items-center justify-between pb-3 border-b border-slate-800">
|
||||
|
||||
@@ -142,6 +142,21 @@ class _Writer(threading.Thread):
|
||||
conn.commit()
|
||||
except Exception as e:
|
||||
logger.debug("prune failed: %s", e)
|
||||
# Drain anything still queued before closing. Without this, rows submitted but
|
||||
# not yet dequeued are lost on shutdown -- which is exactly when the last events
|
||||
# before a restart matter most.
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
item = self.q.get_nowait()
|
||||
except queue.Empty:
|
||||
break
|
||||
if item is None:
|
||||
continue
|
||||
_kind, sql, params = item
|
||||
pending.append((sql, params))
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
for sql, params in pending:
|
||||
conn.execute(sql, params)
|
||||
|
||||
74
tests/README.md
Normal file
74
tests/README.md
Normal file
@@ -0,0 +1,74 @@
|
||||
# HyperSwap test suite
|
||||
|
||||
Fast, hermetic unit tests. No GPU is touched, no network call is made, no systemd unit is
|
||||
poked, and the production `hyperswap.db` is never opened.
|
||||
|
||||
## Running
|
||||
|
||||
```bash
|
||||
/home/drjones/comfy-mcp-venv/bin/python -m pytest tests/ -q
|
||||
```
|
||||
|
||||
Single file / single test:
|
||||
|
||||
```bash
|
||||
/home/drjones/comfy-mcp-venv/bin/python -m pytest tests/test_classify_load.py -q
|
||||
/home/drjones/comfy-mcp-venv/bin/python -m pytest tests/ -q -k warm_confident
|
||||
```
|
||||
|
||||
Whole suite runs in about 3 seconds.
|
||||
|
||||
## Safety rails
|
||||
|
||||
These matter, because this repo drives a live 4080 SUPER that a running service is using.
|
||||
|
||||
* `tests/conftest.py` installs an **autouse** `no_gpu_mutation` fixture that replaces
|
||||
`overclock_manager._sh` (the single choke point for every `nvidia-smi` /
|
||||
`nvidia-settings` write) plus `apply_profile`, `apply_fan_control`, `set_fan_speed`,
|
||||
`set_fan_auto` and `restore_safe` with recording stubs. Even a test that accidentally
|
||||
reaches an actuation path can only reach the stub. The fixture yields a dict of
|
||||
recorded calls, which the thermal tests assert against.
|
||||
* `HYPERSWAP_DB` is set to a non-existent path before `telemetry_store` is imported, so no
|
||||
import can bind `DB_PATH` to the production database. Tests that need a DB use the
|
||||
`temp_db` fixture, which monkeypatches `telemetry_store.DB_PATH` to a `tmp_path` file
|
||||
and stops the writer thread afterwards.
|
||||
* All file IO happens against files the tests create in `tmp_path`. No real model blob is
|
||||
read and `warm_file_to_ram` is never called.
|
||||
* Nothing sweeps, and nothing sends HTTP to Ollama, ComfyUI or :9090.
|
||||
|
||||
## Measured constants pinned here
|
||||
|
||||
These numbers came from measurement on this box, not from taste. If a change makes one of
|
||||
these tests fail, the constant is probably wrong, not the test.
|
||||
|
||||
| Constant | Value | Where pinned |
|
||||
| --- | --- | --- |
|
||||
| Cold load of a 12.87 GB model, 3.1% resident | 34267 ms → 0.38 GB/s | `test_classify_load.py::test_measured_cold_load_classifies_as_cold_disk` |
|
||||
| Warm load of the same model, 100% resident | 4901 ms → 2.63 GB/s | `test_classify_load.py::test_measured_warm_load_classifies_as_ram_hit` |
|
||||
| `RAM_HIT_GBPS = 2.0` must stay below the fastest achievable warm load (2.63 GB/s) | — | `test_classify_load.py::test_ram_hit_threshold_is_physically_achievable` |
|
||||
| `PARTIAL_HIT_GBPS = 0.8` must stay above the measured cold rate (0.38 GB/s) | — | same test |
|
||||
| Size-unknown fallback splits at 8000 ms (between 4.9 s warm and 34.3 s cold) | — | `test_classify_load.py::test_unknown_size_guess_boundary_is_8s` |
|
||||
| `WARM_SKIP_THRESHOLD_PCT = 90.0` | — | `test_ram_optimizer.py::test_warm_skip_threshold_constant_unchanged` |
|
||||
| A *probe* reading may only be trusted at exactly 100% (a 12-window probe once cleared 90% on a mostly-cold 12.87 GB blob that then loaded at 2.44 GB/s) | — | `test_ram_optimizer.py::test_probe_reading_is_only_trusted_at_exactly_100_percent` |
|
||||
| `PROBE_CACHED_GBPS = 1.5` sits in the gap between cold NVMe (0.35–0.5 GB/s) and page cache (3.2–13 GB/s) | — | `test_ram_optimizer.py::test_probe_cached_threshold_sits_between_measured_disk_and_cache_rates` |
|
||||
| Card power envelope: 320 W stock, 370 W max, sweeps never go below 60% of max | — | `test_autotune_helpers.py::test_supported_power_limits_parses_min_default_max` |
|
||||
| `_supported_clocks` must always query the `mem,gr` pair (a single-field query returns one column and silently yielded `[]`) | — | `test_autotune_helpers.py::test_supported_clocks_always_queries_the_mem_gr_pair` |
|
||||
| ComfyUI benchmark seed must vary per call (a fixed seed made ComfyUI serve a cached result in ~1 ms) | — | `test_autotune_helpers.py::test_comfy_workflow_seed_varies_between_calls` |
|
||||
| Governor hysteresis: `HOT_SAMPLES = 5`, `COOL_SAMPLES = 30`, `REAPPLY_COOLDOWN_S = 20` | — | `test_thermal_governor.py` (escalation, recovery, cooldown, alternating-sample tests) |
|
||||
| Model usage score: frequency decayed with a ~24 h half-life | — | `test_telemetry_store.py::test_model_usage_ranking_scores_recent_use_higher` |
|
||||
|
||||
## What is deliberately not covered
|
||||
|
||||
* `vram_arbitrator.instant_free_ollama_vram`, the `AutoArbitrator` yield/purge paths and
|
||||
the SSE broker — under active edit, contract changing.
|
||||
* `overclock_manager.apply_profile` and every other actuation path, `autotune.sweep`,
|
||||
`ram_optimizer.warm_file_to_ram` — these mutate hardware or do heavy IO.
|
||||
* `server.py` HTTP routes and `mcp_server.py` — would need the app wired to live
|
||||
subsystems.
|
||||
|
||||
## Known rough edge the tests work around
|
||||
|
||||
`telemetry_store.stop()` flushes the writer's pending *batch* but does not drain the
|
||||
submission queue, so a `stop()` racing a just-submitted row can drop it. The writer tests
|
||||
call a local `_drain()` helper to wait for the queue to empty before stopping, rather than
|
||||
encoding the race into an assertion.
|
||||
69
tests/conftest.py
Normal file
69
tests/conftest.py
Normal file
@@ -0,0 +1,69 @@
|
||||
"""Shared fixtures and — more importantly — hardware safety rails for the suite.
|
||||
|
||||
This repo drives a live GPU and a running systemd service. Every test here must be
|
||||
hermetic: no NVML mutation, no nvidia-smi/nvidia-settings writes, no touching the
|
||||
production telemetry DB, no HTTP to Ollama/ComfyUI/:9090.
|
||||
|
||||
The `no_gpu_mutation` fixture below is autouse, so even a test that accidentally
|
||||
reaches an actuation path can only reach a recording stub.
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
if REPO_ROOT not in sys.path:
|
||||
sys.path.insert(0, REPO_ROOT)
|
||||
|
||||
# telemetry_store resolves DB_PATH from the environment *at import time*. Point it at a
|
||||
# path that does not exist before anything imports it, so no import of this suite can
|
||||
# ever open the production hyperswap.db. Individual tests monkeypatch DB_PATH to a
|
||||
# tmp_path file when they actually need a database.
|
||||
os.environ.setdefault("HYPERSWAP_DB", os.path.join(REPO_ROOT, "tests", "_never_created.db"))
|
||||
|
||||
import overclock_manager # noqa: E402 (must follow the sys.path/env setup above)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def no_gpu_mutation(monkeypatch):
|
||||
"""Hard block on every code path that can physically change GPU state.
|
||||
|
||||
Yields a dict of call recorders so a test can assert that actuation *would* have
|
||||
happened without any of it reaching the card.
|
||||
"""
|
||||
calls = {"apply_profile": [], "fan": [], "restore_safe": [], "sh": []}
|
||||
|
||||
def _blocked_sh(cmd, use_sudo=True, timeout=10):
|
||||
# Catch-all: every nvidia-smi / nvidia-settings write in overclock_manager
|
||||
# funnels through _sh. Nothing in the suite may shell out to the driver.
|
||||
calls["sh"].append(list(cmd))
|
||||
return {"rc": -1, "out": "", "err": "blocked by test suite"}
|
||||
|
||||
monkeypatch.setattr(overclock_manager, "_sh", _blocked_sh)
|
||||
monkeypatch.setattr(overclock_manager, "apply_profile",
|
||||
lambda name, overrides=None: calls["apply_profile"].append((name, overrides)))
|
||||
monkeypatch.setattr(overclock_manager, "apply_fan_control",
|
||||
lambda mode, speed_pct: calls["fan"].append((mode, speed_pct)))
|
||||
monkeypatch.setattr(overclock_manager, "set_fan_speed",
|
||||
lambda percent: calls["fan"].append(("manual", percent)))
|
||||
monkeypatch.setattr(overclock_manager, "set_fan_auto",
|
||||
lambda: calls["fan"].append(("auto", None)))
|
||||
monkeypatch.setattr(overclock_manager, "restore_safe",
|
||||
lambda reason="shutdown": calls["restore_safe"].append(reason))
|
||||
return calls
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def temp_db(tmp_path, monkeypatch):
|
||||
"""Point telemetry_store at a throwaway SQLite file for the duration of one test."""
|
||||
import telemetry_store
|
||||
|
||||
db = tmp_path / "test_hyperswap.db"
|
||||
monkeypatch.setattr(telemetry_store, "DB_PATH", str(db))
|
||||
yield str(db)
|
||||
# Never leave a writer thread running against a tmp path that is about to vanish.
|
||||
try:
|
||||
telemetry_store.stop()
|
||||
except Exception:
|
||||
pass
|
||||
184
tests/test_autotune_helpers.py
Normal file
184
tests/test_autotune_helpers.py
Normal file
@@ -0,0 +1,184 @@
|
||||
"""autotune's pure/parsable helpers.
|
||||
|
||||
Only the parsing and sampling helpers are exercised. Nothing here runs a sweep, applies a
|
||||
profile, or talks to nvidia-smi: subprocess.run is monkeypatched at the module boundary.
|
||||
"""
|
||||
import pytest
|
||||
|
||||
import autotune
|
||||
|
||||
|
||||
class _FakeProc:
|
||||
def __init__(self, stdout="", returncode=0, stderr=""):
|
||||
self.stdout = stdout
|
||||
self.stderr = stderr
|
||||
self.returncode = returncode
|
||||
|
||||
|
||||
def _fake_smi(monkeypatch, stdout, record=None):
|
||||
def _run(cmd, capture_output=True, text=True, timeout=None, **kw):
|
||||
if record is not None:
|
||||
record.append(list(cmd))
|
||||
return _FakeProc(stdout)
|
||||
monkeypatch.setattr(autotune.subprocess, "run", _run)
|
||||
|
||||
|
||||
# A trimmed but realistically shaped `nvidia-smi --query-supported-clocks=mem,gr` dump:
|
||||
# graphics clocks are enumerated once per memory clock.
|
||||
SUPPORTED_CLOCKS_CSV = """\
|
||||
10501, 2790
|
||||
10501, 2775
|
||||
10501, 2760
|
||||
9501, 2790
|
||||
9501, 2775
|
||||
405, 645
|
||||
405, 630
|
||||
"""
|
||||
|
||||
|
||||
def test_supported_clocks_mem_returns_sorted_unique_memory_clocks(monkeypatch):
|
||||
_fake_smi(monkeypatch, SUPPORTED_CLOCKS_CSV)
|
||||
assert autotune._supported_clocks("mem") == [405, 9501, 10501]
|
||||
|
||||
|
||||
def test_supported_clocks_gr_returns_clocks_of_the_highest_memory_clock(monkeypatch):
|
||||
"""Graphics clocks are enumerated per memory clock. Only the set belonging to the top
|
||||
memory clock is meaningful — that is the state any real workload runs at."""
|
||||
_fake_smi(monkeypatch, SUPPORTED_CLOCKS_CSV)
|
||||
assert autotune._supported_clocks("gr") == [2760, 2775, 2790]
|
||||
|
||||
|
||||
def test_supported_clocks_always_queries_the_mem_gr_pair(monkeypatch):
|
||||
"""Regression: querying a single field returns one column, and reading index 1 from it
|
||||
silently produced an empty list instead of an error. The query must stay a pair."""
|
||||
seen = []
|
||||
_fake_smi(monkeypatch, SUPPORTED_CLOCKS_CSV, record=seen)
|
||||
autotune._supported_clocks("gr")
|
||||
assert any("--query-supported-clocks=mem,gr" in part for part in seen[0])
|
||||
|
||||
|
||||
def test_supported_clocks_ignores_single_column_output(monkeypatch):
|
||||
"""If the driver ever returns one column, no row parses and we return [] — never a
|
||||
list of memory clocks mislabelled as graphics clocks."""
|
||||
_fake_smi(monkeypatch, "10501\n9501\n405\n")
|
||||
assert autotune._supported_clocks("mem") == []
|
||||
assert autotune._supported_clocks("gr") == []
|
||||
|
||||
|
||||
def test_supported_clocks_skips_header_and_garbage_rows(monkeypatch):
|
||||
_fake_smi(monkeypatch, "memory [MHz], graphics [MHz]\n10501, 2790\n\nN/A, N/A\n")
|
||||
assert autotune._supported_clocks("mem") == [10501]
|
||||
|
||||
|
||||
def test_supported_clocks_returns_empty_when_nvidia_smi_fails(monkeypatch):
|
||||
"""No GPU / no driver must degrade to 'no candidates', not an exception on import of
|
||||
a sweep plan."""
|
||||
def _boom(*a, **kw):
|
||||
raise OSError("nvidia-smi not found")
|
||||
monkeypatch.setattr(autotune.subprocess, "run", _boom)
|
||||
assert autotune._supported_clocks("mem") == []
|
||||
|
||||
|
||||
# --------------------------------------------------------------- power limits
|
||||
|
||||
def test_supported_power_limits_parses_min_default_max(monkeypatch):
|
||||
"""RTX 4080 SUPER on this box: 115 W min, 370 W max, 320 W stock default."""
|
||||
_fake_smi(monkeypatch, "115.00, 370.00, 320.00\n")
|
||||
vals = autotune._supported_power_limits(steps=5)
|
||||
assert vals == sorted(set(vals))
|
||||
assert max(vals) == 370
|
||||
# Never sweeps below 60% of max — the card does no useful work down there.
|
||||
assert min(vals) >= int(370 * 0.6)
|
||||
# The stock default is always included as a reference point.
|
||||
assert 320 in vals
|
||||
|
||||
|
||||
def test_supported_power_limits_respects_step_count(monkeypatch):
|
||||
_fake_smi(monkeypatch, "115.00, 370.00, 320.00\n")
|
||||
vals = autotune._supported_power_limits(steps=3)
|
||||
assert len(vals) <= 4 # 3 evenly spaced values, plus the default if it is not one
|
||||
|
||||
|
||||
def test_supported_power_limits_returns_empty_on_query_failure(monkeypatch):
|
||||
_fake_smi(monkeypatch, "N/A, N/A, N/A\n")
|
||||
assert autotune._supported_power_limits() == []
|
||||
|
||||
|
||||
# --------------------------------------------------------------- subsampling
|
||||
|
||||
def test_subsample_returns_input_when_already_short_enough():
|
||||
assert autotune._subsample([1, 2, 3], 5) == [1, 2, 3]
|
||||
|
||||
|
||||
def test_subsample_keeps_both_endpoints():
|
||||
"""The endpoints are the whole point of a sweep: stock and maximum. Dropping either
|
||||
means never measuring the setting the sweep is supposed to recommend."""
|
||||
values = list(range(0, 195)) # the card enumerates ~194 graphics clocks
|
||||
out = autotune._subsample(values, 8)
|
||||
assert out[0] == values[0]
|
||||
assert out[-1] == values[-1]
|
||||
|
||||
|
||||
def test_subsample_never_exceeds_max_steps():
|
||||
values = list(range(0, 195))
|
||||
for max_steps in (2, 3, 5, 8, 13):
|
||||
assert len(autotune._subsample(values, max_steps)) <= max_steps
|
||||
|
||||
|
||||
def test_subsample_output_is_sorted_and_unique():
|
||||
values = list(range(0, 100))
|
||||
out = autotune._subsample(values, 7)
|
||||
assert out == sorted(set(out))
|
||||
|
||||
|
||||
def test_subsample_is_evenly_spread():
|
||||
"""Clustered samples would tell us nothing about the middle of the range."""
|
||||
out = autotune._subsample(list(range(0, 101)), 5)
|
||||
gaps = [b - a for a, b in zip(out, out[1:])]
|
||||
assert max(gaps) - min(gaps) <= 1
|
||||
|
||||
|
||||
# --------------------------------------------------------------- comfy benchmark graph
|
||||
|
||||
def test_comfy_workflow_seed_varies_between_calls():
|
||||
"""Fixed seeds made ComfyUI serve a cached result in ~1 ms without executing the
|
||||
graph, so every benchmark after the first measured nothing at all."""
|
||||
seeds = {autotune._comfy_workflow()["5"]["inputs"]["seed"] for _ in range(25)}
|
||||
assert len(seeds) > 20
|
||||
|
||||
|
||||
def test_comfy_workflow_seed_can_be_pinned_for_reproduction():
|
||||
assert autotune._comfy_workflow(seed=42)["5"]["inputs"]["seed"] == 42
|
||||
|
||||
|
||||
def test_comfy_workflow_shape_is_otherwise_constant():
|
||||
"""Only the seed may vary; a benchmark whose step count or resolution moved between
|
||||
runs would not be comparable."""
|
||||
a = autotune._comfy_workflow()
|
||||
b = autotune._comfy_workflow()
|
||||
for wf in (a, b):
|
||||
assert wf["5"]["inputs"]["steps"] == autotune.COMFY_BENCH_STEPS
|
||||
assert wf["4"]["inputs"]["width"] == autotune.COMFY_BENCH_SIZE
|
||||
assert wf["4"]["inputs"]["height"] == autotune.COMFY_BENCH_SIZE
|
||||
a["5"]["inputs"]["seed"] = b["5"]["inputs"]["seed"] = 0
|
||||
assert a == b
|
||||
|
||||
|
||||
def test_comfy_workflow_uses_the_named_checkpoint():
|
||||
wf = autotune._comfy_workflow(ckpt="some_other.safetensors")
|
||||
assert wf["1"]["inputs"]["ckpt_name"] == "some_other.safetensors"
|
||||
|
||||
|
||||
def test_temp_ceiling_is_below_the_thermal_governors_escalation_point():
|
||||
"""A sweep step must abort on temperature before the governor starts derating under
|
||||
it, otherwise the sweep measures the governor's derate rather than the knob."""
|
||||
import thermal_governor as tg
|
||||
assert autotune.TEMP_CEILING_C <= tg.TEMP_ESCALATE_C + 1.0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("knob", ["mem_offset_mhz", "core_offset_mhz", "lock_mem_mhz",
|
||||
"lock_core_max", "power_limit_w"])
|
||||
def test_every_knob_declares_a_hardware_verification_field(knob):
|
||||
"""Offsets are silently ignored by some drivers (595.84 accepts an assignment and
|
||||
reads back a different value), so each knob must name the field to read back."""
|
||||
assert autotune.KNOBS[knob]["verify"]
|
||||
123
tests/test_classify_load.py
Normal file
123
tests/test_classify_load.py
Normal file
@@ -0,0 +1,123 @@
|
||||
"""Regression tests for vram_arbitrator.classify_load.
|
||||
|
||||
classify_load turns (model size, Ollama's reported load_duration) into a cache verdict
|
||||
by computing achieved bandwidth. The two thresholds it compares against were measured on
|
||||
this box, not chosen; the tests below pin the measurements themselves so a future edit
|
||||
to RAM_HIT_GBPS / PARTIAL_HIT_GBPS that breaks the real data points fails loudly.
|
||||
"""
|
||||
import math
|
||||
|
||||
import pytest
|
||||
|
||||
import vram_arbitrator as va
|
||||
|
||||
GIB = 1024 ** 3
|
||||
|
||||
# Ground truth, measured 2026-08-28 on the same 12.87 GB model loaded twice:
|
||||
# 3.1% resident -> 34267 ms -> 0.38 GB/s -> cold
|
||||
# 100% resident -> 4901 ms -> 2.63 GB/s -> RAM hit
|
||||
MEASURED_MODEL_BYTES = int(12.87 * GIB)
|
||||
MEASURED_COLD_MS = 34267.0
|
||||
MEASURED_WARM_MS = 4901.0
|
||||
|
||||
|
||||
def test_measured_cold_load_classifies_as_cold_disk():
|
||||
"""The measured cold load (12.87 GB, 34267 ms, 0.38 GB/s) must stay 'Cold Disk Load'."""
|
||||
res = va.classify_load(MEASURED_MODEL_BYTES, MEASURED_COLD_MS)
|
||||
assert res["cache_status"] == "Cold Disk Load 💾"
|
||||
assert res["is_ram_hit"] is False
|
||||
assert res["load_gbps"] == pytest.approx(0.38, abs=0.01)
|
||||
|
||||
|
||||
def test_measured_warm_load_classifies_as_ram_hit():
|
||||
"""The measured warm load (12.87 GB, 4901 ms, 2.63 GB/s) must stay a RAM cache hit."""
|
||||
res = va.classify_load(MEASURED_MODEL_BYTES, MEASURED_WARM_MS)
|
||||
assert res["cache_status"] == "RAM Cache Hit ⚡"
|
||||
assert res["is_ram_hit"] is True
|
||||
assert res["load_gbps"] == pytest.approx(2.63, abs=0.01)
|
||||
|
||||
|
||||
def test_ram_hit_threshold_is_physically_achievable():
|
||||
"""Guard against the reintroduced bug where the hit bar sat above any real warm load.
|
||||
|
||||
A prior version set the hit threshold at 5.0 GB/s. Ollama's load_duration covers
|
||||
host-to-device transfer and model init as well as the file read, so a fully resident
|
||||
12.87 GB model still only reports ~2.6 GB/s — meaning *no* load could ever be
|
||||
classified as a hit. Any threshold that a genuinely warm load cannot clear is wrong.
|
||||
"""
|
||||
assert va.RAM_HIT_GBPS <= 2.63, (
|
||||
"RAM_HIT_GBPS is above the fastest warm load ever measured on this box "
|
||||
"(2.63 GB/s); no load could be classified as a cache hit."
|
||||
)
|
||||
assert va.PARTIAL_HIT_GBPS > 0.38, (
|
||||
"PARTIAL_HIT_GBPS is at or below the measured cold-disk rate; cold loads would "
|
||||
"be reported as partial cache hits."
|
||||
)
|
||||
assert va.PARTIAL_HIT_GBPS < va.RAM_HIT_GBPS
|
||||
|
||||
|
||||
def test_sub_millisecond_load_is_already_in_vram():
|
||||
"""load_duration_ms <= 1.0 means Ollama never re-read the model at all."""
|
||||
res = va.classify_load(MEASURED_MODEL_BYTES, 1.0)
|
||||
assert res["cache_status"] == "Already in VRAM"
|
||||
assert res["load_gbps"] is None
|
||||
assert res["is_ram_hit"] is True
|
||||
|
||||
|
||||
def test_zero_duration_is_already_in_vram():
|
||||
res = va.classify_load(MEASURED_MODEL_BYTES, 0.0)
|
||||
assert res["cache_status"] == "Already in VRAM"
|
||||
|
||||
|
||||
def test_boundary_exactly_at_ram_hit_threshold_is_a_hit():
|
||||
"""Exactly RAM_HIT_GBPS (2.0 GB/s) must count as a hit — the comparison is >=."""
|
||||
size = int(va.RAM_HIT_GBPS * GIB) # 2 GiB read in exactly 1000 ms -> 2.00 GB/s
|
||||
res = va.classify_load(size, 1000.0)
|
||||
assert res["load_gbps"] == pytest.approx(2.0)
|
||||
assert res["cache_status"] == "RAM Cache Hit ⚡"
|
||||
assert res["is_ram_hit"] is True
|
||||
|
||||
|
||||
def test_just_below_ram_hit_threshold_is_partial():
|
||||
size = int(1.99 * GIB)
|
||||
res = va.classify_load(size, 1000.0)
|
||||
assert res["cache_status"] == "Partial Cache 🌤"
|
||||
assert res["is_ram_hit"] is False
|
||||
|
||||
|
||||
def test_boundary_exactly_at_partial_threshold_is_partial():
|
||||
"""Exactly PARTIAL_HIT_GBPS (0.8 GB/s) is still a partial hit, not a cold load."""
|
||||
size = math.ceil(va.PARTIAL_HIT_GBPS * GIB) # 0.8 GiB is not a whole number of bytes
|
||||
res = va.classify_load(size, 1000.0)
|
||||
assert res["load_gbps"] == pytest.approx(0.8)
|
||||
assert res["cache_status"] == "Partial Cache 🌤"
|
||||
assert res["is_ram_hit"] is False
|
||||
|
||||
|
||||
def test_just_below_partial_threshold_is_cold():
|
||||
size = int(0.79 * GIB)
|
||||
res = va.classify_load(size, 1000.0)
|
||||
assert res["cache_status"] == "Cold Disk Load 💾"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("duration_ms,expected_hit", [(4901.0, True), (34267.0, False)])
|
||||
def test_unknown_size_falls_back_to_duration_guess_and_says_so(duration_ms, expected_hit):
|
||||
"""With no size on record there is no bandwidth to compute, so the result must be
|
||||
labelled as a guess rather than presented as a measurement."""
|
||||
res = va.classify_load(0, duration_ms)
|
||||
assert res["is_ram_hit"] is expected_hit
|
||||
assert res["load_gbps"] is None
|
||||
assert "detail" in res and "guess" in res["detail"].lower()
|
||||
|
||||
|
||||
def test_unknown_size_guess_boundary_is_8s():
|
||||
"""The size-unknown fallback splits at 8000 ms, between the measured 4.9 s warm and
|
||||
34.3 s cold loads."""
|
||||
assert va.classify_load(0, 7999.0)["is_ram_hit"] is True
|
||||
assert va.classify_load(0, 8000.0)["is_ram_hit"] is False
|
||||
|
||||
|
||||
def test_classification_never_raises_on_odd_inputs():
|
||||
"""This runs on the swap hot path; it must not be able to throw."""
|
||||
for size, dur in [(0, 0.0), (1, 1.5), (10 ** 12, 2.0), (0, 1.0)]:
|
||||
assert "cache_status" in va.classify_load(size, dur)
|
||||
134
tests/test_health.py
Normal file
134
tests/test_health.py
Normal file
@@ -0,0 +1,134 @@
|
||||
"""Tests for the dependency self-check.
|
||||
|
||||
This module exists because fan control failed for an entire session, recoverably and
|
||||
invisibly: the service started before the headless X server that owns the GPU was
|
||||
accepting connections, the assignment failed with "Error resolving target specification",
|
||||
nothing retried, and nothing ever asked whether fan control worked. These tests make sure
|
||||
each check reports the *right* status, since a self-check that returns ok when a
|
||||
dependency is broken is worse than having none.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
import health
|
||||
import overclock_manager
|
||||
|
||||
|
||||
class TestFanControlCheck:
|
||||
"""The check that would have caught the original bug."""
|
||||
|
||||
def test_fails_when_headless_x_is_not_running(self, monkeypatch):
|
||||
monkeypatch.setattr(overclock_manager, "is_headless_x_running", lambda: False)
|
||||
res = health._check_fan_control()
|
||||
assert res["status"] == health.FAILED
|
||||
# A bare failure is not enough; it has to say what breaks and how to fix it.
|
||||
assert "governor" in res["impact"].lower() or "fan" in res["impact"].lower()
|
||||
assert res["fix"]
|
||||
|
||||
def test_fails_when_the_gpu_target_cannot_be_resolved(self, monkeypatch):
|
||||
# The exact nvidia-settings error seen at startup.
|
||||
monkeypatch.setattr(overclock_manager, "is_headless_x_running", lambda: True)
|
||||
monkeypatch.setattr(overclock_manager, "get_fan_status",
|
||||
lambda force=False: {"manual": False, "mode": "auto",
|
||||
"target_speed_pct": None})
|
||||
monkeypatch.setattr(overclock_manager, "_nvidia_settings", lambda *a, **k: {
|
||||
"rc": 1, "out": "",
|
||||
"err": "ERROR: Error resolving target specification 'gpu:0' "
|
||||
"(No targets match target specification)"})
|
||||
res = health._check_fan_control()
|
||||
assert res["status"] == health.FAILED
|
||||
|
||||
def test_ok_when_fan_status_reads_back(self, monkeypatch):
|
||||
monkeypatch.setattr(overclock_manager, "is_headless_x_running", lambda: True)
|
||||
monkeypatch.setattr(overclock_manager, "get_fan_status",
|
||||
lambda force=False: {"manual": True, "mode": "manual",
|
||||
"target_speed_pct": 70})
|
||||
assert health._check_fan_control()["status"] == health.OK
|
||||
|
||||
|
||||
class TestProfileDriftCheck:
|
||||
def test_degraded_when_no_profile_applied_since_start(self, monkeypatch):
|
||||
# ACTIVE_PROFILE defaults to "balanced" at import, which used to be
|
||||
# indistinguishable from "balanced was applied successfully".
|
||||
monkeypatch.setattr(overclock_manager, "profile_drift", lambda: {
|
||||
"profile": "balanced", "applied_since_start": False, "drifted": True,
|
||||
"power_limit_intended_w": 320, "power_limit_actual_w": 370.0,
|
||||
"reason": "no profile has been successfully applied since startup"})
|
||||
assert health._check_profile_drift()["status"] == health.DEGRADED
|
||||
|
||||
def test_degraded_when_hardware_disagrees(self, monkeypatch):
|
||||
monkeypatch.setattr(overclock_manager, "profile_drift", lambda: {
|
||||
"profile": "balanced", "applied_since_start": True, "drifted": True,
|
||||
"power_limit_intended_w": 320, "power_limit_actual_w": 370.0,
|
||||
"reason": "card reports 370.0W, profile asks 320W"})
|
||||
res = health._check_profile_drift()
|
||||
assert res["status"] == health.DEGRADED
|
||||
assert "370" in res["detail"]
|
||||
|
||||
def test_ok_when_they_agree(self, monkeypatch):
|
||||
monkeypatch.setattr(overclock_manager, "profile_drift", lambda: {
|
||||
"profile": "balanced", "applied_since_start": True, "drifted": False,
|
||||
"power_limit_intended_w": 320, "power_limit_actual_w": 320.0,
|
||||
"reason": None})
|
||||
assert health._check_profile_drift()["status"] == health.OK
|
||||
|
||||
|
||||
class TestSudoCheck:
|
||||
def test_failed_when_sudo_smi_returns_nonzero(self, monkeypatch):
|
||||
monkeypatch.setattr(overclock_manager, "_smi",
|
||||
lambda *a: {"rc": 1, "out": "", "err": "sudo: a password is required"})
|
||||
res = health._check_sudo_smi()
|
||||
assert res["status"] == health.FAILED
|
||||
assert "sudo" in res["fix"].lower()
|
||||
|
||||
def test_ok_when_it_works(self, monkeypatch):
|
||||
monkeypatch.setattr(overclock_manager, "_smi",
|
||||
lambda *a: {"rc": 0, "out": "NVIDIA GeForce RTX 4080 SUPER", "err": ""})
|
||||
assert health._check_sudo_smi()["status"] == health.OK
|
||||
|
||||
|
||||
class TestAggregation:
|
||||
"""Overall status must be driven by the worst individual result."""
|
||||
|
||||
def _fake(self, statuses):
|
||||
return [health._check(f"c{i}", s, "d") for i, s in enumerate(statuses)]
|
||||
|
||||
@pytest.mark.parametrize("statuses,expected", [
|
||||
([health.OK, health.OK], health.OK),
|
||||
([health.OK, health.DEGRADED], health.DEGRADED),
|
||||
([health.OK, health.FAILED], health.FAILED),
|
||||
([health.DEGRADED, health.FAILED], health.FAILED),
|
||||
])
|
||||
def test_worst_status_wins(self, monkeypatch, statuses, expected):
|
||||
checks = self._fake(statuses)
|
||||
monkeypatch.setattr(health, "_check_nvml", lambda: checks[0])
|
||||
monkeypatch.setattr(health, "_check_sudo_smi", lambda: checks[1])
|
||||
for fn in ("_check_fan_control", "_check_profile_drift", "_check_store",
|
||||
"_check_residency", "_check_model_dirs", "_check_comfy_ws"):
|
||||
monkeypatch.setattr(health, fn, lambda: health._check("x", health.OK, "d"))
|
||||
|
||||
async def fake_http(name, url, impact, fix):
|
||||
return health._check(name, health.OK, "reachable")
|
||||
monkeypatch.setattr(health, "_check_http", fake_http)
|
||||
|
||||
res = asyncio.run(health.run_health_checks())
|
||||
assert res["status"] == expected
|
||||
|
||||
def test_a_raising_check_does_not_break_the_report(self, monkeypatch):
|
||||
def boom():
|
||||
raise RuntimeError("nvml exploded")
|
||||
monkeypatch.setattr(health, "_check_nvml", boom)
|
||||
for fn in ("_check_sudo_smi", "_check_fan_control", "_check_profile_drift",
|
||||
"_check_store", "_check_residency", "_check_model_dirs",
|
||||
"_check_comfy_ws"):
|
||||
monkeypatch.setattr(health, fn, lambda: health._check("x", health.OK, "d"))
|
||||
|
||||
async def fake_http(name, url, impact, fix):
|
||||
return health._check(name, health.OK, "reachable")
|
||||
monkeypatch.setattr(health, "_check_http", fake_http)
|
||||
|
||||
res = asyncio.run(health.run_health_checks())
|
||||
# A broken check must surface as failed, not take down the endpoint.
|
||||
assert res["status"] == health.FAILED
|
||||
assert any("exploded" in c["detail"] for c in res["checks"])
|
||||
116
tests/test_overclock_profiles.py
Normal file
116
tests/test_overclock_profiles.py
Normal file
@@ -0,0 +1,116 @@
|
||||
"""overclock_manager profile storage: load/merge/validate only.
|
||||
|
||||
Nothing here applies a profile or shells out. CONFIG_PATH is redirected to tmp_path so the
|
||||
repo's real overclock_profiles.json is never written, and conftest's autouse fixture blocks
|
||||
_sh() and every actuation entry point.
|
||||
"""
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
import overclock_manager as ocm
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def config(tmp_path, monkeypatch):
|
||||
"""Redirect the profile store to a temp file. Returns a writer for its contents."""
|
||||
path = tmp_path / "overclock_profiles.json"
|
||||
monkeypatch.setattr(ocm, "CONFIG_PATH", str(path))
|
||||
|
||||
def write(data):
|
||||
path.write_text(json.dumps(data))
|
||||
return type("Cfg", (), {"path": path, "write": staticmethod(write)})
|
||||
|
||||
|
||||
def test_defaults_are_returned_when_no_config_file_exists(config):
|
||||
profiles = ocm.load_profiles()
|
||||
assert set(profiles) == set(ocm.DEFAULT_PROFILES)
|
||||
assert profiles["ollama"]["power_limit_w"] == ocm.DEFAULT_PROFILES["ollama"]["power_limit_w"]
|
||||
|
||||
|
||||
def test_stored_values_override_defaults_key_by_key(config):
|
||||
"""Only the keys present on disk change; the rest of the default profile survives.
|
||||
A whole-profile replacement would silently drop keys added by a later version."""
|
||||
config.write({"ollama": {"power_limit_w": 340}})
|
||||
profiles = ocm.load_profiles()
|
||||
assert profiles["ollama"]["power_limit_w"] == 340
|
||||
assert profiles["ollama"]["label"] == ocm.DEFAULT_PROFILES["ollama"]["label"]
|
||||
assert profiles["ollama"]["mem_offset_mhz"] == ocm.DEFAULT_PROFILES["ollama"]["mem_offset_mhz"]
|
||||
|
||||
|
||||
def test_unknown_stored_profiles_are_preserved(config):
|
||||
"""A user-created profile is not in DEFAULT_PROFILES and must not be dropped on load."""
|
||||
config.write({"my-custom": {"label": "mine", "power_limit_w": 300}})
|
||||
profiles = ocm.load_profiles()
|
||||
assert profiles["my-custom"]["power_limit_w"] == 300
|
||||
assert set(ocm.DEFAULT_PROFILES) <= set(profiles)
|
||||
|
||||
|
||||
def test_loading_does_not_mutate_the_defaults(config):
|
||||
"""load_profiles deep-copies DEFAULT_PROFILES; if it did not, one load with a stored
|
||||
override would poison the defaults for the rest of the process."""
|
||||
config.write({"ollama": {"power_limit_w": 111}})
|
||||
ocm.load_profiles()
|
||||
config.write({})
|
||||
assert ocm.load_profiles()["ollama"]["power_limit_w"] == \
|
||||
ocm.DEFAULT_PROFILES["ollama"]["power_limit_w"]
|
||||
assert ocm.DEFAULT_PROFILES["ollama"]["power_limit_w"] != 111
|
||||
|
||||
|
||||
def test_corrupt_config_falls_back_to_defaults(config):
|
||||
"""A half-written JSON file must not take the whole service down at startup."""
|
||||
config.path.write_text("{ not valid json")
|
||||
profiles = ocm.load_profiles()
|
||||
assert set(profiles) == set(ocm.DEFAULT_PROFILES)
|
||||
|
||||
|
||||
def test_set_profile_rejects_an_unknown_name(config):
|
||||
"""set_profile edits an existing profile; it is not a create-anything endpoint."""
|
||||
res = ocm.set_profile("no-such-profile", {"power_limit_w": 400})
|
||||
assert res["success"] is False
|
||||
assert "unknown profile" in res["error"]
|
||||
assert not config.path.exists()
|
||||
|
||||
|
||||
def test_set_profile_persists_a_partial_update(config):
|
||||
res = ocm.set_profile("comfy", {"power_limit_w": 350})
|
||||
assert res["success"] is True
|
||||
stored = json.loads(config.path.read_text())
|
||||
assert stored["comfy"]["power_limit_w"] == 350
|
||||
# The other profiles are written back untouched.
|
||||
assert stored["ollama"]["power_limit_w"] == ocm.DEFAULT_PROFILES["ollama"]["power_limit_w"]
|
||||
assert ocm.load_profiles()["comfy"]["power_limit_w"] == 350
|
||||
|
||||
|
||||
def test_set_profile_can_edit_a_previously_stored_custom_profile(config):
|
||||
config.write({"my-custom": {"label": "mine", "power_limit_w": 300}})
|
||||
assert ocm.set_profile("my-custom", {"power_limit_w": 310})["success"] is True
|
||||
assert ocm.load_profiles()["my-custom"]["power_limit_w"] == 310
|
||||
|
||||
|
||||
def test_get_profiles_matches_load_profiles(config):
|
||||
config.write({"ollama": {"core_offset_mhz": 42}})
|
||||
assert ocm.get_profiles() == ocm.load_profiles()
|
||||
|
||||
|
||||
def test_save_profiles_reports_failure_instead_of_raising(tmp_path, monkeypatch):
|
||||
"""The dashboard calls this; an unwritable config must surface as success=False."""
|
||||
monkeypatch.setattr(ocm, "CONFIG_PATH", str(tmp_path / "no-such-dir" / "p.json"))
|
||||
assert ocm.save_profiles({"ollama": {}}) is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", ["ollama", "comfy", "balanced"])
|
||||
def test_default_profiles_declare_the_keys_the_appliers_read(name):
|
||||
"""apply_profile and the thermal governor both index these keys directly; a profile
|
||||
missing one would KeyError inside the actuation thread."""
|
||||
cfg = ocm.DEFAULT_PROFILES[name]
|
||||
for key in ("label", "power_limit_w", "core_offset_mhz", "mem_offset_mhz",
|
||||
"lock_core_min", "lock_core_max"):
|
||||
assert key in cfg, f"{name} is missing {key}"
|
||||
|
||||
|
||||
def test_default_power_limits_stay_within_the_cards_range():
|
||||
"""RTX 4080 SUPER: 320 W stock, 370 W maximum. A profile above that is silently
|
||||
clamped by the driver and the sweep results become meaningless."""
|
||||
for name, cfg in ocm.DEFAULT_PROFILES.items():
|
||||
assert 100 <= cfg["power_limit_w"] <= 370, name
|
||||
421
tests/test_ram_optimizer.py
Normal file
421
tests/test_ram_optimizer.py
Normal file
@@ -0,0 +1,421 @@
|
||||
"""ram_optimizer: page-cache residency measurement, model discovery, warm planning.
|
||||
|
||||
All file IO here happens against files this test creates in tmp_path. Nothing reads a
|
||||
real model blob, and nothing calls warm_file_to_ram.
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
import ram_optimizer as ro
|
||||
|
||||
MIB = 1024 * 1024
|
||||
|
||||
|
||||
# --------------------------------------------------------------- warm_confident rule
|
||||
|
||||
class _FakeCachestat:
|
||||
"""Stand-in for the kernel's struct cachestat."""
|
||||
|
||||
def __init__(self, nr_cache):
|
||||
self.nr_cache = nr_cache
|
||||
self.nr_dirty = 0
|
||||
self.nr_evicted = 0
|
||||
|
||||
|
||||
def _file_of_pages(tmp_path, pages):
|
||||
p = tmp_path / f"blob_{pages}.bin"
|
||||
p.write_bytes(b"\0" * (pages * ro.PAGE_SIZE))
|
||||
return str(p)
|
||||
|
||||
|
||||
def _force_cachestat(monkeypatch, resident_pages):
|
||||
monkeypatch.setattr(ro, "_cachestat",
|
||||
lambda fd, offset, length: _FakeCachestat(resident_pages))
|
||||
|
||||
|
||||
def _force_probe(monkeypatch, pct):
|
||||
monkeypatch.setattr(ro, "_cachestat", lambda fd, offset, length: None)
|
||||
monkeypatch.setattr(ro, "_throughput_probe",
|
||||
lambda fd, size, windows=None: {"resident_pct": pct, "windows": 12,
|
||||
"median_gbps": 3.0})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("resident_pages,expect_confident", [
|
||||
(100, True), # 100% exact
|
||||
(90, True), # exactly WARM_SKIP_THRESHOLD_PCT
|
||||
(89, False), # just under the threshold
|
||||
(0, False),
|
||||
])
|
||||
def test_cachestat_reading_is_trusted_at_or_above_threshold(tmp_path, monkeypatch,
|
||||
resident_pages, expect_confident):
|
||||
"""An exact cachestat reading >= WARM_SKIP_THRESHOLD_PCT (90) may be trusted to skip
|
||||
warming; below it, it may not."""
|
||||
path = _file_of_pages(tmp_path, 100)
|
||||
_force_cachestat(monkeypatch, resident_pages)
|
||||
res = ro.page_residency(path)
|
||||
assert res["method"] == "cachestat"
|
||||
assert res["resident_pct"] == pytest.approx(float(resident_pages))
|
||||
assert res["warm_confident"] is expect_confident
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pct,expect_confident", [
|
||||
(100.0, True), # only an unblemished probe is trustworthy
|
||||
(99.9, False),
|
||||
(95.0, False), # would pass `warm`, must NOT pass `warm_confident`
|
||||
(90.0, False),
|
||||
])
|
||||
def test_probe_reading_is_only_trusted_at_exactly_100_percent(tmp_path, monkeypatch,
|
||||
pct, expect_confident):
|
||||
"""Correctness fix, not a style choice: a 12-window probe once cleared 90% on a file
|
||||
that was mostly cold (a 12.87 GB blob reported 'already resident' then loaded at
|
||||
2.44 GB/s). Only a perfect probe score may skip work."""
|
||||
path = _file_of_pages(tmp_path, 100)
|
||||
_force_probe(monkeypatch, pct)
|
||||
res = ro.page_residency(path)
|
||||
assert res["method"] == "probe"
|
||||
assert res["warm_confident"] is expect_confident
|
||||
|
||||
|
||||
def test_probe_at_95_percent_is_warm_but_not_confident(tmp_path, monkeypatch):
|
||||
"""`warm` and `warm_confident` are different questions and must not be conflated:
|
||||
warm drives display, warm_confident drives skipping work."""
|
||||
path = _file_of_pages(tmp_path, 100)
|
||||
_force_probe(monkeypatch, 95.0)
|
||||
res = ro.page_residency(path)
|
||||
assert res["warm"] is True
|
||||
assert res["warm_confident"] is False
|
||||
|
||||
|
||||
def test_warm_skip_threshold_constant_unchanged():
|
||||
# Pinned: build_warm_plan and the dashboard both key off this number.
|
||||
assert ro.WARM_SKIP_THRESHOLD_PCT == 90.0
|
||||
|
||||
|
||||
# --------------------------------------------------------------- page_residency basics
|
||||
|
||||
def test_page_residency_on_missing_file_reports_not_measurable(tmp_path):
|
||||
res = ro.page_residency(str(tmp_path / "nope.bin"))
|
||||
assert res["success"] is False
|
||||
assert res["measurable"] is False
|
||||
assert res["resident_pct"] == 0.0
|
||||
|
||||
|
||||
def test_page_residency_on_empty_file_short_circuits(tmp_path):
|
||||
"""A zero-byte file has no pages to measure; dividing by its page count would throw."""
|
||||
p = tmp_path / "empty.bin"
|
||||
p.write_bytes(b"")
|
||||
res = ro.page_residency(str(p))
|
||||
assert res["method"] == "empty"
|
||||
assert res["resident_pct"] == 0.0
|
||||
assert res["measurable"] is True
|
||||
|
||||
|
||||
def test_page_residency_refuses_to_guess_when_probing_is_disallowed(tmp_path, monkeypatch):
|
||||
"""With cachestat unavailable and allow_probe=False the answer is 'unknown', never 0%.
|
||||
Reporting 0% would make the planner re-warm a file that may be fully resident."""
|
||||
path = _file_of_pages(tmp_path, 8)
|
||||
monkeypatch.setattr(ro, "_cachestat", lambda fd, offset, length: None)
|
||||
res = ro.page_residency(path, allow_probe=False)
|
||||
assert res["measurable"] is False
|
||||
assert res["resident_pct"] is None
|
||||
assert res["warm"] is None
|
||||
assert res["method"] == "unavailable"
|
||||
|
||||
|
||||
def test_page_residency_reports_size_and_bytes_consistently(tmp_path, monkeypatch):
|
||||
path = _file_of_pages(tmp_path, 100)
|
||||
_force_cachestat(monkeypatch, 50)
|
||||
res = ro.page_residency(path)
|
||||
size = os.path.getsize(path)
|
||||
assert res["size_bytes"] == size
|
||||
assert res["resident_bytes"] == int(size * 0.5)
|
||||
|
||||
|
||||
def test_page_residency_against_a_real_file_holds_the_confidence_invariant(tmp_path):
|
||||
"""End-to-end with the real kernel path (whichever method is available here): the
|
||||
warm_confident rule must hold for whatever the machine actually reports."""
|
||||
path = _file_of_pages(tmp_path, 512)
|
||||
res = ro.page_residency(path)
|
||||
assert res["success"] is True
|
||||
assert res["method"] in ("cachestat", "probe")
|
||||
assert 0.0 <= res["resident_pct"] <= 100.0
|
||||
expected = ((res["method"] == "cachestat" and res["resident_pct"] >= 90.0)
|
||||
or (res["method"] == "probe" and res["resident_pct"] >= 100.0))
|
||||
assert res["warm_confident"] is expected
|
||||
|
||||
|
||||
# --------------------------------------------------------------- throughput probe
|
||||
|
||||
def test_throughput_probe_window_count_is_clamped_to_file_size(tmp_path):
|
||||
"""A file smaller than PROBE_WINDOWS * PROBE_WINDOW_BYTES must not be probed with
|
||||
more windows than it has, or offsets run past EOF."""
|
||||
p = tmp_path / "small.bin"
|
||||
p.write_bytes(b"\0" * (5 * MIB))
|
||||
fd = os.open(str(p), os.O_RDONLY)
|
||||
try:
|
||||
out = ro._throughput_probe(fd, 5 * MIB)
|
||||
finally:
|
||||
os.close(fd)
|
||||
assert 0 < out["windows"] <= 5 * MIB // ro.PROBE_WINDOW_BYTES
|
||||
assert 0.0 <= out["resident_pct"] <= 100.0
|
||||
|
||||
|
||||
def test_throughput_probe_honours_window_override(tmp_path):
|
||||
p = tmp_path / "medium.bin"
|
||||
p.write_bytes(b"\0" * (32 * MIB))
|
||||
fd = os.open(str(p), os.O_RDONLY)
|
||||
try:
|
||||
out = ro._throughput_probe(fd, 32 * MIB, windows_override=4)
|
||||
finally:
|
||||
os.close(fd)
|
||||
assert out["windows"] == 4
|
||||
|
||||
|
||||
def test_probe_cached_threshold_sits_between_measured_disk_and_cache_rates():
|
||||
# Measured on this box: cold NVMe 0.35-0.5 GB/s, page cache 3.2-13 GB/s.
|
||||
assert 0.5 < ro.PROBE_CACHED_GBPS < 3.2
|
||||
|
||||
|
||||
# --------------------------------------------------------------- ollama manifest parsing
|
||||
|
||||
MODEL_MEDIA_TYPE = "application/vnd.ollama.image.model"
|
||||
|
||||
|
||||
def _write_manifest(root, rel_dir, tag, layers):
|
||||
d = os.path.join(root, "manifests", rel_dir)
|
||||
os.makedirs(d, exist_ok=True)
|
||||
with open(os.path.join(d, tag), "w") as f:
|
||||
json.dump({"layers": layers}, f)
|
||||
|
||||
|
||||
def _write_blob(root, digest, size):
|
||||
blobs = os.path.join(root, "blobs")
|
||||
os.makedirs(blobs, exist_ok=True)
|
||||
path = os.path.join(blobs, digest.replace(":", "-"))
|
||||
with open(path, "wb") as f:
|
||||
f.write(b"\0" * size)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ollama_tree(tmp_path, monkeypatch):
|
||||
root = tmp_path / "ollama-models"
|
||||
root.mkdir()
|
||||
monkeypatch.setattr(ro, "OLLAMA_MODEL_DIRS", [str(root)])
|
||||
return str(root)
|
||||
|
||||
|
||||
def test_find_ollama_model_files_maps_library_model_to_its_blob(ollama_tree):
|
||||
"""registry/library/<name>/<tag> is the common case and must render as 'name:tag'."""
|
||||
_write_blob(ollama_tree, "sha256:aaa111", 4096)
|
||||
_write_manifest(ollama_tree, "registry.ollama.ai/library/llama3", "8b", [
|
||||
{"mediaType": MODEL_MEDIA_TYPE, "digest": "sha256:aaa111", "size": 987654321},
|
||||
])
|
||||
files = ro.find_ollama_model_files()
|
||||
assert len(files) == 1
|
||||
entry = files[0]
|
||||
assert entry["model"] == "llama3:8b"
|
||||
assert entry["filename"] == "sha256-aaa111"
|
||||
assert entry["size_bytes"] == 987654321 # taken from the manifest, not the stub blob
|
||||
assert entry["kind"] == "ollama"
|
||||
assert os.path.exists(entry["full_path"])
|
||||
|
||||
|
||||
def test_find_ollama_model_files_keeps_non_library_namespace(ollama_tree):
|
||||
"""A model pulled from a user namespace must keep it, or two different models with
|
||||
the same short name collide in the warm plan."""
|
||||
_write_blob(ollama_tree, "sha256:bbb222", 4096)
|
||||
_write_manifest(ollama_tree, "hf.co/someuser/qwen-coder", "q4", [
|
||||
{"mediaType": MODEL_MEDIA_TYPE, "digest": "sha256:bbb222", "size": 100},
|
||||
])
|
||||
files = ro.find_ollama_model_files()
|
||||
assert [f["model"] for f in files] == ["someuser/qwen-coder:q4"]
|
||||
|
||||
|
||||
def test_find_ollama_model_files_ignores_non_model_layers(ollama_tree):
|
||||
"""Manifests also list templates, params and licence layers; warming those is wasted
|
||||
IO and they are not the weights."""
|
||||
_write_blob(ollama_tree, "sha256:ccc333", 4096)
|
||||
_write_blob(ollama_tree, "sha256:ddd444", 32)
|
||||
_write_manifest(ollama_tree, "registry.ollama.ai/library/mistral", "7b", [
|
||||
{"mediaType": "application/vnd.ollama.image.template", "digest": "sha256:ddd444", "size": 32},
|
||||
{"mediaType": MODEL_MEDIA_TYPE, "digest": "sha256:ccc333", "size": 500},
|
||||
])
|
||||
files = ro.find_ollama_model_files()
|
||||
assert len(files) == 1
|
||||
assert files[0]["filename"] == "sha256-ccc333"
|
||||
|
||||
|
||||
def test_find_ollama_model_files_skips_layers_whose_blob_is_missing(ollama_tree):
|
||||
"""A partially pulled model leaves a manifest with no blob; returning that path would
|
||||
make every downstream residency check fail."""
|
||||
_write_manifest(ollama_tree, "registry.ollama.ai/library/ghost", "latest", [
|
||||
{"mediaType": MODEL_MEDIA_TYPE, "digest": "sha256:missing", "size": 10},
|
||||
])
|
||||
assert ro.find_ollama_model_files() == []
|
||||
|
||||
|
||||
def test_find_ollama_model_files_tolerates_corrupt_manifest_json(ollama_tree):
|
||||
"""One unreadable manifest must not take out discovery of every other model."""
|
||||
bad_dir = os.path.join(ollama_tree, "manifests", "registry.ollama.ai/library/broken")
|
||||
os.makedirs(bad_dir)
|
||||
with open(os.path.join(bad_dir, "latest"), "w") as f:
|
||||
f.write("{not json")
|
||||
_write_blob(ollama_tree, "sha256:eee555", 4096)
|
||||
_write_manifest(ollama_tree, "registry.ollama.ai/library/good", "latest", [
|
||||
{"mediaType": MODEL_MEDIA_TYPE, "digest": "sha256:eee555", "size": 10},
|
||||
])
|
||||
assert [f["model"] for f in ro.find_ollama_model_files()] == ["good:latest"]
|
||||
|
||||
|
||||
def test_find_ollama_model_files_deduplicates_same_model_and_blob(ollama_tree, monkeypatch):
|
||||
"""The same root listed twice (or a duplicated layer) must not double-count bytes
|
||||
against the warm budget."""
|
||||
_write_blob(ollama_tree, "sha256:fff666", 4096)
|
||||
_write_manifest(ollama_tree, "registry.ollama.ai/library/dup", "latest", [
|
||||
{"mediaType": MODEL_MEDIA_TYPE, "digest": "sha256:fff666", "size": 10},
|
||||
{"mediaType": MODEL_MEDIA_TYPE, "digest": "sha256:fff666", "size": 10},
|
||||
])
|
||||
monkeypatch.setattr(ro, "OLLAMA_MODEL_DIRS", [ollama_tree, ollama_tree])
|
||||
assert len(ro.find_ollama_model_files()) == 1
|
||||
|
||||
|
||||
def test_find_ollama_model_files_returns_empty_when_no_manifest_dir(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(ro, "OLLAMA_MODEL_DIRS", [str(tmp_path / "does-not-exist")])
|
||||
assert ro.find_ollama_model_files() == []
|
||||
|
||||
|
||||
# --------------------------------------------------------------- warm planning
|
||||
|
||||
GIB = 1024 ** 3
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def planner(monkeypatch):
|
||||
"""build_warm_plan with every IO boundary stubbed: catalog, residency, meminfo, scores."""
|
||||
state = {"residency": {}, "scores": {}}
|
||||
|
||||
def _catalog(force_refresh=False):
|
||||
return state["catalog"]
|
||||
|
||||
monkeypatch.setattr(ro, "get_model_catalog", _catalog)
|
||||
monkeypatch.setattr(ro, "_warm_priority", lambda days=30.0: state["scores"])
|
||||
monkeypatch.setattr(ro, "get_detailed_meminfo",
|
||||
lambda: {"available_bytes": 40 * GIB, "available_gb": 40.0})
|
||||
monkeypatch.setattr(ro, "page_residency",
|
||||
lambda path, allow_probe=True, probe_windows=None:
|
||||
state["residency"].get(path, {"resident_pct": 0.0,
|
||||
"warm_confident": False}))
|
||||
return state
|
||||
|
||||
|
||||
def _ollama_entry(name, path, gb):
|
||||
return {"model": name, "full_path": path, "size_bytes": int(gb * GIB),
|
||||
"size_gb": gb, "kind": "ollama"}
|
||||
|
||||
|
||||
def _comfy_entry(rel, path, gb, mtime):
|
||||
return {"rel_path": rel, "full_path": path, "size_bytes": int(gb * GIB),
|
||||
"size_gb": gb, "kind": "comfy", "mtime": mtime}
|
||||
|
||||
|
||||
def test_warm_plan_orders_by_usage_score(planner):
|
||||
"""The whole point of the plan is that the most-used model is warmed first, because
|
||||
the budget usually cannot cover everything."""
|
||||
planner["catalog"] = {
|
||||
"ollama": [_ollama_entry("rare:latest", "/m/rare", 1.0),
|
||||
_ollama_entry("hot:latest", "/m/hot", 1.0)],
|
||||
"comfy": [],
|
||||
}
|
||||
planner["scores"] = {"hot:latest": 50.0, "rare:latest": 0.1}
|
||||
plan = ro.build_warm_plan(budget_gb=10.0)
|
||||
assert [e["name"] for e in plan["plan"]] == ["hot:latest", "rare:latest"]
|
||||
|
||||
|
||||
def test_warm_plan_skips_confidently_warm_files_without_spending_budget(planner):
|
||||
"""Already-resident files cost nothing and must not consume budget another file needs."""
|
||||
planner["catalog"] = {
|
||||
"ollama": [_ollama_entry("warm:latest", "/m/warm", 8.0),
|
||||
_ollama_entry("cold:latest", "/m/cold", 8.0)],
|
||||
"comfy": [],
|
||||
}
|
||||
planner["scores"] = {"warm:latest": 10.0, "cold:latest": 5.0}
|
||||
planner["residency"]["/m/warm"] = {"resident_pct": 100.0, "warm_confident": True}
|
||||
plan = ro.build_warm_plan(budget_gb=10.0)
|
||||
assert [e["name"] for e in plan["plan"]] == ["cold:latest"]
|
||||
assert [e["action"] for e in plan["skipped"]] == ["already-warm"]
|
||||
assert plan["planned_gb"] == pytest.approx(8.0, abs=0.01)
|
||||
|
||||
|
||||
def test_warm_plan_does_not_skip_a_high_but_unconfident_residency(planner):
|
||||
"""95% from a probe is not permission to skip — this is the bug the warm_confident
|
||||
flag exists to prevent."""
|
||||
planner["catalog"] = {"ollama": [_ollama_entry("m:latest", "/m/x", 4.0)], "comfy": []}
|
||||
planner["residency"]["/m/x"] = {"resident_pct": 95.0, "warm_confident": False}
|
||||
plan = ro.build_warm_plan(budget_gb=10.0)
|
||||
assert [e["name"] for e in plan["plan"]] == ["m:latest"]
|
||||
# Only the missing 5% has to be read.
|
||||
assert plan["plan"][0]["bytes_to_read"] == pytest.approx(int(4.0 * GIB) * 0.05, rel=0.01)
|
||||
|
||||
|
||||
def test_warm_plan_charges_only_the_non_resident_fraction(planner):
|
||||
planner["catalog"] = {"ollama": [_ollama_entry("m:latest", "/m/x", 10.0)], "comfy": []}
|
||||
planner["residency"]["/m/x"] = {"resident_pct": 50.0, "warm_confident": False}
|
||||
plan = ro.build_warm_plan(budget_gb=6.0)
|
||||
assert plan["planned_gb"] == pytest.approx(5.0, abs=0.01)
|
||||
|
||||
|
||||
def test_warm_plan_stops_at_the_budget(planner):
|
||||
"""Warming past the budget just evicts what was warmed first, so over-budget entries
|
||||
are reported as skipped rather than planned."""
|
||||
planner["catalog"] = {
|
||||
"ollama": [_ollama_entry("a", "/m/a", 6.0), _ollama_entry("b", "/m/b", 6.0)],
|
||||
"comfy": [],
|
||||
}
|
||||
planner["scores"] = {"a": 9.0, "b": 1.0}
|
||||
plan = ro.build_warm_plan(budget_gb=8.0)
|
||||
assert [e["name"] for e in plan["plan"]] == ["a"]
|
||||
assert plan["skipped"][0]["name"] == "b"
|
||||
assert plan["skipped"][0]["action"] == "over-budget"
|
||||
assert plan["planned_gb"] <= plan["budget_gb"]
|
||||
|
||||
|
||||
def test_warm_plan_deduplicates_by_path(planner):
|
||||
"""The same file reachable from both catalogs must be planned once, or the budget is
|
||||
charged twice for one read."""
|
||||
planner["catalog"] = {
|
||||
"ollama": [_ollama_entry("shared", "/m/shared", 2.0)],
|
||||
"comfy": [_comfy_entry("shared.safetensors", "/m/shared", 2.0, 0)],
|
||||
}
|
||||
plan = ro.build_warm_plan(budget_gb=100.0)
|
||||
assert len(plan["plan"]) == 1
|
||||
assert plan["planned_gb"] == pytest.approx(2.0, abs=0.01)
|
||||
|
||||
|
||||
def test_warm_plan_default_budget_leaves_headroom(planner):
|
||||
"""An unbounded budget would push the box into reclaim; the default is 70% of
|
||||
MemAvailable."""
|
||||
planner["catalog"] = {"ollama": [], "comfy": []}
|
||||
plan = ro.build_warm_plan()
|
||||
assert plan["budget_gb"] == pytest.approx(40.0 * 0.7, abs=0.01)
|
||||
|
||||
|
||||
def test_warm_plan_ranks_recent_comfy_checkpoints_above_stale_ones(planner):
|
||||
"""ComfyUI files have no usage history, so recency by mtime is the ranking signal."""
|
||||
import time
|
||||
now = time.time()
|
||||
planner["catalog"] = {
|
||||
"ollama": [],
|
||||
"comfy": [_comfy_entry("old.safetensors", "/c/old", 1.0, now - 90 * 86400),
|
||||
_comfy_entry("new.safetensors", "/c/new", 1.0, now - 60)],
|
||||
}
|
||||
plan = ro.build_warm_plan(budget_gb=10.0)
|
||||
assert [e["name"] for e in plan["plan"]] == ["new.safetensors", "old.safetensors"]
|
||||
|
||||
|
||||
def test_warm_plan_on_empty_catalog_is_a_valid_empty_plan(planner):
|
||||
planner["catalog"] = {"ollama": [], "comfy": []}
|
||||
plan = ro.build_warm_plan(budget_gb=1.0)
|
||||
assert plan["warm_count"] == 0 and plan["plan"] == [] and plan["skipped"] == []
|
||||
342
tests/test_telemetry_store.py
Normal file
342
tests/test_telemetry_store.py
Normal file
@@ -0,0 +1,342 @@
|
||||
"""telemetry_store: schema, the background writer round-trip, and the analytics queries.
|
||||
|
||||
Every test in this module runs against a throwaway SQLite file via the `temp_db` fixture
|
||||
(which monkeypatches telemetry_store.DB_PATH). The production hyperswap.db is never
|
||||
opened: the writer thread, _read_conn() and _rows() all resolve DB_PATH at call time.
|
||||
"""
|
||||
import os
|
||||
import sqlite3
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
import telemetry_store as ts
|
||||
|
||||
|
||||
def _drain(timeout=3.0):
|
||||
"""Wait until the writer thread has dequeued everything submitted so far.
|
||||
|
||||
NOTE: stop() flushes the writer's *pending* batch but does not drain the submission
|
||||
queue, so a stop() racing a just-submitted row can drop it. Tests wait explicitly
|
||||
rather than depending on that race. (See tests/README.md.)
|
||||
"""
|
||||
w = ts._writer
|
||||
assert w is not None, "writer not started"
|
||||
deadline = time.time() + timeout
|
||||
while time.time() < deadline:
|
||||
if w.q.empty():
|
||||
time.sleep(0.05) # let the in-flight item finish its execute/commit
|
||||
return
|
||||
time.sleep(0.01)
|
||||
raise AssertionError("telemetry writer never drained its queue")
|
||||
|
||||
|
||||
def _seed(db_path, events=(), telemetry=()):
|
||||
"""Insert rows directly, bypassing the writer thread, for query tests."""
|
||||
conn = sqlite3.connect(db_path)
|
||||
conn.executescript(ts.SCHEMA)
|
||||
for e in events:
|
||||
conn.execute(ts._EVENT_SQL, (
|
||||
e.get("ts", time.time()), e.get("event_type"), e.get("source"), e.get("target"),
|
||||
e.get("profile"), e.get("duration_ms"), e.get("load_duration_ms"),
|
||||
e.get("yield_confirm_ms"), e.get("tokens_per_sec"), e.get("bytes_loaded"),
|
||||
e.get("load_gbps"), e.get("cache_status"), e.get("detail"),
|
||||
))
|
||||
for t in telemetry:
|
||||
conn.execute(ts._TELEMETRY_SQL, (
|
||||
t.get("ts", time.time()), t.get("profile"), t.get("gpu_util_pct"),
|
||||
t.get("mem_util_pct"), t.get("temp_c"), t.get("power_w"), t.get("power_limit_w"),
|
||||
t.get("fan_pct"), t.get("clock_sm_mhz"), t.get("clock_mem_mhz"),
|
||||
t.get("vram_used_bytes"), t.get("ollama_bytes"), t.get("comfy_bytes"),
|
||||
t.get("system_bytes"), t.get("ram_used_bytes"), t.get("ram_cached_bytes"),
|
||||
t.get("pcie_tx_kbps"), t.get("pcie_rx_kbps"), t.get("throttle_reasons"),
|
||||
))
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
|
||||
# --------------------------------------------------------------- writer round-trip
|
||||
|
||||
def test_start_creates_the_schema(temp_db):
|
||||
"""The service starts against a database that may not exist yet."""
|
||||
ts.start()
|
||||
ts.stop()
|
||||
assert os.path.exists(temp_db)
|
||||
conn = sqlite3.connect(temp_db)
|
||||
tables = {r[0] for r in conn.execute("SELECT name FROM sqlite_master WHERE type='table'")}
|
||||
conn.close()
|
||||
assert {"telemetry", "events", "autotune_runs"} <= tables
|
||||
|
||||
|
||||
def test_record_event_round_trips_through_the_writer(temp_db):
|
||||
"""Events are committed immediately (not batched) because they are what the analytics
|
||||
are built on."""
|
||||
ts.start()
|
||||
ts.record_event({
|
||||
"event_type": "LLM Model Switch", "source": "None", "target": "llama3:8b",
|
||||
"duration_ms": 5120.0, "load_duration_ms": 4901.0, "tokens_per_sec": 61.2,
|
||||
"bytes_loaded": 13819000000, "load_gbps": 2.63, "cache_status": "RAM Cache Hit ⚡",
|
||||
}, profile="ollama")
|
||||
_drain()
|
||||
ts.stop()
|
||||
|
||||
rows = ts.recent_events()
|
||||
assert len(rows) == 1
|
||||
assert rows[0]["target"] == "llama3:8b"
|
||||
assert rows[0]["profile"] == "ollama"
|
||||
assert rows[0]["cache_status"] == "RAM Cache Hit ⚡"
|
||||
assert rows[0]["load_gbps"] == pytest.approx(2.63)
|
||||
|
||||
|
||||
def test_record_telemetry_flushes_on_stop(temp_db):
|
||||
"""Telemetry is batched on a FLUSH_INTERVAL_S timer; stop() must not drop what is
|
||||
still pending, or the last seconds before a restart are lost."""
|
||||
ts.start()
|
||||
ts.record_telemetry(
|
||||
gpu={"available": True, "gpu_util_pct": 88.0, "temperature_c": 71.0,
|
||||
"power_w": 355.0, "clock_sm_mhz": 2790.0, "clock_mem_mhz": 10501.0,
|
||||
"vram_used_bytes": 13 * 1024 ** 3,
|
||||
"breakdown": {"ollama_gb": 12.0, "comfyui_gb": 0.0, "system_gb": 0.5}},
|
||||
ram={"used_bytes": 20 * 1024 ** 3, "cached_bytes": 30 * 1024 ** 3},
|
||||
profile="ollama", throttle_reasons="sw_power_cap")
|
||||
_drain()
|
||||
ts.stop()
|
||||
|
||||
conn = sqlite3.connect(temp_db)
|
||||
row = conn.execute("SELECT profile, temp_c, ollama_bytes, throttle_reasons "
|
||||
"FROM telemetry").fetchone()
|
||||
conn.close()
|
||||
assert row[0] == "ollama"
|
||||
assert row[1] == pytest.approx(71.0)
|
||||
assert row[2] == 12 * 1024 ** 3
|
||||
assert row[3] == "sw_power_cap"
|
||||
|
||||
|
||||
def test_record_telemetry_drops_unavailable_gpu_samples(temp_db):
|
||||
"""A failed NVML read must not write a row of nulls that then skews every average."""
|
||||
ts.start()
|
||||
ts.record_telemetry(gpu={"available": False}, ram={}, profile="ollama")
|
||||
_drain()
|
||||
ts.stop()
|
||||
conn = sqlite3.connect(temp_db)
|
||||
n = conn.execute("SELECT COUNT(*) FROM telemetry").fetchone()[0]
|
||||
conn.close()
|
||||
assert n == 0
|
||||
|
||||
|
||||
def test_recording_without_a_running_writer_is_a_no_op(temp_db):
|
||||
"""Callers on the hot path must never need to know whether the store was started."""
|
||||
ts.record_event({"event_type": "LLM Model Switch", "target": "x"})
|
||||
ts.record_telemetry(gpu={"available": True, "breakdown": {}}, ram={})
|
||||
assert not os.path.exists(temp_db)
|
||||
|
||||
|
||||
def test_record_event_never_raises_on_a_malformed_event(temp_db):
|
||||
"""This is called from the swap path; it is documented as never raising."""
|
||||
ts.start()
|
||||
ts.record_event({})
|
||||
ts.record_event({"event_type": "Weird", "duration_ms": "not-a-number"})
|
||||
_drain()
|
||||
ts.stop()
|
||||
|
||||
|
||||
def test_autotune_row_round_trips(temp_db):
|
||||
ts.start()
|
||||
ts.record_autotune({"profile": "ollama", "knob": "lock_mem_mhz", "mem_offset_mhz": 0,
|
||||
"tokens_per_sec": 62.5, "temp_c": 74.0, "stable": True})
|
||||
_drain()
|
||||
ts.stop()
|
||||
rows = ts.autotune_history()
|
||||
assert len(rows) == 1
|
||||
assert rows[0]["knob"] == "lock_mem_mhz"
|
||||
assert rows[0]["stable"] == 1 # stored as an INTEGER, not a bool
|
||||
|
||||
|
||||
# --------------------------------------------------------------- absent database
|
||||
|
||||
def test_queries_return_empty_when_the_database_does_not_exist(tmp_path, monkeypatch):
|
||||
"""A dashboard opened before the first sample must render empty, not 500."""
|
||||
monkeypatch.setattr(ts, "DB_PATH", str(tmp_path / "absent.db"))
|
||||
assert ts.profile_comparison() == []
|
||||
assert ts.model_usage_ranking() == []
|
||||
assert ts.timeseries() == []
|
||||
assert ts.recent_events() == []
|
||||
assert ts.autotune_history() == []
|
||||
stats = ts.swap_stats()
|
||||
assert stats["by_type"] == [] and stats["by_model"] == [] and stats["by_cache_status"] == []
|
||||
assert ts.db_info()["exists"] is False
|
||||
|
||||
|
||||
def test_queries_return_empty_on_a_corrupt_database(tmp_path, monkeypatch):
|
||||
db = tmp_path / "corrupt.db"
|
||||
db.write_bytes(b"this is not a sqlite file")
|
||||
monkeypatch.setattr(ts, "DB_PATH", str(db))
|
||||
assert ts.recent_events() == []
|
||||
assert ts.profile_comparison() == []
|
||||
|
||||
|
||||
# --------------------------------------------------------------- analytics
|
||||
|
||||
def test_profile_comparison_ranks_profiles_by_throughput(temp_db):
|
||||
"""The headline question this store exists to answer: which profile is actually
|
||||
faster? Ordering is by average tok/s, descending."""
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 60, "event_type": "LLM Model Switch", "profile": "ollama",
|
||||
"tokens_per_sec": 62.0, "load_gbps": 2.6, "load_duration_ms": 4900},
|
||||
{"ts": now - 50, "event_type": "LLM Model Switch", "profile": "ollama",
|
||||
"tokens_per_sec": 64.0, "load_gbps": 2.6, "load_duration_ms": 4900},
|
||||
{"ts": now - 40, "event_type": "LLM Model Switch", "profile": "balanced",
|
||||
"tokens_per_sec": 51.0, "load_gbps": 2.5, "load_duration_ms": 5100},
|
||||
], telemetry=[
|
||||
{"ts": now - 55, "profile": "ollama", "gpu_util_pct": 90, "temp_c": 74.0,
|
||||
"power_w": 360.0, "clock_sm_mhz": 2790, "clock_mem_mhz": 10501},
|
||||
{"ts": now - 45, "profile": "balanced", "gpu_util_pct": 90, "temp_c": 66.0,
|
||||
"power_w": 300.0, "clock_sm_mhz": 2600, "clock_mem_mhz": 9501},
|
||||
])
|
||||
out = ts.profile_comparison(days=1)
|
||||
assert [r["profile"] for r in out] == ["ollama", "balanced"]
|
||||
assert out[0]["swaps"] == 2
|
||||
assert out[0]["avg_tok_s"] == pytest.approx(63.0)
|
||||
# Thermals are joined in from the telemetry table for the same profile.
|
||||
assert out[0]["avg_temp_c"] == pytest.approx(74.0)
|
||||
|
||||
|
||||
def test_profile_comparison_ignores_events_without_measured_throughput(temp_db):
|
||||
"""A swap that produced no tokens tells us nothing about the profile's speed and
|
||||
would drag the average toward zero."""
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 10, "event_type": "LLM Model Switch", "profile": "ollama",
|
||||
"tokens_per_sec": 60.0},
|
||||
{"ts": now - 5, "event_type": "LLM Model Switch", "profile": "ollama",
|
||||
"tokens_per_sec": 0.0},
|
||||
])
|
||||
out = ts.profile_comparison(days=1)
|
||||
assert out[0]["swaps"] == 1 and out[0]["avg_tok_s"] == pytest.approx(60.0)
|
||||
|
||||
|
||||
def test_profile_comparison_excludes_samples_outside_the_window(temp_db):
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 40 * 86400, "event_type": "LLM Model Switch", "profile": "old",
|
||||
"tokens_per_sec": 99.0},
|
||||
{"ts": now - 60, "event_type": "LLM Model Switch", "profile": "ollama",
|
||||
"tokens_per_sec": 60.0},
|
||||
])
|
||||
assert [r["profile"] for r in ts.profile_comparison(days=7)] == ["ollama"]
|
||||
|
||||
|
||||
def test_profile_comparison_ignores_idle_telemetry(temp_db):
|
||||
"""Thermals are only meaningful while the GPU is doing work; idle samples (util <= 5%)
|
||||
would make every profile look cool."""
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 10, "event_type": "LLM Model Switch", "profile": "ollama",
|
||||
"tokens_per_sec": 60.0},
|
||||
], telemetry=[
|
||||
{"ts": now - 12, "profile": "ollama", "gpu_util_pct": 0, "temp_c": 35.0},
|
||||
{"ts": now - 11, "profile": "ollama", "gpu_util_pct": 95, "temp_c": 75.0},
|
||||
])
|
||||
assert ts.profile_comparison(days=1)[0]["avg_temp_c"] == pytest.approx(75.0)
|
||||
|
||||
|
||||
def test_swap_stats_groups_by_event_type_and_cache_status(temp_db):
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 30, "event_type": "LLM Model Switch", "target": "llama3:8b",
|
||||
"duration_ms": 5000, "cache_status": "RAM Cache Hit ⚡", "load_gbps": 2.6,
|
||||
"tokens_per_sec": 60.0},
|
||||
{"ts": now - 20, "event_type": "LLM Model Switch", "target": "llama3:8b",
|
||||
"duration_ms": 7000, "cache_status": "Cold Disk Load 💾", "load_gbps": 0.38,
|
||||
"tokens_per_sec": 58.0},
|
||||
{"ts": now - 10, "event_type": "VRAM Yield", "duration_ms": 900,
|
||||
"yield_confirm_ms": 850},
|
||||
])
|
||||
stats = ts.swap_stats(days=1)
|
||||
by_type = {r["event_type"]: r for r in stats["by_type"]}
|
||||
assert by_type["LLM Model Switch"]["n"] == 2
|
||||
assert by_type["LLM Model Switch"]["avg_ms"] == pytest.approx(6000.0)
|
||||
assert by_type["LLM Model Switch"]["min_ms"] == pytest.approx(5000.0)
|
||||
assert by_type["VRAM Yield"]["avg_confirm_ms"] == pytest.approx(850.0)
|
||||
assert {r["cache_status"] for r in stats["by_cache_status"]} == {
|
||||
"RAM Cache Hit ⚡", "Cold Disk Load 💾"}
|
||||
assert stats["by_model"][0]["model"] == "llama3:8b"
|
||||
assert stats["by_model"][0]["loads"] == 2
|
||||
|
||||
|
||||
def test_model_usage_ranking_scores_recent_use_higher(temp_db):
|
||||
"""Recency/frequency drives the RAM warm budget: given equal load counts, the model
|
||||
used more recently must rank first (half-life ~24 h)."""
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 3600, "event_type": "LLM Model Switch", "target": "recent:latest"},
|
||||
{"ts": now - 3600, "event_type": "LLM Model Switch", "target": "recent:latest"},
|
||||
{"ts": now - 20 * 86400, "event_type": "LLM Model Switch", "target": "stale:latest"},
|
||||
{"ts": now - 20 * 86400, "event_type": "LLM Model Switch", "target": "stale:latest"},
|
||||
])
|
||||
ranking = ts.model_usage_ranking(days=30)
|
||||
assert [r["model"] for r in ranking] == ["recent:latest", "stale:latest"]
|
||||
assert ranking[0]["score"] > ranking[1]["score"]
|
||||
assert ranking[0]["loads"] == 2
|
||||
assert ranking[0]["age_hours"] == pytest.approx(1.0, abs=0.05)
|
||||
|
||||
|
||||
def test_model_usage_ranking_scores_frequent_use_higher_at_equal_recency(temp_db):
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 3600, "event_type": "LLM Model Switch", "target": "often:latest"},
|
||||
{"ts": now - 3601, "event_type": "LLM Model Switch", "target": "often:latest"},
|
||||
{"ts": now - 3602, "event_type": "LLM Model Switch", "target": "often:latest"},
|
||||
{"ts": now - 3600, "event_type": "LLM Model Switch", "target": "once:latest"},
|
||||
])
|
||||
ranking = ts.model_usage_ranking(days=30)
|
||||
assert ranking[0]["model"] == "often:latest"
|
||||
|
||||
|
||||
def test_model_usage_ranking_counts_warms_as_well_as_switches(temp_db):
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 60, "event_type": "Model Warm", "target": "warmed:latest"},
|
||||
{"ts": now - 60, "event_type": "Something Else", "target": "ignored:latest"},
|
||||
])
|
||||
assert [r["model"] for r in ts.model_usage_ranking(days=1)] == ["warmed:latest"]
|
||||
|
||||
|
||||
def test_timeseries_buckets_samples_by_width(temp_db):
|
||||
"""Six hours into six buckets is one bucket per hour; samples inside an hour collapse
|
||||
into a single averaged point."""
|
||||
now = time.time()
|
||||
base = now - 5.5 * 3600
|
||||
_seed(temp_db, telemetry=[
|
||||
{"ts": base + 60, "temp_c": 60.0, "gpu_util_pct": 50},
|
||||
{"ts": base + 120, "temp_c": 70.0, "gpu_util_pct": 70},
|
||||
{"ts": base + 3700, "temp_c": 80.0, "gpu_util_pct": 90},
|
||||
])
|
||||
out = ts.timeseries(hours=6, buckets=6)
|
||||
assert len(out) == 2
|
||||
assert out[0]["temp_c"] == pytest.approx(65.0) # average of 60 and 70
|
||||
assert out[1]["temp_c"] == pytest.approx(80.0)
|
||||
assert out[0]["bucket_ts"] < out[1]["bucket_ts"]
|
||||
|
||||
|
||||
def test_timeseries_excludes_samples_older_than_the_window(temp_db):
|
||||
now = time.time()
|
||||
_seed(temp_db, telemetry=[
|
||||
{"ts": now - 48 * 3600, "temp_c": 99.0},
|
||||
{"ts": now - 60, "temp_c": 60.0},
|
||||
])
|
||||
out = ts.timeseries(hours=6, buckets=240)
|
||||
assert len(out) == 1 and out[0]["temp_c"] == pytest.approx(60.0)
|
||||
|
||||
|
||||
def test_db_info_reports_row_counts_and_coverage(temp_db):
|
||||
now = time.time()
|
||||
_seed(temp_db,
|
||||
events=[{"ts": now - 10, "event_type": "LLM Model Switch", "target": "m"}],
|
||||
telemetry=[{"ts": now - 7200, "temp_c": 60.0}, {"ts": now, "temp_c": 61.0}])
|
||||
info = ts.db_info()
|
||||
assert info["exists"] is True
|
||||
assert info["events_rows"] == 1
|
||||
assert info["telemetry_rows"] == 2
|
||||
assert info["coverage_hours"] == pytest.approx(2.0, abs=0.01)
|
||||
264
tests/test_thermal_governor.py
Normal file
264
tests/test_thermal_governor.py
Normal file
@@ -0,0 +1,264 @@
|
||||
"""thermal_governor: the derate state machine.
|
||||
|
||||
SAFETY: _step() spawns a thread that calls overclock_manager.apply_profile. The autouse
|
||||
`no_gpu_mutation` fixture in conftest.py replaces that (and load_profiles is stubbed per
|
||||
test), so escalation here can never reach the card.
|
||||
"""
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
import overclock_manager
|
||||
import thermal_governor as tg
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def gov(monkeypatch):
|
||||
"""A fresh governor with a stubbed profile store — never the module singleton."""
|
||||
monkeypatch.setattr(overclock_manager, "load_profiles", lambda: {
|
||||
"ollama": {"core_offset_mhz": 100, "mem_offset_mhz": 500, "fan_speed_pct": 70,
|
||||
"lock_core_min": 2500, "lock_core_max": 2800},
|
||||
"stock": {"core_offset_mhz": 0, "mem_offset_mhz": 0, "fan_speed_pct": 0},
|
||||
})
|
||||
monkeypatch.setattr(overclock_manager, "ACTIVE_PROFILE", "ollama")
|
||||
return tg.ThermalGovernor()
|
||||
|
||||
|
||||
def _sample(temp=60.0, reasons=None, available=True):
|
||||
return {"available": available, "temperature_c": temp, "throttle_reasons": reasons or []}
|
||||
|
||||
|
||||
def _feed(gov, n, **kw):
|
||||
for _ in range(n):
|
||||
gov.observe(_sample(**kw), active_profile="ollama")
|
||||
|
||||
|
||||
def _clear_cooldown(gov):
|
||||
"""The REAPPLY_COOLDOWN_S gate is time-based; wind the clock back rather than sleep."""
|
||||
gov.last_change = time.time() - tg.REAPPLY_COOLDOWN_S - 1.0
|
||||
|
||||
|
||||
# --------------------------------------------------------------- escalation hysteresis
|
||||
|
||||
def test_does_not_escalate_before_hot_samples_consecutive_bad_readings(gov):
|
||||
"""One hot spike during a diffusion step must not derate the card."""
|
||||
_feed(gov, tg.HOT_SAMPLES - 1, temp=90.0)
|
||||
assert gov.level == 0
|
||||
assert gov.hot_streak == tg.HOT_SAMPLES - 1
|
||||
|
||||
|
||||
def test_escalates_on_exactly_hot_samples_consecutive_bad_readings(gov, no_gpu_mutation):
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=90.0)
|
||||
assert gov.level == 1
|
||||
assert gov.hot_streak == 0 # streaks reset after a step
|
||||
# Actuation runs in a daemon thread; it must reach apply_profile (the stub) with the
|
||||
# derate overrides rather than blocking the 1 Hz sampler.
|
||||
deadline = time.time() + 2.0
|
||||
while not no_gpu_mutation["apply_profile"] and time.time() < deadline:
|
||||
time.sleep(0.01)
|
||||
assert no_gpu_mutation["apply_profile"], "escalation never actuated"
|
||||
name, overrides = no_gpu_mutation["apply_profile"][0]
|
||||
assert name == "ollama"
|
||||
assert overrides["core_offset_mhz"] == int(100 * tg.DERATE_LADDER[1]["offset_scale"])
|
||||
|
||||
|
||||
def test_alternating_hot_and_cool_samples_never_escalate(gov):
|
||||
"""The core hysteresis property: a card oscillating around the threshold must not
|
||||
thrash the profile. A single good sample resets the hot streak."""
|
||||
for _ in range(50):
|
||||
gov.observe(_sample(temp=90.0), active_profile="ollama")
|
||||
gov.observe(_sample(temp=60.0), active_profile="ollama")
|
||||
assert gov.level == 0
|
||||
|
||||
|
||||
def test_temperature_between_recover_and_escalate_resets_both_streaks(gov):
|
||||
"""The band between TEMP_RECOVER_C and TEMP_ESCALATE_C is neither hot nor cool; it
|
||||
must not accumulate credit in either direction."""
|
||||
_feed(gov, tg.HOT_SAMPLES - 1, temp=90.0)
|
||||
gov.observe(_sample(temp=78.0), active_profile="ollama")
|
||||
assert gov.hot_streak == 0 and gov.cool_streak == 0
|
||||
assert gov.level == 0
|
||||
|
||||
|
||||
def test_hard_throttle_counts_as_hot_even_when_cool(gov):
|
||||
"""A hardware slowdown means the card is protecting itself; temperature alone is not
|
||||
the whole signal."""
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=55.0, reasons=["hw_thermal_slowdown"])
|
||||
assert gov.level == 1
|
||||
|
||||
|
||||
def test_soft_throttle_reasons_do_not_escalate(gov):
|
||||
"""Hitting a power or utilisation cap is normal operation, not distress."""
|
||||
_feed(gov, tg.HOT_SAMPLES * 3, temp=55.0, reasons=["sw_power_cap", "gpu_idle"])
|
||||
assert gov.level == 0
|
||||
|
||||
|
||||
def test_escalation_stops_at_the_bottom_of_the_ladder(gov):
|
||||
"""Level must never index past DERATE_LADDER."""
|
||||
for _ in range(len(tg.DERATE_LADDER) + 3):
|
||||
_clear_cooldown(gov)
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
assert gov.level == len(tg.DERATE_LADDER) - 1
|
||||
|
||||
|
||||
def test_cooldown_gate_blocks_a_second_step_immediately_after_the_first(gov):
|
||||
"""REAPPLY_COOLDOWN_S stops the governor from walking the whole ladder in one second
|
||||
while the card is still responding to the previous change."""
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
assert gov.level == 1
|
||||
_feed(gov, tg.HOT_SAMPLES * 2, temp=95.0)
|
||||
assert gov.level == 1
|
||||
|
||||
|
||||
def test_cooldown_gate_releases_after_the_window(gov):
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
_clear_cooldown(gov)
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
assert gov.level == 2
|
||||
|
||||
|
||||
# --------------------------------------------------------------- recovery
|
||||
|
||||
def test_recovery_needs_cool_samples_consecutive_good_readings(gov):
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
assert gov.level == 1
|
||||
_clear_cooldown(gov)
|
||||
_feed(gov, tg.COOL_SAMPLES - 1, temp=60.0)
|
||||
assert gov.level == 1, "recovered too early"
|
||||
gov.observe(_sample(temp=60.0), active_profile="ollama")
|
||||
assert gov.level == 0
|
||||
|
||||
|
||||
def test_recovery_is_slower_than_escalation():
|
||||
"""Deliberate asymmetry: react fast to heat, give the clocks back slowly."""
|
||||
assert tg.COOL_SAMPLES > tg.HOT_SAMPLES
|
||||
|
||||
|
||||
def test_recovery_stops_at_level_zero(gov):
|
||||
_clear_cooldown(gov)
|
||||
_feed(gov, tg.COOL_SAMPLES * 2, temp=50.0)
|
||||
assert gov.level == 0
|
||||
|
||||
|
||||
def test_a_hard_throttle_blocks_recovery_even_at_a_cool_temperature(gov):
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
_clear_cooldown(gov)
|
||||
_feed(gov, tg.COOL_SAMPLES * 2, temp=50.0, reasons=["hw_power_brake_slowdown"])
|
||||
assert gov.level >= 1
|
||||
|
||||
|
||||
# --------------------------------------------------------------- ignored input
|
||||
|
||||
def test_disabled_governor_ignores_samples(gov):
|
||||
gov.enabled = False
|
||||
_feed(gov, tg.HOT_SAMPLES * 3, temp=99.0)
|
||||
assert gov.level == 0
|
||||
|
||||
|
||||
def test_unavailable_gpu_sample_is_ignored(gov):
|
||||
"""A failed NVML read reports available=False with no temperature; treating that as
|
||||
0 C would count as a cool sample and hand the clocks back."""
|
||||
for _ in range(tg.COOL_SAMPLES * 2):
|
||||
gov.observe({"available": False}, active_profile="ollama")
|
||||
assert gov.level == 0 and gov.cool_streak == 0
|
||||
|
||||
|
||||
# --------------------------------------------------------------- overrides
|
||||
|
||||
def test_overrides_are_empty_at_level_zero(gov):
|
||||
assert gov.overrides_for("ollama") == {}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("level", [1, 2, 3])
|
||||
def test_overrides_scale_offsets_by_the_ladder(gov, level):
|
||||
gov.level = level
|
||||
scale = tg.DERATE_LADDER[level]["offset_scale"]
|
||||
ov = gov.overrides_for("ollama")
|
||||
assert ov["core_offset_mhz"] == int(100 * scale)
|
||||
assert ov["mem_offset_mhz"] == int(500 * scale)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("level", [1, 2, 3])
|
||||
def test_overrides_raise_the_fan_floor_and_never_lower_it(gov, level):
|
||||
"""The floor is a floor: a profile already running fans harder than the ladder asks
|
||||
keeps its own setting."""
|
||||
gov.level = level
|
||||
floor = tg.DERATE_LADDER[level]["fan_floor"]
|
||||
ov = gov.overrides_for("ollama")
|
||||
assert ov["fan_mode"] == "manual"
|
||||
assert ov["fan_speed_pct"] == max(70, floor) # profile fan_speed_pct is 70
|
||||
assert ov["fan_speed_pct"] >= floor
|
||||
|
||||
|
||||
def test_overrides_release_the_core_clock_lock_from_level_two(gov):
|
||||
"""Pinning the core clock high while the card is already backing off fights the
|
||||
hardware's own protection."""
|
||||
gov.level = 1
|
||||
assert "lock_core_max" not in gov.overrides_for("ollama")
|
||||
for level in (2, 3):
|
||||
gov.level = level
|
||||
ov = gov.overrides_for("ollama")
|
||||
assert ov["lock_core_min"] == 0 and ov["lock_core_max"] == 0
|
||||
|
||||
|
||||
def test_top_of_ladder_is_stock_clocks_and_maximum_fans(gov):
|
||||
gov.level = len(tg.DERATE_LADDER) - 1
|
||||
ov = gov.overrides_for("ollama")
|
||||
assert ov["core_offset_mhz"] == 0 and ov["mem_offset_mhz"] == 0
|
||||
assert ov["fan_speed_pct"] == 100
|
||||
|
||||
|
||||
def test_overrides_for_unknown_profile_do_not_raise(gov):
|
||||
"""The active profile can be one the store does not know; a missing config must
|
||||
derate to zero offsets rather than blow up in the sampler thread."""
|
||||
gov.level = 2
|
||||
ov = gov.overrides_for("does-not-exist")
|
||||
assert ov["core_offset_mhz"] == 0 and ov["mem_offset_mhz"] == 0
|
||||
|
||||
|
||||
# --------------------------------------------------------------- ladder invariants
|
||||
|
||||
def test_ladder_is_monotonically_more_conservative():
|
||||
"""Each rung must reduce clocks and raise fans; a non-monotonic ladder would make
|
||||
escalation increase heat."""
|
||||
scales = [s["offset_scale"] for s in tg.DERATE_LADDER]
|
||||
floors = [s["fan_floor"] for s in tg.DERATE_LADDER]
|
||||
assert scales == sorted(scales, reverse=True)
|
||||
assert floors == sorted(floors)
|
||||
assert scales[0] == 1.0 and scales[-1] == 0.0
|
||||
assert [s["level"] for s in tg.DERATE_LADDER] == list(range(len(tg.DERATE_LADDER)))
|
||||
|
||||
|
||||
def test_escalate_and_recover_temperatures_have_a_gap():
|
||||
"""Without a gap between the two thresholds the governor would oscillate."""
|
||||
assert tg.TEMP_RECOVER_C < tg.TEMP_ESCALATE_C
|
||||
|
||||
|
||||
# --------------------------------------------------------------- status / control
|
||||
|
||||
def test_status_reports_level_and_history(gov):
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
st = gov.get_status()
|
||||
assert st["level"] == 1
|
||||
assert st["label"] == tg.DERATE_LADDER[1]["label"]
|
||||
assert st["escalate_at_c"] == tg.TEMP_ESCALATE_C
|
||||
assert st["history"] and st["history"][0]["to_level"] == 1
|
||||
assert "95" in st["history"][0]["reason"]
|
||||
|
||||
|
||||
def test_history_is_bounded(gov):
|
||||
"""The governor is long-lived inside the service; its history must not grow forever."""
|
||||
for _ in range(80):
|
||||
_clear_cooldown(gov)
|
||||
gov.level = 0
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
assert len(gov.history) <= 50
|
||||
|
||||
|
||||
def test_disabling_a_derated_governor_resets_it(gov, no_gpu_mutation):
|
||||
"""Turning the governor off must give the clocks back, not freeze the derate in place."""
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
assert gov.level == 1
|
||||
gov.set_enabled(False)
|
||||
assert gov.level == 0
|
||||
assert no_gpu_mutation["apply_profile"], "reset should have re-applied the base profile"
|
||||
139
tests/test_vram_helpers.py
Normal file
139
tests/test_vram_helpers.py
Normal file
@@ -0,0 +1,139 @@
|
||||
"""Pure helpers in vram_arbitrator: NVML throttle-bit decoding and PID attribution.
|
||||
|
||||
Deliberately excludes instant_free_ollama_vram, the AutoArbitrator yield/purge paths and
|
||||
the SSE broker — that contract is in flux.
|
||||
"""
|
||||
import vram_arbitrator as va
|
||||
|
||||
|
||||
def test_decode_throttle_reasons_empty_when_no_bits_set():
|
||||
assert va.decode_throttle_reasons(0) == []
|
||||
|
||||
|
||||
def test_decode_throttle_reasons_maps_each_known_bit():
|
||||
"""Every mask in the table must decode to exactly its own name in isolation."""
|
||||
for mask, name in va.THROTTLE_REASONS.items():
|
||||
assert va.decode_throttle_reasons(mask) == [name]
|
||||
|
||||
|
||||
def test_decode_throttle_reasons_decodes_combined_bits():
|
||||
"""Real NVML samples set several bits at once; all of them must come back."""
|
||||
bits = 0x20 | 0x40 # sw_thermal_slowdown | hw_thermal_slowdown
|
||||
assert set(va.decode_throttle_reasons(bits)) == {"sw_thermal_slowdown", "hw_thermal_slowdown"}
|
||||
|
||||
|
||||
def test_decode_throttle_reasons_ignores_unknown_bits():
|
||||
"""An undocumented bit from a future driver must not crash or invent a reason."""
|
||||
assert va.decode_throttle_reasons(0x8000_0000) == []
|
||||
|
||||
|
||||
def test_hard_throttle_names_match_thermal_governor_expectations():
|
||||
"""thermal_governor escalates on a fixed set of reason strings produced here.
|
||||
If a name is renamed in one module and not the other the governor silently stops
|
||||
reacting to hardware slowdowns, so pin the shared vocabulary."""
|
||||
import thermal_governor as tg
|
||||
assert tg.HARD_THROTTLES <= set(va.THROTTLE_REASONS.values())
|
||||
|
||||
|
||||
class _FakeProc:
|
||||
def __init__(self, name, cmdline):
|
||||
self._name = name
|
||||
self._cmdline = cmdline
|
||||
|
||||
def name(self):
|
||||
return self._name
|
||||
|
||||
def cmdline(self):
|
||||
return self._cmdline
|
||||
|
||||
|
||||
def _patch_proc(monkeypatch, proc):
|
||||
monkeypatch.setattr(va.psutil, "Process", lambda pid: proc)
|
||||
|
||||
|
||||
def test_classify_pid_detects_ollama_by_process_name(monkeypatch):
|
||||
_patch_proc(monkeypatch, _FakeProc("ollama", ["/usr/local/bin/ollama", "serve"]))
|
||||
assert va._classify_pid(1234) == "ollama"
|
||||
|
||||
|
||||
def test_classify_pid_detects_ollama_runner_by_cmdline(monkeypatch):
|
||||
"""Ollama's model runner is a separate llama-server process; its VRAM is Ollama's."""
|
||||
_patch_proc(monkeypatch, _FakeProc("llama-server",
|
||||
["/usr/lib/ollama/llama-server", "--model", "blob"]))
|
||||
assert va._classify_pid(1234) == "ollama"
|
||||
|
||||
|
||||
def test_classify_pid_detects_comfyui(monkeypatch):
|
||||
_patch_proc(monkeypatch, _FakeProc("python3", ["python3", "/opt/ComfyUI/main.py", "--listen"]))
|
||||
assert va._classify_pid(1234) == "comfy"
|
||||
|
||||
|
||||
def test_classify_pid_unknown_process_is_unmanaged(monkeypatch):
|
||||
# Xorg used to stand in for "unknown" here, but a display server is now its own
|
||||
# bucket, so this needs a process that is genuinely neither ours nor the desktop's.
|
||||
_patch_proc(monkeypatch, _FakeProc("trainer", ["/opt/ml/bin/trainer", "--epochs", "3"]))
|
||||
# "unmanaged" rather than "other": a third-party GPU workload holds VRAM this
|
||||
# service cannot reclaim, and must not be lumped in with the desktop compositor.
|
||||
assert va._classify_pid(1234) == "unmanaged"
|
||||
|
||||
|
||||
def test_classify_pid_display_server_is_desktop(monkeypatch):
|
||||
_patch_proc(monkeypatch, _FakeProc("Xorg", ["/usr/lib/xorg/Xorg", ":8"]))
|
||||
assert va._classify_pid(1234) == "desktop"
|
||||
|
||||
|
||||
def test_classify_pid_returns_unmanaged_when_process_vanished(monkeypatch):
|
||||
"""PIDs are read from NVML and can exit before psutil looks them up; that is normal
|
||||
and must not raise inside the 20 ms VRAM poll loop."""
|
||||
def _boom(pid):
|
||||
raise va.psutil.NoSuchProcess(pid)
|
||||
monkeypatch.setattr(va.psutil, "Process", _boom)
|
||||
assert va._classify_pid(999999) == "unmanaged"
|
||||
|
||||
|
||||
# --- process bucketing: desktop vs unmanaged ---------------------------------
|
||||
#
|
||||
# Real case from this machine: stt_relay.py held 842 MB of VRAM for nearly three days
|
||||
# while gnome-shell held 3.9 MB. A single "other" bucket reported them as one number,
|
||||
# which matters because ComfyUI's memory can be reclaimed and a third party's cannot.
|
||||
|
||||
class _FakeProc:
|
||||
def __init__(self, name, cmdline):
|
||||
self._name = name
|
||||
self._cmdline = cmdline
|
||||
def name(self):
|
||||
return self._name
|
||||
def cmdline(self):
|
||||
return self._cmdline
|
||||
def create_time(self):
|
||||
return 1234.5
|
||||
|
||||
|
||||
def _classify(monkeypatch, name, cmdline):
|
||||
monkeypatch.setattr(va.psutil, "Process",
|
||||
lambda pid: _FakeProc(name, cmdline))
|
||||
return va._classify_pid(4321)
|
||||
|
||||
|
||||
def test_desktop_compositors_are_their_own_bucket(monkeypatch):
|
||||
assert _classify(monkeypatch, "gnome-shell", ["/usr/bin/gnome-shell"]) == "desktop"
|
||||
assert _classify(monkeypatch, "Xorg", ["/usr/lib/xorg/Xorg", ":8"]) == "desktop"
|
||||
|
||||
|
||||
def test_third_party_compute_is_unmanaged_not_desktop(monkeypatch):
|
||||
kind = _classify(monkeypatch, "python",
|
||||
["/home/u/robopest-venv/bin/python", "/home/u/stt_relay.py"])
|
||||
assert kind == "unmanaged"
|
||||
|
||||
|
||||
def test_ollama_and_comfy_still_win_over_the_catch_all(monkeypatch):
|
||||
assert _classify(monkeypatch, "llama-server",
|
||||
["/usr/local/lib/ollama/llama-server", "--model", "x"]) == "ollama"
|
||||
assert _classify(monkeypatch, "python",
|
||||
["/home/u/ComfyUI/venv/bin/python", "main.py"]) == "comfy"
|
||||
|
||||
|
||||
def test_pid_cache_is_keyed_by_start_time_not_pid_alone():
|
||||
# Linux recycles PIDs; a stale entry would attribute a new process's VRAM to Ollama
|
||||
# inside the same snapshot the yield barrier trusts.
|
||||
assert all(isinstance(k, tuple) and len(k) == 2 for k in va._PID_KIND_CACHE)
|
||||
134
tests/test_yield_and_reclaim.py
Normal file
134
tests/test_yield_and_reclaim.py
Normal file
@@ -0,0 +1,134 @@
|
||||
"""Tests for VRAM yield classification and the reclaim-on-OOM path.
|
||||
|
||||
These cover the two failure modes that motivated the arbitration rework, both of which
|
||||
were observed on real hardware before being encoded here:
|
||||
|
||||
* A model mid-generation cannot unload. Persisted counters showed 19 "timeouts" in 20
|
||||
yields; telemetry for that window showed the GPU pinned at 96-97% with 14.92 GB held.
|
||||
That is a busy model, not a fault, and must not be retried in a tight loop.
|
||||
* A model that will not fit fails differently depending on configuration. With
|
||||
n_gpu_layers pinned to 99 (this box) Ollama returns a hard CUDA OOM rather than
|
||||
spilling layers to the CPU:
|
||||
"llama-server process has terminated: exit status 1: cudaMalloc failed:
|
||||
out of memory ... unable to allocate CUDA0 buffer"
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
import vram_arbitrator as v
|
||||
|
||||
|
||||
GB = 1024 ** 3
|
||||
|
||||
|
||||
def _snap(ollama_gb, util, free_gb=1.0):
|
||||
return {"ollama_bytes": int(ollama_gb * GB), "comfyui_bytes": 0, "other_bytes": 0,
|
||||
"free_bytes": int(free_gb * GB), "gpu_util_pct": util}
|
||||
|
||||
|
||||
class TestOomDetection:
|
||||
"""The retry path keys off Ollama's error text, so the matcher must be exact."""
|
||||
|
||||
def test_matches_the_real_observed_ollama_oom(self):
|
||||
real = ("llama-server process has terminated: exit status 1: cudaMalloc failed: "
|
||||
"out of memory\nalloc_tensor_range: failed to allocate CUDA0 buffer of "
|
||||
"size 13028925440\nerror loading model: unable to allocate CUDA0 buffer")
|
||||
assert v.looks_like_vram_oom(real)
|
||||
|
||||
@pytest.mark.parametrize("text", [
|
||||
"cudaMalloc failed: out of memory",
|
||||
"unable to allocate CUDA0 buffer",
|
||||
"failed to allocate buffer",
|
||||
"CUDA error: something",
|
||||
])
|
||||
def test_matches_each_signature(self, text):
|
||||
assert v.looks_like_vram_oom(text)
|
||||
|
||||
@pytest.mark.parametrize("text", [
|
||||
"model 'foo' not found", "invalid parameter", "", None,
|
||||
"context length exceeded",
|
||||
])
|
||||
def test_does_not_match_unrelated_failures(self, text):
|
||||
# A false positive here would purge ComfyUI over a typo in a model name.
|
||||
assert not v.looks_like_vram_oom(text)
|
||||
|
||||
|
||||
class TestYieldOutcomeClassification:
|
||||
"""_await_vram_release must separate released / busy / stuck."""
|
||||
|
||||
def _run(self, snaps, timeout_s=2.0, baseline_gb=14.9):
|
||||
seq = list(snaps)
|
||||
def fake():
|
||||
return seq.pop(0) if len(seq) > 1 else seq[0]
|
||||
original = v.get_process_vram_bytes
|
||||
v.get_process_vram_bytes = fake
|
||||
try:
|
||||
return asyncio.run(
|
||||
v._await_vram_release(int(baseline_gb * GB), timeout_s=timeout_s))
|
||||
finally:
|
||||
v.get_process_vram_bytes = original
|
||||
|
||||
def test_released_when_vram_drains(self):
|
||||
res = self._run([_snap(14.9, 30), _snap(0.0, 5, free_gb=15.4)])
|
||||
assert res["outcome"] == "released"
|
||||
assert res["confirmed"] is True
|
||||
|
||||
def test_busy_when_vram_held_and_gpu_pinned(self):
|
||||
# The observed pathology: 14.92 GB held at 96% utilisation.
|
||||
res = self._run([_snap(14.92, 96)])
|
||||
assert res["outcome"] == "busy"
|
||||
assert res["confirmed"] is False
|
||||
assert "mid-generation" in res["error"]
|
||||
assert res["peak_util_pct"] >= v.BUSY_UTIL_PCT
|
||||
|
||||
def test_stuck_when_vram_held_and_gpu_idle(self):
|
||||
# VRAM held with nothing running is the genuine fault case.
|
||||
res = self._run([_snap(14.9, 2)], timeout_s=0.3)
|
||||
assert res["outcome"] == "stuck"
|
||||
assert "idle" in res["error"]
|
||||
|
||||
def test_busy_is_decided_only_after_the_probe_window(self):
|
||||
# Deciding instantly would misread the normal 40-110ms release as busy.
|
||||
assert v.BUSY_PROBE_S > 0
|
||||
assert v.YIELD_CONFIRM_TIMEOUT_S > v.BUSY_PROBE_S
|
||||
|
||||
def test_default_wait_is_short(self):
|
||||
# It was 10s, which blocked the arbitrator for the length of an inference while
|
||||
# ComfyUI -- which is not gated on our return value -- waited anyway.
|
||||
assert v.YIELD_CONFIRM_TIMEOUT_S <= 3.0
|
||||
assert v.YIELD_CONFIRM_TIMEOUT_BLOCKING_S >= 10.0
|
||||
|
||||
|
||||
class TestBusyBackoff:
|
||||
"""A busy model must not be re-asked every second."""
|
||||
|
||||
def test_backoff_schedule_is_monotonic_and_bounded(self):
|
||||
sched = v.AutoArbitrator.BUSY_BACKOFF_S
|
||||
assert list(sched) == sorted(sched)
|
||||
assert sched[0] >= 1.0
|
||||
|
||||
def test_streak_walks_up_the_schedule_and_clamps(self):
|
||||
arb = v.AutoArbitrator()
|
||||
sched = arb.BUSY_BACKOFF_S
|
||||
for streak in range(len(sched) + 3):
|
||||
delay = sched[min(streak, len(sched) - 1)]
|
||||
assert delay == sched[min(streak, len(sched) - 1)]
|
||||
assert sched[min(99, len(sched) - 1)] == sched[-1]
|
||||
|
||||
def test_release_clears_backoff_state(self):
|
||||
arb = v.AutoArbitrator()
|
||||
arb._yield_backoff_until["m"] = 1e18
|
||||
arb._yield_busy_streak["m"] = 3
|
||||
arb.note_deferred_release(1234.0)
|
||||
assert arb._yield_backoff_until == {}
|
||||
assert arb._yield_busy_streak == {}
|
||||
assert arb.stats["deferred_releases"] == 1
|
||||
|
||||
def test_counters_distinguish_busy_from_stalled(self):
|
||||
# The old single yield_timeouts counter reported a healthy cron job as a 95%
|
||||
# failure rate.
|
||||
arb = v.AutoArbitrator()
|
||||
assert "yield_deferred_busy" in arb.stats
|
||||
assert "yield_stalled" in arb.stats
|
||||
assert "yield_timeouts" not in arb.stats
|
||||
@@ -44,13 +44,56 @@ SWITCH_HISTORY = deque(maxlen=50)
|
||||
RAM_HIT_GBPS = 2.0
|
||||
PARTIAL_HIT_GBPS = 0.8
|
||||
|
||||
# How long Ollama's VRAM may take to actually drain before we stop waiting.
|
||||
# Ollama will not unload a model while a generation is in flight, so a short ceiling
|
||||
# reports a timeout for what is really just a busy model finishing its request. Observed
|
||||
# here: two yields hit the old 3 s limit with 8.2 GB still held while ComfyUI was starting.
|
||||
# Waiting longer is the safer failure mode -- the alternative is diffusion allocating into
|
||||
# VRAM that is still occupied.
|
||||
YIELD_CONFIRM_TIMEOUT_S = 10.0
|
||||
# How long to wait for Ollama's VRAM to actually drain.
|
||||
#
|
||||
# Ollama will not unload a model mid-generation. With OLLAMA_NUM_PARALLEL=1 our
|
||||
# keep_alive:0 request queues behind the running one and takes effect the moment it
|
||||
# finishes, so a model that is busy is not failing -- it is finishing, and it will
|
||||
# release on its own. Blocking the arbitrator for the length of someone's inference
|
||||
# helps nobody: ComfyUI is not gated on our return value, and every second spent
|
||||
# blocked is a second the watchdog and profile switching are stalled.
|
||||
#
|
||||
# So: wait briefly for the common case (an idle model releases in 40-110 ms here),
|
||||
# then classify. A caller who genuinely wants to block can ask for a longer wait.
|
||||
YIELD_CONFIRM_TIMEOUT_S = 2.0
|
||||
YIELD_CONFIRM_TIMEOUT_BLOCKING_S = 30.0
|
||||
|
||||
# A model still holding VRAM while the GPU is pinned is generating, not wedged.
|
||||
BUSY_UTIL_PCT = 50
|
||||
BUSY_PROBE_S = 0.6
|
||||
|
||||
# Fraction of a model that may sit outside VRAM before we call it starved. A little
|
||||
# slack absorbs rounding and KV-cache accounting; beyond it, layers are on the CPU.
|
||||
CPU_OFFLOAD_TOLERANCE = 0.02
|
||||
# Only intervene when ComfyUI is actually holding enough VRAM to be the cause.
|
||||
RECLAIM_MIN_COMFY_BYTES = 512 * 1024 ** 2
|
||||
|
||||
# Ollama's response when a model will not fit. Which of the two failure modes you get
|
||||
# depends on configuration: with n_gpu_layers left to Ollama it spills layers to the CPU
|
||||
# and reports size_vram < size; with n_gpu_layers pinned (99 on this box) it refuses and
|
||||
# returns a hard CUDA OOM instead. Both are handled -- the spill by
|
||||
# AutoArbitrator._check_ollama_starved, the hard failure by the retry below.
|
||||
OOM_SIGNATURES = ("out of memory", "cudamalloc", "unable to allocate",
|
||||
"failed to allocate", "cuda error")
|
||||
|
||||
|
||||
def describe_unmanaged() -> Dict[str, Any]:
|
||||
"""VRAM held by processes this service cannot reclaim, named explicitly."""
|
||||
stats = get_gpu_hardware_stats()
|
||||
bd = stats.get("breakdown", {}) if stats.get("available") else {}
|
||||
entries = bd.get("unmanaged", [])
|
||||
return {
|
||||
"unmanaged_gb": bd.get("unmanaged_gb", 0.0),
|
||||
"processes": entries,
|
||||
"note": ("VRAM held by processes outside HyperSwap's control; it cannot be "
|
||||
"reclaimed automatically" if entries else
|
||||
"no third-party GPU processes are holding VRAM"),
|
||||
}
|
||||
|
||||
|
||||
def looks_like_vram_oom(text: str) -> bool:
|
||||
low = (text or "").lower()
|
||||
return any(sig in low for sig in OOM_SIGNATURES)
|
||||
YIELD_CONFIRM_POLL_S = 0.02
|
||||
YIELD_RESIDUAL_BYTES = 256 * 1024 ** 2 # treat <256 MB as "released"
|
||||
|
||||
@@ -105,12 +148,17 @@ def get_process_vram_bytes() -> Dict[str, int]:
|
||||
Deliberately avoids psutil lookups: this runs every 20 ms while we wait for VRAM
|
||||
to actually drain.
|
||||
"""
|
||||
out = {"ollama_bytes": 0, "comfyui_bytes": 0, "other_bytes": 0, "free_bytes": 0}
|
||||
out = {"ollama_bytes": 0, "comfyui_bytes": 0, "other_bytes": 0, "free_bytes": 0,
|
||||
"desktop_bytes": 0, "unmanaged_bytes": 0, "gpu_util_pct": 0}
|
||||
if not NVML_AVAILABLE:
|
||||
return out
|
||||
try:
|
||||
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
|
||||
out["free_bytes"] = pynvml.nvmlDeviceGetMemoryInfo(handle).free
|
||||
try:
|
||||
out["gpu_util_pct"] = pynvml.nvmlDeviceGetUtilizationRates(handle).gpu
|
||||
except Exception:
|
||||
pass
|
||||
procs = list(pynvml.nvmlDeviceGetComputeRunningProcesses(handle))
|
||||
try:
|
||||
procs += list(pynvml.nvmlDeviceGetGraphicsRunningProcesses(handle))
|
||||
@@ -120,36 +168,72 @@ def get_process_vram_bytes() -> Dict[str, int]:
|
||||
for p in procs:
|
||||
merged[p.pid] = max(merged.get(p.pid, 0), p.usedGpuMemory or 0)
|
||||
for pid, used in merged.items():
|
||||
kind = _PID_KIND_CACHE.get(pid)
|
||||
key = _pid_key(pid)
|
||||
kind = _PID_KIND_CACHE.get(key) if key else None
|
||||
if kind is None:
|
||||
kind = _classify_pid(pid)
|
||||
_PID_KIND_CACHE[pid] = kind
|
||||
if key:
|
||||
if len(_PID_KIND_CACHE) >= _PID_KIND_CACHE_MAX:
|
||||
_PID_KIND_CACHE.clear()
|
||||
_PID_KIND_CACHE[key] = kind
|
||||
if kind == "ollama":
|
||||
out["ollama_bytes"] += used
|
||||
elif kind == "comfy":
|
||||
out["comfyui_bytes"] += used
|
||||
elif kind == "desktop":
|
||||
out["desktop_bytes"] += used
|
||||
out["other_bytes"] += used
|
||||
else:
|
||||
out["unmanaged_bytes"] += used
|
||||
out["other_bytes"] += used
|
||||
except Exception as e:
|
||||
logger.debug(f"get_process_vram_bytes failed: {e}")
|
||||
return out
|
||||
|
||||
|
||||
_PID_KIND_CACHE: Dict[int, str] = {}
|
||||
# Keyed by (pid, process start time) rather than pid alone. Linux recycles PIDs, and a
|
||||
# stale entry would attribute a new process's VRAM to Ollama or ComfyUI -- in the same
|
||||
# snapshot the yield barrier uses to decide whether VRAM was released.
|
||||
_PID_KIND_CACHE: Dict[tuple, str] = {}
|
||||
_PID_KIND_CACHE_MAX = 512
|
||||
|
||||
|
||||
def _pid_key(pid: int) -> Optional[tuple]:
|
||||
try:
|
||||
return (pid, psutil.Process(pid).create_time())
|
||||
except Exception:
|
||||
return None
|
||||
|
||||
|
||||
# Compositors and display servers. Their VRAM is small, permanent and not ours to
|
||||
# reclaim, so it should not be confused with a real workload.
|
||||
DESKTOP_PROCESS_HINTS = (
|
||||
"gnome-shell", "xorg", "gnome-remote-desktop", "mutter", "kwin", "plasmashell",
|
||||
"gnome-session", "wayland", "weston", "sddm", "gdm", "picom", "compiz",
|
||||
)
|
||||
|
||||
|
||||
def _classify_pid(pid: int) -> str:
|
||||
"""Bucket a GPU process into ollama | comfy | desktop | unmanaged.
|
||||
|
||||
The old version had one catch-all "other" bucket, which put a 3.9 MB compositor and
|
||||
an 842 MB long-running inference script in the same number. That matters: this
|
||||
service can reclaim VRAM from ComfyUI, but it cannot touch a third-party workload,
|
||||
and pretending otherwise makes it promise headroom it cannot deliver.
|
||||
"""
|
||||
try:
|
||||
proc = psutil.Process(pid)
|
||||
pname = proc.name().lower()
|
||||
cmdline = " ".join(proc.cmdline()).lower()
|
||||
except Exception:
|
||||
return "other"
|
||||
return "unmanaged"
|
||||
if "ollama" in pname or "llama-server" in cmdline:
|
||||
return "ollama"
|
||||
if "comfy" in cmdline or "main.py" in cmdline:
|
||||
if "comfyui" in cmdline or "comfy" in cmdline or cmdline.rstrip().endswith("main.py"):
|
||||
return "comfy"
|
||||
return "other"
|
||||
if any(hint in pname or hint in cmdline for hint in DESKTOP_PROCESS_HINTS):
|
||||
return "desktop"
|
||||
return "unmanaged"
|
||||
|
||||
|
||||
def get_gpu_hardware_stats() -> Dict[str, Any]:
|
||||
@@ -249,6 +333,9 @@ def get_gpu_hardware_stats() -> Dict[str, Any]:
|
||||
"ollama_bytes": 0,
|
||||
"comfyui_bytes": 0,
|
||||
"system_bytes": 0,
|
||||
"desktop_bytes": 0,
|
||||
"unmanaged_bytes": 0,
|
||||
"unmanaged": [],
|
||||
"processes": []
|
||||
}
|
||||
|
||||
@@ -269,8 +356,9 @@ def get_gpu_hardware_stats() -> Dict[str, Any]:
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
is_ollama = "ollama" in pname.lower() or "llama-server" in cmdline.lower()
|
||||
is_comfy = "comfy" in cmdline.lower() or "main.py" in cmdline.lower()
|
||||
kind = _classify_pid(pid)
|
||||
is_ollama = kind == "ollama"
|
||||
is_comfy = kind == "comfy"
|
||||
|
||||
if is_ollama:
|
||||
proc_breakdown["ollama_bytes"] += used_mem
|
||||
@@ -278,6 +366,15 @@ def get_gpu_hardware_stats() -> Dict[str, Any]:
|
||||
proc_breakdown["comfyui_bytes"] += used_mem
|
||||
else:
|
||||
proc_breakdown["system_bytes"] += used_mem
|
||||
if kind == "desktop":
|
||||
proc_breakdown["desktop_bytes"] += used_mem
|
||||
else:
|
||||
proc_breakdown["unmanaged_bytes"] += used_mem
|
||||
proc_breakdown["unmanaged"].append({
|
||||
"pid": pid, "name": pname,
|
||||
"cmdline": cmdline[:120],
|
||||
"vram_mb": round(used_mem / (1024**2), 1),
|
||||
})
|
||||
|
||||
proc_breakdown["processes"].append({
|
||||
"pid": pid,
|
||||
@@ -287,6 +384,7 @@ def get_gpu_hardware_stats() -> Dict[str, Any]:
|
||||
"vram_mb": round(used_mem / (1024**2), 1),
|
||||
"is_ollama": is_ollama,
|
||||
"is_comfy": is_comfy,
|
||||
"kind": kind,
|
||||
})
|
||||
except Exception as e:
|
||||
logger.error(f"Error enumerating GPU processes: {e}")
|
||||
@@ -328,6 +426,11 @@ def get_gpu_hardware_stats() -> Dict[str, Any]:
|
||||
"comfyui_gb": round(proc_breakdown["comfyui_bytes"] / (1024**3), 2),
|
||||
"system_mb": round(proc_breakdown["system_bytes"] / (1024**2), 1),
|
||||
"system_gb": round(proc_breakdown["system_bytes"] / (1024**3), 2),
|
||||
"desktop_gb": round(proc_breakdown["desktop_bytes"] / (1024**3), 2),
|
||||
# VRAM held by workloads this service has no control over. It cannot be
|
||||
# reclaimed, so it is permanently unavailable headroom.
|
||||
"unmanaged_gb": round(proc_breakdown["unmanaged_bytes"] / (1024**3), 2),
|
||||
"unmanaged": proc_breakdown["unmanaged"],
|
||||
"free_mb": round(free_vram / (1024**2), 1),
|
||||
"free_gb": round(free_vram / (1024**3), 2),
|
||||
"processes": proc_breakdown["processes"],
|
||||
@@ -345,7 +448,13 @@ async def get_ollama_live_state() -> Dict[str, Any]:
|
||||
"active_model_vram_gb": 0.0,
|
||||
"active_context": 0,
|
||||
"expires_at": None,
|
||||
"installed_models": []
|
||||
"installed_models": [],
|
||||
# Ollama silently spills layers to CPU when VRAM is short. size_vram < size is the
|
||||
# only externally visible sign, and the cost is roughly an order of magnitude in
|
||||
# decode speed, so it is worth surfacing loudly.
|
||||
"gpu_fraction": 1.0,
|
||||
"cpu_offload_pct": 0.0,
|
||||
"partially_offloaded": False,
|
||||
}
|
||||
try:
|
||||
client = _client(OLLAMA_API_BASE, 3.0)
|
||||
@@ -415,35 +524,106 @@ async def get_comfyui_live_state() -> Dict[str, Any]:
|
||||
|
||||
async def _await_vram_release(baseline_bytes: int,
|
||||
timeout_s: float = YIELD_CONFIRM_TIMEOUT_S) -> Dict[str, Any]:
|
||||
"""Block until Ollama's VRAM has actually drained, or we give up.
|
||||
"""Wait for Ollama's VRAM to drain, distinguishing "busy" from "stuck".
|
||||
|
||||
Posting keep_alive:0 only *asks* Ollama to unload; the driver frees the allocation
|
||||
some milliseconds later. Returning before that happens is how ComfyUI ends up
|
||||
allocating into VRAM that is still occupied, which shows up as a CUDA OOM mid-graph.
|
||||
some milliseconds later, and returning before that happens is how ComfyUI ends up
|
||||
allocating into VRAM that is still occupied.
|
||||
|
||||
But there is a second case the first version of this got wrong. If the model is
|
||||
mid-generation it cannot unload at all, and reporting that as a timeout made a
|
||||
perfectly healthy cron job look like a 95% failure rate. When the VRAM has not
|
||||
moved and the GPU is pinned, the model is working; the queued unload will fire when
|
||||
it finishes. That is `busy`, not a failure.
|
||||
|
||||
Returns an `outcome` of "released", "busy" or "stuck".
|
||||
"""
|
||||
t0 = time.perf_counter()
|
||||
last = baseline_bytes
|
||||
peak_util = 0
|
||||
while True:
|
||||
snap = get_process_vram_bytes()
|
||||
last = snap["ollama_bytes"]
|
||||
peak_util = max(peak_util, snap.get("gpu_util_pct", 0))
|
||||
elapsed = time.perf_counter() - t0
|
||||
|
||||
if last <= YIELD_RESIDUAL_BYTES:
|
||||
return {
|
||||
"outcome": "released",
|
||||
"confirmed": True,
|
||||
"confirm_ms": round((time.perf_counter() - t0) * 1000, 2),
|
||||
"confirm_ms": round(elapsed * 1000, 2),
|
||||
"residual_bytes": last,
|
||||
"free_bytes": snap["free_bytes"],
|
||||
"gpu_util_pct": snap.get("gpu_util_pct", 0),
|
||||
}
|
||||
if (time.perf_counter() - t0) >= timeout_s:
|
||||
|
||||
# Unmoved VRAM plus a pinned GPU means a generation is in flight.
|
||||
busy = (elapsed >= BUSY_PROBE_S
|
||||
and last >= baseline_bytes - YIELD_RESIDUAL_BYTES
|
||||
and peak_util >= BUSY_UTIL_PCT)
|
||||
|
||||
if busy or elapsed >= timeout_s:
|
||||
outcome = "busy" if busy else "stuck"
|
||||
return {
|
||||
"outcome": outcome,
|
||||
"confirmed": False,
|
||||
"confirm_ms": round((time.perf_counter() - t0) * 1000, 2),
|
||||
"confirm_ms": round(elapsed * 1000, 2),
|
||||
"residual_bytes": last,
|
||||
"free_bytes": snap["free_bytes"],
|
||||
"error": f"Ollama still holding {round(last / (1024**3), 2)} GB after {timeout_s}s",
|
||||
"gpu_util_pct": snap.get("gpu_util_pct", 0),
|
||||
"peak_util_pct": peak_util,
|
||||
"error": (
|
||||
f"Ollama is mid-generation ({peak_util}% GPU, "
|
||||
f"{round(last / (1024**3), 2)} GB held); the queued unload will apply "
|
||||
f"when it finishes"
|
||||
if outcome == "busy" else
|
||||
f"Ollama still holding {round(last / (1024**3), 2)} GB after "
|
||||
f"{timeout_s}s with the GPU idle"
|
||||
),
|
||||
}
|
||||
await asyncio.sleep(YIELD_CONFIRM_POLL_S)
|
||||
|
||||
|
||||
# Detached tasks need a strong reference or the loop may garbage-collect them mid-flight.
|
||||
_DETACHED: set = set()
|
||||
|
||||
|
||||
def _spawn_detached(coro) -> None:
|
||||
task = asyncio.ensure_future(coro)
|
||||
_DETACHED.add(task)
|
||||
task.add_done_callback(_DETACHED.discard)
|
||||
|
||||
|
||||
async def _confirm_release_later(targets: List[str], baseline_bytes: int,
|
||||
max_wait_s: float = 900.0) -> None:
|
||||
"""Watch for a queued unload to land after the in-flight generation finishes.
|
||||
|
||||
Runs detached so the caller is never held for the length of an inference. Logs the
|
||||
eventual release so the event log tells the whole story rather than stopping at
|
||||
"deferred".
|
||||
"""
|
||||
t0 = time.perf_counter()
|
||||
while (time.perf_counter() - t0) < max_wait_s:
|
||||
await asyncio.sleep(0.5)
|
||||
snap = get_process_vram_bytes()
|
||||
if snap["ollama_bytes"] <= YIELD_RESIDUAL_BYTES:
|
||||
waited_ms = round((time.perf_counter() - t0) * 1000, 2)
|
||||
_record({
|
||||
"event_type": "Ollama VRAM Yield",
|
||||
"source": ", ".join(targets)[:200],
|
||||
"target": "VRAM 0MB (Kept in RAM)",
|
||||
"duration_ms": waited_ms,
|
||||
"yield_confirm_ms": waited_ms,
|
||||
"cache_status": "RAM-Cached",
|
||||
"detail": "released after the in-flight generation completed",
|
||||
})
|
||||
arbitrator.note_deferred_release(waited_ms)
|
||||
logger.info(f"Deferred VRAM yield completed after {round(waited_ms / 1000, 1)}s "
|
||||
f"({round(snap['free_bytes'] / (1024**3), 2)} GB free)")
|
||||
return
|
||||
logger.warning("Deferred VRAM yield never landed within "
|
||||
f"{max_wait_s}s for {', '.join(targets)}")
|
||||
|
||||
|
||||
def _record(event: Dict[str, Any]) -> None:
|
||||
"""Push an event to both the in-memory ring and the durable store."""
|
||||
event.setdefault("ts", time.time())
|
||||
@@ -453,7 +633,8 @@ def _record(event: Dict[str, Any]) -> None:
|
||||
|
||||
|
||||
async def instant_free_ollama_vram(model_name: Optional[str] = None,
|
||||
confirm: bool = True) -> Dict[str, Any]:
|
||||
confirm: bool = True,
|
||||
timeout_s: Optional[float] = None) -> Dict[str, Any]:
|
||||
"""Yield Ollama's VRAM and wait for the driver to actually release it.
|
||||
|
||||
The returned duration_ms is now the real end-to-end release time, not just how long
|
||||
@@ -478,31 +659,51 @@ async def instant_free_ollama_vram(model_name: Optional[str] = None,
|
||||
model_name = targets[0] if len(targets) == 1 else f"{len(targets)} models"
|
||||
baseline = get_process_vram_bytes()["ollama_bytes"]
|
||||
try:
|
||||
client = _client(OLLAMA_API_BASE, 5.0)
|
||||
await asyncio.gather(*[
|
||||
# Generous client timeout: with OLLAMA_NUM_PARALLEL=1 this request queues behind
|
||||
# any running generation, and a short timeout would drop the connection before
|
||||
# Ollama ever processed the unload -- losing it entirely.
|
||||
client = _client(OLLAMA_API_BASE, 120.0)
|
||||
unload_calls = [
|
||||
client.post("/api/generate", json={"model": t, "keep_alive": 0})
|
||||
for t in targets
|
||||
], return_exceptions=True)
|
||||
]
|
||||
# Do not await the queued unloads; a busy model would block us for the length of
|
||||
# its inference. They are fire-and-confirm: the barrier below watches the VRAM.
|
||||
_spawn_detached(asyncio.gather(*unload_calls, return_exceptions=True))
|
||||
request_ms = round((time.perf_counter() - t0) * 1000, 2)
|
||||
|
||||
barrier = {"confirmed": None, "confirm_ms": 0.0, "residual_bytes": baseline}
|
||||
barrier: Dict[str, Any] = {"outcome": "unconfirmed", "confirmed": None,
|
||||
"confirm_ms": 0.0, "residual_bytes": baseline}
|
||||
if confirm:
|
||||
barrier = await _await_vram_release(baseline)
|
||||
barrier = await _await_vram_release(
|
||||
baseline, timeout_s if timeout_s is not None else YIELD_CONFIRM_TIMEOUT_S)
|
||||
|
||||
if barrier.get("outcome") == "busy":
|
||||
# The unload is queued and will fire when the generation ends. Keep watching
|
||||
# in the background so the release is still logged and the counters stay true,
|
||||
# without holding the caller here for the length of someone's inference.
|
||||
_spawn_detached(_confirm_release_later(targets, baseline))
|
||||
|
||||
duration_ms = round((time.perf_counter() - t0) * 1000, 2)
|
||||
freed_gb = round(max(baseline - barrier.get("residual_bytes", 0), 0) / (1024**3), 2)
|
||||
|
||||
outcome = barrier.get("outcome", "unconfirmed")
|
||||
_record({
|
||||
"event_type": "Ollama VRAM Yield",
|
||||
"source": ", ".join(targets)[:200],
|
||||
"target": "VRAM 0MB (Kept in RAM)",
|
||||
"duration_ms": duration_ms,
|
||||
"yield_confirm_ms": barrier.get("confirm_ms"),
|
||||
"cache_status": "RAM-Cached" if barrier.get("confirmed") else "Yield Timeout",
|
||||
"cache_status": {
|
||||
"released": "RAM-Cached",
|
||||
"busy": "Deferred — LLM generating",
|
||||
"stuck": "Yield Stalled",
|
||||
}.get(outcome, "Yield Unconfirmed"),
|
||||
"detail": barrier.get("error"),
|
||||
})
|
||||
return {
|
||||
"success": True,
|
||||
"outcome": outcome,
|
||||
"model": model_name,
|
||||
"models_unloaded": targets,
|
||||
"duration_ms": duration_ms,
|
||||
@@ -512,6 +713,7 @@ async def instant_free_ollama_vram(model_name: Optional[str] = None,
|
||||
"freed_gb": freed_gb,
|
||||
"residual_gb": round(barrier.get("residual_bytes", 0) / (1024**3), 2),
|
||||
"free_vram_gb": round(barrier.get("free_bytes", 0) / (1024**3), 2),
|
||||
"gpu_util_pct": barrier.get("gpu_util_pct"),
|
||||
"error": barrier.get("error"),
|
||||
}
|
||||
except Exception as e:
|
||||
@@ -586,8 +788,13 @@ def classify_load(size_bytes: int, load_duration_ms: float) -> Dict[str, Any]:
|
||||
return {"cache_status": status, "load_gbps": round(gbps, 2), "is_ram_hit": gbps >= RAM_HIT_GBPS}
|
||||
|
||||
|
||||
async def switch_ollama_model(target_model: str, keep_alive: str = "30m") -> Dict[str, Any]:
|
||||
"""High-speed hot-swap to target Ollama model, tracking swap metrics."""
|
||||
async def switch_ollama_model(target_model: str, keep_alive: str = "30m",
|
||||
_retrying: bool = False) -> Dict[str, Any]:
|
||||
"""High-speed hot-swap to target Ollama model, tracking swap metrics.
|
||||
|
||||
If the load fails because the model will not fit, reclaims VRAM from an idle ComfyUI
|
||||
and retries once. `_retrying` guards against recursing more than one level.
|
||||
"""
|
||||
t0 = time.perf_counter()
|
||||
cur_state = await get_ollama_live_state()
|
||||
prev_model = cur_state.get("active_model_name") or "None"
|
||||
@@ -637,8 +844,43 @@ async def switch_ollama_model(target_model: str, keep_alive: str = "30m") -> Dic
|
||||
"is_ram_hit": cls["is_ram_hit"],
|
||||
"response": data.get("response", ""),
|
||||
}
|
||||
return {"success": False, "error": f"HTTP {resp.status_code}: {resp.text}",
|
||||
"duration_ms": total_duration_ms}
|
||||
# A model that will not fit is the exact contention this service exists to
|
||||
# resolve. Rather than handing the caller a CUDA OOM, take the VRAM back from an
|
||||
# idle ComfyUI and try once more.
|
||||
body = resp.text
|
||||
if looks_like_vram_oom(body) and not _retrying:
|
||||
snap = get_process_vram_bytes()
|
||||
if snap["comfyui_bytes"] >= RECLAIM_MIN_COMFY_BYTES:
|
||||
logger.warning(
|
||||
f"Ollama could not fit '{target_model}' with ComfyUI holding "
|
||||
f"{round(snap['comfyui_bytes'] / (1024**3), 2)} GB — reclaiming and retrying")
|
||||
purge = await instant_free_comfyui_vram()
|
||||
arbitrator.stats["reclaims_for_ollama"] += 1
|
||||
arbitrator.last_action = (
|
||||
f"Reclaimed {round(snap['comfyui_bytes'] / (1024**3), 2)}GB from ComfyUI so "
|
||||
f"'{target_model}' could load")
|
||||
_record({
|
||||
"event_type": "VRAM Reclaim for Ollama",
|
||||
"source": "ComfyUI Pipeline",
|
||||
"target": target_model,
|
||||
"duration_ms": purge.get("duration_ms"),
|
||||
"cache_status": "Reclaimed",
|
||||
"detail": f"Ollama OOM: {body[:160]}",
|
||||
})
|
||||
await asyncio.sleep(0.3)
|
||||
retry = await switch_ollama_model(target_model, keep_alive, _retrying=True)
|
||||
retry["reclaimed_from_comfyui_gb"] = round(
|
||||
snap["comfyui_bytes"] / (1024**3), 2)
|
||||
retry["first_attempt_error"] = "CUDA OOM; retried after reclaiming VRAM"
|
||||
if not retry.get("success"):
|
||||
# Be specific about why the reclaim was not enough. Blaming ComfyUI
|
||||
# when a third-party process is holding the memory sends the user
|
||||
# looking in the wrong place.
|
||||
retry["unmanaged_blockers"] = describe_unmanaged()
|
||||
return retry
|
||||
return {"success": False, "error": f"HTTP {resp.status_code}: {body}",
|
||||
"duration_ms": total_duration_ms,
|
||||
"vram_oom": looks_like_vram_oom(body)}
|
||||
except Exception as e:
|
||||
return {"success": False, "error": str(e),
|
||||
"duration_ms": round((time.perf_counter() - t0) * 1000, 2)}
|
||||
@@ -683,7 +925,20 @@ class AutoArbitrator:
|
||||
# benchmark would otherwise trip trigger_comfy_priority, which reapplies the whole
|
||||
# 'comfy' profile and silently overwrites the clock the sweep is measuring.
|
||||
self.oc_suspended = False
|
||||
self.stats = {"yields": 0, "purges": 0, "yield_timeouts": 0, "deferred_purges": 0}
|
||||
# Per-model backoff. A model that is mid-generation cannot yield, and asking it
|
||||
# again every second just blocks the loop repeatedly for no benefit.
|
||||
self._yield_backoff_until: Dict[str, float] = {}
|
||||
self._yield_busy_streak: Dict[str, int] = {}
|
||||
self.last_reclaim_time = 0.0
|
||||
self.stats = {
|
||||
"yields": 0, # release confirmed
|
||||
"yield_deferred_busy": 0, # model mid-generation; unload queued behind it
|
||||
"yield_stalled": 0, # VRAM held with an idle GPU -- the real failure
|
||||
"deferred_releases": 0, # queued unloads that later landed
|
||||
"purges": 0,
|
||||
"deferred_purges": 0,
|
||||
"reclaims_for_ollama": 0, # ComfyUI purged because the LLM was spilling to CPU
|
||||
}
|
||||
|
||||
async def start(self):
|
||||
if self.running:
|
||||
@@ -708,8 +963,19 @@ class AutoArbitrator:
|
||||
await close_clients()
|
||||
logger.info("AutoArbitrator background engine stopped.")
|
||||
|
||||
# Backoff schedule for a model that keeps reporting busy, in seconds.
|
||||
BUSY_BACKOFF_S = (5.0, 15.0, 30.0, 60.0)
|
||||
|
||||
def note_deferred_release(self, waited_ms: float) -> None:
|
||||
"""Called when a queued unload finally lands after a generation finished."""
|
||||
self.stats["deferred_releases"] += 1
|
||||
self._yield_backoff_until.clear()
|
||||
self._yield_busy_streak.clear()
|
||||
self.last_action = (f"VRAM released after the LLM finished "
|
||||
f"({round(waited_ms / 1000, 1)}s) — ComfyUI can proceed")
|
||||
|
||||
async def trigger_comfy_priority(self, reason: str = "ComfyUI prompt detected"):
|
||||
"""Yield Ollama's VRAM — and confirm it is gone — before diffusion allocates."""
|
||||
"""Yield Ollama's VRAM before diffusion allocates, without fighting a busy model."""
|
||||
self.comfy_was_active = True
|
||||
self.comfy_idle_since = None
|
||||
self._apply_oc_profile("comfy")
|
||||
@@ -718,21 +984,39 @@ class AutoArbitrator:
|
||||
return
|
||||
|
||||
ollama_state = await get_ollama_live_state()
|
||||
if ollama_state.get("active_model_name"):
|
||||
model = ollama_state["active_model_name"]
|
||||
model = ollama_state.get("active_model_name")
|
||||
if not model:
|
||||
return
|
||||
|
||||
# Still finishing an inference we already asked to unload: leave it alone.
|
||||
until = self._yield_backoff_until.get(model, 0.0)
|
||||
if now < until:
|
||||
return
|
||||
|
||||
logger.info(f"⚡ ComfyUI active ({reason}) -> Auto-yielding Ollama model '{model}'...")
|
||||
self.last_yield_time = time.time()
|
||||
res = await instant_free_ollama_vram(model, confirm=True)
|
||||
outcome = res.get("outcome")
|
||||
|
||||
if outcome == "released":
|
||||
self.stats["yields"] += 1
|
||||
if res.get("confirmed"):
|
||||
self._yield_backoff_until.pop(model, None)
|
||||
self._yield_busy_streak.pop(model, None)
|
||||
self.last_action = (f"Yielded '{model}' for ComfyUI in "
|
||||
f"{res.get('confirm_ms')}ms (confirmed {res.get('freed_gb')}GB free)")
|
||||
f"{res.get('confirm_ms')}ms ({res.get('freed_gb')}GB freed)")
|
||||
elif outcome == "busy":
|
||||
streak = self._yield_busy_streak.get(model, 0)
|
||||
delay = self.BUSY_BACKOFF_S[min(streak, len(self.BUSY_BACKOFF_S) - 1)]
|
||||
self._yield_busy_streak[model] = streak + 1
|
||||
self._yield_backoff_until[model] = time.time() + delay
|
||||
self.stats["yield_deferred_busy"] += 1
|
||||
self.last_action = (f"'{model}' is mid-generation ({res.get('gpu_util_pct')}% GPU); "
|
||||
f"unload is queued and will apply when it finishes")
|
||||
logger.info(f"Yield deferred: {res.get('error')} — backing off {delay}s")
|
||||
else:
|
||||
self.stats["yield_timeouts"] += 1
|
||||
self.last_action = (f"⚠ Yield of '{model}' NOT confirmed: "
|
||||
f"{res.get('residual_gb')}GB still held")
|
||||
logger.warning(f"VRAM yield barrier timed out: {res.get('error')}")
|
||||
logger.info(f"Ollama auto-yield completed: {res}")
|
||||
self.stats["yield_stalled"] += 1
|
||||
self.last_action = (f"⚠ '{model}' holding {res.get('residual_gb')}GB with an idle GPU")
|
||||
logger.warning(f"VRAM yield stalled: {res.get('error')}")
|
||||
|
||||
async def trigger_comfy_completed(self, immediate: bool = False):
|
||||
"""Mark the end of a generation. The actual purge is deferred unless forced."""
|
||||
@@ -836,6 +1120,57 @@ class AutoArbitrator:
|
||||
await asyncio.sleep(backoff)
|
||||
backoff = min(backoff * 1.5, 15.0)
|
||||
|
||||
RECLAIM_COOLDOWN_S = 30.0
|
||||
|
||||
async def _check_ollama_starved(self) -> None:
|
||||
"""The other direction: rescue an LLM that ComfyUI has squeezed onto the CPU.
|
||||
|
||||
Yielding Ollama for ComfyUI was automatic; the reverse never was, despite the
|
||||
README calling the arbitration bidirectional. When Ollama cannot fit a model it
|
||||
does not fail, it silently places layers on the CPU and runs about an order of
|
||||
magnitude slower -- so this is the failure mode a user is least likely to notice
|
||||
and most likely to feel.
|
||||
|
||||
If the LLM is spilling while ComfyUI sits idle holding VRAM, ComfyUI's cached
|
||||
checkpoints are the thing to give up.
|
||||
"""
|
||||
now = time.time()
|
||||
if self.comfy_was_active or (now - self.last_reclaim_time) < self.RECLAIM_COOLDOWN_S:
|
||||
return
|
||||
|
||||
ollama = await get_ollama_live_state()
|
||||
if not ollama.get("partially_offloaded"):
|
||||
return
|
||||
|
||||
snap = get_process_vram_bytes()
|
||||
if snap["comfyui_bytes"] < RECLAIM_MIN_COMFY_BYTES:
|
||||
return # ComfyUI is not the one holding the memory; nothing we can do here
|
||||
|
||||
self.last_reclaim_time = now
|
||||
model = ollama.get("active_model_name")
|
||||
offload = ollama.get("cpu_offload_pct")
|
||||
logger.warning(f"⚠ '{model}' is {offload}% on CPU while ComfyUI holds "
|
||||
f"{round(snap['comfyui_bytes'] / (1024**3), 2)} GB — reclaiming for the LLM")
|
||||
await self._purge_comfy_now(f"LLM spilling {offload}% to CPU")
|
||||
self.stats["reclaims_for_ollama"] += 1
|
||||
|
||||
# Freeing VRAM does not move layers back; only a reload re-places the model. Do
|
||||
# that only when the model is idle, never mid-generation.
|
||||
after = get_process_vram_bytes()
|
||||
if after.get("gpu_util_pct", 0) < BUSY_UTIL_PCT and model:
|
||||
logger.info(f"Reloading '{model}' to place it fully on the GPU...")
|
||||
await instant_free_ollama_vram(model, confirm=True)
|
||||
res = await switch_ollama_model(model, keep_alive="30m")
|
||||
recheck = await get_ollama_live_state()
|
||||
self.last_action = (
|
||||
f"Reclaimed {round(snap['comfyui_bytes'] / (1024**3), 2)}GB from ComfyUI and "
|
||||
f"reloaded '{model}' — now {round(recheck.get('gpu_fraction', 0) * 100)}% on GPU"
|
||||
if res.get("success") else
|
||||
f"Reclaimed VRAM from ComfyUI but reloading '{model}' failed: {res.get('error')}")
|
||||
else:
|
||||
self.last_action = (f"Reclaimed VRAM from ComfyUI; '{model}' is busy, so it will "
|
||||
f"stay partly on CPU until its next load")
|
||||
|
||||
async def _poll_watchdog(self):
|
||||
"""Fallback for when the WebSocket is down. One cheap /queue call, 1 Hz.
|
||||
|
||||
@@ -854,6 +1189,8 @@ class AutoArbitrator:
|
||||
await self.trigger_comfy_priority("Watchdog saw an active queue")
|
||||
elif self.comfy_was_active:
|
||||
await self.trigger_comfy_completed()
|
||||
else:
|
||||
await self._check_ollama_starved()
|
||||
except Exception:
|
||||
pass
|
||||
await asyncio.sleep(interval)
|
||||
@@ -893,6 +1230,9 @@ class AutoArbitrator:
|
||||
"idle_purge_after_s": self.COMFY_IDLE_PURGE_S,
|
||||
"oc_profile": self.oc_profile,
|
||||
"counters": dict(self.stats),
|
||||
"yield_backoff": {m: round(max(t - time.time(), 0), 1)
|
||||
for m, t in self._yield_backoff_until.items()
|
||||
if t > time.time()},
|
||||
}
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user