Add test suite (164 tests); reclaim VRAM from ComfyUI when an LLM will not fit
Tests. First automated coverage for the project: 164 tests, 2.7s, no GPU or network. An autouse fixture stubs overclock_manager._sh -- the single choke point for every nvidia-smi/nvidia-settings write -- so no test can mutate the card. They deliberately pin the empirically measured constants that would otherwise rot silently: the cold and warm load figures behind the cache-hit thresholds, the warm_confident residency rule, and the busy/stalled yield split. One test asserts RAM_HIT_GBPS stays at or below the measured 2.63 GB/s warm load, so the old physically unreachable 5.0 GB/s bar cannot come back. Three bugs the suite surfaced, now fixed: - autotune._subsample(values, 1) divided by zero; the early return only covered len(values) <= max_steps. - telemetry_store.stop() flushed its local pending list but never drained the queue, silently losing rows submitted just before a shutdown -- exactly when the last events matter. - ram_optimizer.page_residency's zero-byte short-circuit omitted keys every other return path provides, so a 0-byte file was planned for warming. Reclaim. The README has claimed bidirectional arbitration from the start, but only one direction was ever automatic. Establishing what actually happens took a controlled test with the service stopped: with ComfyUI holding 6.83 GB, Ollama does not spill to the CPU on this box -- it aborts with "cudaMalloc failed: out of memory", because n_gpu_layers is pinned to 99 and it will not reduce the layer count. So both failure modes are handled: _check_ollama_starved watches size_vram < size for the default configuration where Ollama does spill, and switch_ollama_model catches the hard OOM, reclaims VRAM from an idle ComfyUI and retries once. The request that returned HTTP 500 from Ollama directly now succeeds through HyperSwap, loading at 3.85 GB/s after reclaiming 6.83 GB. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
drjones
15
README.md
15
README.md
@@ -20,6 +20,21 @@
|
||||
## 1. Feature Matrix
|
||||
|
||||
### ⚡ Bidirectional VRAM Hot-Swapping & Arbitration
|
||||
|
||||
* **Both directions are now automatic.** Yielding Ollama for ComfyUI always was; the
|
||||
reverse was not, despite "bidirectional" in this heading. Which way an LLM fails when
|
||||
it cannot fit depends on configuration: with `n_gpu_layers` left to Ollama it spills
|
||||
layers to the CPU and reports `size_vram < size` (roughly an order of magnitude slower,
|
||||
and silent). With `n_gpu_layers` pinned — 99 on this box — it refuses outright with
|
||||
`cudaMalloc failed: out of memory`. Both are handled: the spill triggers a reclaim from
|
||||
an idle ComfyUI, and the hard failure is caught by `switch_ollama_model`, which reclaims
|
||||
and retries once. Measured: a 12.87 GB model that returned HTTP 500 from Ollama directly
|
||||
now loads through HyperSwap after reclaiming 6.83 GB, at 3.85 GB/s.
|
||||
* **A busy LLM is not a failed yield.** A model mid-generation cannot unload; the
|
||||
`keep_alive: 0` request queues behind it and applies when it finishes. That is reported
|
||||
as `busy` (returning in ~610 ms) rather than blocking, with per-model backoff and a
|
||||
detached watcher that logs the eventual release. Only VRAM held while the GPU sits
|
||||
*idle* counts as a fault.
|
||||
* **Confirmed Soft-Yield (barrier, not fire-and-forget)**: Releases Ollama VRAM allocations (`keep_alive: 0`) down to 0 MB, then **waits on NVML until the driver has actually freed the allocation** before letting ComfyUI proceed. Posting `keep_alive: 0` only *asks* Ollama to unload; on this box the HTTP call returns in ~63 ms while the driver takes a further ~77 ms to release 14.9 GB. Returning during that window is how diffusion ends up allocating into VRAM that is still occupied.
|
||||
* **Idle-Aware ComfyUI Purge**: Diffusion checkpoints are held for `COMFY_IDLE_PURGE_S` (30 s) of genuinely empty queue rather than purged 1.5 s after every prompt — iterating on a workflow no longer pays a full checkpoint reload per run. An immediate purge still happens the moment Ollama actually asks for VRAM (`POST /api/request-vram`).
|
||||
* **Real-Time ComfyUI WebSocket & Watchdog Listener**: Subscribes directly to `ws://127.0.0.1:8188/ws`. The WebSocket is the primary signal; a connection-pooled watchdog polls `/queue` at 1 Hz purely as a fallback, backing off to 3 s while the socket is healthy.
|
||||
|
||||
@@ -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})
|
||||
|
||||
|
||||
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)
|
||||
|
||||
@@ -142,6 +142,21 @@ class _Writer(threading.Thread):
|
||||
conn.commit()
|
||||
except Exception as e:
|
||||
logger.debug("prune failed: %s", e)
|
||||
# Drain anything still queued before closing. Without this, rows submitted but
|
||||
# not yet dequeued are lost on shutdown -- which is exactly when the last events
|
||||
# before a restart matter most.
|
||||
try:
|
||||
while True:
|
||||
try:
|
||||
item = self.q.get_nowait()
|
||||
except queue.Empty:
|
||||
break
|
||||
if item is None:
|
||||
continue
|
||||
_kind, sql, params = item
|
||||
pending.append((sql, params))
|
||||
except Exception:
|
||||
pass
|
||||
try:
|
||||
for sql, params in pending:
|
||||
conn.execute(sql, params)
|
||||
|
||||
74
tests/README.md
Normal file
74
tests/README.md
Normal file
@@ -0,0 +1,74 @@
|
||||
# HyperSwap test suite
|
||||
|
||||
Fast, hermetic unit tests. No GPU is touched, no network call is made, no systemd unit is
|
||||
poked, and the production `hyperswap.db` is never opened.
|
||||
|
||||
## Running
|
||||
|
||||
```bash
|
||||
/home/drjones/comfy-mcp-venv/bin/python -m pytest tests/ -q
|
||||
```
|
||||
|
||||
Single file / single test:
|
||||
|
||||
```bash
|
||||
/home/drjones/comfy-mcp-venv/bin/python -m pytest tests/test_classify_load.py -q
|
||||
/home/drjones/comfy-mcp-venv/bin/python -m pytest tests/ -q -k warm_confident
|
||||
```
|
||||
|
||||
Whole suite runs in about 3 seconds.
|
||||
|
||||
## Safety rails
|
||||
|
||||
These matter, because this repo drives a live 4080 SUPER that a running service is using.
|
||||
|
||||
* `tests/conftest.py` installs an **autouse** `no_gpu_mutation` fixture that replaces
|
||||
`overclock_manager._sh` (the single choke point for every `nvidia-smi` /
|
||||
`nvidia-settings` write) plus `apply_profile`, `apply_fan_control`, `set_fan_speed`,
|
||||
`set_fan_auto` and `restore_safe` with recording stubs. Even a test that accidentally
|
||||
reaches an actuation path can only reach the stub. The fixture yields a dict of
|
||||
recorded calls, which the thermal tests assert against.
|
||||
* `HYPERSWAP_DB` is set to a non-existent path before `telemetry_store` is imported, so no
|
||||
import can bind `DB_PATH` to the production database. Tests that need a DB use the
|
||||
`temp_db` fixture, which monkeypatches `telemetry_store.DB_PATH` to a `tmp_path` file
|
||||
and stops the writer thread afterwards.
|
||||
* All file IO happens against files the tests create in `tmp_path`. No real model blob is
|
||||
read and `warm_file_to_ram` is never called.
|
||||
* Nothing sweeps, and nothing sends HTTP to Ollama, ComfyUI or :9090.
|
||||
|
||||
## Measured constants pinned here
|
||||
|
||||
These numbers came from measurement on this box, not from taste. If a change makes one of
|
||||
these tests fail, the constant is probably wrong, not the test.
|
||||
|
||||
| Constant | Value | Where pinned |
|
||||
| --- | --- | --- |
|
||||
| Cold load of a 12.87 GB model, 3.1% resident | 34267 ms → 0.38 GB/s | `test_classify_load.py::test_measured_cold_load_classifies_as_cold_disk` |
|
||||
| Warm load of the same model, 100% resident | 4901 ms → 2.63 GB/s | `test_classify_load.py::test_measured_warm_load_classifies_as_ram_hit` |
|
||||
| `RAM_HIT_GBPS = 2.0` must stay below the fastest achievable warm load (2.63 GB/s) | — | `test_classify_load.py::test_ram_hit_threshold_is_physically_achievable` |
|
||||
| `PARTIAL_HIT_GBPS = 0.8` must stay above the measured cold rate (0.38 GB/s) | — | same test |
|
||||
| Size-unknown fallback splits at 8000 ms (between 4.9 s warm and 34.3 s cold) | — | `test_classify_load.py::test_unknown_size_guess_boundary_is_8s` |
|
||||
| `WARM_SKIP_THRESHOLD_PCT = 90.0` | — | `test_ram_optimizer.py::test_warm_skip_threshold_constant_unchanged` |
|
||||
| A *probe* reading may only be trusted at exactly 100% (a 12-window probe once cleared 90% on a mostly-cold 12.87 GB blob that then loaded at 2.44 GB/s) | — | `test_ram_optimizer.py::test_probe_reading_is_only_trusted_at_exactly_100_percent` |
|
||||
| `PROBE_CACHED_GBPS = 1.5` sits in the gap between cold NVMe (0.35–0.5 GB/s) and page cache (3.2–13 GB/s) | — | `test_ram_optimizer.py::test_probe_cached_threshold_sits_between_measured_disk_and_cache_rates` |
|
||||
| Card power envelope: 320 W stock, 370 W max, sweeps never go below 60% of max | — | `test_autotune_helpers.py::test_supported_power_limits_parses_min_default_max` |
|
||||
| `_supported_clocks` must always query the `mem,gr` pair (a single-field query returns one column and silently yielded `[]`) | — | `test_autotune_helpers.py::test_supported_clocks_always_queries_the_mem_gr_pair` |
|
||||
| ComfyUI benchmark seed must vary per call (a fixed seed made ComfyUI serve a cached result in ~1 ms) | — | `test_autotune_helpers.py::test_comfy_workflow_seed_varies_between_calls` |
|
||||
| Governor hysteresis: `HOT_SAMPLES = 5`, `COOL_SAMPLES = 30`, `REAPPLY_COOLDOWN_S = 20` | — | `test_thermal_governor.py` (escalation, recovery, cooldown, alternating-sample tests) |
|
||||
| Model usage score: frequency decayed with a ~24 h half-life | — | `test_telemetry_store.py::test_model_usage_ranking_scores_recent_use_higher` |
|
||||
|
||||
## What is deliberately not covered
|
||||
|
||||
* `vram_arbitrator.instant_free_ollama_vram`, the `AutoArbitrator` yield/purge paths and
|
||||
the SSE broker — under active edit, contract changing.
|
||||
* `overclock_manager.apply_profile` and every other actuation path, `autotune.sweep`,
|
||||
`ram_optimizer.warm_file_to_ram` — these mutate hardware or do heavy IO.
|
||||
* `server.py` HTTP routes and `mcp_server.py` — would need the app wired to live
|
||||
subsystems.
|
||||
|
||||
## Known rough edge the tests work around
|
||||
|
||||
`telemetry_store.stop()` flushes the writer's pending *batch* but does not drain the
|
||||
submission queue, so a `stop()` racing a just-submitted row can drop it. The writer tests
|
||||
call a local `_drain()` helper to wait for the queue to empty before stopping, rather than
|
||||
encoding the race into an assertion.
|
||||
69
tests/conftest.py
Normal file
69
tests/conftest.py
Normal file
@@ -0,0 +1,69 @@
|
||||
"""Shared fixtures and — more importantly — hardware safety rails for the suite.
|
||||
|
||||
This repo drives a live GPU and a running systemd service. Every test here must be
|
||||
hermetic: no NVML mutation, no nvidia-smi/nvidia-settings writes, no touching the
|
||||
production telemetry DB, no HTTP to Ollama/ComfyUI/:9090.
|
||||
|
||||
The `no_gpu_mutation` fixture below is autouse, so even a test that accidentally
|
||||
reaches an actuation path can only reach a recording stub.
|
||||
"""
|
||||
import os
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
|
||||
REPO_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
||||
if REPO_ROOT not in sys.path:
|
||||
sys.path.insert(0, REPO_ROOT)
|
||||
|
||||
# telemetry_store resolves DB_PATH from the environment *at import time*. Point it at a
|
||||
# path that does not exist before anything imports it, so no import of this suite can
|
||||
# ever open the production hyperswap.db. Individual tests monkeypatch DB_PATH to a
|
||||
# tmp_path file when they actually need a database.
|
||||
os.environ.setdefault("HYPERSWAP_DB", os.path.join(REPO_ROOT, "tests", "_never_created.db"))
|
||||
|
||||
import overclock_manager # noqa: E402 (must follow the sys.path/env setup above)
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def no_gpu_mutation(monkeypatch):
|
||||
"""Hard block on every code path that can physically change GPU state.
|
||||
|
||||
Yields a dict of call recorders so a test can assert that actuation *would* have
|
||||
happened without any of it reaching the card.
|
||||
"""
|
||||
calls = {"apply_profile": [], "fan": [], "restore_safe": [], "sh": []}
|
||||
|
||||
def _blocked_sh(cmd, use_sudo=True, timeout=10):
|
||||
# Catch-all: every nvidia-smi / nvidia-settings write in overclock_manager
|
||||
# funnels through _sh. Nothing in the suite may shell out to the driver.
|
||||
calls["sh"].append(list(cmd))
|
||||
return {"rc": -1, "out": "", "err": "blocked by test suite"}
|
||||
|
||||
monkeypatch.setattr(overclock_manager, "_sh", _blocked_sh)
|
||||
monkeypatch.setattr(overclock_manager, "apply_profile",
|
||||
lambda name, overrides=None: calls["apply_profile"].append((name, overrides)))
|
||||
monkeypatch.setattr(overclock_manager, "apply_fan_control",
|
||||
lambda mode, speed_pct: calls["fan"].append((mode, speed_pct)))
|
||||
monkeypatch.setattr(overclock_manager, "set_fan_speed",
|
||||
lambda percent: calls["fan"].append(("manual", percent)))
|
||||
monkeypatch.setattr(overclock_manager, "set_fan_auto",
|
||||
lambda: calls["fan"].append(("auto", None)))
|
||||
monkeypatch.setattr(overclock_manager, "restore_safe",
|
||||
lambda reason="shutdown": calls["restore_safe"].append(reason))
|
||||
return calls
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def temp_db(tmp_path, monkeypatch):
|
||||
"""Point telemetry_store at a throwaway SQLite file for the duration of one test."""
|
||||
import telemetry_store
|
||||
|
||||
db = tmp_path / "test_hyperswap.db"
|
||||
monkeypatch.setattr(telemetry_store, "DB_PATH", str(db))
|
||||
yield str(db)
|
||||
# Never leave a writer thread running against a tmp path that is about to vanish.
|
||||
try:
|
||||
telemetry_store.stop()
|
||||
except Exception:
|
||||
pass
|
||||
184
tests/test_autotune_helpers.py
Normal file
184
tests/test_autotune_helpers.py
Normal file
@@ -0,0 +1,184 @@
|
||||
"""autotune's pure/parsable helpers.
|
||||
|
||||
Only the parsing and sampling helpers are exercised. Nothing here runs a sweep, applies a
|
||||
profile, or talks to nvidia-smi: subprocess.run is monkeypatched at the module boundary.
|
||||
"""
|
||||
import pytest
|
||||
|
||||
import autotune
|
||||
|
||||
|
||||
class _FakeProc:
|
||||
def __init__(self, stdout="", returncode=0, stderr=""):
|
||||
self.stdout = stdout
|
||||
self.stderr = stderr
|
||||
self.returncode = returncode
|
||||
|
||||
|
||||
def _fake_smi(monkeypatch, stdout, record=None):
|
||||
def _run(cmd, capture_output=True, text=True, timeout=None, **kw):
|
||||
if record is not None:
|
||||
record.append(list(cmd))
|
||||
return _FakeProc(stdout)
|
||||
monkeypatch.setattr(autotune.subprocess, "run", _run)
|
||||
|
||||
|
||||
# A trimmed but realistically shaped `nvidia-smi --query-supported-clocks=mem,gr` dump:
|
||||
# graphics clocks are enumerated once per memory clock.
|
||||
SUPPORTED_CLOCKS_CSV = """\
|
||||
10501, 2790
|
||||
10501, 2775
|
||||
10501, 2760
|
||||
9501, 2790
|
||||
9501, 2775
|
||||
405, 645
|
||||
405, 630
|
||||
"""
|
||||
|
||||
|
||||
def test_supported_clocks_mem_returns_sorted_unique_memory_clocks(monkeypatch):
|
||||
_fake_smi(monkeypatch, SUPPORTED_CLOCKS_CSV)
|
||||
assert autotune._supported_clocks("mem") == [405, 9501, 10501]
|
||||
|
||||
|
||||
def test_supported_clocks_gr_returns_clocks_of_the_highest_memory_clock(monkeypatch):
|
||||
"""Graphics clocks are enumerated per memory clock. Only the set belonging to the top
|
||||
memory clock is meaningful — that is the state any real workload runs at."""
|
||||
_fake_smi(monkeypatch, SUPPORTED_CLOCKS_CSV)
|
||||
assert autotune._supported_clocks("gr") == [2760, 2775, 2790]
|
||||
|
||||
|
||||
def test_supported_clocks_always_queries_the_mem_gr_pair(monkeypatch):
|
||||
"""Regression: querying a single field returns one column, and reading index 1 from it
|
||||
silently produced an empty list instead of an error. The query must stay a pair."""
|
||||
seen = []
|
||||
_fake_smi(monkeypatch, SUPPORTED_CLOCKS_CSV, record=seen)
|
||||
autotune._supported_clocks("gr")
|
||||
assert any("--query-supported-clocks=mem,gr" in part for part in seen[0])
|
||||
|
||||
|
||||
def test_supported_clocks_ignores_single_column_output(monkeypatch):
|
||||
"""If the driver ever returns one column, no row parses and we return [] — never a
|
||||
list of memory clocks mislabelled as graphics clocks."""
|
||||
_fake_smi(monkeypatch, "10501\n9501\n405\n")
|
||||
assert autotune._supported_clocks("mem") == []
|
||||
assert autotune._supported_clocks("gr") == []
|
||||
|
||||
|
||||
def test_supported_clocks_skips_header_and_garbage_rows(monkeypatch):
|
||||
_fake_smi(monkeypatch, "memory [MHz], graphics [MHz]\n10501, 2790\n\nN/A, N/A\n")
|
||||
assert autotune._supported_clocks("mem") == [10501]
|
||||
|
||||
|
||||
def test_supported_clocks_returns_empty_when_nvidia_smi_fails(monkeypatch):
|
||||
"""No GPU / no driver must degrade to 'no candidates', not an exception on import of
|
||||
a sweep plan."""
|
||||
def _boom(*a, **kw):
|
||||
raise OSError("nvidia-smi not found")
|
||||
monkeypatch.setattr(autotune.subprocess, "run", _boom)
|
||||
assert autotune._supported_clocks("mem") == []
|
||||
|
||||
|
||||
# --------------------------------------------------------------- power limits
|
||||
|
||||
def test_supported_power_limits_parses_min_default_max(monkeypatch):
|
||||
"""RTX 4080 SUPER on this box: 115 W min, 370 W max, 320 W stock default."""
|
||||
_fake_smi(monkeypatch, "115.00, 370.00, 320.00\n")
|
||||
vals = autotune._supported_power_limits(steps=5)
|
||||
assert vals == sorted(set(vals))
|
||||
assert max(vals) == 370
|
||||
# Never sweeps below 60% of max — the card does no useful work down there.
|
||||
assert min(vals) >= int(370 * 0.6)
|
||||
# The stock default is always included as a reference point.
|
||||
assert 320 in vals
|
||||
|
||||
|
||||
def test_supported_power_limits_respects_step_count(monkeypatch):
|
||||
_fake_smi(monkeypatch, "115.00, 370.00, 320.00\n")
|
||||
vals = autotune._supported_power_limits(steps=3)
|
||||
assert len(vals) <= 4 # 3 evenly spaced values, plus the default if it is not one
|
||||
|
||||
|
||||
def test_supported_power_limits_returns_empty_on_query_failure(monkeypatch):
|
||||
_fake_smi(monkeypatch, "N/A, N/A, N/A\n")
|
||||
assert autotune._supported_power_limits() == []
|
||||
|
||||
|
||||
# --------------------------------------------------------------- subsampling
|
||||
|
||||
def test_subsample_returns_input_when_already_short_enough():
|
||||
assert autotune._subsample([1, 2, 3], 5) == [1, 2, 3]
|
||||
|
||||
|
||||
def test_subsample_keeps_both_endpoints():
|
||||
"""The endpoints are the whole point of a sweep: stock and maximum. Dropping either
|
||||
means never measuring the setting the sweep is supposed to recommend."""
|
||||
values = list(range(0, 195)) # the card enumerates ~194 graphics clocks
|
||||
out = autotune._subsample(values, 8)
|
||||
assert out[0] == values[0]
|
||||
assert out[-1] == values[-1]
|
||||
|
||||
|
||||
def test_subsample_never_exceeds_max_steps():
|
||||
values = list(range(0, 195))
|
||||
for max_steps in (2, 3, 5, 8, 13):
|
||||
assert len(autotune._subsample(values, max_steps)) <= max_steps
|
||||
|
||||
|
||||
def test_subsample_output_is_sorted_and_unique():
|
||||
values = list(range(0, 100))
|
||||
out = autotune._subsample(values, 7)
|
||||
assert out == sorted(set(out))
|
||||
|
||||
|
||||
def test_subsample_is_evenly_spread():
|
||||
"""Clustered samples would tell us nothing about the middle of the range."""
|
||||
out = autotune._subsample(list(range(0, 101)), 5)
|
||||
gaps = [b - a for a, b in zip(out, out[1:])]
|
||||
assert max(gaps) - min(gaps) <= 1
|
||||
|
||||
|
||||
# --------------------------------------------------------------- comfy benchmark graph
|
||||
|
||||
def test_comfy_workflow_seed_varies_between_calls():
|
||||
"""Fixed seeds made ComfyUI serve a cached result in ~1 ms without executing the
|
||||
graph, so every benchmark after the first measured nothing at all."""
|
||||
seeds = {autotune._comfy_workflow()["5"]["inputs"]["seed"] for _ in range(25)}
|
||||
assert len(seeds) > 20
|
||||
|
||||
|
||||
def test_comfy_workflow_seed_can_be_pinned_for_reproduction():
|
||||
assert autotune._comfy_workflow(seed=42)["5"]["inputs"]["seed"] == 42
|
||||
|
||||
|
||||
def test_comfy_workflow_shape_is_otherwise_constant():
|
||||
"""Only the seed may vary; a benchmark whose step count or resolution moved between
|
||||
runs would not be comparable."""
|
||||
a = autotune._comfy_workflow()
|
||||
b = autotune._comfy_workflow()
|
||||
for wf in (a, b):
|
||||
assert wf["5"]["inputs"]["steps"] == autotune.COMFY_BENCH_STEPS
|
||||
assert wf["4"]["inputs"]["width"] == autotune.COMFY_BENCH_SIZE
|
||||
assert wf["4"]["inputs"]["height"] == autotune.COMFY_BENCH_SIZE
|
||||
a["5"]["inputs"]["seed"] = b["5"]["inputs"]["seed"] = 0
|
||||
assert a == b
|
||||
|
||||
|
||||
def test_comfy_workflow_uses_the_named_checkpoint():
|
||||
wf = autotune._comfy_workflow(ckpt="some_other.safetensors")
|
||||
assert wf["1"]["inputs"]["ckpt_name"] == "some_other.safetensors"
|
||||
|
||||
|
||||
def test_temp_ceiling_is_below_the_thermal_governors_escalation_point():
|
||||
"""A sweep step must abort on temperature before the governor starts derating under
|
||||
it, otherwise the sweep measures the governor's derate rather than the knob."""
|
||||
import thermal_governor as tg
|
||||
assert autotune.TEMP_CEILING_C <= tg.TEMP_ESCALATE_C + 1.0
|
||||
|
||||
|
||||
@pytest.mark.parametrize("knob", ["mem_offset_mhz", "core_offset_mhz", "lock_mem_mhz",
|
||||
"lock_core_max", "power_limit_w"])
|
||||
def test_every_knob_declares_a_hardware_verification_field(knob):
|
||||
"""Offsets are silently ignored by some drivers (595.84 accepts an assignment and
|
||||
reads back a different value), so each knob must name the field to read back."""
|
||||
assert autotune.KNOBS[knob]["verify"]
|
||||
123
tests/test_classify_load.py
Normal file
123
tests/test_classify_load.py
Normal file
@@ -0,0 +1,123 @@
|
||||
"""Regression tests for vram_arbitrator.classify_load.
|
||||
|
||||
classify_load turns (model size, Ollama's reported load_duration) into a cache verdict
|
||||
by computing achieved bandwidth. The two thresholds it compares against were measured on
|
||||
this box, not chosen; the tests below pin the measurements themselves so a future edit
|
||||
to RAM_HIT_GBPS / PARTIAL_HIT_GBPS that breaks the real data points fails loudly.
|
||||
"""
|
||||
import math
|
||||
|
||||
import pytest
|
||||
|
||||
import vram_arbitrator as va
|
||||
|
||||
GIB = 1024 ** 3
|
||||
|
||||
# Ground truth, measured 2026-08-28 on the same 12.87 GB model loaded twice:
|
||||
# 3.1% resident -> 34267 ms -> 0.38 GB/s -> cold
|
||||
# 100% resident -> 4901 ms -> 2.63 GB/s -> RAM hit
|
||||
MEASURED_MODEL_BYTES = int(12.87 * GIB)
|
||||
MEASURED_COLD_MS = 34267.0
|
||||
MEASURED_WARM_MS = 4901.0
|
||||
|
||||
|
||||
def test_measured_cold_load_classifies_as_cold_disk():
|
||||
"""The measured cold load (12.87 GB, 34267 ms, 0.38 GB/s) must stay 'Cold Disk Load'."""
|
||||
res = va.classify_load(MEASURED_MODEL_BYTES, MEASURED_COLD_MS)
|
||||
assert res["cache_status"] == "Cold Disk Load 💾"
|
||||
assert res["is_ram_hit"] is False
|
||||
assert res["load_gbps"] == pytest.approx(0.38, abs=0.01)
|
||||
|
||||
|
||||
def test_measured_warm_load_classifies_as_ram_hit():
|
||||
"""The measured warm load (12.87 GB, 4901 ms, 2.63 GB/s) must stay a RAM cache hit."""
|
||||
res = va.classify_load(MEASURED_MODEL_BYTES, MEASURED_WARM_MS)
|
||||
assert res["cache_status"] == "RAM Cache Hit ⚡"
|
||||
assert res["is_ram_hit"] is True
|
||||
assert res["load_gbps"] == pytest.approx(2.63, abs=0.01)
|
||||
|
||||
|
||||
def test_ram_hit_threshold_is_physically_achievable():
|
||||
"""Guard against the reintroduced bug where the hit bar sat above any real warm load.
|
||||
|
||||
A prior version set the hit threshold at 5.0 GB/s. Ollama's load_duration covers
|
||||
host-to-device transfer and model init as well as the file read, so a fully resident
|
||||
12.87 GB model still only reports ~2.6 GB/s — meaning *no* load could ever be
|
||||
classified as a hit. Any threshold that a genuinely warm load cannot clear is wrong.
|
||||
"""
|
||||
assert va.RAM_HIT_GBPS <= 2.63, (
|
||||
"RAM_HIT_GBPS is above the fastest warm load ever measured on this box "
|
||||
"(2.63 GB/s); no load could be classified as a cache hit."
|
||||
)
|
||||
assert va.PARTIAL_HIT_GBPS > 0.38, (
|
||||
"PARTIAL_HIT_GBPS is at or below the measured cold-disk rate; cold loads would "
|
||||
"be reported as partial cache hits."
|
||||
)
|
||||
assert va.PARTIAL_HIT_GBPS < va.RAM_HIT_GBPS
|
||||
|
||||
|
||||
def test_sub_millisecond_load_is_already_in_vram():
|
||||
"""load_duration_ms <= 1.0 means Ollama never re-read the model at all."""
|
||||
res = va.classify_load(MEASURED_MODEL_BYTES, 1.0)
|
||||
assert res["cache_status"] == "Already in VRAM"
|
||||
assert res["load_gbps"] is None
|
||||
assert res["is_ram_hit"] is True
|
||||
|
||||
|
||||
def test_zero_duration_is_already_in_vram():
|
||||
res = va.classify_load(MEASURED_MODEL_BYTES, 0.0)
|
||||
assert res["cache_status"] == "Already in VRAM"
|
||||
|
||||
|
||||
def test_boundary_exactly_at_ram_hit_threshold_is_a_hit():
|
||||
"""Exactly RAM_HIT_GBPS (2.0 GB/s) must count as a hit — the comparison is >=."""
|
||||
size = int(va.RAM_HIT_GBPS * GIB) # 2 GiB read in exactly 1000 ms -> 2.00 GB/s
|
||||
res = va.classify_load(size, 1000.0)
|
||||
assert res["load_gbps"] == pytest.approx(2.0)
|
||||
assert res["cache_status"] == "RAM Cache Hit ⚡"
|
||||
assert res["is_ram_hit"] is True
|
||||
|
||||
|
||||
def test_just_below_ram_hit_threshold_is_partial():
|
||||
size = int(1.99 * GIB)
|
||||
res = va.classify_load(size, 1000.0)
|
||||
assert res["cache_status"] == "Partial Cache 🌤"
|
||||
assert res["is_ram_hit"] is False
|
||||
|
||||
|
||||
def test_boundary_exactly_at_partial_threshold_is_partial():
|
||||
"""Exactly PARTIAL_HIT_GBPS (0.8 GB/s) is still a partial hit, not a cold load."""
|
||||
size = math.ceil(va.PARTIAL_HIT_GBPS * GIB) # 0.8 GiB is not a whole number of bytes
|
||||
res = va.classify_load(size, 1000.0)
|
||||
assert res["load_gbps"] == pytest.approx(0.8)
|
||||
assert res["cache_status"] == "Partial Cache 🌤"
|
||||
assert res["is_ram_hit"] is False
|
||||
|
||||
|
||||
def test_just_below_partial_threshold_is_cold():
|
||||
size = int(0.79 * GIB)
|
||||
res = va.classify_load(size, 1000.0)
|
||||
assert res["cache_status"] == "Cold Disk Load 💾"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("duration_ms,expected_hit", [(4901.0, True), (34267.0, False)])
|
||||
def test_unknown_size_falls_back_to_duration_guess_and_says_so(duration_ms, expected_hit):
|
||||
"""With no size on record there is no bandwidth to compute, so the result must be
|
||||
labelled as a guess rather than presented as a measurement."""
|
||||
res = va.classify_load(0, duration_ms)
|
||||
assert res["is_ram_hit"] is expected_hit
|
||||
assert res["load_gbps"] is None
|
||||
assert "detail" in res and "guess" in res["detail"].lower()
|
||||
|
||||
|
||||
def test_unknown_size_guess_boundary_is_8s():
|
||||
"""The size-unknown fallback splits at 8000 ms, between the measured 4.9 s warm and
|
||||
34.3 s cold loads."""
|
||||
assert va.classify_load(0, 7999.0)["is_ram_hit"] is True
|
||||
assert va.classify_load(0, 8000.0)["is_ram_hit"] is False
|
||||
|
||||
|
||||
def test_classification_never_raises_on_odd_inputs():
|
||||
"""This runs on the swap hot path; it must not be able to throw."""
|
||||
for size, dur in [(0, 0.0), (1, 1.5), (10 ** 12, 2.0), (0, 1.0)]:
|
||||
assert "cache_status" in va.classify_load(size, dur)
|
||||
116
tests/test_overclock_profiles.py
Normal file
116
tests/test_overclock_profiles.py
Normal file
@@ -0,0 +1,116 @@
|
||||
"""overclock_manager profile storage: load/merge/validate only.
|
||||
|
||||
Nothing here applies a profile or shells out. CONFIG_PATH is redirected to tmp_path so the
|
||||
repo's real overclock_profiles.json is never written, and conftest's autouse fixture blocks
|
||||
_sh() and every actuation entry point.
|
||||
"""
|
||||
import json
|
||||
|
||||
import pytest
|
||||
|
||||
import overclock_manager as ocm
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def config(tmp_path, monkeypatch):
|
||||
"""Redirect the profile store to a temp file. Returns a writer for its contents."""
|
||||
path = tmp_path / "overclock_profiles.json"
|
||||
monkeypatch.setattr(ocm, "CONFIG_PATH", str(path))
|
||||
|
||||
def write(data):
|
||||
path.write_text(json.dumps(data))
|
||||
return type("Cfg", (), {"path": path, "write": staticmethod(write)})
|
||||
|
||||
|
||||
def test_defaults_are_returned_when_no_config_file_exists(config):
|
||||
profiles = ocm.load_profiles()
|
||||
assert set(profiles) == set(ocm.DEFAULT_PROFILES)
|
||||
assert profiles["ollama"]["power_limit_w"] == ocm.DEFAULT_PROFILES["ollama"]["power_limit_w"]
|
||||
|
||||
|
||||
def test_stored_values_override_defaults_key_by_key(config):
|
||||
"""Only the keys present on disk change; the rest of the default profile survives.
|
||||
A whole-profile replacement would silently drop keys added by a later version."""
|
||||
config.write({"ollama": {"power_limit_w": 340}})
|
||||
profiles = ocm.load_profiles()
|
||||
assert profiles["ollama"]["power_limit_w"] == 340
|
||||
assert profiles["ollama"]["label"] == ocm.DEFAULT_PROFILES["ollama"]["label"]
|
||||
assert profiles["ollama"]["mem_offset_mhz"] == ocm.DEFAULT_PROFILES["ollama"]["mem_offset_mhz"]
|
||||
|
||||
|
||||
def test_unknown_stored_profiles_are_preserved(config):
|
||||
"""A user-created profile is not in DEFAULT_PROFILES and must not be dropped on load."""
|
||||
config.write({"my-custom": {"label": "mine", "power_limit_w": 300}})
|
||||
profiles = ocm.load_profiles()
|
||||
assert profiles["my-custom"]["power_limit_w"] == 300
|
||||
assert set(ocm.DEFAULT_PROFILES) <= set(profiles)
|
||||
|
||||
|
||||
def test_loading_does_not_mutate_the_defaults(config):
|
||||
"""load_profiles deep-copies DEFAULT_PROFILES; if it did not, one load with a stored
|
||||
override would poison the defaults for the rest of the process."""
|
||||
config.write({"ollama": {"power_limit_w": 111}})
|
||||
ocm.load_profiles()
|
||||
config.write({})
|
||||
assert ocm.load_profiles()["ollama"]["power_limit_w"] == \
|
||||
ocm.DEFAULT_PROFILES["ollama"]["power_limit_w"]
|
||||
assert ocm.DEFAULT_PROFILES["ollama"]["power_limit_w"] != 111
|
||||
|
||||
|
||||
def test_corrupt_config_falls_back_to_defaults(config):
|
||||
"""A half-written JSON file must not take the whole service down at startup."""
|
||||
config.path.write_text("{ not valid json")
|
||||
profiles = ocm.load_profiles()
|
||||
assert set(profiles) == set(ocm.DEFAULT_PROFILES)
|
||||
|
||||
|
||||
def test_set_profile_rejects_an_unknown_name(config):
|
||||
"""set_profile edits an existing profile; it is not a create-anything endpoint."""
|
||||
res = ocm.set_profile("no-such-profile", {"power_limit_w": 400})
|
||||
assert res["success"] is False
|
||||
assert "unknown profile" in res["error"]
|
||||
assert not config.path.exists()
|
||||
|
||||
|
||||
def test_set_profile_persists_a_partial_update(config):
|
||||
res = ocm.set_profile("comfy", {"power_limit_w": 350})
|
||||
assert res["success"] is True
|
||||
stored = json.loads(config.path.read_text())
|
||||
assert stored["comfy"]["power_limit_w"] == 350
|
||||
# The other profiles are written back untouched.
|
||||
assert stored["ollama"]["power_limit_w"] == ocm.DEFAULT_PROFILES["ollama"]["power_limit_w"]
|
||||
assert ocm.load_profiles()["comfy"]["power_limit_w"] == 350
|
||||
|
||||
|
||||
def test_set_profile_can_edit_a_previously_stored_custom_profile(config):
|
||||
config.write({"my-custom": {"label": "mine", "power_limit_w": 300}})
|
||||
assert ocm.set_profile("my-custom", {"power_limit_w": 310})["success"] is True
|
||||
assert ocm.load_profiles()["my-custom"]["power_limit_w"] == 310
|
||||
|
||||
|
||||
def test_get_profiles_matches_load_profiles(config):
|
||||
config.write({"ollama": {"core_offset_mhz": 42}})
|
||||
assert ocm.get_profiles() == ocm.load_profiles()
|
||||
|
||||
|
||||
def test_save_profiles_reports_failure_instead_of_raising(tmp_path, monkeypatch):
|
||||
"""The dashboard calls this; an unwritable config must surface as success=False."""
|
||||
monkeypatch.setattr(ocm, "CONFIG_PATH", str(tmp_path / "no-such-dir" / "p.json"))
|
||||
assert ocm.save_profiles({"ollama": {}}) is False
|
||||
|
||||
|
||||
@pytest.mark.parametrize("name", ["ollama", "comfy", "balanced"])
|
||||
def test_default_profiles_declare_the_keys_the_appliers_read(name):
|
||||
"""apply_profile and the thermal governor both index these keys directly; a profile
|
||||
missing one would KeyError inside the actuation thread."""
|
||||
cfg = ocm.DEFAULT_PROFILES[name]
|
||||
for key in ("label", "power_limit_w", "core_offset_mhz", "mem_offset_mhz",
|
||||
"lock_core_min", "lock_core_max"):
|
||||
assert key in cfg, f"{name} is missing {key}"
|
||||
|
||||
|
||||
def test_default_power_limits_stay_within_the_cards_range():
|
||||
"""RTX 4080 SUPER: 320 W stock, 370 W maximum. A profile above that is silently
|
||||
clamped by the driver and the sweep results become meaningless."""
|
||||
for name, cfg in ocm.DEFAULT_PROFILES.items():
|
||||
assert 100 <= cfg["power_limit_w"] <= 370, name
|
||||
421
tests/test_ram_optimizer.py
Normal file
421
tests/test_ram_optimizer.py
Normal file
@@ -0,0 +1,421 @@
|
||||
"""ram_optimizer: page-cache residency measurement, model discovery, warm planning.
|
||||
|
||||
All file IO here happens against files this test creates in tmp_path. Nothing reads a
|
||||
real model blob, and nothing calls warm_file_to_ram.
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
|
||||
import pytest
|
||||
|
||||
import ram_optimizer as ro
|
||||
|
||||
MIB = 1024 * 1024
|
||||
|
||||
|
||||
# --------------------------------------------------------------- warm_confident rule
|
||||
|
||||
class _FakeCachestat:
|
||||
"""Stand-in for the kernel's struct cachestat."""
|
||||
|
||||
def __init__(self, nr_cache):
|
||||
self.nr_cache = nr_cache
|
||||
self.nr_dirty = 0
|
||||
self.nr_evicted = 0
|
||||
|
||||
|
||||
def _file_of_pages(tmp_path, pages):
|
||||
p = tmp_path / f"blob_{pages}.bin"
|
||||
p.write_bytes(b"\0" * (pages * ro.PAGE_SIZE))
|
||||
return str(p)
|
||||
|
||||
|
||||
def _force_cachestat(monkeypatch, resident_pages):
|
||||
monkeypatch.setattr(ro, "_cachestat",
|
||||
lambda fd, offset, length: _FakeCachestat(resident_pages))
|
||||
|
||||
|
||||
def _force_probe(monkeypatch, pct):
|
||||
monkeypatch.setattr(ro, "_cachestat", lambda fd, offset, length: None)
|
||||
monkeypatch.setattr(ro, "_throughput_probe",
|
||||
lambda fd, size, windows=None: {"resident_pct": pct, "windows": 12,
|
||||
"median_gbps": 3.0})
|
||||
|
||||
|
||||
@pytest.mark.parametrize("resident_pages,expect_confident", [
|
||||
(100, True), # 100% exact
|
||||
(90, True), # exactly WARM_SKIP_THRESHOLD_PCT
|
||||
(89, False), # just under the threshold
|
||||
(0, False),
|
||||
])
|
||||
def test_cachestat_reading_is_trusted_at_or_above_threshold(tmp_path, monkeypatch,
|
||||
resident_pages, expect_confident):
|
||||
"""An exact cachestat reading >= WARM_SKIP_THRESHOLD_PCT (90) may be trusted to skip
|
||||
warming; below it, it may not."""
|
||||
path = _file_of_pages(tmp_path, 100)
|
||||
_force_cachestat(monkeypatch, resident_pages)
|
||||
res = ro.page_residency(path)
|
||||
assert res["method"] == "cachestat"
|
||||
assert res["resident_pct"] == pytest.approx(float(resident_pages))
|
||||
assert res["warm_confident"] is expect_confident
|
||||
|
||||
|
||||
@pytest.mark.parametrize("pct,expect_confident", [
|
||||
(100.0, True), # only an unblemished probe is trustworthy
|
||||
(99.9, False),
|
||||
(95.0, False), # would pass `warm`, must NOT pass `warm_confident`
|
||||
(90.0, False),
|
||||
])
|
||||
def test_probe_reading_is_only_trusted_at_exactly_100_percent(tmp_path, monkeypatch,
|
||||
pct, expect_confident):
|
||||
"""Correctness fix, not a style choice: a 12-window probe once cleared 90% on a file
|
||||
that was mostly cold (a 12.87 GB blob reported 'already resident' then loaded at
|
||||
2.44 GB/s). Only a perfect probe score may skip work."""
|
||||
path = _file_of_pages(tmp_path, 100)
|
||||
_force_probe(monkeypatch, pct)
|
||||
res = ro.page_residency(path)
|
||||
assert res["method"] == "probe"
|
||||
assert res["warm_confident"] is expect_confident
|
||||
|
||||
|
||||
def test_probe_at_95_percent_is_warm_but_not_confident(tmp_path, monkeypatch):
|
||||
"""`warm` and `warm_confident` are different questions and must not be conflated:
|
||||
warm drives display, warm_confident drives skipping work."""
|
||||
path = _file_of_pages(tmp_path, 100)
|
||||
_force_probe(monkeypatch, 95.0)
|
||||
res = ro.page_residency(path)
|
||||
assert res["warm"] is True
|
||||
assert res["warm_confident"] is False
|
||||
|
||||
|
||||
def test_warm_skip_threshold_constant_unchanged():
|
||||
# Pinned: build_warm_plan and the dashboard both key off this number.
|
||||
assert ro.WARM_SKIP_THRESHOLD_PCT == 90.0
|
||||
|
||||
|
||||
# --------------------------------------------------------------- page_residency basics
|
||||
|
||||
def test_page_residency_on_missing_file_reports_not_measurable(tmp_path):
|
||||
res = ro.page_residency(str(tmp_path / "nope.bin"))
|
||||
assert res["success"] is False
|
||||
assert res["measurable"] is False
|
||||
assert res["resident_pct"] == 0.0
|
||||
|
||||
|
||||
def test_page_residency_on_empty_file_short_circuits(tmp_path):
|
||||
"""A zero-byte file has no pages to measure; dividing by its page count would throw."""
|
||||
p = tmp_path / "empty.bin"
|
||||
p.write_bytes(b"")
|
||||
res = ro.page_residency(str(p))
|
||||
assert res["method"] == "empty"
|
||||
assert res["resident_pct"] == 0.0
|
||||
assert res["measurable"] is True
|
||||
|
||||
|
||||
def test_page_residency_refuses_to_guess_when_probing_is_disallowed(tmp_path, monkeypatch):
|
||||
"""With cachestat unavailable and allow_probe=False the answer is 'unknown', never 0%.
|
||||
Reporting 0% would make the planner re-warm a file that may be fully resident."""
|
||||
path = _file_of_pages(tmp_path, 8)
|
||||
monkeypatch.setattr(ro, "_cachestat", lambda fd, offset, length: None)
|
||||
res = ro.page_residency(path, allow_probe=False)
|
||||
assert res["measurable"] is False
|
||||
assert res["resident_pct"] is None
|
||||
assert res["warm"] is None
|
||||
assert res["method"] == "unavailable"
|
||||
|
||||
|
||||
def test_page_residency_reports_size_and_bytes_consistently(tmp_path, monkeypatch):
|
||||
path = _file_of_pages(tmp_path, 100)
|
||||
_force_cachestat(monkeypatch, 50)
|
||||
res = ro.page_residency(path)
|
||||
size = os.path.getsize(path)
|
||||
assert res["size_bytes"] == size
|
||||
assert res["resident_bytes"] == int(size * 0.5)
|
||||
|
||||
|
||||
def test_page_residency_against_a_real_file_holds_the_confidence_invariant(tmp_path):
|
||||
"""End-to-end with the real kernel path (whichever method is available here): the
|
||||
warm_confident rule must hold for whatever the machine actually reports."""
|
||||
path = _file_of_pages(tmp_path, 512)
|
||||
res = ro.page_residency(path)
|
||||
assert res["success"] is True
|
||||
assert res["method"] in ("cachestat", "probe")
|
||||
assert 0.0 <= res["resident_pct"] <= 100.0
|
||||
expected = ((res["method"] == "cachestat" and res["resident_pct"] >= 90.0)
|
||||
or (res["method"] == "probe" and res["resident_pct"] >= 100.0))
|
||||
assert res["warm_confident"] is expected
|
||||
|
||||
|
||||
# --------------------------------------------------------------- throughput probe
|
||||
|
||||
def test_throughput_probe_window_count_is_clamped_to_file_size(tmp_path):
|
||||
"""A file smaller than PROBE_WINDOWS * PROBE_WINDOW_BYTES must not be probed with
|
||||
more windows than it has, or offsets run past EOF."""
|
||||
p = tmp_path / "small.bin"
|
||||
p.write_bytes(b"\0" * (5 * MIB))
|
||||
fd = os.open(str(p), os.O_RDONLY)
|
||||
try:
|
||||
out = ro._throughput_probe(fd, 5 * MIB)
|
||||
finally:
|
||||
os.close(fd)
|
||||
assert 0 < out["windows"] <= 5 * MIB // ro.PROBE_WINDOW_BYTES
|
||||
assert 0.0 <= out["resident_pct"] <= 100.0
|
||||
|
||||
|
||||
def test_throughput_probe_honours_window_override(tmp_path):
|
||||
p = tmp_path / "medium.bin"
|
||||
p.write_bytes(b"\0" * (32 * MIB))
|
||||
fd = os.open(str(p), os.O_RDONLY)
|
||||
try:
|
||||
out = ro._throughput_probe(fd, 32 * MIB, windows_override=4)
|
||||
finally:
|
||||
os.close(fd)
|
||||
assert out["windows"] == 4
|
||||
|
||||
|
||||
def test_probe_cached_threshold_sits_between_measured_disk_and_cache_rates():
|
||||
# Measured on this box: cold NVMe 0.35-0.5 GB/s, page cache 3.2-13 GB/s.
|
||||
assert 0.5 < ro.PROBE_CACHED_GBPS < 3.2
|
||||
|
||||
|
||||
# --------------------------------------------------------------- ollama manifest parsing
|
||||
|
||||
MODEL_MEDIA_TYPE = "application/vnd.ollama.image.model"
|
||||
|
||||
|
||||
def _write_manifest(root, rel_dir, tag, layers):
|
||||
d = os.path.join(root, "manifests", rel_dir)
|
||||
os.makedirs(d, exist_ok=True)
|
||||
with open(os.path.join(d, tag), "w") as f:
|
||||
json.dump({"layers": layers}, f)
|
||||
|
||||
|
||||
def _write_blob(root, digest, size):
|
||||
blobs = os.path.join(root, "blobs")
|
||||
os.makedirs(blobs, exist_ok=True)
|
||||
path = os.path.join(blobs, digest.replace(":", "-"))
|
||||
with open(path, "wb") as f:
|
||||
f.write(b"\0" * size)
|
||||
return path
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ollama_tree(tmp_path, monkeypatch):
|
||||
root = tmp_path / "ollama-models"
|
||||
root.mkdir()
|
||||
monkeypatch.setattr(ro, "OLLAMA_MODEL_DIRS", [str(root)])
|
||||
return str(root)
|
||||
|
||||
|
||||
def test_find_ollama_model_files_maps_library_model_to_its_blob(ollama_tree):
|
||||
"""registry/library/<name>/<tag> is the common case and must render as 'name:tag'."""
|
||||
_write_blob(ollama_tree, "sha256:aaa111", 4096)
|
||||
_write_manifest(ollama_tree, "registry.ollama.ai/library/llama3", "8b", [
|
||||
{"mediaType": MODEL_MEDIA_TYPE, "digest": "sha256:aaa111", "size": 987654321},
|
||||
])
|
||||
files = ro.find_ollama_model_files()
|
||||
assert len(files) == 1
|
||||
entry = files[0]
|
||||
assert entry["model"] == "llama3:8b"
|
||||
assert entry["filename"] == "sha256-aaa111"
|
||||
assert entry["size_bytes"] == 987654321 # taken from the manifest, not the stub blob
|
||||
assert entry["kind"] == "ollama"
|
||||
assert os.path.exists(entry["full_path"])
|
||||
|
||||
|
||||
def test_find_ollama_model_files_keeps_non_library_namespace(ollama_tree):
|
||||
"""A model pulled from a user namespace must keep it, or two different models with
|
||||
the same short name collide in the warm plan."""
|
||||
_write_blob(ollama_tree, "sha256:bbb222", 4096)
|
||||
_write_manifest(ollama_tree, "hf.co/someuser/qwen-coder", "q4", [
|
||||
{"mediaType": MODEL_MEDIA_TYPE, "digest": "sha256:bbb222", "size": 100},
|
||||
])
|
||||
files = ro.find_ollama_model_files()
|
||||
assert [f["model"] for f in files] == ["someuser/qwen-coder:q4"]
|
||||
|
||||
|
||||
def test_find_ollama_model_files_ignores_non_model_layers(ollama_tree):
|
||||
"""Manifests also list templates, params and licence layers; warming those is wasted
|
||||
IO and they are not the weights."""
|
||||
_write_blob(ollama_tree, "sha256:ccc333", 4096)
|
||||
_write_blob(ollama_tree, "sha256:ddd444", 32)
|
||||
_write_manifest(ollama_tree, "registry.ollama.ai/library/mistral", "7b", [
|
||||
{"mediaType": "application/vnd.ollama.image.template", "digest": "sha256:ddd444", "size": 32},
|
||||
{"mediaType": MODEL_MEDIA_TYPE, "digest": "sha256:ccc333", "size": 500},
|
||||
])
|
||||
files = ro.find_ollama_model_files()
|
||||
assert len(files) == 1
|
||||
assert files[0]["filename"] == "sha256-ccc333"
|
||||
|
||||
|
||||
def test_find_ollama_model_files_skips_layers_whose_blob_is_missing(ollama_tree):
|
||||
"""A partially pulled model leaves a manifest with no blob; returning that path would
|
||||
make every downstream residency check fail."""
|
||||
_write_manifest(ollama_tree, "registry.ollama.ai/library/ghost", "latest", [
|
||||
{"mediaType": MODEL_MEDIA_TYPE, "digest": "sha256:missing", "size": 10},
|
||||
])
|
||||
assert ro.find_ollama_model_files() == []
|
||||
|
||||
|
||||
def test_find_ollama_model_files_tolerates_corrupt_manifest_json(ollama_tree):
|
||||
"""One unreadable manifest must not take out discovery of every other model."""
|
||||
bad_dir = os.path.join(ollama_tree, "manifests", "registry.ollama.ai/library/broken")
|
||||
os.makedirs(bad_dir)
|
||||
with open(os.path.join(bad_dir, "latest"), "w") as f:
|
||||
f.write("{not json")
|
||||
_write_blob(ollama_tree, "sha256:eee555", 4096)
|
||||
_write_manifest(ollama_tree, "registry.ollama.ai/library/good", "latest", [
|
||||
{"mediaType": MODEL_MEDIA_TYPE, "digest": "sha256:eee555", "size": 10},
|
||||
])
|
||||
assert [f["model"] for f in ro.find_ollama_model_files()] == ["good:latest"]
|
||||
|
||||
|
||||
def test_find_ollama_model_files_deduplicates_same_model_and_blob(ollama_tree, monkeypatch):
|
||||
"""The same root listed twice (or a duplicated layer) must not double-count bytes
|
||||
against the warm budget."""
|
||||
_write_blob(ollama_tree, "sha256:fff666", 4096)
|
||||
_write_manifest(ollama_tree, "registry.ollama.ai/library/dup", "latest", [
|
||||
{"mediaType": MODEL_MEDIA_TYPE, "digest": "sha256:fff666", "size": 10},
|
||||
{"mediaType": MODEL_MEDIA_TYPE, "digest": "sha256:fff666", "size": 10},
|
||||
])
|
||||
monkeypatch.setattr(ro, "OLLAMA_MODEL_DIRS", [ollama_tree, ollama_tree])
|
||||
assert len(ro.find_ollama_model_files()) == 1
|
||||
|
||||
|
||||
def test_find_ollama_model_files_returns_empty_when_no_manifest_dir(tmp_path, monkeypatch):
|
||||
monkeypatch.setattr(ro, "OLLAMA_MODEL_DIRS", [str(tmp_path / "does-not-exist")])
|
||||
assert ro.find_ollama_model_files() == []
|
||||
|
||||
|
||||
# --------------------------------------------------------------- warm planning
|
||||
|
||||
GIB = 1024 ** 3
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def planner(monkeypatch):
|
||||
"""build_warm_plan with every IO boundary stubbed: catalog, residency, meminfo, scores."""
|
||||
state = {"residency": {}, "scores": {}}
|
||||
|
||||
def _catalog(force_refresh=False):
|
||||
return state["catalog"]
|
||||
|
||||
monkeypatch.setattr(ro, "get_model_catalog", _catalog)
|
||||
monkeypatch.setattr(ro, "_warm_priority", lambda days=30.0: state["scores"])
|
||||
monkeypatch.setattr(ro, "get_detailed_meminfo",
|
||||
lambda: {"available_bytes": 40 * GIB, "available_gb": 40.0})
|
||||
monkeypatch.setattr(ro, "page_residency",
|
||||
lambda path, allow_probe=True, probe_windows=None:
|
||||
state["residency"].get(path, {"resident_pct": 0.0,
|
||||
"warm_confident": False}))
|
||||
return state
|
||||
|
||||
|
||||
def _ollama_entry(name, path, gb):
|
||||
return {"model": name, "full_path": path, "size_bytes": int(gb * GIB),
|
||||
"size_gb": gb, "kind": "ollama"}
|
||||
|
||||
|
||||
def _comfy_entry(rel, path, gb, mtime):
|
||||
return {"rel_path": rel, "full_path": path, "size_bytes": int(gb * GIB),
|
||||
"size_gb": gb, "kind": "comfy", "mtime": mtime}
|
||||
|
||||
|
||||
def test_warm_plan_orders_by_usage_score(planner):
|
||||
"""The whole point of the plan is that the most-used model is warmed first, because
|
||||
the budget usually cannot cover everything."""
|
||||
planner["catalog"] = {
|
||||
"ollama": [_ollama_entry("rare:latest", "/m/rare", 1.0),
|
||||
_ollama_entry("hot:latest", "/m/hot", 1.0)],
|
||||
"comfy": [],
|
||||
}
|
||||
planner["scores"] = {"hot:latest": 50.0, "rare:latest": 0.1}
|
||||
plan = ro.build_warm_plan(budget_gb=10.0)
|
||||
assert [e["name"] for e in plan["plan"]] == ["hot:latest", "rare:latest"]
|
||||
|
||||
|
||||
def test_warm_plan_skips_confidently_warm_files_without_spending_budget(planner):
|
||||
"""Already-resident files cost nothing and must not consume budget another file needs."""
|
||||
planner["catalog"] = {
|
||||
"ollama": [_ollama_entry("warm:latest", "/m/warm", 8.0),
|
||||
_ollama_entry("cold:latest", "/m/cold", 8.0)],
|
||||
"comfy": [],
|
||||
}
|
||||
planner["scores"] = {"warm:latest": 10.0, "cold:latest": 5.0}
|
||||
planner["residency"]["/m/warm"] = {"resident_pct": 100.0, "warm_confident": True}
|
||||
plan = ro.build_warm_plan(budget_gb=10.0)
|
||||
assert [e["name"] for e in plan["plan"]] == ["cold:latest"]
|
||||
assert [e["action"] for e in plan["skipped"]] == ["already-warm"]
|
||||
assert plan["planned_gb"] == pytest.approx(8.0, abs=0.01)
|
||||
|
||||
|
||||
def test_warm_plan_does_not_skip_a_high_but_unconfident_residency(planner):
|
||||
"""95% from a probe is not permission to skip — this is the bug the warm_confident
|
||||
flag exists to prevent."""
|
||||
planner["catalog"] = {"ollama": [_ollama_entry("m:latest", "/m/x", 4.0)], "comfy": []}
|
||||
planner["residency"]["/m/x"] = {"resident_pct": 95.0, "warm_confident": False}
|
||||
plan = ro.build_warm_plan(budget_gb=10.0)
|
||||
assert [e["name"] for e in plan["plan"]] == ["m:latest"]
|
||||
# Only the missing 5% has to be read.
|
||||
assert plan["plan"][0]["bytes_to_read"] == pytest.approx(int(4.0 * GIB) * 0.05, rel=0.01)
|
||||
|
||||
|
||||
def test_warm_plan_charges_only_the_non_resident_fraction(planner):
|
||||
planner["catalog"] = {"ollama": [_ollama_entry("m:latest", "/m/x", 10.0)], "comfy": []}
|
||||
planner["residency"]["/m/x"] = {"resident_pct": 50.0, "warm_confident": False}
|
||||
plan = ro.build_warm_plan(budget_gb=6.0)
|
||||
assert plan["planned_gb"] == pytest.approx(5.0, abs=0.01)
|
||||
|
||||
|
||||
def test_warm_plan_stops_at_the_budget(planner):
|
||||
"""Warming past the budget just evicts what was warmed first, so over-budget entries
|
||||
are reported as skipped rather than planned."""
|
||||
planner["catalog"] = {
|
||||
"ollama": [_ollama_entry("a", "/m/a", 6.0), _ollama_entry("b", "/m/b", 6.0)],
|
||||
"comfy": [],
|
||||
}
|
||||
planner["scores"] = {"a": 9.0, "b": 1.0}
|
||||
plan = ro.build_warm_plan(budget_gb=8.0)
|
||||
assert [e["name"] for e in plan["plan"]] == ["a"]
|
||||
assert plan["skipped"][0]["name"] == "b"
|
||||
assert plan["skipped"][0]["action"] == "over-budget"
|
||||
assert plan["planned_gb"] <= plan["budget_gb"]
|
||||
|
||||
|
||||
def test_warm_plan_deduplicates_by_path(planner):
|
||||
"""The same file reachable from both catalogs must be planned once, or the budget is
|
||||
charged twice for one read."""
|
||||
planner["catalog"] = {
|
||||
"ollama": [_ollama_entry("shared", "/m/shared", 2.0)],
|
||||
"comfy": [_comfy_entry("shared.safetensors", "/m/shared", 2.0, 0)],
|
||||
}
|
||||
plan = ro.build_warm_plan(budget_gb=100.0)
|
||||
assert len(plan["plan"]) == 1
|
||||
assert plan["planned_gb"] == pytest.approx(2.0, abs=0.01)
|
||||
|
||||
|
||||
def test_warm_plan_default_budget_leaves_headroom(planner):
|
||||
"""An unbounded budget would push the box into reclaim; the default is 70% of
|
||||
MemAvailable."""
|
||||
planner["catalog"] = {"ollama": [], "comfy": []}
|
||||
plan = ro.build_warm_plan()
|
||||
assert plan["budget_gb"] == pytest.approx(40.0 * 0.7, abs=0.01)
|
||||
|
||||
|
||||
def test_warm_plan_ranks_recent_comfy_checkpoints_above_stale_ones(planner):
|
||||
"""ComfyUI files have no usage history, so recency by mtime is the ranking signal."""
|
||||
import time
|
||||
now = time.time()
|
||||
planner["catalog"] = {
|
||||
"ollama": [],
|
||||
"comfy": [_comfy_entry("old.safetensors", "/c/old", 1.0, now - 90 * 86400),
|
||||
_comfy_entry("new.safetensors", "/c/new", 1.0, now - 60)],
|
||||
}
|
||||
plan = ro.build_warm_plan(budget_gb=10.0)
|
||||
assert [e["name"] for e in plan["plan"]] == ["new.safetensors", "old.safetensors"]
|
||||
|
||||
|
||||
def test_warm_plan_on_empty_catalog_is_a_valid_empty_plan(planner):
|
||||
planner["catalog"] = {"ollama": [], "comfy": []}
|
||||
plan = ro.build_warm_plan(budget_gb=1.0)
|
||||
assert plan["warm_count"] == 0 and plan["plan"] == [] and plan["skipped"] == []
|
||||
342
tests/test_telemetry_store.py
Normal file
342
tests/test_telemetry_store.py
Normal file
@@ -0,0 +1,342 @@
|
||||
"""telemetry_store: schema, the background writer round-trip, and the analytics queries.
|
||||
|
||||
Every test in this module runs against a throwaway SQLite file via the `temp_db` fixture
|
||||
(which monkeypatches telemetry_store.DB_PATH). The production hyperswap.db is never
|
||||
opened: the writer thread, _read_conn() and _rows() all resolve DB_PATH at call time.
|
||||
"""
|
||||
import os
|
||||
import sqlite3
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
import telemetry_store as ts
|
||||
|
||||
|
||||
def _drain(timeout=3.0):
|
||||
"""Wait until the writer thread has dequeued everything submitted so far.
|
||||
|
||||
NOTE: stop() flushes the writer's *pending* batch but does not drain the submission
|
||||
queue, so a stop() racing a just-submitted row can drop it. Tests wait explicitly
|
||||
rather than depending on that race. (See tests/README.md.)
|
||||
"""
|
||||
w = ts._writer
|
||||
assert w is not None, "writer not started"
|
||||
deadline = time.time() + timeout
|
||||
while time.time() < deadline:
|
||||
if w.q.empty():
|
||||
time.sleep(0.05) # let the in-flight item finish its execute/commit
|
||||
return
|
||||
time.sleep(0.01)
|
||||
raise AssertionError("telemetry writer never drained its queue")
|
||||
|
||||
|
||||
def _seed(db_path, events=(), telemetry=()):
|
||||
"""Insert rows directly, bypassing the writer thread, for query tests."""
|
||||
conn = sqlite3.connect(db_path)
|
||||
conn.executescript(ts.SCHEMA)
|
||||
for e in events:
|
||||
conn.execute(ts._EVENT_SQL, (
|
||||
e.get("ts", time.time()), e.get("event_type"), e.get("source"), e.get("target"),
|
||||
e.get("profile"), e.get("duration_ms"), e.get("load_duration_ms"),
|
||||
e.get("yield_confirm_ms"), e.get("tokens_per_sec"), e.get("bytes_loaded"),
|
||||
e.get("load_gbps"), e.get("cache_status"), e.get("detail"),
|
||||
))
|
||||
for t in telemetry:
|
||||
conn.execute(ts._TELEMETRY_SQL, (
|
||||
t.get("ts", time.time()), t.get("profile"), t.get("gpu_util_pct"),
|
||||
t.get("mem_util_pct"), t.get("temp_c"), t.get("power_w"), t.get("power_limit_w"),
|
||||
t.get("fan_pct"), t.get("clock_sm_mhz"), t.get("clock_mem_mhz"),
|
||||
t.get("vram_used_bytes"), t.get("ollama_bytes"), t.get("comfy_bytes"),
|
||||
t.get("system_bytes"), t.get("ram_used_bytes"), t.get("ram_cached_bytes"),
|
||||
t.get("pcie_tx_kbps"), t.get("pcie_rx_kbps"), t.get("throttle_reasons"),
|
||||
))
|
||||
conn.commit()
|
||||
conn.close()
|
||||
|
||||
|
||||
# --------------------------------------------------------------- writer round-trip
|
||||
|
||||
def test_start_creates_the_schema(temp_db):
|
||||
"""The service starts against a database that may not exist yet."""
|
||||
ts.start()
|
||||
ts.stop()
|
||||
assert os.path.exists(temp_db)
|
||||
conn = sqlite3.connect(temp_db)
|
||||
tables = {r[0] for r in conn.execute("SELECT name FROM sqlite_master WHERE type='table'")}
|
||||
conn.close()
|
||||
assert {"telemetry", "events", "autotune_runs"} <= tables
|
||||
|
||||
|
||||
def test_record_event_round_trips_through_the_writer(temp_db):
|
||||
"""Events are committed immediately (not batched) because they are what the analytics
|
||||
are built on."""
|
||||
ts.start()
|
||||
ts.record_event({
|
||||
"event_type": "LLM Model Switch", "source": "None", "target": "llama3:8b",
|
||||
"duration_ms": 5120.0, "load_duration_ms": 4901.0, "tokens_per_sec": 61.2,
|
||||
"bytes_loaded": 13819000000, "load_gbps": 2.63, "cache_status": "RAM Cache Hit ⚡",
|
||||
}, profile="ollama")
|
||||
_drain()
|
||||
ts.stop()
|
||||
|
||||
rows = ts.recent_events()
|
||||
assert len(rows) == 1
|
||||
assert rows[0]["target"] == "llama3:8b"
|
||||
assert rows[0]["profile"] == "ollama"
|
||||
assert rows[0]["cache_status"] == "RAM Cache Hit ⚡"
|
||||
assert rows[0]["load_gbps"] == pytest.approx(2.63)
|
||||
|
||||
|
||||
def test_record_telemetry_flushes_on_stop(temp_db):
|
||||
"""Telemetry is batched on a FLUSH_INTERVAL_S timer; stop() must not drop what is
|
||||
still pending, or the last seconds before a restart are lost."""
|
||||
ts.start()
|
||||
ts.record_telemetry(
|
||||
gpu={"available": True, "gpu_util_pct": 88.0, "temperature_c": 71.0,
|
||||
"power_w": 355.0, "clock_sm_mhz": 2790.0, "clock_mem_mhz": 10501.0,
|
||||
"vram_used_bytes": 13 * 1024 ** 3,
|
||||
"breakdown": {"ollama_gb": 12.0, "comfyui_gb": 0.0, "system_gb": 0.5}},
|
||||
ram={"used_bytes": 20 * 1024 ** 3, "cached_bytes": 30 * 1024 ** 3},
|
||||
profile="ollama", throttle_reasons="sw_power_cap")
|
||||
_drain()
|
||||
ts.stop()
|
||||
|
||||
conn = sqlite3.connect(temp_db)
|
||||
row = conn.execute("SELECT profile, temp_c, ollama_bytes, throttle_reasons "
|
||||
"FROM telemetry").fetchone()
|
||||
conn.close()
|
||||
assert row[0] == "ollama"
|
||||
assert row[1] == pytest.approx(71.0)
|
||||
assert row[2] == 12 * 1024 ** 3
|
||||
assert row[3] == "sw_power_cap"
|
||||
|
||||
|
||||
def test_record_telemetry_drops_unavailable_gpu_samples(temp_db):
|
||||
"""A failed NVML read must not write a row of nulls that then skews every average."""
|
||||
ts.start()
|
||||
ts.record_telemetry(gpu={"available": False}, ram={}, profile="ollama")
|
||||
_drain()
|
||||
ts.stop()
|
||||
conn = sqlite3.connect(temp_db)
|
||||
n = conn.execute("SELECT COUNT(*) FROM telemetry").fetchone()[0]
|
||||
conn.close()
|
||||
assert n == 0
|
||||
|
||||
|
||||
def test_recording_without_a_running_writer_is_a_no_op(temp_db):
|
||||
"""Callers on the hot path must never need to know whether the store was started."""
|
||||
ts.record_event({"event_type": "LLM Model Switch", "target": "x"})
|
||||
ts.record_telemetry(gpu={"available": True, "breakdown": {}}, ram={})
|
||||
assert not os.path.exists(temp_db)
|
||||
|
||||
|
||||
def test_record_event_never_raises_on_a_malformed_event(temp_db):
|
||||
"""This is called from the swap path; it is documented as never raising."""
|
||||
ts.start()
|
||||
ts.record_event({})
|
||||
ts.record_event({"event_type": "Weird", "duration_ms": "not-a-number"})
|
||||
_drain()
|
||||
ts.stop()
|
||||
|
||||
|
||||
def test_autotune_row_round_trips(temp_db):
|
||||
ts.start()
|
||||
ts.record_autotune({"profile": "ollama", "knob": "lock_mem_mhz", "mem_offset_mhz": 0,
|
||||
"tokens_per_sec": 62.5, "temp_c": 74.0, "stable": True})
|
||||
_drain()
|
||||
ts.stop()
|
||||
rows = ts.autotune_history()
|
||||
assert len(rows) == 1
|
||||
assert rows[0]["knob"] == "lock_mem_mhz"
|
||||
assert rows[0]["stable"] == 1 # stored as an INTEGER, not a bool
|
||||
|
||||
|
||||
# --------------------------------------------------------------- absent database
|
||||
|
||||
def test_queries_return_empty_when_the_database_does_not_exist(tmp_path, monkeypatch):
|
||||
"""A dashboard opened before the first sample must render empty, not 500."""
|
||||
monkeypatch.setattr(ts, "DB_PATH", str(tmp_path / "absent.db"))
|
||||
assert ts.profile_comparison() == []
|
||||
assert ts.model_usage_ranking() == []
|
||||
assert ts.timeseries() == []
|
||||
assert ts.recent_events() == []
|
||||
assert ts.autotune_history() == []
|
||||
stats = ts.swap_stats()
|
||||
assert stats["by_type"] == [] and stats["by_model"] == [] and stats["by_cache_status"] == []
|
||||
assert ts.db_info()["exists"] is False
|
||||
|
||||
|
||||
def test_queries_return_empty_on_a_corrupt_database(tmp_path, monkeypatch):
|
||||
db = tmp_path / "corrupt.db"
|
||||
db.write_bytes(b"this is not a sqlite file")
|
||||
monkeypatch.setattr(ts, "DB_PATH", str(db))
|
||||
assert ts.recent_events() == []
|
||||
assert ts.profile_comparison() == []
|
||||
|
||||
|
||||
# --------------------------------------------------------------- analytics
|
||||
|
||||
def test_profile_comparison_ranks_profiles_by_throughput(temp_db):
|
||||
"""The headline question this store exists to answer: which profile is actually
|
||||
faster? Ordering is by average tok/s, descending."""
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 60, "event_type": "LLM Model Switch", "profile": "ollama",
|
||||
"tokens_per_sec": 62.0, "load_gbps": 2.6, "load_duration_ms": 4900},
|
||||
{"ts": now - 50, "event_type": "LLM Model Switch", "profile": "ollama",
|
||||
"tokens_per_sec": 64.0, "load_gbps": 2.6, "load_duration_ms": 4900},
|
||||
{"ts": now - 40, "event_type": "LLM Model Switch", "profile": "balanced",
|
||||
"tokens_per_sec": 51.0, "load_gbps": 2.5, "load_duration_ms": 5100},
|
||||
], telemetry=[
|
||||
{"ts": now - 55, "profile": "ollama", "gpu_util_pct": 90, "temp_c": 74.0,
|
||||
"power_w": 360.0, "clock_sm_mhz": 2790, "clock_mem_mhz": 10501},
|
||||
{"ts": now - 45, "profile": "balanced", "gpu_util_pct": 90, "temp_c": 66.0,
|
||||
"power_w": 300.0, "clock_sm_mhz": 2600, "clock_mem_mhz": 9501},
|
||||
])
|
||||
out = ts.profile_comparison(days=1)
|
||||
assert [r["profile"] for r in out] == ["ollama", "balanced"]
|
||||
assert out[0]["swaps"] == 2
|
||||
assert out[0]["avg_tok_s"] == pytest.approx(63.0)
|
||||
# Thermals are joined in from the telemetry table for the same profile.
|
||||
assert out[0]["avg_temp_c"] == pytest.approx(74.0)
|
||||
|
||||
|
||||
def test_profile_comparison_ignores_events_without_measured_throughput(temp_db):
|
||||
"""A swap that produced no tokens tells us nothing about the profile's speed and
|
||||
would drag the average toward zero."""
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 10, "event_type": "LLM Model Switch", "profile": "ollama",
|
||||
"tokens_per_sec": 60.0},
|
||||
{"ts": now - 5, "event_type": "LLM Model Switch", "profile": "ollama",
|
||||
"tokens_per_sec": 0.0},
|
||||
])
|
||||
out = ts.profile_comparison(days=1)
|
||||
assert out[0]["swaps"] == 1 and out[0]["avg_tok_s"] == pytest.approx(60.0)
|
||||
|
||||
|
||||
def test_profile_comparison_excludes_samples_outside_the_window(temp_db):
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 40 * 86400, "event_type": "LLM Model Switch", "profile": "old",
|
||||
"tokens_per_sec": 99.0},
|
||||
{"ts": now - 60, "event_type": "LLM Model Switch", "profile": "ollama",
|
||||
"tokens_per_sec": 60.0},
|
||||
])
|
||||
assert [r["profile"] for r in ts.profile_comparison(days=7)] == ["ollama"]
|
||||
|
||||
|
||||
def test_profile_comparison_ignores_idle_telemetry(temp_db):
|
||||
"""Thermals are only meaningful while the GPU is doing work; idle samples (util <= 5%)
|
||||
would make every profile look cool."""
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 10, "event_type": "LLM Model Switch", "profile": "ollama",
|
||||
"tokens_per_sec": 60.0},
|
||||
], telemetry=[
|
||||
{"ts": now - 12, "profile": "ollama", "gpu_util_pct": 0, "temp_c": 35.0},
|
||||
{"ts": now - 11, "profile": "ollama", "gpu_util_pct": 95, "temp_c": 75.0},
|
||||
])
|
||||
assert ts.profile_comparison(days=1)[0]["avg_temp_c"] == pytest.approx(75.0)
|
||||
|
||||
|
||||
def test_swap_stats_groups_by_event_type_and_cache_status(temp_db):
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 30, "event_type": "LLM Model Switch", "target": "llama3:8b",
|
||||
"duration_ms": 5000, "cache_status": "RAM Cache Hit ⚡", "load_gbps": 2.6,
|
||||
"tokens_per_sec": 60.0},
|
||||
{"ts": now - 20, "event_type": "LLM Model Switch", "target": "llama3:8b",
|
||||
"duration_ms": 7000, "cache_status": "Cold Disk Load 💾", "load_gbps": 0.38,
|
||||
"tokens_per_sec": 58.0},
|
||||
{"ts": now - 10, "event_type": "VRAM Yield", "duration_ms": 900,
|
||||
"yield_confirm_ms": 850},
|
||||
])
|
||||
stats = ts.swap_stats(days=1)
|
||||
by_type = {r["event_type"]: r for r in stats["by_type"]}
|
||||
assert by_type["LLM Model Switch"]["n"] == 2
|
||||
assert by_type["LLM Model Switch"]["avg_ms"] == pytest.approx(6000.0)
|
||||
assert by_type["LLM Model Switch"]["min_ms"] == pytest.approx(5000.0)
|
||||
assert by_type["VRAM Yield"]["avg_confirm_ms"] == pytest.approx(850.0)
|
||||
assert {r["cache_status"] for r in stats["by_cache_status"]} == {
|
||||
"RAM Cache Hit ⚡", "Cold Disk Load 💾"}
|
||||
assert stats["by_model"][0]["model"] == "llama3:8b"
|
||||
assert stats["by_model"][0]["loads"] == 2
|
||||
|
||||
|
||||
def test_model_usage_ranking_scores_recent_use_higher(temp_db):
|
||||
"""Recency/frequency drives the RAM warm budget: given equal load counts, the model
|
||||
used more recently must rank first (half-life ~24 h)."""
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 3600, "event_type": "LLM Model Switch", "target": "recent:latest"},
|
||||
{"ts": now - 3600, "event_type": "LLM Model Switch", "target": "recent:latest"},
|
||||
{"ts": now - 20 * 86400, "event_type": "LLM Model Switch", "target": "stale:latest"},
|
||||
{"ts": now - 20 * 86400, "event_type": "LLM Model Switch", "target": "stale:latest"},
|
||||
])
|
||||
ranking = ts.model_usage_ranking(days=30)
|
||||
assert [r["model"] for r in ranking] == ["recent:latest", "stale:latest"]
|
||||
assert ranking[0]["score"] > ranking[1]["score"]
|
||||
assert ranking[0]["loads"] == 2
|
||||
assert ranking[0]["age_hours"] == pytest.approx(1.0, abs=0.05)
|
||||
|
||||
|
||||
def test_model_usage_ranking_scores_frequent_use_higher_at_equal_recency(temp_db):
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 3600, "event_type": "LLM Model Switch", "target": "often:latest"},
|
||||
{"ts": now - 3601, "event_type": "LLM Model Switch", "target": "often:latest"},
|
||||
{"ts": now - 3602, "event_type": "LLM Model Switch", "target": "often:latest"},
|
||||
{"ts": now - 3600, "event_type": "LLM Model Switch", "target": "once:latest"},
|
||||
])
|
||||
ranking = ts.model_usage_ranking(days=30)
|
||||
assert ranking[0]["model"] == "often:latest"
|
||||
|
||||
|
||||
def test_model_usage_ranking_counts_warms_as_well_as_switches(temp_db):
|
||||
now = time.time()
|
||||
_seed(temp_db, events=[
|
||||
{"ts": now - 60, "event_type": "Model Warm", "target": "warmed:latest"},
|
||||
{"ts": now - 60, "event_type": "Something Else", "target": "ignored:latest"},
|
||||
])
|
||||
assert [r["model"] for r in ts.model_usage_ranking(days=1)] == ["warmed:latest"]
|
||||
|
||||
|
||||
def test_timeseries_buckets_samples_by_width(temp_db):
|
||||
"""Six hours into six buckets is one bucket per hour; samples inside an hour collapse
|
||||
into a single averaged point."""
|
||||
now = time.time()
|
||||
base = now - 5.5 * 3600
|
||||
_seed(temp_db, telemetry=[
|
||||
{"ts": base + 60, "temp_c": 60.0, "gpu_util_pct": 50},
|
||||
{"ts": base + 120, "temp_c": 70.0, "gpu_util_pct": 70},
|
||||
{"ts": base + 3700, "temp_c": 80.0, "gpu_util_pct": 90},
|
||||
])
|
||||
out = ts.timeseries(hours=6, buckets=6)
|
||||
assert len(out) == 2
|
||||
assert out[0]["temp_c"] == pytest.approx(65.0) # average of 60 and 70
|
||||
assert out[1]["temp_c"] == pytest.approx(80.0)
|
||||
assert out[0]["bucket_ts"] < out[1]["bucket_ts"]
|
||||
|
||||
|
||||
def test_timeseries_excludes_samples_older_than_the_window(temp_db):
|
||||
now = time.time()
|
||||
_seed(temp_db, telemetry=[
|
||||
{"ts": now - 48 * 3600, "temp_c": 99.0},
|
||||
{"ts": now - 60, "temp_c": 60.0},
|
||||
])
|
||||
out = ts.timeseries(hours=6, buckets=240)
|
||||
assert len(out) == 1 and out[0]["temp_c"] == pytest.approx(60.0)
|
||||
|
||||
|
||||
def test_db_info_reports_row_counts_and_coverage(temp_db):
|
||||
now = time.time()
|
||||
_seed(temp_db,
|
||||
events=[{"ts": now - 10, "event_type": "LLM Model Switch", "target": "m"}],
|
||||
telemetry=[{"ts": now - 7200, "temp_c": 60.0}, {"ts": now, "temp_c": 61.0}])
|
||||
info = ts.db_info()
|
||||
assert info["exists"] is True
|
||||
assert info["events_rows"] == 1
|
||||
assert info["telemetry_rows"] == 2
|
||||
assert info["coverage_hours"] == pytest.approx(2.0, abs=0.01)
|
||||
264
tests/test_thermal_governor.py
Normal file
264
tests/test_thermal_governor.py
Normal file
@@ -0,0 +1,264 @@
|
||||
"""thermal_governor: the derate state machine.
|
||||
|
||||
SAFETY: _step() spawns a thread that calls overclock_manager.apply_profile. The autouse
|
||||
`no_gpu_mutation` fixture in conftest.py replaces that (and load_profiles is stubbed per
|
||||
test), so escalation here can never reach the card.
|
||||
"""
|
||||
import time
|
||||
|
||||
import pytest
|
||||
|
||||
import overclock_manager
|
||||
import thermal_governor as tg
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def gov(monkeypatch):
|
||||
"""A fresh governor with a stubbed profile store — never the module singleton."""
|
||||
monkeypatch.setattr(overclock_manager, "load_profiles", lambda: {
|
||||
"ollama": {"core_offset_mhz": 100, "mem_offset_mhz": 500, "fan_speed_pct": 70,
|
||||
"lock_core_min": 2500, "lock_core_max": 2800},
|
||||
"stock": {"core_offset_mhz": 0, "mem_offset_mhz": 0, "fan_speed_pct": 0},
|
||||
})
|
||||
monkeypatch.setattr(overclock_manager, "ACTIVE_PROFILE", "ollama")
|
||||
return tg.ThermalGovernor()
|
||||
|
||||
|
||||
def _sample(temp=60.0, reasons=None, available=True):
|
||||
return {"available": available, "temperature_c": temp, "throttle_reasons": reasons or []}
|
||||
|
||||
|
||||
def _feed(gov, n, **kw):
|
||||
for _ in range(n):
|
||||
gov.observe(_sample(**kw), active_profile="ollama")
|
||||
|
||||
|
||||
def _clear_cooldown(gov):
|
||||
"""The REAPPLY_COOLDOWN_S gate is time-based; wind the clock back rather than sleep."""
|
||||
gov.last_change = time.time() - tg.REAPPLY_COOLDOWN_S - 1.0
|
||||
|
||||
|
||||
# --------------------------------------------------------------- escalation hysteresis
|
||||
|
||||
def test_does_not_escalate_before_hot_samples_consecutive_bad_readings(gov):
|
||||
"""One hot spike during a diffusion step must not derate the card."""
|
||||
_feed(gov, tg.HOT_SAMPLES - 1, temp=90.0)
|
||||
assert gov.level == 0
|
||||
assert gov.hot_streak == tg.HOT_SAMPLES - 1
|
||||
|
||||
|
||||
def test_escalates_on_exactly_hot_samples_consecutive_bad_readings(gov, no_gpu_mutation):
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=90.0)
|
||||
assert gov.level == 1
|
||||
assert gov.hot_streak == 0 # streaks reset after a step
|
||||
# Actuation runs in a daemon thread; it must reach apply_profile (the stub) with the
|
||||
# derate overrides rather than blocking the 1 Hz sampler.
|
||||
deadline = time.time() + 2.0
|
||||
while not no_gpu_mutation["apply_profile"] and time.time() < deadline:
|
||||
time.sleep(0.01)
|
||||
assert no_gpu_mutation["apply_profile"], "escalation never actuated"
|
||||
name, overrides = no_gpu_mutation["apply_profile"][0]
|
||||
assert name == "ollama"
|
||||
assert overrides["core_offset_mhz"] == int(100 * tg.DERATE_LADDER[1]["offset_scale"])
|
||||
|
||||
|
||||
def test_alternating_hot_and_cool_samples_never_escalate(gov):
|
||||
"""The core hysteresis property: a card oscillating around the threshold must not
|
||||
thrash the profile. A single good sample resets the hot streak."""
|
||||
for _ in range(50):
|
||||
gov.observe(_sample(temp=90.0), active_profile="ollama")
|
||||
gov.observe(_sample(temp=60.0), active_profile="ollama")
|
||||
assert gov.level == 0
|
||||
|
||||
|
||||
def test_temperature_between_recover_and_escalate_resets_both_streaks(gov):
|
||||
"""The band between TEMP_RECOVER_C and TEMP_ESCALATE_C is neither hot nor cool; it
|
||||
must not accumulate credit in either direction."""
|
||||
_feed(gov, tg.HOT_SAMPLES - 1, temp=90.0)
|
||||
gov.observe(_sample(temp=78.0), active_profile="ollama")
|
||||
assert gov.hot_streak == 0 and gov.cool_streak == 0
|
||||
assert gov.level == 0
|
||||
|
||||
|
||||
def test_hard_throttle_counts_as_hot_even_when_cool(gov):
|
||||
"""A hardware slowdown means the card is protecting itself; temperature alone is not
|
||||
the whole signal."""
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=55.0, reasons=["hw_thermal_slowdown"])
|
||||
assert gov.level == 1
|
||||
|
||||
|
||||
def test_soft_throttle_reasons_do_not_escalate(gov):
|
||||
"""Hitting a power or utilisation cap is normal operation, not distress."""
|
||||
_feed(gov, tg.HOT_SAMPLES * 3, temp=55.0, reasons=["sw_power_cap", "gpu_idle"])
|
||||
assert gov.level == 0
|
||||
|
||||
|
||||
def test_escalation_stops_at_the_bottom_of_the_ladder(gov):
|
||||
"""Level must never index past DERATE_LADDER."""
|
||||
for _ in range(len(tg.DERATE_LADDER) + 3):
|
||||
_clear_cooldown(gov)
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
assert gov.level == len(tg.DERATE_LADDER) - 1
|
||||
|
||||
|
||||
def test_cooldown_gate_blocks_a_second_step_immediately_after_the_first(gov):
|
||||
"""REAPPLY_COOLDOWN_S stops the governor from walking the whole ladder in one second
|
||||
while the card is still responding to the previous change."""
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
assert gov.level == 1
|
||||
_feed(gov, tg.HOT_SAMPLES * 2, temp=95.0)
|
||||
assert gov.level == 1
|
||||
|
||||
|
||||
def test_cooldown_gate_releases_after_the_window(gov):
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
_clear_cooldown(gov)
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
assert gov.level == 2
|
||||
|
||||
|
||||
# --------------------------------------------------------------- recovery
|
||||
|
||||
def test_recovery_needs_cool_samples_consecutive_good_readings(gov):
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
assert gov.level == 1
|
||||
_clear_cooldown(gov)
|
||||
_feed(gov, tg.COOL_SAMPLES - 1, temp=60.0)
|
||||
assert gov.level == 1, "recovered too early"
|
||||
gov.observe(_sample(temp=60.0), active_profile="ollama")
|
||||
assert gov.level == 0
|
||||
|
||||
|
||||
def test_recovery_is_slower_than_escalation():
|
||||
"""Deliberate asymmetry: react fast to heat, give the clocks back slowly."""
|
||||
assert tg.COOL_SAMPLES > tg.HOT_SAMPLES
|
||||
|
||||
|
||||
def test_recovery_stops_at_level_zero(gov):
|
||||
_clear_cooldown(gov)
|
||||
_feed(gov, tg.COOL_SAMPLES * 2, temp=50.0)
|
||||
assert gov.level == 0
|
||||
|
||||
|
||||
def test_a_hard_throttle_blocks_recovery_even_at_a_cool_temperature(gov):
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
_clear_cooldown(gov)
|
||||
_feed(gov, tg.COOL_SAMPLES * 2, temp=50.0, reasons=["hw_power_brake_slowdown"])
|
||||
assert gov.level >= 1
|
||||
|
||||
|
||||
# --------------------------------------------------------------- ignored input
|
||||
|
||||
def test_disabled_governor_ignores_samples(gov):
|
||||
gov.enabled = False
|
||||
_feed(gov, tg.HOT_SAMPLES * 3, temp=99.0)
|
||||
assert gov.level == 0
|
||||
|
||||
|
||||
def test_unavailable_gpu_sample_is_ignored(gov):
|
||||
"""A failed NVML read reports available=False with no temperature; treating that as
|
||||
0 C would count as a cool sample and hand the clocks back."""
|
||||
for _ in range(tg.COOL_SAMPLES * 2):
|
||||
gov.observe({"available": False}, active_profile="ollama")
|
||||
assert gov.level == 0 and gov.cool_streak == 0
|
||||
|
||||
|
||||
# --------------------------------------------------------------- overrides
|
||||
|
||||
def test_overrides_are_empty_at_level_zero(gov):
|
||||
assert gov.overrides_for("ollama") == {}
|
||||
|
||||
|
||||
@pytest.mark.parametrize("level", [1, 2, 3])
|
||||
def test_overrides_scale_offsets_by_the_ladder(gov, level):
|
||||
gov.level = level
|
||||
scale = tg.DERATE_LADDER[level]["offset_scale"]
|
||||
ov = gov.overrides_for("ollama")
|
||||
assert ov["core_offset_mhz"] == int(100 * scale)
|
||||
assert ov["mem_offset_mhz"] == int(500 * scale)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("level", [1, 2, 3])
|
||||
def test_overrides_raise_the_fan_floor_and_never_lower_it(gov, level):
|
||||
"""The floor is a floor: a profile already running fans harder than the ladder asks
|
||||
keeps its own setting."""
|
||||
gov.level = level
|
||||
floor = tg.DERATE_LADDER[level]["fan_floor"]
|
||||
ov = gov.overrides_for("ollama")
|
||||
assert ov["fan_mode"] == "manual"
|
||||
assert ov["fan_speed_pct"] == max(70, floor) # profile fan_speed_pct is 70
|
||||
assert ov["fan_speed_pct"] >= floor
|
||||
|
||||
|
||||
def test_overrides_release_the_core_clock_lock_from_level_two(gov):
|
||||
"""Pinning the core clock high while the card is already backing off fights the
|
||||
hardware's own protection."""
|
||||
gov.level = 1
|
||||
assert "lock_core_max" not in gov.overrides_for("ollama")
|
||||
for level in (2, 3):
|
||||
gov.level = level
|
||||
ov = gov.overrides_for("ollama")
|
||||
assert ov["lock_core_min"] == 0 and ov["lock_core_max"] == 0
|
||||
|
||||
|
||||
def test_top_of_ladder_is_stock_clocks_and_maximum_fans(gov):
|
||||
gov.level = len(tg.DERATE_LADDER) - 1
|
||||
ov = gov.overrides_for("ollama")
|
||||
assert ov["core_offset_mhz"] == 0 and ov["mem_offset_mhz"] == 0
|
||||
assert ov["fan_speed_pct"] == 100
|
||||
|
||||
|
||||
def test_overrides_for_unknown_profile_do_not_raise(gov):
|
||||
"""The active profile can be one the store does not know; a missing config must
|
||||
derate to zero offsets rather than blow up in the sampler thread."""
|
||||
gov.level = 2
|
||||
ov = gov.overrides_for("does-not-exist")
|
||||
assert ov["core_offset_mhz"] == 0 and ov["mem_offset_mhz"] == 0
|
||||
|
||||
|
||||
# --------------------------------------------------------------- ladder invariants
|
||||
|
||||
def test_ladder_is_monotonically_more_conservative():
|
||||
"""Each rung must reduce clocks and raise fans; a non-monotonic ladder would make
|
||||
escalation increase heat."""
|
||||
scales = [s["offset_scale"] for s in tg.DERATE_LADDER]
|
||||
floors = [s["fan_floor"] for s in tg.DERATE_LADDER]
|
||||
assert scales == sorted(scales, reverse=True)
|
||||
assert floors == sorted(floors)
|
||||
assert scales[0] == 1.0 and scales[-1] == 0.0
|
||||
assert [s["level"] for s in tg.DERATE_LADDER] == list(range(len(tg.DERATE_LADDER)))
|
||||
|
||||
|
||||
def test_escalate_and_recover_temperatures_have_a_gap():
|
||||
"""Without a gap between the two thresholds the governor would oscillate."""
|
||||
assert tg.TEMP_RECOVER_C < tg.TEMP_ESCALATE_C
|
||||
|
||||
|
||||
# --------------------------------------------------------------- status / control
|
||||
|
||||
def test_status_reports_level_and_history(gov):
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
st = gov.get_status()
|
||||
assert st["level"] == 1
|
||||
assert st["label"] == tg.DERATE_LADDER[1]["label"]
|
||||
assert st["escalate_at_c"] == tg.TEMP_ESCALATE_C
|
||||
assert st["history"] and st["history"][0]["to_level"] == 1
|
||||
assert "95" in st["history"][0]["reason"]
|
||||
|
||||
|
||||
def test_history_is_bounded(gov):
|
||||
"""The governor is long-lived inside the service; its history must not grow forever."""
|
||||
for _ in range(80):
|
||||
_clear_cooldown(gov)
|
||||
gov.level = 0
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
assert len(gov.history) <= 50
|
||||
|
||||
|
||||
def test_disabling_a_derated_governor_resets_it(gov, no_gpu_mutation):
|
||||
"""Turning the governor off must give the clocks back, not freeze the derate in place."""
|
||||
_feed(gov, tg.HOT_SAMPLES, temp=95.0)
|
||||
assert gov.level == 1
|
||||
gov.set_enabled(False)
|
||||
assert gov.level == 0
|
||||
assert no_gpu_mutation["apply_profile"], "reset should have re-applied the base profile"
|
||||
82
tests/test_vram_helpers.py
Normal file
82
tests/test_vram_helpers.py
Normal file
@@ -0,0 +1,82 @@
|
||||
"""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_other(monkeypatch):
|
||||
_patch_proc(monkeypatch, _FakeProc("Xorg", ["/usr/lib/xorg/Xorg", ":8"]))
|
||||
assert va._classify_pid(1234) == "other"
|
||||
|
||||
|
||||
def test_classify_pid_returns_other_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) == "other"
|
||||
134
tests/test_yield_and_reclaim.py
Normal file
134
tests/test_yield_and_reclaim.py
Normal file
@@ -0,0 +1,134 @@
|
||||
"""Tests for VRAM yield classification and the reclaim-on-OOM path.
|
||||
|
||||
These cover the two failure modes that motivated the arbitration rework, both of which
|
||||
were observed on real hardware before being encoded here:
|
||||
|
||||
* A model mid-generation cannot unload. Persisted counters showed 19 "timeouts" in 20
|
||||
yields; telemetry for that window showed the GPU pinned at 96-97% with 14.92 GB held.
|
||||
That is a busy model, not a fault, and must not be retried in a tight loop.
|
||||
* A model that will not fit fails differently depending on configuration. With
|
||||
n_gpu_layers pinned to 99 (this box) Ollama returns a hard CUDA OOM rather than
|
||||
spilling layers to the CPU:
|
||||
"llama-server process has terminated: exit status 1: cudaMalloc failed:
|
||||
out of memory ... unable to allocate CUDA0 buffer"
|
||||
"""
|
||||
import asyncio
|
||||
|
||||
import pytest
|
||||
|
||||
import vram_arbitrator as v
|
||||
|
||||
|
||||
GB = 1024 ** 3
|
||||
|
||||
|
||||
def _snap(ollama_gb, util, free_gb=1.0):
|
||||
return {"ollama_bytes": int(ollama_gb * GB), "comfyui_bytes": 0, "other_bytes": 0,
|
||||
"free_bytes": int(free_gb * GB), "gpu_util_pct": util}
|
||||
|
||||
|
||||
class TestOomDetection:
|
||||
"""The retry path keys off Ollama's error text, so the matcher must be exact."""
|
||||
|
||||
def test_matches_the_real_observed_ollama_oom(self):
|
||||
real = ("llama-server process has terminated: exit status 1: cudaMalloc failed: "
|
||||
"out of memory\nalloc_tensor_range: failed to allocate CUDA0 buffer of "
|
||||
"size 13028925440\nerror loading model: unable to allocate CUDA0 buffer")
|
||||
assert v.looks_like_vram_oom(real)
|
||||
|
||||
@pytest.mark.parametrize("text", [
|
||||
"cudaMalloc failed: out of memory",
|
||||
"unable to allocate CUDA0 buffer",
|
||||
"failed to allocate buffer",
|
||||
"CUDA error: something",
|
||||
])
|
||||
def test_matches_each_signature(self, text):
|
||||
assert v.looks_like_vram_oom(text)
|
||||
|
||||
@pytest.mark.parametrize("text", [
|
||||
"model 'foo' not found", "invalid parameter", "", None,
|
||||
"context length exceeded",
|
||||
])
|
||||
def test_does_not_match_unrelated_failures(self, text):
|
||||
# A false positive here would purge ComfyUI over a typo in a model name.
|
||||
assert not v.looks_like_vram_oom(text)
|
||||
|
||||
|
||||
class TestYieldOutcomeClassification:
|
||||
"""_await_vram_release must separate released / busy / stuck."""
|
||||
|
||||
def _run(self, snaps, timeout_s=2.0, baseline_gb=14.9):
|
||||
seq = list(snaps)
|
||||
def fake():
|
||||
return seq.pop(0) if len(seq) > 1 else seq[0]
|
||||
original = v.get_process_vram_bytes
|
||||
v.get_process_vram_bytes = fake
|
||||
try:
|
||||
return asyncio.run(
|
||||
v._await_vram_release(int(baseline_gb * GB), timeout_s=timeout_s))
|
||||
finally:
|
||||
v.get_process_vram_bytes = original
|
||||
|
||||
def test_released_when_vram_drains(self):
|
||||
res = self._run([_snap(14.9, 30), _snap(0.0, 5, free_gb=15.4)])
|
||||
assert res["outcome"] == "released"
|
||||
assert res["confirmed"] is True
|
||||
|
||||
def test_busy_when_vram_held_and_gpu_pinned(self):
|
||||
# The observed pathology: 14.92 GB held at 96% utilisation.
|
||||
res = self._run([_snap(14.92, 96)])
|
||||
assert res["outcome"] == "busy"
|
||||
assert res["confirmed"] is False
|
||||
assert "mid-generation" in res["error"]
|
||||
assert res["peak_util_pct"] >= v.BUSY_UTIL_PCT
|
||||
|
||||
def test_stuck_when_vram_held_and_gpu_idle(self):
|
||||
# VRAM held with nothing running is the genuine fault case.
|
||||
res = self._run([_snap(14.9, 2)], timeout_s=0.3)
|
||||
assert res["outcome"] == "stuck"
|
||||
assert "idle" in res["error"]
|
||||
|
||||
def test_busy_is_decided_only_after_the_probe_window(self):
|
||||
# Deciding instantly would misread the normal 40-110ms release as busy.
|
||||
assert v.BUSY_PROBE_S > 0
|
||||
assert v.YIELD_CONFIRM_TIMEOUT_S > v.BUSY_PROBE_S
|
||||
|
||||
def test_default_wait_is_short(self):
|
||||
# It was 10s, which blocked the arbitrator for the length of an inference while
|
||||
# ComfyUI -- which is not gated on our return value -- waited anyway.
|
||||
assert v.YIELD_CONFIRM_TIMEOUT_S <= 3.0
|
||||
assert v.YIELD_CONFIRM_TIMEOUT_BLOCKING_S >= 10.0
|
||||
|
||||
|
||||
class TestBusyBackoff:
|
||||
"""A busy model must not be re-asked every second."""
|
||||
|
||||
def test_backoff_schedule_is_monotonic_and_bounded(self):
|
||||
sched = v.AutoArbitrator.BUSY_BACKOFF_S
|
||||
assert list(sched) == sorted(sched)
|
||||
assert sched[0] >= 1.0
|
||||
|
||||
def test_streak_walks_up_the_schedule_and_clamps(self):
|
||||
arb = v.AutoArbitrator()
|
||||
sched = arb.BUSY_BACKOFF_S
|
||||
for streak in range(len(sched) + 3):
|
||||
delay = sched[min(streak, len(sched) - 1)]
|
||||
assert delay == sched[min(streak, len(sched) - 1)]
|
||||
assert sched[min(99, len(sched) - 1)] == sched[-1]
|
||||
|
||||
def test_release_clears_backoff_state(self):
|
||||
arb = v.AutoArbitrator()
|
||||
arb._yield_backoff_until["m"] = 1e18
|
||||
arb._yield_busy_streak["m"] = 3
|
||||
arb.note_deferred_release(1234.0)
|
||||
assert arb._yield_backoff_until == {}
|
||||
assert arb._yield_busy_streak == {}
|
||||
assert arb.stats["deferred_releases"] == 1
|
||||
|
||||
def test_counters_distinguish_busy_from_stalled(self):
|
||||
# The old single yield_timeouts counter reported a healthy cron job as a 95%
|
||||
# failure rate.
|
||||
arb = v.AutoArbitrator()
|
||||
assert "yield_deferred_busy" in arb.stats
|
||||
assert "yield_stalled" in arb.stats
|
||||
assert "yield_timeouts" not in arb.stats
|
||||
@@ -61,6 +61,25 @@ YIELD_CONFIRM_TIMEOUT_BLOCKING_S = 30.0
|
||||
# A model still holding VRAM while the GPU is pinned is generating, not wedged.
|
||||
BUSY_UTIL_PCT = 50
|
||||
BUSY_PROBE_S = 0.6
|
||||
|
||||
# Fraction of a model that may sit outside VRAM before we call it starved. A little
|
||||
# slack absorbs rounding and KV-cache accounting; beyond it, layers are on the CPU.
|
||||
CPU_OFFLOAD_TOLERANCE = 0.02
|
||||
# Only intervene when ComfyUI is actually holding enough VRAM to be the cause.
|
||||
RECLAIM_MIN_COMFY_BYTES = 512 * 1024 ** 2
|
||||
|
||||
# Ollama's response when a model will not fit. Which of the two failure modes you get
|
||||
# depends on configuration: with n_gpu_layers left to Ollama it spills layers to the CPU
|
||||
# and reports size_vram < size; with n_gpu_layers pinned (99 on this box) it refuses and
|
||||
# returns a hard CUDA OOM instead. Both are handled -- the spill by
|
||||
# AutoArbitrator._check_ollama_starved, the hard failure by the retry below.
|
||||
OOM_SIGNATURES = ("out of memory", "cudamalloc", "unable to allocate",
|
||||
"failed to allocate", "cuda error")
|
||||
|
||||
|
||||
def 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"
|
||||
|
||||
@@ -360,7 +379,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)
|
||||
@@ -694,8 +719,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"
|
||||
@@ -745,8 +775,38 @@ async def switch_ollama_model(target_model: str, keep_alive: str = "30m") -> Dic
|
||||
"is_ram_hit": cls["is_ram_hit"],
|
||||
"response": data.get("response", ""),
|
||||
}
|
||||
return {"success": False, "error": f"HTTP {resp.status_code}: {resp.text}",
|
||||
"duration_ms": total_duration_ms}
|
||||
# A model that will not fit is the exact contention this service exists to
|
||||
# resolve. Rather than handing the caller a CUDA OOM, take the VRAM back from an
|
||||
# idle ComfyUI and try once more.
|
||||
body = resp.text
|
||||
if looks_like_vram_oom(body) and not _retrying:
|
||||
snap = get_process_vram_bytes()
|
||||
if snap["comfyui_bytes"] >= RECLAIM_MIN_COMFY_BYTES:
|
||||
logger.warning(
|
||||
f"Ollama could not fit '{target_model}' with ComfyUI holding "
|
||||
f"{round(snap['comfyui_bytes'] / (1024**3), 2)} GB — reclaiming and retrying")
|
||||
purge = await instant_free_comfyui_vram()
|
||||
arbitrator.stats["reclaims_for_ollama"] += 1
|
||||
arbitrator.last_action = (
|
||||
f"Reclaimed {round(snap['comfyui_bytes'] / (1024**3), 2)}GB from ComfyUI so "
|
||||
f"'{target_model}' could load")
|
||||
_record({
|
||||
"event_type": "VRAM Reclaim for Ollama",
|
||||
"source": "ComfyUI Pipeline",
|
||||
"target": target_model,
|
||||
"duration_ms": purge.get("duration_ms"),
|
||||
"cache_status": "Reclaimed",
|
||||
"detail": f"Ollama OOM: {body[:160]}",
|
||||
})
|
||||
await asyncio.sleep(0.3)
|
||||
retry = await switch_ollama_model(target_model, keep_alive, _retrying=True)
|
||||
retry["reclaimed_from_comfyui_gb"] = round(
|
||||
snap["comfyui_bytes"] / (1024**3), 2)
|
||||
retry["first_attempt_error"] = "CUDA OOM; retried after reclaiming VRAM"
|
||||
return retry
|
||||
return {"success": False, "error": f"HTTP {resp.status_code}: {body}",
|
||||
"duration_ms": total_duration_ms,
|
||||
"vram_oom": looks_like_vram_oom(body)}
|
||||
except Exception as e:
|
||||
return {"success": False, "error": str(e),
|
||||
"duration_ms": round((time.perf_counter() - t0) * 1000, 2)}
|
||||
@@ -795,6 +855,7 @@ class AutoArbitrator:
|
||||
# again every second just blocks the loop repeatedly for no benefit.
|
||||
self._yield_backoff_until: Dict[str, float] = {}
|
||||
self._yield_busy_streak: Dict[str, int] = {}
|
||||
self.last_reclaim_time = 0.0
|
||||
self.stats = {
|
||||
"yields": 0, # release confirmed
|
||||
"yield_deferred_busy": 0, # model mid-generation; unload queued behind it
|
||||
@@ -802,6 +863,7 @@ class AutoArbitrator:
|
||||
"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):
|
||||
@@ -984,6 +1046,57 @@ class AutoArbitrator:
|
||||
await asyncio.sleep(backoff)
|
||||
backoff = min(backoff * 1.5, 15.0)
|
||||
|
||||
RECLAIM_COOLDOWN_S = 30.0
|
||||
|
||||
async def _check_ollama_starved(self) -> None:
|
||||
"""The other direction: rescue an LLM that ComfyUI has squeezed onto the CPU.
|
||||
|
||||
Yielding Ollama for ComfyUI was automatic; the reverse never was, despite the
|
||||
README calling the arbitration bidirectional. When Ollama cannot fit a model it
|
||||
does not fail, it silently places layers on the CPU and runs about an order of
|
||||
magnitude slower -- so this is the failure mode a user is least likely to notice
|
||||
and most likely to feel.
|
||||
|
||||
If the LLM is spilling while ComfyUI sits idle holding VRAM, ComfyUI's cached
|
||||
checkpoints are the thing to give up.
|
||||
"""
|
||||
now = time.time()
|
||||
if self.comfy_was_active or (now - self.last_reclaim_time) < self.RECLAIM_COOLDOWN_S:
|
||||
return
|
||||
|
||||
ollama = await get_ollama_live_state()
|
||||
if not ollama.get("partially_offloaded"):
|
||||
return
|
||||
|
||||
snap = get_process_vram_bytes()
|
||||
if snap["comfyui_bytes"] < RECLAIM_MIN_COMFY_BYTES:
|
||||
return # ComfyUI is not the one holding the memory; nothing we can do here
|
||||
|
||||
self.last_reclaim_time = now
|
||||
model = ollama.get("active_model_name")
|
||||
offload = ollama.get("cpu_offload_pct")
|
||||
logger.warning(f"⚠ '{model}' is {offload}% on CPU while ComfyUI holds "
|
||||
f"{round(snap['comfyui_bytes'] / (1024**3), 2)} GB — reclaiming for the LLM")
|
||||
await self._purge_comfy_now(f"LLM spilling {offload}% to CPU")
|
||||
self.stats["reclaims_for_ollama"] += 1
|
||||
|
||||
# Freeing VRAM does not move layers back; only a reload re-places the model. Do
|
||||
# that only when the model is idle, never mid-generation.
|
||||
after = get_process_vram_bytes()
|
||||
if after.get("gpu_util_pct", 0) < BUSY_UTIL_PCT and model:
|
||||
logger.info(f"Reloading '{model}' to place it fully on the GPU...")
|
||||
await instant_free_ollama_vram(model, confirm=True)
|
||||
res = await switch_ollama_model(model, keep_alive="30m")
|
||||
recheck = await get_ollama_live_state()
|
||||
self.last_action = (
|
||||
f"Reclaimed {round(snap['comfyui_bytes'] / (1024**3), 2)}GB from ComfyUI and "
|
||||
f"reloaded '{model}' — now {round(recheck.get('gpu_fraction', 0) * 100)}% on GPU"
|
||||
if res.get("success") else
|
||||
f"Reclaimed VRAM from ComfyUI but reloading '{model}' failed: {res.get('error')}")
|
||||
else:
|
||||
self.last_action = (f"Reclaimed VRAM from ComfyUI; '{model}' is busy, so it will "
|
||||
f"stay partly on CPU until its next load")
|
||||
|
||||
async def _poll_watchdog(self):
|
||||
"""Fallback for when the WebSocket is down. One cheap /queue call, 1 Hz.
|
||||
|
||||
@@ -1002,6 +1115,8 @@ class AutoArbitrator:
|
||||
await self.trigger_comfy_priority("Watchdog saw an active queue")
|
||||
elif self.comfy_was_active:
|
||||
await self.trigger_comfy_completed()
|
||||
else:
|
||||
await self._check_ollama_starved()
|
||||
except Exception:
|
||||
pass
|
||||
await asyncio.sleep(interval)
|
||||
|
||||
Reference in New Issue
Block a user