diff --git a/agents/progress.md b/agents/progress.md index e295818..dc17e9b 100644 --- a/agents/progress.md +++ b/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=, ...)` 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. + +--- diff --git a/src/gen_worker/executor.py b/src/gen_worker/executor.py index a9502b2..d84a197 100644 --- a/src/gen_worker/executor.py +++ b/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(): diff --git a/tests/test_executor_model_ops.py b/tests/test_executor_model_ops.py new file mode 100644 index 0000000..597bb3e --- /dev/null +++ b/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"