diff --git a/docs/environment.md b/docs/environment.md index 051dba28..a603abcc 100644 --- a/docs/environment.md +++ b/docs/environment.md @@ -42,6 +42,7 @@ this page covers the worker itself. | `HF_HOME` | `hf_home` | HF cache root (also read by huggingface_hub itself) | | `TENSORHUB_URL` | `tensorhub_url` | standalone-CLI resolve base URL | | `TENSORHUB_CACHE_DIR` / `TENSORHUB_CAS_DIR` | `tensorhub_cache_dir` / `tensorhub_cas_dir` | move cache/CAS off `/tmp` (cozy local persistence); `CAS_DIR` also isolates the `cli/run.py` standalone CLI in tests | +| `TENSORHUB_FILL_SOURCE_DIR` | `tensorhub_fill_source_dir` | th#850 managed-tier ruling: an endpoint-scoped datacenter-warm CAS mount (RunPod volume), checked before R2 on a blob miss and write-through warmed from R2. Never the CAS root — that always stays `TENSORHUB_CACHE_DIR`/local. tensorhub sets this only when a volume is attached; ismount-guarded, so a plain directory never gets mistaken for it | ## C2PA Content Credentials (Settings fields, th#714) diff --git a/proto/worker_scheduler.proto b/proto/worker_scheduler.proto index 2fdc3f53..a42cc977 100644 --- a/proto/worker_scheduler.proto +++ b/proto/worker_scheduler.proto @@ -550,6 +550,15 @@ message ModelEvent { // one later causal compiled-runtime FAILED; empty on boot-attached cells and // ordinary model events. string target_incarnation_id = 19; + // th#850 managed-tier fill hierarchy (gw#599): bytes actually fetched over + // the network (R2 origin) to produce this ON_DISK transition, as opposed + // to bytes served from a warm local-disk or endpoint-volume cache hit. Set + // on every fresh materialization's ON_DISK event, including an explicit 0 + // — pair with the DOWNLOADING events' bytes_total for this ref to read + // "warm boot ⇒ ~0 R2 bytes". Additive: old workers never set it, which a + // receiver must treat as identity unknown (same convention as + // snapshot_digest/residency_generation above), never as a verified zero. + int64 network_bytes = 20; } enum ModelState { diff --git a/src/gen_worker/config/loader.py b/src/gen_worker/config/loader.py index bfced93a..b750b68d 100644 --- a/src/gen_worker/config/loader.py +++ b/src/gen_worker/config/loader.py @@ -42,6 +42,7 @@ "TENSORHUB_TOKEN": "tensorhub_token", "TENSORHUB_CACHE_DIR": "tensorhub_cache_dir", "TENSORHUB_CAS_DIR": "tensorhub_cas_dir", + "TENSORHUB_FILL_SOURCE_DIR": "tensorhub_fill_source_dir", "CIVITAI_API_KEY": "civitai_api_key", "GEN_WORKER_C2PA_CERT_PATH": "c2pa_cert_path", "GEN_WORKER_C2PA_KEY_PATH": "c2pa_key_path", diff --git a/src/gen_worker/config/settings.py b/src/gen_worker/config/settings.py index e7d48ac6..e4be1d5c 100644 --- a/src/gen_worker/config/settings.py +++ b/src/gen_worker/config/settings.py @@ -66,6 +66,10 @@ class Settings(msgspec.Struct, frozen=True, kw_only=True): # models/provision.py::resolve_local_path). tensorhub_cache_dir: str = "" # TENSORHUB_CACHE_DIR tensorhub_cas_dir: str = "" # TENSORHUB_CAS_DIR + # th#850 managed-tier ruling (gw#599): endpoint-scoped datacenter-warm + # fill source (RunPod network volume mount), tried before R2. Never the + # CAS root itself — see models/cache_paths.py::tensorhub_fill_source_dir. + tensorhub_fill_source_dir: str = "" # TENSORHUB_FILL_SOURCE_DIR # Civitai provider credential (CIVITAI_API_KEY, alias CIVITAI_TOKEN). civitai_api_key: str = "" diff --git a/src/gen_worker/executor.py b/src/gen_worker/executor.py index 21fa2dd0..677ffeed 100644 --- a/src/gen_worker/executor.py +++ b/src/gen_worker/executor.py @@ -53,6 +53,7 @@ InsufficientHostRamError, ) from .input_assets import cleanup_input_assets, materialize_input_assets +from .models import cozy_snapshot from .models import disk_gc from .models import provision from .models import residency as residency_mod @@ -68,7 +69,7 @@ next_offload_rung, release_unused_pinned_host_cache, ) -from .models.cache_paths import tensorhub_cas_dir +from .models.cache_paths import tensorhub_cas_dir, tensorhub_fill_source_dir from .models.download import ensure_local, lookup_provider_for_ref from .models.errors import MissingSnapshotError, UrlExpiredError from .models.residency import Residency @@ -541,11 +542,19 @@ def __init__( cache_dir: Optional[Path] = None, vram_budget_bytes: Optional[int] = None, disk_free_bytes_fn: Optional[Callable[[], int]] = None, + fill_source_dir: Optional[Path] = None, ) -> None: self._emit = emit self._hf_home = hf_home or None self._hf_token = hf_token or None self._cache_dir = cache_dir or tensorhub_cas_dir() + # th#850 managed-tier ruling (gw#599): endpoint-scoped datacenter-warm + # fill source (RunPod volume mount), consulted before R2 on a blob + # miss — resolved once at boot like _cache_dir; never the CAS root. + # Same `or` shape as _cache_dir above: an explicit path (tests) wins, + # otherwise resolve from env (production/tensorhub; unset -> None, + # the cozy-local/no-volume degenerate case). + self._fill_source_dir = fill_source_dir or tensorhub_fill_source_dir() self.residency = Residency( on_event=self._on_residency_event, vram_budget_bytes=vram_budget_bytes, ) @@ -585,6 +594,14 @@ def __init__( # Cold-ref waiters (th#763): ensure_local blocks here until the # hub's re-minted DOWNLOAD banks a snapshot for the ref. self._snapshot_waiters: Dict[str, asyncio.Event] = {} + # th#850 managed-tier ruling (gw#599): network_bytes for the NEXT + # ON_DISK transition of this ref, handed off to + # _on_residency_event so the one authoritative wire event Residency + # emits carries it — set immediately before track_disk(), consumed + # (popped) by _on_residency_event if it fires, cleared defensively + # otherwise. Avoids a second, redundant ON_DISK event and avoids + # widening EventFn's arity (Residency has other direct callers). + self._pending_network_bytes: Dict[str, int] = {} def _default_disk_free(self) -> int: p = Path(self._cache_dir) @@ -625,6 +642,9 @@ def _on_residency_event( ) if identity[0]: self._resident_identities[ref] = identity + pending_network_bytes = self._pending_network_bytes.pop(ref, None) + if pending_network_bytes is not None: + kw["network_bytes"] = int(pending_network_bytes) else: identity = self.resident_identity(ref) coro = self._event(ref, pb_state, identity=identity, **kw) @@ -965,6 +985,10 @@ def _gc_candidates( keep: Tuple[str, ...], keep_rank: Dict[str, int], ) -> List[str]: + """The evictable SET for one gc_disk pass: hard invariants only + (never exclude/in-use — no policy ever overrides these), plus this + pass's keep-membership/grace filter. Ordering within that set is a + separate seam, see ``_disk_eviction_order``.""" now = time.time() out: List[Tuple[float, str]] = [] for ref in self.residency.refs_in(residency_mod.Tier.DISK): @@ -976,11 +1000,28 @@ def _gc_candidates( if honor_grace and (now - last) < _DISK_GC_GRACE_S: continue out.append((last, ref)) + return self._disk_eviction_order(out, include_keep, keep_rank) + + # th#850 managed-tier ruling (gw#599): the eviction POLICY (ranking one + # pass's evictable set) is a distinct seam from the evictable SET itself + # (``_gc_candidates`` above, which owns the hard never-evict invariants). + # Default is the LRU-oldest-first/keep-priority-escape-hatch ordering + # this store has always used — an instance may swap this attribute for a + # scheduler-intent-aware policy without touching gc_disk's free-space + # loop or the invariant filter. Building that policy is a follow-on + # (Paul ruled the seam only here); the default below is exactly today's + # behavior, byte-for-byte. + @staticmethod + def _default_disk_eviction_order( + entries: List[Tuple[float, str]], include_keep: bool, keep_rank: Dict[str, int], + ) -> List[str]: if include_keep: - out.sort(key=lambda item: (-keep_rank[item[1]], item[0], item[1])) + ordered = sorted(entries, key=lambda item: (-keep_rank[item[1]], item[0], item[1])) else: - out.sort() - return [r for _, r in out] + ordered = sorted(entries) + return [ref for _, ref in ordered] + + _disk_eviction_order = _default_disk_eviction_order def _evict_disk_ref(self, ref: str) -> None: path = self.residency.local_path(ref) or self._index.path(ref) @@ -1153,6 +1194,14 @@ def complete(path: Path) -> _MaterializedLocal: operation_identity, ) last_progress = 0.0 + # th#850 managed-tier ruling (gw#599): opened before _progress so + # its DOWNLOADING ticks can read the running total, and entered + # once for the whole retry loop so it accumulates across + # attempts. The hub (tensorhub th#850/PR#493) reads network_bytes + # off the DOWNLOADING events' running value (mirrors + # bytes_done/bytes_total), not just the terminal ON_DISK one — + # both must carry it for the wire contract to actually work. + net_scope = cozy_snapshot.NetworkBytesScope() def _progress(done: int, total: Optional[int]) -> None: nonlocal last_progress @@ -1164,7 +1213,8 @@ def _progress(done: int, total: Optional[int]) -> None: asyncio.run_coroutine_threadsafe( self._event(ref, pb.MODEL_STATE_DOWNLOADING, identity=operation_identity, - bytes_done=int(done), bytes_total=int(total or 0)), + bytes_done=int(done), bytes_total=int(total or 0), + network_bytes=net_scope.network_bytes), self._loop, ) @@ -1177,17 +1227,19 @@ def _progress(done: int, total: Optional[int]) -> None: resolved = None if snapshot is not None and snapshot.digest: resolved = _snapshot_to_resolved(snapshot) - path = await ensure_local( - ref, - provider=getattr(binding, "source", None), - snapshot=resolved, - cache_dir=self._cache_dir, - hf_home=self._hf_home, - hf_token=self._hf_token, - allow_patterns=tuple(getattr(binding, "files", ()) or ()), - components=tuple(getattr(binding, "components", ()) or ()), - progress=_progress, - ) + with net_scope: + path = await ensure_local( + ref, + provider=getattr(binding, "source", None), + snapshot=resolved, + cache_dir=self._cache_dir, + hf_home=self._hf_home, + hf_token=self._hf_token, + allow_patterns=tuple(getattr(binding, "files", ()) or ()), + components=tuple(getattr(binding, "components", ()) or ()), + progress=_progress, + fill_source_dir=self._fill_source_dir, + ) tier_before = self.residency.tier(ref) with self._identity_lock: identity_changed = ( @@ -1198,13 +1250,25 @@ def _progress(done: int, total: Optional[int]) -> None: self._disk_identities[ref] = operation_identity if tier_before in (None, residency_mod.Tier.DISK): self._resident_identities[ref] = operation_identity + # th#850 managed-tier ruling (gw#599): handed off to + # _on_residency_event so the ON_DISK event Residency + # emits for a genuinely fresh registration carries the + # bytes this materialization fetched over the network + # (0 included — pairs with the DOWNLOADING events' + # bytes_total for the "warm boot ⇒ ~0 R2 bytes" signal). + self._pending_network_bytes[ref] = net_scope.network_bytes self.residency.track_disk(ref, path) + self._pending_network_bytes.pop(ref, None) # defensive: unconsumed if no event fired if tier_before is residency_mod.Tier.DISK and identity_changed: - # Residency suppresses same-tier event spam. A digest - # move is nevertheless a semantic ON_DISK transition. + # Residency suppresses same-tier event spam (track_disk + # above did not consume the pending value above). A + # digest move is nevertheless a semantic ON_DISK + # transition — carries network_bytes directly since + # this is our own explicit event, not Residency's. await self._event( ref, pb.MODEL_STATE_ON_DISK, identity=operation_identity, + network_bytes=net_scope.network_bytes, ) # tree_bytes stats every file — off-loop (gw#407: no # multi-GB directory walks on the event loop). diff --git a/src/gen_worker/models/cache_paths.py b/src/gen_worker/models/cache_paths.py index cb60be7a..e94105e6 100644 --- a/src/gen_worker/models/cache_paths.py +++ b/src/gen_worker/models/cache_paths.py @@ -1,5 +1,6 @@ from __future__ import annotations +import os from pathlib import Path from ..config import get_settings @@ -13,14 +14,14 @@ def tensorhub_cache_dir() -> Path: Honors the ``TENSORHUB_CACHE_DIR`` environment variable when set. This is the ONE knob for where the CAS lives: the cozy-local runner points it at - a persistent ``~/.cache/tensorhub`` (weights survive reboots); tensorhub - points it at a mounted RunPod endpoint volume (`/tensorhub-endpoint-cache`) - when the pod has one attached, so the ordinary R2 prefetch that - materializes blobs here on the first boot warms the volume for later - same-endpoint pods too (th#850) — no separate shared/read-through tier. - Falls back to the ``/tmp`` default when unset. The CAS implementation - itself is deliberately backend-agnostic: nothing branches on what's - mounted here. + a persistent ``~/.cache/tensorhub`` (weights survive reboots). The CAS + root ALWAYS stays on local/pod-local storage — a managed, bounded LRU + tier (th#850 managed-tier ruling, gw#599: supersedes the earlier + CAS-root-on-volume shape). A mounted RunPod endpoint volume, when + attached, is a FILL SOURCE consulted before R2 (see + ``tensorhub_fill_source_dir``) — it is never the CAS root itself. Falls + back to the ``/tmp`` default when unset. The CAS implementation itself is + deliberately backend-agnostic: nothing branches on what's mounted here. """ env = get_settings().tensorhub_cache_dir.strip() if env: @@ -29,5 +30,28 @@ def tensorhub_cache_dir() -> Path: def tensorhub_cas_dir() -> Path: - """Worker CAS root: /cas.""" + """Worker CAS root: /cas. Always local/pod storage.""" return tensorhub_cache_dir() / "cas" + + +def tensorhub_fill_source_dir() -> Path | None: + """Endpoint-scoped datacenter-warm fill source (th#850 managed-tier + ruling), or ``None`` when no volume is attached. + + Honors ``TENSORHUB_FILL_SOURCE_DIR``, set by tensorhub only when this + pod's endpoint has a RunPod network volume attached. Guarded by + ``os.path.ismount`` — a plain directory baked into the image or left on + the container disk must never be mistaken for the real per-endpoint + volume (same defensive pattern the removed ``endpoint_shared_cas_dir`` + used). This is FILL SOURCE #1 in the CAS layer's fetch order (volume, + then R2); it is never the CAS root. cozy-local and any pod without a + volume leave this unset, which is the degenerate case: fetch goes + straight to R2, byte-identical to pre-th#850 behavior. + """ + env = get_settings().tensorhub_fill_source_dir.strip() + if not env: + return None + path = Path(env).expanduser() + if not os.path.ismount(path): + return None + return path diff --git a/src/gen_worker/models/cozy_snapshot.py b/src/gen_worker/models/cozy_snapshot.py index fdc24166..0524c693 100644 --- a/src/gen_worker/models/cozy_snapshot.py +++ b/src/gen_worker/models/cozy_snapshot.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import contextvars import hashlib import json import logging @@ -8,12 +9,13 @@ import re import shutil import threading +import time import uuid from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Sequence, Set from .cozy_cas import _blake3_file, _download_one_file as _download_one_file -from .cozy_cas import _norm_rel_path, fsync_dir +from .cozy_cas import _norm_rel_path, fsync_dir, fsync_file from .download import components_present, select_component_paths from .hub_client import WorkerResolvedRepo, WorkerResolvedRepoFile from .refs import TensorhubRef @@ -41,12 +43,15 @@ def _inflight_blob_lock(digest: str) -> "asyncio.Lock": # gw#598 / th#850: verify-on-first-use-per-process for REUSED CAS state. -# The CAS root can now be a persistent volume shared by several pods, so -# blobs and materialized snapshot trees found on disk may be another pod's -# bytes of any age — they must pass a full BLAKE3 check once per process -# before being trusted, exactly like freshly downloaded bytes. Keyed by -# (resolved root, digest/key) so tests with distinct tmp roots never share -# trust; in production the root is constant, so this is once per boot. +# The CAS root persists across pod restarts (cozy-local) and, historically, +# could itself be a volume shared by several pods (superseded by gw#599's +# managed-tier ruling — the root is local-only now, but an operator can +# still point TENSORHUB_CACHE_DIR at a shared path manually). Either way, +# blobs and materialized snapshot trees found on disk may be bytes from a +# different process/era — they must pass a full BLAKE3 check once per +# process before being trusted, exactly like freshly downloaded bytes. Keyed +# by (resolved root, digest/key) so tests with distinct tmp roots never +# share trust; in production the root is constant, so this is once per boot. _TRUST_CAP = 65536 _VERIFIED_BLOBS: Set[tuple] = set() _TRUSTED_SNAPSHOTS: Set[tuple] = set() @@ -59,6 +64,43 @@ def _mark_trusted(s: Set[tuple], item: tuple) -> None: ProgressFn = Callable[[int, Optional[int]], None] +# th#850 managed-tier ruling (gw#599): bytes actually fetched over the +# network (R2 origin), as opposed to bytes served from a warm local/volume +# cache — the signal a "volume-attached boot ⇒ ~0 network bytes" runtime +# assertion needs. Carried via a ContextVar (same idiom as +# provision.py's ArmingScope) scoped to one ensure_local/ensure_snapshot +# call, so no download/stub call site anywhere needs a new parameter to +# forward it — a caller that never opens the scope sees no behavior change. +_NETWORK_BYTES_SINK: "contextvars.ContextVar[Optional[list]]" = contextvars.ContextVar( + "cozy_snapshot_network_bytes_sink", default=None +) + + +class NetworkBytesScope: + """Capture network-fetched bytes for one download call. + + ``with NetworkBytesScope() as scope: path = await ensure_local(...)`` + then ``scope.network_bytes`` is the total bytes this call fetched from + R2 (0 if every blob was already warm locally or on the volume). + """ + + def __init__(self) -> None: + self._sink: List[int] = [0] + self._token: Optional["contextvars.Token"] = None + + def __enter__(self) -> "NetworkBytesScope": + self._token = _NETWORK_BYTES_SINK.set(self._sink) + return self + + def __exit__(self, *exc_info: Any) -> None: + if self._token is not None: + _NETWORK_BYTES_SINK.reset(self._token) + self._token = None + + @property + def network_bytes(self) -> int: + return self._sink[0] + # Free space that must remain after downloading the missing blobs. _DISK_HEADROOM_BYTES = 1 << 30 @@ -109,6 +151,59 @@ def _is_parts_manifest(path: str) -> bool: return path.endswith(".parts.json") +def _copy_verified_blob(src: Path, dst: Path, digest: str, expected_size: int) -> bool: + """Copy one immutable CAS blob through a writer-unique atomic stage, + verifying declared size and BLAKE3 before publishing (th#850 managed-tier + ruling, gw#599). Used both to fill local CAS from the volume fill source + and to write a fresh R2 fetch through to the volume. The final path is + digest-only; readers never observe partial bytes, and racing writers may + replace the same final name only after each independently verifies size + and BLAKE3 (mirrors the multi-writer discipline gw#597 established for + ordinary downloads).""" + try: + if not src.is_file(): + return False + if expected_size and src.stat().st_size != expected_size: + _log.warning( + "blob_fill_corrupt source=%s digest=%s reason=size expected=%d actual=%d", + src, digest[:16], expected_size, src.stat().st_size, + ) + return False + + dst.parent.mkdir(parents=True, exist_ok=True) + tmp = dst.parent / ( + f".{dst.name}.writer-{os.getpid()}-{threading.get_ident()}-{uuid.uuid4().hex}" + ) + try: + with src.open("rb") as source, tmp.open("xb") as staged: + shutil.copyfileobj(source, staged, length=4 * 1024 * 1024) + staged.flush() + os.fsync(staged.fileno()) + if expected_size and tmp.stat().st_size != expected_size: + _log.warning( + "blob_fill_corrupt source=%s digest=%s reason=staged_size expected=%d actual=%d", + src, digest[:16], expected_size, tmp.stat().st_size, + ) + return False + if digest and _blake3_file(tmp).lower() != digest.lower(): + _log.warning( + "blob_fill_corrupt source=%s digest=%s reason=blake3", src, digest[:16], + ) + return False + os.replace(tmp, dst) + fsync_file(dst) + fsync_dir(dst.parent) + return True + finally: + tmp.unlink(missing_ok=True) + except OSError as exc: + _log.warning( + "blob_fill_copy_failed source=%s destination=%s digest=%s error=%s", + src, dst, digest[:16], exc, + ) + return False + + def _try_hardlink_or_copy(src: Path, dst: Path) -> None: if dst.exists(): dst.unlink() @@ -217,11 +312,19 @@ async def ensure_snapshot( resolved: Optional[WorkerResolvedRepo], progress: Optional[ProgressFn] = None, components: Sequence[str] = (), + fill_source_dir: Optional[Path] = None, ) -> Path: + """``fill_source_dir`` (th#850 managed-tier ruling, gw#599): an + endpoint-scoped datacenter-warm CAS root (RunPod volume mount, same + ``blobs/`` layout as ``base_dir``) consulted before R2 on a blob + miss. ``None`` is the degenerate cozy-local/no-volume case — fetch + goes straight to R2, byte-identical to pre-th#850 behavior. Never + the CAS root itself; ``base_dir`` (local disk) always is.""" blobs_root = base_dir / "blobs" snaps_root = base_dir / "snapshots" blobs_root.mkdir(parents=True, exist_ok=True) snaps_root.mkdir(parents=True, exist_ok=True) + fill_blobs_root = fill_source_dir / "blobs" if fill_source_dir is not None else None if resolved is None: # Workers don't resolve via HTTP — the orchestrator pre-resolves @@ -289,6 +392,7 @@ async def ensure_snapshot( blobs_root, res.files, progress=progress, + fill_blobs_root=fill_blobs_root, ) # Materialization copies/concatenates multi-GB trees — strictly # off the event loop (gw#407: a loop blocked for the duration of @@ -326,9 +430,10 @@ def _materialize_snapshot( """Blocking build phase (worker thread): reassemble + hardlink into a writer-unique ``.building-`` dir, then atomically rename into place. Writer-unique (not a fixed ``.building`` name) so two - pods racing to materialize the same snapshot key on a SHARED CAS - root (e.g. one RunPod volume, th#850) never interleave writes into - the same tree — only the atomic rename below, not the build, decides + writers racing to materialize the same snapshot key on the same CAS + root (concurrent tasks within one pod, or an operator-configured + shared root) never interleave writes into the same tree — only the + atomic rename below, not the build, decides the winner.""" writer_id = f"{os.getpid()}-{threading.get_ident()}-{uuid.uuid4().hex}" tmp = snaps_root / f"{snap_dir.name}.building-{writer_id}" @@ -360,6 +465,7 @@ async def _ensure_blobs( files: List[WorkerResolvedRepoFile], *, progress: Optional[ProgressFn] = None, + fill_blobs_root: Optional[Path] = None, ) -> None: # Deduplicate by digest — same blob referenced by multiple paths (e.g. # fp16 and normal variants sharing the same part) is downloaded once. @@ -401,6 +507,17 @@ def _on_bytes(n: int, *, network: bool = False) -> None: done += n if network: network_bytes += n + # th#850 managed-tier ruling (gw#599): update the + # NetworkBytesScope sink LIVE (not just once at the end + # of this call, see below) so a caller's mid-flight + # `progress()` tick — called synchronously right below, + # still holding nothing but this lock — can read a + # genuinely-running total. tensorhub reads network_bytes + # off the DOWNLOADING events' running value, the same + # way it reads bytes_done/bytes_total. + sink = _NETWORK_BYTES_SINK.get() + if sink is not None: + sink[0] += n d = done if total is None else min(done, total) if progress is not None: try: @@ -442,6 +559,45 @@ async def _blob_trusted(dst: Path, f: WorkerResolvedRepoFile, digest: str) -> bo dst.unlink(missing_ok=True) return False + async def _fill_from_volume(f: WorkerResolvedRepoFile, digest: str, dst: Path) -> bool: + """th#850 managed-tier ruling (gw#599): FILL SOURCE #1. A blob + present on the endpoint volume is BLAKE3-verified and copied to + local CAS — never trusted on size alone (digest-verification of + volume-read blobs is mandatory regardless of the local/shared + distinction, per Paul's ruling).""" + if fill_blobs_root is None: + return False + src = _blob_path(fill_blobs_root, digest) + started = time.monotonic() + if await asyncio.to_thread( + _copy_verified_blob, src, dst, digest, int(f.size_bytes or 0) + ): + _on_bytes(int(f.size_bytes or 0)) # not network: volume hit + _mark_trusted(_VERIFIED_BLOBS, (blobs_root_id, digest)) + _log.info( + "blob_cache_hit source=volume digest=%s bytes=%d transfer_ms=%d", + digest[:16], int(f.size_bytes or 0), + int((time.monotonic() - started) * 1000), + ) + return True + _log.info("blob_cache_miss source=volume digest=%s", digest[:16]) + return False + + async def _fill_to_volume(f: WorkerResolvedRepoFile, digest: str, dst: Path) -> None: + """Write-through (gw#599): a fresh R2 fetch warms the volume for + the next same-endpoint pod. Best-effort — a publish failure never + fails the request; the next pod simply falls through to R2 too.""" + if fill_blobs_root is None: + return + published = await asyncio.to_thread( + _copy_verified_blob, dst, _blob_path(fill_blobs_root, digest), + digest, int(f.size_bytes or 0), + ) + _log.info( + "blob_fill_publish source=r2 destination=volume digest=%s bytes=%d published=%s", + digest[:16], int(f.size_bytes or 0), published, + ) + async def _dl(f: WorkerResolvedRepoFile) -> None: digest = f.blake3.strip().lower() dst = _blob_path(blobs_root, digest) @@ -454,7 +610,10 @@ async def _dl(f: WorkerResolvedRepoFile) -> None: _log.info("blob_shared_inflight path=%s digest=%s (sibling ref downloaded it)", f.path, digest[:16]) return + if await _fill_from_volume(f, digest, dst): + return await _dl_locked(f, digest, dst) + await _fill_to_volume(f, digest, dst) async def _dl_locked(f: WorkerResolvedRepoFile, digest: str, dst: Path) -> None: async with sem: @@ -488,11 +647,12 @@ async def _dl_locked(f: WorkerResolvedRepoFile, digest: str, dst: Path) -> None: _log.info("blob_download_done path=%s digest=%s", f.path, digest[:16]) await asyncio.gather(*(_dl(f) for f in unique)) - # th#850 runtime-assertion signal: on a volume-attached boot with - # blobs already warm, network_bytes should land near zero. Log-only - # for now (deliberately not threaded through the ProgressFn/ - # ModelEvent wire contract here — that needs a coordinated proto - # change with the th#850 owner, see gw#597). + # th#850 managed-tier runtime-assertion signal: on a volume-attached + # boot with blobs already warm, network_bytes should land near zero. + # _on_bytes above already streamed this total into the caller's + # NetworkBytesScope live, chunk by chunk (so mid-flight DOWNLOADING + # ticks see a genuine running total, not just the final tally) — this + # is a log line only, not a second write to the sink. cache_hit_pct = 100.0 * (1 - network_bytes / total) if total else 100.0 _log.info( "ensure_blobs_summary total_bytes=%s network_bytes=%d cache_hit_pct=%.1f", @@ -688,8 +848,10 @@ async def ensure_snapshot_async( resolved: Optional[WorkerResolvedRepo], progress: Optional[ProgressFn] = None, components: Sequence[str] = (), + fill_source_dir: Optional[Path] = None, ) -> Path: dl = CozySnapshotDownloader() return await dl.ensure_snapshot( base_dir, ref, resolved=resolved, progress=progress, components=components, + fill_source_dir=fill_source_dir, ) diff --git a/src/gen_worker/models/download.py b/src/gen_worker/models/download.py index 6654dd43..4886136c 100644 --- a/src/gen_worker/models/download.py +++ b/src/gen_worker/models/download.py @@ -197,6 +197,7 @@ async def ensure_local( allow_patterns: Sequence[str] = (), components: Sequence[str] = (), progress: Optional[ProgressFn] = None, + fill_source_dir: Optional[Path] = None, ) -> Path: """Materialize ``ref`` on disk; return its local path. @@ -222,6 +223,12 @@ async def ensure_local( for tensorhub refs lands on the hub-less CLI path instead (``models/provision.py::_fetch_tensorhub_snapshot``), which owns its own resolve+download+materialize loop end to end. + + ``fill_source_dir`` (th#850 managed-tier ruling, gw#599): an + endpoint-scoped datacenter-warm CAS mount (RunPod volume) consulted + before R2 on the tensorhub-snapshot branch only. ``None`` (the default, + and always true for cozy-local / non-tensorhub providers) preserves + today's straight-to-R2 behavior exactly. """ base = Path(cache_dir) if cache_dir is not None else tensorhub_cas_dir() prov = provider or lookup_provider_for_ref(ref) @@ -235,6 +242,7 @@ async def ensure_local( ref=_snapshot_ref(parsed, ref), resolved=snapshot, progress=progress, + fill_source_dir=fill_source_dir, ) if parsed.provider == "tensorhub" and parsed.tensorhub is not None: diff --git a/src/gen_worker/pb/worker_scheduler_pb2.py b/src/gen_worker/pb/worker_scheduler_pb2.py index d6ea9360..749fd6d8 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\"\x9b\x02\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\x12/\n\x0bhost_canary\x18\t \x01(\x0b\x32\x1a.cozy.scheduler.HostCanary\x12\x15\n\rtorch_version\x18\n \x01(\t\x12\x1a\n\x12gen_worker_version\x18\x0b \x01(\t\"\xc9\x01\n\nHostCanary\x12\x13\n\x0bmemcpy_gbps\x18\x01 \x01(\x01\x12\x10\n\x08h2d_gbps\x18\x02 \x01(\x01\x12\x10\n\x08\x64\x32h_gbps\x18\x03 \x01(\x01\x12\x17\n\x0fpinned_alloc_ok\x18\x04 \x01(\x08\x12\x17\n\x0f\x63pu_single_mbps\x18\x05 \x01(\x01\x12\x16\n\x0e\x63pu_multi_mbps\x18\x06 \x01(\x01\x12\r\n\x05vcpus\x18\x07 \x01(\x05\x12\x14\n\x0cram_total_gb\x18\x08 \x01(\x01\x12\x13\n\x0b\x64uration_ms\x18\t \x01(\x03\"\x95\x01\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\x12\x17\n\x0fsnapshot_digest\x18\x04 \x01(\t\x12\x1c\n\x14residency_generation\x18\x05 \x01(\x04\"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\"\xb3\x02\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\x12\x36\n\x0f\x63ompile_targets\x18\x07 \x03(\x0b\x32\x1d.cozy.scheduler.CompileTarget\x12\x30\n\x0c\x63\x65ll_lookups\x18\x08 \x03(\x0b\x32\x1a.cozy.scheduler.CellLookup\".\n\nCellLookup\x12\x0e\n\x06\x66\x61mily\x18\x01 \x01(\t\x12\x10\n\x08\x63\x65ll_key\x18\x02 \x01(\t\"\xc6\x03\n\rCompileTarget\x12\x16\n\x0eincarnation_id\x18\x01 \x01(\t\x12\x0e\n\x06\x66\x61mily\x18\x02 \x01(\t\x12\x1c\n\x14pipeline_weight_lane\x18\x03 \x01(\t\x12\x13\n\x0blora_bucket\x18\x04 \x01(\x05\x12\x17\n\x0f\x63ontract_digest\x18\x05 \x01(\t\x12\x1a\n\x12\x61\x63tive_compile_ref\x18\x06 \x01(\t\x12&\n\x1e\x61\x63tive_compile_snapshot_digest\x18\x07 \x01(\t\x12\x16\n\x0e\x66unction_names\x18\x08 \x03(\t\x12<\n\x0emodel_bindings\x18\t \x03(\x0b\x32$.cozy.scheduler.CompileTargetBinding\x12\x1a\n\x12requested_cell_key\x18\n \x01(\t\x12Q\n\x13requested_cell_axes\x18\x0b \x03(\x0b\x32\x34.cozy.scheduler.CompileTarget.RequestedCellAxesEntry\x1a\x38\n\x16RequestedCellAxesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"J\n\x14\x43ompileTargetBinding\x12\x0c\n\x04slot\x18\x01 \x01(\t\x12\x0b\n\x03ref\x18\x02 \x01(\t\x12\x17\n\x0fsnapshot_digest\x18\x03 \x01(\t\"\x88\x04\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\x12\x42\n\x10required_compile\x18\r \x01(\x0b\x32(.cozy.scheduler.RequiredCompileExecution\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\"\x82\x01\n\x18RequiredCompileExecution\x12\x1d\n\x15target_incarnation_id\x18\x01 \x01(\t\x12\x10\n\x08\x63\x65ll_ref\x18\x02 \x01(\t\x12\x1c\n\x14\x63\x65ll_snapshot_digest\x18\x03 \x01(\t\x12\x17\n\x0f\x63ontract_digest\x18\x04 \x01(\t\"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\"\xa0\x01\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\x12\x14\n\x0coperation_id\x18\x04 \x01(\t\x12\x1d\n\x15target_incarnation_id\x18\x05 \x01(\t\"\x84\x04\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\x12\x1f\n\x17host_ram_required_bytes\x18\x0b \x01(\x03\x12\'\n\x1fhost_ram_available_before_bytes\x18\x0c \x01(\x03\x12&\n\x1ehost_ram_available_after_bytes\x18\r \x01(\x03\x12\x1d\n\x15host_ram_evicted_refs\x18\x0e \x03(\t\x12$\n\x1chost_ram_capacity_generation\x18\x0f \x01(\x04\x12\x17\n\x0fsnapshot_digest\x18\x10 \x01(\t\x12\x1c\n\x14residency_generation\x18\x11 \x01(\x04\x12\x14\n\x0coperation_id\x18\x12 \x01(\t\x12\x1d\n\x15target_incarnation_id\x18\x13 \x01(\t\"\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*\x82\x02\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\x12&\n\"MODEL_STATE_HOST_CAPACITY_PROGRESS\x10\x08\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\"\x9b\x02\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\x12/\n\x0bhost_canary\x18\t \x01(\x0b\x32\x1a.cozy.scheduler.HostCanary\x12\x15\n\rtorch_version\x18\n \x01(\t\x12\x1a\n\x12gen_worker_version\x18\x0b \x01(\t\"\xc9\x01\n\nHostCanary\x12\x13\n\x0bmemcpy_gbps\x18\x01 \x01(\x01\x12\x10\n\x08h2d_gbps\x18\x02 \x01(\x01\x12\x10\n\x08\x64\x32h_gbps\x18\x03 \x01(\x01\x12\x17\n\x0fpinned_alloc_ok\x18\x04 \x01(\x08\x12\x17\n\x0f\x63pu_single_mbps\x18\x05 \x01(\x01\x12\x16\n\x0e\x63pu_multi_mbps\x18\x06 \x01(\x01\x12\r\n\x05vcpus\x18\x07 \x01(\x05\x12\x14\n\x0cram_total_gb\x18\x08 \x01(\x01\x12\x13\n\x0b\x64uration_ms\x18\t \x01(\x03\"\x95\x01\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\x12\x17\n\x0fsnapshot_digest\x18\x04 \x01(\t\x12\x1c\n\x14residency_generation\x18\x05 \x01(\x04\"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\"\xb3\x02\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\x12\x36\n\x0f\x63ompile_targets\x18\x07 \x03(\x0b\x32\x1d.cozy.scheduler.CompileTarget\x12\x30\n\x0c\x63\x65ll_lookups\x18\x08 \x03(\x0b\x32\x1a.cozy.scheduler.CellLookup\".\n\nCellLookup\x12\x0e\n\x06\x66\x61mily\x18\x01 \x01(\t\x12\x10\n\x08\x63\x65ll_key\x18\x02 \x01(\t\"\xc6\x03\n\rCompileTarget\x12\x16\n\x0eincarnation_id\x18\x01 \x01(\t\x12\x0e\n\x06\x66\x61mily\x18\x02 \x01(\t\x12\x1c\n\x14pipeline_weight_lane\x18\x03 \x01(\t\x12\x13\n\x0blora_bucket\x18\x04 \x01(\x05\x12\x17\n\x0f\x63ontract_digest\x18\x05 \x01(\t\x12\x1a\n\x12\x61\x63tive_compile_ref\x18\x06 \x01(\t\x12&\n\x1e\x61\x63tive_compile_snapshot_digest\x18\x07 \x01(\t\x12\x16\n\x0e\x66unction_names\x18\x08 \x03(\t\x12<\n\x0emodel_bindings\x18\t \x03(\x0b\x32$.cozy.scheduler.CompileTargetBinding\x12\x1a\n\x12requested_cell_key\x18\n \x01(\t\x12Q\n\x13requested_cell_axes\x18\x0b \x03(\x0b\x32\x34.cozy.scheduler.CompileTarget.RequestedCellAxesEntry\x1a\x38\n\x16RequestedCellAxesEntry\x12\x0b\n\x03key\x18\x01 \x01(\t\x12\r\n\x05value\x18\x02 \x01(\t:\x02\x38\x01\"J\n\x14\x43ompileTargetBinding\x12\x0c\n\x04slot\x18\x01 \x01(\t\x12\x0b\n\x03ref\x18\x02 \x01(\t\x12\x17\n\x0fsnapshot_digest\x18\x03 \x01(\t\"\x88\x04\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\x12\x42\n\x10required_compile\x18\r \x01(\x0b\x32(.cozy.scheduler.RequiredCompileExecution\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\"\x82\x01\n\x18RequiredCompileExecution\x12\x1d\n\x15target_incarnation_id\x18\x01 \x01(\t\x12\x10\n\x08\x63\x65ll_ref\x18\x02 \x01(\t\x12\x1c\n\x14\x63\x65ll_snapshot_digest\x18\x03 \x01(\t\x12\x17\n\x0f\x63ontract_digest\x18\x04 \x01(\t\"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\"\xa0\x01\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\x12\x14\n\x0coperation_id\x18\x04 \x01(\t\x12\x1d\n\x15target_incarnation_id\x18\x05 \x01(\t\"\x9b\x04\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\x12\x1f\n\x17host_ram_required_bytes\x18\x0b \x01(\x03\x12\'\n\x1fhost_ram_available_before_bytes\x18\x0c \x01(\x03\x12&\n\x1ehost_ram_available_after_bytes\x18\r \x01(\x03\x12\x1d\n\x15host_ram_evicted_refs\x18\x0e \x03(\t\x12$\n\x1chost_ram_capacity_generation\x18\x0f \x01(\x04\x12\x17\n\x0fsnapshot_digest\x18\x10 \x01(\t\x12\x1c\n\x14residency_generation\x18\x11 \x01(\x04\x12\x14\n\x0coperation_id\x18\x12 \x01(\t\x12\x1d\n\x15target_incarnation_id\x18\x13 \x01(\t\x12\x15\n\rnetwork_bytes\x18\x14 \x01(\x03\"\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*\x82\x02\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\x12&\n\"MODEL_STATE_HOST_CAPACITY_PROGRESS\x10\x08\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) @@ -40,20 +40,20 @@ _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=6187 - _globals['_PROTOCOLVERSION']._serialized_end=6268 - _globals['_RESIDENCYTIER']._serialized_start=6270 - _globals['_RESIDENCYTIER']._serialized_end=6391 - _globals['_WORKERPHASE']._serialized_start=6394 - _globals['_WORKERPHASE']._serialized_end=6610 - _globals['_OUTPUTMODE']._serialized_start=6612 - _globals['_OUTPUTMODE']._serialized_end=6698 - _globals['_JOBSTATUS']._serialized_start=6701 - _globals['_JOBSTATUS']._serialized_end=6856 - _globals['_MODELOPKIND']._serialized_start=6859 - _globals['_MODELOPKIND']._serialized_end=7026 - _globals['_MODELSTATE']._serialized_start=7029 - _globals['_MODELSTATE']._serialized_end=7287 + _globals['_PROTOCOLVERSION']._serialized_start=6210 + _globals['_PROTOCOLVERSION']._serialized_end=6291 + _globals['_RESIDENCYTIER']._serialized_start=6293 + _globals['_RESIDENCYTIER']._serialized_end=6414 + _globals['_WORKERPHASE']._serialized_start=6417 + _globals['_WORKERPHASE']._serialized_end=6633 + _globals['_OUTPUTMODE']._serialized_start=6635 + _globals['_OUTPUTMODE']._serialized_end=6721 + _globals['_JOBSTATUS']._serialized_start=6724 + _globals['_JOBSTATUS']._serialized_end=6879 + _globals['_MODELOPKIND']._serialized_start=6882 + _globals['_MODELOPKIND']._serialized_end=7049 + _globals['_MODELSTATE']._serialized_start=7052 + _globals['_MODELSTATE']._serialized_end=7310 _globals['_WORKERMESSAGE']._serialized_start=43 _globals['_WORKERMESSAGE']._serialized_end=470 _globals['_SCHEDULERMESSAGE']._serialized_start=473 @@ -117,17 +117,17 @@ _globals['_MODELOP']._serialized_start=5103 _globals['_MODELOP']._serialized_end=5263 _globals['_MODELEVENT']._serialized_start=5266 - _globals['_MODELEVENT']._serialized_end=5782 - _globals['_FNUNAVAILABLE']._serialized_start=5785 - _globals['_FNUNAVAILABLE']._serialized_end=5955 - _globals['_FNUNAVAILABLE_AXESENTRY']._serialized_start=5912 - _globals['_FNUNAVAILABLE_AXESENTRY']._serialized_end=5955 - _globals['_FNDEGRADED']._serialized_start=5958 - _globals['_FNDEGRADED']._serialized_end=6099 - _globals['_DRAIN']._serialized_start=6101 - _globals['_DRAIN']._serialized_end=6129 - _globals['_TOKENREFRESH']._serialized_start=6131 - _globals['_TOKENREFRESH']._serialized_end=6185 - _globals['_WORKERSCHEDULER']._serialized_start=7289 - _globals['_WORKERSCHEDULER']._serialized_end=7386 + _globals['_MODELEVENT']._serialized_end=5805 + _globals['_FNUNAVAILABLE']._serialized_start=5808 + _globals['_FNUNAVAILABLE']._serialized_end=5978 + _globals['_FNUNAVAILABLE_AXESENTRY']._serialized_start=5935 + _globals['_FNUNAVAILABLE_AXESENTRY']._serialized_end=5978 + _globals['_FNDEGRADED']._serialized_start=5981 + _globals['_FNDEGRADED']._serialized_end=6122 + _globals['_DRAIN']._serialized_start=6124 + _globals['_DRAIN']._serialized_end=6152 + _globals['_TOKENREFRESH']._serialized_start=6154 + _globals['_TOKENREFRESH']._serialized_end=6208 + _globals['_WORKERSCHEDULER']._serialized_start=7312 + _globals['_WORKERSCHEDULER']._serialized_end=7409 # @@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 ebfb4766..5e20308d 100644 --- a/src/gen_worker/pb/worker_scheduler_pb2.pyi +++ b/src/gen_worker/pb/worker_scheduler_pb2.pyi @@ -530,7 +530,7 @@ class ModelOp(_message.Message): def __init__(self, op: _Optional[_Union[ModelOpKind, str]] = ..., ref: _Optional[str] = ..., snapshot: _Optional[_Union[Snapshot, _Mapping]] = ..., operation_id: _Optional[str] = ..., target_incarnation_id: _Optional[str] = ...) -> None: ... class ModelEvent(_message.Message): - __slots__ = ("ref", "state", "vram_bytes", "error", "bytes_done", "bytes_total", "duration_ms", "cache_hits", "cache_misses", "warmup_s", "host_ram_required_bytes", "host_ram_available_before_bytes", "host_ram_available_after_bytes", "host_ram_evicted_refs", "host_ram_capacity_generation", "snapshot_digest", "residency_generation", "operation_id", "target_incarnation_id") + __slots__ = ("ref", "state", "vram_bytes", "error", "bytes_done", "bytes_total", "duration_ms", "cache_hits", "cache_misses", "warmup_s", "host_ram_required_bytes", "host_ram_available_before_bytes", "host_ram_available_after_bytes", "host_ram_evicted_refs", "host_ram_capacity_generation", "snapshot_digest", "residency_generation", "operation_id", "target_incarnation_id", "network_bytes") REF_FIELD_NUMBER: _ClassVar[int] STATE_FIELD_NUMBER: _ClassVar[int] VRAM_BYTES_FIELD_NUMBER: _ClassVar[int] @@ -550,6 +550,7 @@ class ModelEvent(_message.Message): RESIDENCY_GENERATION_FIELD_NUMBER: _ClassVar[int] OPERATION_ID_FIELD_NUMBER: _ClassVar[int] TARGET_INCARNATION_ID_FIELD_NUMBER: _ClassVar[int] + NETWORK_BYTES_FIELD_NUMBER: _ClassVar[int] ref: str state: ModelState vram_bytes: int @@ -569,7 +570,8 @@ class ModelEvent(_message.Message): residency_generation: int operation_id: str target_incarnation_id: str - def __init__(self, ref: _Optional[str] = ..., state: _Optional[_Union[ModelState, str]] = ..., vram_bytes: _Optional[int] = ..., error: _Optional[str] = ..., bytes_done: _Optional[int] = ..., bytes_total: _Optional[int] = ..., duration_ms: _Optional[int] = ..., cache_hits: _Optional[int] = ..., cache_misses: _Optional[int] = ..., warmup_s: _Optional[float] = ..., host_ram_required_bytes: _Optional[int] = ..., host_ram_available_before_bytes: _Optional[int] = ..., host_ram_available_after_bytes: _Optional[int] = ..., host_ram_evicted_refs: _Optional[_Iterable[str]] = ..., host_ram_capacity_generation: _Optional[int] = ..., snapshot_digest: _Optional[str] = ..., residency_generation: _Optional[int] = ..., operation_id: _Optional[str] = ..., target_incarnation_id: _Optional[str] = ...) -> None: ... + network_bytes: int + def __init__(self, ref: _Optional[str] = ..., state: _Optional[_Union[ModelState, str]] = ..., vram_bytes: _Optional[int] = ..., error: _Optional[str] = ..., bytes_done: _Optional[int] = ..., bytes_total: _Optional[int] = ..., duration_ms: _Optional[int] = ..., cache_hits: _Optional[int] = ..., cache_misses: _Optional[int] = ..., warmup_s: _Optional[float] = ..., host_ram_required_bytes: _Optional[int] = ..., host_ram_available_before_bytes: _Optional[int] = ..., host_ram_available_after_bytes: _Optional[int] = ..., host_ram_evicted_refs: _Optional[_Iterable[str]] = ..., host_ram_capacity_generation: _Optional[int] = ..., snapshot_digest: _Optional[str] = ..., residency_generation: _Optional[int] = ..., operation_id: _Optional[str] = ..., target_incarnation_id: _Optional[str] = ..., network_bytes: _Optional[int] = ...) -> None: ... class FnUnavailable(_message.Message): __slots__ = ("function_name", "reason", "detail", "axes") diff --git a/tests/test_managed_tier_fill_source.py b/tests/test_managed_tier_fill_source.py new file mode 100644 index 00000000..e648922e --- /dev/null +++ b/tests/test_managed_tier_fill_source.py @@ -0,0 +1,308 @@ +"""th#850 managed-tier ruling (gw#599): the CAS root stays on local/pod disk +as a managed, bounded LRU tier. A RunPod endpoint volume, when attached, is +FILL SOURCE #1 (checked before R2, FILL SOURCE #2); an R2 fill writes +through to the volume so the next same-endpoint pod finds it warm. This +supersedes the CAS-root-on-volume shape +(test_shared_cas_root_multiwriter.py covers that mechanism's multi-writer +temp-file safety, which write-through publishing still relies on). + +Outcome-level tests only, against the real ``ensure_snapshot_async`` CAS +path with the R2 transport stubbed — no mocks of the fill-source mechanism +itself, since it is just filesystem copy+verify. +""" + +from __future__ import annotations + +import asyncio +import os +from pathlib import Path + +from blake3 import blake3 + +import gen_worker.executor as executor_mod +import gen_worker.models.cozy_snapshot as snap_mod +from gen_worker.executor import ModelStore +from gen_worker.models.cache_paths import tensorhub_fill_source_dir +from gen_worker.models.cozy_snapshot import NetworkBytesScope, ensure_snapshot_async +from gen_worker.models.hub_client import WorkerResolvedRepo, WorkerResolvedRepoFile +from gen_worker.models.refs import TensorhubRef +from gen_worker.pb import worker_scheduler_pb2 as pb + +_PAYLOAD = b"managed-tier-fill-source-payload" * 4096 # ~128KB +_BLAKE3 = blake3(_PAYLOAD).hexdigest() +_SNAPSHOT = "c7" * 32 + + +def _resolved() -> WorkerResolvedRepo: + return WorkerResolvedRepo( + snapshot_digest=_SNAPSHOT, + files=[ + WorkerResolvedRepoFile( + path="model.safetensors", + size_bytes=len(_PAYLOAD), + blake3=_BLAKE3, + url="https://tensorhub.invalid/authorized-blob", + ) + ], + ) + + +def _blob_at(cas_root: Path, digest: str) -> Path: + return cas_root / "blobs" / "blake3" / digest[:2] / digest[2:4] / digest + + +def _blob(cas_root: Path) -> Path: + return _blob_at(cas_root, _BLAKE3) + + +def _stub_r2(monkeypatch, calls: list) -> None: + async def _public_get( + _url: str, dst: Path, expected_size: int, expected_blake3: str, + on_bytes=None, + ) -> None: + del expected_size, expected_blake3 + calls.append(1) + dst.write_bytes(_PAYLOAD) + if on_bytes is not None: + on_bytes(len(_PAYLOAD)) + + monkeypatch.setattr(snap_mod, "_download_one_file", _public_get) + + +# --------------------------------------------------------------------------- +# Fill-source ordering (cozy_snapshot layer) +# --------------------------------------------------------------------------- + +def test_volume_blob_preferred_over_r2(tmp_path: Path, monkeypatch) -> None: + calls: list = [] + _stub_r2(monkeypatch, calls) + volume = tmp_path / "volume" + local = tmp_path / "local" + blob = _blob(volume) + blob.parent.mkdir(parents=True, exist_ok=True) + blob.write_bytes(_PAYLOAD) + + ref = TensorhubRef(owner="org", repo="model") + with NetworkBytesScope() as scope: + snap = asyncio.run(ensure_snapshot_async( + base_dir=local, ref=ref, resolved=_resolved(), fill_source_dir=volume, + )) + assert (snap / "model.safetensors").read_bytes() == _PAYLOAD + assert calls == [] # no R2 fetch — the volume already had it + assert scope.network_bytes == 0 + assert _blob(local).read_bytes() == _PAYLOAD # copied into local CAS + + +def test_r2_fetch_writes_through_to_volume(tmp_path: Path, monkeypatch) -> None: + calls: list = [] + _stub_r2(monkeypatch, calls) + volume = tmp_path / "volume" + local = tmp_path / "local" + ref = TensorhubRef(owner="org", repo="model") + + with NetworkBytesScope() as scope: + snap = asyncio.run(ensure_snapshot_async( + base_dir=local, ref=ref, resolved=_resolved(), fill_source_dir=volume, + )) + assert (snap / "model.safetensors").read_bytes() == _PAYLOAD + assert calls == [1] # exactly one R2 fetch + assert scope.network_bytes == len(_PAYLOAD) + assert _blob(volume).read_bytes() == _PAYLOAD # warmed for the next pod + + +def test_no_fill_source_is_byte_identical_to_pre_th850( + tmp_path: Path, monkeypatch, +) -> None: + """cozy-local / no-volume degenerate case: straight to R2, no new branch + taken, no volume path ever touched.""" + calls: list = [] + _stub_r2(monkeypatch, calls) + local = tmp_path / "local" + ref = TensorhubRef(owner="org", repo="model") + + with NetworkBytesScope() as scope: + snap = asyncio.run(ensure_snapshot_async( + base_dir=local, ref=ref, resolved=_resolved(), + )) + assert (snap / "model.safetensors").read_bytes() == _PAYLOAD + assert calls == [1] + assert scope.network_bytes == len(_PAYLOAD) + + +def test_corrupt_volume_blob_falls_through_to_r2(tmp_path: Path, monkeypatch) -> None: + """Revert-turns-red guard: digest-verification of volume-read blobs is + mandatory (Paul's ruling) — a same-SIZE, wrong-content volume blob must + never be silently trusted just because it's the right length.""" + calls: list = [] + _stub_r2(monkeypatch, calls) + volume = tmp_path / "volume" + local = tmp_path / "local" + blob = _blob(volume) + blob.parent.mkdir(parents=True, exist_ok=True) + corrupt = bytes(b ^ 0xFF for b in _PAYLOAD) # same length, different bytes + assert len(corrupt) == len(_PAYLOAD) + blob.write_bytes(corrupt) + + ref = TensorhubRef(owner="org", repo="model") + with NetworkBytesScope() as scope: + snap = asyncio.run(ensure_snapshot_async( + base_dir=local, ref=ref, resolved=_resolved(), fill_source_dir=volume, + )) + assert (snap / "model.safetensors").read_bytes() == _PAYLOAD # real bytes + assert calls == [1] # fell through to R2, not the corrupt volume copy + assert scope.network_bytes == len(_PAYLOAD) + + +# --------------------------------------------------------------------------- +# tensorhub_fill_source_dir(): ismount-guarded, env-driven +# --------------------------------------------------------------------------- + +def test_fill_source_dir_unset_is_none(monkeypatch) -> None: + monkeypatch.delenv("TENSORHUB_FILL_SOURCE_DIR", raising=False) + assert tensorhub_fill_source_dir() is None + + +def test_fill_source_dir_requires_a_real_mount(tmp_path: Path, monkeypatch) -> None: + """A plain directory (baked into the image, or a stray path) must never + be mistaken for the real per-endpoint volume.""" + plain_dir = tmp_path / "not-a-mount" + plain_dir.mkdir() + monkeypatch.setenv("TENSORHUB_FILL_SOURCE_DIR", str(plain_dir)) + assert tensorhub_fill_source_dir() is None # ismount() is False -> rejected + + monkeypatch.setattr(os.path, "ismount", lambda p: str(p) == str(plain_dir)) + assert tensorhub_fill_source_dir() == plain_dir + + +# --------------------------------------------------------------------------- +# Disk-residency network_bytes reaches the wire (executor layer) +# --------------------------------------------------------------------------- + +def test_network_bytes_reaches_on_disk_model_event(tmp_path: Path, monkeypatch) -> None: + calls: list = [] + _stub_r2(monkeypatch, calls) + volume = tmp_path / "volume" + local = tmp_path / "local" + sent: list = [] + + async def _emit(msg: pb.WorkerMessage) -> None: + sent.append(msg) + + store = ModelStore(_emit, cache_dir=local, fill_source_dir=volume) + + async def _run() -> None: + await store.ensure_local( + "org/model", + pb.Snapshot(digest=_SNAPSHOT, files=[ + pb.SnapshotFile( + path="model.safetensors", size_bytes=len(_PAYLOAD), + blake3=_BLAKE3, url="https://tensorhub.invalid/authorized-blob", + ), + ]), + ) + + asyncio.run(_run()) + on_disk = [ + m.model_event for m in sent + if m.WhichOneof("msg") == "model_event" + and m.model_event.state == pb.MODEL_STATE_ON_DISK + ] + assert on_disk, "expected at least one ON_DISK ModelEvent" + # Residency's own generic transition event (network_bytes-blind) and the + # executor's explicit evidence event (network_bytes-carrying) may land + # in either order — protocol-v3 events are observation, not + # convergence, so a receiver reads the most informative one it saw. + assert max(e.network_bytes for e in on_disk) == len(_PAYLOAD) # fetched from R2 + assert _blob(volume).read_bytes() == _PAYLOAD # write-through happened + + # A second, fresh ref whose blob is already warm on the volume reports + # network_bytes == 0 — the "warm boot ⇒ ~0 R2 bytes" signal. + calls.clear() + sent.clear() + payload2 = _PAYLOAD + b"-2" + digest2 = blake3(payload2).hexdigest() + dst2 = _blob_at(volume, digest2) + dst2.parent.mkdir(parents=True, exist_ok=True) + dst2.write_bytes(payload2) + + async def _run2() -> None: + await store.ensure_local( + "org/model2", + pb.Snapshot(digest="d2" * 32, files=[ + pb.SnapshotFile( + path="model.safetensors", size_bytes=len(payload2), + blake3=digest2, url="https://tensorhub.invalid/authorized-blob-2", + ), + ]), + ) + + asyncio.run(_run2()) + on_disk2 = [ + m.model_event for m in sent + if m.WhichOneof("msg") == "model_event" + and m.model_event.state == pb.MODEL_STATE_ON_DISK + ] + assert on_disk2 + assert max(e.network_bytes for e in on_disk2) == 0 + assert calls == [] # no R2 fetch at all — warm from the volume + + +def test_network_bytes_is_a_running_total_on_downloading_ticks( + tmp_path: Path, monkeypatch, +) -> None: + """tensorhub th#850/PR#493 reads network_bytes off the DOWNLOADING + events' running value (WorkerModelDownloadState.NetworkBytes, populated + the same way as BytesDownloaded/BytesTotal), not only the terminal + ON_DISK event — both must carry it or the hub's accounting silently + stays zero. Disable the progress-event debounce so every chunk reaches + the wire, and stream the fake download in chunks (not one write) so a + mid-flight DOWNLOADING event genuinely sees a PARTIAL network_bytes.""" + monkeypatch.setattr(executor_mod, "_PROGRESS_EVENT_MIN_INTERVAL_S", 0.0) + chunk = _PAYLOAD[: len(_PAYLOAD) // 4] + n_chunks = 4 + assert chunk * n_chunks == _PAYLOAD + + async def _chunked_get( + _url: str, dst: Path, expected_size: int, expected_blake3: str, + on_bytes=None, + ) -> None: + del expected_size, expected_blake3 + with open(dst, "wb") as f: + for _ in range(n_chunks): + f.write(chunk) + if on_bytes is not None: + on_bytes(len(chunk)) + await asyncio.sleep(0) # let the executor's callback run + + monkeypatch.setattr(snap_mod, "_download_one_file", _chunked_get) + + local = tmp_path / "local" + sent: list = [] + + async def _emit(msg: pb.WorkerMessage) -> None: + sent.append(msg) + + store = ModelStore(_emit, cache_dir=local) + + async def _run() -> None: + await store.ensure_local( + "org/model", + pb.Snapshot(digest=_SNAPSHOT, files=[ + pb.SnapshotFile( + path="model.safetensors", size_bytes=len(_PAYLOAD), + blake3=_BLAKE3, url="https://tensorhub.invalid/authorized-blob", + ), + ]), + ) + + asyncio.run(_run()) + downloading = [ + m.model_event for m in sent + if m.WhichOneof("msg") == "model_event" + and m.model_event.state == pb.MODEL_STATE_DOWNLOADING + ] + partial = [e.network_bytes for e in downloading if 0 < e.network_bytes < len(_PAYLOAD)] + assert partial, ( + "expected at least one DOWNLOADING event with a PARTIAL running " + f"network_bytes total; got {[e.network_bytes for e in downloading]}" + )