fix(executor): ModelOp OOM vocabulary + planner-audit issues #369-#371

agents/progress.md

@@ -7,7 +7,7 @@
 > Source for all findings: /home/fidika/cozy/python-gen-worker/AUDIT.md (2026-07-03 full-stack audit,
 > file:line evidence for every claim). Counterpart protocol/orchestrator issues: tensorhub #503-#510.
 
-next_id: 369
+next_id: 372
 
 ---
 
@@ -319,3 +319,63 @@ NOTE: implementation lives in gen-orchestrator (separate repo); tracked here for
 - [ ] Audit gen-orchestrator for a user-facing cancel-request API; confirm it emits `interrupt_job_cmd` to the owning worker.
 - [ ] Auto-cancel on client / response-stream disconnect in the orchestrator (production analog of #352's disconnect backstop).
 - [x] Write the single end-to-end cancellation contract doc (sources → request_id → ctx.cancel() → tenant idiom).
+
+# #369: residency registry must own the executor's pipelines — honest vram_bytes, no load races
+
+**Completed:** no
+
+Planner-audit finding (tensorhub docs/planner.md; counterpart hub issues #533-#537). The wire executor and the `Residency` registry are two disconnected worlds: loaded pipelines live in `Executor._classes[key].instance` while `Residency` entries get `obj=None` (`ModelStore.mark_in_vram(ref, vram_delta)` never passes the object). Consequences: `make_room`/`promote`/`demote` are inert in production (only the CLI serve path uses them); `demote()` on an executor-loaded ref would flip the hub's view to ON_DISK while the pipeline still sits in VRAM; and the reported footprint is wrong in the common cases below. The hub's cost-based router (tensorhub #535) and pre-positioner (#537) both consume `vram_bytes` — garbage in, garbage placement out.
+
+Bugs, file:line:
+- **Multi-model setups report 0**: `ensure_setup` books `vram_delta if len(refs)==1 else 0` per ref (executor.py:439) — every ref of a multi-binding endpoint lands IN_VRAM with `vram_bytes=0`, and with `obj=None` the `estimate_cuda_resident_gb` fallback in `track_vram` measures nothing.
+- **Concurrent setup cross-contamination**: `ensure_setup` runs BEFORE the GPU semaphore acquire (executor.py:720 vs 745), so two jobs for different specs can interleave `from_pretrained` and both allocator deltas double-count each other's weights. Loads must serialize (a load lock or the GPU semaphore) — also protects `place_pipeline`'s free-VRAM read.
+- **Device-0-only probes**: `lifecycle.free_vram_bytes()` and `Residency._default_free_vram_bytes` read `mem_get_info(0)`; multi-GPU workers report GPU0 only while `ResolvedCompute.gpu_index` spreads jobs across cards.
+
+Tasks:
+- [ ] Register the actual owner in Residency: `track_vram(ref, obj=<instance-or-teardown-handle>, ...)` from `ensure_setup`; `_unload_ref` becomes `Residency`-driven so registry state and instance state cannot diverge.
+- [ ] Per-ref attribution for multi-binding setups: measure each slot's load separately (allocator delta around each `load_from_pretrained`) instead of one blob delta around the whole setup.
+- [ ] Serialize model loads under a lock shared with the GPU semaphore path; take the allocator-delta measurement inside it.
+- [ ] Multi-GPU: sum `mem_get_info` across devices for StateDelta.free_vram_bytes (or report per-device once the hub consumes it); make Residency's probe device-aware.
+- [ ] (shipped separately) LOAD failures now emit `error="oom"` when OOM so the hub's UNLOAD-for-headroom reaction (connect_worker.handleModelFailure) is reachable — was dead code because the worker only ever said "load_failed".
+
+## Acceptance
+Hub's `WorkerInfo.Models[ref].VRAMBytes` matches `torch.cuda.memory_allocated` attribution per ref on a two-model endpoint; concurrent RunJobs for two cold specs produce non-overlapping measurements.
+
+---
+
+# #370: keep-set disk retention — `keep` is write-only, no disk GC exists
+
+**Completed:** no
+
+Contract §7 promises: "W's eviction policy never removes a `keep` ref from disk while headroom allows; if disk pressure forces it, W emits ModelEvent{EVICTED} so O knows to re-download later." None of that exists. `ModelStore.keep` is assigned from HelloAck (lifecycle.py:141) and never read again; nothing ever deletes model bytes from disk; a long-lived worker serving rotating tenant fine-tunes fills its disk until the pod dies. Downloads also start with no disk-headroom check (snapshot file sizes are known up front — `SnapshotFile.size_bytes`).
+
+Tasks:
+- [ ] Pre-download headroom gate in `ModelStore.ensure_local`: `statvfs` free vs `sum(size_bytes)` + margin; on shortfall run disk GC first; if still short emit `ModelEvent{FAILED, error="insufficient_disk"}` (vocab already exists).
+- [ ] Disk GC: LRU over Residency DISK-tier entries, never evicting `keep` refs or refs of in-flight jobs; delete bytes + `evict()` + emit EVICTED.
+- [ ] Keep-pressure escape hatch per contract: when GC of non-keep refs is insufficient, evict LRU keep refs too (still emitting EVICTED — the hub's 15s prefetch refresh re-downloads when demand returns).
+- [ ] Boot-time disk rescan: after restart, Residency starts empty while the CAS dir is full — Hello.models omits refs that are actually local until first touch. Walk the CAS dir at startup and `track_disk` complete snapshots so the hub's baseline (and GC's view) is honest.
+- [ ] e2e: fill a small TENSORHUB_CACHE_DIR, download over budget → LRU non-keep ref evicted with EVICTED event; keep ref survives.
+
+## Acceptance
+A worker with a bounded cache dir serves an unbounded rotation of refs without manual cleanup; hub residency view matches disk truth across worker restarts.
+
+---
+
+# #371: worker-side VRAM juggling — make_room before load; UNLOAD demotes instead of destroying
+
+**Completed:** no
+
+Depends on #369. Today a new pipeline load never evicts an idle endpoint's pipeline: `place_pipeline` reads free VRAM and degrades ITSELF down the offload ladder (vae_only → model_offload → ...) even when demoting an idle LRU pipeline would let both run at full speed. And hub-initiated `UNLOAD` is a full instance teardown (`_unload_ref` → shutdown + del): re-serving the model later pays a complete disk reload + setup, making the hub's future UNLOAD/LOAD juggling (tensorhub #537 pre-positioner) needlessly expensive. The warm RAM tier — designed, reported on the wire, scored by the hub's router (localityRAM) — is unreachable in the wire path.
+
+Tasks:
+- [ ] `ensure_setup` calls `store.residency.make_room(estimated_bytes)` before loading (estimate: binding resources.vram_gb, or the ref's `vram_hint` from a prior load); LRU victims demote per Residency policy. Requires #369's obj ownership so demote actually frees VRAM.
+- [ ] Demote = pipeline `.to("cpu")` keeping the `_ClassRecord` instance alive (RAM tier, honest IN_RAM event) when host RAM allows (`_RAM_FLOOR_GB` logic exists); teardown only on RAM pressure or explicit evict.
+- [ ] `_unload_ref` (hub UNLOAD) demotes to RAM/disk through the same path instead of destroying the instance; next LOAD/RunJob promotes RAM→VRAM (`.to("cuda")`) in seconds instead of a cold setup.
+- [ ] Promotion path: RunJob for a RAM-tier instance re-promotes before dispatch to the handler (executor checks tier in `ensure_setup`).
+- [ ] Re-evaluate `place_pipeline` interplay: offload ladder remains the fallback when make_room cannot reach the requirement (single model bigger than the card), not the first response to a busy card.
+- [ ] e2e (CPU-tier simulable with vram_budget_bytes): two endpoints, budget for one — alternating jobs demote/promote instead of tearing down; ModelEvents follow IN_VRAM ⇄ IN_RAM.
+
+## Acceptance
+Alternating requests across two endpoints on one GPU show demote/promote transitions (no full reloads) and both run unoffloaded when resident alone.
+
+---

src/gen_worker/executor.py

@@ -105,6 +105,16 @@ def _snapshot_to_resolved(snap: pb.Snapshot) -> Dict[str, Any]:
     }
 
 
+def _model_op_error_vocab(exc: BaseException) -> str:
+    """Contract §9 ModelEvent.error vocabulary for LOAD/UNLOAD failures."""
+    if type(exc).__name__ in ("OutOfMemoryError", "CUDAOutOfMemoryError"):
+        return "oom"
+    text = str(exc).lower()
+    if "out of memory" in text or "cuda oom" in text:
+        return "oom"
+    return "load_failed"
+
+
 def _is_terminal_download_error(exc: BaseException) -> bool:
     status = getattr(exc, "status_code", None)
     if isinstance(status, int) and 400 <= status < 500 and status != 429:
@@ -569,10 +579,13 @@ async def handle_model_op(self, op: pb.ModelOp) -> None:
         except Exception as exc:
             logger.warning("ModelOp %s on %s failed: %s", op.op, ref, exc)
             # ensure_local already emitted FAILED for download errors; emit for
-            # load/unload paths that failed outside it.
+            # load/unload paths that failed outside it. OOM must say "oom" —
+            # it is the orchestrator's trigger to UNLOAD a resident model for
+            # headroom (contract §9 vocabulary).
             if op.op != pb.MODEL_OP_KIND_DOWNLOAD:
                 await self._send(pb.WorkerMessage(model_event=pb.ModelEvent(
-                    ref=ref, state=pb.MODEL_STATE_FAILED, error="load_failed")))
+                    ref=ref, state=pb.MODEL_STATE_FAILED,
+                    error=_model_op_error_vocab(exc))))
 
     def _ref_in_use(self, ref: str) -> bool:
         for job in self.jobs.values():

tests/test_executor_model_ops.py

@@ -0,0 +1,78 @@
+"""ModelOp failure vocabulary: an OOM during LOAD must surface as
+ModelEvent{FAILED, error="oom"} — the orchestrator's trigger to UNLOAD a
+resident model for headroom. Everything else stays "load_failed"."""
+
+from __future__ import annotations
+
+import asyncio
+from pathlib import Path
+
+import msgspec
+import pytest
+
+from gen_worker.api.binding import HF
+from gen_worker.executor import Executor, _model_op_error_vocab
+from gen_worker.pb import worker_scheduler_pb2 as pb
+from gen_worker.registry import EndpointSpec
+
+
+class OutOfMemoryError(Exception):  # torch.cuda.OutOfMemoryError stand-in
+    pass
+
+
+class _In(msgspec.Struct):
+    x: str
+
+
+class _Out(msgspec.Struct):
+    y: str
+
+
+def test_model_op_error_vocab_classification() -> None:
+    assert _model_op_error_vocab(OutOfMemoryError("CUDA out of memory")) == "oom"
+    assert _model_op_error_vocab(RuntimeError("CUDA out of memory. Tried to allocate")) == "oom"
+    assert _model_op_error_vocab(RuntimeError("shape mismatch")) == "load_failed"
+
+
+@pytest.mark.parametrize(
+    "exc, expected",
+    [(OutOfMemoryError("CUDA out of memory"), "oom"),
+     (RuntimeError("weights corrupt"), "load_failed")],
+)
+def test_model_op_load_failure_emits_vocab(tmp_path, exc, expected) -> None:
+    binding = HF("acme/tiny")
+
+    class Endpoint:
+        def setup(self, model: str) -> None:
+            raise exc
+
+        def run(self, ctx, payload: _In) -> _Out:  # pragma: no cover
+            return _Out(y=payload.x)
+
+    spec = EndpointSpec(
+        name="ep", method=Endpoint.run, kind="inference",
+        payload_type=_In, output_mode="single", cls=Endpoint,
+        attr_name="run", models={"model": binding},
+    )
+
+    sent: list[pb.WorkerMessage] = []
+
+    async def _send(msg: pb.WorkerMessage) -> None:
+        sent.append(msg)
+
+    async def _run() -> None:
+        ex = Executor([spec], _send)
+
+        async def _fake_ensure_local(ref, snapshot=None, *, binding=None) -> Path:
+            return tmp_path
+
+        ex.store.ensure_local = _fake_ensure_local  # type: ignore[method-assign]
+        await ex.handle_model_op(pb.ModelOp(op=pb.MODEL_OP_KIND_LOAD, ref="acme/tiny"))
+
+    asyncio.run(_run())
+
+    failed = [m for m in sent if m.WhichOneof("msg") == "model_event"
+              and m.model_event.state == pb.MODEL_STATE_FAILED]
+    assert failed, f"no FAILED ModelEvent emitted; sent={sent}"
+    assert failed[-1].model_event.error == expected
+    assert failed[-1].model_event.ref == "acme/tiny"