gpu-program-swapper/tests/test_yield_and_reclaim.py

"""Tests for VRAM yield classification and the reclaim-on-OOM path.

These cover the two failure modes that motivated the arbitration rework, both of which
were observed on real hardware before being encoded here:

  * A model mid-generation cannot unload. Persisted counters showed 19 "timeouts" in 20
    yields; telemetry for that window showed the GPU pinned at 96-97% with 14.92 GB held.
    That is a busy model, not a fault, and must not be retried in a tight loop.
  * A model that will not fit fails differently depending on configuration. With
    n_gpu_layers pinned to 99 (this box) Ollama returns a hard CUDA OOM rather than
    spilling layers to the CPU:
        "llama-server process has terminated: exit status 1: cudaMalloc failed:
         out of memory ... unable to allocate CUDA0 buffer"
"""
import asyncio
import time

import pytest

import vram_arbitrator as v


GB = 1024 ** 3


def _snap(ollama_gb, util, free_gb=1.0):
    return {"ollama_bytes": int(ollama_gb * GB), "comfyui_bytes": 0, "other_bytes": 0,
            "free_bytes": int(free_gb * GB), "gpu_util_pct": util}


class TestOomDetection:
    """The retry path keys off Ollama's error text, so the matcher must be exact."""

    def test_matches_the_real_observed_ollama_oom(self):
        real = ("llama-server process has terminated: exit status 1: cudaMalloc failed: "
                "out of memory\nalloc_tensor_range: failed to allocate CUDA0 buffer of "
                "size 13028925440\nerror loading model: unable to allocate CUDA0 buffer")
        assert v.looks_like_vram_oom(real)

    @pytest.mark.parametrize("text", [
        "cudaMalloc failed: out of memory",
        "unable to allocate CUDA0 buffer",
        "failed to allocate buffer",
        "CUDA error: something",
    ])
    def test_matches_each_signature(self, text):
        assert v.looks_like_vram_oom(text)

    @pytest.mark.parametrize("text", [
        "model 'foo' not found", "invalid parameter", "", None,
        "context length exceeded",
    ])
    def test_does_not_match_unrelated_failures(self, text):
        # A false positive here would purge ComfyUI over a typo in a model name.
        assert not v.looks_like_vram_oom(text)


class TestYieldOutcomeClassification:
    """_await_vram_release must separate released / busy / stuck."""

    def _run(self, snaps, timeout_s=2.0, baseline_gb=14.9):
        seq = list(snaps)
        def fake():
            return seq.pop(0) if len(seq) > 1 else seq[0]
        original = v.get_process_vram_bytes
        v.get_process_vram_bytes = fake
        try:
            return asyncio.run(
                v._await_vram_release(int(baseline_gb * GB), timeout_s=timeout_s))
        finally:
            v.get_process_vram_bytes = original

    def test_released_when_vram_drains(self):
        res = self._run([_snap(14.9, 30), _snap(0.0, 5, free_gb=15.4)])
        assert res["outcome"] == "released"
        assert res["confirmed"] is True

    def test_busy_when_vram_held_and_gpu_pinned(self):
        # The observed pathology: 14.92 GB held at 96% utilisation.
        res = self._run([_snap(14.92, 96)])
        assert res["outcome"] == "busy"
        assert res["confirmed"] is False
        assert "mid-generation" in res["error"]
        assert res["peak_util_pct"] >= v.BUSY_UTIL_PCT

    def test_stuck_when_vram_held_and_gpu_idle(self):
        # VRAM held with nothing running is the genuine fault case.
        res = self._run([_snap(14.9, 2)], timeout_s=0.3)
        assert res["outcome"] == "stuck"
        assert "idle" in res["error"]

    def test_busy_is_decided_only_after_the_probe_window(self):
        # Deciding instantly would misread the normal 40-110ms release as busy.
        assert v.BUSY_PROBE_S > 0
        assert v.YIELD_CONFIRM_TIMEOUT_S > v.BUSY_PROBE_S

    def test_default_wait_is_short(self):
        # It was 10s, which blocked the arbitrator for the length of an inference while
        # ComfyUI -- which is not gated on our return value -- waited anyway.
        assert v.YIELD_CONFIRM_TIMEOUT_S <= 3.0
        assert v.YIELD_CONFIRM_TIMEOUT_BLOCKING_S >= 10.0


class TestBusyBackoff:
    """A busy model must not be re-asked every second."""

    def test_backoff_schedule_is_monotonic_and_bounded(self):
        sched = v.AutoArbitrator.BUSY_BACKOFF_S
        assert list(sched) == sorted(sched)
        assert sched[0] >= 1.0

    def test_streak_walks_up_the_schedule_and_clamps(self):
        arb = v.AutoArbitrator()
        sched = arb.BUSY_BACKOFF_S
        for streak in range(len(sched) + 3):
            delay = sched[min(streak, len(sched) - 1)]
            assert delay == sched[min(streak, len(sched) - 1)]
        assert sched[min(99, len(sched) - 1)] == sched[-1]

    def test_release_clears_backoff_state(self):
        arb = v.AutoArbitrator()
        arb._yield_backoff_until["m"] = 1e18
        arb._yield_busy_streak["m"] = 3
        arb.note_deferred_release(1234.0)
        assert arb._yield_backoff_until == {}
        assert arb._yield_busy_streak == {}
        assert arb.stats["deferred_releases"] == 1

    def test_counters_distinguish_busy_from_stalled(self):
        # The old single yield_timeouts counter reported a healthy cron job as a 95%
        # failure rate.
        arb = v.AutoArbitrator()
        assert "yield_deferred_busy" in arb.stats
        assert "yield_stalled" in arb.stats
        assert "yield_timeouts" not in arb.stats


class TestComfyStaleQueueDetection:
    """ComfyUI can leave a dead job in queue_running forever.

    Observed on this machine: a WAN 2.1 i2v entry sat in queue_running while the GPU was
    idle and ComfyUI held 0.56 GB. Trusting that flag made the watchdog believe ComfyUI
    was permanently busy, so it evicted the LLM on every poll, never ran the idle purge,
    and never checked whether the LLM had been pushed onto the CPU. Instrumenting the
    watchdog showed busy=6, idle_check=0 -- one stale row had disabled half the logic.
    """

    def _arb(self, comfy_bytes):
        arb = v.AutoArbitrator()
        v.get_process_vram_bytes = lambda: {
            "ollama_bytes": 0, "comfyui_bytes": int(comfy_bytes), "other_bytes": 0,
            "desktop_bytes": 0, "unmanaged_bytes": 0, "free_bytes": 0, "gpu_util_pct": 0}
        return arb

    def teardown_method(self):
        import importlib
        importlib.reload(v)

    def test_empty_queue_is_not_busy(self):
        arb = self._arb(0)
        assert arb._comfy_genuinely_busy({"queue_running": [], "queue_pending": []}) is False

    def test_pending_work_is_always_busy(self):
        arb = self._arb(0)
        assert arb._comfy_genuinely_busy(
            {"queue_running": [], "queue_pending": [[1, "p"]]}) is True

    def test_a_running_job_is_believed_at_first(self):
        # It must not be called stale before it has had time to load anything.
        arb = self._arb(0.1 * GB)
        assert arb._comfy_genuinely_busy(
            {"queue_running": [[1, "abc"]], "queue_pending": []}) is True

    def test_long_running_job_holding_no_vram_is_stale(self):
        arb = self._arb(0.56 * GB)   # the observed CUDA-context floor
        q = {"queue_running": [[1, "abc"]], "queue_pending": []}
        arb._comfy_genuinely_busy(q)
        arb._running_since = time.time() - (arb.STALE_RUNNING_S + 5)
        assert arb._comfy_genuinely_busy(q) is False
        assert arb.comfy_stale_job == "abc"

    def test_long_running_job_holding_a_checkpoint_is_real(self):
        # 6.8 GB is a loaded SDXL checkpoint; slow is not the same as stuck.
        arb = self._arb(6.8 * GB)
        q = {"queue_running": [[1, "abc"]], "queue_pending": []}
        arb._comfy_genuinely_busy(q)
        arb._running_since = time.time() - (arb.STALE_RUNNING_S + 5)
        assert arb._comfy_genuinely_busy(q) is True
        assert arb.comfy_stale_job is None

    def test_a_new_prompt_id_resets_the_staleness_clock(self):
        arb = self._arb(0.5 * GB)
        arb._comfy_genuinely_busy({"queue_running": [[1, "old"]], "queue_pending": []})
        arb._running_since = time.time() - 1000
        assert arb._comfy_genuinely_busy(
            {"queue_running": [[1, "new"]], "queue_pending": []}) is True

    def test_vram_not_utilisation_is_the_signal(self):
        # Utilisation is shared with Ollama and any third-party process, so it stayed
        # above every sensible threshold and a stuck entry never looked stale.
        assert hasattr(v.AutoArbitrator, "STALE_COMFY_BYTES")
        assert not hasattr(v.AutoArbitrator, "STALE_UTIL_PCT")