From c5ed1521a31e044e2a640545ce1c4884e72db929 Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Tue, 14 Jul 2026 14:03:59 -0600 Subject: [PATCH 1/3] feat: make worker model residency declarative --- CHANGELOG.md | 13 + proto/CONTRACT.md | 91 ++++--- proto/worker_scheduler.proto | 43 ++- pyproject.toml | 2 +- src/gen_worker/discovery/discover.py | 2 +- src/gen_worker/executor.py | 207 +++++--------- src/gen_worker/lifecycle.py | 108 ++++++-- src/gen_worker/models/download.py | 2 +- src/gen_worker/models/provision.py | 2 +- src/gen_worker/models/residency.py | 2 +- src/gen_worker/pb/worker_scheduler_pb2.py | 128 +++++---- src/gen_worker/pb/worker_scheduler_pb2.pyi | 45 ++- tests/e2e_endpoints.py | 4 +- tests/test_declarative_residency.py | 256 ++++++++++++++++++ ...est_dynamic_slot_materialization_pgw532.py | 43 +-- tests/test_executor_model_ops.py | 78 ------ tests/test_executor_prefetch_binding.py | 23 +- tests/test_executor_residency.py | 70 ----- tests/test_model_op_prefetch_no_cascade.py | 252 ----------------- tests/test_snapshot_remint.py | 100 +++++++ tests/test_startup_hub_snapshot_wait.py | 5 +- tests/test_worker_grpc_e2e.py | 89 +++--- uv.lock | 2 +- 23 files changed, 784 insertions(+), 783 deletions(-) create mode 100644 tests/test_declarative_residency.py delete mode 100644 tests/test_executor_model_ops.py delete mode 100644 tests/test_model_op_prefetch_no_cascade.py create mode 100644 tests/test_snapshot_remint.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 7b487049..32913d70 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,5 +1,18 @@ # Changelog +## 0.26.0 (2026-07-14) + +- **Model residency is declarative.** Protocol v3 replaces ordinary + download/load/unload commands with a full-replace per-worker desired disk + set and ordered hot runnable instances. Workers report the accepted + generation separately from actual residency events. +- **Tenant work preempts background reconciliation.** A `RunJob` cancels + unrelated desired-state work before request setup, then resumes the current + desired generation when the executor becomes idle. +- Hot dynamic-slot instances reuse the request binding and setup path, so the + function plus complete slot-to-immutable-ref map identifies exactly one + runnable instance without a second loader. + ## 0.25.2 (2026-07-14) - **Mixed CPU/GPU releases probe the device owned by their concrete image.** diff --git a/proto/CONTRACT.md b/proto/CONTRACT.md index 582a9922..981e36fe 100644 --- a/proto/CONTRACT.md +++ b/proto/CONTRACT.md @@ -3,7 +3,7 @@ Proto: `worker_scheduler.proto`, package `cozy.scheduler`, service `WorkerScheduler`, one RPC: `Connect(stream WorkerMessage) returns (stream SchedulerMessage)`. This REPLACES the old `scheduler.v1` proto entirely. No compat, no negotiation: -`ProtocolVersion.PROTOCOL_VERSION_CURRENT = 2` is the only accepted version. +`ProtocolVersion.PROTOCOL_VERSION_CURRENT = 3` is the only accepted version. Roles: **W** = gen-worker (Python), **O** = orchestrator (tensorhub Go). Every field below names its producer and consumer; anything without both was deleted. @@ -83,8 +83,10 @@ Immediately after `HelloAck`, W flushes its buffered unsent `JobResult`s (§4 results are never dropped). **HelloAck re-send.** O re-sends `HelloAck` mid-connection whenever release -config changes (e.g. the `keep` set). Full-replace semantics: W overwrites -`file_base_url` and its keep set wholesale. +config or desired model residency changes. Full-replace semantics: W +overwrites `file_base_url`, resolutions, and desired residency wholesale. +The same residency generation MAY be resent with refreshed snapshot URLs; +only a generation lower than the latest observed generation is stale. **Send-queue policy (W).** One bounded outbound queue. `JobResult` is NEVER dropped: results persist in the queue across reconnects until written to a @@ -200,7 +202,20 @@ Reply to Hello; re-sent on config change (full replace). |---|---|---|---| | `protocol_version` | O constant | W version gate | symmetric validation | | `file_base_url` | O config (`FileAPIBaseURL`) | W upload stack / blob-ref PUT flow (§8) | tensorhub HTTP base for capability-token calls | -| `keep` | O release prefetch plan | W disk-eviction policy | refs the worker must not evict from disk; NOT a download order (see §7) | +| `keep` | legacy v2 release config | none in v3 | retained field number only; v3 W ignores it | +| `resolutions` | O precision resolver | W endpoint bindings | full-replace declared-ref to worker-specific ref/cast picks | +| `desired_residency` | O scheduler/controller | W model reconciler | full-replace per-worker disk and hot-instance goal (§7) | + +### DesiredResidency / DesiredInstance (embedded in HelloAck) + +| field | producer | consumer | semantics | +|---|---|---|---| +| `generation` | O desired-state store | W observation fence | monotonic desired revision; observation does not imply convergence | +| `disk_refs` | O locality scheduler | W downloader + disk GC | ordered refs to materialize and protect from eviction while headroom allows | +| `hot` | O locality scheduler | W endpoint setup | ordered runnable instances; each names a function and its complete `ModelBinding` slot map | +| `snapshots` | O resolver | W downloader | snapshots for every disk/hot ref that needs one; same semantics as `RunJob.snapshots` | +| `hot[].function_name` | O release manifest | W registry lookup | endpoint function whose persistent instance should be warm | +| `hot[].models` | O binding resolver | W existing RunJob binding path | complete slot to immutable-ref map; incomplete or mismatched maps are rejected | ### StateDelta (W → O) Full-replace snapshot of ALL dynamic worker state. Edge-triggered: sent @@ -215,6 +230,7 @@ to avoid chatter). Never periodic. O overwrites its copy wholesale. | `loading_functions` | W registry | O supply (counts 0 capacity) + display | present but models still materializing; disjoint from available | | `free_vram_bytes` | W CUDA probe | O placement (free-VRAM ladder) | measured free VRAM | | `finalizing_jobs` | W executor (gw#516) | O drain/retire gating + worker status display | jobs past the decode→finalize handoff: GPU slot released, encode/upload tail running, `JobResult` unshipped. GPU-idle alone is NOT work-idle | +| `observed_residency_generation` | W desired-state receiver | O controller status | latest non-stale generation accepted; not a convergence claim — `ModelEvent`/`Hello.models` report actual state | ### RunJob (O → W) @@ -288,7 +304,7 @@ pick with `{"steps":4,"guidance":0}` — the merged `ctx.slots["pipeline"].defaults` reads `steps=4, guidance=0`, scheduler/ max_guidance untouched (the lora left them null). -### Snapshot / SnapshotFile (embedded in RunJob.snapshots and ModelOp) +### Snapshot / SnapshotFile (embedded in RunJob, DesiredResidency, and ModelOp) | field | producer | consumer | semantics | |---|---|---|---| @@ -296,7 +312,7 @@ max_guidance untouched (the lora left them null). | `files[].path` | O | W file placement | repo-relative path | | `files[].size_bytes` | O | W disk-headroom check + progress totals | file size | | `files[].blake3` | O | W post-download verification (digest-poisoning guard) | content hash | -| `files[].url` | O presigner | W downloader | presigned GET; expiry ⇒ `ModelEvent{FAILED, error:"url_expired"}` and O re-sends the op with fresh URLs | +| `files[].url` | O presigner | W downloader | presigned GET; expiry ⇒ `ModelEvent{FAILED, error:"url_expired"}` and O re-sends the same desired generation or job with fresh URLs | Standalone CLIENTS (cozy CLI, `gen-worker prefetch`) pull the SAME shape over HTTP instead: `GET /api/v1/repos/:tenant/:name/resolve?tag=&flavor=` (#560 — @@ -414,19 +430,18 @@ and the cancel is a no-op. Cancel of an unknown pair is silently ignored. There is no queued-only/item-level cancel. ### ModelOp (O → W) -Producers: prefetch planner (after Hello, for `keep` refs missing from -`Hello.models`), cold-locality download affinity, VRAM packer (LOAD ahead of -predicted demand, UNLOAD for headroom). +Producer: compile-cache adoption only. Model download/load/unload lifecycle is +declarative through `HelloAck.desired_residency`; v3 O MUST NOT send the legacy +DOWNLOAD/LOAD/UNLOAD enum values. | field | producer | consumer | semantics | |---|---|---|---| -| `op` | O | W model layer | DOWNLOAD = to disk only; LOAD = promote to VRAM (fetching first if needed); UNLOAD = demote out of VRAM (stays on disk/RAM per cache policy); ADOPT_COMPILE_CACHE = hot-adopt a torch.compile cache (below) | +| `op` | O | W model layer | ADOPT_COMPILE_CACHE only | | `ref` | O | W | canonical ref | -| `snapshot` | O resolver | W downloader | required for tensorhub-CAS refs not already local; unset for hf/civitai; always set on ADOPT_COMPILE_CACHE | +| `snapshot` | O resolver | W downloader | required for ADOPT_COMPILE_CACHE | Every ModelOp is answered by ≥1 `ModelEvent` for the ref (success path emits -the new state; failure path emits FAILED). W MAY refuse UNLOAD for a model -serving in-flight jobs: `ModelEvent{FAILED, error:"model_in_use"}`. +ADOPTED; failure path emits FAILED). **ADOPT_COMPILE_CACHE** (hot adoption, #567): `ref` is a compile-cache flavor ref — `_system/family-#inductor--torch`. W downloads the @@ -450,17 +465,17 @@ defensively. Workers predating this kind ignore it silently (unknown enum, no ModelEvent); O treats the absence of a reply as not-adopted, never an error. ### ModelEvent (W → O) -The single model-residency channel. Emitted for ModelOp outcomes AND -worker-initiated transitions (LRU eviction, demotion, on-demand load during a -job). Replaces load/unload results, model-ready signals, and the JSON -download-event fabric. +The single model-residency channel. Emitted for desired-state outcomes, +compile-cache adoption, AND worker-initiated transitions (LRU eviction, +demotion, on-demand load during a job). Replaces load/unload results, +model-ready signals, and the JSON download-event fabric. | field | producer | consumer | semantics | |---|---|---|---| | `ref` | W | O residency index | | | `state` | W model cache | O: DOWNLOADING/ON_DISK feed pending-download affinity (parked jobs dispatch on ON_DISK); IN_RAM/IN_VRAM/ON_DISK update the placement tier index; EVICTED removes the ref; FAILED fails the pending op | current state machine position | | `vram_bytes` | W measured (allocator delta across load) | O model-size cache → VRAM packing | set with IN_VRAM | -| `error` | W | O op-failure handling (e.g. OOM ⇒ pick LRU vram model, send UNLOAD then retry LOAD; url_expired ⇒ re-mint snapshot and re-send) + triage log | set with FAILED | +| `error` | W | O reconcile handling (e.g. OOM ⇒ revise desired hot set; url_expired ⇒ re-mint snapshots and re-send the same generation) + triage log | set with FAILED | | `bytes_done`/`bytes_total` | W downloader (emit ≤1 per 5s per ref) | O boot/capacity progress display | set with DOWNLOADING | | `duration_ms` | W adopt handler | O adoption bookkeeping + fleet-adoption-latency metric | set with ADOPTED: wall time of download+seed+re-wrap+warmup | | `cache_hits`/`cache_misses` | W adopt handler (inductor FX-graph cache counter delta across the warmup) | O adoption observability: hits ≥ 1 on every ADOPTED; misses > 0 = partial shape coverage | set with ADOPTED (gw#391) | @@ -483,8 +498,8 @@ when a function's model download / pipeline setup fails terminally (`reason: setup_failed` — the function is dropped from BOTH `available_functions` and `loading_functions` instead of sitting in loading forever under a READY phase). One message per function, -deduplicated. A later hub-directed `ModelOp{LOAD}` retries the setup; on -success the function reappears in `StateDelta.available_functions`. +deduplicated. A later desired-state reconciliation or `RunJob` retries setup; +on success the function reappears in `StateDelta.available_functions`. | field | producer | consumer | semantics | |---|---|---|---| @@ -552,14 +567,21 @@ The stream is HTTP/2: reliable, ordered per direction. Consequences relied on: --- -## 7. Model prefetch / keep +## 7. Declarative model residency + +`HelloAck.desired_residency` is the full per-worker goal. W accepts each +non-stale generation without blocking the receive loop, replaces its disk +keep set from `disk_refs`, cancels obsolete background reconciliation, and +works the ordered disk refs then ordered hot instances. Reconciliation starts +only while the executor is job-idle; an active `RunJob` and its exact bindings +always take precedence over pending background work. Hot instances reuse the +same binding derivation and setup path as `RunJob`; there is no second loader. -`HelloAck.keep` is a retention set, not a download order. After HelloAck, O -compares `keep` against its residency view (`Hello.models` + events) and -issues `ModelOp{DOWNLOAD}` (with snapshots where needed) for missing refs — -the ONE prefetch mechanism. 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. +`observed_residency_generation` means only that W accepted the desired state. +Actual progress and failures are `ModelEvent`s, with `Hello.models` as the +reconnect baseline. W MAY evict desired refs under real disk/VRAM pressure and +reports that honestly; O then revises or re-sends desired state. Same-generation +re-sends refresh expired snapshots and retry failed work. **Compile-cache snapshots (#569).** When a release's endpoint declares `compile=Compile(family=...)` and boot-attach is enabled (opt-in; default OFF @@ -569,8 +591,7 @@ the resolved `_system/family-#inductor--torch` snapshot to recognizes the key by the `_system/family-#inductor-` prefix for the declared family, downloads it like any snapshot, and seeds it before pipeline load; verification failure ⇒ eager, never an error. The same ref may arrive -via `ModelOp{DOWNLOAD}` (pre-positioning) or `ModelOp{ADOPT_COMPILE_CACHE}` -(hot adoption, §4). +via `ModelOp{ADOPT_COMPILE_CACHE}` (hot adoption, §4). --- @@ -658,9 +679,9 @@ ADOPT_COMPILE_CACHE only — `adopt_failed:` with reason ∈ {`bad_ref`, per-function capability metadata belongs to deploy-time manifests. - `WorkerKVPrefixCache` — producer never existed. - `EndpointConfig` (repo allowlists, per-function model keyspaces, disabled - functions, payload-ref availability) — O gates dispatch itself; W learns - its retention set from `HelloAck.keep` and gets snapshots on - `ModelOp`/`RunJob`. Nothing else was consumed by the worker it will keep. + functions, payload-ref availability) — O gates dispatch itself; W gets its + model goal from `HelloAck.desired_residency` and request bindings from + `RunJob`. Nothing else was consumed by the worker it will keep. - `InterruptJobCommand.item_ids`/`cancel_queued_only` — item batching is dead; cancel is whole-request. - `JobExecutionObservation` fields: identity/context fields O already knows @@ -676,6 +697,8 @@ ADOPT_COMPILE_CACHE only — `adopt_failed:` with reason ∈ {`bad_ref`, and re-added deliberately as `ModelBinding.loras` once BYOM (gw#393 / th#585) gave it a producer. - `WorkerResources.gpu_is_busy` — derivable from O's own assignment map. -- `DownloadModelCommand.priority` — never read; retention lives in `keep`. +- `DownloadModelCommand.priority` — never read; ordered desired-state lists + express controller preference without a second priority system. - `WorkerDrainResult` — stream close is the drain ack. -- Reserved-tag graveyard — renumbered from 1, no reservations, one version. +- Reserved-tag graveyard — messages were renumbered from 1; only removed v3 + `ModelOpKind` values 1–3 stay reserved to prevent accidental wire reuse. diff --git a/proto/worker_scheduler.proto b/proto/worker_scheduler.proto index 26ee666b..4cf9c322 100644 --- a/proto/worker_scheduler.proto +++ b/proto/worker_scheduler.proto @@ -28,7 +28,7 @@ option go_package = "github.com/cozy-creator/tensorhub/internal/orchestrator/grp // FAILED_PRECONDITION. No minor versions, no feature gates, no negotiation. enum ProtocolVersion { PROTOCOL_VERSION_UNSPECIFIED = 0; - PROTOCOL_VERSION_CURRENT = 2; + PROTOCOL_VERSION_CURRENT = 3; } service WorkerScheduler { @@ -118,13 +118,31 @@ message InFlightJob { message HelloAck { ProtocolVersion protocol_version = 1; // orchestrator's version, for symmetric validation string file_base_url = 2; // tensorhub base URL for capability-token HTTP calls (uploads, blob refs) - repeated string keep = 3; // refs the worker's disk eviction must protect + repeated string keep = 3; // legacy v2 field; ignored by protocol v3 workers // th#697 precision ladder: per-model flavor resolutions for THIS worker's // card. The worker applies each entry to bindings whose wire ref equals // `ref` (a bare declared ref): download/load `resolved_ref` and apply the // `cast` transform at load. Full-replace with the ack. Additive: old // workers ignore it and load their declared bindings. repeated ModelResolution resolutions = 4; + // Full-replace desired model residency for this worker (protocol v3). + DesiredResidency desired_residency = 5; +} + +// Declarative worker model state. Lists are ordered by orchestrator priority; +// snapshots cover every disk/hot ref the worker may need. +message DesiredResidency { + uint64 generation = 1; + repeated string disk_refs = 2; + repeated DesiredInstance hot = 3; + map snapshots = 4; +} + +// One complete runnable endpoint instance. `models` is the same complete +// slot-to-ref binding set carried by RunJob. +message DesiredInstance { + string function_name = 1; + repeated ModelBinding models = 2; } // One precision-ladder resolution (th#697): declared bare ref -> the concrete @@ -148,6 +166,9 @@ message StateDelta { // Drain/retire must wait for these — GPU-idle alone is NOT work-idle. // Additive: old workers never set it; old orchestrators ignore it. int32 finalizing_jobs = 5; + // Latest desired-residency generation accepted by this worker. This is + // observation, not convergence; ModelEvent reports best-effort outcomes. + uint64 observed_residency_generation = 6; } enum WorkerPhase { @@ -241,7 +262,8 @@ message LoraOverlay { // A pinned tensorhub repo snapshot with presigned download URLs. The worker // never calls tensorhub for ref resolution; the orchestrator is the only -// resolver and ships snapshots on ModelOp / RunJob. +// resolver and ships snapshots on DesiredResidency / RunJob (or compile-cache +// adoption ModelOp). message Snapshot { string digest = 1; // snapshot digest the download is pinned to repeated SnapshotFile files = 2; @@ -352,18 +374,15 @@ message CancelJob { message ModelOp { ModelOpKind op = 1; string ref = 2; - // Presigned snapshot for tensorhub-CAS refs; unset for hf/civitai refs the - // worker fetches from the upstream provider directly. Set on DOWNLOAD (and - // on LOAD when the ref may not be on disk yet), and always on - // ADOPT_COMPILE_CACHE. + // Presigned compile-cache snapshot. Ordinary model residency is declared in + // HelloAck.desired_residency; ModelOp is only the compile-cache escape hatch. Snapshot snapshot = 3; } enum ModelOpKind { MODEL_OP_KIND_UNSPECIFIED = 0; - MODEL_OP_KIND_DOWNLOAD = 1; // fetch to disk; do NOT promote to VRAM - MODEL_OP_KIND_LOAD = 2; // promote to VRAM (downloading first if needed) - MODEL_OP_KIND_UNLOAD = 3; // demote out of VRAM + reserved 1, 2, 3; + reserved "MODEL_OP_KIND_DOWNLOAD", "MODEL_OP_KIND_LOAD", "MODEL_OP_KIND_UNLOAD"; // Hot-adopt a torch.compile cache artifact: ref is a compile-cache flavor // ref (`_system/family-#inductor--torch`). The worker // downloads+verifies+seeds the cache, re-wraps the already-resident modules @@ -376,8 +395,8 @@ enum ModelOpKind { } // Worker -> orchestrator model residency transition. The single channel for -// download progress, load/unload outcomes, VRAM footprints, and evictions — -// whether triggered by a ModelOp or by the worker's own cache policy. +// transfer/setup outcomes, VRAM footprints, and evictions — whether triggered +// by desired-state reconciliation, compile-cache adoption, or local policy. message ModelEvent { string ref = 1; ModelState state = 2; diff --git a/pyproject.toml b/pyproject.toml index 5bc07365..c68e5f6c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "gen-worker" -version = "0.25.2" +version = "0.26.0" description = "A library used to build custom functions in Cozy Creator's serverless function platform." readme = "README.md" license = "MIT" diff --git a/src/gen_worker/discovery/discover.py b/src/gen_worker/discovery/discover.py index 0673702c..281a9d97 100644 --- a/src/gen_worker/discovery/discover.py +++ b/src/gen_worker/discovery/discover.py @@ -291,7 +291,7 @@ def _binding_to_manifest(binding: Binding, param_name: str = "") -> Dict[str, An if binding.flavor: out["flavor"] = binding.flavor if binding.components: - # pgw#505: the hub's ModelOp DOWNLOAD scoping (platform-side, + # pgw#505: the hub's desired-snapshot scoping (platform-side, # not yet built) reads this to resolve only the named pipeline # component subfolders instead of the whole repo. out["components"] = list(binding.components) diff --git a/src/gen_worker/executor.py b/src/gen_worker/executor.py index 8a26884e..8c1b6b34 100644 --- a/src/gen_worker/executor.py +++ b/src/gen_worker/executor.py @@ -1,6 +1,7 @@ """Job execution: intake, GPU semaphore, deadline + cancellation watchdog, sync-on-thread / async-on-loop, JobProgress deltas, result send, and the -worker-side model seam (ensure-local + setup injection + ModelOp handling). +worker-side model seam (ensure-local, setup injection, declarative residency, +and compile-cache adoption). One dispatch path for every endpoint kind. Everything runs on the single asyncio loop; sync tenant code runs in threads via asyncio.to_thread. @@ -217,7 +218,7 @@ def _resolve_slots_kwargs(spec: EndpointSpec, run: "pb.RunJob") -> Dict[str, Any def _hub_binding_for_wire_ref(ref: str) -> ModelRef: """A tensorhub-source binding for a hub-named wire ref (pgw#532). - ``RunJob.models`` / ModelOp refs name hub-CAS repos in the canonical + ``RunJob.models`` / desired-instance refs name hub-CAS repos in the canonical ``owner/repo[:tag][#flavor]`` grammar; this mints the binding the executor materializes them through (``ensure_local`` then follows the tensorhub lane: orchestrator snapshots or the th#763 missing_snapshot @@ -343,8 +344,8 @@ def _snapshot_to_resolved(snap: pb.Snapshot) -> "WorkerResolvedRepo": ) -def _model_op_error_vocab(exc: BaseException) -> str: - """Contract §9 ModelEvent.error vocabulary for LOAD/UNLOAD failures.""" +def _model_failure_vocab(exc: BaseException) -> str: + """Contract §9 ModelEvent.error vocabulary for residency failures.""" if is_cuda_oom(exc): return "oom" if isinstance(exc, MissingSnapshotError): @@ -435,11 +436,10 @@ def __init__( # (gw#408): a cached snapshot is re-verified on first use per process # so pod-churn corruption can never be trusted forever. self._verified: set[str] = set() - # Last digest-carrying snapshot seen per ref (gw#465): ModelOp{LOAD} - # and companion-slot setups arrive snapshot-less; without memory of - # the hub's earlier DOWNLOAD snapshot they cannot materialize - # tensorhub refs. Stale URLs self-heal: they fail url_expired and the - # hub re-mints. + # Last digest-carrying snapshot seen per ref (gw#465): companion-slot + # setups may arrive snapshot-less; without memory of the hub's desired + # state / RunJob snapshot they cannot materialize tensorhub refs. Stale + # URLs self-heal: they fail url_expired and the hub re-mints. self._snapshots: Dict[str, pb.Snapshot] = {} # Cold-ref waiters (th#763): ensure_local blocks here until the # hub's re-minted DOWNLOAD banks a snapshot for the ref. @@ -633,16 +633,16 @@ def _lock(self, ref: str) -> asyncio.Lock: def register_binding(self, ref: str, binding: Any) -> None: """Endpoint-spec binding for ``ref`` — supplies files/provider on - download paths that only carry the bare ref (ModelOp, startup - prefetch), so ``files=`` selections apply everywhere (#377).""" + download paths that only carry the bare ref (DesiredResidency or + startup prefetch), so ``files=`` selections apply everywhere (#377).""" self._bindings.setdefault(ref, binding) async def _await_hub_snapshot(self, ref: str) -> pb.Snapshot: """Cold tensorhub ref with no orchestrator-resolved snapshot: emit - ``missing_snapshot`` (the hub re-mints a DOWNLOAD op with fresh - URLs on seeing it — connect_worker handleModelFailure) and block + ``missing_snapshot`` (the hub refreshes desired state with fresh URLs + on seeing it — connect_worker handleModelFailure) and block until that snapshot is banked (th#763). The bank site runs OUTSIDE - the per-ref lock this coroutine holds, so the re-minted DOWNLOAD's + the per-ref lock this coroutine holds, so the refreshed reconcile's ensure_local wakes us and then queues behind the lock. Raises :class:`MissingSnapshotError` when nothing arrives in ``_MISSING_SNAPSHOT_WAIT_S``.""" @@ -1194,7 +1194,7 @@ def apply_model_resolutions(self, resolutions: Dict[str, Tuple[str, str]]) -> No # gw#494: transactional application — (1) a ready instance whose # loaded refs diverged from the rebound refs is stale: vacate it # so nothing stays booked under the old resolved refs (pins, - # promotes, adapters and UNLOAD all key off the CURRENT wire + # promotions, adapters, and eviction all key off the CURRENT wire # refs; a divergent record orphans its VRAM forever). stale: List[_ClassRecord] = [] seen: set = set() @@ -1541,6 +1541,51 @@ def _effective_spec(self, spec: EndpointSpec, run: "pb.RunJob") -> EndpointSpec: return spec return dc_replace(spec, models=effective) + async def ensure_desired_instance( + self, + desired: "pb.DesiredInstance", + snapshots: Dict[str, "pb.Snapshot"], + ) -> None: + """Best-effort warm of one declarative, fully bound instance.""" + spec = self.specs.get(desired.function_name) + if spec is None: + raise ValidationError(f"unknown function {desired.function_name!r}") + if spec.cls is None: + raise ValidationError( + f"function {desired.function_name!r} has no persistent instance to warm" + ) + + pairs = [(m.slot.strip(), m.ref.strip()) for m in desired.models] + bindings = dict(pairs) + expected = set(spec.models) + if any(not slot or not ref for slot, ref in pairs): + raise ValidationError( + f"desired instance {desired.function_name!r} has an empty slot or ref" + ) + if len(bindings) != len(pairs) or set(bindings) != expected: + raise ValidationError( + f"desired instance {desired.function_name!r} must bind exactly " + f"{sorted(expected)!r}; got {sorted(bindings)!r}" + ) + + run = pb.RunJob(function_name=desired.function_name, models=desired.models) + effective = self._effective_spec(spec, run) + resolved = {slot: wire_ref(binding) for slot, binding in effective.models.items()} + if resolved != bindings: + raise ValidationError( + f"desired instance {desired.function_name!r} does not match the " + "worker's resolved bindings" + ) + try: + await self.ensure_setup(effective, snapshots) + except Exception as exc: + error = _model_failure_vocab(exc) + for ref in dict.fromkeys(bindings.values()): + await self._send(pb.WorkerMessage(model_event=pb.ModelEvent( + ref=ref, state=pb.MODEL_STATE_FAILED, error=error, + ))) + raise + def _class_record(self, spec: EndpointSpec) -> _ClassRecord: """Instance-group record for ``spec``, created on first sight for DERIVED (per-pick) specs. Never removed: records are tiny and the @@ -1585,7 +1630,7 @@ async def ensure_setup( self._mark_setup_failed(rec, exc) raise if rec.failed is not None: - # Recovery (hub-directed LOAD retry succeeded): lift the + # Recovery (desired-state retry succeeded): lift the # per-function disable; the next StateDelta re-advertises. rec.failed = None for s in rec.specs: @@ -1860,24 +1905,6 @@ async def _make_room_for(self, spec: EndpointSpec, setup_slots: List[str]) -> No break self._on_state_change() - def _missing_slot_refs( - self, spec: EndpointSpec, snapshots: Optional[Dict[str, pb.Snapshot]] - ) -> set: - """Tensorhub setup-slot refs of ``spec`` the worker cannot materialize - right now: not on disk, no snapshot in this op, none remembered - (gw#465). Non-tensorhub slots self-fetch and never block.""" - missing: set = set() - for slot in self._setup_slots(spec): - binding = spec.models[slot] - if getattr(binding, "source", "tensorhub") != "tensorhub": - continue - r = wire_ref(binding) - if (self.store.local_path(r) is None - and not (snapshots and r in snapshots) - and not self.store.has_snapshot(r)): - missing.add(r) - return missing - @staticmethod def _setup_slots(spec: EndpointSpec) -> List[str]: """Model slots loaded once at setup time. Classes without setup() @@ -2242,97 +2269,24 @@ async def shutdown_instances(self) -> None: if server is not None: await asyncio.to_thread(server.stop) - # ---- ModelOp ----------------------------------------------------------- + # ---- Compile-cache adoption ------------------------------------------- async def handle_model_op(self, op: pb.ModelOp) -> None: + """Handle the sole v3 ModelOp: hot adoption of a compile cache.""" + if op.op != pb.MODEL_OP_KIND_ADOPT_COMPILE_CACHE: + return self.store.bind_loop() ref = op.ref snap = op.snapshot if op.HasField("snapshot") else None try: - if op.op == pb.MODEL_OP_KIND_DOWNLOAD: - await self.store.ensure_local(ref, snap) - elif op.op == pb.MODEL_OP_KIND_LOAD: - await self.store.ensure_local(ref, snap) - snapshots = {ref: snap} if snap is not None else None - specs = [s for s in self.specs.values() - if ref in (wire_ref(b) for b in s.models.values())] - if not specs: - # pgw#532: derived (per-pick) specs bind refs the static - # decls never name — a LOAD for a pick this worker has - # dispatched before must promote/re-set-up its instance. - seen: set = set() - for rec in self._classes.values(): - if id(rec) in seen: - continue - seen.add(id(rec)) - for s in rec.specs: - if s not in specs and ref in ( - wire_ref(b) for b in s.models.values()): - specs.append(s) - if not specs: - # No endpoint binds this ref; nothing owns a VRAM load - # for it. (For a Slot-declared endpoint the bytes+snapshot - # ARE banked above, so the next dispatch naming this pick - # materializes instantly — the pre-warm degrades to a - # download, never to an upstream fetch.) - await self._send(pb.WorkerMessage(model_event=pb.ModelEvent( - ref=ref, state=pb.MODEL_STATE_FAILED, error="load_failed"))) - return - ready = [ - s for s in specs - if (owner := self._classes.get(s.instance_key)) is not None - and owner.ready - ] - if ready: - # A resident instance already owns the ref: promote/touch it. - # Never cold-set-up the OTHER specs sharing the ref — a - # shared companion (one vae bound to every variant of a - # family) would cascade one LOAD into loading every sibling - # checkpoint (gw#465). - for s in ready: - await self.ensure_setup(s, snapshots) - else: - target = next( - (s for s in specs - if not self._missing_slot_refs(s, snapshots)), None) - if target is None: - # Every candidate needs a sibling slot the worker cannot - # materialize. Name the blockers so the hub re-mints and - # re-sends DOWNLOAD for them, and fail THIS op with the - # same vocabulary — never a phantom download_failed. - blockers: set = set() - for s in specs: - blockers |= self._missing_slot_refs(s, snapshots) - for m in sorted(blockers): - logger.warning( - "ModelOp LOAD %s deferred: sibling slot %s has " - "no snapshot; reporting missing_snapshot", ref, m) - await self._send(pb.WorkerMessage(model_event=pb.ModelEvent( - ref=m, state=pb.MODEL_STATE_FAILED, - error="missing_snapshot"))) - await self._send(pb.WorkerMessage(model_event=pb.ModelEvent( - ref=ref, state=pb.MODEL_STATE_FAILED, - error="missing_snapshot"))) - return - await self.ensure_setup(target, snapshots) - elif op.op == pb.MODEL_OP_KIND_UNLOAD: - if self._ref_in_use(ref): - await self._send(pb.WorkerMessage(model_event=pb.ModelEvent( - ref=ref, state=pb.MODEL_STATE_FAILED, error="model_in_use"))) - return - await self._unload_ref(ref) - elif op.op == pb.MODEL_OP_KIND_ADOPT_COMPILE_CACHE: - await self._adopt_compile_cache(ref, snap) + await self._adopt_compile_cache(ref, snap) 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. 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=_model_op_error_vocab(exc)))) + logger.warning("compile-cache adoption on %s failed: %s", ref, exc) + await self._send(pb.WorkerMessage(model_event=pb.ModelEvent( + ref=ref, + state=pb.MODEL_STATE_FAILED, + error=f"adopt_failed:{type(exc).__name__.lower()}", + ))) async def _adopt_compile_cache(self, ref: str, snap: Optional[pb.Snapshot]) -> None: """Hot adoption (th#567): download+verify a compiled artifact and @@ -2570,23 +2524,6 @@ async def _vacate_record(self, rec: _ClassRecord) -> None: self.store.residency.drain_shared() 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 -------------------------------------------------------- async def handle_run_job(self, run: pb.RunJob) -> None: diff --git a/src/gen_worker/lifecycle.py b/src/gen_worker/lifecycle.py index d10c33b4..f5d05ea2 100644 --- a/src/gen_worker/lifecycle.py +++ b/src/gen_worker/lifecycle.py @@ -96,6 +96,9 @@ def __init__(self, settings: Settings, executor: Executor) -> None: self._emitted_degraded: dict[str, str] = {} self._drain_task: Optional[asyncio.Task] = None self._drain_deadline_at: Optional[float] = None + self._desired_residency: Optional[pb.DesiredResidency] = None + self._residency_task: Optional[asyncio.Task] = None + self._observed_residency_generation = 0 # ---- snapshots ----------------------------------------------------------- @@ -111,6 +114,7 @@ def _state_delta(self) -> pb.StateDelta: # gw#516: encode/upload tails past the GPU-slot release. The hub # must not drain/retire this worker on GPU-idleness alone. finalizing_jobs=self.executor.finalizing_jobs(), + observed_residency_generation=self._observed_residency_generation, ) def build_resources(self) -> pb.WorkerResources: @@ -149,14 +153,25 @@ def build_hello(self) -> pb.Hello: ) async def on_hello_ack(self, ack: pb.HelloAck) -> None: - # Full-replace config: file base URL + disk-retention keep set. + # Full-replace config: file base URL + desired model residency. self.executor.file_base_url = ack.file_base_url or "" - self.executor.store.keep = set(ack.keep) # th#697: apply the hub's precision-ladder picks for THIS card # (full-replace: refs absent from the map revert to declared). self.executor.apply_model_resolutions({ r.ref: (r.resolved_ref, r.cast) for r in ack.resolutions }) + desired = pb.DesiredResidency() + desired.CopyFrom(ack.desired_residency) + generation = int(desired.generation) + if generation < self._observed_residency_generation: + logger.info( + "ignoring stale desired residency generation %d (observed %d)", + generation, self._observed_residency_generation, + ) + else: + self._observed_residency_generation = generation + self.executor.store.keep = {ref for ref in desired.disk_refs if ref} + self._replace_residency_reconcile(desired) # New connection: per-worker fn disables/degradations were wiped by # Hello; re-emit any that still hold, then re-baseline dynamic state. self._emitted_unavailable.clear() @@ -165,19 +180,79 @@ async def on_hello_ack(self, ack: pb.HelloAck) -> None: self._last_delta = None await self.maybe_send_state_delta() + def _replace_residency_reconcile(self, desired: "pb.DesiredResidency") -> None: + self._cancel_residency_reconcile() + self._desired_residency = desired + self._resume_residency_reconcile() + + def _cancel_residency_reconcile(self) -> None: + task = getattr(self, "_residency_task", None) + if task is not None: + task.cancel() + self._residency_task = None + + def _resume_residency_reconcile(self) -> None: + desired = getattr(self, "_desired_residency", None) + if desired is None or self.draining: + return + self._residency_task = asyncio.create_task( + self._reconcile_residency(desired), + name=f"residency-{int(desired.generation)}", + ) + + async def _reconcile_residency(self, desired: "pb.DesiredResidency") -> None: + """Converge in declared order while tenant work has first claim.""" + snapshots = dict(desired.snapshots) + try: + for ref in desired.disk_refs: + if not ref: + continue + await self.executor.wait_idle() + if self.draining: + return + try: + await self.executor.store.ensure_local(ref, snapshots.get(ref)) + except asyncio.CancelledError: + raise + except Exception as exc: + logger.warning("desired disk residency failed for %s: %s", ref, exc) + + for instance in desired.hot: + await self.executor.wait_idle() + if self.draining: + return + try: + await self.executor.ensure_desired_instance(instance, snapshots) + except asyncio.CancelledError: + raise + except Exception as exc: + logger.warning( + "desired hot residency failed for %s: %s", + instance.function_name, exc, + ) + except asyncio.CancelledError: + return + async def on_message(self, msg: pb.SchedulerMessage) -> None: which = msg.WhichOneof("msg") if which == "run_job": - await self.executor.handle_run_job(msg.run_job) + # A tenant request preempts unrelated background transfer/setup. + # Re-running desired state after idle is cheap: local refs and + # ready instances short-circuit through the existing dedupe paths. + self._cancel_residency_reconcile() + try: + await self.executor.handle_run_job(msg.run_job) + finally: + self._resume_residency_reconcile() elif which == "cancel_job": self.executor.handle_cancel(msg.cancel_job) elif which == "model_op": - # Model work must never block the receive path. - asyncio.create_task(self._model_op_then_delta(msg.model_op)) + # Compile-cache adoption must never block the receive path. + asyncio.create_task(self._adopt_compile_cache_then_delta(msg.model_op)) elif which == "drain": self.start_drain(int(msg.drain.deadline_ms or 0)) - async def _model_op_then_delta(self, op: pb.ModelOp) -> None: + async def _adopt_compile_cache_then_delta(self, op: pb.ModelOp) -> None: try: await self.executor.handle_model_op(op) finally: @@ -285,7 +360,7 @@ async def startup(self) -> None: ref = wire_ref(binding) if binding.source != "tensorhub" and ref not in prefetch_refs: # hf/civitai refs need no orchestrator snapshot; tensorhub - # refs arrive via ModelOp{DOWNLOAD} after HelloAck (§7). + # refs arrive via DesiredResidency / RunJob (§7). prefetch_refs.append(ref) if prefetch_refs: @@ -304,7 +379,7 @@ async def startup(self) -> None: continue if spec.slots: # pgw#532: dynamic slots materialize the HUB-resolved ref - # (ModelOp pre-drives / RunJob resolved_models + snapshots), + # (DesiredResidency pre-warms / RunJob supplies snapshots), # per dispatch, instance-per-pick. Setting up eagerly here # would load the code seed — the exact fc157 setup-failure # bug (raw civitai default -> civitai_not_found -> every @@ -317,12 +392,11 @@ async def startup(self) -> None: and self.executor.store.local_path(wire_ref(b)) is None }) if missing: - # Waits for ModelOp / RunJob snapshots (§7). This wait is - # UNBOUNDED and hub-driven: if the release was registered - # without bindings the hub's keep set is empty, no - # ModelOp{DOWNLOAD} ever arrives, and the function sits in - # loading_functions forever (ie#455 z-image fns=[]) — so say - # so loudly instead of dropping the function in silence. + # Waits for DesiredResidency / RunJob snapshots (§7). This wait + # is unbounded and hub-driven: a release without resolved + # desired bindings leaves the function in loading_functions + # forever (ie#455 z-image fns=[]), so say so loudly instead of + # dropping the function in silence. awaiting_hub[spec.name] = missing continue try: @@ -335,10 +409,9 @@ async def startup(self) -> None: "(pgw#532; no boot-time setup): %s", ", ".join(sorted(dynamic))) if awaiting_hub: logger.warning( - "functions waiting on hub-supplied snapshots (ModelOp{DOWNLOAD}, " + "functions waiting on hub-supplied snapshots (DesiredResidency, " "contract §7) and NOT yet servable: %s — if these never arrive, " - "check that the release was registered WITH function bindings " - "for these refs (hub keep set must not be empty)", + "check that the release has resolved desired bindings for these refs", "; ".join(f"{fn} <- {', '.join(refs)}" for fn, refs in sorted(awaiting_hub.items())), ) @@ -364,6 +437,7 @@ def _begin_drain(self, deadline_ms: int) -> None: """Synchronously stop admission and anchor the deadline at receipt.""" self.draining = True self.executor.draining = True + self._cancel_residency_reconcile() logger.info("drain started (deadline_ms=%d)", deadline_ms) deadline_s = (deadline_ms / 1000.0) if deadline_ms > 0 else None loop = asyncio.get_running_loop() diff --git a/src/gen_worker/models/download.py b/src/gen_worker/models/download.py index 3b68e29c..e5cc41af 100644 --- a/src/gen_worker/models/download.py +++ b/src/gen_worker/models/download.py @@ -217,7 +217,7 @@ async def ensure_local( — filtering what gets written here without the orchestrator also filtering what it verifies would turn every boot into a spurious corruption/quarantine loop. Selective fetch for tensorhub-sourced - snapshots is therefore the hub's ModelOp DOWNLOAD scoping (the manifest + snapshots is therefore the hub's desired-snapshot scoping (the manifest now carries ``components`` for it to read) — worker-side selective fetch for tensorhub refs lands on the hub-less CLI path instead (``models/provision.py::_fetch_tensorhub_snapshot``), which owns its own diff --git a/src/gen_worker/models/provision.py b/src/gen_worker/models/provision.py index 1faaa94e..c266b9d4 100644 --- a/src/gen_worker/models/provision.py +++ b/src/gen_worker/models/provision.py @@ -353,7 +353,7 @@ def _fetch_tensorhub_snapshot( ``components`` (pgw#505): th#560's resolve route always returns the FULL repo manifest today (selective CAS resolve is the hub-side - ModelOp DOWNLOAD scoping — a separate, not-yet-built platform change). + desired-snapshot scoping — a separate, not-yet-built platform change). Until then this narrows client-side: the worker fully owns this resolve+download+materialize loop (unlike the production executor path, which digest-verifies against an orchestrator-issued file list), so it diff --git a/src/gen_worker/models/residency.py b/src/gen_worker/models/residency.py index 939e776c..b11c3bf7 100644 --- a/src/gen_worker/models/residency.py +++ b/src/gen_worker/models/residency.py @@ -475,7 +475,7 @@ def promote(self, ref: str, device: str = "cuda") -> bool: return True def release_to_disk(self, ref: str) -> bool: - """Drop the loaded object entirely (UNLOAD); disk snapshot kept. + """Drop the loaded object entirely; disk snapshot kept. Refuses executing/shared-held entries.""" with self._lock: e = self._entries.get(ref) diff --git a/src/gen_worker/pb/worker_scheduler_pb2.py b/src/gen_worker/pb/worker_scheduler_pb2.py index 3a99b877..85289baf 100644 --- a/src/gen_worker/pb/worker_scheduler_pb2.py +++ b/src/gen_worker/pb/worker_scheduler_pb2.py @@ -24,7 +24,7 @@ -DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x16worker_scheduler.proto\x12\x0e\x63ozy.scheduler\"\xab\x03\n\rWorkerMessage\x12&\n\x05hello\x18\x01 \x01(\x0b\x32\x15.cozy.scheduler.HelloH\x00\x12\x31\n\x0bstate_delta\x18\x02 \x01(\x0b\x32\x1a.cozy.scheduler.StateDeltaH\x00\x12\x33\n\x0cjob_accepted\x18\x03 \x01(\x0b\x32\x1b.cozy.scheduler.JobAcceptedH\x00\x12/\n\njob_result\x18\x04 \x01(\x0b\x32\x19.cozy.scheduler.JobResultH\x00\x12\x33\n\x0cjob_progress\x18\x05 \x01(\x0b\x32\x1b.cozy.scheduler.JobProgressH\x00\x12\x31\n\x0bmodel_event\x18\x06 \x01(\x0b\x32\x1a.cozy.scheduler.ModelEventH\x00\x12\x37\n\x0e\x66n_unavailable\x18\x07 \x01(\x0b\x32\x1d.cozy.scheduler.FnUnavailableH\x00\x12\x31\n\x0b\x66n_degraded\x18\x08 \x01(\x0b\x32\x1a.cozy.scheduler.FnDegradedH\x00\x42\x05\n\x03msg\"\xb0\x02\n\x10SchedulerMessage\x12-\n\thello_ack\x18\x01 \x01(\x0b\x32\x18.cozy.scheduler.HelloAckH\x00\x12)\n\x07run_job\x18\x02 \x01(\x0b\x32\x16.cozy.scheduler.RunJobH\x00\x12/\n\ncancel_job\x18\x03 \x01(\x0b\x32\x19.cozy.scheduler.CancelJobH\x00\x12+\n\x08model_op\x18\x04 \x01(\x0b\x32\x17.cozy.scheduler.ModelOpH\x00\x12&\n\x05\x64rain\x18\x05 \x01(\x0b\x32\x15.cozy.scheduler.DrainH\x00\x12\x35\n\rtoken_refresh\x18\x06 \x01(\x0b\x32\x1c.cozy.scheduler.TokenRefreshH\x00\x42\x05\n\x03msg\"\xa8\x02\n\x05Hello\x12\x39\n\x10protocol_version\x18\x01 \x01(\x0e\x32\x1f.cozy.scheduler.ProtocolVersion\x12\x11\n\tworker_id\x18\x02 \x01(\t\x12\x12\n\nrelease_id\x18\x03 \x01(\t\x12\x32\n\tresources\x18\x04 \x01(\x0b\x32\x1f.cozy.scheduler.WorkerResources\x12)\n\x05state\x18\x05 \x01(\x0b\x32\x1a.cozy.scheduler.StateDelta\x12.\n\x06models\x18\x06 \x03(\x0b\x32\x1e.cozy.scheduler.ModelResidency\x12.\n\tin_flight\x18\x07 \x03(\x0b\x32\x1b.cozy.scheduler.InFlightJob\"\xb7\x01\n\x0fWorkerResources\x12\x11\n\tgpu_count\x18\x01 \x01(\x05\x12\x18\n\x10vram_total_bytes\x18\x02 \x01(\x03\x12\x10\n\x08gpu_name\x18\x03 \x01(\t\x12\x0e\n\x06gpu_sm\x18\x04 \x01(\t\x12\x16\n\x0einstalled_libs\x18\x05 \x03(\t\x12\x14\n\x0cimage_digest\x18\x06 \x01(\t\x12\x12\n\ngit_commit\x18\x07 \x01(\t\x12\x13\n\x0binstance_id\x18\x08 \x01(\t\"^\n\x0eModelResidency\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12+\n\x04tier\x18\x02 \x01(\x0e\x32\x1d.cozy.scheduler.ResidencyTier\x12\x12\n\nvram_bytes\x18\x03 \x01(\x03\"2\n\x0bInFlightJob\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\"\xa0\x01\n\x08HelloAck\x12\x39\n\x10protocol_version\x18\x01 \x01(\x0e\x32\x1f.cozy.scheduler.ProtocolVersion\x12\x15\n\rfile_base_url\x18\x02 \x01(\t\x12\x0c\n\x04keep\x18\x03 \x03(\t\x12\x34\n\x0bresolutions\x18\x04 \x03(\x0b\x32\x1f.cozy.scheduler.ModelResolution\"B\n\x0fModelResolution\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12\x14\n\x0cresolved_ref\x18\x02 \x01(\t\x12\x0c\n\x04\x63\x61st\x18\x03 \x01(\t\"\xa2\x01\n\nStateDelta\x12*\n\x05phase\x18\x01 \x01(\x0e\x32\x1b.cozy.scheduler.WorkerPhase\x12\x1b\n\x13\x61vailable_functions\x18\x02 \x03(\t\x12\x19\n\x11loading_functions\x18\x03 \x03(\t\x12\x17\n\x0f\x66ree_vram_bytes\x18\x04 \x01(\x03\x12\x17\n\x0f\x66inalizing_jobs\x18\x05 \x01(\x05\"\xc4\x03\n\x06RunJob\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\x12\x15\n\rfunction_name\x18\x03 \x01(\t\x12\x15\n\rinput_payload\x18\x04 \x01(\x0c\x12\x12\n\ntimeout_ms\x18\x05 \x01(\x03\x12\x0e\n\x06tenant\x18\x06 \x01(\t\x12\x12\n\ninvoker_id\x18\x07 \x01(\t\x12\x18\n\x10\x63\x61pability_token\x18\x08 \x01(\t\x12/\n\x0boutput_mode\x18\t \x01(\x0e\x32\x1a.cozy.scheduler.OutputMode\x12\x30\n\x07\x63ompute\x18\n \x01(\x0b\x32\x1f.cozy.scheduler.ResolvedCompute\x12,\n\x06models\x18\x0b \x03(\x0b\x32\x1c.cozy.scheduler.ModelBinding\x12\x38\n\tsnapshots\x18\x0c \x03(\x0b\x32%.cozy.scheduler.RunJob.SnapshotsEntry\x1aJ\n\x0eSnapshotsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\'\n\x05value\x18\x02 \x01(\x0b\x32\x18.cozy.scheduler.Snapshot:\x02\x38\x01\"Y\n\x0fResolvedCompute\x12\x13\n\x0b\x61\x63\x63\x65lerator\x18\x01 \x01(\t\x12\x11\n\tgpu_index\x18\x02 \x01(\x05J\x04\x08\x03\x10\x04J\x04\x08\x04\x10\x05R\tgpu_countR\x07vram_gb\"q\n\x0cModelBinding\x12\x0c\n\x04slot\x18\x01 \x01(\t\x12\x0b\n\x03ref\x18\x02 \x01(\t\x12*\n\x05loras\x18\x03 \x03(\x0b\x32\x1b.cozy.scheduler.LoraOverlay\x12\x1a\n\x12inference_defaults\x18\x04 \x01(\t\"F\n\x0bLoraOverlay\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12\x0e\n\x06weight\x18\x02 \x01(\x01\x12\x1a\n\x12inference_defaults\x18\x03 \x01(\t\"G\n\x08Snapshot\x12\x0e\n\x06\x64igest\x18\x01 \x01(\t\x12+\n\x05\x66iles\x18\x02 \x03(\x0b\x32\x1c.cozy.scheduler.SnapshotFile\"M\n\x0cSnapshotFile\x12\x0c\n\x04path\x18\x01 \x01(\t\x12\x12\n\nsize_bytes\x18\x02 \x01(\x03\x12\x0e\n\x06\x62lake3\x18\x03 \x01(\t\x12\x0b\n\x03url\x18\x04 \x01(\t\"2\n\x0bJobAccepted\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\"\xce\x01\n\tJobResult\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\x12)\n\x06status\x18\x03 \x01(\x0e\x32\x19.cozy.scheduler.JobStatus\x12\x10\n\x06inline\x18\x04 \x01(\x0cH\x00\x12\x12\n\x08\x62lob_ref\x18\x05 \x01(\tH\x00\x12\x14\n\x0csafe_message\x18\x06 \x01(\t\x12+\n\x07metrics\x18\x07 \x01(\x0b\x32\x1a.cozy.scheduler.JobMetricsB\x08\n\x06output\"\xb4\x02\n\nJobMetrics\x12\x12\n\nruntime_ms\x18\x01 \x01(\x03\x12\x10\n\x08queue_ms\x18\x02 \x01(\x03\x12\x18\n\x10rss_at_end_bytes\x18\x03 \x01(\x03\x12\x17\n\x0fpeak_vram_bytes\x18\x04 \x01(\x03\x12\x1c\n\x14\x63oncurrency_at_start\x18\x05 \x01(\x05\x12\x1f\n\x17output_media_duration_s\x18\x06 \x01(\x01\x12\x14\n\x0cinput_tokens\x18\x07 \x01(\x03\x12\x1b\n\x13input_cached_tokens\x18\x08 \x01(\x03\x12\x15\n\routput_tokens\x18\t \x01(\x03\x12\x14\n\x0coutput_count\x18\n \x01(\x03\x12\x14\n\x0cslot_held_ms\x18\x0b \x01(\x03\x12\x18\n\x10\x66inalize_wall_ms\x18\x0c \x01(\x03\"c\n\x0bJobProgress\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\x12\x0b\n\x03seq\x18\x03 \x01(\x03\x12\x0c\n\x04\x64\x61ta\x18\x04 \x01(\x0c\x12\x14\n\x0c\x63ontent_type\x18\x05 \x01(\t\"0\n\tCancelJob\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\"k\n\x07ModelOp\x12\'\n\x02op\x18\x01 \x01(\x0e\x32\x1b.cozy.scheduler.ModelOpKind\x12\x0b\n\x03ref\x18\x02 \x01(\t\x12*\n\x08snapshot\x18\x03 \x01(\x0b\x32\x18.cozy.scheduler.Snapshot\"\xe1\x01\n\nModelEvent\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12)\n\x05state\x18\x02 \x01(\x0e\x32\x1a.cozy.scheduler.ModelState\x12\x12\n\nvram_bytes\x18\x03 \x01(\x03\x12\r\n\x05\x65rror\x18\x04 \x01(\t\x12\x12\n\nbytes_done\x18\x05 \x01(\x03\x12\x13\n\x0b\x62ytes_total\x18\x06 \x01(\x03\x12\x13\n\x0b\x64uration_ms\x18\x07 \x01(\x03\x12\x12\n\ncache_hits\x18\x08 \x01(\x03\x12\x14\n\x0c\x63\x61\x63he_misses\x18\t \x01(\x03\x12\x10\n\x08warmup_s\x18\n \x01(\x01\"\xaa\x01\n\rFnUnavailable\x12\x15\n\rfunction_name\x18\x01 \x01(\t\x12\x0e\n\x06reason\x18\x02 \x01(\t\x12\x0e\n\x06\x64\x65tail\x18\x03 \x01(\t\x12\x35\n\x04\x61xes\x18\x04 \x03(\x0b\x32\'.cozy.scheduler.FnUnavailable.AxesEntry\x1a+\n\tAxesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"\x8d\x01\n\nFnDegraded\x12\x15\n\rfunction_name\x18\x01 \x01(\t\x12\x0e\n\x06wanted\x18\x02 \x01(\t\x12\x0b\n\x03ran\x18\x03 \x01(\t\x12\x0e\n\x06reason\x18\x04 \x01(\t\x12\x1e\n\x16\x65st_latency_multiplier\x18\x05 \x01(\x01\x12\x1b\n\x13recommended_vram_gb\x18\x06 \x01(\x01\"\x1c\n\x05\x44rain\x12\x13\n\x0b\x64\x65\x61\x64line_ms\x18\x01 \x01(\x03\"6\n\x0cTokenRefresh\x12\r\n\x05token\x18\x01 \x01(\t\x12\x17\n\x0f\x65xpires_at_unix\x18\x02 \x01(\x03*Q\n\x0fProtocolVersion\x12 \n\x1cPROTOCOL_VERSION_UNSPECIFIED\x10\x00\x12\x1c\n\x18PROTOCOL_VERSION_CURRENT\x10\x02*y\n\rResidencyTier\x12\x1e\n\x1aRESIDENCY_TIER_UNSPECIFIED\x10\x00\x12\x17\n\x13RESIDENCY_TIER_DISK\x10\x01\x12\x16\n\x12RESIDENCY_TIER_RAM\x10\x02\x12\x17\n\x13RESIDENCY_TIER_VRAM\x10\x03*\xd8\x01\n\x0bWorkerPhase\x12\x1c\n\x18WORKER_PHASE_UNSPECIFIED\x10\x00\x12\x18\n\x14WORKER_PHASE_BOOTING\x10\x01\x12#\n\x1fWORKER_PHASE_DOWNLOADING_MODELS\x10\x02\x12\"\n\x1eWORKER_PHASE_LOADING_PIPELINES\x10\x03\x12\x18\n\x14WORKER_PHASE_WARMING\x10\x04\x12\x16\n\x12WORKER_PHASE_READY\x10\x05\x12\x16\n\x12WORKER_PHASE_ERROR\x10\x06*V\n\nOutputMode\x12\x1b\n\x17OUTPUT_MODE_UNSPECIFIED\x10\x00\x12\x13\n\x0fOUTPUT_MODE_URL\x10\x01\x12\x16\n\x12OUTPUT_MODE_INLINE\x10\x02*\x9b\x01\n\tJobStatus\x12\x1a\n\x16JOB_STATUS_UNSPECIFIED\x10\x00\x12\x11\n\rJOB_STATUS_OK\x10\x01\x12\x16\n\x12JOB_STATUS_INVALID\x10\x02\x12\x18\n\x14JOB_STATUS_RETRYABLE\x10\x03\x12\x14\n\x10JOB_STATUS_FATAL\x10\x04\x12\x17\n\x13JOB_STATUS_CANCELED\x10\x05*\xa1\x01\n\x0bModelOpKind\x12\x1d\n\x19MODEL_OP_KIND_UNSPECIFIED\x10\x00\x12\x1a\n\x16MODEL_OP_KIND_DOWNLOAD\x10\x01\x12\x16\n\x12MODEL_OP_KIND_LOAD\x10\x02\x12\x18\n\x14MODEL_OP_KIND_UNLOAD\x10\x03\x12%\n!MODEL_OP_KIND_ADOPT_COMPILE_CACHE\x10\x04*\xda\x01\n\nModelState\x12\x1b\n\x17MODEL_STATE_UNSPECIFIED\x10\x00\x12\x1b\n\x17MODEL_STATE_DOWNLOADING\x10\x01\x12\x17\n\x13MODEL_STATE_ON_DISK\x10\x02\x12\x16\n\x12MODEL_STATE_IN_RAM\x10\x03\x12\x17\n\x13MODEL_STATE_IN_VRAM\x10\x04\x12\x17\n\x13MODEL_STATE_EVICTED\x10\x05\x12\x16\n\x12MODEL_STATE_FAILED\x10\x06\x12\x17\n\x13MODEL_STATE_ADOPTED\x10\x07\x32\x61\n\x0fWorkerScheduler\x12N\n\x07\x43onnect\x12\x1d.cozy.scheduler.WorkerMessage\x1a .cozy.scheduler.SchedulerMessage(\x01\x30\x01\x42\x61Z_github.com/cozy-creator/tensorhub/internal/orchestrator/grpc/pb/workerscheduler;workerschedulerb\x06proto3') +DESCRIPTOR = _descriptor_pool.Default().AddSerializedFile(b'\n\x16worker_scheduler.proto\x12\x0e\x63ozy.scheduler\"\xab\x03\n\rWorkerMessage\x12&\n\x05hello\x18\x01 \x01(\x0b\x32\x15.cozy.scheduler.HelloH\x00\x12\x31\n\x0bstate_delta\x18\x02 \x01(\x0b\x32\x1a.cozy.scheduler.StateDeltaH\x00\x12\x33\n\x0cjob_accepted\x18\x03 \x01(\x0b\x32\x1b.cozy.scheduler.JobAcceptedH\x00\x12/\n\njob_result\x18\x04 \x01(\x0b\x32\x19.cozy.scheduler.JobResultH\x00\x12\x33\n\x0cjob_progress\x18\x05 \x01(\x0b\x32\x1b.cozy.scheduler.JobProgressH\x00\x12\x31\n\x0bmodel_event\x18\x06 \x01(\x0b\x32\x1a.cozy.scheduler.ModelEventH\x00\x12\x37\n\x0e\x66n_unavailable\x18\x07 \x01(\x0b\x32\x1d.cozy.scheduler.FnUnavailableH\x00\x12\x31\n\x0b\x66n_degraded\x18\x08 \x01(\x0b\x32\x1a.cozy.scheduler.FnDegradedH\x00\x42\x05\n\x03msg\"\xb0\x02\n\x10SchedulerMessage\x12-\n\thello_ack\x18\x01 \x01(\x0b\x32\x18.cozy.scheduler.HelloAckH\x00\x12)\n\x07run_job\x18\x02 \x01(\x0b\x32\x16.cozy.scheduler.RunJobH\x00\x12/\n\ncancel_job\x18\x03 \x01(\x0b\x32\x19.cozy.scheduler.CancelJobH\x00\x12+\n\x08model_op\x18\x04 \x01(\x0b\x32\x17.cozy.scheduler.ModelOpH\x00\x12&\n\x05\x64rain\x18\x05 \x01(\x0b\x32\x15.cozy.scheduler.DrainH\x00\x12\x35\n\rtoken_refresh\x18\x06 \x01(\x0b\x32\x1c.cozy.scheduler.TokenRefreshH\x00\x42\x05\n\x03msg\"\xa8\x02\n\x05Hello\x12\x39\n\x10protocol_version\x18\x01 \x01(\x0e\x32\x1f.cozy.scheduler.ProtocolVersion\x12\x11\n\tworker_id\x18\x02 \x01(\t\x12\x12\n\nrelease_id\x18\x03 \x01(\t\x12\x32\n\tresources\x18\x04 \x01(\x0b\x32\x1f.cozy.scheduler.WorkerResources\x12)\n\x05state\x18\x05 \x01(\x0b\x32\x1a.cozy.scheduler.StateDelta\x12.\n\x06models\x18\x06 \x03(\x0b\x32\x1e.cozy.scheduler.ModelResidency\x12.\n\tin_flight\x18\x07 \x03(\x0b\x32\x1b.cozy.scheduler.InFlightJob\"\xb7\x01\n\x0fWorkerResources\x12\x11\n\tgpu_count\x18\x01 \x01(\x05\x12\x18\n\x10vram_total_bytes\x18\x02 \x01(\x03\x12\x10\n\x08gpu_name\x18\x03 \x01(\t\x12\x0e\n\x06gpu_sm\x18\x04 \x01(\t\x12\x16\n\x0einstalled_libs\x18\x05 \x03(\t\x12\x14\n\x0cimage_digest\x18\x06 \x01(\t\x12\x12\n\ngit_commit\x18\x07 \x01(\t\x12\x13\n\x0binstance_id\x18\x08 \x01(\t\"^\n\x0eModelResidency\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12+\n\x04tier\x18\x02 \x01(\x0e\x32\x1d.cozy.scheduler.ResidencyTier\x12\x12\n\nvram_bytes\x18\x03 \x01(\x03\"2\n\x0bInFlightJob\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\"\xdd\x01\n\x08HelloAck\x12\x39\n\x10protocol_version\x18\x01 \x01(\x0e\x32\x1f.cozy.scheduler.ProtocolVersion\x12\x15\n\rfile_base_url\x18\x02 \x01(\t\x12\x0c\n\x04keep\x18\x03 \x03(\t\x12\x34\n\x0bresolutions\x18\x04 \x03(\x0b\x32\x1f.cozy.scheduler.ModelResolution\x12;\n\x11\x64\x65sired_residency\x18\x05 \x01(\x0b\x32 .cozy.scheduler.DesiredResidency\"\xf7\x01\n\x10\x44\x65siredResidency\x12\x12\n\ngeneration\x18\x01 \x01(\x04\x12\x11\n\tdisk_refs\x18\x02 \x03(\t\x12,\n\x03hot\x18\x03 \x03(\x0b\x32\x1f.cozy.scheduler.DesiredInstance\x12\x42\n\tsnapshots\x18\x04 \x03(\x0b\x32/.cozy.scheduler.DesiredResidency.SnapshotsEntry\x1aJ\n\x0eSnapshotsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\'\n\x05value\x18\x02 \x01(\x0b\x32\x18.cozy.scheduler.Snapshot:\x02\x38\x01\"V\n\x0f\x44\x65siredInstance\x12\x15\n\rfunction_name\x18\x01 \x01(\t\x12,\n\x06models\x18\x02 \x03(\x0b\x32\x1c.cozy.scheduler.ModelBinding\"B\n\x0fModelResolution\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12\x14\n\x0cresolved_ref\x18\x02 \x01(\t\x12\x0c\n\x04\x63\x61st\x18\x03 \x01(\t\"\xc9\x01\n\nStateDelta\x12*\n\x05phase\x18\x01 \x01(\x0e\x32\x1b.cozy.scheduler.WorkerPhase\x12\x1b\n\x13\x61vailable_functions\x18\x02 \x03(\t\x12\x19\n\x11loading_functions\x18\x03 \x03(\t\x12\x17\n\x0f\x66ree_vram_bytes\x18\x04 \x01(\x03\x12\x17\n\x0f\x66inalizing_jobs\x18\x05 \x01(\x05\x12%\n\x1dobserved_residency_generation\x18\x06 \x01(\x04\"\xc4\x03\n\x06RunJob\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\x12\x15\n\rfunction_name\x18\x03 \x01(\t\x12\x15\n\rinput_payload\x18\x04 \x01(\x0c\x12\x12\n\ntimeout_ms\x18\x05 \x01(\x03\x12\x0e\n\x06tenant\x18\x06 \x01(\t\x12\x12\n\ninvoker_id\x18\x07 \x01(\t\x12\x18\n\x10\x63\x61pability_token\x18\x08 \x01(\t\x12/\n\x0boutput_mode\x18\t \x01(\x0e\x32\x1a.cozy.scheduler.OutputMode\x12\x30\n\x07\x63ompute\x18\n \x01(\x0b\x32\x1f.cozy.scheduler.ResolvedCompute\x12,\n\x06models\x18\x0b \x03(\x0b\x32\x1c.cozy.scheduler.ModelBinding\x12\x38\n\tsnapshots\x18\x0c \x03(\x0b\x32%.cozy.scheduler.RunJob.SnapshotsEntry\x1aJ\n\x0eSnapshotsEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\'\n\x05value\x18\x02 \x01(\x0b\x32\x18.cozy.scheduler.Snapshot:\x02\x38\x01\"Y\n\x0fResolvedCompute\x12\x13\n\x0b\x61\x63\x63\x65lerator\x18\x01 \x01(\t\x12\x11\n\tgpu_index\x18\x02 \x01(\x05J\x04\x08\x03\x10\x04J\x04\x08\x04\x10\x05R\tgpu_countR\x07vram_gb\"q\n\x0cModelBinding\x12\x0c\n\x04slot\x18\x01 \x01(\t\x12\x0b\n\x03ref\x18\x02 \x01(\t\x12*\n\x05loras\x18\x03 \x03(\x0b\x32\x1b.cozy.scheduler.LoraOverlay\x12\x1a\n\x12inference_defaults\x18\x04 \x01(\t\"F\n\x0bLoraOverlay\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12\x0e\n\x06weight\x18\x02 \x01(\x01\x12\x1a\n\x12inference_defaults\x18\x03 \x01(\t\"G\n\x08Snapshot\x12\x0e\n\x06\x64igest\x18\x01 \x01(\t\x12+\n\x05\x66iles\x18\x02 \x03(\x0b\x32\x1c.cozy.scheduler.SnapshotFile\"M\n\x0cSnapshotFile\x12\x0c\n\x04path\x18\x01 \x01(\t\x12\x12\n\nsize_bytes\x18\x02 \x01(\x03\x12\x0e\n\x06\x62lake3\x18\x03 \x01(\t\x12\x0b\n\x03url\x18\x04 \x01(\t\"2\n\x0bJobAccepted\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\"\xce\x01\n\tJobResult\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\x12)\n\x06status\x18\x03 \x01(\x0e\x32\x19.cozy.scheduler.JobStatus\x12\x10\n\x06inline\x18\x04 \x01(\x0cH\x00\x12\x12\n\x08\x62lob_ref\x18\x05 \x01(\tH\x00\x12\x14\n\x0csafe_message\x18\x06 \x01(\t\x12+\n\x07metrics\x18\x07 \x01(\x0b\x32\x1a.cozy.scheduler.JobMetricsB\x08\n\x06output\"\xb4\x02\n\nJobMetrics\x12\x12\n\nruntime_ms\x18\x01 \x01(\x03\x12\x10\n\x08queue_ms\x18\x02 \x01(\x03\x12\x18\n\x10rss_at_end_bytes\x18\x03 \x01(\x03\x12\x17\n\x0fpeak_vram_bytes\x18\x04 \x01(\x03\x12\x1c\n\x14\x63oncurrency_at_start\x18\x05 \x01(\x05\x12\x1f\n\x17output_media_duration_s\x18\x06 \x01(\x01\x12\x14\n\x0cinput_tokens\x18\x07 \x01(\x03\x12\x1b\n\x13input_cached_tokens\x18\x08 \x01(\x03\x12\x15\n\routput_tokens\x18\t \x01(\x03\x12\x14\n\x0coutput_count\x18\n \x01(\x03\x12\x14\n\x0cslot_held_ms\x18\x0b \x01(\x03\x12\x18\n\x10\x66inalize_wall_ms\x18\x0c \x01(\x03\"c\n\x0bJobProgress\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\x12\x0b\n\x03seq\x18\x03 \x01(\x03\x12\x0c\n\x04\x64\x61ta\x18\x04 \x01(\x0c\x12\x14\n\x0c\x63ontent_type\x18\x05 \x01(\t\"0\n\tCancelJob\x12\x12\n\nrequest_id\x18\x01 \x01(\t\x12\x0f\n\x07\x61ttempt\x18\x02 \x01(\x03\"k\n\x07ModelOp\x12\'\n\x02op\x18\x01 \x01(\x0e\x32\x1b.cozy.scheduler.ModelOpKind\x12\x0b\n\x03ref\x18\x02 \x01(\t\x12*\n\x08snapshot\x18\x03 \x01(\x0b\x32\x18.cozy.scheduler.Snapshot\"\xe1\x01\n\nModelEvent\x12\x0b\n\x03ref\x18\x01 \x01(\t\x12)\n\x05state\x18\x02 \x01(\x0e\x32\x1a.cozy.scheduler.ModelState\x12\x12\n\nvram_bytes\x18\x03 \x01(\x03\x12\r\n\x05\x65rror\x18\x04 \x01(\t\x12\x12\n\nbytes_done\x18\x05 \x01(\x03\x12\x13\n\x0b\x62ytes_total\x18\x06 \x01(\x03\x12\x13\n\x0b\x64uration_ms\x18\x07 \x01(\x03\x12\x12\n\ncache_hits\x18\x08 \x01(\x03\x12\x14\n\x0c\x63\x61\x63he_misses\x18\t \x01(\x03\x12\x10\n\x08warmup_s\x18\n \x01(\x01\"\xaa\x01\n\rFnUnavailable\x12\x15\n\rfunction_name\x18\x01 \x01(\t\x12\x0e\n\x06reason\x18\x02 \x01(\t\x12\x0e\n\x06\x64\x65tail\x18\x03 \x01(\t\x12\x35\n\x04\x61xes\x18\x04 \x03(\x0b\x32\'.cozy.scheduler.FnUnavailable.AxesEntry\x1a+\n\tAxesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"\x8d\x01\n\nFnDegraded\x12\x15\n\rfunction_name\x18\x01 \x01(\t\x12\x0e\n\x06wanted\x18\x02 \x01(\t\x12\x0b\n\x03ran\x18\x03 \x01(\t\x12\x0e\n\x06reason\x18\x04 \x01(\t\x12\x1e\n\x16\x65st_latency_multiplier\x18\x05 \x01(\x01\x12\x1b\n\x13recommended_vram_gb\x18\x06 \x01(\x01\"\x1c\n\x05\x44rain\x12\x13\n\x0b\x64\x65\x61\x64line_ms\x18\x01 \x01(\x03\"6\n\x0cTokenRefresh\x12\r\n\x05token\x18\x01 \x01(\t\x12\x17\n\x0f\x65xpires_at_unix\x18\x02 \x01(\x03*Q\n\x0fProtocolVersion\x12 \n\x1cPROTOCOL_VERSION_UNSPECIFIED\x10\x00\x12\x1c\n\x18PROTOCOL_VERSION_CURRENT\x10\x03*y\n\rResidencyTier\x12\x1e\n\x1aRESIDENCY_TIER_UNSPECIFIED\x10\x00\x12\x17\n\x13RESIDENCY_TIER_DISK\x10\x01\x12\x16\n\x12RESIDENCY_TIER_RAM\x10\x02\x12\x17\n\x13RESIDENCY_TIER_VRAM\x10\x03*\xd8\x01\n\x0bWorkerPhase\x12\x1c\n\x18WORKER_PHASE_UNSPECIFIED\x10\x00\x12\x18\n\x14WORKER_PHASE_BOOTING\x10\x01\x12#\n\x1fWORKER_PHASE_DOWNLOADING_MODELS\x10\x02\x12\"\n\x1eWORKER_PHASE_LOADING_PIPELINES\x10\x03\x12\x18\n\x14WORKER_PHASE_WARMING\x10\x04\x12\x16\n\x12WORKER_PHASE_READY\x10\x05\x12\x16\n\x12WORKER_PHASE_ERROR\x10\x06*V\n\nOutputMode\x12\x1b\n\x17OUTPUT_MODE_UNSPECIFIED\x10\x00\x12\x13\n\x0fOUTPUT_MODE_URL\x10\x01\x12\x16\n\x12OUTPUT_MODE_INLINE\x10\x02*\x9b\x01\n\tJobStatus\x12\x1a\n\x16JOB_STATUS_UNSPECIFIED\x10\x00\x12\x11\n\rJOB_STATUS_OK\x10\x01\x12\x16\n\x12JOB_STATUS_INVALID\x10\x02\x12\x18\n\x14JOB_STATUS_RETRYABLE\x10\x03\x12\x14\n\x10JOB_STATUS_FATAL\x10\x04\x12\x17\n\x13JOB_STATUS_CANCELED\x10\x05*\xa7\x01\n\x0bModelOpKind\x12\x1d\n\x19MODEL_OP_KIND_UNSPECIFIED\x10\x00\x12%\n!MODEL_OP_KIND_ADOPT_COMPILE_CACHE\x10\x04\"\x04\x08\x01\x10\x01\"\x04\x08\x02\x10\x02\"\x04\x08\x03\x10\x03*\x16MODEL_OP_KIND_DOWNLOAD*\x12MODEL_OP_KIND_LOAD*\x14MODEL_OP_KIND_UNLOAD*\xda\x01\n\nModelState\x12\x1b\n\x17MODEL_STATE_UNSPECIFIED\x10\x00\x12\x1b\n\x17MODEL_STATE_DOWNLOADING\x10\x01\x12\x17\n\x13MODEL_STATE_ON_DISK\x10\x02\x12\x16\n\x12MODEL_STATE_IN_RAM\x10\x03\x12\x17\n\x13MODEL_STATE_IN_VRAM\x10\x04\x12\x17\n\x13MODEL_STATE_EVICTED\x10\x05\x12\x16\n\x12MODEL_STATE_FAILED\x10\x06\x12\x17\n\x13MODEL_STATE_ADOPTED\x10\x07\x32\x61\n\x0fWorkerScheduler\x12N\n\x07\x43onnect\x12\x1d.cozy.scheduler.WorkerMessage\x1a .cozy.scheduler.SchedulerMessage(\x01\x30\x01\x42\x61Z_github.com/cozy-creator/tensorhub/internal/orchestrator/grpc/pb/workerscheduler;workerschedulerb\x06proto3') _globals = globals() _builder.BuildMessageAndEnumDescriptors(DESCRIPTOR, _globals) @@ -32,24 +32,26 @@ if not _descriptor._USE_C_DESCRIPTORS: _globals['DESCRIPTOR']._loaded_options = None _globals['DESCRIPTOR']._serialized_options = b'Z_github.com/cozy-creator/tensorhub/internal/orchestrator/grpc/pb/workerscheduler;workerscheduler' + _globals['_DESIREDRESIDENCY_SNAPSHOTSENTRY']._loaded_options = None + _globals['_DESIREDRESIDENCY_SNAPSHOTSENTRY']._serialized_options = b'8\001' _globals['_RUNJOB_SNAPSHOTSENTRY']._loaded_options = None _globals['_RUNJOB_SNAPSHOTSENTRY']._serialized_options = b'8\001' _globals['_FNUNAVAILABLE_AXESENTRY']._loaded_options = None _globals['_FNUNAVAILABLE_AXESENTRY']._serialized_options = b'8\001' - _globals['_PROTOCOLVERSION']._serialized_start=4156 - _globals['_PROTOCOLVERSION']._serialized_end=4237 - _globals['_RESIDENCYTIER']._serialized_start=4239 - _globals['_RESIDENCYTIER']._serialized_end=4360 - _globals['_WORKERPHASE']._serialized_start=4363 - _globals['_WORKERPHASE']._serialized_end=4579 - _globals['_OUTPUTMODE']._serialized_start=4581 - _globals['_OUTPUTMODE']._serialized_end=4667 - _globals['_JOBSTATUS']._serialized_start=4670 - _globals['_JOBSTATUS']._serialized_end=4825 - _globals['_MODELOPKIND']._serialized_start=4828 - _globals['_MODELOPKIND']._serialized_end=4989 - _globals['_MODELSTATE']._serialized_start=4992 - _globals['_MODELSTATE']._serialized_end=5210 + _globals['_PROTOCOLVERSION']._serialized_start=4594 + _globals['_PROTOCOLVERSION']._serialized_end=4675 + _globals['_RESIDENCYTIER']._serialized_start=4677 + _globals['_RESIDENCYTIER']._serialized_end=4798 + _globals['_WORKERPHASE']._serialized_start=4801 + _globals['_WORKERPHASE']._serialized_end=5017 + _globals['_OUTPUTMODE']._serialized_start=5019 + _globals['_OUTPUTMODE']._serialized_end=5105 + _globals['_JOBSTATUS']._serialized_start=5108 + _globals['_JOBSTATUS']._serialized_end=5263 + _globals['_MODELOPKIND']._serialized_start=5266 + _globals['_MODELOPKIND']._serialized_end=5433 + _globals['_MODELSTATE']._serialized_start=5436 + _globals['_MODELSTATE']._serialized_end=5654 _globals['_WORKERMESSAGE']._serialized_start=43 _globals['_WORKERMESSAGE']._serialized_end=470 _globals['_SCHEDULERMESSAGE']._serialized_start=473 @@ -63,49 +65,55 @@ _globals['_INFLIGHTJOB']._serialized_start=1360 _globals['_INFLIGHTJOB']._serialized_end=1410 _globals['_HELLOACK']._serialized_start=1413 - _globals['_HELLOACK']._serialized_end=1573 - _globals['_MODELRESOLUTION']._serialized_start=1575 - _globals['_MODELRESOLUTION']._serialized_end=1641 - _globals['_STATEDELTA']._serialized_start=1644 - _globals['_STATEDELTA']._serialized_end=1806 - _globals['_RUNJOB']._serialized_start=1809 - _globals['_RUNJOB']._serialized_end=2261 - _globals['_RUNJOB_SNAPSHOTSENTRY']._serialized_start=2187 - _globals['_RUNJOB_SNAPSHOTSENTRY']._serialized_end=2261 - _globals['_RESOLVEDCOMPUTE']._serialized_start=2263 - _globals['_RESOLVEDCOMPUTE']._serialized_end=2352 - _globals['_MODELBINDING']._serialized_start=2354 - _globals['_MODELBINDING']._serialized_end=2467 - _globals['_LORAOVERLAY']._serialized_start=2469 - _globals['_LORAOVERLAY']._serialized_end=2539 - _globals['_SNAPSHOT']._serialized_start=2541 - _globals['_SNAPSHOT']._serialized_end=2612 - _globals['_SNAPSHOTFILE']._serialized_start=2614 - _globals['_SNAPSHOTFILE']._serialized_end=2691 - _globals['_JOBACCEPTED']._serialized_start=2693 - _globals['_JOBACCEPTED']._serialized_end=2743 - _globals['_JOBRESULT']._serialized_start=2746 - _globals['_JOBRESULT']._serialized_end=2952 - _globals['_JOBMETRICS']._serialized_start=2955 - _globals['_JOBMETRICS']._serialized_end=3263 - _globals['_JOBPROGRESS']._serialized_start=3265 - _globals['_JOBPROGRESS']._serialized_end=3364 - _globals['_CANCELJOB']._serialized_start=3366 - _globals['_CANCELJOB']._serialized_end=3414 - _globals['_MODELOP']._serialized_start=3416 - _globals['_MODELOP']._serialized_end=3523 - _globals['_MODELEVENT']._serialized_start=3526 - _globals['_MODELEVENT']._serialized_end=3751 - _globals['_FNUNAVAILABLE']._serialized_start=3754 - _globals['_FNUNAVAILABLE']._serialized_end=3924 - _globals['_FNUNAVAILABLE_AXESENTRY']._serialized_start=3881 - _globals['_FNUNAVAILABLE_AXESENTRY']._serialized_end=3924 - _globals['_FNDEGRADED']._serialized_start=3927 - _globals['_FNDEGRADED']._serialized_end=4068 - _globals['_DRAIN']._serialized_start=4070 - _globals['_DRAIN']._serialized_end=4098 - _globals['_TOKENREFRESH']._serialized_start=4100 - _globals['_TOKENREFRESH']._serialized_end=4154 - _globals['_WORKERSCHEDULER']._serialized_start=5212 - _globals['_WORKERSCHEDULER']._serialized_end=5309 + _globals['_HELLOACK']._serialized_end=1634 + _globals['_DESIREDRESIDENCY']._serialized_start=1637 + _globals['_DESIREDRESIDENCY']._serialized_end=1884 + _globals['_DESIREDRESIDENCY_SNAPSHOTSENTRY']._serialized_start=1810 + _globals['_DESIREDRESIDENCY_SNAPSHOTSENTRY']._serialized_end=1884 + _globals['_DESIREDINSTANCE']._serialized_start=1886 + _globals['_DESIREDINSTANCE']._serialized_end=1972 + _globals['_MODELRESOLUTION']._serialized_start=1974 + _globals['_MODELRESOLUTION']._serialized_end=2040 + _globals['_STATEDELTA']._serialized_start=2043 + _globals['_STATEDELTA']._serialized_end=2244 + _globals['_RUNJOB']._serialized_start=2247 + _globals['_RUNJOB']._serialized_end=2699 + _globals['_RUNJOB_SNAPSHOTSENTRY']._serialized_start=1810 + _globals['_RUNJOB_SNAPSHOTSENTRY']._serialized_end=1884 + _globals['_RESOLVEDCOMPUTE']._serialized_start=2701 + _globals['_RESOLVEDCOMPUTE']._serialized_end=2790 + _globals['_MODELBINDING']._serialized_start=2792 + _globals['_MODELBINDING']._serialized_end=2905 + _globals['_LORAOVERLAY']._serialized_start=2907 + _globals['_LORAOVERLAY']._serialized_end=2977 + _globals['_SNAPSHOT']._serialized_start=2979 + _globals['_SNAPSHOT']._serialized_end=3050 + _globals['_SNAPSHOTFILE']._serialized_start=3052 + _globals['_SNAPSHOTFILE']._serialized_end=3129 + _globals['_JOBACCEPTED']._serialized_start=3131 + _globals['_JOBACCEPTED']._serialized_end=3181 + _globals['_JOBRESULT']._serialized_start=3184 + _globals['_JOBRESULT']._serialized_end=3390 + _globals['_JOBMETRICS']._serialized_start=3393 + _globals['_JOBMETRICS']._serialized_end=3701 + _globals['_JOBPROGRESS']._serialized_start=3703 + _globals['_JOBPROGRESS']._serialized_end=3802 + _globals['_CANCELJOB']._serialized_start=3804 + _globals['_CANCELJOB']._serialized_end=3852 + _globals['_MODELOP']._serialized_start=3854 + _globals['_MODELOP']._serialized_end=3961 + _globals['_MODELEVENT']._serialized_start=3964 + _globals['_MODELEVENT']._serialized_end=4189 + _globals['_FNUNAVAILABLE']._serialized_start=4192 + _globals['_FNUNAVAILABLE']._serialized_end=4362 + _globals['_FNUNAVAILABLE_AXESENTRY']._serialized_start=4319 + _globals['_FNUNAVAILABLE_AXESENTRY']._serialized_end=4362 + _globals['_FNDEGRADED']._serialized_start=4365 + _globals['_FNDEGRADED']._serialized_end=4506 + _globals['_DRAIN']._serialized_start=4508 + _globals['_DRAIN']._serialized_end=4536 + _globals['_TOKENREFRESH']._serialized_start=4538 + _globals['_TOKENREFRESH']._serialized_end=4592 + _globals['_WORKERSCHEDULER']._serialized_start=5656 + _globals['_WORKERSCHEDULER']._serialized_end=5753 # @@protoc_insertion_point(module_scope) diff --git a/src/gen_worker/pb/worker_scheduler_pb2.pyi b/src/gen_worker/pb/worker_scheduler_pb2.pyi index de6698c1..906a1471 100644 --- a/src/gen_worker/pb/worker_scheduler_pb2.pyi +++ b/src/gen_worker/pb/worker_scheduler_pb2.pyi @@ -47,9 +47,6 @@ class JobStatus(int, metaclass=_enum_type_wrapper.EnumTypeWrapper): class ModelOpKind(int, metaclass=_enum_type_wrapper.EnumTypeWrapper): __slots__ = () MODEL_OP_KIND_UNSPECIFIED: _ClassVar[ModelOpKind] - MODEL_OP_KIND_DOWNLOAD: _ClassVar[ModelOpKind] - MODEL_OP_KIND_LOAD: _ClassVar[ModelOpKind] - MODEL_OP_KIND_UNLOAD: _ClassVar[ModelOpKind] MODEL_OP_KIND_ADOPT_COMPILE_CACHE: _ClassVar[ModelOpKind] class ModelState(int, metaclass=_enum_type_wrapper.EnumTypeWrapper): @@ -85,9 +82,6 @@ JOB_STATUS_RETRYABLE: JobStatus JOB_STATUS_FATAL: JobStatus JOB_STATUS_CANCELED: JobStatus MODEL_OP_KIND_UNSPECIFIED: ModelOpKind -MODEL_OP_KIND_DOWNLOAD: ModelOpKind -MODEL_OP_KIND_LOAD: ModelOpKind -MODEL_OP_KIND_UNLOAD: ModelOpKind MODEL_OP_KIND_ADOPT_COMPILE_CACHE: ModelOpKind MODEL_STATE_UNSPECIFIED: ModelState MODEL_STATE_DOWNLOADING: ModelState @@ -191,16 +185,45 @@ class InFlightJob(_message.Message): def __init__(self, request_id: _Optional[str] = ..., attempt: _Optional[int] = ...) -> None: ... class HelloAck(_message.Message): - __slots__ = ("protocol_version", "file_base_url", "keep", "resolutions") + __slots__ = ("protocol_version", "file_base_url", "keep", "resolutions", "desired_residency") PROTOCOL_VERSION_FIELD_NUMBER: _ClassVar[int] FILE_BASE_URL_FIELD_NUMBER: _ClassVar[int] KEEP_FIELD_NUMBER: _ClassVar[int] RESOLUTIONS_FIELD_NUMBER: _ClassVar[int] + DESIRED_RESIDENCY_FIELD_NUMBER: _ClassVar[int] protocol_version: ProtocolVersion file_base_url: str keep: _containers.RepeatedScalarFieldContainer[str] resolutions: _containers.RepeatedCompositeFieldContainer[ModelResolution] - def __init__(self, protocol_version: _Optional[_Union[ProtocolVersion, str]] = ..., file_base_url: _Optional[str] = ..., keep: _Optional[_Iterable[str]] = ..., resolutions: _Optional[_Iterable[_Union[ModelResolution, _Mapping]]] = ...) -> None: ... + desired_residency: DesiredResidency + def __init__(self, protocol_version: _Optional[_Union[ProtocolVersion, str]] = ..., file_base_url: _Optional[str] = ..., keep: _Optional[_Iterable[str]] = ..., resolutions: _Optional[_Iterable[_Union[ModelResolution, _Mapping]]] = ..., desired_residency: _Optional[_Union[DesiredResidency, _Mapping]] = ...) -> None: ... + +class DesiredResidency(_message.Message): + __slots__ = ("generation", "disk_refs", "hot", "snapshots") + class SnapshotsEntry(_message.Message): + __slots__ = ("key", "value") + KEY_FIELD_NUMBER: _ClassVar[int] + VALUE_FIELD_NUMBER: _ClassVar[int] + key: str + value: Snapshot + def __init__(self, key: _Optional[str] = ..., value: _Optional[_Union[Snapshot, _Mapping]] = ...) -> None: ... + GENERATION_FIELD_NUMBER: _ClassVar[int] + DISK_REFS_FIELD_NUMBER: _ClassVar[int] + HOT_FIELD_NUMBER: _ClassVar[int] + SNAPSHOTS_FIELD_NUMBER: _ClassVar[int] + generation: int + disk_refs: _containers.RepeatedScalarFieldContainer[str] + hot: _containers.RepeatedCompositeFieldContainer[DesiredInstance] + snapshots: _containers.MessageMap[str, Snapshot] + def __init__(self, generation: _Optional[int] = ..., disk_refs: _Optional[_Iterable[str]] = ..., hot: _Optional[_Iterable[_Union[DesiredInstance, _Mapping]]] = ..., snapshots: _Optional[_Mapping[str, Snapshot]] = ...) -> None: ... + +class DesiredInstance(_message.Message): + __slots__ = ("function_name", "models") + FUNCTION_NAME_FIELD_NUMBER: _ClassVar[int] + MODELS_FIELD_NUMBER: _ClassVar[int] + function_name: str + models: _containers.RepeatedCompositeFieldContainer[ModelBinding] + def __init__(self, function_name: _Optional[str] = ..., models: _Optional[_Iterable[_Union[ModelBinding, _Mapping]]] = ...) -> None: ... class ModelResolution(_message.Message): __slots__ = ("ref", "resolved_ref", "cast") @@ -213,18 +236,20 @@ class ModelResolution(_message.Message): def __init__(self, ref: _Optional[str] = ..., resolved_ref: _Optional[str] = ..., cast: _Optional[str] = ...) -> None: ... class StateDelta(_message.Message): - __slots__ = ("phase", "available_functions", "loading_functions", "free_vram_bytes", "finalizing_jobs") + __slots__ = ("phase", "available_functions", "loading_functions", "free_vram_bytes", "finalizing_jobs", "observed_residency_generation") PHASE_FIELD_NUMBER: _ClassVar[int] AVAILABLE_FUNCTIONS_FIELD_NUMBER: _ClassVar[int] LOADING_FUNCTIONS_FIELD_NUMBER: _ClassVar[int] FREE_VRAM_BYTES_FIELD_NUMBER: _ClassVar[int] FINALIZING_JOBS_FIELD_NUMBER: _ClassVar[int] + OBSERVED_RESIDENCY_GENERATION_FIELD_NUMBER: _ClassVar[int] phase: WorkerPhase available_functions: _containers.RepeatedScalarFieldContainer[str] loading_functions: _containers.RepeatedScalarFieldContainer[str] free_vram_bytes: int finalizing_jobs: int - def __init__(self, phase: _Optional[_Union[WorkerPhase, str]] = ..., available_functions: _Optional[_Iterable[str]] = ..., loading_functions: _Optional[_Iterable[str]] = ..., free_vram_bytes: _Optional[int] = ..., finalizing_jobs: _Optional[int] = ...) -> None: ... + observed_residency_generation: int + def __init__(self, phase: _Optional[_Union[WorkerPhase, str]] = ..., available_functions: _Optional[_Iterable[str]] = ..., loading_functions: _Optional[_Iterable[str]] = ..., free_vram_bytes: _Optional[int] = ..., finalizing_jobs: _Optional[int] = ..., observed_residency_generation: _Optional[int] = ...) -> None: ... class RunJob(_message.Message): __slots__ = ("request_id", "attempt", "function_name", "input_payload", "timeout_ms", "tenant", "invoker_id", "capability_token", "output_mode", "compute", "models", "snapshots") diff --git a/tests/e2e_endpoints.py b/tests/e2e_endpoints.py index fdf94a74..68259579 100644 --- a/tests/e2e_endpoints.py +++ b/tests/e2e_endpoints.py @@ -80,7 +80,7 @@ def finalize_peer(self, ctx: RequestContext, data: EchoIn) -> EchoOut: @endpoint(model=Hub("e2e/tiny")) class ModelBoundEndpoint: - """Tensorhub-bound endpoint: only becomes available after ModelOp LOAD.""" + """Tensorhub-bound endpoint: available after its desired instance warms.""" def setup(self, model: str) -> None: self.model_path = model @@ -98,7 +98,7 @@ def model_echo(self, ctx: RequestContext, data: EchoIn) -> EchoOut: @endpoint(model=Hub("e2e/broken")) class BrokenSetupEndpoint: - """Setup raises while BREAK_SETUP is set; a later LOAD retry recovers.""" + """Setup raises while BREAK_SETUP is set; desired-state retry recovers.""" def setup(self, model: str) -> None: if BREAK_SETUP.is_set(): diff --git a/tests/test_declarative_residency.py b/tests/test_declarative_residency.py new file mode 100644 index 00000000..b0a1cec9 --- /dev/null +++ b/tests/test_declarative_residency.py @@ -0,0 +1,256 @@ +"""Declarative model residency: full replace, exact hot picks, and observation.""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace +from typing import Any + +import msgspec + +from gen_worker.api.binding import Civitai, wire_ref +from gen_worker.api.slot import Slot +from gen_worker.executor import Executor +from gen_worker.families.base import FamilyDefaults, family +from gen_worker.lifecycle import Lifecycle +from gen_worker.pb import worker_scheduler_pb2 as pb +from gen_worker.registry import EndpointSpec + + +@family("declarative-residency-test") +class _Defaults(FamilyDefaults): + steps: int = 1 + + +class _Pipeline: + pass + + +class _Input(msgspec.Struct): + prompt: str = "" + + +def _spec() -> EndpointSpec: + default = Civitai("827184", version="2883731") + + class Endpoint: + def setup(self, pipeline: str) -> None: # pragma: no cover - replaced in tests + self.pipeline = pipeline + + def generate(self, ctx: Any, payload: _Input) -> dict: # pragma: no cover + return {} + + return EndpointSpec( + name="generate", + method=Endpoint.generate, + kind="inference", + payload_type=_Input, + output_mode="single", + cls=Endpoint, + attr_name="generate", + models={"pipeline": default}, + slots={ + "pipeline": Slot( + _Pipeline, + default_checkpoint=default, + default_config=_Defaults(), + ) + }, + slot_family={"pipeline": "declarative-residency-test"}, + ) + + +async def _noop_send(msg: pb.WorkerMessage) -> None: + pass + + +class _Transport: + def __init__(self) -> None: + self.connected = True + self.sent: list[pb.WorkerMessage] = [] + self.queue = SimpleNamespace(pending_result_keys=set()) + + async def send(self, msg: pb.WorkerMessage) -> None: + self.sent.append(msg) + + +def _lifecycle() -> tuple[Lifecycle, Executor, _Transport]: + executor = Executor([_spec()], _noop_send) + lifecycle = Lifecycle( + SimpleNamespace(worker_jwt="", worker_id="worker", runpod_pod_id=""), + executor, + ) + transport = _Transport() + lifecycle.transport = transport # type: ignore[assignment] + return lifecycle, executor, transport + + +def _snapshot(url: str) -> pb.Snapshot: + return pb.Snapshot( + digest="ab" * 32, + files=[pb.SnapshotFile( + path="model.safetensors", + size_bytes=1, + blake3="cd" * 32, + url=url, + )], + ) + + +def test_proto_field_numbers_match_tensorhub_contract() -> None: + assert pb.PROTOCOL_VERSION_CURRENT == 3 + assert pb.HelloAck.DESIRED_RESIDENCY_FIELD_NUMBER == 5 + assert pb.DesiredResidency.GENERATION_FIELD_NUMBER == 1 + assert pb.DesiredResidency.DISK_REFS_FIELD_NUMBER == 2 + assert pb.DesiredResidency.HOT_FIELD_NUMBER == 3 + assert pb.DesiredResidency.SNAPSHOTS_FIELD_NUMBER == 4 + assert pb.DesiredInstance.FUNCTION_NAME_FIELD_NUMBER == 1 + assert pb.DesiredInstance.MODELS_FIELD_NUMBER == 2 + assert pb.StateDelta.OBSERVED_RESIDENCY_GENERATION_FIELD_NUMBER == 6 + assert not hasattr(pb, "MODEL_OP_KIND_DOWNLOAD") + assert not hasattr(pb, "MODEL_OP_KIND_LOAD") + assert not hasattr(pb, "MODEL_OP_KIND_UNLOAD") + + +def test_full_replace_supersedes_and_same_generation_refreshes_urls(monkeypatch) -> None: + async def run() -> None: + lifecycle, executor, transport = _lifecycle() + calls: list[tuple[str, str]] = [] + old_started = asyncio.Event() + refreshed = asyncio.Event() + + async def ensure_local(ref: str, snapshot=None, *, binding=None): + url = snapshot.files[0].url if snapshot and snapshot.files else "" + calls.append((ref, url)) + if ref == "acme/old": + old_started.set() + await asyncio.Event().wait() + if url == "https://r2/new-2": + refreshed.set() + + monkeypatch.setattr(executor.store, "ensure_local", ensure_local) + + # Tenant work owns the lane: reconciliation does not start while the + # executor is busy. + executor._idle.clear() + await lifecycle.on_hello_ack(pb.HelloAck(desired_residency=pb.DesiredResidency( + generation=1, + disk_refs=["acme/old"], + snapshots={"acme/old": _snapshot("https://r2/old")}, + ))) + await asyncio.sleep(0) + assert calls == [] + assert transport.sent[-1].state_delta.observed_residency_generation == 1 + + executor._idle.set() + await old_started.wait() + + # A newer full replacement cancels the obsolete blocked reconcile. + await lifecycle.on_hello_ack(pb.HelloAck(desired_residency=pb.DesiredResidency( + generation=2, + disk_refs=["acme/new"], + snapshots={"acme/new": _snapshot("https://r2/new-1")}, + ))) + if lifecycle._residency_task is not None: + await lifecycle._residency_task + + # Same generation is not stale: it may carry refreshed presigned URLs. + await lifecycle.on_hello_ack(pb.HelloAck(desired_residency=pb.DesiredResidency( + generation=2, + disk_refs=["acme/new"], + snapshots={"acme/new": _snapshot("https://r2/new-2")}, + ))) + await refreshed.wait() + + # An older generation cannot replace keep state or regress observation. + await lifecycle.on_hello_ack(pb.HelloAck(desired_residency=pb.DesiredResidency( + generation=1, + disk_refs=["acme/stale"], + snapshots={"acme/stale": _snapshot("https://r2/stale")}, + ))) + assert executor.store.keep == {"acme/new"} + assert lifecycle._state_delta().observed_residency_generation == 2 + assert calls == [ + ("acme/old", "https://r2/old"), + ("acme/new", "https://r2/new-1"), + ("acme/new", "https://r2/new-2"), + ] + + asyncio.run(run()) + + +def test_hot_instance_uses_exact_dynamic_slot_binding(monkeypatch) -> None: + async def run() -> None: + executor = Executor([_spec()], _noop_send) + base = executor.specs["generate"] + captured: list[tuple[EndpointSpec, dict[str, pb.Snapshot]]] = [] + + async def ensure_setup(spec, snapshots=None, promote_slots=None): + captured.append((spec, snapshots or {})) + + monkeypatch.setattr(executor, "ensure_setup", ensure_setup) + picked = "tensorhub/cyberrealistic-pony:prod" + snapshot = _snapshot("https://r2/picked") + await executor.ensure_desired_instance( + pb.DesiredInstance( + function_name="generate", + models=[pb.ModelBinding(slot="pipeline", ref=picked)], + ), + {picked: snapshot}, + ) + + assert len(captured) == 1 + effective, snapshots = captured[0] + assert effective.instance_key != base.instance_key + assert wire_ref(effective.models["pipeline"]) == picked + assert snapshots[picked].files[0].url == "https://r2/picked" + + asyncio.run(run()) + + +def test_run_job_preempts_then_resumes_current_desired_state(monkeypatch) -> None: + async def run() -> None: + lifecycle, executor, _ = _lifecycle() + started = asyncio.Event() + canceled = asyncio.Event() + resumed = asyncio.Event() + calls = 0 + + async def ensure_local(ref: str, snapshot=None, *, binding=None): + nonlocal calls + calls += 1 + if calls == 1: + started.set() + try: + await asyncio.Event().wait() + except asyncio.CancelledError: + canceled.set() + raise + resumed.set() + + async def handle_run_job(run_job: pb.RunJob) -> None: + executor._idle.clear() + await asyncio.sleep(0) + assert canceled.is_set(), "background work must cancel before request setup" + + monkeypatch.setattr(executor.store, "ensure_local", ensure_local) + monkeypatch.setattr(executor, "handle_run_job", handle_run_job) + + await lifecycle.on_hello_ack(pb.HelloAck(desired_residency=pb.DesiredResidency( + generation=7, + disk_refs=["acme/background"], + snapshots={"acme/background": _snapshot("https://r2/background")}, + ))) + await started.wait() + + await lifecycle.on_message(pb.SchedulerMessage(run_job=pb.RunJob( + request_id="request", attempt=1, function_name="generate", + ))) + await asyncio.sleep(0) + assert calls == 1, "background work must stay paused while a job is active" + + executor._idle.set() + await resumed.wait() + assert calls == 2, "current desired state must resume after the job becomes idle" + + asyncio.run(run()) diff --git a/tests/test_dynamic_slot_materialization_pgw532.py b/tests/test_dynamic_slot_materialization_pgw532.py index 644a8d97..695fc3d7 100644 --- a/tests/test_dynamic_slot_materialization_pgw532.py +++ b/tests/test_dynamic_slot_materialization_pgw532.py @@ -15,9 +15,8 @@ 2. dispatch pick != default: the executor loads the PICKED checkpoint and ``ctx.slots[name].ref`` reflects the pick, not the code default. 3. instance-per-pick residency: two picks = two resident instances (one - ``setup()`` per (class, resolved pick)); UNLOAD evicts one whole - instance while the other stays warm; a later LOAD for the evicted pick - matches its derived record and re-sets it up. + ``setup()`` per (class, resolved pick)); a repeated pick reuses its + existing instance. 4. default-only path with a CAS default_checkpoint: unchanged (one instance, base record, ctx.slots ref = declared default). 5. an unusable (non-CAS) resolved_models stamp with a raw upstream @@ -245,12 +244,11 @@ async def _run() -> pb.JobResult: # --------------------------------------------------------------------------- -# 3. instance-per-pick: two picks = two instances; UNLOAD evicts one whole -# instance; LOAD re-materializes the evicted pick's derived record +# 3. instance-per-pick: two picks = two instances; repeated picks reuse them # --------------------------------------------------------------------------- -def test_two_picks_two_instances_then_lru_evicts_whole_instance(tmp_path, monkeypatch) -> None: +def test_two_picks_create_two_reusable_instances(tmp_path, monkeypatch) -> None: setup_calls: List[Tuple[str, str]] = [] ex, sent, downloads = _harness(tmp_path, monkeypatch, setup_calls) @@ -284,17 +282,6 @@ async def _run() -> None: assert r3.status == pb.JOB_STATUS_OK, r3.safe_message assert len(setup_calls) == 2, "resident pick must not re-run setup()" - # Hub-directed UNLOAD of pick A vacates THAT instance only. - await ex.handle_model_op(pb.ModelOp(op=pb.MODEL_OP_KIND_UNLOAD, ref=PICK_A)) - ready = _ready_records(ex) - assert len(ready) == 1, "UNLOAD must evict exactly one pick's instance" - held = {r for rec in ready for r in rec.held_refs} - assert PICK_B in held and PICK_A not in held - - # LOAD for the evicted pick matches its derived record and re-sets up. - await ex.handle_model_op(pb.ModelOp( - op=pb.MODEL_OP_KIND_LOAD, ref=PICK_A, snapshot=_snapshot("aa" * 32))) - assert len(setup_calls) == 3 assert len(_ready_records(ex)) == 2 asyncio.run(_run()) @@ -371,28 +358,6 @@ async def _run() -> pb.JobResult: assert "generate" not in ex.unavailable, "a per-request refusal must not disable the function" -# --------------------------------------------------------------------------- -# LOAD for a never-dispatched pick: bytes banked, typed failure, no setup -# --------------------------------------------------------------------------- - - -def test_load_of_unknown_pick_banks_bytes_and_fails_typed(tmp_path, monkeypatch) -> None: - setup_calls: List[Tuple[str, str]] = [] - ex, sent, downloads = _harness(tmp_path, monkeypatch, setup_calls) - - asyncio.run(ex.handle_model_op(pb.ModelOp( - op=pb.MODEL_OP_KIND_LOAD, ref=PICK_A, snapshot=_snapshot("aa" * 32)))) - - # Pre-warm degraded to a download: bytes + snapshot banked for the next - # dispatch, no instance guessed into existence, typed failure reported. - assert {d["ref"] for d in downloads} == {PICK_A} - assert ex.store.has_snapshot(PICK_A) - assert setup_calls == [] - failed = [m.model_event for m in sent if m.WhichOneof("msg") == "model_event" - and m.model_event.state == pb.MODEL_STATE_FAILED] - assert [(e.ref, e.error) for e in failed] == [(PICK_A, "load_failed")] - - # --------------------------------------------------------------------------- # 6. hub-less path: default_checkpoint raw source resolves upstream, unchanged # --------------------------------------------------------------------------- diff --git a/tests/test_executor_model_ops.py b/tests/test_executor_model_ops.py deleted file mode 100644 index 597bb3ee..00000000 --- a/tests/test_executor_model_ops.py +++ /dev/null @@ -1,78 +0,0 @@ -"""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" diff --git a/tests/test_executor_prefetch_binding.py b/tests/test_executor_prefetch_binding.py index e2621419..6cc80847 100644 --- a/tests/test_executor_prefetch_binding.py +++ b/tests/test_executor_prefetch_binding.py @@ -1,8 +1,4 @@ -"""#377: binding files=/provider must apply on EVERY download path. - -The ModelOp/startup-prefetch paths only carry the bare ref; the executor -registers each endpoint binding by wire ref so ModelStore.ensure_local can -recover files/provider at the download-layer boundary (no network).""" +"""#377: binding files/provider apply to bare-ref residency downloads.""" from __future__ import annotations @@ -65,23 +61,8 @@ async def _fake_ensure_local(ref: str, **kwargs: Any) -> Path: return calls -def test_model_op_download_applies_binding_files(monkeypatch, tmp_path) -> None: - calls = _capture_download(monkeypatch, tmp_path) - - async def _run() -> None: - ex = Executor([_spec()], _noop_send) - await ex.handle_model_op(pb.ModelOp(op=pb.MODEL_OP_KIND_DOWNLOAD, ref=_REF)) - - asyncio.run(_run()) - assert len(calls) == 1 - assert calls[0]["ref"] == _REF - assert calls[0]["provider"] == "huggingface" - assert calls[0]["allow_patterns"] == _BINDING.files - - def test_bare_ref_store_ensure_local_applies_binding_files(monkeypatch, tmp_path) -> None: - # The lifecycle startup-prefetch shape: store.ensure_local(ref) with - # nothing else. + # DesiredResidency carries a ref and snapshot, not the declared binding. calls = _capture_download(monkeypatch, tmp_path) async def _run() -> None: diff --git a/tests/test_executor_residency.py b/tests/test_executor_residency.py index 746a8fce..5d0e55b4 100644 --- a/tests/test_executor_residency.py +++ b/tests/test_executor_residency.py @@ -197,76 +197,6 @@ async def _run() -> None: 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, snapshots) -> None: - from diffusers import DDPMPipeline - - class Endpoint: - def setup(self, m: DDPMPipeline) -> 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 = [] - refs = {"acme/unet": snapshots["big-a"]} - - async def _run() -> None: - ex = _executor([spec], tmp_path, 4 * _GiB, sent, refs) - inst = await ex.ensure_setup(spec) - base = torch.cuda.memory_allocated() - 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 _device(inst.m) == "cpu" # actually moved, not just booked - assert torch.cuda.memory_allocated() <= base - 100 * _MiB # VRAM freed - 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 _device(inst.m) == "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, snapshots) -> None: - class Endpoint: - def setup(self, model: str) -> None: - self.buf = torch.zeros(64 * _MiB, dtype=torch.uint8, device="cuda") - - def shutdown(self) -> None: - self.buf = None - - def run(self, ctx, payload: _In): # pragma: no cover - return payload - - spec = _spec("ep", Endpoint, {"model": HF("acme/opaque")}) - sent: list = [] - refs = {"acme/opaque": snapshots["small"]} - - async def _run() -> None: - base = torch.cuda.memory_allocated() - ex = _executor([spec], tmp_path, 4 * _GiB, sent, refs) - await ex.ensure_setup(spec) - assert torch.cuda.memory_allocated() >= base + 64 * _MiB - 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 torch.cuda.memory_allocated() == base # shutdown really freed it - - asyncio.run(_run()) - assert _events(sent, pb.MODEL_STATE_ON_DISK)[-1] == "acme/opaque" - - def test_alternating_endpoints_swap_via_demote_promote(tmp_path, snapshots) -> None: """Budget fits one ~148MB pipeline (+2GiB make_room margin): alternating setups must demote/promote, never tear down or degrade to offload.""" diff --git a/tests/test_model_op_prefetch_no_cascade.py b/tests/test_model_op_prefetch_no_cascade.py deleted file mode 100644 index a0fbe66d..00000000 --- a/tests/test_model_op_prefetch_no_cascade.py +++ /dev/null @@ -1,252 +0,0 @@ -"""gw#465: boot-prefetch/preposition ModelOps must never cascade or conflate. - -The J23 failure: a snapshot-less ModelOp{LOAD} for a SHARED companion ref (one -vae bound to every variant of an SDXL family) cold-set-up the first non-ready -variant spec; its checkpoint slot had no snapshot, the worker can't resolve -tensorhub refs itself, and the deterministic local miss burned the 1s+4s retry -loop before being mislabeled ``download_failed`` (paired with the companion's -``load_failed``). Demand ops for the same refs — which carry snapshots — -succeeded seconds later. -""" - -from __future__ import annotations - -import asyncio -import time -from pathlib import Path -from typing import Any, Dict, List, Optional - -import msgspec -import pytest - -from gen_worker.api.binding import Hub, ModelRef -from gen_worker.executor import Executor -from gen_worker.models.errors import MissingSnapshotError -from gen_worker.pb import worker_scheduler_pb2 as pb -from gen_worker.registry import EndpointSpec - - -class _In(msgspec.Struct): - x: str - - -class _Out(msgspec.Struct): - y: str - - -VAE = Hub("tensorhub/sdxl-vae-fp16-fix") -VARIANT_A = Hub("acme/variant-a") -VARIANT_B = Hub("acme/variant-b") - - -def _spec(name: str, checkpoint: ModelRef, calls: List[str]) -> EndpointSpec: - class Endpoint: - def setup(self, model: str, vae: str) -> None: - calls.append(name) - - def run(self, ctx, payload: _In) -> _Out: # pragma: no cover - return _Out(y=payload.x) - - return EndpointSpec( - name=name, method=Endpoint.run, kind="inference", - payload_type=_In, output_mode="single", cls=Endpoint, - attr_name="run", models={"model": checkpoint, "vae": VAE}, - ) - - -def _snapshot(digest: str = "ab" * 32) -> pb.Snapshot: - return pb.Snapshot(digest=digest, files=[pb.SnapshotFile( - path="model.safetensors", size_bytes=5, blake3="cd" * 32, - url="http://r2.invalid/presigned")]) - - -def _harness(tmp_path: Path, calls: List[str], *, on_disk: List[str] = ()): - """Executor with the REAL ModelStore/handle_model_op orchestration; only - the network-touching download primitive is faked (existing test pattern). - """ - sent: List[pb.WorkerMessage] = [] - - async def _send(msg: pb.WorkerMessage) -> None: - sent.append(msg) - - ex = Executor([_spec("gen-a", VARIANT_A, calls), _spec("gen-b", VARIANT_B, calls)], _send) - - downloads: List[Dict[str, Any]] = [] - - async def _fake_download(ref: str, **kwargs: Any) -> Path: - downloads.append({"ref": ref, **kwargs}) - if kwargs.get("snapshot") is None: - raise MissingSnapshotError(f"tensorhub ref {ref!r} needs a snapshot") - p = tmp_path / ref.replace("/", "_") - p.mkdir(parents=True, exist_ok=True) - return p - - import gen_worker.executor as ex_mod - for ref in on_disk: - p = tmp_path / ref.replace("/", "_") - p.mkdir(parents=True, exist_ok=True) - ex.store.residency.track_disk(ref, p) - ex.store._verified.add(ref) - return ex, sent, downloads, ex_mod, _fake_download - - -def _failed_events(sent: List[pb.WorkerMessage]) -> List[pb.ModelEvent]: - return [m.model_event for m in sent - if m.WhichOneof("msg") == "model_event" - and m.model_event.state == pb.MODEL_STATE_FAILED] - - -def test_load_of_shared_companion_does_not_cascade(tmp_path, monkeypatch) -> None: - """A LOAD for the shared vae while instance A is ready must touch/promote - A only — NOT cold-set-up variant B (whose checkpoint has no snapshot).""" - calls: List[str] = [] - ex, sent, downloads, ex_mod, fake = _harness( - tmp_path, calls, on_disk=["acme/variant-a", "tensorhub/sdxl-vae-fp16-fix"]) - monkeypatch.setattr(ex_mod, "ensure_local", fake) - - async def _run() -> None: - spec_a = ex.specs["gen-a"] - await ex.ensure_setup(spec_a) # instance A resident - assert calls == ["gen-a"] - await ex.handle_model_op(pb.ModelOp( - op=pb.MODEL_OP_KIND_LOAD, ref="tensorhub/sdxl-vae-fp16-fix")) - - asyncio.run(_run()) - assert calls == ["gen-a"], f"LOAD(vae) cascaded into cold setups: {calls}" - assert not _failed_events(sent), ( - f"LOAD(vae) with a ready owner emitted failures: {_failed_events(sent)}") - - -def test_cold_load_missing_sibling_snapshot_fails_fast_and_typed(tmp_path, monkeypatch) -> None: - """Cold worker, LOAD(vae): every candidate spec needs a checkpoint the - worker cannot materialize. The op must fail FAST (no 1s+4s local retry - burn) with ``missing_snapshot`` naming the blocking refs — never a phantom - ``download_failed`` — and must not disable any function.""" - calls: List[str] = [] - ex, sent, downloads, ex_mod, fake = _harness( - tmp_path, calls, on_disk=["tensorhub/sdxl-vae-fp16-fix"]) - monkeypatch.setattr(ex_mod, "ensure_local", fake) - - t0 = time.monotonic() - asyncio.run(ex.handle_model_op(pb.ModelOp( - op=pb.MODEL_OP_KIND_LOAD, ref="tensorhub/sdxl-vae-fp16-fix"))) - elapsed = time.monotonic() - t0 - - assert elapsed < 1.0, f"deterministic local miss burned retries ({elapsed:.1f}s)" - assert calls == [], f"unmaterializable specs were cold-set-up: {calls}" - failed = _failed_events(sent) - errors = {(e.ref, e.error) for e in failed} - assert ("acme/variant-a", "missing_snapshot") in errors - assert ("acme/variant-b", "missing_snapshot") in errors - assert ("tensorhub/sdxl-vae-fp16-fix", "missing_snapshot") in errors - assert all(e.error != "download_failed" for e in failed), errors - assert not ex.unavailable, f"functions disabled by a transient miss: {ex.unavailable}" - - -def test_snapshot_memory_enables_companion_setup(tmp_path, monkeypatch) -> None: - """The store remembers digest-carrying snapshots (gw#465): after the hub's - DOWNLOAD(variant-b, snapshot), a snapshot-less LOAD(vae) can set up the - variant-b spec — the boot-prefetch batch lands instead of wedging.""" - calls: List[str] = [] - ex, sent, downloads, ex_mod, fake = _harness( - tmp_path, calls, on_disk=["tensorhub/sdxl-vae-fp16-fix"]) - monkeypatch.setattr(ex_mod, "ensure_local", fake) - - async def _run() -> None: - await ex.handle_model_op(pb.ModelOp( - op=pb.MODEL_OP_KIND_DOWNLOAD, ref="acme/variant-b", - snapshot=_snapshot())) - assert ex.store.has_snapshot("acme/variant-b") - await ex.handle_model_op(pb.ModelOp( - op=pb.MODEL_OP_KIND_LOAD, ref="tensorhub/sdxl-vae-fp16-fix")) - - asyncio.run(_run()) - assert calls == ["gen-b"], f"companion LOAD did not set up variant-b: {calls}" - assert not _failed_events(sent), _failed_events(sent) - - -def test_store_missing_snapshot_waits_then_raises_typed(tmp_path, monkeypatch) -> None: - """ModelStore.ensure_local on a snapshot-less tensorhub ref (th#763): - emits FAILED ``missing_snapshot`` immediately (the hub's re-mint - trigger), then BLOCKS for the re-minted snapshot; when nothing arrives - it raises the typed error — no DOWNLOADING ghost, never a phantom - ``download_failed``.""" - import gen_worker.executor as ex_mod - from gen_worker.executor import ModelStore - - monkeypatch.setattr(ex_mod, "_MISSING_SNAPSHOT_WAIT_S", 0.2) - sent: List[pb.WorkerMessage] = [] - - async def _send(msg: pb.WorkerMessage) -> None: - sent.append(msg) - - store = ModelStore(_send, cache_dir=tmp_path) - # endpoint refs are registered at boot (#377) — that's what makes the - # provider classification confident enough to classify the miss - store.register_binding("acme/never-resolved", Hub("acme/never-resolved")) - - async def _run() -> None: - with pytest.raises(MissingSnapshotError): - await store.ensure_local("acme/never-resolved") - - t0 = time.monotonic() - asyncio.run(_run()) - elapsed = time.monotonic() - t0 - assert 0.2 <= elapsed < 2.0, f"must block for the re-mint window, took {elapsed:.2f}s" - states = [m.model_event.state for m in sent if m.WhichOneof("msg") == "model_event"] - assert pb.MODEL_STATE_DOWNLOADING not in states - failed = _failed_events(sent) - assert failed and failed[-1].error == "missing_snapshot" - - -def test_store_cold_ref_blocks_until_remint_then_serves(tmp_path, monkeypatch) -> None: - """th#763 block-and-serve: a cold ensure_local reports missing_snapshot - and WAITS; the hub's re-minted DOWNLOAD (a concurrent ensure_local WITH - the snapshot) wakes it and the original call returns the materialized - path — the first user request per unseen ref succeeds instead of - fataling as the sacrificial cache warmer.""" - import gen_worker.executor as ex_mod - from gen_worker.executor import ModelStore - - monkeypatch.setattr(ex_mod, "_MISSING_SNAPSHOT_WAIT_S", 10.0) - sent: List[pb.WorkerMessage] = [] - - async def _send(msg: pb.WorkerMessage) -> None: - sent.append(msg) - - async def _fake_download(ref: str, **kwargs: Any) -> Path: - assert kwargs.get("snapshot") is not None, "download must carry the re-minted snapshot" - p = tmp_path / ref.replace("/", "_") - p.mkdir(parents=True, exist_ok=True) - return p - - monkeypatch.setattr(ex_mod, "ensure_local", _fake_download) - store = ModelStore(_send, cache_dir=tmp_path) - store.register_binding("acme/cold-ref", Hub("acme/cold-ref")) - - async def _run() -> None: - cold = asyncio.create_task(store.ensure_local("acme/cold-ref")) - # Hub reacts to the FAILED(missing_snapshot) event with a DOWNLOAD op. - for _ in range(100): - if _failed_events(sent): - break - await asyncio.sleep(0.01) - assert _failed_events(sent) and _failed_events(sent)[-1].error == "missing_snapshot" - assert not cold.done(), "cold caller must still be waiting on the re-mint" - remint = asyncio.create_task(store.ensure_local("acme/cold-ref", _snapshot())) - path = await asyncio.wait_for(cold, 5.0) - assert path == tmp_path / "acme_cold-ref" - assert await asyncio.wait_for(remint, 5.0) == path - - asyncio.run(_run()) - - -def test_missing_snapshot_maps_retryable_never_fatal() -> None: - """th#763: a job that hits MissingSnapshotError must come back - JOB_STATUS_RETRYABLE — a cold worker mid-resolution must never fatal a - user request.""" - from gen_worker.executor import _map_exception - - status, msg = _map_exception(MissingSnapshotError("tensorhub ref needs a snapshot")) - assert status == pb.JOB_STATUS_RETRYABLE - assert "snapshot" in msg diff --git a/tests/test_snapshot_remint.py b/tests/test_snapshot_remint.py new file mode 100644 index 00000000..7b4d1dae --- /dev/null +++ b/tests/test_snapshot_remint.py @@ -0,0 +1,100 @@ +"""Snapshot-bank and re-mint behavior for cold Tensorhub refs.""" + +from __future__ import annotations + +import asyncio +import time +from pathlib import Path +from typing import Any, List + +import pytest + +from gen_worker.api.binding import Hub +from gen_worker.models.errors import MissingSnapshotError +from gen_worker.pb import worker_scheduler_pb2 as pb + + +def _snapshot(digest: str = "ab" * 32) -> pb.Snapshot: + return pb.Snapshot(digest=digest, files=[pb.SnapshotFile( + path="model.safetensors", size_bytes=5, blake3="cd" * 32, + url="http://r2.invalid/presigned")]) + + +def _failed_events(sent: List[pb.WorkerMessage]) -> List[pb.ModelEvent]: + return [m.model_event for m in sent + if m.WhichOneof("msg") == "model_event" + and m.model_event.state == pb.MODEL_STATE_FAILED] + + +def test_store_missing_snapshot_waits_then_raises_typed(tmp_path, monkeypatch) -> None: + """A snapshot-less Tensorhub ref waits for a refresh, then fails typed.""" + import gen_worker.executor as ex_mod + from gen_worker.executor import ModelStore + + monkeypatch.setattr(ex_mod, "_MISSING_SNAPSHOT_WAIT_S", 0.2) + sent: List[pb.WorkerMessage] = [] + + async def _send(msg: pb.WorkerMessage) -> None: + sent.append(msg) + + store = ModelStore(_send, cache_dir=tmp_path) + store.register_binding("acme/never-resolved", Hub("acme/never-resolved")) + + async def _run() -> None: + with pytest.raises(MissingSnapshotError): + await store.ensure_local("acme/never-resolved") + + t0 = time.monotonic() + asyncio.run(_run()) + elapsed = time.monotonic() - t0 + assert 0.2 <= elapsed < 2.0, f"must block for the refresh window, took {elapsed:.2f}s" + states = [m.model_event.state for m in sent if m.WhichOneof("msg") == "model_event"] + assert pb.MODEL_STATE_DOWNLOADING not in states + failed = _failed_events(sent) + assert failed and failed[-1].error == "missing_snapshot" + + +def test_store_cold_ref_blocks_until_remint_then_serves(tmp_path, monkeypatch) -> None: + """A refreshed desired/job snapshot wakes and serves the cold caller.""" + import gen_worker.executor as ex_mod + from gen_worker.executor import ModelStore + + monkeypatch.setattr(ex_mod, "_MISSING_SNAPSHOT_WAIT_S", 10.0) + sent: List[pb.WorkerMessage] = [] + + async def _send(msg: pb.WorkerMessage) -> None: + sent.append(msg) + + async def _fake_download(ref: str, **kwargs: Any) -> Path: + assert kwargs.get("snapshot") is not None, "download must carry the refreshed snapshot" + p = tmp_path / ref.replace("/", "_") + p.mkdir(parents=True, exist_ok=True) + return p + + monkeypatch.setattr(ex_mod, "ensure_local", _fake_download) + store = ModelStore(_send, cache_dir=tmp_path) + store.register_binding("acme/cold-ref", Hub("acme/cold-ref")) + + async def _run() -> None: + cold = asyncio.create_task(store.ensure_local("acme/cold-ref")) + for _ in range(100): + if _failed_events(sent): + break + await asyncio.sleep(0.01) + assert _failed_events(sent) and _failed_events(sent)[-1].error == "missing_snapshot" + assert not cold.done(), "cold caller must still be waiting on the refresh" + refresh = asyncio.create_task(store.ensure_local("acme/cold-ref", _snapshot())) + path = await asyncio.wait_for(cold, 5.0) + assert path == tmp_path / "acme_cold-ref" + assert await asyncio.wait_for(refresh, 5.0) == path + + asyncio.run(_run()) + + +def test_missing_snapshot_maps_retryable_never_fatal() -> None: + """A cold worker mid-resolution asks the scheduler to retry the job.""" + from gen_worker.executor import _map_exception + + status, msg = _map_exception(MissingSnapshotError("tensorhub ref needs a snapshot")) + assert status == pb.JOB_STATUS_RETRYABLE + assert "snapshot" in msg diff --git a/tests/test_startup_hub_snapshot_wait.py b/tests/test_startup_hub_snapshot_wait.py index 4386d9d6..39bcf06e 100644 --- a/tests/test_startup_hub_snapshot_wait.py +++ b/tests/test_startup_hub_snapshot_wait.py @@ -69,7 +69,10 @@ def test_startup_logs_functions_waiting_on_hub_snapshots(caplog, monkeypatch) -> # (c) the wait is LOGGED, naming the function and the hub ref msgs = [r.getMessage() for r in caplog.records if r.levelno >= logging.WARNING] - assert any("generate" in m and "tensorhub/z-image" in m and "ModelOp" in m for m in msgs), msgs + assert any( + "generate" in m and "tensorhub/z-image" in m and "DesiredResidency" in m + for m in msgs + ), msgs def test_startup_does_not_log_when_nothing_waits_on_hub(caplog, tmp_path, monkeypatch) -> None: diff --git a/tests/test_worker_grpc_e2e.py b/tests/test_worker_grpc_e2e.py index 3b69d6a4..f63fb32a 100644 --- a/tests/test_worker_grpc_e2e.py +++ b/tests/test_worker_grpc_e2e.py @@ -749,8 +749,8 @@ async def _run() -> None: # --------------------------------------------------------------------------- -# ModelOp DOWNLOAD / LOAD / UNLOAD round-trip (#366): a tensorhub-bound -# function is gated until LOAD; ModelEvents follow the contract state machine. +# Desired-residency round-trip: a Tensorhub-bound function is gated until its +# exact instance is warm; ModelEvents follow the contract state machine. # --------------------------------------------------------------------------- @@ -762,7 +762,7 @@ def _is_model_event(ref: str, state: int): ) -def test_model_op_download_load_unload_round_trip(tmp_path, monkeypatch) -> None: +def test_desired_residency_downloads_and_warms_round_trip(tmp_path, monkeypatch) -> None: import http.server from blake3 import blake3 @@ -812,20 +812,27 @@ def log_message(self, *a): assert "model-echo" not in ready.state_delta.available_functions assert "model-echo" in ready.state_delta.loading_functions - # DOWNLOAD -> DOWNLOADING then ON_DISK. - conn.send(model_op=pb.ModelOp( - op=pb.MODEL_OP_KIND_DOWNLOAD, ref="e2e/tiny", snapshot=snapshot)) + # One full desired state downloads the ref, warms the exact endpoint + # instance, and reports the accepted generation. + conn.send(hello_ack=pb.HelloAck( + protocol_version=pb.PROTOCOL_VERSION_CURRENT, + desired_residency=pb.DesiredResidency( + generation=1, + disk_refs=["e2e/tiny"], + hot=[pb.DesiredInstance( + function_name="model-echo", + models=[pb.ModelBinding(slot="model", ref="e2e/tiny")], + )], + snapshots={"e2e/tiny": snapshot}, + ), + )) conn.wait_for(_is_model_event("e2e/tiny", pb.MODEL_STATE_DOWNLOADING)) conn.wait_for(_is_model_event("e2e/tiny", pb.MODEL_STATE_ON_DISK)) - - # LOAD -> setup runs (typed str path injection); CPU host => IN_RAM; - # the function flips to available. - conn.send(model_op=pb.ModelOp( - op=pb.MODEL_OP_KIND_LOAD, ref="e2e/tiny", snapshot=snapshot)) conn.wait_for(_is_model_event("e2e/tiny", pb.MODEL_STATE_IN_RAM)) conn.wait_for( lambda m: m.WhichOneof("msg") == "state_delta" and "model-echo" in m.state_delta.available_functions + and m.state_delta.observed_residency_generation == 1 ) # The handler sees the materialized snapshot content. @@ -840,18 +847,6 @@ def log_message(self, *a): snap_refs = {r.ref for r in harness.worker.executor.store.residency_snapshot()} assert "e2e/tiny" in snap_refs - # UNLOAD -> teardown; residency falls back to ON_DISK (a SECOND - # on-disk event, after the download's); function gated again. - conn.send(model_op=pb.ModelOp(op=pb.MODEL_OP_KIND_UNLOAD, ref="e2e/tiny")) - deadline = time.monotonic() + _TIMEOUT - while conn.count(_is_model_event("e2e/tiny", pb.MODEL_STATE_ON_DISK)) < 2: - assert time.monotonic() < deadline, "no ON_DISK event after UNLOAD" - time.sleep(0.02) - conn.wait_for( - lambda m: m.WhichOneof("msg") == "state_delta" - and "model-echo" not in m.state_delta.available_functions - and "model-echo" in m.state_delta.loading_functions - ) finally: harness.stop() server.stop(grace=0) @@ -862,8 +857,8 @@ def log_message(self, *a): # Setup failure surfaces FnUnavailable (th#581 worker-side / ernie roster # find): a function whose pipeline setup raises must NOT sit in # loading_functions forever under a READY phase — it leaves BOTH lists and a -# terminal FnUnavailable{setup_failed} reaches the hub. A later LOAD retry -# that succeeds re-advertises it; a re-failure re-emits. +# terminal FnUnavailable{setup_failed} reaches the hub. Re-sending the same +# desired generation retries setup and re-advertises it after recovery. # --------------------------------------------------------------------------- @@ -923,10 +918,22 @@ def _is_fn_unavailable(m: pb.WorkerMessage) -> bool: ) assert "broken-echo" in ready.state_delta.loading_functions - # LOAD -> setup raises -> terminal per-function signal, and the fn + desired = pb.DesiredResidency( + generation=1, + disk_refs=["e2e/broken"], + hot=[pb.DesiredInstance( + function_name="broken-echo", + models=[pb.ModelBinding(slot="model", ref="e2e/broken")], + )], + snapshots={"e2e/broken": snapshot}, + ) + + # Desired setup raises -> terminal per-function signal, and the fn # leaves BOTH available and loading (no more silent-ready limbo). - conn.send(model_op=pb.ModelOp( - op=pb.MODEL_OP_KIND_LOAD, ref="e2e/broken", snapshot=snapshot)) + conn.send(hello_ack=pb.HelloAck( + protocol_version=pb.PROTOCOL_VERSION_CURRENT, + desired_residency=desired, + )) sig = conn.wait_for(_is_fn_unavailable).fn_unavailable assert "pipeline exploded" in sig.detail conn.wait_for( @@ -934,17 +941,20 @@ def _is_fn_unavailable(m: pb.WorkerMessage) -> bool: and "broken-echo" not in m.state_delta.available_functions and "broken-echo" not in m.state_delta.loading_functions ) - # The load path also reports the model op failure itself. + # The residency path also reports the model failure itself. conn.wait_for(_is_model_event("e2e/broken", pb.MODEL_STATE_FAILED)) - # Hub-directed retry succeeds -> the disable lifts and the function - # is advertised (this is the capability-change event the hub needs). + # Same-generation full replacement is a retry (and URL refresh), so + # setup succeeds without inventing an imperative command. e2e_endpoints.BREAK_SETUP.clear() - conn.send(model_op=pb.ModelOp( - op=pb.MODEL_OP_KIND_LOAD, ref="e2e/broken", snapshot=snapshot)) + conn.send(hello_ack=pb.HelloAck( + protocol_version=pb.PROTOCOL_VERSION_CURRENT, + desired_residency=desired, + )) conn.wait_for( lambda m: m.WhichOneof("msg") == "state_delta" and "broken-echo" in m.state_delta.available_functions + and m.state_delta.observed_residency_generation == 1 ) assert "broken-echo" not in harness.worker.executor.unavailable @@ -956,19 +966,6 @@ def _is_fn_unavailable(m: pb.WorkerMessage) -> bool: assert res.status == pb.JOB_STATUS_OK assert _decode_out(res.inline).response == payload.decode() - # Re-failure after recovery re-emits (dedupe set is pruned). - e2e_endpoints.BREAK_SETUP.set() - conn.send(model_op=pb.ModelOp(op=pb.MODEL_OP_KIND_UNLOAD, ref="e2e/broken")) - deadline = time.monotonic() + _TIMEOUT - while "broken-echo" in harness.worker.executor.available_functions(): - assert time.monotonic() < deadline, "UNLOAD did not gate broken-echo" - time.sleep(0.02) - conn.send(model_op=pb.ModelOp( - op=pb.MODEL_OP_KIND_LOAD, ref="e2e/broken", snapshot=snapshot)) - deadline = time.monotonic() + _TIMEOUT - while conn.count(_is_fn_unavailable) < 2: - assert time.monotonic() < deadline, "no second fn_unavailable" - time.sleep(0.02) finally: e2e_endpoints.BREAK_SETUP.set() harness.stop() diff --git a/uv.lock b/uv.lock index 87da0435..c907dc23 100644 --- a/uv.lock +++ b/uv.lock @@ -583,7 +583,7 @@ wheels = [ [[package]] name = "gen-worker" -version = "0.25.2" +version = "0.26.0" source = { editable = "." } dependencies = [ { name = "blake3" }, From 40d9912b0e7655ff2cc91ec32c05956e6b583de7 Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Tue, 14 Jul 2026 14:49:34 -0600 Subject: [PATCH 2/3] fix: close declarative residency gaps --- pyproject.toml | 2 +- src/gen_worker/executor.py | 22 +++++++++---- src/gen_worker/lifecycle.py | 4 ++- tests/test_declarative_residency.py | 48 ++++++++++++++++++++++++++++- tests/test_disk_gc.py | 21 +++++++++++-- uv.lock | 2 +- 6 files changed, 87 insertions(+), 12 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index c68e5f6c..51d15e85 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -23,7 +23,7 @@ requires-python = ">=3.12" dependencies = [ "grpcio>=1.82.1", "msgspec>=0.18.6", - "protobuf>=6.30.0", + "protobuf>=7.35.0", "requests>=2.32.0", "boto3>=1.41.0", "psutil>=7.0.0", diff --git a/src/gen_worker/executor.py b/src/gen_worker/executor.py index 8c1b6b34..fa9a5317 100644 --- a/src/gen_worker/executor.py +++ b/src/gen_worker/executor.py @@ -428,7 +428,7 @@ def __init__( ) self._locks: Dict[str, asyncio.Lock] = {} self._bindings: Dict[str, Any] = {} - self.keep: set[str] = set() + self.keep: list[str] = [] 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 @@ -507,6 +507,15 @@ def has_snapshot(self, ref: str) -> bool: (gw#465): snapshot-less ops for it can still materialize the bytes.""" return ref in self._snapshots + def bank_snapshot(self, ref: str, snapshot: pb.Snapshot) -> None: + """Make hub metadata available without starting a download.""" + if not ref or not snapshot.digest or not snapshot.files: + return + self._snapshots[ref] = snapshot + waiter = self._snapshot_waiters.get(ref) + if waiter is not None: + waiter.set() + def component_digests(self, ref: str, local_path: Optional[Path] = None) -> Dict[str, str]: """Per-component content identity of ``ref``'s snapshot (gw#479): ``{top_level_subfolder: content_set_digest}``. Weight/data files use @@ -600,7 +609,11 @@ def _gc_candidates( if honor_grace and (now - last) < _DISK_GC_GRACE_S: continue out.append((last, ref)) - out.sort() + if include_keep: + rank = {ref: index for index, ref in enumerate(self.keep)} + out.sort(key=lambda item: (-rank[item[1]], item[0], item[1])) + else: + out.sort() return [r for _, r in out] def _evict_disk_ref(self, ref: str) -> None: @@ -684,10 +697,7 @@ async def ensure_local( if binding is None: binding = self._bindings.get(ref) if snapshot is not None and snapshot.digest and snapshot.files: - self._snapshots[ref] = snapshot - waiter = self._snapshot_waiters.get(ref) - if waiter is not None: - waiter.set() + self.bank_snapshot(ref, snapshot) elif snapshot is None: snapshot = self._snapshots.get(ref) async with self._lock(ref): diff --git a/src/gen_worker/lifecycle.py b/src/gen_worker/lifecycle.py index f5d05ea2..6e061612 100644 --- a/src/gen_worker/lifecycle.py +++ b/src/gen_worker/lifecycle.py @@ -170,7 +170,9 @@ async def on_hello_ack(self, ack: pb.HelloAck) -> None: ) else: self._observed_residency_generation = generation - self.executor.store.keep = {ref for ref in desired.disk_refs if ref} + self.executor.store.keep = list(dict.fromkeys(ref for ref in desired.disk_refs if ref)) + for ref, snapshot in desired.snapshots.items(): + self.executor.store.bank_snapshot(ref, snapshot) self._replace_residency_reconcile(desired) # New connection: per-worker fn disables/degradations were wiped by # Hello; re-emit any that still hold, then re-baseline dynamic state. diff --git a/tests/test_declarative_residency.py b/tests/test_declarative_residency.py index b0a1cec9..cc9a92d5 100644 --- a/tests/test_declarative_residency.py +++ b/tests/test_declarative_residency.py @@ -3,6 +3,9 @@ from __future__ import annotations import asyncio +import re +import tomllib +from pathlib import Path from types import SimpleNamespace from typing import Any @@ -112,6 +115,22 @@ def test_proto_field_numbers_match_tensorhub_contract() -> None: assert not hasattr(pb, "MODEL_OP_KIND_UNLOAD") +def test_declared_protobuf_floor_imports_generated_code() -> None: + root = Path(__file__).parents[1] + project = tomllib.loads((root / "pyproject.toml").read_text(encoding="utf-8")) + requirement = next( + dep for dep in project["project"]["dependencies"] if dep.startswith("protobuf>=") + ) + floor = tuple(map(int, requirement.removeprefix("protobuf>=").split("."))) + source = (root / "src/gen_worker/pb/worker_scheduler_pb2.py").read_text( + encoding="utf-8" + ) + match = re.search(r"Protobuf Python Version: (\d+)\.(\d+)\.(\d+)", source) + assert match is not None + assert floor >= tuple(map(int, match.groups())) + assert pb.DESCRIPTOR.name == "worker_scheduler.proto" + + def test_full_replace_supersedes_and_same_generation_refreshes_urls(monkeypatch) -> None: async def run() -> None: lifecycle, executor, transport = _lifecycle() @@ -168,7 +187,7 @@ async def ensure_local(ref: str, snapshot=None, *, binding=None): disk_refs=["acme/stale"], snapshots={"acme/stale": _snapshot("https://r2/stale")}, ))) - assert executor.store.keep == {"acme/new"} + assert executor.store.keep == ["acme/new"] assert lifecycle._state_delta().observed_residency_generation == 2 assert calls == [ ("acme/old", "https://r2/old"), @@ -179,6 +198,33 @@ async def ensure_local(ref: str, snapshot=None, *, binding=None): asyncio.run(run()) +def test_non_idle_hello_ack_banks_snapshot_before_reconcile() -> None: + async def run() -> None: + lifecycle, executor, _ = _lifecycle() + ref = "tensorhub/active-request" + executor._idle.clear() + waiting = asyncio.create_task(executor.store._await_hub_snapshot(ref)) + await asyncio.sleep(0) + assert not waiting.done() + + await lifecycle.on_hello_ack(pb.HelloAck( + desired_residency=pb.DesiredResidency( + generation=1, + disk_refs=[ref], + snapshots={ref: _snapshot("https://r2/reminted")}, + ) + )) + + snapshot = await asyncio.wait_for(waiting, 0.1) + assert snapshot.files[0].url == "https://r2/reminted" + assert not executor._idle.is_set() + assert lifecycle._residency_task is not None + assert not lifecycle._residency_task.done() + lifecycle._cancel_residency_reconcile() + + asyncio.run(run()) + + def test_hot_instance_uses_exact_dynamic_slot_binding(monkeypatch) -> None: async def run() -> None: executor = Executor([_spec()], _noop_send) diff --git a/tests/test_disk_gc.py b/tests/test_disk_gc.py index 655bfaeb..36103094 100644 --- a/tests/test_disk_gc.py +++ b/tests/test_disk_gc.py @@ -76,7 +76,7 @@ def _failed(sent) -> list: 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"} + store.keep = ["t/keep"] async def _run() -> None: await store.ensure_local("t/a", _snapshot("da", 3000)) @@ -97,7 +97,7 @@ async def _run() -> None: 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"} + store.keep = ["t/keep"] async def _run() -> None: await store.ensure_local("t/keep", _snapshot("dk", 3000)) @@ -108,6 +108,23 @@ async def _run() -> None: assert store.residency.tier("t/big") is Tier.DISK +def test_keep_pressure_evicts_lowest_controller_priority_first( + tmp_path, _fake_download +) -> None: + sent: list = [] + store = _store(tmp_path, sent) + store.keep = ["t/high", "t/mid", "t/low"] + + async def _run() -> None: + await store.ensure_local("t/high", _snapshot("dh", 3000)) + await store.ensure_local("t/mid", _snapshot("dm", 3000)) + await store.ensure_local("t/low", _snapshot("dl", 3000)) + await store.ensure_local("t/job", _snapshot("dj", 2000)) + + asyncio.run(_run()) + assert _evicted(sent) == ["t/low"] + + def test_insufficient_disk_fails_fast_with_event(tmp_path, _fake_download) -> None: sent: list = [] store = _store(tmp_path, sent) diff --git a/uv.lock b/uv.lock index c907dc23..335f3b7e 100644 --- a/uv.lock +++ b/uv.lock @@ -664,7 +664,7 @@ requires-dist = [ { name = "numpy", marker = "extra == 'video'", specifier = ">=1.24" }, { name = "pillow", marker = "extra == 'dev'", specifier = ">=10.0" }, { name = "pillow", marker = "extra == 'images'", specifier = ">=10.0" }, - { name = "protobuf", specifier = ">=6.30.0" }, + { name = "protobuf", specifier = ">=7.35.0" }, { name = "psutil", specifier = ">=7.0.0" }, { name = "pyarrow", marker = "extra == 'datasets'", specifier = ">=17.0.0" }, { name = "pyarrow", marker = "extra == 'dev'", specifier = ">=17.0.0" }, From ded3dd2ae1199779fd5bfdc3829008cdb6d2b9b3 Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Tue, 14 Jul 2026 14:55:25 -0600 Subject: [PATCH 3/3] fix: close residency reconciliation races --- src/gen_worker/executor.py | 28 +++++++++++++------ tests/test_disk_gc.py | 52 +++++++++++++++++++++++++++++++++++ tests/test_snapshot_remint.py | 35 +++++++++++++++++++++++ 3 files changed, 107 insertions(+), 8 deletions(-) diff --git a/src/gen_worker/executor.py b/src/gen_worker/executor.py index fa9a5317..bdf29b92 100644 --- a/src/gen_worker/executor.py +++ b/src/gen_worker/executor.py @@ -584,34 +584,43 @@ def rescan_disk(self) -> None: 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). + keep-pressure the escape hatch evicts lowest-priority `keep` refs too + (contract §7 — EVICTED is emitted so the hub re-downloads when demand + returns). In-use / loaded refs are never touched.""" + keep = tuple(self.keep) + keep_rank = {ref: index for index, ref in enumerate(keep)} for include_keep, honor_grace in ( - (False, True), (False, False), (True, True), (True, False), + (False, True), (False, False), (True, False), ): - for ref in self._gc_candidates(include_keep, honor_grace, exclude): + for ref in self._gc_candidates( + include_keep, honor_grace, exclude, keep, keep_rank + ): 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, ...] + self, + include_keep: bool, + honor_grace: bool, + exclude: Tuple[str, ...], + keep: Tuple[str, ...], + keep_rank: Dict[str, int], ) -> 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: + if (ref in 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)) if include_keep: - rank = {ref: index for index, ref in enumerate(self.keep)} - out.sort(key=lambda item: (-rank[item[1]], item[0], item[1])) + out.sort(key=lambda item: (-keep_rank[item[1]], item[0], item[1])) else: out.sort() return [r for _, r in out] @@ -659,6 +668,9 @@ async def _await_hub_snapshot(self, ref: str) -> pb.Snapshot: ensure_local wakes us and then queues behind the lock. Raises :class:`MissingSnapshotError` when nothing arrives in ``_MISSING_SNAPSHOT_WAIT_S``.""" + snapshot = self._snapshots.get(ref) + if snapshot is not None and snapshot.digest and snapshot.files: + return snapshot waiter = self._snapshot_waiters.get(ref) if waiter is None: waiter = self._snapshot_waiters[ref] = asyncio.Event() diff --git a/tests/test_disk_gc.py b/tests/test_disk_gc.py index 36103094..0b613c53 100644 --- a/tests/test_disk_gc.py +++ b/tests/test_disk_gc.py @@ -9,6 +9,7 @@ from __future__ import annotations import asyncio +import time from pathlib import Path import pytest @@ -125,6 +126,57 @@ async def _run() -> None: assert _evicted(sent) == ["t/low"] +def test_keep_priority_outranks_recent_use(tmp_path, _fake_download, monkeypatch) -> None: + sent: list = [] + store = _store(tmp_path, sent) + store.keep = ["t/high", "t/low"] + + async def _run() -> None: + await store.ensure_local("t/high", _snapshot("dh", 3000)) + await store.ensure_local("t/low", _snapshot("dl", 3000)) + now = time.time() + monkeypatch.setattr( + store._index, + "last_used", + lambda ref: now - 7200 if ref == "t/high" else now, + ) + monkeypatch.setattr(executor_mod, "_DISK_GC_GRACE_S", 3600.0) + await asyncio.to_thread(store.gc_disk, 6000) + await asyncio.sleep(0) + + asyncio.run(_run()) + assert _evicted(sent) == ["t/low"] + + +def test_gc_uses_one_keep_snapshot_during_replace( + tmp_path, _fake_download, monkeypatch +) -> None: + sent: list = [] + store = _store(tmp_path, sent) + store.keep = ["t/high", "t/low"] + + async def _run() -> None: + await store.ensure_local("t/high", _snapshot("dh", 3000)) + await store.ensure_local("t/low", _snapshot("dl", 3000)) + refs_in = store.residency.refs_in + replaced = False + + def replace_during_scan(tier): + nonlocal replaced + refs = refs_in(tier) + if not replaced: + replaced = True + store.keep = ["t/replacement"] + return refs + + monkeypatch.setattr(store.residency, "refs_in", replace_during_scan) + await asyncio.to_thread(store.gc_disk, 6000) + await asyncio.sleep(0) + + asyncio.run(_run()) + assert _evicted(sent) == ["t/low"] + + def test_insufficient_disk_fails_fast_with_event(tmp_path, _fake_download) -> None: sent: list = [] store = _store(tmp_path, sent) diff --git a/tests/test_snapshot_remint.py b/tests/test_snapshot_remint.py index 7b4d1dae..63d57d5f 100644 --- a/tests/test_snapshot_remint.py +++ b/tests/test_snapshot_remint.py @@ -91,6 +91,41 @@ async def _run() -> None: asyncio.run(_run()) +def test_store_uses_snapshot_banked_while_waiting_for_ref_lock(tmp_path, monkeypatch) -> None: + """A remint between the pre-lock lookup and snapshot wait is not lost.""" + import gen_worker.executor as ex_mod + from gen_worker.executor import ModelStore + + sent: List[pb.WorkerMessage] = [] + + async def _send(msg: pb.WorkerMessage) -> None: + sent.append(msg) + + async def _fake_download(ref: str, **kwargs: Any) -> Path: + assert kwargs.get("snapshot") is not None + path = tmp_path / ref.replace("/", "_") + path.mkdir(parents=True, exist_ok=True) + return path + + monkeypatch.setattr(ex_mod, "ensure_local", _fake_download) + ref = "acme/banked-before-wait" + store = ModelStore(_send, cache_dir=tmp_path) + store.register_binding(ref, Hub(ref)) + + async def _run() -> None: + lock = store._lock(ref) + await lock.acquire() + cold = asyncio.create_task(store.ensure_local(ref)) + await asyncio.sleep(0) + assert not cold.done() + store.bank_snapshot(ref, _snapshot()) + lock.release() + assert await asyncio.wait_for(cold, 0.5) == tmp_path / "acme_banked-before-wait" + + asyncio.run(_run()) + assert _failed_events(sent) == [] + + def test_missing_snapshot_maps_retryable_never_fatal() -> None: """A cold worker mid-resolution asks the scheduler to retry the job.""" from gen_worker.executor import _map_exception