diff --git a/README.md b/README.md index 8adc719..08c121e 100644 --- a/README.md +++ b/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. diff --git a/autotune.py b/autotune.py index 15e68fb..b8ce214 100644 --- a/autotune.py +++ b/autotune.py @@ -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}) diff --git a/pytest.ini b/pytest.ini new file mode 100644 index 0000000..922dcf7 --- /dev/null +++ b/pytest.ini @@ -0,0 +1,5 @@ +[pytest] +testpaths = tests +addopts = -ra +filterwarnings = + ignore::DeprecationWarning diff --git a/ram_optimizer.py b/ram_optimizer.py index e29b782..c1c4445 100644 --- a/ram_optimizer.py +++ b/ram_optimizer.py @@ -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) diff --git a/telemetry_store.py b/telemetry_store.py index e846250..3c807e8 100644 --- a/telemetry_store.py +++ b/telemetry_store.py @@ -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) diff --git a/tests/README.md b/tests/README.md new file mode 100644 index 0000000..d5a12f0 --- /dev/null +++ b/tests/README.md @@ -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. diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..9b6c399 --- /dev/null +++ b/tests/conftest.py @@ -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 diff --git a/tests/test_autotune_helpers.py b/tests/test_autotune_helpers.py new file mode 100644 index 0000000..92c84aa --- /dev/null +++ b/tests/test_autotune_helpers.py @@ -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"] diff --git a/tests/test_classify_load.py b/tests/test_classify_load.py new file mode 100644 index 0000000..63c5013 --- /dev/null +++ b/tests/test_classify_load.py @@ -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) diff --git a/tests/test_overclock_profiles.py b/tests/test_overclock_profiles.py new file mode 100644 index 0000000..7c73a68 --- /dev/null +++ b/tests/test_overclock_profiles.py @@ -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 diff --git a/tests/test_ram_optimizer.py b/tests/test_ram_optimizer.py new file mode 100644 index 0000000..b00ae27 --- /dev/null +++ b/tests/test_ram_optimizer.py @@ -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// 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"] == [] diff --git a/tests/test_telemetry_store.py b/tests/test_telemetry_store.py new file mode 100644 index 0000000..6f6f58c --- /dev/null +++ b/tests/test_telemetry_store.py @@ -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) diff --git a/tests/test_thermal_governor.py b/tests/test_thermal_governor.py new file mode 100644 index 0000000..8c245b9 --- /dev/null +++ b/tests/test_thermal_governor.py @@ -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" diff --git a/tests/test_vram_helpers.py b/tests/test_vram_helpers.py new file mode 100644 index 0000000..4c6416f --- /dev/null +++ b/tests/test_vram_helpers.py @@ -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" diff --git a/tests/test_yield_and_reclaim.py b/tests/test_yield_and_reclaim.py new file mode 100644 index 0000000..54f9069 --- /dev/null +++ b/tests/test_yield_and_reclaim.py @@ -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 diff --git a/vram_arbitrator.py b/vram_arbitrator.py index 968e771..19af0ba 100644 --- a/vram_arbitrator.py +++ b/vram_arbitrator.py @@ -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)