Compare commits

..

7 Commits

Author SHA1 Message Date
drjones
b53b026ced Stop the dashboard reporting things that are not true
A screenshot of the running dashboard showed four claims that the backend does not
support. Each is the same failure the rest of this branch has been correcting:
displaying a number without checking what it measures.

"Models in RAM: 54.31 GB" was ram.cached_gb -- the entire Linux page cache, every
file the kernel has cached, not models. Measured model residency at the same moment
was 42.87 GB of a 331 GB catalogue. The two are different quantities that happened
to look similar. The legend now says "Page Cache (all files)" for what that number
is, and shows measured model residency separately beside it.

"LAST SWAP TIME 1.65 ms" and "RAM HIT STATUS: Cleaned" were read from history[0],
the most recent event of any type. Both tiles were describing a ComfyUI VRAM purge:
its duration under a swap-time label, its status under a cache-hit label. They now
select the most recent LLM Model Switch, and fall back to the persisted log when the
in-memory ring is empty, so a restart no longer blanks them.

"Host Pinned Memory: 53.6 GB Staging Buffer", "Async PCIe Offloading: Enabled (2
Streams)" and "Fast Disk RAM Mmap: Active" were hardcoded in the HTML and measured
from nothing at all. Replaced with values the service actually has: VRAM currently
held by ComfyUI, live PCIe throughput from NVML, measured checkpoint residency, and
the real idle-purge countdown.

Also: the VRAM legend's "System" bucket hid an 842 MB third-party process behind the
same label as a 3.9 MB compositor, so it is now split into Desktop and Unmanaged;
and the version badge read a hardcoded "v1.0-DEPLOY" while the API served 2.0.0, so
it now reads the version from /openapi.json.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-05 23:48:37 -07:00
drjones
2948b0b440 Detect fan-mode drift; document health checks, VRAM accounting and tests
The reconciler only compared the power limit, so a profile whose fan setting never
applied stayed wrong indefinitely. The headless X server that owns the GPU can be
starting when this unit does; in-process retries cover a short delay, but if X
arrives later nothing noticed that the fan mode had never been set. profile_drift()
now compares fan mode too, gated on fan control having worked at least once so the
check does not fire forever on a machine without it.

README documents the self-check endpoint, the ollama/comfy/desktop/unmanaged
bucketing and why unreclaimable VRAM is reported separately, the SSE trimming, and
a tests section listing the measured constants the suite pins.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-05 19:17:16 -07:00
drjones
fb104ac9c0 Add a dependency self-check; cut the SSE payload by 74%
Self-check. Fan control failed for an entire session -- recoverably, and completely
invisibly. It appeared once, inside one field of one log line, and nothing ever
asked whether fan control worked. health.py now checks everything this service
depends on (NVML, passwordless sudo for nvidia-smi, fan control via the headless X
server, overclock drift, the telemetry store, residency measurement capability,
model directories, the ComfyUI websocket, and both upstream HTTP services) and
reports for each one what is broken, what that breaks, and how to fix it. Exposed at
GET /api/health, as an MCP tool, and as a dashboard panel that collapses to a badge
when healthy and expands to impact-and-fix when not. Current state: 9 ok, 1 degraded
(the known cachestat permission limit on Ollama's blobs).

A self-check that returns ok while a dependency is broken is worse than none, so the
tests drive each check to its failure state -- including the exact "Error resolving
target specification 'gpu:0'" string from the original incident -- and assert that a
check which raises surfaces as failed rather than taking down the endpoint.

SSE payload. The installed-model catalog was 10.6 KB of a 13.1 KB frame, 81% of the
stream, re-sent to every subscriber every second despite changing only when a model
is pulled or removed: 135 MB/hour across three tabs. It is now sent on a
subscriber's first frame and whenever the set changes; the client keeps the last
known list. Steady-state frames dropped from 14041 to 3664 bytes, a 74% reduction,
and /api/stats still returns the complete snapshot for API consumers.

Tests: 182 (was 169).

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-05 19:15:52 -07:00
drjones
25b601e24a Fix two more stale-cache bugs; account for VRAM this service cannot reclaim
The stale-readback bug fixed in f5917a0 was a class, not an instance. Two more:

- ram_optimizer's residency report cached for 15s and was never invalidated when
  anything warmed a file, so warming a model and then looking at residency showed
  the state from before the warm. warm_file_to_ram now invalidates it.
- _PID_KIND_CACHE was keyed on pid alone and never expired. Linux recycles PIDs, so
  a stale entry could attribute a new process's VRAM to Ollama or ComfyUI -- inside
  the very snapshot the yield barrier trusts to decide whether VRAM was released.
  Now keyed by (pid, process start time) and bounded.

Unmanaged VRAM. Investigating a persistence-mode warning turned up a third GPU
consumer this service does not model: stt_relay.py, holding 842 MB for nearly three
days. It was bucketed as "system" alongside gnome-shell's 3.9 MB. That conflation
matters, because ComfyUI's memory can be reclaimed and a third party's cannot, and
the reclaim path assumed ComfyUI was always to blame for missing headroom.

Processes are now bucketed ollama | comfy | desktop | unmanaged. The breakdown
reports desktop_gb and unmanaged_gb separately and names the unmanaged processes;
when a reclaim-and-retry still fails, the error identifies them rather than
implying ComfyUI was at fault; and the dashboard shows the unreclaimable total, so
headroom the arbitrator can never give back is visible rather than inferred.

Checked and deliberately not changed: persistence mode reads Disabled, but
nvidia-persistenced is active and two clients hold the GPU open continuously, so
the driver never unloads. The nvidia-smi warning is legacy noise here and is not a
source of the profile drift.

Tests: 169 (was 164). The new ones cover the bucketing, and one existing test used
Xorg as its "unknown process" fixture -- correct before a display server had its own
bucket, wrong after.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-05 19:11:40 -07:00
drjones
f5917a0464 Report the GPU state that is real, and reconcile when it drifts
The API claimed the card was at 320W while nvidia-smi reported 370W. Three
separate defects, all introduced by me in this branch.

The readback was stale. get_gpu_state()/get_fan_status() gained a 2s cache so the
dashboard's polling would stop forking sudo every few seconds, but apply_profile
read back through that cache and its invalidation ran afterwards. A profile that
had just moved the card 370W -> 320W therefore returned a payload whose detail
string said "set to 320.00 W from 370.00 W" next to a power_limit_w of 370.0.
Caches are now cleared before the readback, which is forced.

Fan control could fail for an entire session. On boot this unit can start before
the headless X server on :8 that owns the GPU accepts connections, and the fan
assignment fails with "Error resolving target specification 'gpu:0'". Nothing
retried and nothing surfaced it, so the fans were left unconfigured with the
failure visible only inside one log line. apply_fan_control now recognises that
specific error and retries up to 5 times.

Nothing verified the result. ACTIVE_PROFILE defaults to "balanced" at import,
which is indistinguishable from "balanced was successfully applied" -- so a failed
startup apply left the app confidently reporting a profile it had never put on the
hardware. apply_profile now returns a `verified` block comparing intent against
readback and logs a warning on mismatch; profile_drift() exposes the comparison
plus whether any profile has actually been applied since startup; and the 1Hz
sampler calls reconcile_profile() once a minute to re-apply on drift.

Verified by setting 370W externally behind the service's back: the drift was
reported immediately and corrected automatically 40s later.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-02 19:51:23 -07:00
drjones
868d82794d Add test suite (164 tests); reclaim VRAM from ComfyUI when an LLM will not fit
Tests. First automated coverage for the project: 164 tests, 2.7s, no GPU or network.
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. They deliberately
pin the empirically measured constants that would otherwise rot silently: the cold and
warm load figures behind the cache-hit thresholds, the warm_confident residency rule,
and the busy/stalled yield split. One test asserts RAM_HIT_GBPS stays at or below the
measured 2.63 GB/s warm load, so the old physically unreachable 5.0 GB/s bar cannot
come back.

Three bugs the suite surfaced, now fixed:
- autotune._subsample(values, 1) divided by zero; the early return only covered
  len(values) <= max_steps.
- telemetry_store.stop() flushed its local pending list but never drained the queue,
  silently losing rows submitted just before a shutdown -- exactly when the last
  events matter.
- ram_optimizer.page_residency's zero-byte short-circuit omitted keys every other
  return path provides, so a 0-byte file was planned for warming.

Reclaim. The README has claimed bidirectional arbitration from the start, but only one
direction was ever automatic. Establishing what actually happens took a controlled test
with the service stopped: with ComfyUI holding 6.83 GB, Ollama does not spill to the CPU
on this box -- it aborts with "cudaMalloc failed: out of memory", because n_gpu_layers is
pinned to 99 and it will not reduce the layer count. So both failure modes are handled:
_check_ollama_starved watches size_vram < size for the default configuration where Ollama
does spill, and switch_ollama_model catches the hard OOM, reclaims VRAM from an idle
ComfyUI and retries once. The request that returned HTTP 500 from Ollama directly now
succeeds through HyperSwap, loading at 3.85 GB/s after reclaiming 6.83 GB.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-01 13:40:30 -07:00
drjones
bacaf50713 Treat a mid-generation LLM as busy, not as a failed yield
The persisted counters showed 19 timeouts in 20 yields. All 19 were one model,
ornith-1.5:9b-cron, in two bursts at 06:25 and 06:33. Telemetry for that window
shows the GPU pinned at 96-97% with Ollama holding 14.92 GB throughout: the model
was mid-generation. Ollama will not unload a model that is inferencing, so every
request failed, and with a 1s trigger debounce against a 10s blocking wait the
arbitrator simply asked again, four times per burst, blocking the loop for 40s.

Ollama's behaviour is correct. Ours was wrong in three ways.

Busy is now a distinct outcome. _await_vram_release returns "released", "busy" or
"stuck": VRAM that has not moved while the GPU is pinned means a generation is in
flight, which is not a failure. With OLLAMA_NUM_PARALLEL=1 our keep_alive:0 request
queues behind the running one and applies the moment it finishes, so the correct
response is to stop waiting, not to retry. Only "stuck" -- VRAM held with an idle
GPU -- is a real fault.

The wait is short again (2s, from 10s) because blocking helps nobody: ComfyUI is
not gated on our return value, and every blocked second stalls the watchdog and
profile switching. Callers who genuinely want to wait out an inference can pass
wait_for_generation=true. The unload POST itself now gets a 120s client timeout,
since a 5s one could drop the connection before Ollama ever processed a request
queued behind a long generation, losing the unload entirely.

A busy model gets per-model backoff (5s, 15s, 30s, 60s) instead of being asked
again every second, and a detached watcher confirms and logs the release when the
generation ends, so the event log tells the whole story rather than stopping at
"deferred". Measured: a mid-generation yield now returns busy in 610ms instead of
blocking 10s, and the queued unload lands on its own 3s later when the generation
completes.

Counters are honest: yields (released), yield_deferred_busy, deferred_releases,
yield_stalled. The old yield_timeouts conflated a healthy cron job with a fault
and implied a 95% failure rate.

Also adds a VRAM Arbitration panel to the dashboard. The arbitrator is the core of
this application and its state was not displayed anywhere -- there was no way to
see whether handoffs were working, which is why this went unnoticed until the
persisted counters were read by hand.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
2026-09-01 13:32:17 -07:00
23 changed files with 3166 additions and 93 deletions

View File

@@ -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

View File

@@ -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
View 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,
}

View File

@@ -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.

View File

@@ -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
View File

@@ -0,0 +1,5 @@
[pytest]
testpaths = tests
addopts = -ra
filterwarnings =
ignore::DeprecationWarning

View File

@@ -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,

View File

@@ -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():

View File

@@ -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' ? '&#10003;' : c.status === 'degraded' ? '!' : '&#10007;';
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);
});

View File

@@ -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">

View File

@@ -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
View 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
View 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

View 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
View 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
View 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"])

View 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
View 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"] == []

View 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)

View 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
View 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)

View 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

View File

@@ -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()},
}