From 4b26f57c56faf195698b5541adf8e79f70058f36 Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Sat, 4 Jul 2026 02:40:14 -0600 Subject: [PATCH 1/2] feat(worker): residency owns executor pipelines; disk GC; VRAM demote/promote juggling (#369, #370, #371) - #369: per-slot allocator-delta measurement registers each ref with its own measured vram_bytes AND the pipeline object; tenant-loaded refs split the residual. Loads serialize under a load lock. demote()/promote() only perform transitions they can execute (movable obj + RAM headroom); record teardown books every held ref back to disk. Free-VRAM probes sum across all CUDA devices. ModelStore.bind_loop fixes dropped thread-side residency events. - #370: pre-download disk headroom gate + LRU disk GC over DISK-tier refs (keep + in-use pins + grace, keep-pressure escape hatch emits EVICTED), persisted ref-index, boot-time rescan, fail-fast insufficient_disk. - #371: make_room before load (idle LRU pipelines demote to warm RAM instead of the new load degrading down the offload ladder); hub UNLOAD demotes instead of destroying; next RunJob/LOAD promotes RAM->VRAM. --- CHANGELOG.md | 17 ++ src/gen_worker/executor.py | 334 ++++++++++++++++++++++++----- src/gen_worker/lifecycle.py | 10 +- src/gen_worker/models/disk_gc.py | 159 ++++++++++++++ src/gen_worker/models/residency.py | 77 ++++--- tests/test_disk_gc.py | 155 +++++++++++++ tests/test_executor_residency.py | 292 +++++++++++++++++++++++++ tests/test_residency.py | 16 +- 8 files changed, 968 insertions(+), 92 deletions(-) create mode 100644 src/gen_worker/models/disk_gc.py create mode 100644 tests/test_disk_gc.py create mode 100644 tests/test_executor_residency.py diff --git a/CHANGELOG.md b/CHANGELOG.md index d0c5c83..f6a96e9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,23 @@ ## Unreleased +- **Residency unification + worker-side VRAM juggling + disk GC (#369, + #370, #371).** The `Residency` registry now owns the executor's pipelines: + worker-built pipelines register per ref with their own measured allocator + delta (multi-model endpoints no longer report `vram_bytes=0`); tenant-loaded + refs carry the residual. Model loads serialize under a load lock; free-VRAM + probes sum across all CUDA devices. `ensure_setup` runs `make_room` before + loading — idle LRU pipelines demote to the warm CPU-RAM tier instead of the + new load degrading down the offload ladder — and hub `UNLOAD` demotes + instead of destroying; the next RunJob/LOAD promotes RAM→VRAM in seconds. + `demote()` only performs transitions it can actually execute (movable + object + RAM headroom); otherwise the executor tears the owning record down + and books every ref back to disk. Disk retention exists now: a persisted + ref index, a pre-download headroom gate, LRU disk GC honoring `keep` + + in-use pins + a grace window (keep-pressure escape still emits EVICTED), + fail-fast `insufficient_disk`, and a boot-time rescan so Hello.models + matches disk truth across restarts. + - **Conversion ETL split out as `cozy-convert` (#367) — breaking.** `gen_worker.clone` and `gen_worker.conversion` are gone; the mirror / convert / publish ETL and the conversion tenant SDK (`Source`, `Dataset`, diff --git a/src/gen_worker/executor.py b/src/gen_worker/executor.py index a8f2e17..29aeb98 100644 --- a/src/gen_worker/executor.py +++ b/src/gen_worker/executor.py @@ -11,6 +11,7 @@ import asyncio import logging import os +import shutil import time import typing from dataclasses import dataclass, field as dc_field @@ -29,7 +30,9 @@ from .api.streaming import BatchItemDelta, Done, Error, IncrementalTokenDelta from .api.types import Compute from .capability import HardwareUnmetError, InsufficientDiskError +from .models import disk_gc from .models import residency as residency_mod +from .models.memory import estimate_cuda_resident_gb from .models.cache_paths import tensorhub_cas_dir from .models.download import ensure_local from .models.errors import UrlExpiredError @@ -57,6 +60,11 @@ _STUCK_THREAD_RECYCLE_S = 30.0 _DOWNLOAD_RETRIES = 3 _PROGRESS_EVENT_MIN_INTERVAL_S = 5.0 +_GiB = 1024 ** 3 +# Disk headroom preserved beyond a download's known size (#370). +_DISK_GC_MARGIN_BYTES = 2 * _GiB +# Refs used within the grace window are not disk-GC candidates. +_DISK_GC_GRACE_S = 300.0 try: # torch is optional at import time; the executor works without it. import torch @@ -146,9 +154,10 @@ def _is_terminal_download_error(exc: BaseException) -> bool: class ModelStore: - """The worker's model seam: ensure-local with retries + the residency map. - All tier transitions flow through :class:`~gen_worker.models.residency. - Residency`, whose events this store forwards as wire ``ModelEvent``s.""" + """The worker's model seam: ensure-local with retries, the residency map, + and disk retention (#370). All tier transitions flow through + :class:`~gen_worker.models.residency.Residency`, whose events this store + forwards as wire ``ModelEvent``s.""" def __init__( self, @@ -157,18 +166,41 @@ def __init__( hf_home: str = "", hf_token: str = "", cache_dir: Optional[Path] = None, + vram_budget_bytes: Optional[int] = None, + disk_free_bytes_fn: Optional[Callable[[], int]] = None, ) -> None: self._emit = emit self._hf_home = hf_home or None self._hf_token = hf_token or None self._cache_dir = cache_dir or tensorhub_cas_dir() - self.residency = Residency(on_event=self._on_residency_event) + self.residency = Residency( + on_event=self._on_residency_event, vram_budget_bytes=vram_budget_bytes, + ) self._locks: Dict[str, asyncio.Lock] = {} self.keep: set[str] = set() self._loop: Optional[asyncio.AbstractEventLoop] = None + self._index = disk_gc.RefIndex(self._cache_dir) + self._disk_free = disk_free_bytes_fn or self._default_disk_free + + def _default_disk_free(self) -> int: + p = Path(self._cache_dir) + for candidate in (p, *p.parents): # cache dir may not exist yet + try: + return int(shutil.disk_usage(candidate).free) + except OSError: + continue + return 0 # ---- events ------------------------------------------------------------ + def bind_loop(self) -> None: + """Capture the running loop so residency events raised from worker + threads (demote/promote via to_thread) still reach the wire.""" + try: + self._loop = asyncio.get_running_loop() + except RuntimeError: + pass + def _on_residency_event(self, ref: str, state: str, vram_bytes: int) -> None: pb_state = _RESIDENCY_STATE_TO_PB.get(state) if pb_state is None: @@ -201,14 +233,72 @@ def residency_snapshot(self) -> List[pb.ModelResidency]: def local_path(self, ref: str) -> Optional[Path]: return self.residency.local_path(ref) - def mark_in_vram(self, ref: str, vram_bytes: int, obj: Any = None) -> None: - self.residency.track_vram(ref, obj, vram_bytes=int(vram_bytes)) + # ---- disk retention (#370) ------------------------------------------------ - def mark_in_ram(self, ref: str, obj: Any = None) -> None: - self.residency.track_ram(ref, obj) + def rescan_disk(self) -> None: + """Boot-time truth: re-register still-present downloads from the + persisted ref index so Hello.models and GC see what disk holds.""" + for ref, ent in self._index.entries().items(): + p = Path(str(ent.get("path") or "")) + if p.exists(): + if self.residency.tier(ref) is None: + self.residency.track_disk(ref, p) + else: + self._index.remove(ref) + + def gc_disk(self, target_free_bytes: int, *, exclude: Tuple[str, ...] = ()) -> None: + """Evict LRU disk-tier refs until free disk reaches the target. + Non-keep refs go first (grace-honoring, then grace-ignoring); under + keep-pressure the escape hatch evicts LRU `keep` refs too (contract §7 + — EVICTED is emitted so the hub re-downloads when demand returns). + In-use / loaded refs are never touched.""" + for include_keep, honor_grace in ( + (False, True), (False, False), (True, True), (True, False), + ): + for ref in self._gc_candidates(include_keep, honor_grace, exclude): + if self._disk_free() >= target_free_bytes: + return + self._evict_disk_ref(ref) + + def _gc_candidates( + self, include_keep: bool, honor_grace: bool, exclude: Tuple[str, ...] + ) -> List[str]: + now = time.time() + out: List[Tuple[float, str]] = [] + for ref in self.residency.refs_in(residency_mod.Tier.DISK): + if ref in exclude or self.residency.in_use(ref): + continue + if (ref in self.keep) != include_keep: + continue + last = self._index.last_used(ref) + if honor_grace and (now - last) < _DISK_GC_GRACE_S: + continue + out.append((last, ref)) + out.sort() + return [r for _, r in out] - def mark_unloaded(self, ref: str) -> None: - self.residency.release_to_disk(ref) + def _evict_disk_ref(self, ref: str) -> None: + path = self.residency.local_path(ref) or self._index.path(ref) + if not self.residency.evict(ref): # refuses in-use entries; emits EVICTED + return + if path is not None: + disk_gc.delete_ref_bytes(ref, path, self._cache_dir) + disk_gc.sweep_orphan_blobs(self._cache_dir) + self._index.remove(ref) + + async def _ensure_disk_headroom(self, ref: str, needed_bytes: int) -> None: + target = int(needed_bytes) + _DISK_GC_MARGIN_BYTES + if self._disk_free() >= target: + return + await asyncio.to_thread(self.gc_disk, target, exclude=(ref,)) + free = self._disk_free() + if free < target: + await self._event(ref, pb.MODEL_STATE_FAILED, error="insufficient_disk") + raise InsufficientDiskError( + f"need {needed_bytes} bytes for {ref}; {free} free after disk GC", + available_bytes=free, required_bytes=needed_bytes, + path=str(self._cache_dir), + ) # ---- ensure-local ---------------------------------------------------------- @@ -225,11 +315,18 @@ async def ensure_local( """Materialize `ref` on disk. Transient failures retry with backoff; terminal (4xx-class) failures raise immediately. Emits ModelEvents. ``binding`` (when known) supplies provider + file-selection metadata.""" + self.bind_loop() async with self._lock(ref): cached = self.residency.local_path(ref) if cached is not None and cached.exists(): + self._index.touch(ref) return cached - self._loop = asyncio.get_running_loop() + if snapshot is not None and snapshot.files: + # Sizes are known up front for tensorhub snapshots: gate on + # disk headroom, GC-ing LRU refs first (#370). + await self._ensure_disk_headroom( + ref, sum(int(f.size_bytes) for f in snapshot.files) + ) last_progress = 0.0 def _progress(done: int, total: Optional[int]) -> None: @@ -263,6 +360,7 @@ def _progress(done: int, total: Optional[int]) -> None: progress=_progress, ) self.residency.track_disk(ref, path) + self._index.record(ref, path, disk_gc.tree_bytes(path)) return path except Exception as exc: terminal = _is_terminal_download_error(exc) or attempt >= _DOWNLOAD_RETRIES @@ -337,6 +435,9 @@ def __init__( self._settings = settings self.store = store or ModelStore(send) self._gpu_semaphore = asyncio.Semaphore(max(1, gpu_slots)) + # Model loads/promotions serialize so allocator-delta measurements + # and free-VRAM reads don't cross-contaminate (#369). + self._load_lock = asyncio.Lock() self._on_state_change = on_state_change or (lambda: None) self.file_base_url: str = "" self.draining = False @@ -424,9 +525,11 @@ def in_flight_keys(self) -> List[Tuple[str, int]]: async def ensure_setup(self, spec: EndpointSpec, snapshots: Optional[Dict[str, pb.Snapshot]] = None) -> Any: if spec.cls is None: return None # function-shaped endpoint: no instance, no setup + self.store.bind_loop() rec = self._classes[spec.instance_key] async with rec.lock: if rec.ready: + await self._promote_setup_refs(spec) return rec.instance setup_slots = self._setup_slots(spec) paths: Dict[str, str] = {} @@ -436,37 +539,106 @@ async def ensure_setup(self, spec: EndpointSpec, snapshots: Optional[Dict[str, p snap = (snapshots or {}).get(ref) path = await self.store.ensure_local(ref, snap, binding=binding) paths[slot] = str(path) - instance = spec.cls() - setup = getattr(instance, "setup", None) - vram_before = self._vram_allocated() - if spec.runtime: - rec.server = await self._boot_engine_server(spec, paths) - if callable(setup): - kwargs = await self._injection_kwargs(spec, setup, paths, server=rec.server) - if asyncio.iscoroutinefunction(setup): - await setup(**kwargs) - else: - await asyncio.to_thread(setup, **kwargs) - warmup = getattr(instance, "warmup", None) - if callable(warmup): - if asyncio.iscoroutinefunction(warmup): - await warmup() - else: - await asyncio.to_thread(warmup) - # Measured VRAM (allocator delta across the load) — the number the - # orchestrator's VRAM packer sees; never an estimate constant. - vram_delta = max(0, self._vram_allocated() - vram_before) - refs = [wire_ref(spec.models[s]) for s in setup_slots] - for ref in refs: - if vram_delta > 0: - self.store.mark_in_vram(ref, vram_delta if len(refs) == 1 else 0) - else: - self.store.mark_in_ram(ref) + # Loads serialize: concurrent setups would cross-contaminate each + # other's allocator deltas and place_pipeline's free-VRAM reads. + async with self._load_lock: + await self._make_room_for(spec, setup_slots) + instance = spec.cls() + setup = getattr(instance, "setup", None) + loaded: Dict[str, Tuple[Any, int]] = {} + vram_before = self._vram_allocated() + if spec.runtime: + rec.server = await self._boot_engine_server(spec, paths) + if callable(setup): + kwargs, loaded = await self._injection_kwargs( + spec, setup, paths, server=rec.server) + if asyncio.iscoroutinefunction(setup): + await setup(**kwargs) + else: + await asyncio.to_thread(setup, **kwargs) + warmup = getattr(instance, "warmup", None) + if callable(warmup): + if asyncio.iscoroutinefunction(warmup): + await warmup() + else: + await asyncio.to_thread(warmup) + vram_delta = max(0, self._vram_allocated() - vram_before) + self._register_residency(spec, setup_slots, loaded, vram_delta) rec.instance = instance rec.ready = True self._on_state_change() return instance + def _register_residency( + self, + spec: EndpointSpec, + setup_slots: List[str], + loaded: Dict[str, Tuple[Any, int]], + total_delta: int, + ) -> None: + """Honest per-ref residency after a setup (#369). Worker-constructed + pipelines carry their own measured allocator delta AND the object + (Residency owns it: demote/promote actually move memory). Refs the + tenant loaded inside setup() split the residual delta — no object, + so their VRAM is only reclaimable by record teardown.""" + res = self.store.residency + per_ref: Dict[str, Tuple[Any, int]] = {} + for slot in setup_slots: + ref = wire_ref(spec.models[slot]) + obj, measured = loaded.get(slot, (None, 0)) + prev_obj, prev_bytes = per_ref.get(ref, (None, 0)) + per_ref[ref] = (obj or prev_obj, prev_bytes + measured) + residual = max(0, total_delta - sum(b for _, b in per_ref.values())) + opaque = [r for r, (obj, _) in per_ref.items() if obj is None] + share = residual // len(opaque) if opaque else 0 + for ref, (obj, measured) in per_ref.items(): + vram = measured + (share if obj is None else 0) + if vram > 0: + res.track_vram(ref, obj, vram_bytes=vram) + elif obj is not None and int(estimate_cuda_resident_gb(obj) * _GiB) > 0: + res.track_vram(ref, obj) # measured via cuda-resident estimate + else: + res.track_ram(ref, obj) # CPU-only host / offloaded load + + async def _promote_setup_refs(self, spec: EndpointSpec) -> None: + """RunJob/LOAD for a demoted (RAM-tier) instance: swap the pipelines + back into VRAM instead of a cold reload (#371).""" + res = self.store.residency + refs = [wire_ref(spec.models[s]) for s in self._setup_slots(spec)] + if any(res.tier(r) is residency_mod.Tier.RAM for r in refs): + async with self._load_lock: + for ref in refs: + if res.tier(ref) is residency_mod.Tier.RAM: + await asyncio.to_thread(res.promote, ref) + self._on_state_change() + for ref in refs: + res.touch(ref) + + async def _make_room_for(self, spec: EndpointSpec, setup_slots: List[str]) -> None: + """Evict idle LRU pipelines before loading instead of degrading the + new load down the offload ladder (#371). Estimate: per-ref vram_hint + from a prior load, else the endpoint's declared vram_gb.""" + res = self.store.residency + refs = [wire_ref(spec.models[s]) for s in setup_slots] + needed = sum(res.vram_hint(r) for r in refs) + if needed <= 0 and spec.resources.vram_gb: + needed = int(float(spec.resources.vram_gb) * _GiB) + if needed <= 0: + return + if await asyncio.to_thread(res.make_room, needed): + self._on_state_change() + return + # Movable demotions weren't enough: tear down idle records holding + # non-movable LRU victims (tenant-loaded refs). + for ref in res.lru_vram_victims(): + rec = self._record_holding(ref) + if rec is None or self._record_in_use(rec): + continue + await self._vacate_record(rec) + if await asyncio.to_thread(res.make_room, needed): + break + self._on_state_change() + @staticmethod def _setup_slots(spec: EndpointSpec) -> List[str]: """Model slots loaded once at setup time. Classes without setup() @@ -497,13 +669,17 @@ async def _injection_kwargs( paths: Dict[str, str], *, server: Any = None, - ) -> Dict[str, Any]: + ) -> Tuple[Dict[str, Any], Dict[str, Tuple[Any, int]]]: """Typed injection: each slot receives exactly what its ``setup`` annotation says — a ``str``/``Path`` local path, or a constructed pipeline for a class annotation exposing ``from_pretrained`` (built off the loop; the binding dtype is honored and the worker applies its placement/offload policy to the result). A parameter annotated - ``ServerHandle`` receives the booted engine server.""" + ``ServerHandle`` receives the booted engine server. + + Returns ``(kwargs, loaded)`` where ``loaded`` maps each pipeline slot + to ``(pipeline, allocator_delta)`` — the per-ref measurement residency + books (#369; caller holds the load lock).""" from .runtimes.server import ServerHandle try: @@ -511,6 +687,7 @@ async def _injection_kwargs( except Exception: hints = {} kwargs: Dict[str, Any] = {} + loaded: Dict[str, Tuple[Any, int]] = {} if server is not None: for pname, ann in hints.items(): if ann is ServerHandle: @@ -527,16 +704,18 @@ async def _injection_kwargs( binding = spec.models.get(slot) dtype = str(getattr(binding, "dtype", "") or "") + before = self._vram_allocated() pipe = await asyncio.to_thread( load_from_pretrained, ann, path, dtype=dtype ) # Worker-owned placement/offload policy: one decider for the # whole worker; endpoints never write device/offload code. await asyncio.to_thread(place_pipeline, pipe) + loaded[slot] = (pipe, max(0, self._vram_allocated() - before)) kwargs[slot] = pipe else: kwargs[slot] = path - return kwargs + return kwargs, loaded @staticmethod def _vram_allocated() -> int: @@ -566,6 +745,7 @@ async def shutdown_instances(self) -> None: # ---- ModelOp ----------------------------------------------------------- async def handle_model_op(self, op: pb.ModelOp) -> None: + self.store.bind_loop() ref = op.ref snap = op.snapshot if op.HasField("snapshot") else None try: @@ -608,30 +788,64 @@ def _ref_in_use(self, ref: str) -> bool: return True return False - async def _unload_ref(self, ref: str) -> None: - """Demote a ref out of VRAM by tearing down the classes that hold it.""" + def _record_holding(self, ref: str) -> Optional[_ClassRecord]: for rec in self._classes.values(): - holds = any( - ref in (wire_ref(b) for b in s.models.values()) for s in rec.specs - ) - if holds and rec.ready: - inst, rec.instance, rec.ready = rec.instance, None, False - shutdown = getattr(inst, "shutdown", None) - if callable(shutdown): - try: - await asyncio.to_thread(shutdown) - except Exception: - logger.exception("shutdown() during UNLOAD failed") - del inst - server, rec.server = rec.server, None - if server is not None: - await asyncio.to_thread(server.stop) + if not rec.ready: + continue + if any(ref in (wire_ref(b) for b in s.models.values()) for s in rec.specs): + return rec + return None + + def _record_in_use(self, rec: _ClassRecord) -> bool: + for s in rec.specs: + for b in s.models.values(): + ref = wire_ref(b) + if self._ref_in_use(ref) or self.store.residency.in_use(ref): + return True + return False + + async def _vacate_record(self, rec: _ClassRecord) -> None: + """Tear an instance down and book every ref it held back to disk — + registry state and instance state move together (#369).""" + inst, rec.instance, rec.ready = rec.instance, None, False + shutdown = getattr(inst, "shutdown", None) + if inst is not None and callable(shutdown): + try: + if asyncio.iscoroutinefunction(shutdown): + await shutdown() + else: + await asyncio.to_thread(shutdown) + except Exception: + logger.exception("shutdown() during vacate failed") + del inst + server, rec.server = rec.server, None + if server is not None: + await asyncio.to_thread(server.stop) if torch is not None and torch.cuda.is_available(): try: await asyncio.to_thread(torch.cuda.empty_cache) except Exception: pass - self.store.mark_unloaded(ref) + for s in rec.specs: + for b in s.models.values(): + self.store.residency.release_to_disk(wire_ref(b)) + self._on_state_change() + + async def _unload_ref(self, ref: str) -> None: + """Hub UNLOAD: free the ref's VRAM. Worker-owned pipelines demote to + the warm RAM tier (instance stays ready; the next LOAD/RunJob promotes + back in seconds); tenant-loaded refs require record teardown (#371).""" + async with self._load_lock: + res = self.store.residency + if res.tier(ref) is residency_mod.Tier.VRAM: + if await asyncio.to_thread(res.demote, ref): + self._on_state_change() + return + rec = self._record_holding(ref) + if rec is not None: + await self._vacate_record(rec) + else: + res.release_to_disk(ref) self._on_state_change() # ---- job intake -------------------------------------------------------- @@ -749,9 +963,11 @@ async def _run_job(self, job: _Job, run: pb.RunJob) -> None: await self._finish(job, pb.JOB_STATUS_CANCELED, safe_message="canceled") return except Exception as exc: - if isinstance(exc, HardwareUnmetError): + if isinstance(exc, HardwareUnmetError) and not isinstance(exc, InsufficientDiskError): # Self-disable the function on this worker; lifecycle emits # FnUnavailable and drops it from available_functions. + # (Disk pressure is transient — GC frees space — so it only + # fails the job RETRYABLE, never disables the function.) axes = exc.axes() if hasattr(exc, "axes") else {} self.unavailable[spec.name] = ( getattr(exc, "reason", "hardware_unmet"), _sanitize(str(exc)), diff --git a/src/gen_worker/lifecycle.py b/src/gen_worker/lifecycle.py index 4503d48..71451fe 100644 --- a/src/gen_worker/lifecycle.py +++ b/src/gen_worker/lifecycle.py @@ -52,12 +52,15 @@ def probe_hardware() -> Dict[str, Any]: def free_vram_bytes() -> int: + """Free VRAM summed across ALL CUDA devices (StateDelta.free_vram_bytes).""" try: import torch if torch.cuda.is_available(): - free_mem, _total = torch.cuda.mem_get_info(0) - return int(free_mem) + return sum( + int(torch.cuda.mem_get_info(i)[0]) + for i in range(torch.cuda.device_count()) + ) except Exception: pass return 0 @@ -211,6 +214,9 @@ async def startup(self) -> None: """Gate functions, prefetch worker-fetchable models with retry/backoff, set up endpoints, advance phases. Never raises: failures gate individual functions, not the process.""" + # Disk truth first: after a restart the CAS dir is full while Residency + # starts empty — rescan so Hello.models (and disk GC) see reality. + self.executor.store.rescan_disk() self.executor.gate_functions(self.hardware) from .api.binding import wire_ref diff --git a/src/gen_worker/models/disk_gc.py b/src/gen_worker/models/disk_gc.py new file mode 100644 index 0000000..4aaa3f3 --- /dev/null +++ b/src/gen_worker/models/disk_gc.py @@ -0,0 +1,159 @@ +"""Disk retention (#370): a persistent ref->bytes index + deletion helpers. + +The CAS stores tensorhub models by snapshot digest, HF models under the HF +cache, civitai under version dirs — none of which map back to wire refs on +their own. ``RefIndex`` persists {ref: path, bytes, last_used} at +``/ref-index.json`` so disk GC and the boot-time rescan can reason +in refs (the vocabulary of `keep`, Residency, and ModelEvents). +""" + +from __future__ import annotations + +import json +import logging +import os +import shutil +import threading +import time +from pathlib import Path +from typing import Any, Dict, Optional + +logger = logging.getLogger(__name__) + +_INDEX_NAME = "ref-index.json" + + +class RefIndex: + """Persistent {ref: {path, bytes, last_used}}. Thread-safe; every mutation + is written through (small file, rare writes).""" + + def __init__(self, cache_dir: Path) -> None: + self._path = Path(cache_dir) / _INDEX_NAME + self._lock = threading.Lock() + self._data: Dict[str, Dict[str, Any]] = {} + try: + raw = json.loads(self._path.read_text("utf-8")) + if isinstance(raw, dict): + self._data = { + str(k): v for k, v in raw.items() + if isinstance(v, dict) and v.get("path") + } + except FileNotFoundError: + pass + except Exception as exc: + logger.warning("ref-index unreadable (%s); starting empty", exc) + + def _save_locked(self) -> None: + try: + self._path.parent.mkdir(parents=True, exist_ok=True) + tmp = self._path.with_suffix(".tmp") + tmp.write_text(json.dumps(self._data), encoding="utf-8") + tmp.replace(self._path) + except Exception as exc: + logger.warning("ref-index write failed: %s", exc) + + def record(self, ref: str, path: Path, size_bytes: int) -> None: + with self._lock: + self._data[ref] = { + "path": str(path), "bytes": int(size_bytes), "last_used": time.time(), + } + self._save_locked() + + def touch(self, ref: str) -> None: + with self._lock: + e = self._data.get(ref) + if e is not None: + e["last_used"] = time.time() + self._save_locked() + + def remove(self, ref: str) -> None: + with self._lock: + if self._data.pop(ref, None) is not None: + self._save_locked() + + def path(self, ref: str) -> Optional[Path]: + with self._lock: + e = self._data.get(ref) + return Path(e["path"]) if e else None + + def last_used(self, ref: str) -> float: + with self._lock: + e = self._data.get(ref) + return float(e.get("last_used") or 0.0) if e else 0.0 + + def entries(self) -> Dict[str, Dict[str, Any]]: + with self._lock: + return {k: dict(v) for k, v in self._data.items()} + + +def tree_bytes(path: Path) -> int: + """Bytes under ``path`` (file or tree), hardlinked inodes counted once.""" + p = Path(path) + try: + if p.is_file(): + return int(p.stat().st_size) + total = 0 + seen: set = set() + for dirpath, _dirs, names in os.walk(p): + for name in names: + try: + st = os.stat(os.path.join(dirpath, name)) + except OSError: + continue + key = (st.st_dev, st.st_ino) + if key in seen: + continue + seen.add(key) + total += int(st.st_size) + return total + except OSError: + return 0 + + +def _retention_unit(path: Path, cas_dir: Path) -> Path: + """The directory/file that must be deleted to reclaim a ref's bytes: + the snapshot dir for CAS refs, the ``models--*`` repo dir for HF refs, + the tracked path otherwise.""" + p = Path(path) + snaps_root = Path(cas_dir) / "snapshots" + try: + rel = p.resolve().relative_to(snaps_root.resolve()) + return snaps_root / rel.parts[0] + except (ValueError, OSError): + pass + for parent in (p, *p.parents): + if parent.name.startswith("models--"): + return parent + return p + + +def delete_ref_bytes(ref: str, path: Path, cas_dir: Path) -> None: + unit = _retention_unit(path, cas_dir) + logger.info("disk-gc: deleting %s (%s)", ref, unit) + if unit.is_dir(): + shutil.rmtree(unit, ignore_errors=True) + else: + unit.unlink(missing_ok=True) + + +def sweep_orphan_blobs(cas_dir: Path) -> int: + """Delete CAS blobs no snapshot links anymore (st_nlink == 1). Snapshot + trees hardlink into blobs/, so link count is the reference count.""" + blobs = Path(cas_dir) / "blobs" + freed = 0 + if not blobs.is_dir(): + return 0 + for dirpath, _dirs, names in os.walk(blobs): + for name in names: + fp = os.path.join(dirpath, name) + try: + st = os.stat(fp) + if st.st_nlink <= 1: + os.unlink(fp) + freed += int(st.st_size) + except OSError: + continue + return freed + + +__all__ = ["RefIndex", "tree_bytes", "delete_ref_bytes", "sweep_orphan_blobs"] diff --git a/src/gen_worker/models/residency.py b/src/gen_worker/models/residency.py index 3b228a7..9b9e75f 100644 --- a/src/gen_worker/models/residency.py +++ b/src/gen_worker/models/residency.py @@ -41,8 +41,8 @@ _GiB = 1024 ** 3 # Free-VRAM slack preserved beyond the requested headroom (activations). _VRAM_MARGIN_BYTES = 2 * _GiB -# Host-RAM floor below which VRAM demotions skip the warm tier and drop -# straight to disk (don't push the host into swap). +# Host-RAM floor below which the warm RAM tier is refused (don't push the +# host into swap); demote() then fails and the owner tears down instead. _RAM_FLOOR_GB = 8.0 # Residency event states (mirrors the wire ModelEvent vocabulary). @@ -72,14 +72,28 @@ class _Entry: refcount: int = 0 # live executions / shared holders last_used: float = field(default_factory=time.monotonic) + @property + def movable(self) -> bool: + """True when the registry can actually move this object between + devices (``.to()``); offload-hooked pipelines own their placement.""" + return ( + self.obj is not None + and callable(getattr(self.obj, "to", None)) + and not _obj_manages_own_device(self.obj) + ) + def _default_free_vram_bytes() -> int: + """Free VRAM summed across ALL CUDA devices (multi-GPU workers spread + jobs across cards; a device-0-only probe under-reports).""" try: import torch if torch.cuda.is_available(): - free, _total = torch.cuda.mem_get_info(0) - return int(free) + return sum( + int(torch.cuda.mem_get_info(i)[0]) + for i in range(torch.cuda.device_count()) + ) except Exception: pass return 0 @@ -166,6 +180,18 @@ def vram_bytes(self, ref: str) -> int: e = self._entries.get(ref) return e.vram_bytes if e else 0 + def vram_hint(self, ref: str) -> int: + """Last measured VRAM footprint (survives demotion) — the load-size + estimate for make_room before a re-load/promotion.""" + with self._lock: + e = self._entries.get(ref) + return e.vram_hint if e else 0 + + def movable(self, ref: str) -> bool: + with self._lock: + e = self._entries.get(ref) + return bool(e and e.movable) + def refs_in(self, tier: Tier) -> List[str]: with self._lock: return [r for r, e in self._entries.items() if e.tier is tier] @@ -239,34 +265,29 @@ def track_vram( self._emit(ref, IN_VRAM, max(0, measured)) def demote(self, ref: str) -> bool: - """VRAM -> RAM warm tier (or straight to disk when host RAM is tight). - Refuses pinned / executing entries. Returns True when VRAM was freed.""" + """VRAM -> RAM warm tier. Only performs transitions it can actually + execute: the entry must hold a movable object and host RAM must have + headroom — otherwise it refuses (False) and the OWNER of the memory + (executor record teardown) must free it and book the result. Never + books a state it didn't produce. Refuses pinned / executing entries.""" with self._lock: e = self._entries.get(ref) if e is None or e.tier is not Tier.VRAM or e.pinned or e.refcount > 0: return False - keep_warm = e.obj is not None and get_available_ram_gb() >= _RAM_FLOOR_GB + if not e.movable or get_available_ram_gb() < _RAM_FLOOR_GB: + return False + self._move(e.obj, "cpu") + e.tier = Tier.RAM e.vram_bytes = 0 - if keep_warm: - self._move(e.obj, "cpu") - e.tier = Tier.RAM - state = IN_RAM - elif e.path is not None: - e.obj = None - e.tier = Tier.DISK - state = ON_DISK - else: - del self._entries[ref] - state = EVICTED flush_memory() - self._emit(ref, state) + self._emit(ref, IN_RAM) return True def promote(self, ref: str, device: str = "cuda") -> bool: """RAM -> VRAM (makes room first). True when resident afterward.""" with self._lock: e = self._entries.get(ref) - if e is None or e.obj is None: + if e is None or not e.movable: return False if e.tier is Tier.VRAM: e.last_used = time.monotonic() @@ -275,7 +296,7 @@ def promote(self, ref: str, device: str = "cuda") -> bool: self.make_room(hint) with self._lock: e = self._entries.get(ref) - if e is None or e.obj is None: + if e is None or not e.movable: return False self._move(e.obj, device) e.tier = Tier.VRAM @@ -367,7 +388,8 @@ def in_use(self, ref: str) -> bool: # ---- pressure ------------------------------------------------------------- - def _lru_vram_victims(self) -> List[str]: + def lru_vram_victims(self) -> List[str]: + """Evictable VRAM refs, LRU first (pinned/executing excluded).""" with self._lock: candidates = [ e for e in self._entries.values() @@ -378,13 +400,16 @@ def _lru_vram_victims(self) -> List[str]: def make_room(self, needed_bytes: int) -> bool: """Demote LRU VRAM entries until measured free VRAM covers - ``needed_bytes`` + margin. True when the headroom was reached.""" + ``needed_bytes`` + margin. True when the headroom was reached. + Only movable entries are demoted here; when this returns False the + caller (executor) tears down non-movable LRU victims itself.""" target = int(needed_bytes) + _VRAM_MARGIN_BYTES if self.free_vram_bytes() >= target: return True - for ref in self._lru_vram_victims(): - logger.info("residency: demoting LRU %s for %d bytes headroom", ref, needed_bytes) - self.demote(ref) + for ref in self.lru_vram_victims(): + if not self.demote(ref): + continue + logger.info("residency: demoted LRU %s for %d bytes headroom", ref, needed_bytes) if self.free_vram_bytes() >= target: return True return self.free_vram_bytes() >= target diff --git a/tests/test_disk_gc.py b/tests/test_disk_gc.py new file mode 100644 index 0000000..123c6eb --- /dev/null +++ b/tests/test_disk_gc.py @@ -0,0 +1,155 @@ +"""Disk retention (#370) — real files in a budgeted tmp cache dir. + +Drives the REAL ModelStore.ensure_local path (download layer stubbed to write +actual bytes) with an injected free-disk probe: LRU non-keep eviction with +EVICTED events, keep-pressure escape hatch, in-use pins, fail-fast +insufficient_disk, and the boot-time rescan baseline. +""" + +from __future__ import annotations + +import asyncio +from pathlib import Path + +import pytest + +import gen_worker.executor as executor_mod +from gen_worker.capability import InsufficientDiskError +from gen_worker.executor import ModelStore +from gen_worker.models.residency import Tier +from gen_worker.pb import worker_scheduler_pb2 as pb + +_BUDGET = 10_000 + + +def _snapshot(digest: str, size: int) -> pb.Snapshot: + return pb.Snapshot(digest=digest, files=[ + pb.SnapshotFile(path="w.bin", size_bytes=size, blake3="ab" * 16, + url="http://example.invalid/w.bin"), + ]) + + +@pytest.fixture(autouse=True) +def _tight_budget(monkeypatch): + monkeypatch.setattr(executor_mod, "_DISK_GC_MARGIN_BYTES", 0) + monkeypatch.setattr(executor_mod, "_DISK_GC_GRACE_S", 0.0) + + +def _store(tmp_path: Path, sent: list) -> ModelStore: + async def _emit(msg: pb.WorkerMessage) -> None: + sent.append(msg) + + def _free() -> int: + from gen_worker.models.disk_gc import tree_bytes + + return max(0, _BUDGET - tree_bytes(tmp_path)) + + return ModelStore(_emit, cache_dir=tmp_path, disk_free_bytes_fn=_free) + + +@pytest.fixture() +def _fake_download(monkeypatch, tmp_path): + """Stub the provider download layer: write real bytes into the CAS + snapshots layout so GC deletes real files.""" + + async def fake_ensure_local(ref, *, snapshot=None, cache_dir=None, **kw) -> Path: + size = sum(int(f["size_bytes"]) for f in (snapshot or {}).get("files", [])) + d = Path(cache_dir) / "snapshots" / ref.replace("/", "--") + d.mkdir(parents=True, exist_ok=True) + (d / "w.bin").write_bytes(b"\0" * size) + return d + + monkeypatch.setattr(executor_mod, "ensure_local", fake_ensure_local) + + +def _evicted(sent) -> list: + return [m.model_event.ref for m in sent if m.WhichOneof("msg") == "model_event" + and m.model_event.state == pb.MODEL_STATE_EVICTED] + + +def _failed(sent) -> list: + return [(m.model_event.ref, m.model_event.error) for m in sent + if m.WhichOneof("msg") == "model_event" + and m.model_event.state == pb.MODEL_STATE_FAILED] + + +def test_gc_evicts_lru_nonkeep_and_spares_recent_and_keep(tmp_path, _fake_download) -> None: + sent: list = [] + store = _store(tmp_path, sent) + store.keep = {"t/keep"} + + async def _run() -> None: + await store.ensure_local("t/a", _snapshot("da", 3000)) + await store.ensure_local("t/b", _snapshot("db", 3000)) + await store.ensure_local("t/keep", _snapshot("dk", 1000)) + await store.ensure_local("t/a", _snapshot("da", 3000)) # touch: b is LRU + # 7000 used; 4000 needed -> GC must evict exactly the LRU non-keep ref. + await store.ensure_local("t/c", _snapshot("dc", 4000)) + + asyncio.run(_run()) + assert _evicted(sent) == ["t/b"] + assert store.residency.tier("t/b") is None + assert not (tmp_path / "snapshots" / "t--b").exists() # bytes actually gone + for ref in ("t/a", "t/keep", "t/c"): + assert store.residency.tier(ref) is Tier.DISK, ref + + +def test_keep_pressure_escape_hatch_evicts_keep_with_event(tmp_path, _fake_download) -> None: + sent: list = [] + store = _store(tmp_path, sent) + store.keep = {"t/keep"} + + async def _run() -> None: + await store.ensure_local("t/keep", _snapshot("dk", 3000)) + await store.ensure_local("t/big", _snapshot("dbig", 8000)) + + asyncio.run(_run()) + assert _evicted(sent) == ["t/keep"] # EVICTED still emitted -> hub re-downloads + assert store.residency.tier("t/big") is Tier.DISK + + +def test_insufficient_disk_fails_fast_with_event(tmp_path, _fake_download) -> None: + sent: list = [] + store = _store(tmp_path, sent) + + async def _run() -> None: + with pytest.raises(InsufficientDiskError): + await store.ensure_local("t/huge", _snapshot("dh", 50_000)) + + asyncio.run(_run()) + assert ("t/huge", "insufficient_disk") in _failed(sent) + + +def test_gc_never_evicts_in_use_refs(tmp_path, _fake_download) -> None: + sent: list = [] + store = _store(tmp_path, sent) + + async def _run() -> None: + await store.ensure_local("t/a", _snapshot("da", 6000)) + with store.residency.executing("t/a"): + with pytest.raises(InsufficientDiskError): + await store.ensure_local("t/b", _snapshot("db", 6000)) + + asyncio.run(_run()) + assert _evicted(sent) == [] + assert store.residency.tier("t/a") is Tier.DISK + + +def test_rescan_restores_disk_baseline_after_restart(tmp_path, _fake_download) -> None: + sent: list = [] + store = _store(tmp_path, sent) + + async def _run() -> None: + await store.ensure_local("t/a", _snapshot("da", 2000)) + await store.ensure_local("t/b", _snapshot("db", 2000)) + + asyncio.run(_run()) + (tmp_path / "snapshots" / "t--b" / "w.bin").unlink() # b lost outside our control + (tmp_path / "snapshots" / "t--b").rmdir() + + restarted = _store(tmp_path, []) # fresh process: Residency starts empty + assert restarted.residency.tier("t/a") is None + restarted.rescan_disk() + assert restarted.residency.tier("t/a") is Tier.DISK + assert restarted.residency.local_path("t/a") == tmp_path / "snapshots" / "t--a" + assert restarted.residency.tier("t/b") is None # stale entry dropped diff --git a/tests/test_executor_residency.py b/tests/test_executor_residency.py new file mode 100644 index 0000000..cf34873 --- /dev/null +++ b/tests/test_executor_residency.py @@ -0,0 +1,292 @@ +"""Executor <-> Residency unification (#369/#371) — CPU-only, deterministic. + +A fake CUDA allocator (module counter) + fake pipelines with ``from_pretrained`` +drive the REAL ensure_setup / handle_model_op paths against a budgeted +Residency: per-ref measured vram_bytes, serialized loads, UNLOAD demoting to +the warm RAM tier, and make_room-before-load demote/promote swaps. +""" + +import asyncio +import time +from pathlib import Path + +import msgspec +import pytest + +from gen_worker.api.binding import HF +from gen_worker.api.decorators import Resources +from gen_worker.executor import Executor, ModelStore +from gen_worker.models import residency as residency_mod +from gen_worker.models.residency import Tier +from gen_worker.pb import worker_scheduler_pb2 as pb +from gen_worker.registry import EndpointSpec + +_GiB = 1024 ** 3 + + +class _In(msgspec.Struct): + x: str + + +class _Alloc: + bytes = 0 + + +class _FakePipe: + """Stands in for a diffusers pipeline: from_pretrained allocates, .to() + moves the same object the endpoint instance holds.""" + + size = 1 * _GiB + load_sleep_s = 0.0 + + def __init__(self) -> None: + self.device = "cuda" + + @classmethod + def from_pretrained(cls, path, **kwargs): + _Alloc.bytes += cls.size + if cls.load_sleep_s: + time.sleep(cls.load_sleep_s) + return cls() + + def to(self, device: str) -> "_FakePipe": + if device == "cpu" and self.device == "cuda": + _Alloc.bytes -= self.size + elif device == "cuda" and self.device == "cpu": + _Alloc.bytes += self.size + self.device = device + return self + + +class _Unet(_FakePipe): + size = 3 * _GiB + + +class _Vae(_FakePipe): + size = 1 * _GiB + + +@pytest.fixture(autouse=True) +def _fake_gpu(monkeypatch): + _Alloc.bytes = 0 + monkeypatch.setattr(Executor, "_vram_allocated", staticmethod(lambda: _Alloc.bytes)) + monkeypatch.setattr(residency_mod, "get_available_ram_gb", lambda: 64.0) + + +def _spec(name: str, cls: type, models: dict, vram_gb: float | None = None) -> EndpointSpec: + return EndpointSpec( + name=name, method=cls.run, kind="inference", payload_type=_In, + output_mode="single", cls=cls, attr_name="run", models=models, + resources=Resources(vram_gb=vram_gb) if vram_gb else Resources(), + ) + + +def _executor(specs, tmp_path: Path, budget_gb: int, sent: list) -> Executor: + async def _send(msg: pb.WorkerMessage) -> None: + sent.append(msg) + + store = ModelStore(_send, cache_dir=tmp_path, vram_budget_bytes=budget_gb * _GiB) + + async def _fake_ensure_local(ref, snapshot=None, *, binding=None) -> Path: + store.residency.track_disk(ref, tmp_path) + return tmp_path + + store.ensure_local = _fake_ensure_local # type: ignore[method-assign] + return Executor(specs, _send, store=store) + + +def _events(sent, state) -> list: + return [m.model_event.ref for m in sent + if m.WhichOneof("msg") == "model_event" and m.model_event.state == state] + + +# --------------------------------------------------------------------------- # +# #369: per-ref measured vram_bytes; residency owns the pipeline objects +# --------------------------------------------------------------------------- # + + +def test_two_model_setup_books_per_ref_measured_bytes(tmp_path: Path) -> None: + class Endpoint: + def setup(self, unet: _Unet, vae: _Vae) -> None: + self.unet, self.vae = unet, vae + + def run(self, ctx, payload: _In): # pragma: no cover + return payload + + spec = _spec("ep", Endpoint, {"unet": HF("acme/unet"), "vae": HF("acme/vae")}) + sent: list = [] + + async def _run() -> None: + ex = _executor([spec], tmp_path, budget_gb=10, sent=sent) + inst = await ex.ensure_setup(spec) + res = ex.store.residency + assert res.vram_bytes("acme/unet") == 3 * _GiB + assert res.vram_bytes("acme/vae") == 1 * _GiB + # Residency owns the very objects the instance uses. + assert res.obj("acme/unet") is inst.unet + assert res.obj("acme/vae") is inst.vae + + asyncio.run(_run()) + vram_events = [(m.model_event.ref, m.model_event.vram_bytes) for m in sent + if m.WhichOneof("msg") == "model_event" + and m.model_event.state == pb.MODEL_STATE_IN_VRAM] + assert ("acme/unet", 3 * _GiB) in vram_events + assert ("acme/vae", 1 * _GiB) in vram_events + + +def test_tenant_loaded_ref_gets_residual_delta(tmp_path: Path) -> None: + class Endpoint: + def setup(self, model: str) -> None: + _Alloc.bytes += 2 * _GiB # tenant loads weights itself + + def run(self, ctx, payload: _In): # pragma: no cover + return payload + + spec = _spec("ep", Endpoint, {"model": HF("acme/opaque")}) + + async def _run() -> None: + ex = _executor([spec], tmp_path, budget_gb=10, sent=[]) + await ex.ensure_setup(spec) + assert ex.store.residency.vram_bytes("acme/opaque") == 2 * _GiB + assert ex.store.residency.movable("acme/opaque") is False + + asyncio.run(_run()) + + +def test_concurrent_cold_setups_do_not_cross_contaminate(tmp_path: Path, monkeypatch) -> None: + monkeypatch.setattr(_FakePipe, "load_sleep_s", 0.05) + + class A: + def setup(self, m: _Unet) -> None: + self.m = m + + def run(self, ctx, payload: _In): # pragma: no cover + return payload + + class B: + def setup(self, m: _Vae) -> None: + self.m = m + + def run(self, ctx, payload: _In): # pragma: no cover + return payload + + spec_a = _spec("a", A, {"m": HF("acme/a")}) + spec_b = _spec("b", B, {"m": HF("acme/b")}) + + async def _run() -> None: + ex = _executor([spec_a, spec_b], tmp_path, budget_gb=20, sent=[]) + await asyncio.gather(ex.ensure_setup(spec_a), ex.ensure_setup(spec_b)) + # Interleaved loads would double-count each other's allocations. + assert ex.store.residency.vram_bytes("acme/a") == 3 * _GiB + assert ex.store.residency.vram_bytes("acme/b") == 1 * _GiB + + asyncio.run(_run()) + + +# --------------------------------------------------------------------------- # +# #371: UNLOAD demotes to warm RAM; next setup promotes; make_room swaps LRU +# --------------------------------------------------------------------------- # + + +def test_unload_demotes_to_ram_and_next_setup_promotes(tmp_path: Path) -> None: + class Endpoint: + def setup(self, m: _Unet) -> None: + self.m = m + + def run(self, ctx, payload: _In): # pragma: no cover + return payload + + spec = _spec("ep", Endpoint, {"m": HF("acme/unet")}) + sent: list = [] + + async def _run() -> None: + ex = _executor([spec], tmp_path, budget_gb=10, sent=sent) + inst = await ex.ensure_setup(spec) + await ex.handle_model_op(pb.ModelOp(op=pb.MODEL_OP_KIND_UNLOAD, ref="acme/unet")) + res = ex.store.residency + assert res.tier("acme/unet") is Tier.RAM + assert inst.m.device == "cpu" # actually moved, not just booked + rec = ex._classes[spec.instance_key] + assert rec.ready and rec.instance is inst # no teardown + + again = await ex.ensure_setup(spec) # RunJob path: promote, no reload + assert again is inst + assert res.tier("acme/unet") is Tier.VRAM + assert inst.m.device == "cuda" + + asyncio.run(_run()) + assert _events(sent, pb.MODEL_STATE_IN_RAM) == ["acme/unet"] + assert _events(sent, pb.MODEL_STATE_IN_VRAM) == ["acme/unet", "acme/unet"] + + +def test_unload_of_tenant_loaded_ref_tears_down_record(tmp_path: Path) -> None: + class Endpoint: + def setup(self, model: str) -> None: + _Alloc.bytes += 2 * _GiB + + def shutdown(self) -> None: + _Alloc.bytes -= 2 * _GiB + + def run(self, ctx, payload: _In): # pragma: no cover + return payload + + spec = _spec("ep", Endpoint, {"model": HF("acme/opaque")}) + sent: list = [] + + async def _run() -> None: + ex = _executor([spec], tmp_path, budget_gb=10, sent=sent) + await ex.ensure_setup(spec) + await ex.handle_model_op(pb.ModelOp(op=pb.MODEL_OP_KIND_UNLOAD, ref="acme/opaque")) + assert ex.store.residency.tier("acme/opaque") is Tier.DISK + assert ex._classes[spec.instance_key].ready is False + assert _Alloc.bytes == 0 # shutdown() actually freed the memory + + asyncio.run(_run()) + assert _events(sent, pb.MODEL_STATE_ON_DISK)[-1] == "acme/opaque" + + +def test_alternating_endpoints_swap_via_demote_promote(tmp_path: Path) -> None: + """Budget fits one 3GiB pipeline (+2GiB make_room margin): alternating + setups must demote/promote, never tear down or degrade to offload.""" + + class A: + def setup(self, m: _Unet) -> None: + self.m = m + + def run(self, ctx, payload: _In): # pragma: no cover + return payload + + class B: + def setup(self, m: _Unet) -> None: + self.m = m + + def run(self, ctx, payload: _In): # pragma: no cover + return payload + + spec_a = _spec("a", A, {"m": HF("acme/a")}, vram_gb=3.0) + spec_b = _spec("b", B, {"m": HF("acme/b")}, vram_gb=3.0) + sent: list = [] + + async def _run() -> None: + ex = _executor([spec_a, spec_b], tmp_path, budget_gb=6, sent=sent) + res = ex.store.residency + + inst_a = await ex.ensure_setup(spec_a) + assert res.tier("acme/a") is Tier.VRAM + + inst_b = await ex.ensure_setup(spec_b) # make_room demotes idle A first + assert res.tier("acme/b") is Tier.VRAM + assert res.tier("acme/a") is Tier.RAM + assert inst_a.m.device == "cpu" + + again_a = await ex.ensure_setup(spec_a) # promote A back; B demoted + assert again_a is inst_a # instance survived throughout + assert res.tier("acme/a") is Tier.VRAM + assert res.tier("acme/b") is Tier.RAM + assert inst_b.m.device == "cpu" + for rec in ex._classes.values(): + assert rec.ready + + asyncio.run(_run()) + assert _events(sent, pb.MODEL_STATE_IN_RAM) == ["acme/a", "acme/b"] + assert _events(sent, pb.MODEL_STATE_IN_VRAM) == ["acme/a", "acme/b", "acme/a"] diff --git a/tests/test_residency.py b/tests/test_residency.py index 6c1b0de..25dec86 100644 --- a/tests/test_residency.py +++ b/tests/test_residency.py @@ -94,13 +94,19 @@ def test_track_disk_is_idempotent_no_event_spam(tmp_path: Path) -> None: assert [s for _, s, _ in events] == ["on_disk"] -def test_demote_without_disk_path_or_obj_evicts(tmp_path: Path) -> None: +def test_demote_refuses_unmovable_entries(monkeypatch) -> None: + """demote() never books a transition it can't perform: entries without a + movable object (tenant-loaded weights) and RAM-tight hosts are refused — + freeing them is the owner's (executor teardown) job.""" events: list = [] res = _res(events) - res.track_vram("no-obj", None, vram_bytes=1 * _GiB) # nothing to keep warm - assert res.demote("no-obj") is True - assert res.tier("no-obj") is None - assert events[-1] == ("no-obj", "evicted", 0) + res.track_vram("no-obj", None, vram_bytes=1 * _GiB) + assert res.demote("no-obj") is False + assert res.tier("no-obj") is Tier.VRAM # still the truth + + res.track_vram("movable", _Pipe(), vram_bytes=1 * _GiB) + monkeypatch.setattr(residency_mod, "get_available_ram_gb", lambda: 1.0) + assert res.demote("movable") is False # RAM floor: no warm tier # --------------------------------------------------------------------------- # From ac4c9e351f9ae2aa9095e9ecccec3fc9e1a3de66 Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Sat, 4 Jul 2026 02:40:14 -0600 Subject: [PATCH 2/2] tracker: complete #369-#371 (residency unification, disk GC, VRAM juggling) --- agents/completed.md | 63 +++++++++++++++++++++++++++++++++++++++++++++ agents/progress.md | 57 ---------------------------------------- 2 files changed, 63 insertions(+), 57 deletions(-) diff --git a/agents/completed.md b/agents/completed.md index 1ddee7a..a0094d4 100644 --- a/agents/completed.md +++ b/agents/completed.md @@ -8,6 +8,69 @@ --- +# #369: residency registry must own the executor's pipelines — honest vram_bytes, no load races + +**Completed:** yes +**Status:** DONE (2026-07-04, Claude) — Residency now owns the executor's pipelines: per-slot allocator-delta measurement in `_injection_kwargs` (worker-built pipelines registered per ref with obj+measured bytes; tenant-loaded refs split the residual delta), loads serialized under a dedicated `Executor._load_lock` (same exclusion the GPU-semaphore option would give, without coupling loads to job slots), `demote`/`promote` only perform transitions they can actually execute (movable obj + RAM headroom) and `_vacate_record` books every held ref back to disk on teardown — registry and instance state cannot diverge. Free-VRAM probes sum across ALL CUDA devices (residency probe + lifecycle StateDelta). Also fixed: residency events from worker threads were dropped when no loop was captured (`ModelStore.bind_loop`). Tests: tests/test_executor_residency.py. + +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: +- [x] 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. +- [x] 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. +- [x] Serialize model loads under a lock shared with the GPU semaphore path; take the allocator-delta measurement inside it. +- [x] 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. +- [x] (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:** yes +**Status:** DONE (2026-07-04, Claude) — pre-download headroom gate in `ModelStore.ensure_local` (snapshot sizes vs disk-free probe → GC first → `ModelEvent{FAILED, insufficient_disk}` + typed `InsufficientDiskError`, which fails the job RETRYABLE without self-disabling the function), LRU disk GC over Residency DISK-tier refs honoring keep + in-use pins + 300s grace, keep-pressure escape hatch still emits EVICTED; deletion units in models/disk_gc.py (CAS snapshot dir + hardlink-count orphan-blob sweep; HF `models--*` repo dir), boot-time rescan wired into lifecycle.startup(). Deliberate: rescan reads the persisted `ref-index.json` instead of walking bare CAS digests (digests don't map back to wire refs); hf/civitai downloads have no up-front size and skip the pre-gate. Tests: tests/test_disk_gc.py (real files, budgeted probe). + +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: +- [x] 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). +- [x] Disk GC: LRU over Residency DISK-tier entries, never evicting `keep` refs or refs of in-flight jobs; delete bytes + `evict()` + emit EVICTED. +- [x] 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). +- [x] 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. +- [x] 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:** yes +**Status:** DONE (2026-07-04, Claude) — `ensure_setup` runs make_room before loading (estimate: per-ref `vram_hint` from a prior load, else declared `resources.vram_gb`); movable LRU victims demote to the warm RAM tier, non-movable victims vacate via record teardown; hub UNLOAD demotes instead of destroying (instance stays ready) and the next RunJob/LOAD promotes RAM→VRAM under the load lock; the offload ladder remains the fallback when make_room can't reach the requirement. Simulated-VRAM e2e (vram_budget_bytes): alternating endpoints swap via demote/promote with IN_VRAM⇄IN_RAM ModelEvents and zero teardowns (tests/test_executor_residency.py). + +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: +- [x] `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. +- [x] 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. +- [x] `_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. +- [x] Promotion path: RunJob for a RAM-tier instance re-promotes before dispatch to the handler (executor checks tier in `ensure_setup`). +- [x] 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. +- [x] 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. + +--- + # #375: `yield ProducedFlavor` conversion contract is gone — from-scratch example (and training-endpoints quant endpoints) silently publish nothing **Completed:** yes diff --git a/agents/progress.md b/agents/progress.md index 73fbc76..b4ac069 100644 --- a/agents/progress.md +++ b/agents/progress.md @@ -320,62 +320,5 @@ NOTE: implementation lives in gen-orchestrator (separate repo); tracked here for - [ ] 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. ---