Compare commits
20 Commits
172d812820
...
main
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
48fbde040c | ||
|
|
0fbc3963b9 | ||
|
|
48bd0096ff | ||
|
|
ca97f18be6 | ||
|
|
c6455d7c6e | ||
|
|
2b00ab4e12 | ||
|
|
4a38cd68b3 | ||
|
|
1bcfbb2335 | ||
|
|
c3d9b36035 | ||
|
|
aeba1b47fd | ||
|
|
81e5d88426 | ||
|
|
aed1c360f0 | ||
|
|
043d61722b | ||
|
|
b53b026ced | ||
|
|
2948b0b440 | ||
|
|
fb104ac9c0 | ||
|
|
25b601e24a | ||
|
|
f5917a0464 | ||
|
|
868d82794d | ||
|
|
bacaf50713 |
201
README.md
201
README.md
@@ -20,6 +20,30 @@
|
||||
## 1. Feature Matrix
|
||||
|
||||
### ⚡ Bidirectional VRAM Hot-Swapping & Arbitration
|
||||
|
||||
* **A stale ComfyUI queue entry no longer disables arbitration.** ComfyUI can leave a
|
||||
dead job in `queue_running` indefinitely; one was found sitting there with the GPU idle
|
||||
and ComfyUI holding 0.56 GB. Trusting that flag made this service believe ComfyUI was
|
||||
permanently busy — so it evicted the LLM on every poll, never ran the idle purge, and
|
||||
never checked for CPU spill. Instrumenting the watchdog showed `busy=6, idle_check=0`.
|
||||
A running entry is now corroborated against ComfyUI's own VRAM (a real job loads
|
||||
gigabytes; a dead one holds only its CUDA context) before it is believed, and a stale
|
||||
entry is reported by `/api/health`. Utilisation is deliberately *not* the signal — it is
|
||||
shared with Ollama and any third-party process.
|
||||
* **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 +120,40 @@ 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.
|
||||
|
||||
### 🔧 Live Engine Configuration (`engines.py`)
|
||||
* `GET /api/engines` reports the **real, current** configuration of both engines and what
|
||||
each setting implies for arbitration — because the settings that dictate this service's
|
||||
behaviour live outside its own codebase.
|
||||
* `OLLAMA_NUM_PARALLEL=1` is why a `keep_alive: 0` unload queues behind a running
|
||||
generation and is reported as *deferred* rather than failed. `OLLAMA_MAX_LOADED_MODELS=1`
|
||||
is why every swap evicts the previous model. Working these out originally meant reading
|
||||
journald and the systemd unit by hand.
|
||||
* The dashboard's engine subtitles now come from this endpoint. They were previously
|
||||
hardcoded — and happened to be accurate, which is worse than being wrong, since they
|
||||
would have kept looking accurate after the configuration changed.
|
||||
|
||||
### 🩺 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 +162,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 +174,145 @@ only if the card actually needs it.
|
||||
|
||||
---
|
||||
|
||||
## 1c. Lining Work Up
|
||||
|
||||
Until now this service only *reacted*: it noticed an application had started and
|
||||
scrambled to free memory. Nothing could be queued. Each application has its own queue,
|
||||
but they cannot see each other, so work submitted to one has no way to wait for the other.
|
||||
|
||||
```bash
|
||||
curl -X POST localhost:9090/api/jobs -H 'Content-Type: application/json' -d '{
|
||||
"tenant": "ollama", "label": "nightly-summary",
|
||||
"payload": {"model": "qwen3.8fast:latest", "prompt": "..."}
|
||||
}'
|
||||
```
|
||||
|
||||
Jobs live in SQLite, so the queue is bounded by disk rather than memory and survives a
|
||||
restart. The scheduler takes the highest-priority pending job, arbitrates VRAM for it with
|
||||
the same `plan_release`, runs it, and moves on. One at a time by design — the GPU is the
|
||||
scarce resource this service exists to hand between applications, and overlapping jobs
|
||||
would just recreate the contention it resolves.
|
||||
|
||||
`GET /api/jobs` · `GET /api/jobs/{id}` · `DELETE /api/jobs/{id}` (pending only — running
|
||||
work is never killed) · `DELETE /api/jobs` to clear the queue. Agents get the same through
|
||||
MCP (`queue_job`, `get_job_queue`, `cancel_job`), and the dashboard's **Job Queue** panel
|
||||
shows what is running, what it released to get there, and why the scheduler is waiting if
|
||||
it is.
|
||||
|
||||
**A job that cannot run yet waits; a job that can never run fails with the reason.**
|
||||
Dispatching into insufficient VRAM does not fail gracefully — it kills `llama-server`
|
||||
with a CUDA OOM. The requirement is computed per job (an LLM job needs the size of *its*
|
||||
model, not a tenant-wide figure), and if the memory can never be assembled the job fails
|
||||
naming what stands in the way rather than blocking the queue forever.
|
||||
|
||||
## 1b. Any Application, Not Just These Two
|
||||
|
||||
The purpose is fast handoff of one GPU between applications. It grew up around the two on
|
||||
this box, and their names ended up compiled into process matching, VRAM attribution, busy
|
||||
detection and release calls alike — about 385 references. That made it a script for Ollama
|
||||
and ComfyUI rather than a GPU arbitrator.
|
||||
|
||||
A tenant is now **described as data** in `tenants.json`:
|
||||
|
||||
```json
|
||||
{
|
||||
"name": "trainer",
|
||||
"kind": "other",
|
||||
"priority": 80,
|
||||
"match": { "cmdline": ["train.py"] },
|
||||
"busy": { "type": "vram", "vram_busy_gb": 1.0 },
|
||||
"release": { "type": "http_post", "url": "http://localhost:9999/release" }
|
||||
}
|
||||
```
|
||||
|
||||
| Field | What it answers |
|
||||
| :--- | :--- |
|
||||
| `match` | Which GPU processes belong to this application (name, cmdline substring, or suffix — ComfyUI is a bare `python main.py`) |
|
||||
| `busy` | Whether it is *genuinely* working. `http_count` sums queue lists; `vram` needs no API at all. `vram_floor_gb` catches a queue that claims work while nothing is loaded |
|
||||
| `release` | How to ask for VRAM back — `http_post` with a body, `per_model` for Ollama's per-model unload, or `none` |
|
||||
| `priority` | Who is asked to yield **first** among idle tenants — it never protects idle memory, and never interrupts work |
|
||||
| `overclock_profile` | GPU profile applied while this tenant is the active workload |
|
||||
| `events` | Optional stream (e.g. a websocket) used purely as a wake-up, so reaction is sub-second rather than waiting for the next poll |
|
||||
|
||||
Two more fields drive the decision loop: **`needs_vram_gb`** (how much free memory the
|
||||
application needs before it can work) and **`idle_release_after_s`** (how long it may sit
|
||||
idle holding VRAM before being asked for it back — deliberately not immediate, so
|
||||
iterating on a ComfyUI workflow does not reload the checkpoint between every run).
|
||||
|
||||
**Priority orders, it does not veto.** An idle tenant is not using its VRAM, so
|
||||
outranking the demander is no reason to keep it; busy tenants are never interrupted
|
||||
whatever their rank. Getting this wrong broke both directions in turn — with the LLM
|
||||
ranked above diffusion, ComfyUI could never preempt Ollama (the service's central
|
||||
behaviour), and once the ranks were swapped, a starved Ollama could no longer reclaim
|
||||
from an idle ComfyUI. Diffusion now outranks the LLM, whose weights reload from page
|
||||
cache in seconds.
|
||||
|
||||
`plan_release()` then arbitrates generically: a busy tenant that cannot reach
|
||||
`needs_vram_gb` *even counting what it already holds* is starved, and the memory is taken
|
||||
from idle reclaimable tenants below it in priority, lowest first, stopping as soon as
|
||||
enough is freed. Tenants that cannot be released are named as blockers rather than
|
||||
ignored, so `possible: false` comes with the reason. The plan is returned before it is
|
||||
acted on, which makes the decision testable and loggable.
|
||||
|
||||
Ollama, ComfyUI and the desktop compositor ship as defaults, so behaviour is unchanged —
|
||||
but nothing in the arbitration logic knows their names, and three applications can
|
||||
contend for the card as easily as two. The dashboard's **GPU Tenants** panel lists all of
|
||||
them ordered by the priority arbitration actually considers, with the last decision and
|
||||
why it could or could not be satisfied. Endpoints are generic:
|
||||
`GET /api/tenants`, `GET /api/tenants/{name}`, `POST /api/tenants/{name}/release`.
|
||||
|
||||
A tenant with `"release": {"type": "none"}` is still worth declaring. The 842 MB speech
|
||||
relay on this box cannot be reclaimed, and naming it turns anonymous "unmanaged VRAM" into
|
||||
"held by stt-relay, which exposes no release API" — and a release request returns **409**
|
||||
explaining that, rather than silently doing nothing.
|
||||
|
||||
## 1a. Tests
|
||||
|
||||
```bash
|
||||
/home/drjones/comfy-mcp-venv/bin/python -m pytest tests/ -q # 271 passed in ~4.0s
|
||||
```
|
||||
|
||||
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 |
|
||||
|
||||
## 1b. End-to-End Verification
|
||||
|
||||
```bash
|
||||
python verify_arbitration.py # full cycle, a few minutes
|
||||
python verify_arbitration.py --quick # skip the diffusion stages
|
||||
```
|
||||
|
||||
The unit suite covers logic in isolation. This exercises the promise the service exists
|
||||
to make — an LLM and a diffusion pipeline sharing one 16 GB card — against real hardware,
|
||||
and reports what actually happened at each stage. It restores what it changes and refuses
|
||||
to start if ComfyUI is busy.
|
||||
|
||||
A representative run on this machine:
|
||||
|
||||
| Stage | Result |
|
||||
| :--- | :--- |
|
||||
| LLM load, classified by achieved bandwidth | 1.96 GB in 1327 ms → 1.47 GB/s → Partial Cache |
|
||||
| VRAM yield confirmed against NVML | released in 43 ms, 2.39 GB freed |
|
||||
| Diffusion, cold (includes checkpoint load) | 17863 ms → 1.12 it/s |
|
||||
| Diffusion, warm | 3645 ms → **5.49 it/s** |
|
||||
| ComfyUI retains its checkpoint | 7.03 GB held through the idle window |
|
||||
| VRAM attribution adds up | 15.58 GB attributed vs 15.80 GB NVML — Ollama 7.71 + ComfyUI 7.03 coexisting |
|
||||
| Reported GPU state matches hardware | profile asks 320 W, card reports 320 W |
|
||||
|
||||
The stages report warnings rather than passes when they did not actually prove anything —
|
||||
a reclaim that was never needed is not evidence that reclaiming works.
|
||||
|
||||
## 2. Architectural Overview
|
||||
|
||||
```mermaid
|
||||
@@ -207,6 +404,8 @@ 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/engines` | `GET` | Live Ollama and ComfyUI configuration, with what each setting implies for arbitration. |
|
||||
| `/api/health` | `GET` | Dependency self-check: NVML, sudo, fan control, drift, store, upstreams — each with impact and remediation. |
|
||||
|
||||
### Governor & Autotune Endpoints
|
||||
|
||||
|
||||
@@ -111,6 +111,10 @@ def _subsample(values: List[int], max_steps: int) -> List[int]:
|
||||
"""
|
||||
if len(values) <= max_steps:
|
||||
return values
|
||||
if max_steps < 2:
|
||||
# One step cannot span a range; take the top, which is what a caller asking for
|
||||
# a single sample almost certainly wants. Guards a ZeroDivisionError below.
|
||||
return values[-1:]
|
||||
idx = [round(i * (len(values) - 1) / (max_steps - 1)) for i in range(max_steps)]
|
||||
return sorted({values[i] for i in idx})
|
||||
|
||||
|
||||
117
engines.py
Normal file
117
engines.py
Normal file
@@ -0,0 +1,117 @@
|
||||
"""Live configuration of the two engines HyperSwap arbitrates between.
|
||||
|
||||
Arbitration behaviour is largely dictated by settings that live outside this codebase.
|
||||
Working out why a yield behaved the way it did meant reading journald and the ollama
|
||||
unit by hand: OLLAMA_NUM_PARALLEL decides whether an unload queues behind a running
|
||||
generation, OLLAMA_MAX_LOADED_MODELS decides whether more than one model can be
|
||||
resident, and a pinned n_gpu_layers decides whether a model that will not fit spills to
|
||||
the CPU or fails outright. Those are worth reading and explaining rather than hardcoding
|
||||
into a dashboard subtitle that silently goes stale.
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
import subprocess
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
import vram_arbitrator
|
||||
|
||||
logger = logging.getLogger("engines")
|
||||
|
||||
# Settings that change how the arbitrator must behave, with what they imply.
|
||||
OLLAMA_SETTING_NOTES = {
|
||||
"OLLAMA_NUM_PARALLEL": (
|
||||
"Requests per model. At 1, a keep_alive:0 unload queues behind any running "
|
||||
"generation and applies when it finishes — which is why a busy model is "
|
||||
"reported as deferred rather than failed."),
|
||||
"OLLAMA_MAX_LOADED_MODELS": (
|
||||
"How many models may be resident at once. At 1, Ollama evicts the previous "
|
||||
"model on every swap."),
|
||||
"OLLAMA_KEEP_ALIVE": (
|
||||
"Default residency after a request. Long values keep VRAM occupied and make "
|
||||
"ComfyUI wait for an explicit yield."),
|
||||
"OLLAMA_FLASH_ATTENTION": "FlashAttention kernels for attention.",
|
||||
"OLLAMA_KV_CACHE_TYPE": "KV cache quantisation; smaller types cut VRAM per context.",
|
||||
"OLLAMA_NUM_BATCH": "Prompt-evaluation batch size.",
|
||||
}
|
||||
|
||||
|
||||
def _ollama_unit_environment() -> Dict[str, str]:
|
||||
"""Read the ollama service's environment. Empty if it is not a systemd unit."""
|
||||
env: Dict[str, str] = {}
|
||||
try:
|
||||
proc = subprocess.run(["systemctl", "show", "ollama", "-p", "Environment",
|
||||
"--value"], capture_output=True, text=True, timeout=8)
|
||||
for token in proc.stdout.split():
|
||||
if "=" in token and token.startswith("OLLAMA"):
|
||||
k, _, v = token.partition("=")
|
||||
env[k] = v
|
||||
except Exception as e:
|
||||
logger.debug(f"could not read ollama unit environment: {e}")
|
||||
return env
|
||||
|
||||
|
||||
async def get_engine_config() -> Dict[str, Any]:
|
||||
"""Real, live configuration of both engines, with arbitration implications."""
|
||||
ollama_env = _ollama_unit_environment()
|
||||
ollama_settings = [
|
||||
{"key": k, "value": v, "means": OLLAMA_SETTING_NOTES.get(k, "")}
|
||||
for k, v in sorted(ollama_env.items())
|
||||
]
|
||||
|
||||
# A short, honest summary line to replace the dashboard's hardcoded subtitle.
|
||||
feature_bits: List[str] = []
|
||||
if ollama_env.get("OLLAMA_FLASH_ATTENTION") == "1":
|
||||
feature_bits.append("FlashAttention")
|
||||
kv = ollama_env.get("OLLAMA_KV_CACHE_TYPE")
|
||||
if kv:
|
||||
feature_bits.append(f"{kv} KV cache")
|
||||
host = ollama_env.get("OLLAMA_HOST", "")
|
||||
port = host.rsplit(":", 1)[-1] if ":" in host else "11434"
|
||||
|
||||
ollama = {
|
||||
"port": port,
|
||||
"settings": ollama_settings,
|
||||
"summary": " + ".join(feature_bits) if feature_bits else "default configuration",
|
||||
"max_loaded_models": ollama_env.get("OLLAMA_MAX_LOADED_MODELS"),
|
||||
"num_parallel": ollama_env.get("OLLAMA_NUM_PARALLEL"),
|
||||
"keep_alive": ollama_env.get("OLLAMA_KEEP_ALIVE"),
|
||||
"config_source": "systemd unit environment" if ollama_env else "unavailable",
|
||||
}
|
||||
|
||||
comfy: Dict[str, Any] = {"online": False}
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=4.0) as c:
|
||||
r = await c.get(f"{vram_arbitrator.COMFY_API_BASE}/system_stats")
|
||||
if r.status_code == 200:
|
||||
data = r.json()
|
||||
system = data.get("system", {})
|
||||
argv = system.get("argv") or []
|
||||
devices = data.get("devices") or []
|
||||
dev = devices[0] if devices else {}
|
||||
# The allocator is named in the device string; it is the closest thing
|
||||
# ComfyUI reports to the "async offload" the old subtitle asserted.
|
||||
dev_name = dev.get("name", "")
|
||||
allocator = ("cudaMallocAsync" if "cudaMallocAsync" in dev_name
|
||||
else "cudaMalloc" if "cudaMalloc" in dev_name else "unknown")
|
||||
vram_flags = [a for a in argv
|
||||
if a in ("--lowvram", "--novram", "--highvram", "--normalvram",
|
||||
"--gpu-only", "--cpu")]
|
||||
comfy = {
|
||||
"online": True,
|
||||
"version": system.get("comfyui_version"),
|
||||
"pytorch": system.get("pytorch_version"),
|
||||
# split()[0] on an absent version raises IndexError, which would have
|
||||
# been swallowed by the except below and reported ComfyUI as offline.
|
||||
"python": ((system.get("python_version") or "").split() or [None])[0],
|
||||
"argv": argv,
|
||||
"vram_mode": vram_flags[0] if vram_flags else "default (auto)",
|
||||
"allocator": allocator,
|
||||
"device": dev_name,
|
||||
"summary": f"{allocator}, {vram_flags[0] if vram_flags else 'auto VRAM'}",
|
||||
}
|
||||
except Exception as e:
|
||||
comfy = {"online": False, "error": str(e)[:120]}
|
||||
|
||||
return {"ollama": ollama, "comfyui": comfy}
|
||||
277
health.py
Normal file
277
health.py
Normal file
@@ -0,0 +1,277 @@
|
||||
"""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 time
|
||||
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_comfy_queue() -> Dict[str, Any]:
|
||||
"""A stale ComfyUI queue entry disables half of this service's logic."""
|
||||
arb = vram_arbitrator.arbitrator
|
||||
if arb.comfy_stale_job:
|
||||
return _check("comfyui queue", DEGRADED,
|
||||
f"prompt {arb.comfy_stale_job} claims to be running but the GPU is idle",
|
||||
"ComfyUI looks permanently busy, so the LLM is evicted repeatedly, "
|
||||
"the idle purge never runs and CPU-spill is never checked",
|
||||
"Clear it from the ComfyUI queue, or POST /queue with "
|
||||
"{\"clear\": true} to ComfyUI")
|
||||
branches = arb.watchdog_branches
|
||||
if branches.get("idle_check", 0) == 0 and branches.get("busy", 0) > 20:
|
||||
return _check("comfyui queue", DEGRADED,
|
||||
"the watchdog has only ever seen ComfyUI as busy",
|
||||
"The idle purge and starvation check are not running",
|
||||
"Check the ComfyUI queue for a stuck entry")
|
||||
return _check("comfyui queue", OK, "queue state corroborated against GPU activity")
|
||||
|
||||
|
||||
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 _comfy_vram_floor_gb(default: float = 0.0, days: float = 1.0) -> float:
|
||||
"""Lowest VRAM ComfyUI has been observed holding while alive.
|
||||
|
||||
A purge frees checkpoints but not the CUDA context, so ComfyUI keeps a few hundred
|
||||
MB for as long as the process runs. The minimum seen in recent telemetry is a better
|
||||
estimate of that floor than whatever it happens to hold right now, which could be a
|
||||
7 GB checkpoint mid-generation.
|
||||
"""
|
||||
try:
|
||||
rows = telemetry_store._rows(
|
||||
"SELECT MIN(comfy_bytes) AS floor FROM telemetry "
|
||||
"WHERE ts > ? AND comfy_bytes > 0",
|
||||
(time.time() - days * 86400,))
|
||||
if rows and rows[0].get("floor"):
|
||||
return round(rows[0]["floor"] / (1024 ** 3), 2)
|
||||
except Exception as e:
|
||||
logger.debug(f"comfy floor lookup failed: {e}")
|
||||
return default
|
||||
|
||||
|
||||
def _check_unmanaged_vram() -> Dict[str, Any]:
|
||||
"""Report unreclaimable VRAM in terms of what it actually costs.
|
||||
|
||||
"0.82 GB unmanaged" is a number. "0.82 GB unmanaged, which is why three of your
|
||||
models can no longer fit" is something you can act on.
|
||||
"""
|
||||
stats = vram_arbitrator.get_gpu_hardware_stats()
|
||||
if not stats.get("available"):
|
||||
return _check("unmanaged VRAM", DEGRADED, "GPU unavailable")
|
||||
bd = stats.get("breakdown", {})
|
||||
unmanaged_gb = bd.get("unmanaged_gb", 0.0)
|
||||
procs = bd.get("unmanaged", [])
|
||||
if not procs:
|
||||
return _check("unmanaged VRAM", OK, "no third-party GPU processes")
|
||||
|
||||
total_gb = stats.get("vram_total_gb", 0)
|
||||
# What HyperSwap could offer at best. Three things are never available: the desktop,
|
||||
# processes it cannot touch, and ComfyUI's own CUDA context, which survives a purge.
|
||||
# Omitting that last one made this check claim a 14.93 GB model would fit against a
|
||||
# real ceiling of 14.60 GB -- the model that had just returned 507.
|
||||
comfy_floor_gb = _comfy_vram_floor_gb(default=bd.get("comfyui_gb", 0.0))
|
||||
ceiling_gb = total_gb - unmanaged_gb - bd.get("desktop_gb", 0.0) - comfy_floor_gb
|
||||
try:
|
||||
blobs = ram_optimizer.find_ollama_model_files()
|
||||
except Exception:
|
||||
blobs = []
|
||||
# Measured on this box: a 12.87 GB blob occupies 14.9 GB once context and KV cache
|
||||
# are allocated.
|
||||
VRAM_OVERHEAD = 1.16
|
||||
blocked = sorted(
|
||||
{b["model"]: b for b in blobs
|
||||
if b["size_gb"] * VRAM_OVERHEAD > ceiling_gb
|
||||
and b["size_gb"] * VRAM_OVERHEAD <= ceiling_gb + unmanaged_gb}.values(),
|
||||
key=lambda b: -b["size_gb"])
|
||||
|
||||
names = ", ".join(b["model"] for b in blocked[:3])
|
||||
who = ", ".join(f"{p['name']} ({p['vram_mb']} MB)" for p in procs[:2])
|
||||
if blocked:
|
||||
return _check("unmanaged VRAM", DEGRADED,
|
||||
f"{unmanaged_gb} GB held by {who}",
|
||||
f"{len(blocked)} model(s) fit within {ceiling_gb + unmanaged_gb:.2f} GB "
|
||||
f"but not the {ceiling_gb:.2f} GB actually available: {names}",
|
||||
"Stop that process to reclaim the difference, or accept that "
|
||||
"these models cannot load")
|
||||
return _check("unmanaged VRAM", OK,
|
||||
f"{unmanaged_gb} GB held by {who}; no model is blocked by it",
|
||||
"", "")
|
||||
|
||||
|
||||
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, _check_unmanaged_vram]
|
||||
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,
|
||||
}
|
||||
445
jobs.py
Normal file
445
jobs.py
Normal file
@@ -0,0 +1,445 @@
|
||||
"""A durable, unbounded job queue across every GPU tenant.
|
||||
|
||||
Until now this service only reacted: it noticed an application had started working and
|
||||
scrambled to free memory. Nothing could be *lined up*. Each application has its own queue
|
||||
(ComfyUI's prompt queue, Ollama's serialised requests), but they cannot see each other, so
|
||||
work submitted to one has no way to wait politely for the other.
|
||||
|
||||
Jobs submitted here are stored in SQLite, so the queue is limited by disk rather than
|
||||
memory and survives a restart. The scheduler takes the highest-priority pending job,
|
||||
makes sure its tenant actually has the VRAM to run it -- reusing the same plan_release
|
||||
arbitration -- dispatches it, and moves on.
|
||||
"""
|
||||
import asyncio
|
||||
import json
|
||||
import logging
|
||||
import sqlite3
|
||||
import time
|
||||
import uuid
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
import telemetry_store
|
||||
import tenants as tenants_mod
|
||||
|
||||
logger = logging.getLogger("jobs")
|
||||
|
||||
PENDING, RUNNING, DONE, FAILED, CANCELLED = (
|
||||
"pending", "running", "done", "failed", "cancelled")
|
||||
|
||||
SCHEMA = """
|
||||
CREATE TABLE IF NOT EXISTS jobs (
|
||||
id TEXT PRIMARY KEY,
|
||||
tenant TEXT NOT NULL,
|
||||
priority INTEGER NOT NULL DEFAULT 50,
|
||||
payload TEXT NOT NULL,
|
||||
state TEXT NOT NULL,
|
||||
submitted_at REAL NOT NULL,
|
||||
started_at REAL,
|
||||
finished_at REAL,
|
||||
error TEXT,
|
||||
result TEXT,
|
||||
label TEXT
|
||||
);
|
||||
CREATE INDEX IF NOT EXISTS idx_jobs_state ON jobs(state, priority DESC, submitted_at);
|
||||
"""
|
||||
|
||||
|
||||
def _conn() -> sqlite3.Connection:
|
||||
c = sqlite3.connect(telemetry_store.DB_PATH, timeout=10.0)
|
||||
c.row_factory = sqlite3.Row
|
||||
return c
|
||||
|
||||
|
||||
def init() -> None:
|
||||
with _conn() as c:
|
||||
c.executescript(SCHEMA)
|
||||
|
||||
|
||||
def submit(tenant: str, payload: Dict[str, Any], priority: Optional[int] = None,
|
||||
label: Optional[str] = None) -> Dict[str, Any]:
|
||||
"""Queue a job. There is no depth limit: the queue lives on disk."""
|
||||
t = tenants_mod.get_tenant(tenant)
|
||||
if not t:
|
||||
return {"success": False, "error": f"no tenant named '{tenant}'"}
|
||||
job_id = uuid.uuid4().hex[:12]
|
||||
row = {
|
||||
"id": job_id, "tenant": tenant,
|
||||
"priority": t.priority if priority is None else int(priority),
|
||||
"payload": json.dumps(payload), "state": PENDING,
|
||||
"submitted_at": time.time(), "label": label,
|
||||
}
|
||||
with _conn() as c:
|
||||
c.execute("INSERT INTO jobs (id, tenant, priority, payload, state, submitted_at,"
|
||||
" label) VALUES (:id,:tenant,:priority,:payload,:state,:submitted_at,"
|
||||
":label)", row)
|
||||
logger.info(f"queued job {job_id} for '{tenant}' at priority {row['priority']}")
|
||||
return {"success": True, "id": job_id, "tenant": tenant,
|
||||
"priority": row["priority"], "state": PENDING}
|
||||
|
||||
|
||||
def cancel(job_id: str) -> Dict[str, Any]:
|
||||
with _conn() as c:
|
||||
cur = c.execute("UPDATE jobs SET state=?, finished_at=? WHERE id=? AND state=?",
|
||||
(CANCELLED, time.time(), job_id, PENDING))
|
||||
if cur.rowcount:
|
||||
return {"success": True, "id": job_id, "state": CANCELLED}
|
||||
return {"success": False, "error": "job is not pending (already running or finished)"}
|
||||
|
||||
|
||||
def clear_pending() -> Dict[str, Any]:
|
||||
with _conn() as c:
|
||||
cur = c.execute("UPDATE jobs SET state=?, finished_at=? WHERE state=?",
|
||||
(CANCELLED, time.time(), PENDING))
|
||||
return {"success": True, "cancelled": cur.rowcount}
|
||||
|
||||
|
||||
def get(job_id: str) -> Optional[Dict[str, Any]]:
|
||||
with _conn() as c:
|
||||
r = c.execute("SELECT * FROM jobs WHERE id=?", (job_id,)).fetchone()
|
||||
return _row(r) if r else None
|
||||
|
||||
|
||||
def _row(r: sqlite3.Row) -> Dict[str, Any]:
|
||||
d = dict(r)
|
||||
for key in ("payload", "result"):
|
||||
if d.get(key):
|
||||
try:
|
||||
d[key] = json.loads(d[key])
|
||||
except Exception:
|
||||
pass
|
||||
if d.get("started_at") and d.get("finished_at"):
|
||||
d["duration_s"] = round(d["finished_at"] - d["started_at"], 2)
|
||||
if d.get("state") == PENDING:
|
||||
d["waiting_s"] = round(time.time() - d["submitted_at"], 1)
|
||||
return d
|
||||
|
||||
|
||||
def listing(state: Optional[str] = None, limit: int = 100) -> List[Dict[str, Any]]:
|
||||
q = "SELECT * FROM jobs"
|
||||
args: List[Any] = []
|
||||
if state:
|
||||
q += " WHERE state=?"
|
||||
args.append(state)
|
||||
# Pending jobs in the order the scheduler will take them; everything else newest first.
|
||||
q += (" ORDER BY priority DESC, submitted_at ASC" if state == PENDING
|
||||
else " ORDER BY submitted_at DESC")
|
||||
q += " LIMIT ?"
|
||||
args.append(limit)
|
||||
with _conn() as c:
|
||||
return [_row(r) for r in c.execute(q, args).fetchall()]
|
||||
|
||||
|
||||
def stats() -> Dict[str, Any]:
|
||||
with _conn() as c:
|
||||
rows = c.execute("SELECT state, COUNT(*) n FROM jobs GROUP BY state").fetchall()
|
||||
by_state = {r["state"]: r["n"] for r in rows}
|
||||
pend = c.execute(
|
||||
"SELECT tenant, COUNT(*) n FROM jobs WHERE state=? GROUP BY tenant",
|
||||
(PENDING,)).fetchall()
|
||||
oldest = c.execute(
|
||||
"SELECT MIN(submitted_at) t FROM jobs WHERE state=?", (PENDING,)).fetchone()
|
||||
return {
|
||||
"by_state": by_state,
|
||||
"pending_by_tenant": {r["tenant"]: r["n"] for r in pend},
|
||||
"queue_depth": by_state.get(PENDING, 0),
|
||||
"oldest_pending_s": (round(time.time() - oldest["t"], 1)
|
||||
if oldest and oldest["t"] else None),
|
||||
}
|
||||
|
||||
|
||||
def requeue_orphans() -> int:
|
||||
"""Return jobs abandoned mid-run to the queue.
|
||||
|
||||
RUNNING means "this process is working on it". If no process is, that is untrue, and
|
||||
the job would otherwise never finish and never retry.
|
||||
"""
|
||||
with _conn() as c:
|
||||
cur = c.execute("UPDATE jobs SET state=?, started_at=NULL WHERE state=?",
|
||||
(PENDING, RUNNING))
|
||||
return cur.rowcount
|
||||
|
||||
|
||||
def _next_job() -> Optional[Dict[str, Any]]:
|
||||
with _conn() as c:
|
||||
r = c.execute(
|
||||
"SELECT * FROM jobs WHERE state=? ORDER BY priority DESC, submitted_at ASC"
|
||||
" LIMIT 1", (PENDING,)).fetchone()
|
||||
return _row(r) if r else None
|
||||
|
||||
|
||||
def _mark(job_id: str, state: str, **fields) -> None:
|
||||
sets = ", ".join(f"{k}=?" for k in fields)
|
||||
args = list(fields.values()) + [state, job_id]
|
||||
with _conn() as c:
|
||||
c.execute(f"UPDATE jobs SET {sets + ', ' if sets else ''}state=? WHERE id=?", args)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------- dispatch
|
||||
|
||||
async def _dispatch_comfy(payload: Dict[str, Any]) -> Dict[str, Any]:
|
||||
"""Hand a workflow to ComfyUI and wait for it to finish."""
|
||||
base = tenants_mod.get_tenant("comfyui").busy.url.rsplit("/", 1)[0]
|
||||
async with httpx.AsyncClient(timeout=30.0) as c:
|
||||
r = await c.post(f"{base}/prompt", json={"prompt": payload.get("prompt", payload),
|
||||
"client_id": "hyperswap-jobs"})
|
||||
if r.status_code != 200:
|
||||
return {"ok": False, "error": f"HTTP {r.status_code}: {r.text[:200]}"}
|
||||
prompt_id = r.json().get("prompt_id")
|
||||
deadline = time.time() + payload.get("timeout_s", 1800)
|
||||
while time.time() < deadline:
|
||||
await asyncio.sleep(0.5)
|
||||
h = await c.get(f"{base}/history/{prompt_id}")
|
||||
entry = (h.json() or {}).get(prompt_id) if h.status_code == 200 else None
|
||||
if not entry:
|
||||
continue
|
||||
status = entry.get("status", {})
|
||||
if status.get("status_str") == "error":
|
||||
return {"ok": False, "error": "ComfyUI reported an execution error"}
|
||||
if status.get("completed"):
|
||||
return {"ok": True, "prompt_id": prompt_id}
|
||||
return {"ok": False, "error": "timed out waiting for ComfyUI"}
|
||||
|
||||
|
||||
async def _dispatch_ollama(payload: Dict[str, Any]) -> Dict[str, Any]:
|
||||
async with httpx.AsyncClient(timeout=payload.get("timeout_s", 1800)) as c:
|
||||
body = {"stream": False, "keep_alive": payload.get("keep_alive", "5m"), **payload}
|
||||
body.pop("timeout_s", None)
|
||||
r = await c.post("http://localhost:11434/api/generate", json=body)
|
||||
if r.status_code != 200:
|
||||
return {"ok": False, "error": f"HTTP {r.status_code}: {r.text[:200]}"}
|
||||
d = r.json()
|
||||
return {"ok": True, "response": (d.get("response") or "")[:2000],
|
||||
"eval_count": d.get("eval_count"),
|
||||
"tokens_per_sec": (round(d.get("eval_count", 0)
|
||||
/ (d.get("eval_duration", 1) / 1e9), 2)
|
||||
if d.get("eval_duration") else None)}
|
||||
|
||||
|
||||
DISPATCHERS = {
|
||||
tenants_mod.KIND_DIFFUSION: _dispatch_comfy,
|
||||
tenants_mod.KIND_LLM: _dispatch_ollama,
|
||||
}
|
||||
|
||||
|
||||
class Scheduler:
|
||||
"""Drains the queue, making room for each job before it runs.
|
||||
|
||||
One job at a time by design. The GPU is the scarce resource this whole service
|
||||
exists to hand between applications; running two jobs concurrently would just
|
||||
recreate the contention it is meant to resolve. Throughput comes from swapping
|
||||
quickly, not from overlapping.
|
||||
"""
|
||||
|
||||
def __init__(self) -> None:
|
||||
self.running = False
|
||||
self.task: Optional[asyncio.Task] = None
|
||||
self.current: Optional[Dict[str, Any]] = None
|
||||
self.last_finished: Optional[Dict[str, Any]] = None
|
||||
self.completed = 0
|
||||
self.failed = 0
|
||||
self.waits = 0
|
||||
self.blocked: Optional[Dict[str, Any]] = None
|
||||
self.idle_poll_s = 1.0
|
||||
self.blocked_poll_s = 2.0
|
||||
# How long a job may wait for VRAM before it is declared impossible. Long enough
|
||||
# to outlast a normal diffusion run, short enough not to wedge the queue.
|
||||
self.max_block_s = 120.0
|
||||
|
||||
async def start(self) -> None:
|
||||
if self.running:
|
||||
return
|
||||
init()
|
||||
# A job left RUNNING by a crash or a hard restart would sit there forever.
|
||||
requeued = requeue_orphans()
|
||||
if requeued:
|
||||
logger.warning(f"requeued {requeued} job(s) left running by a previous process")
|
||||
self.running = True
|
||||
self.task = asyncio.create_task(self._loop())
|
||||
logger.info("job scheduler started")
|
||||
|
||||
async def stop(self) -> None:
|
||||
self.running = False
|
||||
if self.task:
|
||||
self.task.cancel()
|
||||
|
||||
def _job_vram_requirement(self, tenant_name: str, payload: Dict[str, Any]) -> float:
|
||||
"""How much VRAM *this* job needs, not the tenant's generic figure.
|
||||
|
||||
A tenant-wide needs_vram_gb cannot be right for an LLM: the requirement is a
|
||||
property of the model being loaded. Ollama's generic 4 GB passed the room check
|
||||
with 8 GB free, and then a 14.9 GB model was dispatched into it and killed
|
||||
llama-server with a CUDA OOM -- three queued jobs destroyed in a row.
|
||||
"""
|
||||
import vram_arbitrator
|
||||
|
||||
t = tenants_mod.get_tenant(tenant_name)
|
||||
default = t.needs_vram_gb if t else 0.0
|
||||
model = payload.get("model")
|
||||
if t and t.kind == tenants_mod.KIND_LLM and model:
|
||||
size = vram_arbitrator._model_size_bytes(model)
|
||||
if size:
|
||||
# Measured on this box: a 12.87 GB blob occupies 14.9 GB once context
|
||||
# and KV cache are allocated.
|
||||
return round((size / (1024 ** 3)) * 1.16, 2)
|
||||
return default
|
||||
|
||||
async def _make_room(self, tenant_name: str,
|
||||
payload: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||||
"""Ensure the job's tenant has the VRAM it needs, using the normal arbitration."""
|
||||
import vram_arbitrator # imported late: it imports this module's siblings
|
||||
|
||||
t = tenants_mod.get_tenant(tenant_name)
|
||||
needed = self._job_vram_requirement(tenant_name, payload or {})
|
||||
if not t or not needed:
|
||||
return {"ready": True, "reason": "no VRAM requirement declared"}
|
||||
|
||||
state = await vram_arbitrator.arbitrator._tenant_state()
|
||||
free_gb = vram_arbitrator.arbitrator._last_tenant_state["free_gb"]
|
||||
held = next((s["vram_gb"] for s in state if s["name"] == tenant_name), 0.0)
|
||||
if held + free_gb >= needed:
|
||||
return {"ready": True, "needed_gb": needed,
|
||||
"reason": f"{free_gb:.2f} GB free, job needs {needed:.2f} GB"}
|
||||
|
||||
plan = tenants_mod.plan_release(tenant_name, state, free_gb, needed)
|
||||
for victim in plan["release"]:
|
||||
await vram_arbitrator.arbitrator._release_tenant(
|
||||
victim, f"queued job for '{tenant_name}'")
|
||||
if plan["release"]:
|
||||
# Give the driver a moment to actually hand the memory back.
|
||||
deadline = time.perf_counter() + 30
|
||||
target = int(needed * (1024 ** 3))
|
||||
while time.perf_counter() < deadline:
|
||||
if vram_arbitrator.get_process_vram_bytes()["free_bytes"] >= target:
|
||||
break
|
||||
await asyncio.sleep(0.05)
|
||||
# Only ready once the memory is genuinely there. A plan that *could* work is not
|
||||
# the same as VRAM that *is* free, and dispatching on the former is what OOMs.
|
||||
free_now = (vram_arbitrator.get_process_vram_bytes()["free_bytes"] / (1024 ** 3))
|
||||
# The best this GPU could ever offer this tenant: everything currently free, plus
|
||||
# what it already holds, plus everything that is reclaimable at all.
|
||||
reclaimable_gb = sum(s["vram_gb"] for s in state
|
||||
if s["name"] != tenant_name and s.get("reclaimable"))
|
||||
max_possible = round(held + free_now + reclaimable_gb, 2)
|
||||
return {"ready": plan["possible"] and (held + free_now) >= needed,
|
||||
"max_possible_gb": max_possible,
|
||||
"released": plan["release"], "needed_gb": needed,
|
||||
"free_gb": round(free_now, 2),
|
||||
"reason": f"job needs {needed:.2f} GB; {plan['reason']}",
|
||||
"blockers": plan.get("blockers")}
|
||||
|
||||
async def _loop(self) -> None:
|
||||
while self.running:
|
||||
try:
|
||||
job = _next_job()
|
||||
if not job:
|
||||
await asyncio.sleep(self.idle_poll_s)
|
||||
continue
|
||||
|
||||
tenant = tenants_mod.get_tenant(job["tenant"])
|
||||
dispatcher = DISPATCHERS.get(tenant.kind) if tenant else None
|
||||
if not dispatcher:
|
||||
_mark(job["id"], FAILED, finished_at=time.time(),
|
||||
error=f"no dispatcher for tenant kind "
|
||||
f"'{tenant.kind if tenant else '?'}'")
|
||||
self.failed += 1
|
||||
continue
|
||||
|
||||
# Check for room *before* claiming the job. Dispatching into
|
||||
# insufficient VRAM does not fail gracefully -- it kills llama-server
|
||||
# with a CUDA OOM, which is how three queued LLM jobs were destroyed
|
||||
# while ComfyUI legitimately held the card. A job that cannot run yet
|
||||
# waits; it does not fail.
|
||||
room = await self._make_room(job["tenant"], job.get("payload") or {})
|
||||
if not room.get("ready"):
|
||||
since = (self.blocked.get("since", time.time())
|
||||
if self.blocked and self.blocked.get("id") == job["id"]
|
||||
else time.time())
|
||||
waited = time.time() - since
|
||||
self.blocked = {"id": job["id"], "tenant": job["tenant"],
|
||||
"reason": room.get("reason"),
|
||||
"blockers": room.get("blockers"),
|
||||
"needed_gb": room.get("needed_gb"),
|
||||
"since": since, "waited_s": round(waited, 1)}
|
||||
|
||||
# Waiting is right while the memory might still arrive. It is wrong
|
||||
# when the job can never fit -- three LLM jobs sat pending forever
|
||||
# needing 14.93 GB on a card where only ~14.8 GB can ever be free,
|
||||
# because an unreclaimable process holds 0.82 GB. Say so and move on
|
||||
# rather than blocking the queue behind an impossibility.
|
||||
if waited > self.max_block_s:
|
||||
ceiling = room.get("max_possible_gb")
|
||||
detail = (f"needs {room.get('needed_gb')} GB but at most "
|
||||
f"{ceiling} GB can ever be free on this GPU"
|
||||
if ceiling is not None and room.get("needed_gb", 0) > ceiling
|
||||
else f"waited {int(waited)}s for VRAM: {room.get('reason')}")
|
||||
blockers = ", ".join(
|
||||
f"{b['name']} ({b['vram_gb']} GB, {b['why']})"
|
||||
for b in (room.get("blockers") or []))
|
||||
_mark(job["id"], FAILED, finished_at=time.time(),
|
||||
error=f"{detail}{'; blocked by ' + blockers if blockers else ''}")
|
||||
self.failed += 1
|
||||
self.blocked = None
|
||||
logger.warning(f"job {job['id']} cannot run: {detail}")
|
||||
continue
|
||||
|
||||
self.waits += 1
|
||||
await asyncio.sleep(self.blocked_poll_s)
|
||||
continue
|
||||
self.blocked = None
|
||||
|
||||
t0 = time.time()
|
||||
_mark(job["id"], RUNNING, started_at=t0)
|
||||
self.current = {**job, "state": RUNNING, "room": room,
|
||||
"started_at": t0}
|
||||
logger.info(f"running job {job['id']} for '{job['tenant']}' "
|
||||
f"({room.get('reason')})")
|
||||
|
||||
# Any failure here must land on the job. An exception used to escape to
|
||||
# the loop's handler, leaving the row RUNNING forever while the scheduler
|
||||
# moved on -- an orphan that never completed and never freed its slot.
|
||||
try:
|
||||
res = await dispatcher(job["payload"])
|
||||
except asyncio.CancelledError:
|
||||
_mark(job["id"], PENDING, started_at=None)
|
||||
self.current = None
|
||||
raise
|
||||
except Exception as e:
|
||||
res = {"ok": False, "error": f"dispatch raised: {e}"}
|
||||
|
||||
finished = time.time()
|
||||
if res.get("ok"):
|
||||
_mark(job["id"], DONE, finished_at=finished,
|
||||
result=json.dumps(res))
|
||||
self.completed += 1
|
||||
else:
|
||||
_mark(job["id"], FAILED, finished_at=finished,
|
||||
error=str(res.get("error"))[:500])
|
||||
self.failed += 1
|
||||
self.last_finished = {"id": job["id"], "tenant": job["tenant"],
|
||||
"ok": bool(res.get("ok")),
|
||||
"duration_s": round(finished - t0, 2),
|
||||
"made_room": room.get("released") or []}
|
||||
self.current = None
|
||||
except asyncio.CancelledError:
|
||||
raise
|
||||
except Exception as e:
|
||||
logger.error(f"scheduler error: {e}")
|
||||
self.current = None
|
||||
await asyncio.sleep(1.0)
|
||||
|
||||
def get_status(self) -> Dict[str, Any]:
|
||||
return {
|
||||
"running": self.running,
|
||||
"current": self.current,
|
||||
"last_finished": self.last_finished,
|
||||
"completed": self.completed,
|
||||
"failed": self.failed,
|
||||
"waits": self.waits,
|
||||
"blocked": self.blocked,
|
||||
**stats(),
|
||||
}
|
||||
|
||||
|
||||
scheduler = Scheduler()
|
||||
@@ -8,6 +8,9 @@ from typing import Dict, List, Any, Optional
|
||||
|
||||
from mcp.server import MCPServer
|
||||
import autotune
|
||||
import engines
|
||||
import health
|
||||
import jobs as jobs_mod
|
||||
import overclock_manager
|
||||
import ram_optimizer
|
||||
import telemetry_store
|
||||
@@ -136,6 +139,49 @@ 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 get_engine_config() -> str:
|
||||
"""Live configuration of Ollama and ComfyUI (parallelism, max loaded models,
|
||||
keep-alive, KV cache type, ComfyUI VRAM mode and allocator), with what each setting
|
||||
implies for VRAM arbitration."""
|
||||
return json.dumps(await engines.get_engine_config(), indent=2, default=str)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def queue_job(tenant: str, payload: Dict[str, Any], priority: Optional[int] = None,
|
||||
label: Optional[str] = None) -> str:
|
||||
"""Queue work for a GPU application without waiting for it.
|
||||
|
||||
tenant: 'ollama' (payload: model, prompt, options) or 'comfyui' (payload: {"prompt":
|
||||
<workflow>}). The queue is on disk, so there is no depth limit; jobs run one at a
|
||||
time, highest priority first, with VRAM arbitrated before each starts."""
|
||||
return json.dumps(jobs_mod.submit(tenant, payload, priority, label), indent=2,
|
||||
default=str)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def get_job_queue(state: Optional[str] = None, limit: int = 50) -> str:
|
||||
"""Queued and recent jobs, plus what the scheduler is doing and why it may be
|
||||
waiting. Pending jobs are listed in the order they will run."""
|
||||
return json.dumps({"jobs": jobs_mod.listing(state, limit),
|
||||
"scheduler": jobs_mod.scheduler.get_status()},
|
||||
indent=2, default=str)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def cancel_job(job_id: str) -> str:
|
||||
"""Cancel a job that has not started. Running work is never killed."""
|
||||
return json.dumps(jobs_mod.cancel(job_id), indent=2, default=str)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
async def check_system_health() -> str:
|
||||
"""Check every dependency HyperSwap needs (NVML, sudo nvidia-smi, fan control via the
|
||||
headless X server, Ollama, ComfyUI, the telemetry store, model directories) and report
|
||||
what is broken, what it breaks, and how to fix it."""
|
||||
return json.dumps(await health.run_health_checks(), indent=2, default=str)
|
||||
|
||||
|
||||
@mcp.tool()
|
||||
def get_page_cache_residency(include_files: bool = True) -> str:
|
||||
"""Measure how much of each model on disk is genuinely resident in the Linux page cache.
|
||||
|
||||
@@ -18,6 +18,7 @@ import logging
|
||||
import os
|
||||
import shutil
|
||||
import subprocess
|
||||
import time
|
||||
from typing import Dict, Any, Optional, List
|
||||
|
||||
logger = logging.getLogger("overclock_manager")
|
||||
@@ -278,12 +279,40 @@ def apply_profile(name: str, overrides: Optional[Dict[str, Any]] = None) -> Dict
|
||||
"detail": "skipped: this driver accepts clock offsets and ignores them"}),
|
||||
"fan": apply_fan_control(fan_mode, fan_speed),
|
||||
}
|
||||
result["gpu"] = get_gpu_state()
|
||||
result["fan_status"] = get_fan_status()
|
||||
result["overrides"] = overrides or {}
|
||||
|
||||
# Invalidate BEFORE reading back. These caches exist so the dashboard's polling does
|
||||
# not fork sudo every few seconds, but reading through them here reported the
|
||||
# pre-change value: a profile that had just moved the card 370W -> 320W came back
|
||||
# claiming 370W, so the API contradicted nvidia-smi.
|
||||
_STATE_CACHE["value"] = None
|
||||
_FAN_CACHE["value"] = None
|
||||
result["gpu"] = get_gpu_state(force=True)
|
||||
result["fan_status"] = get_fan_status(force=True)
|
||||
result["overrides"] = overrides or {}
|
||||
|
||||
# Say plainly whether the card ended up where the profile asked it to.
|
||||
intended_w = int(cfg.get("power_limit_w", 0))
|
||||
actual_w = result["gpu"].get("power_limit_w")
|
||||
result["verified"] = {
|
||||
"power_limit_intended_w": intended_w,
|
||||
"power_limit_actual_w": actual_w,
|
||||
"power_limit_ok": (actual_w is None or intended_w == 0
|
||||
or abs(float(actual_w) - intended_w) < 1.0),
|
||||
"fan_mode_intended": fan_mode,
|
||||
"fan_mode_actual": result["fan_status"].get("mode"),
|
||||
"fan_ok": result["fan"].get("applied", False),
|
||||
}
|
||||
if not result["verified"]["power_limit_ok"]:
|
||||
logger.warning(f"Profile '{name}' asked for {intended_w}W but the card reports "
|
||||
f"{actual_w}W")
|
||||
if not result["verified"]["fan_ok"]:
|
||||
logger.warning(f"Profile '{name}' could not set fans: "
|
||||
f"{result['fan'].get('detail')}")
|
||||
|
||||
global _APPLIED_ONCE, _FAN_AVAILABLE
|
||||
if result["verified"]["power_limit_ok"]:
|
||||
_APPLIED_ONCE = True
|
||||
if result["verified"]["fan_ok"]:
|
||||
_FAN_AVAILABLE = True
|
||||
ACTIVE_PROFILE = name
|
||||
_LAST_RESULT = result
|
||||
logger.info(f"Overclock profile applied: {name} -> {json.dumps(result, default=str)}")
|
||||
@@ -332,14 +361,39 @@ def is_headless_x_running() -> bool:
|
||||
return r["rc"] == 0
|
||||
|
||||
|
||||
def apply_fan_control(mode: str, speed_pct: int) -> Dict[str, Any]:
|
||||
FAN_RETRY_ATTEMPTS = 5
|
||||
FAN_RETRY_DELAY_S = 2.0
|
||||
|
||||
|
||||
def _fan_target_missing(result: Dict[str, Any]) -> bool:
|
||||
"""True when nvidia-settings could not see the GPU at all.
|
||||
|
||||
On boot this service can start before the headless X server on :8 that owns the GPU
|
||||
is accepting connections, and the fan assignment fails with 'Error resolving target
|
||||
specification'. Nothing retried, so the fans were simply never configured for the
|
||||
whole session and the failure was only visible deep in a log line.
|
||||
"""
|
||||
text = ((result.get("err") or "") + (result.get("out") or "")).lower()
|
||||
return ("error resolving target" in text or "no targets match" in text
|
||||
or "cannot open display" in text)
|
||||
|
||||
|
||||
def apply_fan_control(mode: str, speed_pct: int, _attempt: int = 0) -> Dict[str, Any]:
|
||||
global FAN_MANUAL
|
||||
if mode == "auto":
|
||||
r = _nvidia_settings("-a", "[gpu:0]/GPUFanControlState=0")
|
||||
ok = r["rc"] == 0
|
||||
if not ok and _fan_target_missing(r) and _attempt < FAN_RETRY_ATTEMPTS:
|
||||
logger.info(f"Fan control target not ready (attempt {_attempt + 1}/"
|
||||
f"{FAN_RETRY_ATTEMPTS}); X on {HEADLESS_DISPLAY} may still be "
|
||||
f"starting — retrying in {FAN_RETRY_DELAY_S}s")
|
||||
time.sleep(FAN_RETRY_DELAY_S)
|
||||
return apply_fan_control(mode, speed_pct, _attempt + 1)
|
||||
if ok:
|
||||
FAN_MANUAL = False
|
||||
return {"applied": ok, "mode": "auto", "detail": r.get("out") or r.get("err")}
|
||||
_FAN_CACHE["value"] = None
|
||||
return {"applied": ok, "mode": "auto", "detail": r.get("out") or r.get("err"),
|
||||
"attempts": _attempt + 1}
|
||||
|
||||
speed_pct = max(30, min(100, int(speed_pct)))
|
||||
r = _nvidia_settings(
|
||||
@@ -422,10 +476,76 @@ def restore_safe(reason: str = "shutdown") -> Dict[str, Any]:
|
||||
return result
|
||||
|
||||
|
||||
_APPLIED_ONCE = False
|
||||
# Set once fan control has worked at least once, so drift checks do not fire forever on
|
||||
# a machine that simply has no fan control available.
|
||||
_FAN_AVAILABLE = False
|
||||
|
||||
|
||||
def profile_drift() -> Dict[str, Any]:
|
||||
"""Compare what the active profile asks for against what the card actually reports.
|
||||
|
||||
ACTIVE_PROFILE defaults to "balanced" at import, which is indistinguishable from
|
||||
"balanced was successfully applied" -- so a startup apply that failed left the app
|
||||
confidently reporting a profile it had never put on the hardware. This makes the
|
||||
difference visible instead.
|
||||
"""
|
||||
profiles = load_profiles()
|
||||
cfg = profiles.get(ACTIVE_PROFILE, {})
|
||||
state = get_gpu_state()
|
||||
intended = int(cfg.get("power_limit_w", 0) or 0)
|
||||
actual = state.get("power_limit_w")
|
||||
power_drift = bool(intended and actual is not None
|
||||
and abs(float(actual) - intended) >= 1.0)
|
||||
|
||||
# Fan mode is checked too. The headless X server that owns the GPU can still be
|
||||
# starting when this unit does, and the fan assignment then fails; the in-process
|
||||
# retries cover a short delay, but if X arrives later nothing else would ever notice
|
||||
# that the profile's fan setting was never applied.
|
||||
fan_intended = cfg.get("fan_mode", "auto")
|
||||
fan_actual = None
|
||||
fan_drift = False
|
||||
if _FAN_AVAILABLE:
|
||||
fan_actual = get_fan_status().get("mode")
|
||||
fan_drift = bool(fan_actual and fan_actual != fan_intended)
|
||||
|
||||
drifted = power_drift or fan_drift or not _APPLIED_ONCE
|
||||
reasons = []
|
||||
if not _APPLIED_ONCE:
|
||||
reasons.append("no profile has been successfully applied since startup")
|
||||
if power_drift:
|
||||
reasons.append(f"card reports {actual}W, profile asks {intended}W")
|
||||
if fan_drift:
|
||||
reasons.append(f"fans are {fan_actual}, profile asks {fan_intended}")
|
||||
return {
|
||||
"profile": ACTIVE_PROFILE,
|
||||
"applied_since_start": _APPLIED_ONCE,
|
||||
"power_limit_intended_w": intended,
|
||||
"power_limit_actual_w": actual,
|
||||
"fan_mode_intended": fan_intended,
|
||||
"fan_mode_actual": fan_actual,
|
||||
"fan_available": _FAN_AVAILABLE,
|
||||
"drifted": drifted,
|
||||
"reason": "; ".join(reasons) or None,
|
||||
}
|
||||
|
||||
|
||||
def reconcile_profile() -> Dict[str, Any]:
|
||||
"""Re-apply the active profile if the hardware has drifted away from it."""
|
||||
drift = profile_drift()
|
||||
if not drift["drifted"]:
|
||||
return {"reconciled": False, "drift": drift}
|
||||
logger.warning(f"Overclock drift detected — re-applying '{ACTIVE_PROFILE}': "
|
||||
f"{drift['reason']}")
|
||||
res = apply_profile(ACTIVE_PROFILE)
|
||||
return {"reconciled": True, "drift": drift, "result": res.get("verified")}
|
||||
|
||||
|
||||
def get_status() -> Dict[str, Any]:
|
||||
"""Full overclock status for the dashboard."""
|
||||
return {
|
||||
"active_profile": ACTIVE_PROFILE,
|
||||
"drift": profile_drift(),
|
||||
"offsets_supported": offsets_supported(),
|
||||
"effective_levers": (["power_limit", "clock_lock", "mem_lock", "fan"]
|
||||
+ (["offsets"] if offsets_supported() else [])),
|
||||
|
||||
@@ -1,7 +1,6 @@
|
||||
{
|
||||
"ollama": {
|
||||
"label": "Ollama — LLM decode (memory-bandwidth bound; measured insensitive to power and clocks)",
|
||||
"measured": "73.0-73.5 tok/s flat from 222W to 370W (qwen3.8long, 2026-08-28). Actual draw never exceeded 224W at any limit. Memory clock lock made no difference (72.6 locked vs 72.7 unlocked).",
|
||||
"label": "Ollama \u2014 LLM decode (memory-bandwidth bound; measured insensitive to power and clocks)",
|
||||
"power_limit_w": 320,
|
||||
"core_offset_mhz": 0,
|
||||
"mem_offset_mhz": 0,
|
||||
@@ -9,11 +8,11 @@
|
||||
"lock_core_max": 0,
|
||||
"lock_mem_mhz": 0,
|
||||
"fan_mode": "auto",
|
||||
"fan_speed_pct": 0
|
||||
"fan_speed_pct": 0,
|
||||
"measured": "73.0-73.5 tok/s flat from 222W to 370W (qwen3.8long, 2026-08-28). Actual draw never exceeded 224W at any limit. Memory clock lock made no difference (72.6 locked vs 72.7 unlocked)."
|
||||
},
|
||||
"comfy": {
|
||||
"label": "ComfyUI — diffusion (compute bound; genuinely power-scaling)",
|
||||
"measured": "SDXL 1024/20-step: 5.48 it/s @222W, 6.22 @259W, 6.50 @296W, 6.52 @320W, 6.63 @333W, 6.71 @370W (2026-08-28). Worth +2.8% over the 320W stock default. Core clock lock made no difference across 2400-3105 MHz.",
|
||||
"label": "ComfyUI \u2014 diffusion (compute bound; genuinely power-scaling)",
|
||||
"power_limit_w": 370,
|
||||
"core_offset_mhz": 0,
|
||||
"mem_offset_mhz": 0,
|
||||
@@ -21,18 +20,19 @@
|
||||
"lock_core_max": 0,
|
||||
"lock_mem_mhz": 0,
|
||||
"fan_mode": "auto",
|
||||
"fan_speed_pct": 0
|
||||
"fan_speed_pct": 0,
|
||||
"measured": "SDXL 1024/20-step: 5.48 it/s @222W, 6.22 @259W, 6.50 @296W, 6.52 @320W, 6.63 @333W, 6.71 @370W (2026-08-28). Worth +2.8% over the 320W stock default. Core clock lock made no difference across 2400-3105 MHz."
|
||||
},
|
||||
"balanced": {
|
||||
"label": "Balanced — stock power and boost, automatic fans",
|
||||
"measured": "Card's own design point. 48k telemetry samples show 67.8C average under load at 39.5% auto fan, 81C all-time max, zero thermal throttle events.",
|
||||
"power_limit_w": 320,
|
||||
"core_offset_mhz": 0,
|
||||
"mem_offset_mhz": 0,
|
||||
"label": "Balanced \u2014 stock power and boost, automatic fans",
|
||||
"power_limit_w": 340,
|
||||
"core_offset_mhz": 10,
|
||||
"mem_offset_mhz": 150,
|
||||
"lock_core_min": 0,
|
||||
"lock_core_max": 0,
|
||||
"lock_mem_mhz": 0,
|
||||
"fan_mode": "auto",
|
||||
"fan_speed_pct": 0
|
||||
"fan_mode": "manual",
|
||||
"fan_speed_pct": 95,
|
||||
"measured": "Card's own design point. 48k telemetry samples show 67.8C average under load at 39.5% auto fan, 81C all-time max, zero thermal throttle events."
|
||||
}
|
||||
}
|
||||
}
|
||||
5
pytest.ini
Normal file
5
pytest.ini
Normal file
@@ -0,0 +1,5 @@
|
||||
[pytest]
|
||||
testpaths = tests
|
||||
addopts = -ra
|
||||
filterwarnings =
|
||||
ignore::DeprecationWarning
|
||||
@@ -203,8 +203,10 @@ def page_residency(filepath: str, allow_probe: bool = True,
|
||||
except OSError as e:
|
||||
return {"success": False, "error": str(e), "resident_pct": 0.0, "measurable": False}
|
||||
if size == 0:
|
||||
return {"success": True, "resident_pct": 0.0, "size_bytes": 0, "measurable": True,
|
||||
"method": "empty"}
|
||||
# Keep the shape identical to every other return path; callers index these.
|
||||
return {"success": True, "filepath": filepath, "size_bytes": 0, "size_gb": 0.0,
|
||||
"resident_pct": 0.0, "resident_bytes": 0, "measurable": True,
|
||||
"method": "empty", "warm": True, "warm_confident": True}
|
||||
|
||||
try:
|
||||
fd = os.open(filepath, os.O_RDONLY)
|
||||
@@ -415,6 +417,16 @@ _report_cache: Dict[str, Any] = {"ts": 0.0, "report": None}
|
||||
REPORT_TTL_S = 15.0
|
||||
|
||||
|
||||
def invalidate_cache_report() -> None:
|
||||
"""Drop the cached residency report.
|
||||
|
||||
Anything that changes what is resident must call this, or the report keeps serving
|
||||
pre-change numbers for up to REPORT_TTL_S -- so warming a model and then looking at
|
||||
residency showed the state from before the warm.
|
||||
"""
|
||||
_report_cache["report"] = None
|
||||
|
||||
|
||||
def get_cache_report(include_files: bool = True, force_refresh: bool = False) -> Dict[str, Any]:
|
||||
"""Measured page-cache residency across the whole model catalog.
|
||||
|
||||
@@ -510,6 +522,7 @@ def warm_file_to_ram(filepath: str, chunk_size: int = 16 * 1024 * 1024,
|
||||
|
||||
duration = time.perf_counter() - t0
|
||||
after = page_residency(filepath, probe_windows=32)
|
||||
invalidate_cache_report()
|
||||
return {
|
||||
"success": True,
|
||||
"filepath": filepath,
|
||||
|
||||
232
server.py
232
server.py
@@ -16,9 +16,13 @@ from fastapi.middleware.cors import CORSMiddleware
|
||||
from pydantic import BaseModel, Field
|
||||
|
||||
import autotune
|
||||
import engines
|
||||
import health
|
||||
import jobs as jobs_mod
|
||||
import overclock_manager
|
||||
import ram_optimizer
|
||||
import telemetry_store
|
||||
import tenants as tenants_mod
|
||||
import thermal_governor
|
||||
import vram_arbitrator
|
||||
|
||||
@@ -49,6 +53,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 +63,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 +94,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 +134,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:
|
||||
@@ -186,10 +221,13 @@ def _install_shutdown_hook() -> None:
|
||||
@asynccontextmanager
|
||||
async def lifespan(app: FastAPI):
|
||||
telemetry_store.start()
|
||||
jobs_mod.init()
|
||||
await broker.start()
|
||||
await jobs_mod.scheduler.start()
|
||||
await vram_arbitrator.arbitrator.start()
|
||||
_install_shutdown_hook()
|
||||
yield
|
||||
await jobs_mod.scheduler.stop()
|
||||
await vram_arbitrator.arbitrator.stop()
|
||||
await broker.stop()
|
||||
# Never leave the card with locked clocks and pinned fans after we exit.
|
||||
@@ -279,6 +317,151 @@ 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()
|
||||
|
||||
|
||||
class TenantReleaseRequest(BaseModel):
|
||||
models: Optional[List[str]] = Field(None, description="For per-model tenants (Ollama), which to unload; defaults to everything resident")
|
||||
confirm: bool = Field(True, description="Wait for NVML to confirm the VRAM was actually released")
|
||||
|
||||
|
||||
class JobRequest(BaseModel):
|
||||
tenant: str = Field(..., description="Which application should run this job", example="comfyui")
|
||||
payload: Dict[str, Any] = Field(..., description="What to run: a ComfyUI workflow under 'prompt', or Ollama generate parameters")
|
||||
priority: Optional[int] = Field(None, description="Defaults to the tenant's priority; higher runs sooner")
|
||||
label: Optional[str] = Field(None, description="Human-readable name for the queue view")
|
||||
|
||||
|
||||
@app.post("/api/jobs", summary="Queue a Job", tags=["Jobs"])
|
||||
async def api_submit_job(req: JobRequest):
|
||||
"""Queue work for any tenant. The queue is on disk, so there is no depth limit and
|
||||
it survives a restart. Jobs run one at a time, highest priority first, with VRAM
|
||||
arbitrated before each one starts."""
|
||||
res = jobs_mod.submit(req.tenant, req.payload, req.priority, req.label)
|
||||
if not res.get("success"):
|
||||
raise HTTPException(status_code=400, detail=res.get("error"))
|
||||
return res
|
||||
|
||||
|
||||
@app.get("/api/jobs", summary="The Job Queue", tags=["Jobs"])
|
||||
async def api_jobs(state: Optional[str] = Query(None, description="pending | running | done | failed | cancelled"),
|
||||
limit: int = Query(100)):
|
||||
"""Queued and recent jobs. Pending jobs are listed in the order they will run."""
|
||||
return {"jobs": jobs_mod.listing(state, limit), "scheduler": jobs_mod.scheduler.get_status()}
|
||||
|
||||
|
||||
@app.get("/api/jobs/{job_id}", summary="One Job", tags=["Jobs"])
|
||||
async def api_job(job_id: str):
|
||||
job = jobs_mod.get(job_id)
|
||||
if not job:
|
||||
raise HTTPException(status_code=404, detail=f"no job '{job_id}'")
|
||||
return job
|
||||
|
||||
|
||||
@app.delete("/api/jobs/{job_id}", summary="Cancel a Pending Job", tags=["Jobs"])
|
||||
async def api_cancel_job(job_id: str):
|
||||
"""Cancel a job that has not started. Running jobs are left alone -- this service
|
||||
frees VRAM by asking, never by killing work in flight."""
|
||||
res = jobs_mod.cancel(job_id)
|
||||
if not res.get("success"):
|
||||
raise HTTPException(status_code=409, detail=res.get("error"))
|
||||
return res
|
||||
|
||||
|
||||
@app.delete("/api/jobs", summary="Cancel All Pending Jobs", tags=["Jobs"])
|
||||
async def api_clear_jobs():
|
||||
return jobs_mod.clear_pending()
|
||||
|
||||
|
||||
@app.get("/api/tenants", summary="GPU Tenants", tags=["Tenants"])
|
||||
async def api_tenants():
|
||||
"""Applications competing for the GPU, as configured.
|
||||
|
||||
Each entry declares how its processes are recognised, how to tell whether it is
|
||||
working, and how to ask it for VRAM back. Adding an application is a config change
|
||||
in tenants.json, not a code change.
|
||||
"""
|
||||
gpu = vram_arbitrator.get_gpu_hardware_stats()
|
||||
by_tenant = (gpu.get("breakdown", {}) or {}).get("by_tenant_gb", {})
|
||||
out = []
|
||||
for t in tenants_mod.describe():
|
||||
name = t["name"]
|
||||
# The two original tenants are reported under the bucket names the API has
|
||||
# always used.
|
||||
bucket = {"comfyui": "comfy"}.get(name, name)
|
||||
t["vram_gb"] = by_tenant.get(bucket, 0.0)
|
||||
out.append(t)
|
||||
return {"tenants": out, "config_path": tenants_mod.CONFIG_PATH,
|
||||
"unmanaged_gb": (gpu.get("breakdown", {}) or {}).get("unmanaged_gb", 0.0)}
|
||||
|
||||
|
||||
@app.get("/api/tenants/{name}", summary="One GPU Tenant", tags=["Tenants"])
|
||||
async def api_tenant(name: str):
|
||||
"""A single tenant's definition, current VRAM, and whether it is genuinely busy."""
|
||||
t = tenants_mod.get_tenant(name)
|
||||
if not t:
|
||||
raise HTTPException(status_code=404, detail=f"no tenant named '{name}'")
|
||||
gpu = vram_arbitrator.get_gpu_hardware_stats()
|
||||
bucket = {"comfyui": "comfy"}.get(name, name)
|
||||
vram_gb = (gpu.get("breakdown", {}) or {}).get("by_tenant_gb", {}).get(bucket, 0.0)
|
||||
busy = await tenants_mod.probe_busy(t, vram_gb=vram_gb)
|
||||
d = t.to_dict()
|
||||
d.update({"vram_gb": vram_gb, "reclaimable": t.reclaimable, "busy": busy})
|
||||
return d
|
||||
|
||||
|
||||
@app.post("/api/tenants/{name}/release", summary="Ask a Tenant for its VRAM", tags=["Tenants"])
|
||||
async def api_tenant_release(name: str, req: Optional[TenantReleaseRequest] = None):
|
||||
"""Release a tenant's VRAM using whatever mechanism that tenant declares.
|
||||
|
||||
This is the generic form of the Ollama soft-yield and the ComfyUI purge: the same
|
||||
request works for any application in the registry, including ones added later.
|
||||
"""
|
||||
t = tenants_mod.get_tenant(name)
|
||||
if not t:
|
||||
raise HTTPException(status_code=404, detail=f"no tenant named '{name}'")
|
||||
if not t.reclaimable:
|
||||
raise HTTPException(status_code=409,
|
||||
detail=f"'{name}' declares no way to release VRAM; its "
|
||||
f"memory cannot be reclaimed by this service")
|
||||
|
||||
models = req.models if req else None
|
||||
if t.release.per_model and not models:
|
||||
state = await vram_arbitrator.get_ollama_live_state()
|
||||
models = [m.get("name") for m in state.get("loaded_models", []) if m.get("name")]
|
||||
|
||||
before = vram_arbitrator.get_process_vram_bytes()
|
||||
res = await tenants_mod.release_vram(t, models=models)
|
||||
|
||||
if (req is None or req.confirm) and res.get("released"):
|
||||
bucket = {"comfyui": "comfy"}.get(name, name)
|
||||
key = {"ollama": "ollama_bytes", "comfy": "comfyui_bytes"}.get(bucket)
|
||||
if key:
|
||||
baseline = before[key]
|
||||
barrier = await vram_arbitrator._await_vram_release(baseline) \
|
||||
if key == "ollama_bytes" else None
|
||||
if barrier:
|
||||
res.update({"outcome": barrier.get("outcome"),
|
||||
"confirm_ms": barrier.get("confirm_ms")})
|
||||
after = vram_arbitrator.get_process_vram_bytes()
|
||||
res["free_vram_gb"] = round(after["free_bytes"] / (1024**3), 2)
|
||||
return res
|
||||
|
||||
|
||||
@app.get("/api/engines", summary="Live Engine Configuration", tags=["Telemetry"])
|
||||
async def api_engines():
|
||||
"""Real configuration of Ollama and ComfyUI, with what each setting implies for
|
||||
arbitration. These live outside this codebase but dictate how it must behave."""
|
||||
return await engines.get_engine_config()
|
||||
|
||||
|
||||
@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 +479,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():
|
||||
@@ -332,13 +516,35 @@ async def api_switch_model(req: SwitchRequest):
|
||||
await vram_arbitrator.arbitrator.request_vram_for_ollama()
|
||||
res = await vram_arbitrator.switch_ollama_model(req.model, keep_alive=req.keep_alive or "30m")
|
||||
if not res.get("success"):
|
||||
raise HTTPException(status_code=500, detail=res.get("error"))
|
||||
# Reflect what actually went wrong. Ollama returns 400 for an unusable request --
|
||||
# asking an embedding model to generate, say -- and reporting that as 500 blames
|
||||
# this service for the caller's mistake. A model that will not fit is neither:
|
||||
# the request is valid and the service is healthy, there is simply no room.
|
||||
upstream = res.get("upstream_status")
|
||||
if res.get("vram_oom"):
|
||||
status = 507 # Insufficient Storage
|
||||
elif isinstance(upstream, int) and 400 <= upstream < 500:
|
||||
status = 400
|
||||
else:
|
||||
status = 502 if upstream else 500
|
||||
raise HTTPException(status_code=status, detail=res.get("error"))
|
||||
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():
|
||||
@@ -541,9 +747,23 @@ app.mount("/static", StaticFiles(directory=f"{BASE_DIR}/static"), name="static")
|
||||
|
||||
@app.get("/", summary="Dashboard Web UI", tags=["UI"])
|
||||
async def root_index():
|
||||
"""Serve the dashboard with cache-busted asset URLs.
|
||||
|
||||
StaticFiles sends an ETag, but browsers were still serving app.js from cache after
|
||||
it changed, so a reload showed the old dashboard against the new API -- a panel that
|
||||
had just been added simply never appeared. Stamping each asset with its mtime means
|
||||
a changed file is always a different URL.
|
||||
"""
|
||||
with open(f"{BASE_DIR}/static/index.html", "r") as f:
|
||||
content = f.read()
|
||||
return HTMLResponse(content=content)
|
||||
for asset in ("app.js", "styles.css"):
|
||||
try:
|
||||
stamp = int(os.path.getmtime(f"{BASE_DIR}/static/{asset}"))
|
||||
except OSError:
|
||||
continue
|
||||
content = content.replace(f"/static/{asset}", f"/static/{asset}?v={stamp}")
|
||||
return HTMLResponse(content=content,
|
||||
headers={"Cache-Control": "no-cache, must-revalidate"})
|
||||
|
||||
if __name__ == "__main__":
|
||||
import uvicorn
|
||||
|
||||
391
static/app.js
391
static/app.js
@@ -33,8 +33,10 @@ 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);
|
||||
if (data.arbitrator) renderTenants(data.arbitrator, data.gpu);
|
||||
|
||||
// 1. GPU VRAM Stats
|
||||
const gpu = data.gpu || {};
|
||||
@@ -153,6 +155,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 +201,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 +249,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 +915,348 @@ document.addEventListener('DOMContentLoaded', () => {
|
||||
if (d.last_result) renderSweep(d.last_result);
|
||||
}).catch(() => {});
|
||||
});
|
||||
|
||||
|
||||
// ---------------------------------------------------------------- arbitration
|
||||
|
||||
function renderArbitrator(arb, gpu) {
|
||||
const el = (id) => document.getElementById(id);
|
||||
if (!el('arb-action')) return;
|
||||
const c = arb.counters || {};
|
||||
|
||||
// VRAM held by processes HyperSwap cannot reclaim. Worth showing: it is headroom the
|
||||
// arbitrator can never give back, no matter how much it purges.
|
||||
const bd = (gpu && gpu.breakdown) || {};
|
||||
const un = el('arb-unmanaged');
|
||||
if (un) {
|
||||
const procs = bd.unmanaged || [];
|
||||
un.innerHTML = procs.length
|
||||
? `<span class="text-amber-400">${bd.unmanaged_gb} GB unreclaimable</span> — ` +
|
||||
procs.map(p => `${p.name} (${p.vram_mb} MB)`).join(', ')
|
||||
: '';
|
||||
}
|
||||
|
||||
el('arb-action').textContent = arb.last_action || 'Idle';
|
||||
el('arb-yields').textContent = c.yields ?? 0;
|
||||
el('arb-busy').textContent = c.yield_deferred_busy ?? 0;
|
||||
el('arb-later').textContent = c.deferred_releases ?? 0;
|
||||
el('arb-stalled').textContent = c.yield_stalled ?? 0;
|
||||
el('arb-purges').textContent = c.purges ?? 0;
|
||||
el('arb-defpurge').textContent = c.deferred_purges ?? 0;
|
||||
|
||||
const ws = el('arb-ws');
|
||||
ws.textContent = arb.connected_ws ? 'ComfyUI WS live' : 'WS down — polling';
|
||||
ws.className = 'text-xs font-mono ' + (arb.connected_ws ? 'text-emerald-400' : 'text-amber-400');
|
||||
|
||||
// Show why we are holding off, and the idle countdown before ComfyUI is purged.
|
||||
const parts = [];
|
||||
const backoff = arb.yield_backoff || {};
|
||||
for (const [model, secs] of Object.entries(backoff)) {
|
||||
parts.push(`waiting ${secs}s before asking '${model}' again`);
|
||||
}
|
||||
if (arb.pending_purge && arb.comfy_idle_s != null) {
|
||||
const left = Math.max((arb.idle_purge_after_s || 0) - arb.comfy_idle_s, 0).toFixed(0);
|
||||
parts.push(`ComfyUI idle ${arb.comfy_idle_s}s — holding its checkpoints ${left}s longer`);
|
||||
}
|
||||
el('arb-backoff').textContent = parts.join(' · ');
|
||||
}
|
||||
|
||||
|
||||
// ---------------------------------------------------------------- health
|
||||
|
||||
async function fetchHealth(verbose = false) {
|
||||
const badge = document.getElementById('health-badge');
|
||||
const body = document.getElementById('health-body');
|
||||
if (!badge) return;
|
||||
try {
|
||||
const d = await (await fetch('/api/health')).json();
|
||||
const style = {
|
||||
ok: 'bg-emerald-950/70 border border-emerald-800 text-emerald-300',
|
||||
degraded: 'bg-amber-950/70 border border-amber-800 text-amber-300',
|
||||
failed: 'bg-rose-950/70 border border-rose-800 text-rose-300',
|
||||
}[d.status];
|
||||
badge.className = `px-2 py-1 text-xs font-bold rounded-lg ${style}`;
|
||||
badge.textContent = `${d.status.toUpperCase()} · ${d.summary}`;
|
||||
|
||||
// Healthy checks collapse to one line; anything wrong is shown in full with the
|
||||
// impact and the fix, because that is the information you actually need.
|
||||
const problems = d.checks.filter(c => c.status !== 'ok');
|
||||
const shown = (verbose || problems.length) ? (verbose ? d.checks : problems) : [];
|
||||
body.innerHTML = shown.map(c => {
|
||||
const colour = c.status === 'ok' ? 'text-emerald-400'
|
||||
: c.status === 'degraded' ? 'text-amber-400' : 'text-rose-400';
|
||||
const mark = c.status === 'ok' ? '✓' : c.status === 'degraded' ? '!' : '✗';
|
||||
let html = `<div><span class="${colour}">${mark} ${c.name}</span>` +
|
||||
`<span class="text-slate-400"> — ${c.detail}</span></div>`;
|
||||
if (c.status !== 'ok') {
|
||||
if (c.impact) html += `<div class="text-slate-500 pl-4">impact: ${c.impact}</div>`;
|
||||
if (c.fix) html += `<div class="text-slate-500 pl-4">fix: ${c.fix}</div>`;
|
||||
}
|
||||
return html;
|
||||
}).join('') || '<div class="text-emerald-500">All dependencies healthy.</div>';
|
||||
} catch (e) {
|
||||
badge.className = 'px-2 py-1 text-xs font-bold rounded-lg bg-rose-950/70 border border-rose-800 text-rose-300';
|
||||
badge.textContent = 'UNREACHABLE';
|
||||
body.innerHTML = `<div class="text-rose-400">${e}</div>`;
|
||||
}
|
||||
}
|
||||
|
||||
document.addEventListener('DOMContentLoaded', () => {
|
||||
fetchHealth(false);
|
||||
setInterval(() => fetchHealth(false), 60000);
|
||||
});
|
||||
|
||||
|
||||
// ---------------------------------------------------------------- honest labels
|
||||
|
||||
async function fetchMeasuredResidency() {
|
||||
const el = document.getElementById('legend-models-resident');
|
||||
if (!el) return;
|
||||
try {
|
||||
const d = await (await fetch('/api/cache/report?files=false')).json();
|
||||
// Deliberately distinct from the page-cache figure beside it: that is every file
|
||||
// the kernel has cached; this is model weights, measured.
|
||||
el.textContent = `${d.resident_total_gb} GB of ${d.catalog_total_gb} GB`;
|
||||
const cr = document.getElementById('comfy-resident');
|
||||
if (cr) cr.textContent = `${d.residency_pct}%`;
|
||||
} catch (e) {
|
||||
el.textContent = 'unavailable';
|
||||
}
|
||||
}
|
||||
|
||||
async function fetchVersion() {
|
||||
const el = document.getElementById('app-version');
|
||||
if (!el) return;
|
||||
try {
|
||||
const d = await (await fetch('/openapi.json')).json();
|
||||
el.textContent = `v${d.info.version}`;
|
||||
} catch (e) {
|
||||
el.textContent = 'v?';
|
||||
}
|
||||
}
|
||||
|
||||
document.addEventListener('DOMContentLoaded', () => {
|
||||
fetchVersion();
|
||||
fetchMeasuredResidency();
|
||||
setInterval(fetchMeasuredResidency, 30000);
|
||||
});
|
||||
|
||||
|
||||
// ---------------------------------------------------------------- durable swap stats
|
||||
|
||||
let lastSwapFromStore = null;
|
||||
|
||||
async function fetchLastSwapFromStore() {
|
||||
// The in-memory ring empties on restart, which blanked the swap tiles even though
|
||||
// every swap is on disk. Fall back to the persisted log so a restart does not erase
|
||||
// what the dashboard knows.
|
||||
try {
|
||||
const rows = await (await fetch('/api/history?durable=true&limit=50')).json();
|
||||
const swap = (rows || []).find(r => r.event_type === 'LLM Model Switch');
|
||||
if (!swap) return;
|
||||
lastSwapFromStore = swap;
|
||||
const t = document.getElementById('ollama-last-swap');
|
||||
const c = document.getElementById('ollama-cache-hit');
|
||||
if (t && (!t.textContent || t.textContent.includes('\u2014') || t.textContent.includes('--'))) {
|
||||
t.textContent = `${Math.round(swap.duration_ms)} ms`;
|
||||
}
|
||||
if (c && (!c.textContent || c.textContent.includes('\u2014') || c.textContent.includes('--')
|
||||
|| c.textContent === 'no swaps yet')) {
|
||||
c.textContent = swap.cache_status || 'OK';
|
||||
}
|
||||
} catch (e) { /* the live stream is the primary source; this is only a fallback */ }
|
||||
}
|
||||
|
||||
document.addEventListener('DOMContentLoaded', () => {
|
||||
setTimeout(fetchLastSwapFromStore, 1500);
|
||||
});
|
||||
|
||||
|
||||
// ---------------------------------------------------------------- engine config
|
||||
|
||||
async function fetchEngineConfig() {
|
||||
// These subtitles used to be hardcoded. They happened to be accurate, which is worse
|
||||
// than being wrong: they would have stayed accurate-looking after the settings changed.
|
||||
try {
|
||||
const d = await (await fetch('/api/engines')).json();
|
||||
const o = document.getElementById('ollama-engine-sub');
|
||||
if (o && d.ollama) {
|
||||
const bits = [`Port :${d.ollama.port}`, d.ollama.summary];
|
||||
if (d.ollama.max_loaded_models) bits.push(`${d.ollama.max_loaded_models} model resident`);
|
||||
if (d.ollama.keep_alive) bits.push(`keep-alive ${d.ollama.keep_alive}`);
|
||||
o.textContent = bits.join(' // ');
|
||||
o.title = (d.ollama.settings || [])
|
||||
.filter(s => s.means)
|
||||
.map(s => `${s.key}=${s.value} — ${s.means}`)
|
||||
.join('\n');
|
||||
}
|
||||
const c = document.getElementById('comfy-engine-sub');
|
||||
if (c && d.comfyui && d.comfyui.online) {
|
||||
c.textContent = `Port :8188 // v${d.comfyui.version} // ${d.comfyui.summary}`;
|
||||
c.title = `torch ${d.comfyui.pytorch}\n${d.comfyui.device || ''}`;
|
||||
}
|
||||
} catch (e) { /* subtitles are cosmetic; never break the page over them */ }
|
||||
}
|
||||
|
||||
document.addEventListener('DOMContentLoaded', () => {
|
||||
fetchEngineConfig();
|
||||
setInterval(fetchEngineConfig, 120000);
|
||||
});
|
||||
|
||||
|
||||
// ---------------------------------------------------------------- chart sizing
|
||||
|
||||
// Chart.js is configured responsive with maintainAspectRatio:false, so it should track
|
||||
// its container on its own. In practice it latched onto a stale size -- the canvas sat
|
||||
// at width:0px, height:288px (the old fixed h-64) while its container had grown to
|
||||
// 988x648 on a large window. An explicit observer makes the chart follow the container
|
||||
// whatever the window does.
|
||||
function watchChartSize() {
|
||||
const canvas = document.getElementById('oc-chart');
|
||||
if (!canvas || !canvas.parentElement) return;
|
||||
const container = canvas.parentElement;
|
||||
|
||||
const resize = () => {
|
||||
if (typeof ocChart === 'undefined' || !ocChart) return;
|
||||
// No explicit dimensions: with responsive + maintainAspectRatio:false, Chart.js
|
||||
// measures the container itself. Passing width/height instead made the canvas grow
|
||||
// but never shrink -- it ended up 988px wide inside a 435px container, overflowing
|
||||
// it, which is also why the canvas is absolutely positioned now.
|
||||
ocChart.resize();
|
||||
};
|
||||
|
||||
if (typeof ResizeObserver !== 'undefined') {
|
||||
new ResizeObserver(resize).observe(container);
|
||||
}
|
||||
window.addEventListener('resize', resize);
|
||||
// Run once after layout settles, in case the chart was built before the container had
|
||||
// a width (which is how it ended up at 0 in the first place).
|
||||
requestAnimationFrame(resize);
|
||||
setTimeout(resize, 300);
|
||||
}
|
||||
|
||||
document.addEventListener('DOMContentLoaded', () => setTimeout(watchChartSize, 200));
|
||||
|
||||
|
||||
// ---------------------------------------------------------------- tenants
|
||||
|
||||
function renderTenants(arb, gpu) {
|
||||
const body = document.getElementById('tenants-body');
|
||||
if (!body) return;
|
||||
const state = arb && arb.tenant_state;
|
||||
if (!state || !state.tenants) return;
|
||||
|
||||
const total = (gpu && gpu.vram_total_gb) || 16;
|
||||
document.getElementById('tenants-free').textContent = `${state.free_gb} GB free`;
|
||||
|
||||
// Sorted by priority, the order arbitration actually considers them in.
|
||||
const rows = [...state.tenants].sort((a, b) => b.priority - a.priority);
|
||||
body.innerHTML = rows.map(t => {
|
||||
const pct = Math.min((t.vram_gb / total) * 100, 100);
|
||||
const bar = t.busy ? 'bg-emerald-500'
|
||||
: t.reclaimable ? 'bg-cyan-600' : 'bg-amber-600';
|
||||
const badge = t.busy
|
||||
? '<span class="text-emerald-400">working</span>'
|
||||
: t.reclaimable
|
||||
? '<span class="text-slate-500">idle · reclaimable</span>'
|
||||
: '<span class="text-amber-400">cannot be reclaimed</span>';
|
||||
return `<div>
|
||||
<div class="flex justify-between text-[11px] font-mono">
|
||||
<span class="text-slate-200">${t.name}
|
||||
<span class="text-slate-600">p${t.priority}</span></span>
|
||||
<span class="text-slate-400">${t.vram_gb.toFixed(2)} GB · ${badge}</span>
|
||||
</div>
|
||||
<div class="w-full bg-slate-950 rounded-full h-1.5 mt-1 overflow-hidden border border-slate-800/60">
|
||||
<div class="${bar} h-full transition-all duration-500" style="width:${pct}%"></div>
|
||||
</div>
|
||||
<div class="text-[10px] text-slate-600 mt-0.5">${t.reason || ''}${
|
||||
t.needs_vram_gb ? ` · needs ${t.needs_vram_gb} GB to work` : ''}</div>
|
||||
</div>`;
|
||||
}).join('');
|
||||
|
||||
// The most recent arbitration decision, including why it could not be satisfied.
|
||||
const dec = document.getElementById('tenants-decision');
|
||||
const a = arb.last_arbitration;
|
||||
if (!a) {
|
||||
dec.innerHTML = '<span class="text-slate-600">No contention — nothing has needed to be released.</span>';
|
||||
return;
|
||||
}
|
||||
const when = new Date(a.ts * 1000).toLocaleTimeString();
|
||||
const blockers = (a.blockers || [])
|
||||
.map(b => `${b.name} (${b.vram_gb} GB, ${b.why})`).join(', ');
|
||||
dec.innerHTML =
|
||||
`<span class="${a.possible ? 'text-cyan-400' : 'text-amber-400'}">${when} · ` +
|
||||
`${a.demanding} short by ${a.shortfall_gb} GB</span> — ${a.reason}` +
|
||||
(blockers ? `<div class="text-slate-600">blocked by: ${blockers}</div>` : '');
|
||||
}
|
||||
|
||||
|
||||
// ---------------------------------------------------------------- job queue
|
||||
|
||||
async function fetchQueue() {
|
||||
const body = document.getElementById('queue-body');
|
||||
if (!body) return;
|
||||
try {
|
||||
const d = await (await fetch('/api/jobs?limit=40')).json();
|
||||
const s = d.scheduler || {};
|
||||
document.getElementById('queue-summary').textContent =
|
||||
`${s.queue_depth ?? 0} queued · ${s.completed ?? 0} done · ${s.failed ?? 0} failed`;
|
||||
|
||||
// What the scheduler is doing right now, including why it is waiting. A job that
|
||||
// cannot get VRAM used to sit silent, which made a stuck queue indistinguishable
|
||||
// from an empty one.
|
||||
const cur = document.getElementById('queue-current');
|
||||
if (s.current) {
|
||||
cur.innerHTML = `<span class="text-emerald-400">running</span> ` +
|
||||
`<span class="text-slate-200">${s.current.tenant}</span>` +
|
||||
`<span class="text-slate-500"> · ${s.current.label || s.current.id}</span>` +
|
||||
(s.current.room && s.current.room.released && s.current.room.released.length
|
||||
? `<span class="text-cyan-400"> · released ${s.current.room.released.join(', ')}</span>` : '');
|
||||
} else if (s.blocked) {
|
||||
cur.innerHTML = `<span class="text-amber-400">waiting ${s.blocked.waited_s}s</span> ` +
|
||||
`<span class="text-slate-200">${s.blocked.tenant}</span>` +
|
||||
`<span class="text-slate-500"> · ${s.blocked.reason || ''}</span>`;
|
||||
} else {
|
||||
cur.innerHTML = '<span class="text-slate-600">scheduler idle</span>';
|
||||
}
|
||||
|
||||
const colour = {
|
||||
pending: 'text-slate-400', running: 'text-emerald-400', done: 'text-cyan-500',
|
||||
failed: 'text-rose-400', cancelled: 'text-slate-600',
|
||||
};
|
||||
body.innerHTML = (d.jobs || []).map(job => {
|
||||
const when = job.state === 'pending'
|
||||
? `waiting ${job.waiting_s}s`
|
||||
: (job.duration_s != null ? `${job.duration_s}s` : '');
|
||||
return `<div class="flex justify-between text-[11px] font-mono gap-2">
|
||||
<span class="truncate">
|
||||
<span class="${colour[job.state] || 'text-slate-400'}">${job.state}</span>
|
||||
<span class="text-slate-600"> p${job.priority}</span>
|
||||
<span class="text-slate-200"> ${job.tenant}</span>
|
||||
<span class="text-slate-500">${job.label ? ' · ' + job.label : ''}</span>
|
||||
</span>
|
||||
<span class="text-slate-500 whitespace-nowrap">${when}${
|
||||
job.state === 'pending'
|
||||
? ` <button onclick="cancelJob('${job.id}')" class="text-rose-500 hover:text-rose-300 ml-1">cancel</button>`
|
||||
: ''}</span>
|
||||
</div>${job.error ? `<div class="text-[10px] text-rose-500/80 pl-2 truncate">${job.error}</div>` : ''}`;
|
||||
}).join('') || '<div class="text-slate-600 text-[11px]">No jobs yet.</div>';
|
||||
} catch (e) {
|
||||
body.innerHTML = `<div class="text-rose-400 text-[11px]">${e}</div>`;
|
||||
}
|
||||
}
|
||||
|
||||
async function cancelJob(id) {
|
||||
await fetch(`/api/jobs/${id}`, { method: 'DELETE' });
|
||||
fetchQueue();
|
||||
}
|
||||
|
||||
async function clearQueue() {
|
||||
await fetch('/api/jobs', { method: 'DELETE' });
|
||||
fetchQueue();
|
||||
}
|
||||
|
||||
document.addEventListener('DOMContentLoaded', () => {
|
||||
fetchQueue();
|
||||
setInterval(fetchQueue, 2000);
|
||||
});
|
||||
|
||||
@@ -5,6 +5,11 @@
|
||||
<meta name="viewport" content="width=device-width, initial-scale=1.0">
|
||||
<title>HYPERSWAP // Dual-Engine Model Orchestrator & Live Telemetry</title>
|
||||
<script src="https://cdn.tailwindcss.com"></script>
|
||||
<script>
|
||||
// The CDN build has no breakpoint above 2xl, so a very wide window kept a
|
||||
// two-column layout with increasingly stretched panels.
|
||||
tailwind.config = { theme: { extend: { screens: { '3xl': '2000px' } } } };
|
||||
</script>
|
||||
<script src="https://cdn.jsdelivr.net/npm/chart.js@4.4.1/dist/chart.umd.min.js"></script>
|
||||
<link rel="stylesheet" href="/static/styles.css">
|
||||
<link rel="stylesheet" href="https://cdnjs.cloudflare.com/ajax/libs/font-awesome/6.4.0/css/all.min.css">
|
||||
@@ -13,7 +18,7 @@
|
||||
|
||||
<!-- TOP HEADER -->
|
||||
<header class="border-b border-slate-800 bg-slate-900/80 backdrop-blur sticky top-0 z-50">
|
||||
<div class="max-w-7xl mx-auto px-4 sm:px-6 lg:px-8 py-3 flex flex-wrap items-center justify-between gap-4">
|
||||
<div class="max-w-[2600px] mx-auto px-4 sm:px-6 lg:px-8 py-3 flex flex-wrap items-center justify-between gap-4">
|
||||
<div class="flex items-center space-x-3">
|
||||
<div class="w-10 h-10 rounded-xl bg-gradient-to-tr from-cyan-500 via-indigo-500 to-purple-500 flex items-center justify-center shadow-lg shadow-cyan-500/20">
|
||||
<i class="fa-solid fa-bolt-lightning text-white text-lg"></i>
|
||||
@@ -24,7 +29,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>
|
||||
@@ -60,7 +65,7 @@
|
||||
</header>
|
||||
|
||||
<!-- MAIN CONTAINER -->
|
||||
<main class="max-w-7xl mx-auto px-4 sm:px-6 lg:px-8 py-6 space-y-6">
|
||||
<main class="max-w-[2600px] mx-auto px-4 sm:px-6 lg:px-8 py-6 space-y-6">
|
||||
|
||||
<!-- HERO MEMORY GAUGES -->
|
||||
<div class="grid grid-cols-1 md:grid-cols-2 gap-6">
|
||||
@@ -115,9 +120,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 +169,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 +177,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>
|
||||
@@ -184,7 +199,7 @@
|
||||
</div>
|
||||
<div>
|
||||
<h3 class="font-bold text-slate-100 text-sm">Ollama LLM Engine</h3>
|
||||
<p class="text-xs text-slate-400">Port :11434 // FlashAttention + Q4 KV Cache</p>
|
||||
<p class="text-xs text-slate-400" id="ollama-engine-sub" title="Read live from the ollama service environment">Port :11434</p>
|
||||
</div>
|
||||
</div>
|
||||
<button onclick="freeOllamaVRAM()" class="px-2.5 py-1 text-xs font-semibold rounded-lg bg-rose-950/70 border border-rose-800 text-rose-300 hover:bg-rose-900 transition flex items-center space-x-1">
|
||||
@@ -255,7 +270,7 @@
|
||||
</div>
|
||||
<div>
|
||||
<h3 class="font-bold text-slate-100 text-sm">ComfyUI Diffusion Engine</h3>
|
||||
<p class="text-xs text-slate-400">Port :8188 // DynamicVRAM + Pinned Async Offload</p>
|
||||
<p class="text-xs text-slate-400" id="comfy-engine-sub" title="Read live from ComfyUI's /system_stats">Port :8188</p>
|
||||
</div>
|
||||
</div>
|
||||
<button onclick="freeComfyVRAM()" class="px-2.5 py-1 text-xs font-semibold rounded-lg bg-rose-950/70 border border-rose-800 text-rose-300 hover:bg-rose-900 transition flex items-center space-x-1">
|
||||
@@ -286,26 +301,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>
|
||||
@@ -425,7 +436,7 @@
|
||||
</div>
|
||||
|
||||
<!-- OVERCLOCK CONTROL PANEL -->
|
||||
<div class="bg-slate-900/80 border border-fuchsia-900/50 rounded-2xl p-5 space-y-4 shadow-lg shadow-fuchsia-950/30">
|
||||
<div class="lg:col-span-2 bg-slate-900/80 border border-fuchsia-900/50 rounded-2xl p-5 space-y-4 shadow-lg shadow-fuchsia-950/30">
|
||||
<div class="flex flex-wrap items-center justify-between gap-3 pb-3 border-b border-slate-800">
|
||||
<div class="flex items-center space-x-2">
|
||||
<div class="p-2 rounded-lg bg-fuchsia-950/80 border border-fuchsia-800 text-fuchsia-400">
|
||||
@@ -593,8 +604,8 @@
|
||||
<span class="flex items-center space-x-1"><span class="w-2 h-2 rounded-full bg-fuchsia-400 inline-block"></span>RAM Cache GB</span>
|
||||
</div>
|
||||
</div>
|
||||
<div class="relative h-64">
|
||||
<canvas id="oc-chart"></canvas>
|
||||
<div class="relative h-[clamp(18rem,45vh,52rem)]">
|
||||
<canvas id="oc-chart" class="absolute inset-0 !w-full !h-full"></canvas>
|
||||
</div>
|
||||
</div>
|
||||
|
||||
@@ -666,7 +677,121 @@
|
||||
</div>
|
||||
|
||||
<!-- ============ NEXT-LEVEL PANELS: governor / residency / analytics / autotune ============ -->
|
||||
<div class="grid grid-cols-1 xl:grid-cols-2 gap-5 mt-5">
|
||||
<div class="grid grid-cols-1 xl:grid-cols-2 3xl:grid-cols-3 gap-5 mt-5">
|
||||
|
||||
|
||||
|
||||
|
||||
<!-- The cross-application job queue -->
|
||||
<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-violet-950/80 border border-violet-800 text-violet-400">
|
||||
<i class="fa-solid fa-list-check text-sm"></i>
|
||||
</div>
|
||||
<div>
|
||||
<h3 class="font-bold text-slate-100 text-sm">Job Queue</h3>
|
||||
<p class="text-xs text-slate-400">Work lined up across every application, highest priority first</p>
|
||||
</div>
|
||||
</div>
|
||||
<div class="flex items-center space-x-2">
|
||||
<span id="queue-summary" class="text-xs font-mono text-slate-500">—</span>
|
||||
<button onclick="clearQueue()" 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">Clear pending</button>
|
||||
</div>
|
||||
</div>
|
||||
<div id="queue-current" class="mt-4 text-[11px] font-mono"></div>
|
||||
<div id="queue-body" class="mt-3 space-y-1.5 max-h-64 overflow-y-auto pr-1"></div>
|
||||
</div>
|
||||
|
||||
<!-- All GPU tenants, however many are configured -->
|
||||
<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-teal-950/80 border border-teal-800 text-teal-400">
|
||||
<i class="fa-solid fa-layer-group text-sm"></i>
|
||||
</div>
|
||||
<div>
|
||||
<h3 class="font-bold text-slate-100 text-sm">GPU Tenants</h3>
|
||||
<p class="text-xs text-slate-400">Every application contending for the card, from <code class="text-teal-400">tenants.json</code></p>
|
||||
</div>
|
||||
</div>
|
||||
<span id="tenants-free" class="text-xs font-mono text-slate-500">—</span>
|
||||
</div>
|
||||
<div id="tenants-body" class="mt-4 space-y-2"></div>
|
||||
<div id="tenants-decision" class="mt-3 text-[11px] font-mono text-slate-400"></div>
|
||||
</div>
|
||||
|
||||
<!-- 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">
|
||||
|
||||
@@ -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)
|
||||
|
||||
108
tenants.json
Normal file
108
tenants.json
Normal file
@@ -0,0 +1,108 @@
|
||||
[
|
||||
{
|
||||
"name": "ollama",
|
||||
"kind": "llm",
|
||||
"priority": 50,
|
||||
"match": {
|
||||
"names": [
|
||||
"ollama"
|
||||
],
|
||||
"cmdline": [
|
||||
"llama-server",
|
||||
"ollama"
|
||||
]
|
||||
},
|
||||
"busy": {
|
||||
"type": "http_count",
|
||||
"url": "http://localhost:11434/api/ps",
|
||||
"count_keys": [
|
||||
"models"
|
||||
]
|
||||
},
|
||||
"release": {
|
||||
"type": "http_post",
|
||||
"url": "http://localhost:11434/api/generate",
|
||||
"body": {
|
||||
"keep_alive": 0
|
||||
},
|
||||
"per_model": true,
|
||||
"timeout_s": 120.0
|
||||
},
|
||||
"notes": "Unloads per model. With OLLAMA_NUM_PARALLEL=1 the request queues behind any running generation and applies when it finishes."
|
||||
},
|
||||
{
|
||||
"name": "comfyui",
|
||||
"kind": "diffusion",
|
||||
"priority": 60,
|
||||
"match": {
|
||||
"cmdline": [
|
||||
"comfyui",
|
||||
"comfy"
|
||||
],
|
||||
"cmdline_endswith": [
|
||||
"main.py"
|
||||
]
|
||||
},
|
||||
"busy": {
|
||||
"type": "http_count",
|
||||
"url": "http://127.0.0.1:8188/queue",
|
||||
"count_keys": [
|
||||
"queue_running",
|
||||
"queue_pending"
|
||||
],
|
||||
"vram_floor_gb": 1.5,
|
||||
"stale_after_s": 90.0
|
||||
},
|
||||
"release": {
|
||||
"type": "http_post",
|
||||
"url": "http://127.0.0.1:8188/free",
|
||||
"body": {
|
||||
"unload_models": true,
|
||||
"free_memory": true
|
||||
},
|
||||
"timeout_s": 30.0
|
||||
},
|
||||
"notes": "Leaves dead jobs in queue_running; the queue flag is corroborated against its own VRAM before being believed."
|
||||
},
|
||||
{
|
||||
"name": "stt-relay",
|
||||
"kind": "other",
|
||||
"priority": 70,
|
||||
"match": {
|
||||
"cmdline": [
|
||||
"stt_relay.py"
|
||||
]
|
||||
},
|
||||
"busy": {
|
||||
"type": "vram",
|
||||
"vram_busy_gb": 1.0
|
||||
},
|
||||
"release": {
|
||||
"type": "none"
|
||||
},
|
||||
"notes": "Long-running speech relay. Holds ~0.8 GB permanently and exposes no release API, so its VRAM is headroom this service can never offer. Declared so it is named rather than lumped into 'unmanaged'."
|
||||
},
|
||||
{
|
||||
"name": "desktop",
|
||||
"kind": "desktop",
|
||||
"priority": 90,
|
||||
"match": {
|
||||
"names": [
|
||||
"gnome-shell",
|
||||
"xorg",
|
||||
"mutter",
|
||||
"kwin",
|
||||
"plasmashell",
|
||||
"gnome-remote-desktop",
|
||||
"sddm",
|
||||
"gdm",
|
||||
"picom",
|
||||
"weston"
|
||||
]
|
||||
},
|
||||
"release": {
|
||||
"type": "none"
|
||||
},
|
||||
"notes": "Compositor and display server. Small, permanent, never reclaimable."
|
||||
}
|
||||
]
|
||||
461
tenants.py
Normal file
461
tenants.py
Normal file
@@ -0,0 +1,461 @@
|
||||
"""GPU tenants: the applications competing for the card, described as data.
|
||||
|
||||
The point of this service is fast handoff of a single GPU between applications. It grew
|
||||
up around the two on this box, and their names ended up compiled into process matching,
|
||||
VRAM attribution, busy detection and release calls alike -- roughly 385 references. That
|
||||
makes it a script for Ollama and ComfyUI rather than a GPU arbitrator.
|
||||
|
||||
A tenant is described here instead:
|
||||
|
||||
* how to recognise its processes (match)
|
||||
* how to tell whether it is actually working (busy probe)
|
||||
* how to ask it to give VRAM back (release strategy)
|
||||
* how much it matters when two want the card (priority)
|
||||
|
||||
Ollama and ComfyUI ship as defaults so behaviour is unchanged, but nothing about the
|
||||
arbitration logic knows their names. A third application -- a training run, a speech
|
||||
model, another inference server -- is a config entry, not a code change. A tenant that
|
||||
cannot be released (no API to ask) is still worth declaring, because naming it turns
|
||||
"unmanaged VRAM" into "held by X, which cannot be reclaimed".
|
||||
"""
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import time
|
||||
from dataclasses import dataclass, field, asdict
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import httpx
|
||||
import psutil
|
||||
|
||||
logger = logging.getLogger("tenants")
|
||||
|
||||
_BASE = os.path.dirname(os.path.abspath(__file__))
|
||||
CONFIG_PATH = os.environ.get("HYPERSWAP_TENANTS", os.path.join(_BASE, "tenants.json"))
|
||||
|
||||
# Kinds are advisory: they drive presentation and sensible defaults, never control flow.
|
||||
KIND_LLM, KIND_DIFFUSION, KIND_DESKTOP, KIND_OTHER = "llm", "diffusion", "desktop", "other"
|
||||
|
||||
|
||||
@dataclass
|
||||
class ProcessMatch:
|
||||
"""How to recognise a tenant's processes among those NVML reports."""
|
||||
names: List[str] = field(default_factory=list) # matched against process name
|
||||
cmdline: List[str] = field(default_factory=list) # substrings of the full cmdline
|
||||
cmdline_endswith: List[str] = field(default_factory=list)
|
||||
|
||||
def matches(self, pname: str, cmdline: str) -> bool:
|
||||
pname, cmdline = pname.lower(), cmdline.lower()
|
||||
if any(n.lower() in pname for n in self.names):
|
||||
return True
|
||||
if any(c.lower() in cmdline for c in self.cmdline):
|
||||
return True
|
||||
return any(cmdline.rstrip().endswith(c.lower()) for c in self.cmdline_endswith)
|
||||
|
||||
|
||||
@dataclass
|
||||
class BusyProbe:
|
||||
"""How to tell whether a tenant is genuinely working.
|
||||
|
||||
`vram_floor_gb` exists because a queue flag can lie: ComfyUI leaves dead jobs in
|
||||
queue_running, and only its VRAM reveals that nothing is loaded. GPU utilisation is
|
||||
deliberately unavailable as a signal -- it is shared by every tenant, so it cannot
|
||||
attribute work to one of them.
|
||||
"""
|
||||
type: str = "none" # none | http_count | vram
|
||||
url: Optional[str] = None
|
||||
count_keys: List[str] = field(default_factory=list) # keys whose lists are summed
|
||||
vram_busy_gb: float = 0.0 # busy when its VRAM exceeds this
|
||||
vram_floor_gb: float = 0.0 # below this it holds no real work
|
||||
stale_after_s: float = 90.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class EventSource:
|
||||
"""A stream that tells us *when* to look, not what to think.
|
||||
|
||||
ComfyUI publishes a websocket, and the original listener parsed its message types to
|
||||
decide what was happening -- which meant understanding one application's schema. Any
|
||||
message is instead treated purely as a wake-up: re-run this tenant's busy probe now
|
||||
rather than waiting for the next poll. That gives sub-second reaction to any
|
||||
application with an event stream, with no knowledge of what it emits.
|
||||
"""
|
||||
type: str = "none" # none | websocket
|
||||
url: Optional[str] = None
|
||||
reconnect_backoff_s: float = 2.0
|
||||
max_backoff_s: float = 15.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class ReleaseStrategy:
|
||||
"""How to ask a tenant to give VRAM back."""
|
||||
type: str = "none" # none | http_post
|
||||
url: Optional[str] = None
|
||||
body: Dict[str, Any] = field(default_factory=dict)
|
||||
# Set when the call must name the loaded model (Ollama unloads per model).
|
||||
per_model: bool = False
|
||||
timeout_s: float = 120.0
|
||||
confirm: bool = True # wait for NVML to show the memory released
|
||||
|
||||
|
||||
@dataclass
|
||||
class GpuTenant:
|
||||
name: str
|
||||
kind: str = KIND_OTHER
|
||||
enabled: bool = True
|
||||
# Higher wins contention; a tenant yields to anything above it.
|
||||
priority: int = 50
|
||||
# How much free VRAM this application needs before it can work. Used to decide
|
||||
# whether a busy tenant is actually being starved, rather than merely busy.
|
||||
needs_vram_gb: float = 0.0
|
||||
# How long a reclaimable tenant may sit idle holding VRAM before it is asked for it
|
||||
# back. Iterating on a ComfyUI workflow should not pay a reload between every run,
|
||||
# so this is deliberately not immediate.
|
||||
idle_release_after_s: float = 30.0
|
||||
# GPU profile to apply while this tenant is the active workload. Clock and power
|
||||
# tuning is workload-specific -- diffusion is compute bound, LLM decode is bandwidth
|
||||
# bound -- and that was previously switched by application name in the arbitrator.
|
||||
overclock_profile: Optional[str] = None
|
||||
# VRAM that survives a release. ComfyUI keeps its CUDA context for as long as the
|
||||
# process lives, so purging it does not return everything it holds. Ignoring this
|
||||
# made plan_release over-promise: it reported that releasing ComfyUI would free
|
||||
# 0.37 GB against a 0.33 GB shortfall, the job was cleared to run, and the memory
|
||||
# never actually arrived.
|
||||
vram_floor_gb: float = 0.0
|
||||
match: ProcessMatch = field(default_factory=ProcessMatch)
|
||||
busy: BusyProbe = field(default_factory=BusyProbe)
|
||||
release: ReleaseStrategy = field(default_factory=ReleaseStrategy)
|
||||
events: EventSource = field(default_factory=EventSource)
|
||||
notes: str = ""
|
||||
|
||||
@property
|
||||
def reclaimable(self) -> bool:
|
||||
return self.release.type != "none"
|
||||
|
||||
def to_dict(self) -> Dict[str, Any]:
|
||||
return asdict(self)
|
||||
|
||||
|
||||
def _tenant_from_dict(d: Dict[str, Any]) -> GpuTenant:
|
||||
return GpuTenant(
|
||||
name=d["name"],
|
||||
kind=d.get("kind", KIND_OTHER),
|
||||
enabled=d.get("enabled", True),
|
||||
priority=int(d.get("priority", 50)),
|
||||
needs_vram_gb=float(d.get("needs_vram_gb", 0.0)),
|
||||
overclock_profile=d.get("overclock_profile"),
|
||||
vram_floor_gb=float(d.get("vram_floor_gb", 0.0)),
|
||||
idle_release_after_s=float(d.get("idle_release_after_s", 30.0)),
|
||||
match=ProcessMatch(**(d.get("match") or {})),
|
||||
busy=BusyProbe(**(d.get("busy") or {})),
|
||||
release=ReleaseStrategy(**(d.get("release") or {})),
|
||||
events=EventSource(**(d.get("events") or {})),
|
||||
notes=d.get("notes", ""),
|
||||
)
|
||||
|
||||
|
||||
# Defaults reproduce today's behaviour exactly; they are data, not special cases.
|
||||
DEFAULT_TENANTS: List[Dict[str, Any]] = [
|
||||
{
|
||||
"name": "ollama",
|
||||
"kind": KIND_LLM,
|
||||
# Lower than ComfyUI on purpose: an interactive diffusion job preempts the LLM,
|
||||
# whose weights stay in the page cache and reload in seconds. Getting this the
|
||||
# wrong way round silently disabled the service's central behaviour -- ComfyUI
|
||||
# could never reclaim from Ollama.
|
||||
"priority": 50,
|
||||
"needs_vram_gb": 4.0,
|
||||
"idle_release_after_s": 0.0,
|
||||
"overclock_profile": "ollama",
|
||||
"match": {"names": ["ollama"], "cmdline": ["llama-server", "ollama"]},
|
||||
"busy": {"type": "http_count", "url": "http://localhost:11434/api/ps",
|
||||
"count_keys": ["models"]},
|
||||
"release": {"type": "http_post", "url": "http://localhost:11434/api/generate",
|
||||
"body": {"keep_alive": 0}, "per_model": True, "timeout_s": 120.0},
|
||||
"notes": "Unloads per model. With OLLAMA_NUM_PARALLEL=1 the request queues "
|
||||
"behind any running generation and applies when it finishes.",
|
||||
},
|
||||
{
|
||||
"name": "comfyui",
|
||||
"kind": KIND_DIFFUSION,
|
||||
"priority": 60,
|
||||
"needs_vram_gb": 6.0,
|
||||
"idle_release_after_s": 30.0,
|
||||
"overclock_profile": "comfy",
|
||||
"vram_floor_gb": 0.45,
|
||||
"match": {"cmdline": ["comfyui", "comfy"], "cmdline_endswith": ["main.py"]},
|
||||
"busy": {"type": "http_count", "url": "http://127.0.0.1:8188/queue",
|
||||
"count_keys": ["queue_running", "queue_pending"],
|
||||
"vram_floor_gb": 1.5, "stale_after_s": 90.0},
|
||||
"release": {"type": "http_post", "url": "http://127.0.0.1:8188/free",
|
||||
"body": {"unload_models": True, "free_memory": True},
|
||||
"timeout_s": 30.0},
|
||||
"events": {"type": "websocket", "url": "ws://127.0.0.1:8188/ws?clientId=hyperswap"},
|
||||
"notes": "Leaves dead jobs in queue_running; the queue flag is corroborated "
|
||||
"against its own VRAM before being believed.",
|
||||
},
|
||||
{
|
||||
"name": "desktop",
|
||||
"kind": KIND_DESKTOP,
|
||||
"priority": 90,
|
||||
"match": {"names": ["gnome-shell", "xorg", "mutter", "kwin", "plasmashell",
|
||||
"gnome-remote-desktop", "sddm", "gdm", "picom", "weston"]},
|
||||
"release": {"type": "none"},
|
||||
"notes": "Compositor and display server. Small, permanent, never reclaimable.",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
_cache: Dict[str, Any] = {"ts": 0.0, "tenants": None, "mtime": None}
|
||||
CACHE_TTL_S = 10.0
|
||||
|
||||
|
||||
def load_tenants(force: bool = False) -> List[GpuTenant]:
|
||||
"""Load tenant definitions, writing the defaults out on first run."""
|
||||
now = time.time()
|
||||
try:
|
||||
mtime = os.path.getmtime(CONFIG_PATH) if os.path.exists(CONFIG_PATH) else None
|
||||
except OSError:
|
||||
mtime = None
|
||||
if (not force and _cache["tenants"] is not None
|
||||
and mtime == _cache["mtime"] and (now - _cache["ts"]) < CACHE_TTL_S):
|
||||
return _cache["tenants"]
|
||||
|
||||
raw: List[Dict[str, Any]]
|
||||
if os.path.exists(CONFIG_PATH):
|
||||
try:
|
||||
with open(CONFIG_PATH) as f:
|
||||
raw = json.load(f)
|
||||
except Exception as e:
|
||||
logger.error(f"could not read {CONFIG_PATH}, using defaults: {e}")
|
||||
raw = DEFAULT_TENANTS
|
||||
else:
|
||||
raw = DEFAULT_TENANTS
|
||||
try:
|
||||
with open(CONFIG_PATH, "w") as f:
|
||||
json.dump(DEFAULT_TENANTS, f, indent=2)
|
||||
logger.info(f"wrote default tenant definitions to {CONFIG_PATH}")
|
||||
except Exception as e:
|
||||
logger.warning(f"could not write {CONFIG_PATH}: {e}")
|
||||
|
||||
# Merge in any fields a shipped default has gained since the config was written.
|
||||
# Without this, adding a field silently disables the behaviour it controls for every
|
||||
# existing install -- needs_vram_gb defaulted to 0, which made starvation
|
||||
# undetectable for the two tenants that had been written out before it existed.
|
||||
defaults_by_name = {d["name"]: d for d in DEFAULT_TENANTS}
|
||||
tenants = []
|
||||
for d in raw:
|
||||
base = defaults_by_name.get(d.get("name"))
|
||||
if base:
|
||||
merged = {**base, **d}
|
||||
for key in ("match", "busy", "release", "events"):
|
||||
if isinstance(base.get(key), dict):
|
||||
merged[key] = {**base[key], **(d.get(key) or {})}
|
||||
d = merged
|
||||
try:
|
||||
tenants.append(_tenant_from_dict(d))
|
||||
except Exception as e:
|
||||
logger.error(f"skipping malformed tenant {d!r}: {e}")
|
||||
_cache.update({"ts": now, "tenants": tenants, "mtime": mtime})
|
||||
return tenants
|
||||
|
||||
|
||||
def get_tenant(name: str) -> Optional[GpuTenant]:
|
||||
return next((t for t in load_tenants() if t.name == name), None)
|
||||
|
||||
|
||||
def save_tenants(tenants: List[Dict[str, Any]]) -> bool:
|
||||
try:
|
||||
with open(CONFIG_PATH, "w") as f:
|
||||
json.dump(tenants, f, indent=2)
|
||||
_cache["tenants"] = None
|
||||
return True
|
||||
except Exception as e:
|
||||
logger.error(f"save_tenants failed: {e}")
|
||||
return False
|
||||
|
||||
|
||||
def classify_process(pname: str, cmdline: str) -> str:
|
||||
"""Return the owning tenant's name, or 'unmanaged'.
|
||||
|
||||
'unmanaged' is meaningful rather than a dumping ground: it is VRAM this service has
|
||||
no way to reclaim, and it is reported as such.
|
||||
"""
|
||||
for t in load_tenants():
|
||||
if t.enabled and t.match.matches(pname, cmdline):
|
||||
return t.name
|
||||
return "unmanaged"
|
||||
|
||||
|
||||
def classify_pid(pid: int) -> str:
|
||||
try:
|
||||
proc = psutil.Process(pid)
|
||||
return classify_process(proc.name(), " ".join(proc.cmdline()))
|
||||
except Exception:
|
||||
return "unmanaged"
|
||||
|
||||
|
||||
async def probe_busy(tenant: GpuTenant, vram_gb: float = 0.0,
|
||||
state: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
|
||||
"""Is this tenant actually working? Returns {busy, reason, stale}."""
|
||||
probe = tenant.busy
|
||||
if probe.type == "vram":
|
||||
busy = vram_gb > probe.vram_busy_gb
|
||||
return {"busy": busy, "reason": f"{vram_gb:.2f} GB held", "stale": False}
|
||||
if probe.type != "http_count" or not probe.url:
|
||||
return {"busy": False, "reason": "no busy probe configured", "stale": False}
|
||||
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=3.0) as c:
|
||||
r = await c.get(probe.url)
|
||||
if r.status_code != 200:
|
||||
return {"busy": False, "reason": f"probe HTTP {r.status_code}", "stale": False}
|
||||
data = r.json()
|
||||
count = sum(len(data.get(k) or []) for k in probe.count_keys)
|
||||
except Exception as e:
|
||||
return {"busy": False, "reason": f"probe failed: {str(e)[:60]}", "stale": False}
|
||||
|
||||
if count == 0:
|
||||
return {"busy": False, "reason": "queue empty", "stale": False}
|
||||
# A queue that claims work while the tenant holds no VRAM is not doing work.
|
||||
if probe.vram_floor_gb and vram_gb < probe.vram_floor_gb:
|
||||
return {"busy": True, "reason": f"{count} queued, holding {vram_gb:.2f} GB",
|
||||
"stale": None, "below_floor": True}
|
||||
return {"busy": True, "reason": f"{count} queued/running", "stale": False}
|
||||
|
||||
|
||||
async def release_vram(tenant: GpuTenant, models: Optional[List[str]] = None
|
||||
) -> Dict[str, Any]:
|
||||
"""Ask a tenant to give its VRAM back, however that tenant expects to be asked."""
|
||||
strategy = tenant.release
|
||||
if strategy.type == "none" or not strategy.url:
|
||||
return {"success": False, "tenant": tenant.name, "released": False,
|
||||
"reason": "this tenant exposes no way to release VRAM"}
|
||||
|
||||
t0 = time.perf_counter()
|
||||
payloads: List[Dict[str, Any]] = []
|
||||
if strategy.per_model:
|
||||
for m in (models or []):
|
||||
payloads.append({**strategy.body, "model": m})
|
||||
if not payloads:
|
||||
return {"success": True, "tenant": tenant.name, "released": False,
|
||||
"reason": "nothing loaded to release"}
|
||||
else:
|
||||
payloads.append(dict(strategy.body))
|
||||
|
||||
errors = []
|
||||
try:
|
||||
async with httpx.AsyncClient(timeout=strategy.timeout_s) as c:
|
||||
for body in payloads:
|
||||
try:
|
||||
await c.post(strategy.url, json=body)
|
||||
except Exception as e:
|
||||
errors.append(str(e)[:80])
|
||||
except Exception as e:
|
||||
errors.append(str(e)[:80])
|
||||
|
||||
return {
|
||||
"success": not errors,
|
||||
"tenant": tenant.name,
|
||||
"released": True,
|
||||
"requests": len(payloads),
|
||||
"duration_ms": round((time.perf_counter() - t0) * 1000, 2),
|
||||
"errors": errors or None,
|
||||
}
|
||||
|
||||
|
||||
def describe() -> List[Dict[str, Any]]:
|
||||
"""Tenant definitions for the API, with what each can and cannot do."""
|
||||
out = []
|
||||
for t in sorted(load_tenants(), key=lambda x: -x.priority):
|
||||
d = t.to_dict()
|
||||
d["reclaimable"] = t.reclaimable
|
||||
d["busy_probe"] = t.busy.type
|
||||
d["release_via"] = t.release.type
|
||||
out.append(d)
|
||||
return out
|
||||
|
||||
|
||||
# ---------------------------------------------------------------- arbitration
|
||||
|
||||
def plan_release(demanding: str, tenants_state: List[Dict[str, Any]],
|
||||
free_gb: float, needed_gb: float) -> Dict[str, Any]:
|
||||
"""Decide who should give up VRAM so a starved tenant can work.
|
||||
|
||||
Generic over any number of applications: candidates are every *reclaimable* tenant
|
||||
that is not itself busy and ranks below the demanding one, taken lowest priority
|
||||
first, until enough would be freed. The two-application version of this was a pair
|
||||
of hardcoded rules -- yield Ollama for ComfyUI, purge ComfyUI for Ollama -- which
|
||||
could not express a third participant at all.
|
||||
|
||||
Returns the plan rather than performing it, so the decision is testable and can be
|
||||
logged before anything is actually released.
|
||||
"""
|
||||
by_name = {s["name"]: s for s in tenants_state}
|
||||
demander = by_name.get(demanding)
|
||||
if not demander:
|
||||
return {"possible": False, "reason": f"unknown tenant '{demanding}'", "release": []}
|
||||
|
||||
# The demander keeps what it already holds; only the remainder must be found.
|
||||
shortfall = needed_gb - free_gb - demander.get("vram_gb", 0.0)
|
||||
if shortfall <= 0:
|
||||
return {"possible": True, "reason": "enough VRAM is already free",
|
||||
"release": [], "shortfall_gb": 0.0}
|
||||
|
||||
# Who may be asked for memory:
|
||||
#
|
||||
# * any idle reclaimable tenant, whatever its rank -- idle memory is not in use;
|
||||
# * a *busy* tenant that ranks strictly below the demander.
|
||||
#
|
||||
# That second clause is the point of the whole service and was nearly lost. Refusing
|
||||
# to touch anything busy looks safe and is not: a diffusion job measured here ran for
|
||||
# 46 s instead of 3 s, squeezed into 1.6 GB, because the LLM reloaded straight after
|
||||
# yielding and was then protected as "busy" while ComfyUI starved. Preempting a
|
||||
# lower-priority tenant is safe precisely because releasing is asynchronous -- an
|
||||
# Ollama unload queues behind its running request and applies when that finishes, so
|
||||
# nothing is killed mid-flight.
|
||||
#
|
||||
# Equal or higher priority is never interrupted, so peers cannot fight.
|
||||
demander_priority = demander.get("priority", 0)
|
||||
candidates = [
|
||||
s for s in tenants_state
|
||||
if s["name"] != demanding
|
||||
and s.get("reclaimable")
|
||||
and s.get("vram_gb", 0) > 0
|
||||
and (not s.get("busy") or s.get("priority", 0) < demander_priority)
|
||||
]
|
||||
# Idle tenants first, then lowest priority: never disturb working software while
|
||||
# something idle still has memory to give.
|
||||
candidates.sort(key=lambda s: (bool(s.get("busy")), s.get("priority", 0),
|
||||
-s.get("vram_gb", 0)))
|
||||
|
||||
plan, freed = [], 0.0
|
||||
for c in candidates:
|
||||
if freed >= shortfall:
|
||||
break
|
||||
# Only what the tenant can actually give back, not everything it holds.
|
||||
releasable = max(c.get("vram_gb", 0.0) - c.get("vram_floor_gb", 0.0), 0.0)
|
||||
if releasable <= 0:
|
||||
continue
|
||||
plan.append(c["name"])
|
||||
freed += releasable
|
||||
|
||||
blockers = [
|
||||
{"name": s["name"], "vram_gb": s.get("vram_gb", 0.0),
|
||||
"why": ("busy and ranks at or above the demander" if s.get("busy") else
|
||||
"declares no release mechanism" if not s.get("reclaimable") else
|
||||
"enough was freed without it")}
|
||||
for s in tenants_state
|
||||
if s["name"] != demanding and s.get("vram_gb", 0) > 0 and s["name"] not in plan
|
||||
]
|
||||
|
||||
return {
|
||||
"possible": freed >= shortfall,
|
||||
"shortfall_gb": round(shortfall, 2),
|
||||
"would_free_gb": round(freed, 2),
|
||||
"release": plan,
|
||||
"blockers": blockers,
|
||||
"reason": (f"releasing {', '.join(plan)} frees {freed:.2f} GB of the "
|
||||
f"{shortfall:.2f} GB shortfall" if plan else
|
||||
"no reclaimable idle tenant holds enough VRAM"),
|
||||
}
|
||||
74
tests/README.md
Normal file
74
tests/README.md
Normal file
@@ -0,0 +1,74 @@
|
||||
# HyperSwap test suite
|
||||
|
||||
Fast, hermetic unit tests. No GPU is touched, no network call is made, no systemd unit is
|
||||
poked, and the production `hyperswap.db` is never opened.
|
||||
|
||||
## Running
|
||||
|
||||
```bash
|
||||
/home/drjones/comfy-mcp-venv/bin/python -m pytest tests/ -q
|
||||
```
|
||||
|
||||
Single file / single test:
|
||||
|
||||
```bash
|
||||
/home/drjones/comfy-mcp-venv/bin/python -m pytest tests/test_classify_load.py -q
|
||||
/home/drjones/comfy-mcp-venv/bin/python -m pytest tests/ -q -k warm_confident
|
||||
```
|
||||
|
||||
Whole suite runs in about 3 seconds.
|
||||
|
||||
## Safety rails
|
||||
|
||||
These matter, because this repo drives a live 4080 SUPER that a running service is using.
|
||||
|
||||
* `tests/conftest.py` installs an **autouse** `no_gpu_mutation` fixture that replaces
|
||||
`overclock_manager._sh` (the single choke point for every `nvidia-smi` /
|
||||
`nvidia-settings` write) plus `apply_profile`, `apply_fan_control`, `set_fan_speed`,
|
||||
`set_fan_auto` and `restore_safe` with recording stubs. Even a test that accidentally
|
||||
reaches an actuation path can only reach the stub. The fixture yields a dict of
|
||||
recorded calls, which the thermal tests assert against.
|
||||
* `HYPERSWAP_DB` is set to a non-existent path before `telemetry_store` is imported, so no
|
||||
import can bind `DB_PATH` to the production database. Tests that need a DB use the
|
||||
`temp_db` fixture, which monkeypatches `telemetry_store.DB_PATH` to a `tmp_path` file
|
||||
and stops the writer thread afterwards.
|
||||
* All file IO happens against files the tests create in `tmp_path`. No real model blob is
|
||||
read and `warm_file_to_ram` is never called.
|
||||
* Nothing sweeps, and nothing sends HTTP to Ollama, ComfyUI or :9090.
|
||||
|
||||
## Measured constants pinned here
|
||||
|
||||
These numbers came from measurement on this box, not from taste. If a change makes one of
|
||||
these tests fail, the constant is probably wrong, not the test.
|
||||
|
||||
| Constant | Value | Where pinned |
|
||||
| --- | --- | --- |
|
||||
| Cold load of a 12.87 GB model, 3.1% resident | 34267 ms → 0.38 GB/s | `test_classify_load.py::test_measured_cold_load_classifies_as_cold_disk` |
|
||||
| Warm load of the same model, 100% resident | 4901 ms → 2.63 GB/s | `test_classify_load.py::test_measured_warm_load_classifies_as_ram_hit` |
|
||||
| `RAM_HIT_GBPS = 2.0` must stay below the fastest achievable warm load (2.63 GB/s) | — | `test_classify_load.py::test_ram_hit_threshold_is_physically_achievable` |
|
||||
| `PARTIAL_HIT_GBPS = 0.8` must stay above the measured cold rate (0.38 GB/s) | — | same test |
|
||||
| Size-unknown fallback splits at 8000 ms (between 4.9 s warm and 34.3 s cold) | — | `test_classify_load.py::test_unknown_size_guess_boundary_is_8s` |
|
||||
| `WARM_SKIP_THRESHOLD_PCT = 90.0` | — | `test_ram_optimizer.py::test_warm_skip_threshold_constant_unchanged` |
|
||||
| A *probe* reading may only be trusted at exactly 100% (a 12-window probe once cleared 90% on a mostly-cold 12.87 GB blob that then loaded at 2.44 GB/s) | — | `test_ram_optimizer.py::test_probe_reading_is_only_trusted_at_exactly_100_percent` |
|
||||
| `PROBE_CACHED_GBPS = 1.5` sits in the gap between cold NVMe (0.35–0.5 GB/s) and page cache (3.2–13 GB/s) | — | `test_ram_optimizer.py::test_probe_cached_threshold_sits_between_measured_disk_and_cache_rates` |
|
||||
| Card power envelope: 320 W stock, 370 W max, sweeps never go below 60% of max | — | `test_autotune_helpers.py::test_supported_power_limits_parses_min_default_max` |
|
||||
| `_supported_clocks` must always query the `mem,gr` pair (a single-field query returns one column and silently yielded `[]`) | — | `test_autotune_helpers.py::test_supported_clocks_always_queries_the_mem_gr_pair` |
|
||||
| ComfyUI benchmark seed must vary per call (a fixed seed made ComfyUI serve a cached result in ~1 ms) | — | `test_autotune_helpers.py::test_comfy_workflow_seed_varies_between_calls` |
|
||||
| Governor hysteresis: `HOT_SAMPLES = 5`, `COOL_SAMPLES = 30`, `REAPPLY_COOLDOWN_S = 20` | — | `test_thermal_governor.py` (escalation, recovery, cooldown, alternating-sample tests) |
|
||||
| Model usage score: frequency decayed with a ~24 h half-life | — | `test_telemetry_store.py::test_model_usage_ranking_scores_recent_use_higher` |
|
||||
|
||||
## What is deliberately not covered
|
||||
|
||||
* `vram_arbitrator.instant_free_ollama_vram`, the `AutoArbitrator` yield/purge paths and
|
||||
the SSE broker — under active edit, contract changing.
|
||||
* `overclock_manager.apply_profile` and every other actuation path, `autotune.sweep`,
|
||||
`ram_optimizer.warm_file_to_ram` — these mutate hardware or do heavy IO.
|
||||
* `server.py` HTTP routes and `mcp_server.py` — would need the app wired to live
|
||||
subsystems.
|
||||
|
||||
## Known rough edge the tests work around
|
||||
|
||||
`telemetry_store.stop()` flushes the writer's pending *batch* but does not drain the
|
||||
submission queue, so a `stop()` racing a just-submitted row can drop it. The writer tests
|
||||
call a local `_drain()` helper to wait for the queue to empty before stopping, rather than
|
||||
encoding the race into an assertion.
|
||||
90
tests/conftest.py
Normal file
90
tests/conftest.py
Normal file
@@ -0,0 +1,90 @@
|
||||
"""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
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def isolated_tenant_registry(tmp_path, monkeypatch):
|
||||
"""Never let tests read the operator's live tenants.json.
|
||||
|
||||
Classification is now configuration, which means a test that reads the real config
|
||||
changes result when someone adds an application to their own machine -- exactly what
|
||||
happened when stt-relay was registered and a "third party is unmanaged" test started
|
||||
seeing it as a named tenant. Every test gets the shipped defaults unless it opts out
|
||||
by pointing CONFIG_PATH somewhere itself.
|
||||
"""
|
||||
import json as _json
|
||||
import tenants as _tenants
|
||||
|
||||
path = tmp_path / "tenants-default.json"
|
||||
path.write_text(_json.dumps(_tenants.DEFAULT_TENANTS))
|
||||
monkeypatch.setattr(_tenants, "CONFIG_PATH", str(path))
|
||||
_tenants._cache.update({"ts": 0.0, "tenants": None, "mtime": None})
|
||||
yield
|
||||
_tenants._cache.update({"ts": 0.0, "tenants": None, "mtime": None})
|
||||
184
tests/test_autotune_helpers.py
Normal file
184
tests/test_autotune_helpers.py
Normal file
@@ -0,0 +1,184 @@
|
||||
"""autotune's pure/parsable helpers.
|
||||
|
||||
Only the parsing and sampling helpers are exercised. Nothing here runs a sweep, applies a
|
||||
profile, or talks to nvidia-smi: subprocess.run is monkeypatched at the module boundary.
|
||||
"""
|
||||
import pytest
|
||||
|
||||
import autotune
|
||||
|
||||
|
||||
class _FakeProc:
|
||||
def __init__(self, stdout="", returncode=0, stderr=""):
|
||||
self.stdout = stdout
|
||||
self.stderr = stderr
|
||||
self.returncode = returncode
|
||||
|
||||
|
||||
def _fake_smi(monkeypatch, stdout, record=None):
|
||||
def _run(cmd, capture_output=True, text=True, timeout=None, **kw):
|
||||
if record is not None:
|
||||
record.append(list(cmd))
|
||||
return _FakeProc(stdout)
|
||||
monkeypatch.setattr(autotune.subprocess, "run", _run)
|
||||
|
||||
|
||||
# A trimmed but realistically shaped `nvidia-smi --query-supported-clocks=mem,gr` dump:
|
||||
# graphics clocks are enumerated once per memory clock.
|
||||
SUPPORTED_CLOCKS_CSV = """\
|
||||
10501, 2790
|
||||
10501, 2775
|
||||
10501, 2760
|
||||
9501, 2790
|
||||
9501, 2775
|
||||
405, 645
|
||||
405, 630
|
||||
"""
|
||||
|
||||
|
||||
def test_supported_clocks_mem_returns_sorted_unique_memory_clocks(monkeypatch):
|
||||
_fake_smi(monkeypatch, SUPPORTED_CLOCKS_CSV)
|
||||
assert autotune._supported_clocks("mem") == [405, 9501, 10501]
|
||||
|
||||
|
||||
def test_supported_clocks_gr_returns_clocks_of_the_highest_memory_clock(monkeypatch):
|
||||
"""Graphics clocks are enumerated per memory clock. Only the set belonging to the top
|
||||
memory clock is meaningful — that is the state any real workload runs at."""
|
||||
_fake_smi(monkeypatch, SUPPORTED_CLOCKS_CSV)
|
||||
assert autotune._supported_clocks("gr") == [2760, 2775, 2790]
|
||||
|
||||
|
||||
def test_supported_clocks_always_queries_the_mem_gr_pair(monkeypatch):
|
||||
"""Regression: querying a single field returns one column, and reading index 1 from it
|
||||
silently produced an empty list instead of an error. The query must stay a pair."""
|
||||
seen = []
|
||||
_fake_smi(monkeypatch, SUPPORTED_CLOCKS_CSV, record=seen)
|
||||
autotune._supported_clocks("gr")
|
||||
assert any("--query-supported-clocks=mem,gr" in part for part in seen[0])
|
||||
|
||||
|
||||
def test_supported_clocks_ignores_single_column_output(monkeypatch):
|
||||
"""If the driver ever returns one column, no row parses and we return [] — never a
|
||||
list of memory clocks mislabelled as graphics clocks."""
|
||||
_fake_smi(monkeypatch, "10501\n9501\n405\n")
|
||||
assert autotune._supported_clocks("mem") == []
|
||||
assert autotune._supported_clocks("gr") == []
|
||||
|
||||
|
||||
def test_supported_clocks_skips_header_and_garbage_rows(monkeypatch):
|
||||
_fake_smi(monkeypatch, "memory [MHz], graphics [MHz]\n10501, 2790\n\nN/A, N/A\n")
|
||||
assert autotune._supported_clocks("mem") == [10501]
|
||||
|
||||
|
||||
def test_supported_clocks_returns_empty_when_nvidia_smi_fails(monkeypatch):
|
||||
"""No GPU / no driver must degrade to 'no candidates', not an exception on import of
|
||||
a sweep plan."""
|
||||
def _boom(*a, **kw):
|
||||
raise OSError("nvidia-smi not found")
|
||||
monkeypatch.setattr(autotune.subprocess, "run", _boom)
|
||||
assert autotune._supported_clocks("mem") == []
|
||||
|
||||
|
||||
# --------------------------------------------------------------- power limits
|
||||
|
||||
def test_supported_power_limits_parses_min_default_max(monkeypatch):
|
||||
"""RTX 4080 SUPER on this box: 115 W min, 370 W max, 320 W stock default."""
|
||||
_fake_smi(monkeypatch, "115.00, 370.00, 320.00\n")
|
||||
vals = autotune._supported_power_limits(steps=5)
|
||||
assert vals == sorted(set(vals))
|
||||
assert max(vals) == 370
|
||||
# Never sweeps below 60% of max — the card does no useful work down there.
|
||||
assert min(vals) >= int(370 * 0.6)
|
||||
# The stock default is always included as a reference point.
|
||||
assert 320 in vals
|
||||
|
||||
|
||||
def test_supported_power_limits_respects_step_count(monkeypatch):
|
||||
_fake_smi(monkeypatch, "115.00, 370.00, 320.00\n")
|
||||
vals = autotune._supported_power_limits(steps=3)
|
||||
assert len(vals) <= 4 # 3 evenly spaced values, plus the default if it is not one
|
||||
|
||||
|
||||
def test_supported_power_limits_returns_empty_on_query_failure(monkeypatch):
|
||||
_fake_smi(monkeypatch, "N/A, N/A, N/A\n")
|
||||
assert autotune._supported_power_limits() == []
|
||||
|
||||
|
||||
# --------------------------------------------------------------- subsampling
|
||||
|
||||
def test_subsample_returns_input_when_already_short_enough():
|
||||
assert autotune._subsample([1, 2, 3], 5) == [1, 2, 3]
|
||||
|
||||
|
||||
def test_subsample_keeps_both_endpoints():
|
||||
"""The endpoints are the whole point of a sweep: stock and maximum. Dropping either
|
||||
means never measuring the setting the sweep is supposed to recommend."""
|
||||
values = list(range(0, 195)) # the card enumerates ~194 graphics clocks
|
||||
out = autotune._subsample(values, 8)
|
||||
assert out[0] == values[0]
|
||||
assert out[-1] == values[-1]
|
||||
|
||||
|
||||
def test_subsample_never_exceeds_max_steps():
|
||||
values = list(range(0, 195))
|
||||
for max_steps in (2, 3, 5, 8, 13):
|
||||
assert len(autotune._subsample(values, max_steps)) <= max_steps
|
||||
|
||||
|
||||
def test_subsample_output_is_sorted_and_unique():
|
||||
values = list(range(0, 100))
|
||||
out = autotune._subsample(values, 7)
|
||||
assert out == sorted(set(out))
|
||||
|
||||
|
||||
def test_subsample_is_evenly_spread():
|
||||
"""Clustered samples would tell us nothing about the middle of the range."""
|
||||
out = autotune._subsample(list(range(0, 101)), 5)
|
||||
gaps = [b - a for a, b in zip(out, out[1:])]
|
||||
assert max(gaps) - min(gaps) <= 1
|
||||
|
||||
|
||||
# --------------------------------------------------------------- comfy benchmark graph
|
||||
|
||||
def test_comfy_workflow_seed_varies_between_calls():
|
||||
"""Fixed seeds made ComfyUI serve a cached result in ~1 ms without executing the
|
||||
graph, so every benchmark after the first measured nothing at all."""
|
||||
seeds = {autotune._comfy_workflow()["5"]["inputs"]["seed"] for _ in range(25)}
|
||||
assert len(seeds) > 20
|
||||
|
||||
|
||||
def test_comfy_workflow_seed_can_be_pinned_for_reproduction():
|
||||
assert autotune._comfy_workflow(seed=42)["5"]["inputs"]["seed"] == 42
|
||||
|
||||
|
||||
def test_comfy_workflow_shape_is_otherwise_constant():
|
||||
"""Only the seed may vary; a benchmark whose step count or resolution moved between
|
||||
runs would not be comparable."""
|
||||
a = autotune._comfy_workflow()
|
||||
b = autotune._comfy_workflow()
|
||||
for wf in (a, b):
|
||||
assert wf["5"]["inputs"]["steps"] == autotune.COMFY_BENCH_STEPS
|
||||
assert wf["4"]["inputs"]["width"] == autotune.COMFY_BENCH_SIZE
|
||||
assert wf["4"]["inputs"]["height"] == autotune.COMFY_BENCH_SIZE
|
||||
a["5"]["inputs"]["seed"] = b["5"]["inputs"]["seed"] = 0
|
||||
assert a == b
|
||||
|
||||
|
||||
def test_comfy_workflow_uses_the_named_checkpoint():
|
||||
wf = autotune._comfy_workflow(ckpt="some_other.safetensors")
|
||||
assert wf["1"]["inputs"]["ckpt_name"] == "some_other.safetensors"
|
||||
|
||||
|
||||
def test_temp_ceiling_is_below_the_thermal_governors_escalation_point():
|
||||
"""A sweep step must abort on temperature before the governor starts derating under
|
||||
it, otherwise the sweep measures the governor's derate rather than the knob."""
|
||||
import thermal_governor as tg
|
||||
assert autotune.TEMP_CEILING_C <= tg.TEMP_ESCALATE_C + 1.0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("knob", ["mem_offset_mhz", "core_offset_mhz", "lock_mem_mhz",
|
||||
"lock_core_max", "power_limit_w"])
|
||||
def test_every_knob_declares_a_hardware_verification_field(knob):
|
||||
"""Offsets are silently ignored by some drivers (595.84 accepts an assignment and
|
||||
reads back a different value), so each knob must name the field to read back."""
|
||||
assert autotune.KNOBS[knob]["verify"]
|
||||
123
tests/test_classify_load.py
Normal file
123
tests/test_classify_load.py
Normal file
@@ -0,0 +1,123 @@
|
||||
"""Regression tests for vram_arbitrator.classify_load.
|
||||
|
||||
classify_load turns (model size, Ollama's reported load_duration) into a cache verdict
|
||||
by computing achieved bandwidth. The two thresholds it compares against were measured on
|
||||
this box, not chosen; the tests below pin the measurements themselves so a future edit
|
||||
to RAM_HIT_GBPS / PARTIAL_HIT_GBPS that breaks the real data points fails loudly.
|
||||
"""
|
||||
import math
|
||||
|
||||
import pytest
|
||||
|
||||
import vram_arbitrator as va
|
||||
|
||||
GIB = 1024 ** 3
|
||||
|
||||
# Ground truth, measured 2026-08-28 on the same 12.87 GB model loaded twice:
|
||||
# 3.1% resident -> 34267 ms -> 0.38 GB/s -> cold
|
||||
# 100% resident -> 4901 ms -> 2.63 GB/s -> RAM hit
|
||||
MEASURED_MODEL_BYTES = int(12.87 * GIB)
|
||||
MEASURED_COLD_MS = 34267.0
|
||||
MEASURED_WARM_MS = 4901.0
|
||||
|
||||
|
||||
def test_measured_cold_load_classifies_as_cold_disk():
|
||||
"""The measured cold load (12.87 GB, 34267 ms, 0.38 GB/s) must stay 'Cold Disk Load'."""
|
||||
res = va.classify_load(MEASURED_MODEL_BYTES, MEASURED_COLD_MS)
|
||||
assert res["cache_status"] == "Cold Disk Load 💾"
|
||||
assert res["is_ram_hit"] is False
|
||||
assert res["load_gbps"] == pytest.approx(0.38, abs=0.01)
|
||||
|
||||
|
||||
def test_measured_warm_load_classifies_as_ram_hit():
|
||||
"""The measured warm load (12.87 GB, 4901 ms, 2.63 GB/s) must stay a RAM cache hit."""
|
||||
res = va.classify_load(MEASURED_MODEL_BYTES, MEASURED_WARM_MS)
|
||||
assert res["cache_status"] == "RAM Cache Hit ⚡"
|
||||
assert res["is_ram_hit"] is True
|
||||
assert res["load_gbps"] == pytest.approx(2.63, abs=0.01)
|
||||
|
||||
|
||||
def test_ram_hit_threshold_is_physically_achievable():
|
||||
"""Guard against the reintroduced bug where the hit bar sat above any real warm load.
|
||||
|
||||
A prior version set the hit threshold at 5.0 GB/s. Ollama's load_duration covers
|
||||
host-to-device transfer and model init as well as the file read, so a fully resident
|
||||
12.87 GB model still only reports ~2.6 GB/s — meaning *no* load could ever be
|
||||
classified as a hit. Any threshold that a genuinely warm load cannot clear is wrong.
|
||||
"""
|
||||
assert va.RAM_HIT_GBPS <= 2.63, (
|
||||
"RAM_HIT_GBPS is above the fastest warm load ever measured on this box "
|
||||
"(2.63 GB/s); no load could be classified as a cache hit."
|
||||
)
|
||||
assert va.PARTIAL_HIT_GBPS > 0.38, (
|
||||
"PARTIAL_HIT_GBPS is at or below the measured cold-disk rate; cold loads would "
|
||||
"be reported as partial cache hits."
|
||||
)
|
||||
assert va.PARTIAL_HIT_GBPS < va.RAM_HIT_GBPS
|
||||
|
||||
|
||||
def test_sub_millisecond_load_is_already_in_vram():
|
||||
"""load_duration_ms <= 1.0 means Ollama never re-read the model at all."""
|
||||
res = va.classify_load(MEASURED_MODEL_BYTES, 1.0)
|
||||
assert res["cache_status"] == "Already in VRAM"
|
||||
assert res["load_gbps"] is None
|
||||
assert res["is_ram_hit"] is True
|
||||
|
||||
|
||||
def test_zero_duration_is_already_in_vram():
|
||||
res = va.classify_load(MEASURED_MODEL_BYTES, 0.0)
|
||||
assert res["cache_status"] == "Already in VRAM"
|
||||
|
||||
|
||||
def test_boundary_exactly_at_ram_hit_threshold_is_a_hit():
|
||||
"""Exactly RAM_HIT_GBPS (2.0 GB/s) must count as a hit — the comparison is >=."""
|
||||
size = int(va.RAM_HIT_GBPS * GIB) # 2 GiB read in exactly 1000 ms -> 2.00 GB/s
|
||||
res = va.classify_load(size, 1000.0)
|
||||
assert res["load_gbps"] == pytest.approx(2.0)
|
||||
assert res["cache_status"] == "RAM Cache Hit ⚡"
|
||||
assert res["is_ram_hit"] is True
|
||||
|
||||
|
||||
def test_just_below_ram_hit_threshold_is_partial():
|
||||
size = int(1.99 * GIB)
|
||||
res = va.classify_load(size, 1000.0)
|
||||
assert res["cache_status"] == "Partial Cache 🌤"
|
||||
assert res["is_ram_hit"] is False
|
||||
|
||||
|
||||
def test_boundary_exactly_at_partial_threshold_is_partial():
|
||||
"""Exactly PARTIAL_HIT_GBPS (0.8 GB/s) is still a partial hit, not a cold load."""
|
||||
size = math.ceil(va.PARTIAL_HIT_GBPS * GIB) # 0.8 GiB is not a whole number of bytes
|
||||
res = va.classify_load(size, 1000.0)
|
||||
assert res["load_gbps"] == pytest.approx(0.8)
|
||||
assert res["cache_status"] == "Partial Cache 🌤"
|
||||
assert res["is_ram_hit"] is False
|
||||
|
||||
|
||||
def test_just_below_partial_threshold_is_cold():
|
||||
size = int(0.79 * GIB)
|
||||
res = va.classify_load(size, 1000.0)
|
||||
assert res["cache_status"] == "Cold Disk Load 💾"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("duration_ms,expected_hit", [(4901.0, True), (34267.0, False)])
|
||||
def test_unknown_size_falls_back_to_duration_guess_and_says_so(duration_ms, expected_hit):
|
||||
"""With no size on record there is no bandwidth to compute, so the result must be
|
||||
labelled as a guess rather than presented as a measurement."""
|
||||
res = va.classify_load(0, duration_ms)
|
||||
assert res["is_ram_hit"] is expected_hit
|
||||
assert res["load_gbps"] is None
|
||||
assert "detail" in res and "guess" in res["detail"].lower()
|
||||
|
||||
|
||||
def test_unknown_size_guess_boundary_is_8s():
|
||||
"""The size-unknown fallback splits at 8000 ms, between the measured 4.9 s warm and
|
||||
34.3 s cold loads."""
|
||||
assert va.classify_load(0, 7999.0)["is_ram_hit"] is True
|
||||
assert va.classify_load(0, 8000.0)["is_ram_hit"] is False
|
||||
|
||||
|
||||
def test_classification_never_raises_on_odd_inputs():
|
||||
"""This runs on the swap hot path; it must not be able to throw."""
|
||||
for size, dur in [(0, 0.0), (1, 1.5), (10 ** 12, 2.0), (0, 1.0)]:
|
||||
assert "cache_status" in va.classify_load(size, dur)
|
||||
133
tests/test_engines.py
Normal file
133
tests/test_engines.py
Normal file
@@ -0,0 +1,133 @@
|
||||
"""Tests for live engine-configuration reporting.
|
||||
|
||||
These settings live outside this codebase but dictate how arbitration must behave, and
|
||||
working out why a yield behaved a certain way once meant reading journald by hand. The
|
||||
dashboard previously asserted them as hardcoded text, which happened to be accurate --
|
||||
worse than being wrong, because it would have stayed accurate-looking after the settings
|
||||
changed.
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
import engines
|
||||
|
||||
|
||||
class _Proc:
|
||||
def __init__(self, stdout=""):
|
||||
self.stdout = stdout
|
||||
self.returncode = 0
|
||||
|
||||
|
||||
class TestOllamaEnvironmentParsing:
|
||||
def test_parses_the_real_unit_environment(self, monkeypatch):
|
||||
# Verbatim from `systemctl show ollama -p Environment --value` on this machine.
|
||||
raw = ("OLLAMA_HOST=0.0.0.0:11434 OLLAMA_FLASH_ATTENTION=1 "
|
||||
"OLLAMA_KV_CACHE_TYPE=q4_0 OLLAMA_KEEP_ALIVE=30m "
|
||||
"OLLAMA_MAX_LOADED_MODELS=1 OLLAMA_NUM_PARALLEL=1 OLLAMA_NUM_BATCH=2048")
|
||||
monkeypatch.setattr(engines.subprocess, "run", lambda *a, **k: _Proc(raw))
|
||||
env = engines._ollama_unit_environment()
|
||||
assert env["OLLAMA_NUM_PARALLEL"] == "1"
|
||||
assert env["OLLAMA_MAX_LOADED_MODELS"] == "1"
|
||||
assert env["OLLAMA_KV_CACHE_TYPE"] == "q4_0"
|
||||
|
||||
def test_ignores_non_ollama_variables(self, monkeypatch):
|
||||
monkeypatch.setattr(engines.subprocess, "run",
|
||||
lambda *a, **k: _Proc("PATH=/usr/bin OLLAMA_HOST=x:1 HOME=/root"))
|
||||
env = engines._ollama_unit_environment()
|
||||
assert set(env) == {"OLLAMA_HOST"}
|
||||
|
||||
def test_returns_empty_rather_than_raising_when_systemctl_fails(self, monkeypatch):
|
||||
def boom(*a, **k):
|
||||
raise FileNotFoundError("systemctl")
|
||||
monkeypatch.setattr(engines.subprocess, "run", boom)
|
||||
assert engines._ollama_unit_environment() == {}
|
||||
|
||||
|
||||
class TestEngineConfigReport:
|
||||
def _run(self, monkeypatch, env, comfy_ok=True):
|
||||
monkeypatch.setattr(engines, "_ollama_unit_environment", lambda: env)
|
||||
|
||||
class _Resp:
|
||||
status_code = 200 if comfy_ok else 500
|
||||
def json(self):
|
||||
return {"system": {"comfyui_version": "0.33.1",
|
||||
"pytorch_version": "2.11.0+cu128",
|
||||
"python_version": "3.14.4 (main)",
|
||||
"argv": ["main.py", "--listen", "0.0.0.0"]},
|
||||
"devices": [{"name": "cuda:0 NVIDIA GeForce RTX 4080 SUPER "
|
||||
": cudaMallocAsync"}]}
|
||||
|
||||
class _Client:
|
||||
async def __aenter__(self): return self
|
||||
async def __aexit__(self, *a): return False
|
||||
async def get(self, url): return _Resp()
|
||||
|
||||
monkeypatch.setattr(engines.httpx, "AsyncClient", lambda **k: _Client())
|
||||
return asyncio.run(engines.get_engine_config())
|
||||
|
||||
def test_surfaces_the_settings_that_drive_arbitration(self, monkeypatch):
|
||||
d = self._run(monkeypatch, {"OLLAMA_NUM_PARALLEL": "1",
|
||||
"OLLAMA_MAX_LOADED_MODELS": "1",
|
||||
"OLLAMA_KEEP_ALIVE": "30m"})
|
||||
assert d["ollama"]["num_parallel"] == "1"
|
||||
assert d["ollama"]["max_loaded_models"] == "1"
|
||||
assert d["ollama"]["keep_alive"] == "30m"
|
||||
|
||||
def test_num_parallel_explains_the_deferred_yield_behaviour(self, monkeypatch):
|
||||
d = self._run(monkeypatch, {"OLLAMA_NUM_PARALLEL": "1"})
|
||||
note = next(s["means"] for s in d["ollama"]["settings"]
|
||||
if s["key"] == "OLLAMA_NUM_PARALLEL")
|
||||
# The explanation is the point: it is why a busy model is deferred, not failed.
|
||||
assert "queue" in note.lower()
|
||||
|
||||
def test_summary_reflects_actual_flags_not_a_fixed_string(self, monkeypatch):
|
||||
on = self._run(monkeypatch, {"OLLAMA_FLASH_ATTENTION": "1",
|
||||
"OLLAMA_KV_CACHE_TYPE": "q4_0"})
|
||||
assert "FlashAttention" in on["ollama"]["summary"]
|
||||
assert "q4_0" in on["ollama"]["summary"]
|
||||
off = self._run(monkeypatch, {})
|
||||
assert "FlashAttention" not in off["ollama"]["summary"]
|
||||
assert off["ollama"]["config_source"] == "unavailable"
|
||||
|
||||
def test_port_comes_from_ollama_host(self, monkeypatch):
|
||||
d = self._run(monkeypatch, {"OLLAMA_HOST": "0.0.0.0:11500"})
|
||||
assert d["ollama"]["port"] == "11500"
|
||||
|
||||
def test_comfy_allocator_and_vram_mode_are_read_not_asserted(self, monkeypatch):
|
||||
d = self._run(monkeypatch, {})
|
||||
assert d["comfyui"]["allocator"] == "cudaMallocAsync"
|
||||
assert d["comfyui"]["vram_mode"] == "default (auto)"
|
||||
assert d["comfyui"]["version"] == "0.33.1"
|
||||
|
||||
def test_comfy_vram_flag_is_detected_when_present(self, monkeypatch):
|
||||
monkeypatch.setattr(engines, "_ollama_unit_environment", lambda: {})
|
||||
|
||||
class _Resp:
|
||||
status_code = 200
|
||||
def json(self):
|
||||
return {"system": {"argv": ["main.py", "--lowvram"]},
|
||||
"devices": [{"name": "cuda:0 X : cudaMalloc"}]}
|
||||
|
||||
class _Client:
|
||||
async def __aenter__(self): return self
|
||||
async def __aexit__(self, *a): return False
|
||||
async def get(self, url): return _Resp()
|
||||
|
||||
monkeypatch.setattr(engines.httpx, "AsyncClient", lambda **k: _Client())
|
||||
d = asyncio.run(engines.get_engine_config())
|
||||
assert d["comfyui"]["vram_mode"] == "--lowvram"
|
||||
assert d["comfyui"]["allocator"] == "cudaMalloc"
|
||||
|
||||
def test_offline_comfy_is_reported_not_raised(self, monkeypatch):
|
||||
monkeypatch.setattr(engines, "_ollama_unit_environment", lambda: {})
|
||||
|
||||
class _Client:
|
||||
async def __aenter__(self): return self
|
||||
async def __aexit__(self, *a): return False
|
||||
async def get(self, url): raise ConnectionError("refused")
|
||||
|
||||
monkeypatch.setattr(engines.httpx, "AsyncClient", lambda **k: _Client())
|
||||
d = asyncio.run(engines.get_engine_config())
|
||||
assert d["comfyui"]["online"] is False
|
||||
assert "error" in d["comfyui"]
|
||||
216
tests/test_health.py
Normal file
216
tests/test_health.py
Normal file
@@ -0,0 +1,216 @@
|
||||
"""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",
|
||||
"_check_unmanaged_vram", "_check_comfy_queue"):
|
||||
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", "_check_unmanaged_vram", "_check_comfy_queue"):
|
||||
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"])
|
||||
|
||||
|
||||
class TestUnmanagedVramCheck:
|
||||
"""Turning an unreclaimable-VRAM number into something actionable.
|
||||
|
||||
The arithmetic here has to be right or the check is worse than useless. A first
|
||||
version omitted ComfyUI's CUDA context -- which survives a purge -- and so reported
|
||||
a 14.93 GB model as fitting against a real ceiling of 14.60 GB. That was the very
|
||||
model the service had just refused with 507 Insufficient Storage.
|
||||
"""
|
||||
|
||||
def _gpu(self, unmanaged_gb=0.82, desktop_gb=0.01, comfy_gb=0.56, total=15.99,
|
||||
procs=None):
|
||||
return {
|
||||
"available": True,
|
||||
"vram_total_gb": total,
|
||||
"breakdown": {
|
||||
"unmanaged_gb": unmanaged_gb, "desktop_gb": desktop_gb,
|
||||
"comfyui_gb": comfy_gb,
|
||||
"unmanaged": procs if procs is not None else
|
||||
[{"pid": 1, "name": "python", "vram_mb": unmanaged_gb * 1024,
|
||||
"cmdline": "stt_relay.py"}],
|
||||
},
|
||||
}
|
||||
|
||||
def _blobs(self, sizes):
|
||||
return [{"model": f"m{i}", "size_gb": s} for i, s in enumerate(sizes)]
|
||||
|
||||
def test_ok_when_nothing_holds_unreclaimable_vram(self, monkeypatch):
|
||||
monkeypatch.setattr(health.vram_arbitrator, "get_gpu_hardware_stats",
|
||||
lambda: self._gpu(unmanaged_gb=0.0, procs=[]))
|
||||
assert health._check_unmanaged_vram()["status"] == health.OK
|
||||
|
||||
def test_comfy_cuda_context_counts_against_the_ceiling(self, monkeypatch):
|
||||
# 15.99 - 0.82 unmanaged - 0.01 desktop - 0.56 comfy floor = 14.60 GB available.
|
||||
# A 12.87 GB blob needs 12.87 * 1.16 = 14.93 GB, so it does not fit -- matching
|
||||
# the observed 507.
|
||||
monkeypatch.setattr(health.vram_arbitrator, "get_gpu_hardware_stats",
|
||||
lambda: self._gpu())
|
||||
monkeypatch.setattr(health, "_comfy_vram_floor_gb", lambda default=0, days=1: 0.56)
|
||||
monkeypatch.setattr(health.ram_optimizer, "find_ollama_model_files",
|
||||
lambda: self._blobs([12.87]))
|
||||
res = health._check_unmanaged_vram()
|
||||
assert res["status"] == health.DEGRADED
|
||||
assert "1 model(s)" in res["impact"]
|
||||
|
||||
def test_model_that_fits_even_without_the_unmanaged_process_is_not_flagged(self, monkeypatch):
|
||||
# A tiny model fits either way, so the unmanaged process is not what blocks it.
|
||||
monkeypatch.setattr(health.vram_arbitrator, "get_gpu_hardware_stats",
|
||||
lambda: self._gpu())
|
||||
monkeypatch.setattr(health, "_comfy_vram_floor_gb", lambda default=0, days=1: 0.56)
|
||||
monkeypatch.setattr(health.ram_optimizer, "find_ollama_model_files",
|
||||
lambda: self._blobs([2.0]))
|
||||
assert health._check_unmanaged_vram()["status"] == health.OK
|
||||
|
||||
def test_model_too_big_to_ever_fit_is_not_blamed_on_the_process(self, monkeypatch):
|
||||
# A 23.7 GB model does not fit on a 16 GB card regardless; saying the 842 MB
|
||||
# process is why would send the user after the wrong thing.
|
||||
monkeypatch.setattr(health.vram_arbitrator, "get_gpu_hardware_stats",
|
||||
lambda: self._gpu())
|
||||
monkeypatch.setattr(health, "_comfy_vram_floor_gb", lambda default=0, days=1: 0.56)
|
||||
monkeypatch.setattr(health.ram_optimizer, "find_ollama_model_files",
|
||||
lambda: self._blobs([23.7]))
|
||||
assert health._check_unmanaged_vram()["status"] == health.OK
|
||||
|
||||
def test_floor_uses_the_minimum_observed_not_the_current_value(self, monkeypatch):
|
||||
# Current VRAM could be a 7 GB checkpoint mid-generation; the floor is what
|
||||
# survives a purge.
|
||||
monkeypatch.setattr(health.telemetry_store, "_rows",
|
||||
lambda *a, **k: [{"floor": int(0.24 * 1024 ** 3)}])
|
||||
assert health._comfy_vram_floor_gb(default=7.0) == 0.24
|
||||
|
||||
def test_floor_falls_back_when_history_is_empty(self, monkeypatch):
|
||||
monkeypatch.setattr(health.telemetry_store, "_rows", lambda *a, **k: [])
|
||||
assert health._comfy_vram_floor_gb(default=0.56) == 0.56
|
||||
|
||||
def test_floor_falls_back_rather_than_raising(self, monkeypatch):
|
||||
def boom(*a, **k):
|
||||
raise RuntimeError("db gone")
|
||||
monkeypatch.setattr(health.telemetry_store, "_rows", boom)
|
||||
assert health._comfy_vram_floor_gb(default=0.5) == 0.5
|
||||
184
tests/test_jobs.py
Normal file
184
tests/test_jobs.py
Normal file
@@ -0,0 +1,184 @@
|
||||
"""Tests for the cross-tenant job queue.
|
||||
|
||||
Every case below corresponds to something that actually went wrong while building this,
|
||||
because the failure modes are not obvious from the code:
|
||||
|
||||
* dispatching without room does not fail gracefully -- a CUDA OOM kills llama-server,
|
||||
and three queued jobs were destroyed in a row;
|
||||
* the VRAM requirement is a property of the job's model, not of the tenant, and a flat
|
||||
4 GB let a 14.9 GB model be dispatched into 8 GB of free memory;
|
||||
* waiting forever is as wrong as failing immediately, when the job can never fit;
|
||||
* an exception during dispatch left the row RUNNING while the scheduler moved on.
|
||||
"""
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
import jobs as J
|
||||
import telemetry_store
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def queue(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(telemetry_store, "DB_PATH", str(tmp_path / "q.db"))
|
||||
J.init()
|
||||
return tmp_path
|
||||
|
||||
|
||||
class TestQueueBasics:
|
||||
def test_submit_and_read_back(self, queue):
|
||||
r = J.submit("ollama", {"model": "m"}, label="first")
|
||||
assert r["success"]
|
||||
job = J.get(r["id"])
|
||||
assert job["state"] == J.PENDING and job["label"] == "first"
|
||||
assert job["payload"] == {"model": "m"}
|
||||
|
||||
def test_unknown_tenant_is_rejected(self, queue):
|
||||
assert J.submit("nope", {})["success"] is False
|
||||
|
||||
def test_priority_defaults_to_the_tenants_own(self, queue):
|
||||
import tenants as T
|
||||
r = J.submit("comfyui", {})
|
||||
assert r["priority"] == T.get_tenant("comfyui").priority
|
||||
|
||||
def test_pending_jobs_are_listed_in_execution_order(self, queue):
|
||||
J.submit("ollama", {}, priority=10, label="low")
|
||||
J.submit("ollama", {}, priority=90, label="high")
|
||||
J.submit("ollama", {}, priority=50, label="mid")
|
||||
order = [j["label"] for j in J.listing(J.PENDING)]
|
||||
assert order == ["high", "mid", "low"]
|
||||
|
||||
def test_equal_priority_is_first_in_first_out(self, queue):
|
||||
a = J.submit("ollama", {}, priority=50, label="a")["id"]
|
||||
time.sleep(0.01)
|
||||
J.submit("ollama", {}, priority=50, label="b")
|
||||
assert J._next_job()["id"] == a
|
||||
|
||||
def test_the_queue_has_no_depth_limit(self, queue):
|
||||
# "Unbounded" is the point; it lives on disk, not in memory.
|
||||
for i in range(500):
|
||||
J.submit("ollama", {}, label=f"j{i}")
|
||||
assert J.stats()["queue_depth"] == 500
|
||||
|
||||
def test_queue_survives_a_restart(self, queue):
|
||||
J.submit("ollama", {}, label="persisted")
|
||||
J._cache = None # nothing in-process is holding it
|
||||
assert [j["label"] for j in J.listing(J.PENDING)] == ["persisted"]
|
||||
|
||||
|
||||
class TestCancellation:
|
||||
def test_pending_jobs_can_be_cancelled(self, queue):
|
||||
jid = J.submit("ollama", {})["id"]
|
||||
assert J.cancel(jid)["success"] is True
|
||||
assert J.get(jid)["state"] == J.CANCELLED
|
||||
|
||||
def test_running_work_is_never_cancelled(self, queue):
|
||||
# This service frees VRAM by asking, never by killing work in flight.
|
||||
jid = J.submit("ollama", {})["id"]
|
||||
J._mark(jid, J.RUNNING, started_at=time.time())
|
||||
assert J.cancel(jid)["success"] is False
|
||||
assert J.get(jid)["state"] == J.RUNNING
|
||||
|
||||
def test_clearing_the_queue_leaves_running_work_alone(self, queue):
|
||||
running = J.submit("ollama", {})["id"]
|
||||
J._mark(running, J.RUNNING, started_at=time.time())
|
||||
J.submit("ollama", {})
|
||||
J.submit("ollama", {})
|
||||
assert J.clear_pending()["cancelled"] == 2
|
||||
assert J.get(running)["state"] == J.RUNNING
|
||||
|
||||
|
||||
class TestOrphanRecovery:
|
||||
def test_jobs_left_running_by_a_dead_process_are_requeued(self, queue):
|
||||
"""RUNNING means "this process is working on it".
|
||||
|
||||
An exception during dispatch left the row RUNNING while the scheduler moved on,
|
||||
so the job never finished and never retried.
|
||||
"""
|
||||
jid = J.submit("ollama", {})["id"]
|
||||
J._mark(jid, J.RUNNING, started_at=time.time())
|
||||
assert J.requeue_orphans() == 1
|
||||
job = J.get(jid)
|
||||
assert job["state"] == J.PENDING and job["started_at"] is None
|
||||
|
||||
def test_finished_jobs_are_untouched_by_recovery(self, queue):
|
||||
done = J.submit("ollama", {})["id"]
|
||||
J._mark(done, J.DONE, finished_at=time.time())
|
||||
assert J.requeue_orphans() == 0
|
||||
assert J.get(done)["state"] == J.DONE
|
||||
|
||||
|
||||
class TestPerJobVramRequirement:
|
||||
"""A tenant-wide figure cannot be right for an LLM."""
|
||||
|
||||
def test_llm_requirement_comes_from_the_model_being_loaded(self, queue, monkeypatch):
|
||||
import vram_arbitrator
|
||||
monkeypatch.setattr(vram_arbitrator, "_model_size_bytes",
|
||||
lambda m: int(12.87 * 1024 ** 3))
|
||||
s = J.Scheduler()
|
||||
# 12.87 GB on disk occupies ~14.9 GB once context and KV cache are allocated.
|
||||
assert 14.5 < s._job_vram_requirement("ollama", {"model": "big"}) < 15.5
|
||||
|
||||
def test_a_small_model_needs_correspondingly_less(self, queue, monkeypatch):
|
||||
import vram_arbitrator
|
||||
monkeypatch.setattr(vram_arbitrator, "_model_size_bytes",
|
||||
lambda m: int(1.96 * 1024 ** 3))
|
||||
s = J.Scheduler()
|
||||
assert s._job_vram_requirement("ollama", {"model": "small"}) < 3.0
|
||||
|
||||
def test_falls_back_to_the_tenant_figure_when_the_model_is_unknown(self, queue, monkeypatch):
|
||||
import vram_arbitrator
|
||||
monkeypatch.setattr(vram_arbitrator, "_model_size_bytes", lambda m: 0)
|
||||
s = J.Scheduler()
|
||||
import tenants as T
|
||||
assert (s._job_vram_requirement("ollama", {"model": "?"})
|
||||
== T.get_tenant("ollama").needs_vram_gb)
|
||||
|
||||
def test_non_llm_tenants_use_their_declared_figure(self, queue):
|
||||
import tenants as T
|
||||
s = J.Scheduler()
|
||||
assert (s._job_vram_requirement("comfyui", {})
|
||||
== T.get_tenant("comfyui").needs_vram_gb)
|
||||
|
||||
|
||||
class TestSchedulerStatus:
|
||||
def test_stats_report_depth_and_the_oldest_wait(self, queue):
|
||||
J.submit("ollama", {})
|
||||
J.submit("comfyui", {})
|
||||
st = J.stats()
|
||||
assert st["queue_depth"] == 2
|
||||
assert st["pending_by_tenant"] == {"ollama": 1, "comfyui": 1}
|
||||
assert st["oldest_pending_s"] is not None
|
||||
|
||||
def test_status_includes_queue_stats(self, queue):
|
||||
J.submit("ollama", {})
|
||||
s = J.Scheduler().get_status()
|
||||
assert s["queue_depth"] == 1 and s["running"] is False
|
||||
|
||||
def test_a_blocked_job_reports_why_it_is_waiting(self, queue):
|
||||
# Silence here is what left three jobs pending indefinitely with no explanation.
|
||||
s = J.Scheduler()
|
||||
s.blocked = {"id": "x", "tenant": "ollama", "reason": "needs 14.93 GB",
|
||||
"waited_s": 42.0, "since": time.time()}
|
||||
assert "14.93" in s.get_status()["blocked"]["reason"]
|
||||
|
||||
|
||||
class TestDispatchers:
|
||||
def test_every_tenant_kind_that_can_run_work_has_a_dispatcher(self):
|
||||
import tenants as T
|
||||
assert T.KIND_LLM in J.DISPATCHERS
|
||||
assert T.KIND_DIFFUSION in J.DISPATCHERS
|
||||
|
||||
def test_ollama_dispatch_reports_failure_rather_than_raising(self, queue, monkeypatch):
|
||||
class _R:
|
||||
status_code = 500
|
||||
text = '{"error":"llama-server process has terminated: cudaMalloc failed"}'
|
||||
class _C:
|
||||
async def __aenter__(self): return self
|
||||
async def __aexit__(self, *a): return False
|
||||
async def post(self, url, json=None): return _R()
|
||||
monkeypatch.setattr(J.httpx, "AsyncClient", lambda **k: _C())
|
||||
res = asyncio.run(J._dispatch_ollama({"model": "m", "prompt": "hi"}))
|
||||
assert res["ok"] is False and "cudaMalloc" in res["error"]
|
||||
116
tests/test_overclock_profiles.py
Normal file
116
tests/test_overclock_profiles.py
Normal file
@@ -0,0 +1,116 @@
|
||||
"""overclock_manager profile storage: load/merge/validate only.
|
||||
|
||||
Nothing here applies a profile or shells out. CONFIG_PATH is redirected to tmp_path so the
|
||||
repo's real overclock_profiles.json is never written, and conftest's autouse fixture blocks
|
||||
_sh() and every actuation entry point.
|
||||
"""
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
import overclock_manager as ocm
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def config(tmp_path, monkeypatch):
|
||||
"""Redirect the profile store to a temp file. Returns a writer for its contents."""
|
||||
path = tmp_path / "overclock_profiles.json"
|
||||
monkeypatch.setattr(ocm, "CONFIG_PATH", str(path))
|
||||
|
||||
def write(data):
|
||||
path.write_text(json.dumps(data))
|
||||
return type("Cfg", (), {"path": path, "write": staticmethod(write)})
|
||||
|
||||
|
||||
def test_defaults_are_returned_when_no_config_file_exists(config):
|
||||
profiles = ocm.load_profiles()
|
||||
assert set(profiles) == set(ocm.DEFAULT_PROFILES)
|
||||
assert profiles["ollama"]["power_limit_w"] == ocm.DEFAULT_PROFILES["ollama"]["power_limit_w"]
|
||||
|
||||
|
||||
def test_stored_values_override_defaults_key_by_key(config):
|
||||
"""Only the keys present on disk change; the rest of the default profile survives.
|
||||
A whole-profile replacement would silently drop keys added by a later version."""
|
||||
config.write({"ollama": {"power_limit_w": 340}})
|
||||
profiles = ocm.load_profiles()
|
||||
assert profiles["ollama"]["power_limit_w"] == 340
|
||||
assert profiles["ollama"]["label"] == ocm.DEFAULT_PROFILES["ollama"]["label"]
|
||||
assert profiles["ollama"]["mem_offset_mhz"] == ocm.DEFAULT_PROFILES["ollama"]["mem_offset_mhz"]
|
||||
|
||||
|
||||
def test_unknown_stored_profiles_are_preserved(config):
|
||||
"""A user-created profile is not in DEFAULT_PROFILES and must not be dropped on load."""
|
||||
config.write({"my-custom": {"label": "mine", "power_limit_w": 300}})
|
||||
profiles = ocm.load_profiles()
|
||||
assert profiles["my-custom"]["power_limit_w"] == 300
|
||||
assert set(ocm.DEFAULT_PROFILES) <= set(profiles)
|
||||
|
||||
|
||||
def test_loading_does_not_mutate_the_defaults(config):
|
||||
"""load_profiles deep-copies DEFAULT_PROFILES; if it did not, one load with a stored
|
||||
override would poison the defaults for the rest of the process."""
|
||||
config.write({"ollama": {"power_limit_w": 111}})
|
||||
ocm.load_profiles()
|
||||
config.write({})
|
||||
assert ocm.load_profiles()["ollama"]["power_limit_w"] == \
|
||||
ocm.DEFAULT_PROFILES["ollama"]["power_limit_w"]
|
||||
assert ocm.DEFAULT_PROFILES["ollama"]["power_limit_w"] != 111
|
||||
|
||||
|
||||
def test_corrupt_config_falls_back_to_defaults(config):
|
||||
"""A half-written JSON file must not take the whole service down at startup."""
|
||||
config.path.write_text("{ not valid json")
|
||||
profiles = ocm.load_profiles()
|
||||
assert set(profiles) == set(ocm.DEFAULT_PROFILES)
|
||||
|
||||
|
||||
def test_set_profile_rejects_an_unknown_name(config):
|
||||
"""set_profile edits an existing profile; it is not a create-anything endpoint."""
|
||||
res = ocm.set_profile("no-such-profile", {"power_limit_w": 400})
|
||||
assert res["success"] is False
|
||||
assert "unknown profile" in res["error"]
|
||||
assert not config.path.exists()
|
||||
|
||||
|
||||
def test_set_profile_persists_a_partial_update(config):
|
||||
res = ocm.set_profile("comfy", {"power_limit_w": 350})
|
||||
assert res["success"] is True
|
||||
stored = json.loads(config.path.read_text())
|
||||
assert stored["comfy"]["power_limit_w"] == 350
|
||||
# The other profiles are written back untouched.
|
||||
assert stored["ollama"]["power_limit_w"] == ocm.DEFAULT_PROFILES["ollama"]["power_limit_w"]
|
||||
assert ocm.load_profiles()["comfy"]["power_limit_w"] == 350
|
||||
|
||||
|
||||
def test_set_profile_can_edit_a_previously_stored_custom_profile(config):
|
||||
config.write({"my-custom": {"label": "mine", "power_limit_w": 300}})
|
||||
assert ocm.set_profile("my-custom", {"power_limit_w": 310})["success"] is True
|
||||
assert ocm.load_profiles()["my-custom"]["power_limit_w"] == 310
|
||||
|
||||
|
||||
def test_get_profiles_matches_load_profiles(config):
|
||||
config.write({"ollama": {"core_offset_mhz": 42}})
|
||||
assert ocm.get_profiles() == ocm.load_profiles()
|
||||
|
||||
|
||||
def test_save_profiles_reports_failure_instead_of_raising(tmp_path, monkeypatch):
|
||||
"""The dashboard calls this; an unwritable config must surface as success=False."""
|
||||
monkeypatch.setattr(ocm, "CONFIG_PATH", str(tmp_path / "no-such-dir" / "p.json"))
|
||||
assert ocm.save_profiles({"ollama": {}}) is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", ["ollama", "comfy", "balanced"])
|
||||
def test_default_profiles_declare_the_keys_the_appliers_read(name):
|
||||
"""apply_profile and the thermal governor both index these keys directly; a profile
|
||||
missing one would KeyError inside the actuation thread."""
|
||||
cfg = ocm.DEFAULT_PROFILES[name]
|
||||
for key in ("label", "power_limit_w", "core_offset_mhz", "mem_offset_mhz",
|
||||
"lock_core_min", "lock_core_max"):
|
||||
assert key in cfg, f"{name} is missing {key}"
|
||||
|
||||
|
||||
def test_default_power_limits_stay_within_the_cards_range():
|
||||
"""RTX 4080 SUPER: 320 W stock, 370 W maximum. A profile above that is silently
|
||||
clamped by the driver and the sweep results become meaningless."""
|
||||
for name, cfg in ocm.DEFAULT_PROFILES.items():
|
||||
assert 100 <= cfg["power_limit_w"] <= 370, name
|
||||
421
tests/test_ram_optimizer.py
Normal file
421
tests/test_ram_optimizer.py
Normal file
@@ -0,0 +1,421 @@
|
||||
"""ram_optimizer: page-cache residency measurement, model discovery, warm planning.
|
||||
|
||||
All file IO here happens against files this test creates in tmp_path. Nothing reads a
|
||||
real model blob, and nothing calls warm_file_to_ram.
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
import ram_optimizer as ro
|
||||
|
||||
MIB = 1024 * 1024
|
||||
|
||||
|
||||
# --------------------------------------------------------------- warm_confident rule
|
||||
|
||||
class _FakeCachestat:
|
||||
"""Stand-in for the kernel's struct cachestat."""
|
||||
|
||||
def __init__(self, nr_cache):
|
||||
self.nr_cache = nr_cache
|
||||
self.nr_dirty = 0
|
||||
self.nr_evicted = 0
|
||||
|
||||
|
||||
def _file_of_pages(tmp_path, pages):
|
||||
p = tmp_path / f"blob_{pages}.bin"
|
||||
p.write_bytes(b"\0" * (pages * ro.PAGE_SIZE))
|
||||
return str(p)
|
||||
|
||||
|
||||
def _force_cachestat(monkeypatch, resident_pages):
|
||||
monkeypatch.setattr(ro, "_cachestat",
|
||||
lambda fd, offset, length: _FakeCachestat(resident_pages))
|
||||
|
||||
|
||||
def _force_probe(monkeypatch, pct):
|
||||
monkeypatch.setattr(ro, "_cachestat", lambda fd, offset, length: None)
|
||||
monkeypatch.setattr(ro, "_throughput_probe",
|
||||
lambda fd, size, windows=None: {"resident_pct": pct, "windows": 12,
|
||||
"median_gbps": 3.0})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("resident_pages,expect_confident", [
|
||||
(100, True), # 100% exact
|
||||
(90, True), # exactly WARM_SKIP_THRESHOLD_PCT
|
||||
(89, False), # just under the threshold
|
||||
(0, False),
|
||||
])
|
||||
def test_cachestat_reading_is_trusted_at_or_above_threshold(tmp_path, monkeypatch,
|
||||
resident_pages, expect_confident):
|
||||
"""An exact cachestat reading >= WARM_SKIP_THRESHOLD_PCT (90) may be trusted to skip
|
||||
warming; below it, it may not."""
|
||||
path = _file_of_pages(tmp_path, 100)
|
||||
_force_cachestat(monkeypatch, resident_pages)
|
||||
res = ro.page_residency(path)
|
||||
assert res["method"] == "cachestat"
|
||||
assert res["resident_pct"] == pytest.approx(float(resident_pages))
|
||||
assert res["warm_confident"] is expect_confident
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pct,expect_confident", [
|
||||
(100.0, True), # only an unblemished probe is trustworthy
|
||||
(99.9, False),
|
||||
(95.0, False), # would pass `warm`, must NOT pass `warm_confident`
|
||||
(90.0, False),
|
||||
])
|
||||
def test_probe_reading_is_only_trusted_at_exactly_100_percent(tmp_path, monkeypatch,
|
||||
pct, expect_confident):
|
||||
"""Correctness fix, not a style choice: a 12-window probe once cleared 90% on a file
|
||||
that was mostly cold (a 12.87 GB blob reported 'already resident' then loaded at
|
||||
2.44 GB/s). Only a perfect probe score may skip work."""
|
||||
path = _file_of_pages(tmp_path, 100)
|
||||
_force_probe(monkeypatch, pct)
|
||||
res = ro.page_residency(path)
|
||||
assert res["method"] == "probe"
|
||||
assert res["warm_confident"] is expect_confident
|
||||
|
||||
|
||||
def test_probe_at_95_percent_is_warm_but_not_confident(tmp_path, monkeypatch):
|
||||
"""`warm` and `warm_confident` are different questions and must not be conflated:
|
||||
warm drives display, warm_confident drives skipping work."""
|
||||
path = _file_of_pages(tmp_path, 100)
|
||||
_force_probe(monkeypatch, 95.0)
|
||||
res = ro.page_residency(path)
|
||||
assert res["warm"] is True
|
||||
assert res["warm_confident"] is False
|
||||
|
||||
|
||||
def test_warm_skip_threshold_constant_unchanged():
|
||||
# Pinned: build_warm_plan and the dashboard both key off this number.
|
||||
assert ro.WARM_SKIP_THRESHOLD_PCT == 90.0
|
||||
|
||||
|
||||
# --------------------------------------------------------------- page_residency basics
|
||||
|
||||
def test_page_residency_on_missing_file_reports_not_measurable(tmp_path):
|
||||
res = ro.page_residency(str(tmp_path / "nope.bin"))
|
||||
assert res["success"] is False
|
||||
assert res["measurable"] is False
|
||||
assert res["resident_pct"] == 0.0
|
||||
|
||||
|
||||
def test_page_residency_on_empty_file_short_circuits(tmp_path):
|
||||
"""A zero-byte file has no pages to measure; dividing by its page count would throw."""
|
||||
p = tmp_path / "empty.bin"
|
||||
p.write_bytes(b"")
|
||||
res = ro.page_residency(str(p))
|
||||
assert res["method"] == "empty"
|
||||
assert res["resident_pct"] == 0.0
|
||||
assert res["measurable"] is True
|
||||
|
||||
|
||||
def test_page_residency_refuses_to_guess_when_probing_is_disallowed(tmp_path, monkeypatch):
|
||||
"""With cachestat unavailable and allow_probe=False the answer is 'unknown', never 0%.
|
||||
Reporting 0% would make the planner re-warm a file that may be fully resident."""
|
||||
path = _file_of_pages(tmp_path, 8)
|
||||
monkeypatch.setattr(ro, "_cachestat", lambda fd, offset, length: None)
|
||||
res = ro.page_residency(path, allow_probe=False)
|
||||
assert res["measurable"] is False
|
||||
assert res["resident_pct"] is None
|
||||
assert res["warm"] is None
|
||||
assert res["method"] == "unavailable"
|
||||
|
||||
|
||||
def test_page_residency_reports_size_and_bytes_consistently(tmp_path, monkeypatch):
|
||||
path = _file_of_pages(tmp_path, 100)
|
||||
_force_cachestat(monkeypatch, 50)
|
||||
res = ro.page_residency(path)
|
||||
size = os.path.getsize(path)
|
||||
assert res["size_bytes"] == size
|
||||
assert res["resident_bytes"] == int(size * 0.5)
|
||||
|
||||
|
||||
def test_page_residency_against_a_real_file_holds_the_confidence_invariant(tmp_path):
|
||||
"""End-to-end with the real kernel path (whichever method is available here): the
|
||||
warm_confident rule must hold for whatever the machine actually reports."""
|
||||
path = _file_of_pages(tmp_path, 512)
|
||||
res = ro.page_residency(path)
|
||||
assert res["success"] is True
|
||||
assert res["method"] in ("cachestat", "probe")
|
||||
assert 0.0 <= res["resident_pct"] <= 100.0
|
||||
expected = ((res["method"] == "cachestat" and res["resident_pct"] >= 90.0)
|
||||
or (res["method"] == "probe" and res["resident_pct"] >= 100.0))
|
||||
assert res["warm_confident"] is expected
|
||||
|
||||
|
||||
# --------------------------------------------------------------- throughput probe
|
||||
|
||||
def test_throughput_probe_window_count_is_clamped_to_file_size(tmp_path):
|
||||
"""A file smaller than PROBE_WINDOWS * PROBE_WINDOW_BYTES must not be probed with
|
||||
more windows than it has, or offsets run past EOF."""
|
||||
p = tmp_path / "small.bin"
|
||||
p.write_bytes(b"\0" * (5 * MIB))
|
||||
fd = os.open(str(p), os.O_RDONLY)
|
||||
try:
|
||||
out = ro._throughput_probe(fd, 5 * MIB)
|
||||
finally:
|
||||
os.close(fd)
|
||||
assert 0 < out["windows"] <= 5 * MIB // ro.PROBE_WINDOW_BYTES
|
||||
assert 0.0 <= out["resident_pct"] <= 100.0
|
||||
|
||||
|
||||
def test_throughput_probe_honours_window_override(tmp_path):
|
||||
p = tmp_path / "medium.bin"
|
||||
p.write_bytes(b"\0" * (32 * MIB))
|
||||
fd = os.open(str(p), os.O_RDONLY)
|
||||
try:
|
||||
out = ro._throughput_probe(fd, 32 * MIB, windows_override=4)
|
||||
finally:
|
||||
os.close(fd)
|
||||
assert out["windows"] == 4
|
||||
|
||||
|
||||
def test_probe_cached_threshold_sits_between_measured_disk_and_cache_rates():
|
||||
# Measured on this box: cold NVMe 0.35-0.5 GB/s, page cache 3.2-13 GB/s.
|
||||
assert 0.5 < ro.PROBE_CACHED_GBPS < 3.2
|
||||
|
||||
|
||||
# --------------------------------------------------------------- ollama manifest parsing
|
||||
|
||||
MODEL_MEDIA_TYPE = "application/vnd.ollama.image.model"
|
||||
|
||||
|
||||
def _write_manifest(root, rel_dir, tag, layers):
|
||||
d = os.path.join(root, "manifests", rel_dir)
|
||||
os.makedirs(d, exist_ok=True)
|
||||
with open(os.path.join(d, tag), "w") as f:
|
||||
json.dump({"layers": layers}, f)
|
||||
|
||||
|
||||
def _write_blob(root, digest, size):
|
||||
blobs = os.path.join(root, "blobs")
|
||||
os.makedirs(blobs, exist_ok=True)
|
||||
path = os.path.join(blobs, digest.replace(":", "-"))
|
||||
with open(path, "wb") as f:
|
||||
f.write(b"\0" * size)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ollama_tree(tmp_path, monkeypatch):
|
||||
root = tmp_path / "ollama-models"
|
||||
root.mkdir()
|
||||
monkeypatch.setattr(ro, "OLLAMA_MODEL_DIRS", [str(root)])
|
||||
return str(root)
|
||||
|
||||
|
||||
def test_find_ollama_model_files_maps_library_model_to_its_blob(ollama_tree):
|
||||
"""registry/library/<name>/<tag> is the common case and must render as 'name:tag'."""
|
||||
_write_blob(ollama_tree, "sha256:aaa111", 4096)
|
||||
_write_manifest(ollama_tree, "registry.ollama.ai/library/llama3", "8b", [
|
||||
{"mediaType": MODEL_MEDIA_TYPE, "digest": "sha256:aaa111", "size": 987654321},
|
||||
])
|
||||
files = ro.find_ollama_model_files()
|
||||
assert len(files) == 1
|
||||
entry = files[0]
|
||||
assert entry["model"] == "llama3:8b"
|
||||
assert entry["filename"] == "sha256-aaa111"
|
||||
assert entry["size_bytes"] == 987654321 # taken from the manifest, not the stub blob
|
||||
assert entry["kind"] == "ollama"
|
||||
assert os.path.exists(entry["full_path"])
|
||||
|
||||
|
||||
def test_find_ollama_model_files_keeps_non_library_namespace(ollama_tree):
|
||||
"""A model pulled from a user namespace must keep it, or two different models with
|
||||
the same short name collide in the warm plan."""
|
||||
_write_blob(ollama_tree, "sha256:bbb222", 4096)
|
||||
_write_manifest(ollama_tree, "hf.co/someuser/qwen-coder", "q4", [
|
||||
{"mediaType": MODEL_MEDIA_TYPE, "digest": "sha256:bbb222", "size": 100},
|
||||
])
|
||||
files = ro.find_ollama_model_files()
|
||||
assert [f["model"] for f in files] == ["someuser/qwen-coder:q4"]
|
||||
|
||||
|
||||
def test_find_ollama_model_files_ignores_non_model_layers(ollama_tree):
|
||||
"""Manifests also list templates, params and licence layers; warming those is wasted
|
||||
IO and they are not the weights."""
|
||||
_write_blob(ollama_tree, "sha256:ccc333", 4096)
|
||||
_write_blob(ollama_tree, "sha256:ddd444", 32)
|
||||
_write_manifest(ollama_tree, "registry.ollama.ai/library/mistral", "7b", [
|
||||
{"mediaType": "application/vnd.ollama.image.template", "digest": "sha256:ddd444", "size": 32},
|
||||
{"mediaType": MODEL_MEDIA_TYPE, "digest": "sha256:ccc333", "size": 500},
|
||||
])
|
||||
files = ro.find_ollama_model_files()
|
||||
assert len(files) == 1
|
||||
assert files[0]["filename"] == "sha256-ccc333"
|
||||
|
||||
|
||||
def test_find_ollama_model_files_skips_layers_whose_blob_is_missing(ollama_tree):
|
||||
"""A partially pulled model leaves a manifest with no blob; returning that path would
|
||||
make every downstream residency check fail."""
|
||||
_write_manifest(ollama_tree, "registry.ollama.ai/library/ghost", "latest", [
|
||||
{"mediaType": MODEL_MEDIA_TYPE, "digest": "sha256:missing", "size": 10},
|
||||
])
|
||||
assert ro.find_ollama_model_files() == []
|
||||
|
||||
|
||||
def test_find_ollama_model_files_tolerates_corrupt_manifest_json(ollama_tree):
|
||||
"""One unreadable manifest must not take out discovery of every other model."""
|
||||
bad_dir = os.path.join(ollama_tree, "manifests", "registry.ollama.ai/library/broken")
|
||||
os.makedirs(bad_dir)
|
||||
with open(os.path.join(bad_dir, "latest"), "w") as f:
|
||||
f.write("{not json")
|
||||
_write_blob(ollama_tree, "sha256:eee555", 4096)
|
||||
_write_manifest(ollama_tree, "registry.ollama.ai/library/good", "latest", [
|
||||
{"mediaType": MODEL_MEDIA_TYPE, "digest": "sha256:eee555", "size": 10},
|
||||
])
|
||||
assert [f["model"] for f in ro.find_ollama_model_files()] == ["good:latest"]
|
||||
|
||||
|
||||
def test_find_ollama_model_files_deduplicates_same_model_and_blob(ollama_tree, monkeypatch):
|
||||
"""The same root listed twice (or a duplicated layer) must not double-count bytes
|
||||
against the warm budget."""
|
||||
_write_blob(ollama_tree, "sha256:fff666", 4096)
|
||||
_write_manifest(ollama_tree, "registry.ollama.ai/library/dup", "latest", [
|
||||
{"mediaType": MODEL_MEDIA_TYPE, "digest": "sha256:fff666", "size": 10},
|
||||
{"mediaType": MODEL_MEDIA_TYPE, "digest": "sha256:fff666", "size": 10},
|
||||
])
|
||||
monkeypatch.setattr(ro, "OLLAMA_MODEL_DIRS", [ollama_tree, ollama_tree])
|
||||
assert len(ro.find_ollama_model_files()) == 1
|
||||
|
||||
|
||||
def test_find_ollama_model_files_returns_empty_when_no_manifest_dir(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(ro, "OLLAMA_MODEL_DIRS", [str(tmp_path / "does-not-exist")])
|
||||
assert ro.find_ollama_model_files() == []
|
||||
|
||||
|
||||
# --------------------------------------------------------------- warm planning
|
||||
|
||||
GIB = 1024 ** 3
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def planner(monkeypatch):
|
||||
"""build_warm_plan with every IO boundary stubbed: catalog, residency, meminfo, scores."""
|
||||
state = {"residency": {}, "scores": {}}
|
||||
|
||||
def _catalog(force_refresh=False):
|
||||
return state["catalog"]
|
||||
|
||||
monkeypatch.setattr(ro, "get_model_catalog", _catalog)
|
||||
monkeypatch.setattr(ro, "_warm_priority", lambda days=30.0: state["scores"])
|
||||
monkeypatch.setattr(ro, "get_detailed_meminfo",
|
||||
lambda: {"available_bytes": 40 * GIB, "available_gb": 40.0})
|
||||
monkeypatch.setattr(ro, "page_residency",
|
||||
lambda path, allow_probe=True, probe_windows=None:
|
||||
state["residency"].get(path, {"resident_pct": 0.0,
|
||||
"warm_confident": False}))
|
||||
return state
|
||||
|
||||
|
||||
def _ollama_entry(name, path, gb):
|
||||
return {"model": name, "full_path": path, "size_bytes": int(gb * GIB),
|
||||
"size_gb": gb, "kind": "ollama"}
|
||||
|
||||
|
||||
def _comfy_entry(rel, path, gb, mtime):
|
||||
return {"rel_path": rel, "full_path": path, "size_bytes": int(gb * GIB),
|
||||
"size_gb": gb, "kind": "comfy", "mtime": mtime}
|
||||
|
||||
|
||||
def test_warm_plan_orders_by_usage_score(planner):
|
||||
"""The whole point of the plan is that the most-used model is warmed first, because
|
||||
the budget usually cannot cover everything."""
|
||||
planner["catalog"] = {
|
||||
"ollama": [_ollama_entry("rare:latest", "/m/rare", 1.0),
|
||||
_ollama_entry("hot:latest", "/m/hot", 1.0)],
|
||||
"comfy": [],
|
||||
}
|
||||
planner["scores"] = {"hot:latest": 50.0, "rare:latest": 0.1}
|
||||
plan = ro.build_warm_plan(budget_gb=10.0)
|
||||
assert [e["name"] for e in plan["plan"]] == ["hot:latest", "rare:latest"]
|
||||
|
||||
|
||||
def test_warm_plan_skips_confidently_warm_files_without_spending_budget(planner):
|
||||
"""Already-resident files cost nothing and must not consume budget another file needs."""
|
||||
planner["catalog"] = {
|
||||
"ollama": [_ollama_entry("warm:latest", "/m/warm", 8.0),
|
||||
_ollama_entry("cold:latest", "/m/cold", 8.0)],
|
||||
"comfy": [],
|
||||
}
|
||||
planner["scores"] = {"warm:latest": 10.0, "cold:latest": 5.0}
|
||||
planner["residency"]["/m/warm"] = {"resident_pct": 100.0, "warm_confident": True}
|
||||
plan = ro.build_warm_plan(budget_gb=10.0)
|
||||
assert [e["name"] for e in plan["plan"]] == ["cold:latest"]
|
||||
assert [e["action"] for e in plan["skipped"]] == ["already-warm"]
|
||||
assert plan["planned_gb"] == pytest.approx(8.0, abs=0.01)
|
||||
|
||||
|
||||
def test_warm_plan_does_not_skip_a_high_but_unconfident_residency(planner):
|
||||
"""95% from a probe is not permission to skip — this is the bug the warm_confident
|
||||
flag exists to prevent."""
|
||||
planner["catalog"] = {"ollama": [_ollama_entry("m:latest", "/m/x", 4.0)], "comfy": []}
|
||||
planner["residency"]["/m/x"] = {"resident_pct": 95.0, "warm_confident": False}
|
||||
plan = ro.build_warm_plan(budget_gb=10.0)
|
||||
assert [e["name"] for e in plan["plan"]] == ["m:latest"]
|
||||
# Only the missing 5% has to be read.
|
||||
assert plan["plan"][0]["bytes_to_read"] == pytest.approx(int(4.0 * GIB) * 0.05, rel=0.01)
|
||||
|
||||
|
||||
def test_warm_plan_charges_only_the_non_resident_fraction(planner):
|
||||
planner["catalog"] = {"ollama": [_ollama_entry("m:latest", "/m/x", 10.0)], "comfy": []}
|
||||
planner["residency"]["/m/x"] = {"resident_pct": 50.0, "warm_confident": False}
|
||||
plan = ro.build_warm_plan(budget_gb=6.0)
|
||||
assert plan["planned_gb"] == pytest.approx(5.0, abs=0.01)
|
||||
|
||||
|
||||
def test_warm_plan_stops_at_the_budget(planner):
|
||||
"""Warming past the budget just evicts what was warmed first, so over-budget entries
|
||||
are reported as skipped rather than planned."""
|
||||
planner["catalog"] = {
|
||||
"ollama": [_ollama_entry("a", "/m/a", 6.0), _ollama_entry("b", "/m/b", 6.0)],
|
||||
"comfy": [],
|
||||
}
|
||||
planner["scores"] = {"a": 9.0, "b": 1.0}
|
||||
plan = ro.build_warm_plan(budget_gb=8.0)
|
||||
assert [e["name"] for e in plan["plan"]] == ["a"]
|
||||
assert plan["skipped"][0]["name"] == "b"
|
||||
assert plan["skipped"][0]["action"] == "over-budget"
|
||||
assert plan["planned_gb"] <= plan["budget_gb"]
|
||||
|
||||
|
||||
def test_warm_plan_deduplicates_by_path(planner):
|
||||
"""The same file reachable from both catalogs must be planned once, or the budget is
|
||||
charged twice for one read."""
|
||||
planner["catalog"] = {
|
||||
"ollama": [_ollama_entry("shared", "/m/shared", 2.0)],
|
||||
"comfy": [_comfy_entry("shared.safetensors", "/m/shared", 2.0, 0)],
|
||||
}
|
||||
plan = ro.build_warm_plan(budget_gb=100.0)
|
||||
assert len(plan["plan"]) == 1
|
||||
assert plan["planned_gb"] == pytest.approx(2.0, abs=0.01)
|
||||
|
||||
|
||||
def test_warm_plan_default_budget_leaves_headroom(planner):
|
||||
"""An unbounded budget would push the box into reclaim; the default is 70% of
|
||||
MemAvailable."""
|
||||
planner["catalog"] = {"ollama": [], "comfy": []}
|
||||
plan = ro.build_warm_plan()
|
||||
assert plan["budget_gb"] == pytest.approx(40.0 * 0.7, abs=0.01)
|
||||
|
||||
|
||||
def test_warm_plan_ranks_recent_comfy_checkpoints_above_stale_ones(planner):
|
||||
"""ComfyUI files have no usage history, so recency by mtime is the ranking signal."""
|
||||
import time
|
||||
now = time.time()
|
||||
planner["catalog"] = {
|
||||
"ollama": [],
|
||||
"comfy": [_comfy_entry("old.safetensors", "/c/old", 1.0, now - 90 * 86400),
|
||||
_comfy_entry("new.safetensors", "/c/new", 1.0, now - 60)],
|
||||
}
|
||||
plan = ro.build_warm_plan(budget_gb=10.0)
|
||||
assert [e["name"] for e in plan["plan"]] == ["new.safetensors", "old.safetensors"]
|
||||
|
||||
|
||||
def test_warm_plan_on_empty_catalog_is_a_valid_empty_plan(planner):
|
||||
planner["catalog"] = {"ollama": [], "comfy": []}
|
||||
plan = ro.build_warm_plan(budget_gb=1.0)
|
||||
assert plan["warm_count"] == 0 and plan["plan"] == [] and plan["skipped"] == []
|
||||
342
tests/test_telemetry_store.py
Normal file
342
tests/test_telemetry_store.py
Normal file
@@ -0,0 +1,342 @@
|
||||
"""telemetry_store: schema, the background writer round-trip, and the analytics queries.
|
||||
|
||||
Every test in this module runs against a throwaway SQLite file via the `temp_db` fixture
|
||||
(which monkeypatches telemetry_store.DB_PATH). The production hyperswap.db is never
|
||||
opened: the writer thread, _read_conn() and _rows() all resolve DB_PATH at call time.
|
||||
"""
|
||||
import os
|
||||
import sqlite3
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
import telemetry_store as ts
|
||||
|
||||
|
||||
def _drain(timeout=3.0):
|
||||
"""Wait until the writer thread has dequeued everything submitted so far.
|
||||
|
||||
NOTE: stop() flushes the writer's *pending* batch but does not drain the submission
|
||||
queue, so a stop() racing a just-submitted row can drop it. Tests wait explicitly
|
||||
rather than depending on that race. (See tests/README.md.)
|
||||
"""
|
||||
w = ts._writer
|
||||
assert w is not None, "writer not started"
|
||||
deadline = time.time() + timeout
|
||||
while time.time() < deadline:
|
||||
if w.q.empty():
|
||||
time.sleep(0.05) # let the in-flight item finish its execute/commit
|
||||
return
|
||||
time.sleep(0.01)
|
||||
raise AssertionError("telemetry writer never drained its queue")
|
||||
|
||||
|
||||
def _seed(db_path, events=(), telemetry=()):
|
||||
"""Insert rows directly, bypassing the writer thread, for query tests."""
|
||||
conn = sqlite3.connect(db_path)
|
||||
conn.executescript(ts.SCHEMA)
|
||||
for e in events:
|
||||
conn.execute(ts._EVENT_SQL, (
|
||||
e.get("ts", time.time()), e.get("event_type"), e.get("source"), e.get("target"),
|
||||
e.get("profile"), e.get("duration_ms"), e.get("load_duration_ms"),
|
||||
e.get("yield_confirm_ms"), e.get("tokens_per_sec"), e.get("bytes_loaded"),
|
||||
e.get("load_gbps"), e.get("cache_status"), e.get("detail"),
|
||||
))
|
||||
for t in telemetry:
|
||||
conn.execute(ts._TELEMETRY_SQL, (
|
||||
t.get("ts", time.time()), t.get("profile"), t.get("gpu_util_pct"),
|
||||
t.get("mem_util_pct"), t.get("temp_c"), t.get("power_w"), t.get("power_limit_w"),
|
||||
t.get("fan_pct"), t.get("clock_sm_mhz"), t.get("clock_mem_mhz"),
|
||||
t.get("vram_used_bytes"), t.get("ollama_bytes"), t.get("comfy_bytes"),
|
||||
t.get("system_bytes"), t.get("ram_used_bytes"), t.get("ram_cached_bytes"),
|
||||
t.get("pcie_tx_kbps"), t.get("pcie_rx_kbps"), t.get("throttle_reasons"),
|
||||
))
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
|
||||
# --------------------------------------------------------------- writer round-trip
|
||||
|
||||
def test_start_creates_the_schema(temp_db):
|
||||
"""The service starts against a database that may not exist yet."""
|
||||
ts.start()
|
||||
ts.stop()
|
||||
assert os.path.exists(temp_db)
|
||||
conn = sqlite3.connect(temp_db)
|
||||
tables = {r[0] for r in conn.execute("SELECT name FROM sqlite_master WHERE type='table'")}
|
||||
conn.close()
|
||||
assert {"telemetry", "events", "autotune_runs"} <= tables
|
||||
|
||||
|
||||
def test_record_event_round_trips_through_the_writer(temp_db):
|
||||
"""Events are committed immediately (not batched) because they are what the analytics
|
||||
are built on."""
|
||||
ts.start()
|
||||
ts.record_event({
|
||||
"event_type": "LLM Model Switch", "source": "None", "target": "llama3:8b",
|
||||
"duration_ms": 5120.0, "load_duration_ms": 4901.0, "tokens_per_sec": 61.2,
|
||||
"bytes_loaded": 13819000000, "load_gbps": 2.63, "cache_status": "RAM Cache Hit ⚡",
|
||||
}, profile="ollama")
|
||||
_drain()
|
||||
ts.stop()
|
||||
|
||||
rows = ts.recent_events()
|
||||
assert len(rows) == 1
|
||||
assert rows[0]["target"] == "llama3:8b"
|
||||
assert rows[0]["profile"] == "ollama"
|
||||
assert rows[0]["cache_status"] == "RAM Cache Hit ⚡"
|
||||
assert rows[0]["load_gbps"] == pytest.approx(2.63)
|
||||
|
||||
|
||||
def test_record_telemetry_flushes_on_stop(temp_db):
|
||||
"""Telemetry is batched on a FLUSH_INTERVAL_S timer; stop() must not drop what is
|
||||
still pending, or the last seconds before a restart are lost."""
|
||||
ts.start()
|
||||
ts.record_telemetry(
|
||||
gpu={"available": True, "gpu_util_pct": 88.0, "temperature_c": 71.0,
|
||||
"power_w": 355.0, "clock_sm_mhz": 2790.0, "clock_mem_mhz": 10501.0,
|
||||
"vram_used_bytes": 13 * 1024 ** 3,
|
||||
"breakdown": {"ollama_gb": 12.0, "comfyui_gb": 0.0, "system_gb": 0.5}},
|
||||
ram={"used_bytes": 20 * 1024 ** 3, "cached_bytes": 30 * 1024 ** 3},
|
||||
profile="ollama", throttle_reasons="sw_power_cap")
|
||||
_drain()
|
||||
ts.stop()
|
||||
|
||||
conn = sqlite3.connect(temp_db)
|
||||
row = conn.execute("SELECT profile, temp_c, ollama_bytes, throttle_reasons "
|
||||
"FROM telemetry").fetchone()
|
||||
conn.close()
|
||||
assert row[0] == "ollama"
|
||||
assert row[1] == pytest.approx(71.0)
|
||||
assert row[2] == 12 * 1024 ** 3
|
||||
assert row[3] == "sw_power_cap"
|
||||
|
||||
|
||||
def test_record_telemetry_drops_unavailable_gpu_samples(temp_db):
|
||||
"""A failed NVML read must not write a row of nulls that then skews every average."""
|
||||
ts.start()
|
||||
ts.record_telemetry(gpu={"available": False}, ram={}, profile="ollama")
|
||||
_drain()
|
||||
ts.stop()
|
||||
conn = sqlite3.connect(temp_db)
|
||||
n = conn.execute("SELECT COUNT(*) FROM telemetry").fetchone()[0]
|
||||
conn.close()
|
||||
assert n == 0
|
||||
|
||||
|
||||
def test_recording_without_a_running_writer_is_a_no_op(temp_db):
|
||||
"""Callers on the hot path must never need to know whether the store was started."""
|
||||
ts.record_event({"event_type": "LLM Model Switch", "target": "x"})
|
||||
ts.record_telemetry(gpu={"available": True, "breakdown": {}}, ram={})
|
||||
assert not os.path.exists(temp_db)
|
||||
|
||||
|
||||
def test_record_event_never_raises_on_a_malformed_event(temp_db):
|
||||
"""This is called from the swap path; it is documented as never raising."""
|
||||
ts.start()
|
||||
ts.record_event({})
|
||||
ts.record_event({"event_type": "Weird", "duration_ms": "not-a-number"})
|
||||
_drain()
|
||||
ts.stop()
|
||||
|
||||
|
||||
def test_autotune_row_round_trips(temp_db):
|
||||
ts.start()
|
||||
ts.record_autotune({"profile": "ollama", "knob": "lock_mem_mhz", "mem_offset_mhz": 0,
|
||||
"tokens_per_sec": 62.5, "temp_c": 74.0, "stable": True})
|
||||
_drain()
|
||||
ts.stop()
|
||||
rows = ts.autotune_history()
|
||||
assert len(rows) == 1
|
||||
assert rows[0]["knob"] == "lock_mem_mhz"
|
||||
assert rows[0]["stable"] == 1 # stored as an INTEGER, not a bool
|
||||
|
||||
|
||||
# --------------------------------------------------------------- absent database
|
||||
|
||||
def test_queries_return_empty_when_the_database_does_not_exist(tmp_path, monkeypatch):
|
||||
"""A dashboard opened before the first sample must render empty, not 500."""
|
||||
monkeypatch.setattr(ts, "DB_PATH", str(tmp_path / "absent.db"))
|
||||
assert ts.profile_comparison() == []
|
||||
assert ts.model_usage_ranking() == []
|
||||
assert ts.timeseries() == []
|
||||
assert ts.recent_events() == []
|
||||
assert ts.autotune_history() == []
|
||||
stats = ts.swap_stats()
|
||||
assert stats["by_type"] == [] and stats["by_model"] == [] and stats["by_cache_status"] == []
|
||||
assert ts.db_info()["exists"] is False
|
||||
|
||||
|
||||
def test_queries_return_empty_on_a_corrupt_database(tmp_path, monkeypatch):
|
||||
db = tmp_path / "corrupt.db"
|
||||
db.write_bytes(b"this is not a sqlite file")
|
||||
monkeypatch.setattr(ts, "DB_PATH", str(db))
|
||||
assert ts.recent_events() == []
|
||||
assert ts.profile_comparison() == []
|
||||
|
||||
|
||||
# --------------------------------------------------------------- analytics
|
||||
|
||||
def test_profile_comparison_ranks_profiles_by_throughput(temp_db):
|
||||
"""The headline question this store exists to answer: which profile is actually
|
||||
faster? Ordering is by average tok/s, descending."""
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 60, "event_type": "LLM Model Switch", "profile": "ollama",
|
||||
"tokens_per_sec": 62.0, "load_gbps": 2.6, "load_duration_ms": 4900},
|
||||
{"ts": now - 50, "event_type": "LLM Model Switch", "profile": "ollama",
|
||||
"tokens_per_sec": 64.0, "load_gbps": 2.6, "load_duration_ms": 4900},
|
||||
{"ts": now - 40, "event_type": "LLM Model Switch", "profile": "balanced",
|
||||
"tokens_per_sec": 51.0, "load_gbps": 2.5, "load_duration_ms": 5100},
|
||||
], telemetry=[
|
||||
{"ts": now - 55, "profile": "ollama", "gpu_util_pct": 90, "temp_c": 74.0,
|
||||
"power_w": 360.0, "clock_sm_mhz": 2790, "clock_mem_mhz": 10501},
|
||||
{"ts": now - 45, "profile": "balanced", "gpu_util_pct": 90, "temp_c": 66.0,
|
||||
"power_w": 300.0, "clock_sm_mhz": 2600, "clock_mem_mhz": 9501},
|
||||
])
|
||||
out = ts.profile_comparison(days=1)
|
||||
assert [r["profile"] for r in out] == ["ollama", "balanced"]
|
||||
assert out[0]["swaps"] == 2
|
||||
assert out[0]["avg_tok_s"] == pytest.approx(63.0)
|
||||
# Thermals are joined in from the telemetry table for the same profile.
|
||||
assert out[0]["avg_temp_c"] == pytest.approx(74.0)
|
||||
|
||||
|
||||
def test_profile_comparison_ignores_events_without_measured_throughput(temp_db):
|
||||
"""A swap that produced no tokens tells us nothing about the profile's speed and
|
||||
would drag the average toward zero."""
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 10, "event_type": "LLM Model Switch", "profile": "ollama",
|
||||
"tokens_per_sec": 60.0},
|
||||
{"ts": now - 5, "event_type": "LLM Model Switch", "profile": "ollama",
|
||||
"tokens_per_sec": 0.0},
|
||||
])
|
||||
out = ts.profile_comparison(days=1)
|
||||
assert out[0]["swaps"] == 1 and out[0]["avg_tok_s"] == pytest.approx(60.0)
|
||||
|
||||
|
||||
def test_profile_comparison_excludes_samples_outside_the_window(temp_db):
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 40 * 86400, "event_type": "LLM Model Switch", "profile": "old",
|
||||
"tokens_per_sec": 99.0},
|
||||
{"ts": now - 60, "event_type": "LLM Model Switch", "profile": "ollama",
|
||||
"tokens_per_sec": 60.0},
|
||||
])
|
||||
assert [r["profile"] for r in ts.profile_comparison(days=7)] == ["ollama"]
|
||||
|
||||
|
||||
def test_profile_comparison_ignores_idle_telemetry(temp_db):
|
||||
"""Thermals are only meaningful while the GPU is doing work; idle samples (util <= 5%)
|
||||
would make every profile look cool."""
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 10, "event_type": "LLM Model Switch", "profile": "ollama",
|
||||
"tokens_per_sec": 60.0},
|
||||
], telemetry=[
|
||||
{"ts": now - 12, "profile": "ollama", "gpu_util_pct": 0, "temp_c": 35.0},
|
||||
{"ts": now - 11, "profile": "ollama", "gpu_util_pct": 95, "temp_c": 75.0},
|
||||
])
|
||||
assert ts.profile_comparison(days=1)[0]["avg_temp_c"] == pytest.approx(75.0)
|
||||
|
||||
|
||||
def test_swap_stats_groups_by_event_type_and_cache_status(temp_db):
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 30, "event_type": "LLM Model Switch", "target": "llama3:8b",
|
||||
"duration_ms": 5000, "cache_status": "RAM Cache Hit ⚡", "load_gbps": 2.6,
|
||||
"tokens_per_sec": 60.0},
|
||||
{"ts": now - 20, "event_type": "LLM Model Switch", "target": "llama3:8b",
|
||||
"duration_ms": 7000, "cache_status": "Cold Disk Load 💾", "load_gbps": 0.38,
|
||||
"tokens_per_sec": 58.0},
|
||||
{"ts": now - 10, "event_type": "VRAM Yield", "duration_ms": 900,
|
||||
"yield_confirm_ms": 850},
|
||||
])
|
||||
stats = ts.swap_stats(days=1)
|
||||
by_type = {r["event_type"]: r for r in stats["by_type"]}
|
||||
assert by_type["LLM Model Switch"]["n"] == 2
|
||||
assert by_type["LLM Model Switch"]["avg_ms"] == pytest.approx(6000.0)
|
||||
assert by_type["LLM Model Switch"]["min_ms"] == pytest.approx(5000.0)
|
||||
assert by_type["VRAM Yield"]["avg_confirm_ms"] == pytest.approx(850.0)
|
||||
assert {r["cache_status"] for r in stats["by_cache_status"]} == {
|
||||
"RAM Cache Hit ⚡", "Cold Disk Load 💾"}
|
||||
assert stats["by_model"][0]["model"] == "llama3:8b"
|
||||
assert stats["by_model"][0]["loads"] == 2
|
||||
|
||||
|
||||
def test_model_usage_ranking_scores_recent_use_higher(temp_db):
|
||||
"""Recency/frequency drives the RAM warm budget: given equal load counts, the model
|
||||
used more recently must rank first (half-life ~24 h)."""
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 3600, "event_type": "LLM Model Switch", "target": "recent:latest"},
|
||||
{"ts": now - 3600, "event_type": "LLM Model Switch", "target": "recent:latest"},
|
||||
{"ts": now - 20 * 86400, "event_type": "LLM Model Switch", "target": "stale:latest"},
|
||||
{"ts": now - 20 * 86400, "event_type": "LLM Model Switch", "target": "stale:latest"},
|
||||
])
|
||||
ranking = ts.model_usage_ranking(days=30)
|
||||
assert [r["model"] for r in ranking] == ["recent:latest", "stale:latest"]
|
||||
assert ranking[0]["score"] > ranking[1]["score"]
|
||||
assert ranking[0]["loads"] == 2
|
||||
assert ranking[0]["age_hours"] == pytest.approx(1.0, abs=0.05)
|
||||
|
||||
|
||||
def test_model_usage_ranking_scores_frequent_use_higher_at_equal_recency(temp_db):
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 3600, "event_type": "LLM Model Switch", "target": "often:latest"},
|
||||
{"ts": now - 3601, "event_type": "LLM Model Switch", "target": "often:latest"},
|
||||
{"ts": now - 3602, "event_type": "LLM Model Switch", "target": "often:latest"},
|
||||
{"ts": now - 3600, "event_type": "LLM Model Switch", "target": "once:latest"},
|
||||
])
|
||||
ranking = ts.model_usage_ranking(days=30)
|
||||
assert ranking[0]["model"] == "often:latest"
|
||||
|
||||
|
||||
def test_model_usage_ranking_counts_warms_as_well_as_switches(temp_db):
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 60, "event_type": "Model Warm", "target": "warmed:latest"},
|
||||
{"ts": now - 60, "event_type": "Something Else", "target": "ignored:latest"},
|
||||
])
|
||||
assert [r["model"] for r in ts.model_usage_ranking(days=1)] == ["warmed:latest"]
|
||||
|
||||
|
||||
def test_timeseries_buckets_samples_by_width(temp_db):
|
||||
"""Six hours into six buckets is one bucket per hour; samples inside an hour collapse
|
||||
into a single averaged point."""
|
||||
now = time.time()
|
||||
base = now - 5.5 * 3600
|
||||
_seed(temp_db, telemetry=[
|
||||
{"ts": base + 60, "temp_c": 60.0, "gpu_util_pct": 50},
|
||||
{"ts": base + 120, "temp_c": 70.0, "gpu_util_pct": 70},
|
||||
{"ts": base + 3700, "temp_c": 80.0, "gpu_util_pct": 90},
|
||||
])
|
||||
out = ts.timeseries(hours=6, buckets=6)
|
||||
assert len(out) == 2
|
||||
assert out[0]["temp_c"] == pytest.approx(65.0) # average of 60 and 70
|
||||
assert out[1]["temp_c"] == pytest.approx(80.0)
|
||||
assert out[0]["bucket_ts"] < out[1]["bucket_ts"]
|
||||
|
||||
|
||||
def test_timeseries_excludes_samples_older_than_the_window(temp_db):
|
||||
now = time.time()
|
||||
_seed(temp_db, telemetry=[
|
||||
{"ts": now - 48 * 3600, "temp_c": 99.0},
|
||||
{"ts": now - 60, "temp_c": 60.0},
|
||||
])
|
||||
out = ts.timeseries(hours=6, buckets=240)
|
||||
assert len(out) == 1 and out[0]["temp_c"] == pytest.approx(60.0)
|
||||
|
||||
|
||||
def test_db_info_reports_row_counts_and_coverage(temp_db):
|
||||
now = time.time()
|
||||
_seed(temp_db,
|
||||
events=[{"ts": now - 10, "event_type": "LLM Model Switch", "target": "m"}],
|
||||
telemetry=[{"ts": now - 7200, "temp_c": 60.0}, {"ts": now, "temp_c": 61.0}])
|
||||
info = ts.db_info()
|
||||
assert info["exists"] is True
|
||||
assert info["events_rows"] == 1
|
||||
assert info["telemetry_rows"] == 2
|
||||
assert info["coverage_hours"] == pytest.approx(2.0, abs=0.01)
|
||||
405
tests/test_tenants.py
Normal file
405
tests/test_tenants.py
Normal file
@@ -0,0 +1,405 @@
|
||||
"""Tests for the GPU tenant registry.
|
||||
|
||||
The point of this service is fast handoff of one GPU between applications, and it should
|
||||
work for any application -- not only the two it grew up around. Their names had ended up
|
||||
compiled into process matching, VRAM attribution, busy detection and release calls alike.
|
||||
These tests pin the properties that make the registry generic: adding an application is
|
||||
configuration, and nothing in the arbitration logic knows a particular name.
|
||||
"""
|
||||
import asyncio
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
import tenants as T
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def cfg(tmp_path, monkeypatch):
|
||||
path = tmp_path / "tenants.json"
|
||||
monkeypatch.setattr(T, "CONFIG_PATH", str(path))
|
||||
T._cache.update({"ts": 0.0, "tenants": None, "mtime": None})
|
||||
return path
|
||||
|
||||
|
||||
class TestProcessMatching:
|
||||
def test_matches_by_process_name(self):
|
||||
m = T.ProcessMatch(names=["ollama"])
|
||||
assert m.matches("ollama", "/usr/bin/ollama serve")
|
||||
assert not m.matches("python", "main.py")
|
||||
|
||||
def test_matches_by_cmdline_substring(self):
|
||||
m = T.ProcessMatch(cmdline=["llama-server"])
|
||||
assert m.matches("python", "/usr/local/lib/ollama/llama-server --model x")
|
||||
|
||||
def test_matches_by_cmdline_suffix(self):
|
||||
# ComfyUI is a bare `python main.py`, with nothing else distinguishing it.
|
||||
m = T.ProcessMatch(cmdline_endswith=["main.py"])
|
||||
assert m.matches("python", "/opt/ComfyUI/venv/bin/python main.py")
|
||||
assert not m.matches("python", "/opt/other/main.py --serve")
|
||||
|
||||
def test_matching_is_case_insensitive(self):
|
||||
assert T.ProcessMatch(names=["Xorg"]).matches("XORG", "")
|
||||
|
||||
|
||||
class TestDefaultsPreserveExistingBehaviour:
|
||||
"""The shipped defaults must classify exactly as the hardcoded version did."""
|
||||
|
||||
@pytest.mark.parametrize("pname,cmdline,expected", [
|
||||
("llama-server", "/usr/local/lib/ollama/llama-server --model x", "ollama"),
|
||||
("ollama", "/usr/bin/ollama serve", "ollama"),
|
||||
("python", "/home/u/ComfyUI/venv/bin/python main.py --listen", "comfyui"),
|
||||
("gnome-shell", "/usr/bin/gnome-shell --mode=ubuntu", "desktop"),
|
||||
("Xorg", "/usr/lib/xorg/Xorg :8", "desktop"),
|
||||
("python", "/home/u/robopest-venv/bin/python /home/u/stt_relay.py", "unmanaged"),
|
||||
("trainer", "/opt/ml/bin/trainer --epochs 3", "unmanaged"),
|
||||
])
|
||||
def test_classification(self, cfg, pname, cmdline, expected):
|
||||
assert T.classify_process(pname, cmdline) == expected
|
||||
|
||||
def test_unknown_process_is_unmanaged_not_silently_owned(self, cfg):
|
||||
# Misattributing a third party's VRAM to a tenant would make this service
|
||||
# promise headroom it cannot deliver.
|
||||
assert T.classify_process("weird", "/opt/x/weird --run") == "unmanaged"
|
||||
|
||||
|
||||
class TestAddingAnApplicationIsConfiguration:
|
||||
def test_a_new_tenant_is_recognised_without_code_changes(self, cfg):
|
||||
cfg.write_text(json.dumps(T.DEFAULT_TENANTS + [{
|
||||
"name": "trainer",
|
||||
"kind": "other",
|
||||
"priority": 80,
|
||||
"match": {"cmdline": ["train.py"]},
|
||||
"release": {"type": "http_post", "url": "http://localhost:9999/release"},
|
||||
}]))
|
||||
assert T.classify_process("python", "/opt/ml/train.py --epochs 3") == "trainer"
|
||||
t = T.get_tenant("trainer")
|
||||
assert t.priority == 80 and t.reclaimable
|
||||
|
||||
def test_first_run_writes_the_defaults(self, cfg):
|
||||
assert not cfg.exists()
|
||||
T.load_tenants(force=True)
|
||||
assert cfg.exists()
|
||||
assert {t["name"] for t in json.loads(cfg.read_text())} == {
|
||||
"ollama", "comfyui", "desktop"}
|
||||
|
||||
def test_a_malformed_entry_is_skipped_not_fatal(self, cfg):
|
||||
cfg.write_text(json.dumps([{"name": "ok", "match": {"names": ["a"]}},
|
||||
{"no_name": True}]))
|
||||
names = [t.name for t in T.load_tenants(force=True)]
|
||||
assert names == ["ok"]
|
||||
|
||||
def test_corrupt_config_falls_back_to_defaults(self, cfg):
|
||||
cfg.write_text("{ not json")
|
||||
assert {t.name for t in T.load_tenants(force=True)} >= {"ollama", "comfyui"}
|
||||
|
||||
|
||||
class TestReclaimability:
|
||||
def test_a_tenant_with_no_release_strategy_is_not_reclaimable(self, cfg):
|
||||
t = T.GpuTenant(name="x", release=T.ReleaseStrategy(type="none"))
|
||||
assert t.reclaimable is False
|
||||
|
||||
def test_release_refuses_rather_than_reporting_success(self, cfg):
|
||||
t = T.GpuTenant(name="x", release=T.ReleaseStrategy(type="none"))
|
||||
res = asyncio.run(T.release_vram(t))
|
||||
assert res["success"] is False and res["released"] is False
|
||||
assert "no way to release" in res["reason"]
|
||||
|
||||
def test_per_model_release_with_nothing_loaded_is_a_no_op(self, cfg):
|
||||
t = T.GpuTenant(name="ollama", release=T.ReleaseStrategy(
|
||||
type="http_post", url="http://x/api", per_model=True))
|
||||
res = asyncio.run(T.release_vram(t, models=[]))
|
||||
assert res["success"] is True and res["released"] is False
|
||||
|
||||
|
||||
class TestBusyProbe:
|
||||
def _probe(self, monkeypatch, payload, status=200):
|
||||
class _R:
|
||||
status_code = status
|
||||
def json(self_inner): return payload
|
||||
class _C:
|
||||
async def __aenter__(self): return self
|
||||
async def __aexit__(self, *a): return False
|
||||
async def get(self, url): return _R()
|
||||
monkeypatch.setattr(T.httpx, "AsyncClient", lambda **k: _C())
|
||||
|
||||
def test_empty_queue_is_not_busy(self, monkeypatch):
|
||||
self._probe(monkeypatch, {"queue_running": [], "queue_pending": []})
|
||||
t = T.GpuTenant(name="c", busy=T.BusyProbe(
|
||||
type="http_count", url="http://x/queue",
|
||||
count_keys=["queue_running", "queue_pending"]))
|
||||
assert asyncio.run(T.probe_busy(t))["busy"] is False
|
||||
|
||||
def test_queued_work_while_holding_no_vram_is_flagged_below_floor(self, monkeypatch):
|
||||
# ComfyUI leaves dead jobs in queue_running; only its VRAM reveals that nothing
|
||||
# is loaded.
|
||||
self._probe(monkeypatch, {"queue_running": [[1, "abc"]], "queue_pending": []})
|
||||
t = T.GpuTenant(name="c", busy=T.BusyProbe(
|
||||
type="http_count", url="http://x/queue", count_keys=["queue_running"],
|
||||
vram_floor_gb=1.5))
|
||||
res = asyncio.run(T.probe_busy(t, vram_gb=0.56))
|
||||
assert res["busy"] is True and res.get("below_floor") is True
|
||||
|
||||
def test_queued_work_with_a_checkpoint_loaded_is_plainly_busy(self, monkeypatch):
|
||||
self._probe(monkeypatch, {"queue_running": [[1, "abc"]], "queue_pending": []})
|
||||
t = T.GpuTenant(name="c", busy=T.BusyProbe(
|
||||
type="http_count", url="http://x/queue", count_keys=["queue_running"],
|
||||
vram_floor_gb=1.5))
|
||||
res = asyncio.run(T.probe_busy(t, vram_gb=6.8))
|
||||
assert res["busy"] is True and not res.get("below_floor")
|
||||
|
||||
def test_vram_probe_needs_no_http_endpoint(self):
|
||||
# An application with no API can still be observed by what it holds.
|
||||
t = T.GpuTenant(name="x", busy=T.BusyProbe(type="vram", vram_busy_gb=1.0))
|
||||
assert asyncio.run(T.probe_busy(t, vram_gb=2.0))["busy"] is True
|
||||
assert asyncio.run(T.probe_busy(t, vram_gb=0.5))["busy"] is False
|
||||
|
||||
def test_an_unreachable_probe_reports_not_busy_rather_than_raising(self, monkeypatch):
|
||||
class _C:
|
||||
async def __aenter__(self): return self
|
||||
async def __aexit__(self, *a): return False
|
||||
async def get(self, url): raise ConnectionError("refused")
|
||||
monkeypatch.setattr(T.httpx, "AsyncClient", lambda **k: _C())
|
||||
t = T.GpuTenant(name="c", busy=T.BusyProbe(type="http_count", url="http://x",
|
||||
count_keys=["q"]))
|
||||
res = asyncio.run(T.probe_busy(t))
|
||||
assert res["busy"] is False and "failed" in res["reason"]
|
||||
|
||||
|
||||
class TestPriority:
|
||||
def test_describe_orders_by_priority(self, cfg):
|
||||
rows = T.describe()
|
||||
prios = [r["priority"] for r in rows]
|
||||
assert prios == sorted(prios, reverse=True)
|
||||
assert all("reclaimable" in r for r in rows)
|
||||
|
||||
|
||||
class TestReleasePlanning:
|
||||
"""Deciding who gives up VRAM, generically over any number of applications.
|
||||
|
||||
The two-application version was a pair of hardcoded rules -- yield Ollama when
|
||||
ComfyUI is busy, purge ComfyUI when Ollama is starved -- which could not express a
|
||||
third participant at all.
|
||||
"""
|
||||
|
||||
def _state(self, **overrides):
|
||||
base = [
|
||||
{"name": "desktop", "priority": 90, "vram_gb": 0.01, "busy": False,
|
||||
"reclaimable": False},
|
||||
{"name": "stt-relay", "priority": 70, "vram_gb": 0.8, "busy": False,
|
||||
"reclaimable": False},
|
||||
# Shipped priorities: diffusion outranks the LLM, whose weights reload
|
||||
# from page cache in seconds.
|
||||
{"name": "comfyui", "priority": 60, "vram_gb": 7.0, "busy": False,
|
||||
"reclaimable": True},
|
||||
{"name": "ollama", "priority": 50, "vram_gb": 0.0, "busy": True,
|
||||
"reclaimable": True},
|
||||
]
|
||||
for s in base:
|
||||
s.update(overrides.get(s["name"], {}))
|
||||
return base
|
||||
|
||||
def test_a_tenant_already_holding_what_it_needs_is_not_starved(self):
|
||||
# A busy GPU has little free by definition. Comparing free VRAM alone flagged a
|
||||
# tenant working fine on 13 GB as demanding, which would have caused pointless
|
||||
# releases from everything else.
|
||||
state = self._state(ollama={"vram_gb": 13.0})
|
||||
plan = T.plan_release("ollama", state, free_gb=1.5, needed_gb=4.0)
|
||||
assert plan["release"] == []
|
||||
assert "already free" in plan["reason"]
|
||||
|
||||
def test_starved_tenant_reclaims_from_the_idle_one_below_it(self):
|
||||
plan = T.plan_release("ollama", self._state(), free_gb=1.5, needed_gb=14.9)
|
||||
assert plan["release"] == ["comfyui"]
|
||||
|
||||
def test_a_busy_tenant_ranking_above_the_demander_is_not_a_victim(self):
|
||||
# comfyui outranks ollama, so ollama may not interrupt it.
|
||||
state = self._state(comfyui={"busy": True})
|
||||
plan = T.plan_release("ollama", state, free_gb=1.5, needed_gb=14.9)
|
||||
assert plan["release"] == []
|
||||
assert any(b["name"] == "comfyui" and "busy" in b["why"]
|
||||
for b in plan["blockers"])
|
||||
|
||||
def test_a_higher_priority_demander_preempts_busy_lower_priority_work(self):
|
||||
"""The measured regression that made this rule necessary.
|
||||
|
||||
Refusing to touch anything busy looks safe and is not. With the LLM protected as
|
||||
"busy", a diffusion job ran 46 s instead of 3 s, squeezed into 1.6 GB, because
|
||||
the LLM reloaded immediately after yielding and was then untouchable. Preempting
|
||||
a lower-priority tenant is safe because releasing is asynchronous: an Ollama
|
||||
unload queues behind its running request rather than killing it.
|
||||
"""
|
||||
state = [
|
||||
{"name": "comfyui", "priority": 60, "vram_gb": 1.65, "busy": True,
|
||||
"reclaimable": True},
|
||||
{"name": "ollama", "priority": 50, "vram_gb": 13.03, "busy": True,
|
||||
"reclaimable": True},
|
||||
]
|
||||
plan = T.plan_release("comfyui", state, free_gb=0.28, needed_gb=6.0)
|
||||
assert plan["release"] == ["ollama"]
|
||||
|
||||
def test_an_idle_tenant_is_preferred_over_preempting_a_busy_one(self):
|
||||
state = [
|
||||
{"name": "d", "priority": 60, "vram_gb": 0.0, "busy": True,
|
||||
"reclaimable": True},
|
||||
{"name": "busy_low", "priority": 10, "vram_gb": 8.0, "busy": True,
|
||||
"reclaimable": True},
|
||||
{"name": "idle_high", "priority": 90, "vram_gb": 8.0, "busy": False,
|
||||
"reclaimable": True},
|
||||
]
|
||||
plan = T.plan_release("d", state, free_gb=0.0, needed_gb=8.0)
|
||||
assert plan["release"] == ["idle_high"]
|
||||
|
||||
def test_unreclaimable_tenants_are_named_as_blockers_not_ignored(self):
|
||||
# The user needs to know a third-party process is what stands in the way.
|
||||
plan = T.plan_release("ollama", self._state(), free_gb=0.0, needed_gb=15.5)
|
||||
blockers = {b["name"]: b["why"] for b in plan["blockers"]}
|
||||
assert blockers["stt-relay"] == "declares no release mechanism"
|
||||
assert plan["possible"] is False
|
||||
|
||||
def test_an_idle_tenant_yields_even_if_it_outranks_the_demander(self):
|
||||
"""Priority orders who is asked first; it does not protect idle memory.
|
||||
|
||||
Filtering candidates by priority broke both directions in turn: with the LLM
|
||||
ranked above diffusion, ComfyUI could never preempt Ollama -- the service's
|
||||
central behaviour -- and once the ranks were swapped, a starved Ollama could no
|
||||
longer reclaim from an idle ComfyUI. An idle tenant is not using its VRAM, so
|
||||
outranking the demander is not a reason to keep it.
|
||||
"""
|
||||
state = self._state(comfyui={"priority": 99, "busy": False})
|
||||
plan = T.plan_release("ollama", state, free_gb=1.0, needed_gb=14.9)
|
||||
assert plan["release"] == ["comfyui"]
|
||||
|
||||
def test_priority_decides_who_is_asked_first(self):
|
||||
state = [
|
||||
{"name": "demander", "priority": 50, "vram_gb": 0.0, "busy": True,
|
||||
"reclaimable": True},
|
||||
{"name": "high", "priority": 90, "vram_gb": 4.0, "busy": False,
|
||||
"reclaimable": True},
|
||||
{"name": "low", "priority": 10, "vram_gb": 4.0, "busy": False,
|
||||
"reclaimable": True},
|
||||
]
|
||||
plan = T.plan_release("demander", state, free_gb=0.0, needed_gb=5.0)
|
||||
# The lowest-priority idle tenant gives up memory first.
|
||||
assert plan["release"][0] == "low"
|
||||
|
||||
def test_peers_cannot_interrupt_each_other(self):
|
||||
# Equal priority is never preempted, so two tenants at the same rank cannot
|
||||
# fight over the card.
|
||||
state = [
|
||||
{"name": "a", "priority": 50, "vram_gb": 0.0, "busy": True,
|
||||
"reclaimable": True},
|
||||
{"name": "b", "priority": 50, "vram_gb": 8.0, "busy": True,
|
||||
"reclaimable": True},
|
||||
]
|
||||
plan = T.plan_release("a", state, free_gb=0.0, needed_gb=8.0)
|
||||
assert plan["release"] == []
|
||||
assert "busy" in plan["blockers"][0]["why"]
|
||||
|
||||
def test_lowest_priority_is_released_first(self):
|
||||
state = self._state() + [
|
||||
{"name": "batch", "priority": 10, "vram_gb": 3.0, "busy": False,
|
||||
"reclaimable": True}]
|
||||
plan = T.plan_release("ollama", state, free_gb=0.0, needed_gb=5.0)
|
||||
assert plan["release"][0] == "batch"
|
||||
|
||||
def test_releases_only_as_many_tenants_as_needed(self):
|
||||
state = self._state() + [
|
||||
{"name": "batch", "priority": 10, "vram_gb": 9.0, "busy": False,
|
||||
"reclaimable": True}]
|
||||
plan = T.plan_release("ollama", state, free_gb=0.0, needed_gb=8.0)
|
||||
assert plan["release"] == ["batch"] # 9 GB covers it; comfyui is left alone
|
||||
|
||||
def test_three_applications_can_all_participate(self):
|
||||
# The property the hardcoded pair of rules could not express.
|
||||
state = [
|
||||
{"name": "llm", "priority": 60, "vram_gb": 0.0, "busy": True,
|
||||
"reclaimable": True},
|
||||
{"name": "diffusion", "priority": 50, "vram_gb": 4.0, "busy": False,
|
||||
"reclaimable": True},
|
||||
{"name": "trainer", "priority": 40, "vram_gb": 5.0, "busy": False,
|
||||
"reclaimable": True},
|
||||
]
|
||||
plan = T.plan_release("llm", state, free_gb=0.0, needed_gb=9.0)
|
||||
assert set(plan["release"]) == {"trainer", "diffusion"}
|
||||
assert plan["possible"] is True
|
||||
|
||||
def test_unknown_tenant_is_rejected_cleanly(self):
|
||||
plan = T.plan_release("nope", self._state(), free_gb=0.0, needed_gb=1.0)
|
||||
assert plan["possible"] is False and plan["release"] == []
|
||||
|
||||
|
||||
class TestConfigUpgrade:
|
||||
def test_fields_added_later_are_merged_into_an_existing_config(self, cfg):
|
||||
# A config written before needs_vram_gb existed must not silently lose the
|
||||
# behaviour that field controls.
|
||||
cfg.write_text(json.dumps([{
|
||||
"name": "ollama",
|
||||
"match": {"names": ["ollama"]},
|
||||
}]))
|
||||
t = T.get_tenant("ollama")
|
||||
assert t.needs_vram_gb > 0
|
||||
assert t.release.type == "http_post"
|
||||
|
||||
def test_explicit_user_values_still_win_over_defaults(self, cfg):
|
||||
cfg.write_text(json.dumps([{
|
||||
"name": "ollama", "priority": 5, "needs_vram_gb": 99.0,
|
||||
"match": {"names": ["ollama"]},
|
||||
}]))
|
||||
t = T.get_tenant("ollama")
|
||||
assert t.priority == 5 and t.needs_vram_gb == 99.0
|
||||
|
||||
|
||||
class TestVramFloor:
|
||||
"""VRAM that survives a release must not be promised to anyone else.
|
||||
|
||||
ComfyUI keeps its CUDA context for as long as the process lives, so a purge does not
|
||||
return everything it holds. Ignoring that made plan_release report it would free
|
||||
0.37 GB against a 0.33 GB shortfall; the job was cleared to run and the memory never
|
||||
arrived, so it waited two minutes and then failed.
|
||||
"""
|
||||
|
||||
def test_only_memory_above_the_floor_counts_as_freeable(self):
|
||||
state = [
|
||||
{"name": "llm", "priority": 50, "vram_gb": 0.0, "busy": True,
|
||||
"reclaimable": True, "vram_floor_gb": 0.0},
|
||||
{"name": "diffusion", "priority": 60, "vram_gb": 0.44, "busy": False,
|
||||
"reclaimable": True, "vram_floor_gb": 0.45},
|
||||
]
|
||||
plan = T.plan_release("llm", state, free_gb=14.6, needed_gb=14.93)
|
||||
assert plan["possible"] is False
|
||||
assert plan["release"] == []
|
||||
|
||||
def test_a_loaded_checkpoint_is_still_freeable_above_its_floor(self):
|
||||
state = [
|
||||
{"name": "llm", "priority": 50, "vram_gb": 0.0, "busy": True,
|
||||
"reclaimable": True, "vram_floor_gb": 0.0},
|
||||
{"name": "diffusion", "priority": 60, "vram_gb": 7.0, "busy": False,
|
||||
"reclaimable": True, "vram_floor_gb": 0.45},
|
||||
]
|
||||
plan = T.plan_release("llm", state, free_gb=7.9, needed_gb=14.0)
|
||||
assert plan["release"] == ["diffusion"]
|
||||
# 7.0 held minus a 0.45 floor.
|
||||
assert abs(plan["would_free_gb"] - 6.55) < 0.01
|
||||
|
||||
def test_a_tenant_at_its_floor_is_not_even_listed_for_release(self):
|
||||
state = [
|
||||
{"name": "llm", "priority": 50, "vram_gb": 0.0, "busy": True,
|
||||
"reclaimable": True, "vram_floor_gb": 0.0},
|
||||
{"name": "at_floor", "priority": 10, "vram_gb": 0.3, "busy": False,
|
||||
"reclaimable": True, "vram_floor_gb": 0.45},
|
||||
{"name": "has_room", "priority": 20, "vram_gb": 5.0, "busy": False,
|
||||
"reclaimable": True, "vram_floor_gb": 0.0},
|
||||
]
|
||||
plan = T.plan_release("llm", state, free_gb=0.0, needed_gb=4.0)
|
||||
assert plan["release"] == ["has_room"]
|
||||
|
||||
def test_default_floor_is_zero_so_existing_configs_are_unchanged(self):
|
||||
state = [
|
||||
{"name": "a", "priority": 50, "vram_gb": 0.0, "busy": True,
|
||||
"reclaimable": True},
|
||||
{"name": "b", "priority": 40, "vram_gb": 5.0, "busy": False,
|
||||
"reclaimable": True},
|
||||
]
|
||||
plan = T.plan_release("a", state, free_gb=0.0, needed_gb=5.0)
|
||||
assert plan["release"] == ["b"] and plan["possible"] is True
|
||||
264
tests/test_thermal_governor.py
Normal file
264
tests/test_thermal_governor.py
Normal file
@@ -0,0 +1,264 @@
|
||||
"""thermal_governor: the derate state machine.
|
||||
|
||||
SAFETY: _step() spawns a thread that calls overclock_manager.apply_profile. The autouse
|
||||
`no_gpu_mutation` fixture in conftest.py replaces that (and load_profiles is stubbed per
|
||||
test), so escalation here can never reach the card.
|
||||
"""
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
import overclock_manager
|
||||
import thermal_governor as tg
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def gov(monkeypatch):
|
||||
"""A fresh governor with a stubbed profile store — never the module singleton."""
|
||||
monkeypatch.setattr(overclock_manager, "load_profiles", lambda: {
|
||||
"ollama": {"core_offset_mhz": 100, "mem_offset_mhz": 500, "fan_speed_pct": 70,
|
||||
"lock_core_min": 2500, "lock_core_max": 2800},
|
||||
"stock": {"core_offset_mhz": 0, "mem_offset_mhz": 0, "fan_speed_pct": 0},
|
||||
})
|
||||
monkeypatch.setattr(overclock_manager, "ACTIVE_PROFILE", "ollama")
|
||||
return tg.ThermalGovernor()
|
||||
|
||||
|
||||
def _sample(temp=60.0, reasons=None, available=True):
|
||||
return {"available": available, "temperature_c": temp, "throttle_reasons": reasons or []}
|
||||
|
||||
|
||||
def _feed(gov, n, **kw):
|
||||
for _ in range(n):
|
||||
gov.observe(_sample(**kw), active_profile="ollama")
|
||||
|
||||
|
||||
def _clear_cooldown(gov):
|
||||
"""The REAPPLY_COOLDOWN_S gate is time-based; wind the clock back rather than sleep."""
|
||||
gov.last_change = time.time() - tg.REAPPLY_COOLDOWN_S - 1.0
|
||||
|
||||
|
||||
# --------------------------------------------------------------- escalation hysteresis
|
||||
|
||||
def test_does_not_escalate_before_hot_samples_consecutive_bad_readings(gov):
|
||||
"""One hot spike during a diffusion step must not derate the card."""
|
||||
_feed(gov, tg.HOT_SAMPLES - 1, temp=90.0)
|
||||
assert gov.level == 0
|
||||
assert gov.hot_streak == tg.HOT_SAMPLES - 1
|
||||
|
||||
|
||||
def test_escalates_on_exactly_hot_samples_consecutive_bad_readings(gov, no_gpu_mutation):
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=90.0)
|
||||
assert gov.level == 1
|
||||
assert gov.hot_streak == 0 # streaks reset after a step
|
||||
# Actuation runs in a daemon thread; it must reach apply_profile (the stub) with the
|
||||
# derate overrides rather than blocking the 1 Hz sampler.
|
||||
deadline = time.time() + 2.0
|
||||
while not no_gpu_mutation["apply_profile"] and time.time() < deadline:
|
||||
time.sleep(0.01)
|
||||
assert no_gpu_mutation["apply_profile"], "escalation never actuated"
|
||||
name, overrides = no_gpu_mutation["apply_profile"][0]
|
||||
assert name == "ollama"
|
||||
assert overrides["core_offset_mhz"] == int(100 * tg.DERATE_LADDER[1]["offset_scale"])
|
||||
|
||||
|
||||
def test_alternating_hot_and_cool_samples_never_escalate(gov):
|
||||
"""The core hysteresis property: a card oscillating around the threshold must not
|
||||
thrash the profile. A single good sample resets the hot streak."""
|
||||
for _ in range(50):
|
||||
gov.observe(_sample(temp=90.0), active_profile="ollama")
|
||||
gov.observe(_sample(temp=60.0), active_profile="ollama")
|
||||
assert gov.level == 0
|
||||
|
||||
|
||||
def test_temperature_between_recover_and_escalate_resets_both_streaks(gov):
|
||||
"""The band between TEMP_RECOVER_C and TEMP_ESCALATE_C is neither hot nor cool; it
|
||||
must not accumulate credit in either direction."""
|
||||
_feed(gov, tg.HOT_SAMPLES - 1, temp=90.0)
|
||||
gov.observe(_sample(temp=78.0), active_profile="ollama")
|
||||
assert gov.hot_streak == 0 and gov.cool_streak == 0
|
||||
assert gov.level == 0
|
||||
|
||||
|
||||
def test_hard_throttle_counts_as_hot_even_when_cool(gov):
|
||||
"""A hardware slowdown means the card is protecting itself; temperature alone is not
|
||||
the whole signal."""
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=55.0, reasons=["hw_thermal_slowdown"])
|
||||
assert gov.level == 1
|
||||
|
||||
|
||||
def test_soft_throttle_reasons_do_not_escalate(gov):
|
||||
"""Hitting a power or utilisation cap is normal operation, not distress."""
|
||||
_feed(gov, tg.HOT_SAMPLES * 3, temp=55.0, reasons=["sw_power_cap", "gpu_idle"])
|
||||
assert gov.level == 0
|
||||
|
||||
|
||||
def test_escalation_stops_at_the_bottom_of_the_ladder(gov):
|
||||
"""Level must never index past DERATE_LADDER."""
|
||||
for _ in range(len(tg.DERATE_LADDER) + 3):
|
||||
_clear_cooldown(gov)
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
assert gov.level == len(tg.DERATE_LADDER) - 1
|
||||
|
||||
|
||||
def test_cooldown_gate_blocks_a_second_step_immediately_after_the_first(gov):
|
||||
"""REAPPLY_COOLDOWN_S stops the governor from walking the whole ladder in one second
|
||||
while the card is still responding to the previous change."""
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
assert gov.level == 1
|
||||
_feed(gov, tg.HOT_SAMPLES * 2, temp=95.0)
|
||||
assert gov.level == 1
|
||||
|
||||
|
||||
def test_cooldown_gate_releases_after_the_window(gov):
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
_clear_cooldown(gov)
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
assert gov.level == 2
|
||||
|
||||
|
||||
# --------------------------------------------------------------- recovery
|
||||
|
||||
def test_recovery_needs_cool_samples_consecutive_good_readings(gov):
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
assert gov.level == 1
|
||||
_clear_cooldown(gov)
|
||||
_feed(gov, tg.COOL_SAMPLES - 1, temp=60.0)
|
||||
assert gov.level == 1, "recovered too early"
|
||||
gov.observe(_sample(temp=60.0), active_profile="ollama")
|
||||
assert gov.level == 0
|
||||
|
||||
|
||||
def test_recovery_is_slower_than_escalation():
|
||||
"""Deliberate asymmetry: react fast to heat, give the clocks back slowly."""
|
||||
assert tg.COOL_SAMPLES > tg.HOT_SAMPLES
|
||||
|
||||
|
||||
def test_recovery_stops_at_level_zero(gov):
|
||||
_clear_cooldown(gov)
|
||||
_feed(gov, tg.COOL_SAMPLES * 2, temp=50.0)
|
||||
assert gov.level == 0
|
||||
|
||||
|
||||
def test_a_hard_throttle_blocks_recovery_even_at_a_cool_temperature(gov):
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
_clear_cooldown(gov)
|
||||
_feed(gov, tg.COOL_SAMPLES * 2, temp=50.0, reasons=["hw_power_brake_slowdown"])
|
||||
assert gov.level >= 1
|
||||
|
||||
|
||||
# --------------------------------------------------------------- ignored input
|
||||
|
||||
def test_disabled_governor_ignores_samples(gov):
|
||||
gov.enabled = False
|
||||
_feed(gov, tg.HOT_SAMPLES * 3, temp=99.0)
|
||||
assert gov.level == 0
|
||||
|
||||
|
||||
def test_unavailable_gpu_sample_is_ignored(gov):
|
||||
"""A failed NVML read reports available=False with no temperature; treating that as
|
||||
0 C would count as a cool sample and hand the clocks back."""
|
||||
for _ in range(tg.COOL_SAMPLES * 2):
|
||||
gov.observe({"available": False}, active_profile="ollama")
|
||||
assert gov.level == 0 and gov.cool_streak == 0
|
||||
|
||||
|
||||
# --------------------------------------------------------------- overrides
|
||||
|
||||
def test_overrides_are_empty_at_level_zero(gov):
|
||||
assert gov.overrides_for("ollama") == {}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("level", [1, 2, 3])
|
||||
def test_overrides_scale_offsets_by_the_ladder(gov, level):
|
||||
gov.level = level
|
||||
scale = tg.DERATE_LADDER[level]["offset_scale"]
|
||||
ov = gov.overrides_for("ollama")
|
||||
assert ov["core_offset_mhz"] == int(100 * scale)
|
||||
assert ov["mem_offset_mhz"] == int(500 * scale)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("level", [1, 2, 3])
|
||||
def test_overrides_raise_the_fan_floor_and_never_lower_it(gov, level):
|
||||
"""The floor is a floor: a profile already running fans harder than the ladder asks
|
||||
keeps its own setting."""
|
||||
gov.level = level
|
||||
floor = tg.DERATE_LADDER[level]["fan_floor"]
|
||||
ov = gov.overrides_for("ollama")
|
||||
assert ov["fan_mode"] == "manual"
|
||||
assert ov["fan_speed_pct"] == max(70, floor) # profile fan_speed_pct is 70
|
||||
assert ov["fan_speed_pct"] >= floor
|
||||
|
||||
|
||||
def test_overrides_release_the_core_clock_lock_from_level_two(gov):
|
||||
"""Pinning the core clock high while the card is already backing off fights the
|
||||
hardware's own protection."""
|
||||
gov.level = 1
|
||||
assert "lock_core_max" not in gov.overrides_for("ollama")
|
||||
for level in (2, 3):
|
||||
gov.level = level
|
||||
ov = gov.overrides_for("ollama")
|
||||
assert ov["lock_core_min"] == 0 and ov["lock_core_max"] == 0
|
||||
|
||||
|
||||
def test_top_of_ladder_is_stock_clocks_and_maximum_fans(gov):
|
||||
gov.level = len(tg.DERATE_LADDER) - 1
|
||||
ov = gov.overrides_for("ollama")
|
||||
assert ov["core_offset_mhz"] == 0 and ov["mem_offset_mhz"] == 0
|
||||
assert ov["fan_speed_pct"] == 100
|
||||
|
||||
|
||||
def test_overrides_for_unknown_profile_do_not_raise(gov):
|
||||
"""The active profile can be one the store does not know; a missing config must
|
||||
derate to zero offsets rather than blow up in the sampler thread."""
|
||||
gov.level = 2
|
||||
ov = gov.overrides_for("does-not-exist")
|
||||
assert ov["core_offset_mhz"] == 0 and ov["mem_offset_mhz"] == 0
|
||||
|
||||
|
||||
# --------------------------------------------------------------- ladder invariants
|
||||
|
||||
def test_ladder_is_monotonically_more_conservative():
|
||||
"""Each rung must reduce clocks and raise fans; a non-monotonic ladder would make
|
||||
escalation increase heat."""
|
||||
scales = [s["offset_scale"] for s in tg.DERATE_LADDER]
|
||||
floors = [s["fan_floor"] for s in tg.DERATE_LADDER]
|
||||
assert scales == sorted(scales, reverse=True)
|
||||
assert floors == sorted(floors)
|
||||
assert scales[0] == 1.0 and scales[-1] == 0.0
|
||||
assert [s["level"] for s in tg.DERATE_LADDER] == list(range(len(tg.DERATE_LADDER)))
|
||||
|
||||
|
||||
def test_escalate_and_recover_temperatures_have_a_gap():
|
||||
"""Without a gap between the two thresholds the governor would oscillate."""
|
||||
assert tg.TEMP_RECOVER_C < tg.TEMP_ESCALATE_C
|
||||
|
||||
|
||||
# --------------------------------------------------------------- status / control
|
||||
|
||||
def test_status_reports_level_and_history(gov):
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
st = gov.get_status()
|
||||
assert st["level"] == 1
|
||||
assert st["label"] == tg.DERATE_LADDER[1]["label"]
|
||||
assert st["escalate_at_c"] == tg.TEMP_ESCALATE_C
|
||||
assert st["history"] and st["history"][0]["to_level"] == 1
|
||||
assert "95" in st["history"][0]["reason"]
|
||||
|
||||
|
||||
def test_history_is_bounded(gov):
|
||||
"""The governor is long-lived inside the service; its history must not grow forever."""
|
||||
for _ in range(80):
|
||||
_clear_cooldown(gov)
|
||||
gov.level = 0
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
assert len(gov.history) <= 50
|
||||
|
||||
|
||||
def test_disabling_a_derated_governor_resets_it(gov, no_gpu_mutation):
|
||||
"""Turning the governor off must give the clocks back, not freeze the derate in place."""
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
assert gov.level == 1
|
||||
gov.set_enabled(False)
|
||||
assert gov.level == 0
|
||||
assert no_gpu_mutation["apply_profile"], "reset should have re-applied the base profile"
|
||||
139
tests/test_vram_helpers.py
Normal file
139
tests/test_vram_helpers.py
Normal file
@@ -0,0 +1,139 @@
|
||||
"""Pure helpers in vram_arbitrator: NVML throttle-bit decoding and PID attribution.
|
||||
|
||||
Deliberately excludes instant_free_ollama_vram, the AutoArbitrator yield/purge paths and
|
||||
the SSE broker — that contract is in flux.
|
||||
"""
|
||||
import vram_arbitrator as va
|
||||
|
||||
|
||||
def test_decode_throttle_reasons_empty_when_no_bits_set():
|
||||
assert va.decode_throttle_reasons(0) == []
|
||||
|
||||
|
||||
def test_decode_throttle_reasons_maps_each_known_bit():
|
||||
"""Every mask in the table must decode to exactly its own name in isolation."""
|
||||
for mask, name in va.THROTTLE_REASONS.items():
|
||||
assert va.decode_throttle_reasons(mask) == [name]
|
||||
|
||||
|
||||
def test_decode_throttle_reasons_decodes_combined_bits():
|
||||
"""Real NVML samples set several bits at once; all of them must come back."""
|
||||
bits = 0x20 | 0x40 # sw_thermal_slowdown | hw_thermal_slowdown
|
||||
assert set(va.decode_throttle_reasons(bits)) == {"sw_thermal_slowdown", "hw_thermal_slowdown"}
|
||||
|
||||
|
||||
def test_decode_throttle_reasons_ignores_unknown_bits():
|
||||
"""An undocumented bit from a future driver must not crash or invent a reason."""
|
||||
assert va.decode_throttle_reasons(0x8000_0000) == []
|
||||
|
||||
|
||||
def test_hard_throttle_names_match_thermal_governor_expectations():
|
||||
"""thermal_governor escalates on a fixed set of reason strings produced here.
|
||||
If a name is renamed in one module and not the other the governor silently stops
|
||||
reacting to hardware slowdowns, so pin the shared vocabulary."""
|
||||
import thermal_governor as tg
|
||||
assert tg.HARD_THROTTLES <= set(va.THROTTLE_REASONS.values())
|
||||
|
||||
|
||||
class _FakeProc:
|
||||
def __init__(self, name, cmdline):
|
||||
self._name = name
|
||||
self._cmdline = cmdline
|
||||
|
||||
def name(self):
|
||||
return self._name
|
||||
|
||||
def cmdline(self):
|
||||
return self._cmdline
|
||||
|
||||
|
||||
def _patch_proc(monkeypatch, proc):
|
||||
monkeypatch.setattr(va.psutil, "Process", lambda pid: proc)
|
||||
|
||||
|
||||
def test_classify_pid_detects_ollama_by_process_name(monkeypatch):
|
||||
_patch_proc(monkeypatch, _FakeProc("ollama", ["/usr/local/bin/ollama", "serve"]))
|
||||
assert va._classify_pid(1234) == "ollama"
|
||||
|
||||
|
||||
def test_classify_pid_detects_ollama_runner_by_cmdline(monkeypatch):
|
||||
"""Ollama's model runner is a separate llama-server process; its VRAM is Ollama's."""
|
||||
_patch_proc(monkeypatch, _FakeProc("llama-server",
|
||||
["/usr/lib/ollama/llama-server", "--model", "blob"]))
|
||||
assert va._classify_pid(1234) == "ollama"
|
||||
|
||||
|
||||
def test_classify_pid_detects_comfyui(monkeypatch):
|
||||
_patch_proc(monkeypatch, _FakeProc("python3", ["python3", "/opt/ComfyUI/main.py", "--listen"]))
|
||||
assert va._classify_pid(1234) == "comfy"
|
||||
|
||||
|
||||
def test_classify_pid_unknown_process_is_unmanaged(monkeypatch):
|
||||
# Xorg used to stand in for "unknown" here, but a display server is now its own
|
||||
# bucket, so this needs a process that is genuinely neither ours nor the desktop's.
|
||||
_patch_proc(monkeypatch, _FakeProc("trainer", ["/opt/ml/bin/trainer", "--epochs", "3"]))
|
||||
# "unmanaged" rather than "other": a third-party GPU workload holds VRAM this
|
||||
# service cannot reclaim, and must not be lumped in with the desktop compositor.
|
||||
assert va._classify_pid(1234) == "unmanaged"
|
||||
|
||||
|
||||
def test_classify_pid_display_server_is_desktop(monkeypatch):
|
||||
_patch_proc(monkeypatch, _FakeProc("Xorg", ["/usr/lib/xorg/Xorg", ":8"]))
|
||||
assert va._classify_pid(1234) == "desktop"
|
||||
|
||||
|
||||
def test_classify_pid_returns_unmanaged_when_process_vanished(monkeypatch):
|
||||
"""PIDs are read from NVML and can exit before psutil looks them up; that is normal
|
||||
and must not raise inside the 20 ms VRAM poll loop."""
|
||||
def _boom(pid):
|
||||
raise va.psutil.NoSuchProcess(pid)
|
||||
monkeypatch.setattr(va.psutil, "Process", _boom)
|
||||
assert va._classify_pid(999999) == "unmanaged"
|
||||
|
||||
|
||||
# --- process bucketing: desktop vs unmanaged ---------------------------------
|
||||
#
|
||||
# Real case from this machine: stt_relay.py held 842 MB of VRAM for nearly three days
|
||||
# while gnome-shell held 3.9 MB. A single "other" bucket reported them as one number,
|
||||
# which matters because ComfyUI's memory can be reclaimed and a third party's cannot.
|
||||
|
||||
class _FakeProc:
|
||||
def __init__(self, name, cmdline):
|
||||
self._name = name
|
||||
self._cmdline = cmdline
|
||||
def name(self):
|
||||
return self._name
|
||||
def cmdline(self):
|
||||
return self._cmdline
|
||||
def create_time(self):
|
||||
return 1234.5
|
||||
|
||||
|
||||
def _classify(monkeypatch, name, cmdline):
|
||||
monkeypatch.setattr(va.psutil, "Process",
|
||||
lambda pid: _FakeProc(name, cmdline))
|
||||
return va._classify_pid(4321)
|
||||
|
||||
|
||||
def test_desktop_compositors_are_their_own_bucket(monkeypatch):
|
||||
assert _classify(monkeypatch, "gnome-shell", ["/usr/bin/gnome-shell"]) == "desktop"
|
||||
assert _classify(monkeypatch, "Xorg", ["/usr/lib/xorg/Xorg", ":8"]) == "desktop"
|
||||
|
||||
|
||||
def test_third_party_compute_is_unmanaged_not_desktop(monkeypatch):
|
||||
kind = _classify(monkeypatch, "python",
|
||||
["/home/u/robopest-venv/bin/python", "/home/u/stt_relay.py"])
|
||||
assert kind == "unmanaged"
|
||||
|
||||
|
||||
def test_ollama_and_comfy_still_win_over_the_catch_all(monkeypatch):
|
||||
assert _classify(monkeypatch, "llama-server",
|
||||
["/usr/local/lib/ollama/llama-server", "--model", "x"]) == "ollama"
|
||||
assert _classify(monkeypatch, "python",
|
||||
["/home/u/ComfyUI/venv/bin/python", "main.py"]) == "comfy"
|
||||
|
||||
|
||||
def test_pid_cache_is_keyed_by_start_time_not_pid_alone():
|
||||
# Linux recycles PIDs; a stale entry would attribute a new process's VRAM to Ollama
|
||||
# inside the same snapshot the yield barrier trusts.
|
||||
assert all(isinstance(k, tuple) and len(k) == 2 for k in va._PID_KIND_CACHE)
|
||||
202
tests/test_yield_and_reclaim.py
Normal file
202
tests/test_yield_and_reclaim.py
Normal file
@@ -0,0 +1,202 @@
|
||||
"""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 time
|
||||
|
||||
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
|
||||
|
||||
|
||||
class TestComfyStaleQueueDetection:
|
||||
"""ComfyUI can leave a dead job in queue_running forever.
|
||||
|
||||
Observed on this machine: a WAN 2.1 i2v entry sat in queue_running while the GPU was
|
||||
idle and ComfyUI held 0.56 GB. Trusting that flag made the watchdog believe ComfyUI
|
||||
was permanently busy, so it evicted the LLM on every poll, never ran the idle purge,
|
||||
and never checked whether the LLM had been pushed onto the CPU. Instrumenting the
|
||||
watchdog showed busy=6, idle_check=0 -- one stale row had disabled half the logic.
|
||||
"""
|
||||
|
||||
def _arb(self, comfy_bytes):
|
||||
arb = v.AutoArbitrator()
|
||||
v.get_process_vram_bytes = lambda: {
|
||||
"ollama_bytes": 0, "comfyui_bytes": int(comfy_bytes), "other_bytes": 0,
|
||||
"desktop_bytes": 0, "unmanaged_bytes": 0, "free_bytes": 0, "gpu_util_pct": 0}
|
||||
return arb
|
||||
|
||||
def teardown_method(self):
|
||||
import importlib
|
||||
importlib.reload(v)
|
||||
|
||||
def test_empty_queue_is_not_busy(self):
|
||||
arb = self._arb(0)
|
||||
assert arb._comfy_genuinely_busy({"queue_running": [], "queue_pending": []}) is False
|
||||
|
||||
def test_pending_work_is_always_busy(self):
|
||||
arb = self._arb(0)
|
||||
assert arb._comfy_genuinely_busy(
|
||||
{"queue_running": [], "queue_pending": [[1, "p"]]}) is True
|
||||
|
||||
def test_a_running_job_is_believed_at_first(self):
|
||||
# It must not be called stale before it has had time to load anything.
|
||||
arb = self._arb(0.1 * GB)
|
||||
assert arb._comfy_genuinely_busy(
|
||||
{"queue_running": [[1, "abc"]], "queue_pending": []}) is True
|
||||
|
||||
def test_long_running_job_holding_no_vram_is_stale(self):
|
||||
arb = self._arb(0.56 * GB) # the observed CUDA-context floor
|
||||
q = {"queue_running": [[1, "abc"]], "queue_pending": []}
|
||||
arb._comfy_genuinely_busy(q)
|
||||
arb._running_since = time.time() - (arb.STALE_RUNNING_S + 5)
|
||||
assert arb._comfy_genuinely_busy(q) is False
|
||||
assert arb.comfy_stale_job == "abc"
|
||||
|
||||
def test_long_running_job_holding_a_checkpoint_is_real(self):
|
||||
# 6.8 GB is a loaded SDXL checkpoint; slow is not the same as stuck.
|
||||
arb = self._arb(6.8 * GB)
|
||||
q = {"queue_running": [[1, "abc"]], "queue_pending": []}
|
||||
arb._comfy_genuinely_busy(q)
|
||||
arb._running_since = time.time() - (arb.STALE_RUNNING_S + 5)
|
||||
assert arb._comfy_genuinely_busy(q) is True
|
||||
assert arb.comfy_stale_job is None
|
||||
|
||||
def test_a_new_prompt_id_resets_the_staleness_clock(self):
|
||||
arb = self._arb(0.5 * GB)
|
||||
arb._comfy_genuinely_busy({"queue_running": [[1, "old"]], "queue_pending": []})
|
||||
arb._running_since = time.time() - 1000
|
||||
assert arb._comfy_genuinely_busy(
|
||||
{"queue_running": [[1, "new"]], "queue_pending": []}) is True
|
||||
|
||||
def test_vram_not_utilisation_is_the_signal(self):
|
||||
# Utilisation is shared with Ollama and any third-party process, so it stayed
|
||||
# above every sensible threshold and a stuck entry never looked stale.
|
||||
assert hasattr(v.AutoArbitrator, "STALE_COMFY_BYTES")
|
||||
assert not hasattr(v.AutoArbitrator, "STALE_UTIL_PCT")
|
||||
338
verify_arbitration.py
Executable file
338
verify_arbitration.py
Executable file
@@ -0,0 +1,338 @@
|
||||
#!/usr/bin/env python
|
||||
"""End-to-end verification of the arbitration cycle, against real hardware.
|
||||
|
||||
The unit suite covers logic in isolation; this exercises the promise the whole service
|
||||
exists to make -- that an LLM and a diffusion pipeline can share one 16 GB card without
|
||||
either failing -- and reports what actually happened at each stage.
|
||||
|
||||
It is deliberately not part of `pytest tests/`: it loads real models, runs a real
|
||||
diffusion graph and moves real VRAM, taking a few minutes. Run it when you want proof
|
||||
the system works on this machine:
|
||||
|
||||
python verify_arbitration.py # full cycle
|
||||
python verify_arbitration.py --quick # skip the diffusion stages
|
||||
|
||||
Every stage restores what it changed, and the script refuses to start if ComfyUI is
|
||||
already busy.
|
||||
"""
|
||||
import argparse
|
||||
import asyncio
|
||||
import sys
|
||||
import time
|
||||
from typing import Any, Dict, List, Optional
|
||||
|
||||
import httpx
|
||||
|
||||
BASE = "http://localhost:9090"
|
||||
|
||||
PASS, FAIL, SKIP, WARN = "PASS", "FAIL", "SKIP", "WARN"
|
||||
_results: List[Dict[str, Any]] = []
|
||||
|
||||
|
||||
def record(stage: str, status: str, detail: str, evidence: str = "") -> None:
|
||||
_results.append({"stage": stage, "status": status, "detail": detail,
|
||||
"evidence": evidence})
|
||||
colour = {"PASS": "\033[32m", "FAIL": "\033[31m",
|
||||
"SKIP": "\033[33m", "WARN": "\033[33m"}[status]
|
||||
print(f" {colour}{status:<4}\033[0m {stage:<38} {detail}")
|
||||
if evidence:
|
||||
print(f" {evidence}")
|
||||
|
||||
|
||||
async def api(client: httpx.AsyncClient, method: str, path: str,
|
||||
allow_error: bool = False, **kw) -> Any:
|
||||
r = await client.request(method, f"{BASE}{path}", **kw)
|
||||
if allow_error:
|
||||
# Some stages deliberately provoke a failure and need to read it.
|
||||
body = r.json() if r.headers.get("content-type", "").startswith("application/json") else {}
|
||||
return {"_status": r.status_code, **(body if isinstance(body, dict) else {})}
|
||||
r.raise_for_status()
|
||||
return r.json()
|
||||
|
||||
|
||||
async def stage_preflight(c: httpx.AsyncClient) -> Optional[Dict[str, Any]]:
|
||||
health = await api(c, "GET", "/api/health")
|
||||
failed = health.get("failed") or []
|
||||
if failed:
|
||||
record("preflight: dependencies", FAIL,
|
||||
f"{len(failed)} dependency check(s) failing", ", ".join(failed))
|
||||
return None
|
||||
record("preflight: dependencies", PASS, health["summary"])
|
||||
|
||||
stats = await api(c, "GET", "/api/stats")
|
||||
if stats["comfyui"].get("executing") or stats["comfyui"].get("queue_remaining"):
|
||||
record("preflight: ComfyUI idle", FAIL, "ComfyUI is busy; refusing to interfere")
|
||||
return None
|
||||
record("preflight: ComfyUI idle", PASS, "queue empty")
|
||||
return stats
|
||||
|
||||
|
||||
async def stage_llm_load(c: httpx.AsyncClient, model: str) -> bool:
|
||||
res = await api(c, "POST", "/api/switch-model",
|
||||
json={"model": model, "keep_alive": "5m"}, timeout=300)
|
||||
if not res.get("success"):
|
||||
record("LLM loads", FAIL, res.get("error", "")[:90])
|
||||
return False
|
||||
gbps, status = res.get("load_gbps"), res.get("cache_status")
|
||||
record("LLM loads", PASS, f"{res['model_size_gb']} GB in {res['load_duration_ms']:.0f} ms",
|
||||
f"{gbps} GB/s -> {status}; {res['tokens_per_sec']} tok/s")
|
||||
|
||||
# The classification must follow from the measured bandwidth, not a fixed duration.
|
||||
if gbps is not None:
|
||||
expected = ("RAM Cache Hit" if gbps >= 2.0
|
||||
else "Partial Cache" if gbps >= 0.8 else "Cold Disk Load")
|
||||
ok = expected.split()[0] in (status or "")
|
||||
record("load classified by bandwidth", PASS if ok else FAIL,
|
||||
f"{gbps} GB/s reported as '{status}'",
|
||||
"" if ok else f"expected something matching '{expected}'")
|
||||
return True
|
||||
|
||||
|
||||
async def stage_yield(c: httpx.AsyncClient) -> bool:
|
||||
before = await api(c, "GET", "/api/gpu")
|
||||
held = before["breakdown"]["ollama_gb"]
|
||||
res = await api(c, "POST", "/api/free-vram", timeout=60)
|
||||
outcome = res.get("outcome")
|
||||
|
||||
if outcome == "released":
|
||||
after = await api(c, "GET", "/api/gpu")
|
||||
freed = held - after["breakdown"]["ollama_gb"]
|
||||
# The barrier's promise: on return, the VRAM is genuinely gone.
|
||||
ok = after["breakdown"]["ollama_gb"] < 0.3
|
||||
record("VRAM yield is confirmed", PASS if ok else FAIL,
|
||||
f"released in {res['confirm_ms']} ms",
|
||||
f"{freed:.2f} GB freed; NVML now reports "
|
||||
f"{after['breakdown']['ollama_gb']} GB held by Ollama")
|
||||
return ok
|
||||
if outcome == "busy":
|
||||
record("VRAM yield is confirmed", WARN,
|
||||
"model was mid-generation, so the unload was deferred",
|
||||
res.get("error", ""))
|
||||
return True
|
||||
record("VRAM yield is confirmed", FAIL, f"outcome={outcome}", res.get("error", ""))
|
||||
return False
|
||||
|
||||
|
||||
async def stage_diffusion(c: httpx.AsyncClient) -> bool:
|
||||
sys.path.insert(0, "/home/drjones/unified-model-manager")
|
||||
import autotune, vram_arbitrator # noqa: E402 (imported late; needs the service's deps)
|
||||
|
||||
# The first run loads the checkpoint from disk. Timing that and calling the result
|
||||
# "it/s" understates throughput by roughly 10x -- 0.67 it/s against a steady-state
|
||||
# 6.7 -- so the load is measured separately and reported as what it is.
|
||||
first = await autotune._diffusion_benchmark()
|
||||
if not first.get("ok"):
|
||||
record("diffusion runs", FAIL, first.get("error", "")[:90])
|
||||
return False
|
||||
record("diffusion runs (cold, includes checkpoint load)", PASS,
|
||||
f"{first['exec_ms']} ms", f"{first['it_per_sec']} it/s including load")
|
||||
|
||||
res = await autotune._diffusion_benchmark()
|
||||
if not res.get("ok"):
|
||||
record("diffusion runs (warm)", FAIL, res.get("error", "")[:90])
|
||||
return False
|
||||
record("diffusion throughput (warm)", PASS,
|
||||
f"SDXL 1024/20 steps in {res['exec_ms']} ms", f"{res['it_per_sec']} it/s")
|
||||
|
||||
snap = vram_arbitrator.get_process_vram_bytes()
|
||||
comfy_gb = snap["comfyui_bytes"] / (1024 ** 3)
|
||||
record("ComfyUI holds its checkpoint", PASS if comfy_gb > 0.5 else WARN,
|
||||
f"{comfy_gb:.2f} GB retained",
|
||||
"held for the idle window rather than purged between iterations")
|
||||
return True
|
||||
|
||||
|
||||
async def stage_idle_purge(c: httpx.AsyncClient) -> bool:
|
||||
# The completion event arrives over the ComfyUI websocket, so the flag is set a
|
||||
# moment after the graph returns. Checking instantly raced it.
|
||||
arb = {}
|
||||
for _ in range(12):
|
||||
arb = (await api(c, "GET", "/api/stats"))["arbitrator"]
|
||||
if arb.get("pending_purge"):
|
||||
break
|
||||
await asyncio.sleep(0.5)
|
||||
if not arb.get("pending_purge"):
|
||||
record("purge is deferred, not immediate", WARN,
|
||||
"no purge pending after 6 s (ComfyUI may already be clean)")
|
||||
return True
|
||||
idle_s = arb.get("comfy_idle_s")
|
||||
record("purge is deferred, not immediate", PASS,
|
||||
f"holding checkpoints for {arb.get('idle_purge_after_s')} s",
|
||||
f"idle {idle_s} s so far" if idle_s is not None
|
||||
else "idle timer just started")
|
||||
return True
|
||||
|
||||
|
||||
async def stage_reclaim(c: httpx.AsyncClient, model: str) -> bool:
|
||||
"""The direction that used to fail outright: an LLM that will not fit.
|
||||
|
||||
This only proves anything if the chosen model genuinely cannot fit in what ComfyUI
|
||||
has left free. A small model fits alongside the checkpoint and the stage passes
|
||||
without exercising the reclaim path at all, so pick the largest model that will not
|
||||
fit and say plainly when no such model exists.
|
||||
"""
|
||||
gpu = await api(c, "GET", "/api/gpu")
|
||||
comfy_gb = gpu["breakdown"]["comfyui_gb"]
|
||||
free_gb = gpu["vram_free_gb"]
|
||||
if comfy_gb < 0.5:
|
||||
record("reclaims VRAM for the LLM", SKIP,
|
||||
f"ComfyUI only holds {comfy_gb} GB; nothing to reclaim")
|
||||
return True
|
||||
|
||||
models = (await api(c, "GET", "/api/models"))["ollama_models"]
|
||||
EMBED = {"bert", "nomic-bert", "gte", "jina-bert"}
|
||||
usable = [m for m in models
|
||||
if (m.get("details", {}).get("family") or "").lower() not in EMBED
|
||||
and "embed" not in m["name"].lower()]
|
||||
# On-disk weight size is not the VRAM footprint: measured on this box, a 12.87 GB
|
||||
# blob occupies 14.9 GB once context and KV cache are allocated. Sizing the test off
|
||||
# disk size picks a model that cannot fit even after a successful reclaim.
|
||||
VRAM_OVERHEAD = 1.18
|
||||
HEADROOM_GB = 0.4
|
||||
|
||||
def vram_need(m):
|
||||
return m.get("size", 0) / (1024 ** 3) * VRAM_OVERHEAD
|
||||
|
||||
# Unmanaged VRAM never comes back, so it is not part of what a reclaim can offer.
|
||||
# Ignoring it picked a model that failed even after a correct reclaim -- on this box
|
||||
# an 842 MB third-party process is the difference between a 14.9 GB model fitting
|
||||
# and not.
|
||||
reclaimable_gb = free_gb + comfy_gb - HEADROOM_GB
|
||||
too_big = [m for m in usable
|
||||
if vram_need(m) > free_gb and vram_need(m) < reclaimable_gb]
|
||||
if too_big:
|
||||
target = max(too_big, key=lambda m: m.get("size", 0))
|
||||
model = target["name"]
|
||||
print(f" using {model} ({target['size'] / (1024**3):.1f} GB on disk, "
|
||||
f"~{vram_need(target):.1f} GB in VRAM) — will not fit in "
|
||||
f"{free_gb:.1f} GB free, should fit after reclaiming {comfy_gb:.1f} GB")
|
||||
else:
|
||||
unmanaged = gpu["breakdown"].get("unmanaged_gb", 0)
|
||||
record("reclaims VRAM for the LLM", SKIP,
|
||||
f"no installed model needs between {free_gb:.1f} and "
|
||||
f"{reclaimable_gb:.1f} GB of VRAM",
|
||||
f"reclaimable ceiling excludes {unmanaged} GB held by processes "
|
||||
f"HyperSwap cannot free")
|
||||
return True
|
||||
|
||||
# Re-run a graph first. The idle purge fires 30 s after ComfyUI goes quiet, and a
|
||||
# large model takes longer than that to load -- so without resetting the timer the
|
||||
# purge frees ComfyUI mid-load and the reclaim path is never reached.
|
||||
sys.path.insert(0, "/home/drjones/unified-model-manager")
|
||||
import autotune # noqa: E402
|
||||
await autotune._diffusion_benchmark()
|
||||
gpu = await api(c, "GET", "/api/gpu")
|
||||
print(f" reset the idle window; ComfyUI holds "
|
||||
f"{gpu['breakdown']['comfyui_gb']} GB, {gpu['vram_free_gb']} GB free")
|
||||
|
||||
res = await api(c, "POST", "/api/switch-model", allow_error=True,
|
||||
json={"model": model, "keep_alive": "2m"}, timeout=600)
|
||||
if res.get("_status") == 507:
|
||||
record("reclaims VRAM for the LLM", FAIL,
|
||||
"reclaim ran but the model still did not fit",
|
||||
str(res.get("detail", ""))[:150])
|
||||
return False
|
||||
if res.get("_status", 200) >= 400 or not res.get("success"):
|
||||
record("reclaims VRAM for the LLM", FAIL,
|
||||
f"HTTP {res.get('_status')}", str(res.get("detail", ""))[:120])
|
||||
return False
|
||||
if res.get("_status") and res.get("_status") != 200:
|
||||
pass
|
||||
if res.get("reclaimed_from_comfyui_gb"):
|
||||
record("reclaims VRAM for the LLM", PASS,
|
||||
f"reclaimed {res['reclaimed_from_comfyui_gb']} GB and retried",
|
||||
f"'{model}' then loaded at {res.get('load_gbps')} GB/s")
|
||||
else:
|
||||
# It fit anyway, so nothing was proven; do not report that as a pass. The usual
|
||||
# cause is the idle purge firing during the load and freeing ComfyUI first.
|
||||
record("reclaims VRAM for the LLM", WARN,
|
||||
"model fit without a reclaim, so the path was not exercised",
|
||||
"the idle purge most likely freed ComfyUI during the load; "
|
||||
f"loaded at {res.get('load_gbps')} GB/s")
|
||||
return True
|
||||
|
||||
|
||||
async def stage_accounting(c: httpx.AsyncClient) -> bool:
|
||||
"""Reported VRAM must add up, and reported settings must match the hardware."""
|
||||
gpu = await api(c, "GET", "/api/gpu")
|
||||
bd = gpu["breakdown"]
|
||||
parts = bd["ollama_gb"] + bd["comfyui_gb"] + bd["system_gb"]
|
||||
used = gpu["vram_used_gb"]
|
||||
# Driver overhead means the parts never sum exactly; a large gap means mis-attribution.
|
||||
ok = abs(parts - used) < 1.5
|
||||
record("VRAM attribution adds up", PASS if ok else FAIL,
|
||||
f"parts {parts:.2f} GB vs NVML used {used:.2f} GB",
|
||||
f"ollama {bd['ollama_gb']} + comfy {bd['comfyui_gb']} + system {bd['system_gb']} "
|
||||
f"(desktop {bd.get('desktop_gb')}, unmanaged {bd.get('unmanaged_gb')})")
|
||||
|
||||
oc = await api(c, "GET", "/api/overclock")
|
||||
drift = oc["drift"]
|
||||
ok2 = not drift["drifted"]
|
||||
record("reported GPU state matches hardware", PASS if ok2 else FAIL,
|
||||
f"profile '{drift['profile']}' asks {drift['power_limit_intended_w']} W, "
|
||||
f"card reports {drift['power_limit_actual_w']} W",
|
||||
drift.get("reason") or "")
|
||||
return ok and ok2
|
||||
|
||||
|
||||
async def main() -> int:
|
||||
ap = argparse.ArgumentParser(description=__doc__,
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
ap.add_argument("--quick", action="store_true",
|
||||
help="skip the diffusion and reclaim stages")
|
||||
ap.add_argument("--model", default=None,
|
||||
help="Ollama model to test with (default: smallest installed)")
|
||||
args = ap.parse_args()
|
||||
|
||||
print("\nHyperSwap arbitration verification\n" + "=" * 62)
|
||||
async with httpx.AsyncClient(timeout=60.0) as c:
|
||||
baseline = await stage_preflight(c)
|
||||
if baseline is None:
|
||||
print("\nPreflight failed; not continuing.\n")
|
||||
return 2
|
||||
|
||||
model = args.model
|
||||
if not model:
|
||||
models = (await api(c, "GET", "/api/models"))["ollama_models"]
|
||||
# Embedding models have no /api/generate endpoint, and the smallest model
|
||||
# installed is often one of them.
|
||||
EMBED_FAMILIES = {"bert", "nomic-bert", "gte", "jina-bert"}
|
||||
usable = [m for m in models
|
||||
if (m.get("details", {}).get("family") or "").lower()
|
||||
not in EMBED_FAMILIES and "embed" not in m["name"].lower()]
|
||||
if not usable:
|
||||
record("choose a test model", FAIL,
|
||||
"no generative Ollama models installed "
|
||||
f"({len(models)} found, all embedding-only)")
|
||||
return 2
|
||||
model = min(usable, key=lambda m: m.get("size", 0))["name"]
|
||||
print(f"\n using model: {model}\n")
|
||||
|
||||
await stage_llm_load(c, model)
|
||||
await stage_yield(c)
|
||||
|
||||
if not args.quick:
|
||||
if await stage_diffusion(c):
|
||||
await stage_idle_purge(c)
|
||||
await stage_reclaim(c, model)
|
||||
else:
|
||||
record("diffusion stages", SKIP, "--quick")
|
||||
|
||||
await stage_accounting(c)
|
||||
|
||||
# Leave the box as we found it.
|
||||
await api(c, "POST", "/api/free-vram", timeout=60)
|
||||
|
||||
failed = [r for r in _results if r["status"] == FAIL]
|
||||
warned = [r for r in _results if r["status"] == WARN]
|
||||
print("=" * 62)
|
||||
print(f" {len(_results) - len(failed) - len(warned)} passed, "
|
||||
f"{len(warned)} warned, {len(failed)} failed")
|
||||
for r in failed:
|
||||
print(f" FAILED: {r['stage']} — {r['detail']}")
|
||||
print()
|
||||
return 1 if failed else 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(asyncio.run(main()))
|
||||
@@ -12,6 +12,7 @@ import websockets
|
||||
|
||||
import overclock_manager
|
||||
import ram_optimizer
|
||||
import tenants as tenants_mod
|
||||
import telemetry_store
|
||||
|
||||
try:
|
||||
@@ -44,13 +45,57 @@ 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._arbitrate (generically, from the tenant registry), 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 +150,18 @@ 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,
|
||||
"by_tenant_bytes": {}}
|
||||
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 +171,76 @@ 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
|
||||
out["by_tenant_bytes"][kind] = out["by_tenant_bytes"].get(kind, 0) + used
|
||||
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
|
||||
|
||||
|
||||
# Tenant names as used by this module's buckets. The tenant registry is the source of
|
||||
# truth for *which* application a process belongs to; these two names are kept because
|
||||
# the REST payloads and the dashboard have used them since the beginning.
|
||||
_BUCKET_ALIASES = {"comfyui": "comfy"}
|
||||
|
||||
|
||||
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:
|
||||
try:
|
||||
proc = psutil.Process(pid)
|
||||
pname = proc.name().lower()
|
||||
cmdline = " ".join(proc.cmdline()).lower()
|
||||
except Exception:
|
||||
return "other"
|
||||
if "ollama" in pname or "llama-server" in cmdline:
|
||||
return "ollama"
|
||||
if "comfy" in cmdline or "main.py" in cmdline:
|
||||
return "comfy"
|
||||
return "other"
|
||||
"""Which tenant owns this GPU process.
|
||||
|
||||
The matching rules used to be substrings compiled into this function, which made the
|
||||
two applications on this box part of the arbitrator rather than input to it. They now
|
||||
come from the tenant registry, so a third application is a config entry.
|
||||
|
||||
"unmanaged" still means something specific and useful: VRAM held by something with no
|
||||
declared way to release it, and therefore headroom this service can never offer.
|
||||
"""
|
||||
return _BUCKET_ALIASES.get(tenants_mod.classify_pid(pid), tenants_mod.classify_pid(pid))
|
||||
|
||||
|
||||
def get_gpu_hardware_stats() -> Dict[str, Any]:
|
||||
@@ -249,7 +340,13 @@ def get_gpu_hardware_stats() -> Dict[str, Any]:
|
||||
"ollama_bytes": 0,
|
||||
"comfyui_bytes": 0,
|
||||
"system_bytes": 0,
|
||||
"processes": []
|
||||
"desktop_bytes": 0,
|
||||
"unmanaged_bytes": 0,
|
||||
"unmanaged": [],
|
||||
"processes": [],
|
||||
# Generic attribution: one entry per tenant, so an application added to the
|
||||
# registry is reported without any change here.
|
||||
"by_tenant": {},
|
||||
}
|
||||
|
||||
try:
|
||||
@@ -269,16 +366,28 @@ 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
|
||||
elif is_comfy:
|
||||
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["by_tenant"][kind] = (
|
||||
proc_breakdown["by_tenant"].get(kind, 0) + used_mem)
|
||||
proc_breakdown["processes"].append({
|
||||
"pid": pid,
|
||||
"name": pname,
|
||||
@@ -287,6 +396,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 +438,13 @@ 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"],
|
||||
"by_tenant_gb": {k: round(b / (1024**3), 2)
|
||||
for k, b in proc_breakdown["by_tenant"].items()},
|
||||
"free_mb": round(free_vram / (1024**2), 1),
|
||||
"free_gb": round(free_vram / (1024**3), 2),
|
||||
"processes": proc_breakdown["processes"],
|
||||
@@ -345,7 +462,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 +538,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 +647,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 +673,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 +727,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 +802,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 +858,52 @@ 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:
|
||||
# Which application should give up memory is a question for the registry,
|
||||
# not something to answer by purging ComfyUI by name. Any reclaimable idle
|
||||
# tenant below Ollama in priority is a candidate.
|
||||
state = await arbitrator._tenant_state()
|
||||
free_gb = arbitrator._last_tenant_state["free_gb"]
|
||||
size_gb = _model_size_bytes(target_model) / (1024**3)
|
||||
needed = size_gb * 1.16 if size_gb else free_gb + 1.0
|
||||
plan = tenants_mod.plan_release("ollama", state, free_gb, needed)
|
||||
if plan["release"]:
|
||||
logger.warning(
|
||||
f"Ollama could not fit '{target_model}' — {plan['reason']}")
|
||||
freed_before = free_gb
|
||||
for victim in plan["release"]:
|
||||
await arbitrator._release_tenant(
|
||||
victim, f"Ollama could not load '{target_model}'")
|
||||
arbitrator.stats["reclaims_for_ollama"] += 1
|
||||
arbitrator.last_action = (
|
||||
f"Released {', '.join(plan['release'])} so '{target_model}' could load")
|
||||
_record({
|
||||
"event_type": "VRAM Reclaim for Ollama",
|
||||
"source": ", ".join(plan["release"]),
|
||||
"target": target_model,
|
||||
"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["released_tenants"] = plan["release"]
|
||||
retry["would_free_gb"] = plan.get("would_free_gb")
|
||||
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 a tenant
|
||||
# when a process nobody can release is holding the memory sends the
|
||||
# user looking in the wrong place.
|
||||
retry["blockers"] = plan.get("blockers")
|
||||
retry["unmanaged_blockers"] = describe_unmanaged()
|
||||
return retry
|
||||
return {"success": False, "error": f"HTTP {resp.status_code}: {body}",
|
||||
"duration_ms": total_duration_ms,
|
||||
"upstream_status": resp.status_code,
|
||||
"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)}
|
||||
@@ -669,6 +934,7 @@ class AutoArbitrator:
|
||||
def __init__(self):
|
||||
self.running = False
|
||||
self.ws_task: Optional[asyncio.Task] = None
|
||||
self.event_tasks: List[asyncio.Task] = []
|
||||
self.poll_task: Optional[asyncio.Task] = None
|
||||
self.idle_task: Optional[asyncio.Task] = None
|
||||
self.last_yield_time = 0.0
|
||||
@@ -683,13 +949,45 @@ 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.watchdog_branches = {"busy": 0, "completed": 0, "idle_check": 0,
|
||||
"bad_status": 0, "error": 0}
|
||||
self._running_id: Optional[str] = None
|
||||
self._running_since: Optional[float] = None
|
||||
self._peak_comfy_bytes = 0
|
||||
self.comfy_stale_job: Optional[str] = None
|
||||
self._idle_since: Dict[str, float] = {}
|
||||
self._last_event_wake = 0.0
|
||||
self.event_sources: Dict[str, str] = {}
|
||||
self._last_tenant_state: Optional[Dict[str, Any]] = None
|
||||
self.last_arbitration: Optional[Dict[str, Any]] = None
|
||||
self.last_handoff: Optional[Dict[str, Any]] = None
|
||||
self.last_watchdog_error: Optional[str] = None
|
||||
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:
|
||||
return
|
||||
self.running = True
|
||||
self.ws_task = asyncio.create_task(self._ws_listener())
|
||||
for t in tenants_mod.load_tenants():
|
||||
if t.enabled and t.events.type == "websocket" and t.events.url:
|
||||
self.event_tasks.append(asyncio.create_task(
|
||||
self._event_listener(t.name, t.events.url,
|
||||
t.events.reconnect_backoff_s,
|
||||
t.events.max_backoff_s)))
|
||||
self.poll_task = asyncio.create_task(self._poll_watchdog())
|
||||
self.idle_task = asyncio.create_task(self._idle_purge_loop())
|
||||
logger.info("AutoArbitrator background engine started (Bidirectional).")
|
||||
@@ -702,14 +1000,26 @@ class AutoArbitrator:
|
||||
|
||||
async def stop(self):
|
||||
self.running = False
|
||||
for task in (self.ws_task, self.poll_task, self.idle_task):
|
||||
for task in [self.ws_task, self.poll_task, self.idle_task, *self.event_tasks]:
|
||||
if task:
|
||||
task.cancel()
|
||||
self.event_tasks.clear()
|
||||
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 +1028,39 @@ class AutoArbitrator:
|
||||
return
|
||||
|
||||
ollama_state = await get_ollama_live_state()
|
||||
if ollama_state.get("active_model_name"):
|
||||
model = ollama_state["active_model_name"]
|
||||
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)
|
||||
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.last_action = (f"Yielded '{model}' for ComfyUI in "
|
||||
f"{res.get('confirm_ms')}ms (confirmed {res.get('freed_gb')}GB free)")
|
||||
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._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 ({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_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."""
|
||||
@@ -785,6 +1113,51 @@ class AutoArbitrator:
|
||||
return {"purged": True, "free_gb": round(snap["free_bytes"] / (1024**3), 2)}
|
||||
return {"purged": False, "free_gb": round(free_gb, 2), "reason": "ComfyUI holds no VRAM"}
|
||||
|
||||
async def _event_listener(self, tenant_name: str, url: str,
|
||||
backoff_s: float, max_backoff_s: float) -> None:
|
||||
"""Wake on a tenant's event stream instead of waiting for the next poll.
|
||||
|
||||
Deliberately does not parse the messages. The previous listener understood
|
||||
ComfyUI's schema -- status/execution_start/executing/execution_success -- which
|
||||
tied the fast path to one application. Treating any message as "look now" and
|
||||
letting the tenant's own busy probe decide gives the same sub-second reaction
|
||||
for any application that emits anything on state change.
|
||||
"""
|
||||
backoff = backoff_s
|
||||
while self.running:
|
||||
try:
|
||||
async with websockets.connect(url, ping_interval=10, ping_timeout=10) as ws:
|
||||
self.event_sources[tenant_name] = "connected"
|
||||
self.connected_ws = True
|
||||
backoff = backoff_s
|
||||
logger.info(f"Event source connected for '{tenant_name}': {url}")
|
||||
while self.running:
|
||||
await ws.recv()
|
||||
# Coalesce bursts: a single graph emits many messages, and one
|
||||
# arbitration pass per burst is enough.
|
||||
now = time.time()
|
||||
if now - self._last_event_wake < 0.05:
|
||||
continue
|
||||
self._last_event_wake = now
|
||||
self.stats["event_wakeups"] = self.stats.get("event_wakeups", 0) + 1
|
||||
try:
|
||||
# A message on this tenant's own stream is live proof it is
|
||||
# working right now, so it is taken as busy rather than asked
|
||||
# over HTTP. A stale queue row could lie; an event arriving
|
||||
# this instant cannot.
|
||||
await self._arbitrate(active_tenant=tenant_name)
|
||||
except Exception as e:
|
||||
logger.debug(f"arbitration from event failed: {e}")
|
||||
except (websockets.exceptions.ConnectionClosed, OSError, asyncio.CancelledError):
|
||||
self.event_sources[tenant_name] = "disconnected"
|
||||
self.connected_ws = False
|
||||
except Exception as e:
|
||||
self.event_sources[tenant_name] = f"error: {str(e)[:60]}"
|
||||
self.connected_ws = False
|
||||
logger.debug(f"event source error for '{tenant_name}': {e}")
|
||||
await asyncio.sleep(backoff)
|
||||
backoff = min(backoff * 1.5, max_backoff_s)
|
||||
|
||||
async def _ws_listener(self):
|
||||
client_id = "hyperswap-arbitrator"
|
||||
ws_url = f"ws://127.0.0.1:8188/ws?clientId={client_id}"
|
||||
@@ -836,6 +1209,211 @@ class AutoArbitrator:
|
||||
await asyncio.sleep(backoff)
|
||||
backoff = min(backoff * 1.5, 15.0)
|
||||
|
||||
RECLAIM_COOLDOWN_S = 30.0
|
||||
# A queue entry that has claimed to be running this long without the GPU ever going
|
||||
# busy is stale, not slow.
|
||||
STALE_RUNNING_S = 90.0
|
||||
# ComfyUI's own VRAM, not GPU utilisation, is what distinguishes a real job from a
|
||||
# stale row. Utilisation is shared: Ollama and any third-party process drive it too,
|
||||
# so peak utilisation stayed above any sensible threshold and a stuck entry never
|
||||
# looked stale. A real diffusion job loads gigabytes of checkpoint; a dead one holds
|
||||
# only the CUDA context.
|
||||
STALE_COMFY_BYTES = 1.5 * 1024 ** 3
|
||||
|
||||
def _comfy_genuinely_busy(self, queue: Dict[str, Any]) -> bool:
|
||||
"""Decide whether ComfyUI is really working, not just claiming to be.
|
||||
|
||||
ComfyUI can leave an entry in queue_running after a job dies -- observed here as
|
||||
a WAN 2.1 i2v entry that sat there with the GPU at 0% and ComfyUI holding 0.56 GB.
|
||||
Trusting that flag alone made this service believe ComfyUI was permanently busy,
|
||||
which meant it evicted the LLM on every poll, never ran the idle purge, and never
|
||||
checked whether the LLM had been squeezed onto the CPU. Half the arbitration was
|
||||
disabled by one stale row.
|
||||
|
||||
A running entry is corroborated against GPU utilisation before it is believed.
|
||||
"""
|
||||
running = queue.get("queue_running") or []
|
||||
pending = queue.get("queue_pending") or []
|
||||
if pending:
|
||||
self._running_since = None
|
||||
self._running_id = None
|
||||
return True
|
||||
if not running:
|
||||
self._running_since = None
|
||||
self._running_id = None
|
||||
self.comfy_stale_job = None
|
||||
return False
|
||||
|
||||
entry = running[0]
|
||||
prompt_id = entry[1] if isinstance(entry, (list, tuple)) and len(entry) > 1 else str(entry)
|
||||
now = time.time()
|
||||
if prompt_id != self._running_id:
|
||||
self._running_id = prompt_id
|
||||
self._running_since = now
|
||||
self._peak_comfy_bytes = 0
|
||||
|
||||
snap = get_process_vram_bytes()
|
||||
self._peak_comfy_bytes = max(self._peak_comfy_bytes, snap.get("comfyui_bytes", 0))
|
||||
|
||||
elapsed = now - (self._running_since or now)
|
||||
if elapsed > self.STALE_RUNNING_S and self._peak_comfy_bytes < self.STALE_COMFY_BYTES:
|
||||
if self.comfy_stale_job != prompt_id:
|
||||
logger.warning(
|
||||
f"ComfyUI reports prompt {prompt_id} running for {int(elapsed)}s while "
|
||||
f"holding only {self._peak_comfy_bytes / (1024**3):.2f} GB — no checkpoint "
|
||||
f"is loaded, so the queue entry is stale. Ignoring it; otherwise ComfyUI "
|
||||
f"looks permanently busy and arbitration stops working.")
|
||||
self.comfy_stale_job = prompt_id
|
||||
return False
|
||||
return True
|
||||
|
||||
async def _tenant_state(self, active_tenant: Optional[str] = None
|
||||
) -> List[Dict[str, Any]]:
|
||||
"""Current VRAM and busy state for every configured tenant.
|
||||
|
||||
`active_tenant` skips the HTTP busy probe for the tenant whose event stream just
|
||||
fired: the event is the evidence. That removes a round trip from the handoff,
|
||||
which is the one path where latency is the entire point.
|
||||
"""
|
||||
# Only the cheap NVML read. The full hardware snapshot also does a psutil lookup
|
||||
# per process, which is wasted work on the handoff path where latency is the
|
||||
# entire point.
|
||||
snap = get_process_vram_bytes()
|
||||
by_tenant = {k: b / (1024 ** 3) for k, b in snap["by_tenant_bytes"].items()}
|
||||
out = []
|
||||
for t in tenants_mod.load_tenants():
|
||||
if not t.enabled:
|
||||
continue
|
||||
bucket = _BUCKET_ALIASES.get(t.name, t.name)
|
||||
vram_gb = by_tenant.get(bucket, 0.0)
|
||||
if t.name == active_tenant:
|
||||
probe = {"busy": True, "reason": "event received from its own stream"}
|
||||
else:
|
||||
probe = await tenants_mod.probe_busy(t, vram_gb=vram_gb)
|
||||
out.append({
|
||||
"name": t.name,
|
||||
"priority": t.priority,
|
||||
"vram_gb": vram_gb,
|
||||
"busy": bool(probe.get("busy")),
|
||||
"below_floor": bool(probe.get("below_floor")),
|
||||
"reclaimable": t.reclaimable,
|
||||
"needs_vram_gb": t.needs_vram_gb,
|
||||
"overclock_profile": t.overclock_profile,
|
||||
"vram_floor_gb": t.vram_floor_gb,
|
||||
"idle_release_after_s": t.idle_release_after_s,
|
||||
"reason": probe.get("reason"),
|
||||
})
|
||||
self._last_tenant_state = {"ts": time.time(), "free_gb":
|
||||
round(snap["free_bytes"] / (1024**3), 2),
|
||||
"tenants": out}
|
||||
return out
|
||||
|
||||
async def _release_tenant(self, name: str, reason: str) -> Dict[str, Any]:
|
||||
"""Release one tenant's VRAM by whatever mechanism it declares."""
|
||||
t = tenants_mod.get_tenant(name)
|
||||
if not t or not t.reclaimable:
|
||||
return {"success": False, "reason": "not reclaimable"}
|
||||
models = None
|
||||
if t.release.per_model:
|
||||
state = await get_ollama_live_state()
|
||||
models = [m.get("name") for m in state.get("loaded_models", []) if m.get("name")]
|
||||
logger.info(f"Releasing VRAM from '{name}': {reason}")
|
||||
res = await tenants_mod.release_vram(t, models=models)
|
||||
self.stats["tenant_releases"] = self.stats.get("tenant_releases", 0) + 1
|
||||
return res
|
||||
|
||||
IDLE_PROFILE = "balanced"
|
||||
|
||||
def _apply_profile_for_active(self, state: List[Dict[str, Any]]) -> None:
|
||||
"""Apply the GPU profile declared by whichever tenant is currently working.
|
||||
|
||||
This used to be two calls naming 'comfy' and 'ollama' directly, so a third
|
||||
application could never get tuned clocks. The highest-priority busy tenant wins;
|
||||
with nothing working the card returns to the idle profile.
|
||||
"""
|
||||
busy = [s for s in state if s["busy"] and s.get("overclock_profile")]
|
||||
if busy:
|
||||
busy.sort(key=lambda s: -s["priority"])
|
||||
self._apply_oc_profile(busy[0]["overclock_profile"])
|
||||
else:
|
||||
self._apply_oc_profile(self.IDLE_PROFILE)
|
||||
|
||||
async def _arbitrate(self, active_tenant: Optional[str] = None) -> None:
|
||||
"""Generic arbitration over any number of tenants.
|
||||
|
||||
The two-application version was a pair of hardcoded rules -- yield Ollama when
|
||||
ComfyUI is busy, purge ComfyUI when Ollama is starved -- which could not express
|
||||
a third participant at all. This works from the registry instead: a busy tenant
|
||||
that lacks the VRAM it declares it needs is starved, and the memory comes from
|
||||
idle reclaimable tenants below it in priority, lowest first.
|
||||
"""
|
||||
t_start = time.perf_counter()
|
||||
state = await self._tenant_state(active_tenant)
|
||||
free_gb = self._last_tenant_state["free_gb"]
|
||||
self._apply_profile_for_active(state)
|
||||
|
||||
# 1. Starvation: highest-priority demanding tenant first.
|
||||
for s in sorted(state, key=lambda x: -x["priority"]):
|
||||
if not s["busy"] or not s["needs_vram_gb"]:
|
||||
continue
|
||||
# Starved means it cannot reach what it needs even counting what it already
|
||||
# holds. Comparing free VRAM alone flagged a tenant that was working
|
||||
# perfectly well on 13 GB as demanding, purely because little was left over
|
||||
# -- which is the normal state of a busy GPU, and would have caused
|
||||
# pointless releases from everyone else.
|
||||
if s["vram_gb"] + free_gb >= s["needs_vram_gb"]:
|
||||
continue
|
||||
plan = tenants_mod.plan_release(s["name"], state, free_gb, s["needs_vram_gb"])
|
||||
self.last_arbitration = {"ts": time.time(), "demanding": s["name"],
|
||||
"free_gb": free_gb, **plan}
|
||||
if not plan["release"]:
|
||||
logger.debug(f"'{s['name']}' is short of VRAM but {plan['reason']}")
|
||||
return
|
||||
if time.time() - self.last_reclaim_time < self.RECLAIM_COOLDOWN_S:
|
||||
return
|
||||
self.last_reclaim_time = time.time()
|
||||
for victim in plan["release"]:
|
||||
await self._release_tenant(
|
||||
victim, f"{s['name']} needs {s['needs_vram_gb']} GB, {free_gb} GB free")
|
||||
# Wait for the memory to actually come back, and record how long the whole
|
||||
# handoff took. Swap speed is the point of this service, so it is measured
|
||||
# rather than assumed.
|
||||
target_bytes = int(s["needs_vram_gb"] * (1024 ** 3))
|
||||
deadline = time.perf_counter() + 30.0
|
||||
while time.perf_counter() < deadline:
|
||||
if get_process_vram_bytes()["free_bytes"] >= target_bytes:
|
||||
break
|
||||
await asyncio.sleep(0.02)
|
||||
handoff_ms = round((time.perf_counter() - t_start) * 1000, 1)
|
||||
self.last_handoff = {"ts": time.time(), "to": s["name"],
|
||||
"released": plan["release"], "handoff_ms": handoff_ms,
|
||||
"triggered_by": "event" if active_tenant else "poll"}
|
||||
self.stats["handoffs"] = self.stats.get("handoffs", 0) + 1
|
||||
logger.info(f"Handoff to '{s['name']}' in {handoff_ms} ms "
|
||||
f"(released {', '.join(plan['release'])})")
|
||||
self.last_action = (f"Released {', '.join(plan['release'])} so "
|
||||
f"'{s['name']}' could work — {handoff_ms} ms")
|
||||
return
|
||||
|
||||
# 2. Idle release: a tenant holding VRAM it is not using, after a grace period.
|
||||
now = time.time()
|
||||
for s in state:
|
||||
if not s["reclaimable"] or s["vram_gb"] <= 0.25:
|
||||
self._idle_since.pop(s["name"], None)
|
||||
continue
|
||||
if s["busy"]:
|
||||
self._idle_since.pop(s["name"], None)
|
||||
continue
|
||||
since = self._idle_since.setdefault(s["name"], now)
|
||||
grace = s["idle_release_after_s"]
|
||||
if grace and (now - since) >= grace:
|
||||
self._idle_since.pop(s["name"], None)
|
||||
await self._release_tenant(
|
||||
s["name"], f"idle {int(now - since)}s holding {s['vram_gb']} GB")
|
||||
self.last_action = (f"Released idle '{s['name']}' after "
|
||||
f"{int(now - since)}s")
|
||||
return
|
||||
|
||||
async def _poll_watchdog(self):
|
||||
"""Fallback for when the WebSocket is down. One cheap /queue call, 1 Hz.
|
||||
|
||||
@@ -849,13 +1427,25 @@ class AutoArbitrator:
|
||||
resp = await client.get("/queue")
|
||||
if resp.status_code == 200:
|
||||
q = resp.json()
|
||||
busy = len(q.get("queue_running", [])) > 0 or len(q.get("queue_pending", [])) > 0
|
||||
busy = self._comfy_genuinely_busy(q)
|
||||
if busy:
|
||||
self.watchdog_branches["busy"] += 1
|
||||
await self.trigger_comfy_priority("Watchdog saw an active queue")
|
||||
elif self.comfy_was_active:
|
||||
self.watchdog_branches["completed"] += 1
|
||||
await self.trigger_comfy_completed()
|
||||
except Exception:
|
||||
pass
|
||||
else:
|
||||
self.watchdog_branches["idle_check"] += 1
|
||||
await self._arbitrate()
|
||||
else:
|
||||
self.watchdog_branches["bad_status"] += 1
|
||||
except Exception as e:
|
||||
# This used to swallow everything silently, including anything raised by
|
||||
# the starvation check, which is why that check could appear to run and
|
||||
# do nothing.
|
||||
self.watchdog_branches["error"] += 1
|
||||
self.last_watchdog_error = str(e)[:200]
|
||||
logger.debug(f"watchdog poll error: {e}")
|
||||
await asyncio.sleep(interval)
|
||||
|
||||
def suspend_oc(self, reason: str = "tuning sweep") -> None:
|
||||
@@ -893,6 +1483,16 @@ class AutoArbitrator:
|
||||
"idle_purge_after_s": self.COMFY_IDLE_PURGE_S,
|
||||
"oc_profile": self.oc_profile,
|
||||
"counters": dict(self.stats),
|
||||
"comfy_stale_job": self.comfy_stale_job,
|
||||
"event_sources": dict(self.event_sources),
|
||||
"last_arbitration": self.last_arbitration,
|
||||
"last_handoff": self.last_handoff,
|
||||
"tenant_state": self._last_tenant_state,
|
||||
"watchdog_branches": dict(self.watchdog_branches),
|
||||
"last_watchdog_error": self.last_watchdog_error,
|
||||
"yield_backoff": {m: round(max(t - time.time(), 0), 1)
|
||||
for m, t in self._yield_backoff_until.items()
|
||||
if t > time.time()},
|
||||
}
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user