diff --git a/tenants.py b/tenants.py index efb4226..785b8f1 100644 --- a/tenants.py +++ b/tenants.py @@ -394,22 +394,32 @@ def plan_release(demanding: str, tenants_state: List[Dict[str, Any]], return {"possible": True, "reason": "enough VRAM is already free", "release": [], "shortfall_gb": 0.0} - # Any idle reclaimable tenant is a candidate, whatever its priority. An idle tenant - # is not using its VRAM, so ranking above the demander should not protect it -- an - # earlier version filtered on priority and thereby broke both directions in turn: - # ComfyUI could not preempt Ollama, and once that was corrected a starved Ollama - # could no longer reclaim from an idle ComfyUI. + # Who may be asked for memory: # - # Priority decides who is asked *first* (lowest gives up memory soonest) and, being - # applied only to idle tenants, never interrupts work. + # * any idle reclaimable tenant, whatever its rank -- idle memory is not in use; + # * a *busy* tenant that ranks strictly below the demander. + # + # That second clause is the point of the whole service and was nearly lost. Refusing + # to touch anything busy looks safe and is not: a diffusion job measured here ran for + # 46 s instead of 3 s, squeezed into 1.6 GB, because the LLM reloaded straight after + # yielding and was then protected as "busy" while ComfyUI starved. Preempting a + # lower-priority tenant is safe precisely because releasing is asynchronous -- an + # Ollama unload queues behind its running request and applies when that finishes, so + # nothing is killed mid-flight. + # + # Equal or higher priority is never interrupted, so peers cannot fight. + demander_priority = demander.get("priority", 0) candidates = [ s for s in tenants_state if s["name"] != demanding and s.get("reclaimable") - and not s.get("busy") and s.get("vram_gb", 0) > 0 + and (not s.get("busy") or s.get("priority", 0) < demander_priority) ] - candidates.sort(key=lambda s: (s.get("priority", 0), -s.get("vram_gb", 0))) + # Idle tenants first, then lowest priority: never disturb working software while + # something idle still has memory to give. + candidates.sort(key=lambda s: (bool(s.get("busy")), s.get("priority", 0), + -s.get("vram_gb", 0))) plan, freed = [], 0.0 for c in candidates: @@ -420,7 +430,7 @@ def plan_release(demanding: str, tenants_state: List[Dict[str, Any]], blockers = [ {"name": s["name"], "vram_gb": s.get("vram_gb", 0.0), - "why": ("busy" if s.get("busy") else + "why": ("busy and ranks at or above the demander" if s.get("busy") else "declares no release mechanism" if not s.get("reclaimable") else "enough was freed without it")} for s in tenants_state diff --git a/tests/test_tenants.py b/tests/test_tenants.py index 0c4e93f..e2fc28e 100644 --- a/tests/test_tenants.py +++ b/tests/test_tenants.py @@ -188,9 +188,11 @@ class TestReleasePlanning: "reclaimable": False}, {"name": "stt-relay", "priority": 70, "vram_gb": 0.8, "busy": False, "reclaimable": False}, - {"name": "ollama", "priority": 60, "vram_gb": 0.0, "busy": True, + # Shipped priorities: diffusion outranks the LLM, whose weights reload + # from page cache in seconds. + {"name": "comfyui", "priority": 60, "vram_gb": 7.0, "busy": False, "reclaimable": True}, - {"name": "comfyui", "priority": 50, "vram_gb": 7.0, "busy": False, + {"name": "ollama", "priority": 50, "vram_gb": 0.0, "busy": True, "reclaimable": True}, ] for s in base: @@ -210,13 +212,44 @@ class TestReleasePlanning: plan = T.plan_release("ollama", self._state(), free_gb=1.5, needed_gb=14.9) assert plan["release"] == ["comfyui"] - def test_a_busy_tenant_is_never_a_victim(self): + def test_a_busy_tenant_ranking_above_the_demander_is_not_a_victim(self): + # comfyui outranks ollama, so ollama may not interrupt it. state = self._state(comfyui={"busy": True}) plan = T.plan_release("ollama", state, free_gb=1.5, needed_gb=14.9) assert plan["release"] == [] - assert any(b["name"] == "comfyui" and b["why"] == "busy" + assert any(b["name"] == "comfyui" and "busy" in b["why"] for b in plan["blockers"]) + def test_a_higher_priority_demander_preempts_busy_lower_priority_work(self): + """The measured regression that made this rule necessary. + + Refusing to touch anything busy looks safe and is not. With the LLM protected as + "busy", a diffusion job ran 46 s instead of 3 s, squeezed into 1.6 GB, because + the LLM reloaded immediately after yielding and was then untouchable. Preempting + a lower-priority tenant is safe because releasing is asynchronous: an Ollama + unload queues behind its running request rather than killing it. + """ + state = [ + {"name": "comfyui", "priority": 60, "vram_gb": 1.65, "busy": True, + "reclaimable": True}, + {"name": "ollama", "priority": 50, "vram_gb": 13.03, "busy": True, + "reclaimable": True}, + ] + plan = T.plan_release("comfyui", state, free_gb=0.28, needed_gb=6.0) + assert plan["release"] == ["ollama"] + + def test_an_idle_tenant_is_preferred_over_preempting_a_busy_one(self): + state = [ + {"name": "d", "priority": 60, "vram_gb": 0.0, "busy": True, + "reclaimable": True}, + {"name": "busy_low", "priority": 10, "vram_gb": 8.0, "busy": True, + "reclaimable": True}, + {"name": "idle_high", "priority": 90, "vram_gb": 8.0, "busy": False, + "reclaimable": True}, + ] + plan = T.plan_release("d", state, free_gb=0.0, needed_gb=8.0) + assert plan["release"] == ["idle_high"] + def test_unreclaimable_tenants_are_named_as_blockers_not_ignored(self): # The user needs to know a third-party process is what stands in the way. plan = T.plan_release("ollama", self._state(), free_gb=0.0, needed_gb=15.5) @@ -250,16 +283,18 @@ class TestReleasePlanning: # The lowest-priority idle tenant gives up memory first. assert plan["release"][0] == "low" - def test_busy_work_is_never_interrupted_whatever_the_priority(self): + def test_peers_cannot_interrupt_each_other(self): + # Equal priority is never preempted, so two tenants at the same rank cannot + # fight over the card. state = [ - {"name": "demander", "priority": 99, "vram_gb": 0.0, "busy": True, + {"name": "a", "priority": 50, "vram_gb": 0.0, "busy": True, "reclaimable": True}, - {"name": "worker", "priority": 1, "vram_gb": 8.0, "busy": True, + {"name": "b", "priority": 50, "vram_gb": 8.0, "busy": True, "reclaimable": True}, ] - plan = T.plan_release("demander", state, free_gb=0.0, needed_gb=8.0) + plan = T.plan_release("a", state, free_gb=0.0, needed_gb=8.0) assert plan["release"] == [] - assert plan["blockers"][0]["why"] == "busy" + assert "busy" in plan["blockers"][0]["why"] def test_lowest_priority_is_released_first(self): state = self._state() + [