From d0fbaf2ce9ee122dd3df0c32f28f68676c3e281f Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Sun, 19 Jul 2026 22:39:58 -0600 Subject: [PATCH 1/2] gw#597: fix multi-writer CAS races; remove now-superseded shared-tier th#850 moves the endpoint-volume design from a separate shared-read-tier (gw#277) to simply pointing the worker's existing CAS root (TENSORHUB_CACHE_DIR) at the mounted RunPod volume when one is attached. The first pod's ordinary R2 prefetch warms it for later same-endpoint pods -- no seed/publish/read-through split needed, so that plumbing is removed from cache_paths.py/download.py/cozy_snapshot.py. That collapse exposed a real bug: the primary blob-download temp path (_download_one_file_sync) staged to a fixed ".part" name, safe only because a pod-local CAS root has exactly one writer. Once the CAS root is a volume shared by concurrent same-endpoint pods, two pods racing on the same missing digest would collide on that name. Same gap in _materialize_snapshot's ".building" tmp dir. Both now stage to a writer-unique temp name + atomic replace/rename, matching the pattern the (now-removed) shared-tier helper already used correctly. Since a persistent volume (unlike ephemeral pod-local /tmp) keeps a crashed writer's orphaned temp forever, added a boot-time sweep (disk_gc.sweep_stale_writer_temp, wired into ModelStore.rescan_disk) and hardened sweep_orphan_blobs to skip dotfile temp artifacts instead of treating an in-flight writer's temp as an orphaned blob. Proof: 4-process spawn tests racing the real download and materialization paths against one shared destination (test_shared_cas_root_multiwriter.py). executor.py's existing gc_disk/_ensure_disk_headroom LRU disk-GC needed no changes -- it already operates on whatever cache_dir is mounted, satisfying th#850's per-volume size-cap design as-is. Full suite: 1518 passed, 35 skipped, 0 failed. ruff + mypy clean. --- src/gen_worker/executor.py | 9 +- src/gen_worker/models/cache_paths.py | 36 ++-- src/gen_worker/models/cozy_cas.py | 26 ++- src/gen_worker/models/cozy_snapshot.py | 134 ++++--------- src/gen_worker/models/disk_gc.py | 61 +++++- src/gen_worker/models/download.py | 3 +- tests/test_disk_gc.py | 34 ++++ tests/test_endpoint_shared_cas.py | 184 ------------------ tests/test_shared_cas_root_multiwriter.py | 224 ++++++++++++++++++++++ 9 files changed, 393 insertions(+), 318 deletions(-) delete mode 100644 tests/test_endpoint_shared_cas.py create mode 100644 tests/test_shared_cas_root_multiwriter.py diff --git a/src/gen_worker/executor.py b/src/gen_worker/executor.py index cae9c75f..b3841dbc 100644 --- a/src/gen_worker/executor.py +++ b/src/gen_worker/executor.py @@ -911,7 +911,11 @@ def component_sizes(self, ref: str) -> Dict[str, int]: def rescan_disk(self) -> None: """Boot-time truth: re-register still-present downloads from the - persisted ref index so Hello.models and GC see what disk holds.""" + persisted ref index so Hello.models and GC see what disk holds. + + Also sweeps abandoned writer-unique CAS temp artifacts (th#850): on + pod-local disk those died with the pod, but a CAS root pointed at a + persistent volume keeps them until swept.""" for ref, ent in self._index.entries().items(): p = Path(str(ent.get("path") or "")) if p.exists(): @@ -919,6 +923,9 @@ def rescan_disk(self) -> None: self.residency.track_disk(ref, p) else: self._index.remove(ref) + removed = disk_gc.sweep_stale_writer_temp(self._cache_dir) + if removed: + logger.info("disk-gc: swept %d abandoned writer temp artifact(s)", removed) def lru_disk_refs(self, *, exclude: Tuple[str, ...] = ()) -> List[str]: """Idle DISK refs in persisted last-use order, oldest first.""" diff --git a/src/gen_worker/models/cache_paths.py b/src/gen_worker/models/cache_paths.py index 5de713ab..cb60be7a 100644 --- a/src/gen_worker/models/cache_paths.py +++ b/src/gen_worker/models/cache_paths.py @@ -1,6 +1,5 @@ from __future__ import annotations -import os from pathlib import Path from ..config import get_settings @@ -8,20 +7,20 @@ TENSORHUB_CACHE_DIR = "/tmp/tensorhub-cache" -# RunPod network volumes are attached at pod creation. Tensorhub reserves this -# mount solely for an endpoint-owned immutable CAS tier; ordinary pod volumes -# continue using their existing path. This is a provider contract, not an -# operator-tunable cache knob. -ENDPOINT_SHARED_CACHE_MOUNT = Path("/tensorhub-endpoint-cache") - def tensorhub_cache_dir() -> Path: - """TensorHub cache root directory. - - Honors the ``TENSORHUB_CACHE_DIR`` environment variable when set, so the - cozy local runner can point the CAS at a persistent ``~/.cache/tensorhub`` - (weights survive reboots) instead of ``/tmp``. Falls back to the ``/tmp`` - default when unset, preserving existing worker/orchestrator behavior. + """TensorHub cache root directory — the worker's CAS root. + + 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. """ env = get_settings().tensorhub_cache_dir.strip() if env: @@ -32,14 +31,3 @@ def tensorhub_cache_dir() -> Path: def tensorhub_cas_dir() -> Path: """Worker CAS root: /cas.""" return tensorhub_cache_dir() / "cas" - - -def endpoint_shared_cas_dir() -> Path | None: - """Return the endpoint-owned shared CAS only when it is really mounted. - - A plain directory in an image or on the container disk must never be - mistaken for the endpoint-isolated provider volume. - """ - if os.path.ismount(ENDPOINT_SHARED_CACHE_MOUNT): - return ENDPOINT_SHARED_CACHE_MOUNT / "tensorhub-cas-v1" - return None diff --git a/src/gen_worker/models/cozy_cas.py b/src/gen_worker/models/cozy_cas.py index 5b804041..0b67df64 100644 --- a/src/gen_worker/models/cozy_cas.py +++ b/src/gen_worker/models/cozy_cas.py @@ -3,8 +3,11 @@ import asyncio import errno import logging +import os import random +import threading import time +import uuid from pathlib import Path from typing import Callable, Dict, Optional @@ -56,12 +59,20 @@ async def _download_one_file( started = time.monotonic() transient_failures = 0 verify_failures = 0 + # Writer-unique for the lifetime of this call (stable across its own + # retries, so HTTP Range resume-on-retry still works) but distinct from + # every OTHER writer of this digest — including a different PROCESS. + # dst's CAS root may now be a network volume shared by concurrent + # same-endpoint pods; a fixed ".part" name would let two pods racing on + # the same missing blob interleave writes into the same file (th#850). + writer_id = f"{os.getpid()}-{threading.get_ident()}-{uuid.uuid4().hex}" while True: try: await loop.run_in_executor( None, lambda: _download_one_file_sync( - url, dst, expected_size, expected_blake3, on_bytes=on_bytes + url, dst, expected_size, expected_blake3, + on_bytes=on_bytes, writer_id=writer_id, ), ) return @@ -92,14 +103,19 @@ def _download_one_file_sync( expected_size: int, expected_blake3: str, on_bytes: Optional[Callable[[int], None]] = None, + writer_id: Optional[str] = None, ) -> None: """Download a single file with HTTP Range resume, size + blake3 validation. Caller runs this in a worker thread; the transfer itself uses the same requests stack as the Tensorhub control plane fallback path. - """ - import logging + ``writer_id`` must be unique per concurrent writer (distinct calls, and + critically distinct PROCESSES when ``dst``'s CAS root is a network volume + shared by several pods) so two writers of the same missing blob never + stage into the same temp path. Falls back to a fresh one-off id when not + supplied, which loses cross-retry resume but is still collision-safe. + """ log = logging.getLogger("gen_worker.download") def _human_size(n: int) -> str: @@ -125,7 +141,9 @@ def _human_size(n: int) -> str: except Exception: pass # re-download - tmp = dst.with_suffix(dst.suffix + ".part") + if not writer_id: + writer_id = f"{os.getpid()}-{threading.get_ident()}-{uuid.uuid4().hex}" + tmp = dst.parent / f".{dst.name}.part-{writer_id}" # Resume from partial download if available. offset = 0 diff --git a/src/gen_worker/models/cozy_snapshot.py b/src/gen_worker/models/cozy_snapshot.py index b88e6040..2bb8d363 100644 --- a/src/gen_worker/models/cozy_snapshot.py +++ b/src/gen_worker/models/cozy_snapshot.py @@ -8,13 +8,12 @@ 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, fsync_file +from .cozy_cas import _download_one_file as _download_one_file +from .cozy_cas import _norm_rel_path, fsync_dir from .download import components_present, select_component_paths from .hub_client import WorkerResolvedRepo, WorkerResolvedRepoFile from .refs import TensorhubRef @@ -197,7 +196,6 @@ async def ensure_snapshot( base_dir: Path, ref: TensorhubRef, *, - shared_base_dir: Optional[Path] = None, resolved: Optional[WorkerResolvedRepo], progress: Optional[ProgressFn] = None, components: Sequence[str] = (), @@ -249,15 +247,9 @@ async def ensure_snapshot( try: _log.info("snapshot_build_start key=%s files=%d", key[:24], len(res.files)) - shared_blobs_root = None - if shared_base_dir is not None: - candidate = Path(shared_base_dir) / "blobs" - if candidate.resolve() != blobs_root.resolve(): - shared_blobs_root = candidate await self._ensure_blobs( blobs_root, res.files, - shared_blobs_root=shared_blobs_root, progress=progress, ) # Materialization copies/concatenates multi-GB trees — strictly @@ -290,16 +282,20 @@ def _materialize_snapshot( res: WorkerResolvedRepo, ) -> None: """Blocking build phase (worker thread): reassemble + hardlink into a - ``.building`` dir, then atomically rename into place.""" - tmp = snaps_root / f"{snap_dir.name}.building" - if tmp.exists(): - shutil.rmtree(tmp) + 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 + the winner.""" + writer_id = f"{os.getpid()}-{threading.get_ident()}-{uuid.uuid4().hex}" + tmp = snaps_root / f"{snap_dir.name}.building-{writer_id}" tmp.mkdir(parents=True, exist_ok=True) self._reassemble_chunked(blobs_root, tmp, res.files) self._materialize_regular(blobs_root, tmp, res.files) - # Atomic rename; handle race with concurrent builder. + # Atomic rename; a concurrent writer may have already published. if snap_dir.exists(): shutil.rmtree(tmp, ignore_errors=True) else: @@ -310,8 +306,6 @@ def _materialize_snapshot( if not snap_dir.exists(): raise else: - from .cozy_cas import fsync_dir - fsync_dir(snaps_root) # persist the rename itself (gw#408) # ------------------------------------------------------------------ @@ -323,7 +317,6 @@ async def _ensure_blobs( blobs_root: Path, files: List[WorkerResolvedRepoFile], *, - shared_blobs_root: Optional[Path] = None, progress: Optional[ProgressFn] = None, ) -> None: # Deduplicate by digest — same blob referenced by multiple paths (e.g. @@ -354,12 +347,18 @@ async def _ensure_blobs( total = sum(int(f.size_bytes or 0) for f in unique) or None done = total - missing_bytes if total else 0 + network_bytes = 0 # th#850: bytes actually fetched over the network + # this call, as opposed to bytes already present under blobs_root — + # the signal a "volume-attached boot ⇒ ~0 network bytes" runtime + # assertion needs. Only _dl_locked's real fetch increments it. done_lock = threading.Lock() - def _on_bytes(n: int) -> None: - nonlocal done + def _on_bytes(n: int, *, network: bool = False) -> None: + nonlocal done, network_bytes with done_lock: done += n + if network: + network_bytes += n d = done if total is None else min(done, total) if progress is not None: try: @@ -391,28 +390,7 @@ 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 shared_blobs_root is not None: - shared = _blob_path(shared_blobs_root, digest) - started = time.monotonic() - if await asyncio.to_thread(self._copy_verified_blob, shared, dst, f): - _on_bytes(int(f.size_bytes or 0)) - _log.info( - "blob_cache_hit source=endpoint_volume digest=%s bytes=%d transfer_ms=%d", - digest[:16], int(f.size_bytes or 0), - int((time.monotonic() - started) * 1000), - ) - return - _log.info("blob_cache_miss source=endpoint_volume digest=%s", digest[:16]) await _dl_locked(f, digest, dst) - if shared_blobs_root is not None: - shared = _blob_path(shared_blobs_root, digest) - published = await asyncio.to_thread( - self._copy_verified_blob, dst, shared, f - ) - _log.info( - "blob_cache_publish source=public destination=endpoint_volume digest=%s bytes=%d published=%s", - digest[:16], int(f.size_bytes or 0), published, - ) async def _dl_locked(f: WorkerResolvedRepoFile, digest: str, dst: Path) -> None: async with sem: @@ -431,7 +409,7 @@ async def _dl_locked(f: WorkerResolvedRepoFile, digest: str, dst: Path) -> None: expected_blake3=digest, ), ) - _on_bytes(int(f.size_bytes or 0)) + _on_bytes(int(f.size_bytes or 0), network=True) else: assert f.url is not None # validated above in _ensure_blobs loop await _download_one_file( @@ -439,67 +417,21 @@ async def _dl_locked(f: WorkerResolvedRepoFile, digest: str, dst: Path) -> None: dst, expected_size=int(f.size_bytes or 0), expected_blake3=digest, - on_bytes=_on_bytes, + on_bytes=lambda n: _on_bytes(n, network=True), ) _log.info("blob_download_done path=%s digest=%s", f.path, digest[:16]) await asyncio.gather(*(_dl(f) for f in unique)) - - @staticmethod - def _copy_verified_blob( - src: Path, dst: Path, f: WorkerResolvedRepoFile, - ) -> bool: - """Copy one immutable blob through a writer-unique atomic stage. - - The final path is digest-only. Readers never observe partial bytes; - racing writers may replace the same final name only after each has - independently verified the declared size and BLAKE3. - """ - expected_size = int(f.size_bytes or 0) - expected_digest = _strip_blake3_prefix(f.blake3) - try: - if not src.is_file(): - return False - if expected_size and src.stat().st_size != expected_size: - _log.warning( - "blob_cache_corrupt source=%s digest=%s reason=size expected=%d actual=%d", - src, expected_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_cache_corrupt source=%s digest=%s reason=staged_size expected=%d actual=%d", - src, expected_digest[:16], expected_size, tmp.stat().st_size, - ) - return False - if expected_digest and _blake3_file(tmp).lower() != expected_digest.lower(): - _log.warning( - "blob_cache_corrupt source=%s digest=%s reason=blake3", - src, expected_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_cache_copy_failed source=%s destination=%s digest=%s error=%s", - src, dst, expected_digest[:16], exc, - ) - return False + # 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). + 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", + total, network_bytes, cache_hit_pct, + ) @staticmethod def _blob_usable(dst: Path, f: WorkerResolvedRepoFile) -> bool: @@ -626,7 +558,6 @@ def delete_blobs(base_dir: Path, digests: Any) -> None: async def ensure_snapshot_async( *, base_dir: Path, - shared_base_dir: Optional[Path] = None, ref: TensorhubRef, resolved: Optional[WorkerResolvedRepo], progress: Optional[ProgressFn] = None, @@ -634,6 +565,5 @@ async def ensure_snapshot_async( ) -> Path: dl = CozySnapshotDownloader() return await dl.ensure_snapshot( - base_dir, ref, shared_base_dir=shared_base_dir, - resolved=resolved, progress=progress, components=components, + base_dir, ref, resolved=resolved, progress=progress, components=components, ) diff --git a/src/gen_worker/models/disk_gc.py b/src/gen_worker/models/disk_gc.py index e7ba35b9..71776b25 100644 --- a/src/gen_worker/models/disk_gc.py +++ b/src/gen_worker/models/disk_gc.py @@ -215,13 +215,21 @@ def delete_ref_bytes(ref: str, path: Path, cas_dir: Path) -> None: def sweep_orphan_blobs(cas_dir: Path) -> int: """Delete CAS blobs no snapshot links anymore (st_nlink == 1). Snapshot - trees hardlink into blobs/, so link count is the reference count.""" + trees hardlink into blobs/, so link count is the reference count. + + Skips dotfiles: writer-unique in-flight temp artifacts + (``..part-``, see cozy_cas.py) live in this same + directory tree and always have nlink==1 while a download is in + progress — treating them as orphans would delete a live writer's bytes + out from under it.""" blobs = Path(cas_dir) / "blobs" freed = 0 if not blobs.is_dir(): return 0 for dirpath, _dirs, names in os.walk(blobs): for name in names: + if name.startswith("."): + continue fp = os.path.join(dirpath, name) try: st = os.stat(fp) @@ -233,10 +241,61 @@ def sweep_orphan_blobs(cas_dir: Path) -> int: return freed +# Generous: the largest blobs can legitimately take hours on a slow link. +# Only genuinely abandoned (crashed/killed writer) temp artifacts are this +# old — a live writer keeps rewriting its temp, advancing its mtime. +_STALE_WRITER_TEMP_AGE_S = 6 * 3600 + + +def sweep_stale_writer_temp( + cas_dir: Path, *, older_than_s: float = _STALE_WRITER_TEMP_AGE_S, +) -> int: + """Remove abandoned writer-unique temp artifacts: blob downloads stage to + ``..part-`` and snapshot builds stage to + ``.building-`` (writer-unique so concurrent writers on a + SHARED CAS root — several pods on one RunPod volume, th#850 — never + collide). A writer that dies mid-transfer leaves its temp behind; on + pod-local disk that vanished with the pod, but a persistent volume keeps + it forever unless swept. Call at boot (`ModelStore.rescan_disk`) — safe + at any time since only artifacts idle past ``older_than_s`` are removed. + """ + removed = 0 + root = Path(cas_dir) + now = time.time() + for base_name in ("blobs", "snapshots"): + base = root / base_name + if not base.is_dir(): + continue + for dirpath, dirnames, filenames in os.walk(base): + for name in list(dirnames): + if ".building-" not in name: + continue + p = Path(dirpath) / name + try: + if now - p.stat().st_mtime > older_than_s: + shutil.rmtree(p, ignore_errors=True) + removed += 1 + dirnames.remove(name) + except OSError: + continue + for name in filenames: + if ".part-" not in name: + continue + p = Path(dirpath) / name + try: + if now - p.stat().st_mtime > older_than_s: + p.unlink(missing_ok=True) + removed += 1 + except OSError: + continue + return removed + + __all__ = [ "RefIndex", "tree_bytes", "reclaim_file_cache", "delete_ref_bytes", "sweep_orphan_blobs", + "sweep_stale_writer_temp", ] diff --git a/src/gen_worker/models/download.py b/src/gen_worker/models/download.py index bb1177b7..6654dd43 100644 --- a/src/gen_worker/models/download.py +++ b/src/gen_worker/models/download.py @@ -29,7 +29,7 @@ from ..api.errors import ValidationError from ..config import get_settings -from .cache_paths import endpoint_shared_cas_dir, tensorhub_cas_dir +from .cache_paths import tensorhub_cas_dir from .refs import HuggingFaceRef, TensorhubRef, fold_ref, parse_model_ref if TYPE_CHECKING: @@ -232,7 +232,6 @@ async def ensure_local( return await ensure_snapshot_async( base_dir=base, - shared_base_dir=endpoint_shared_cas_dir(), ref=_snapshot_ref(parsed, ref), resolved=snapshot, progress=progress, diff --git a/tests/test_disk_gc.py b/tests/test_disk_gc.py index 0b613c53..764801f3 100644 --- a/tests/test_disk_gc.py +++ b/tests/test_disk_gc.py @@ -222,3 +222,37 @@ async def _run() -> None: assert restarted.residency.tier("t/a") is Tier.DISK assert restarted.residency.local_path("t/a") == tmp_path / "snapshots" / "t--a" assert restarted.residency.tier("t/b") is None # stale entry dropped + + +def test_rescan_sweeps_stale_writer_temp_artifacts(tmp_path, _fake_download) -> None: + """th#850: a CAS root on a persistent volume (unlike ephemeral pod-local + disk) keeps a crashed writer's temp artifacts forever unless boot-time + rescan sweeps them. Fresh/live artifacts must survive the sweep.""" + from gen_worker.models import disk_gc + + store = _store(tmp_path, []) + blob_dir = tmp_path / "blobs" / "blake3" / "ab" / "cd" + blob_dir.mkdir(parents=True) + stale_blob_tmp = blob_dir / ".deadbeef.part-stale-writer" + fresh_blob_tmp = blob_dir / ".deadbeef.part-live-writer" + stale_blob_tmp.write_bytes(b"partial") + fresh_blob_tmp.write_bytes(b"partial") + + snaps_root = tmp_path / "snapshots" + stale_building = snaps_root / "abc123.building-stale-writer" + fresh_building = snaps_root / "abc123.building-live-writer" + stale_building.mkdir(parents=True) + fresh_building.mkdir(parents=True) + (stale_building / "w.bin").write_bytes(b"x") + + old = time.time() - disk_gc._STALE_WRITER_TEMP_AGE_S - 1 + import os + os.utime(stale_blob_tmp, (old, old)) + os.utime(stale_building, (old, old)) + + store.rescan_disk() + + assert not stale_blob_tmp.exists() + assert fresh_blob_tmp.exists() + assert not stale_building.exists() + assert fresh_building.exists() diff --git a/tests/test_endpoint_shared_cas.py b/tests/test_endpoint_shared_cas.py deleted file mode 100644 index 44070d0d..00000000 --- a/tests/test_endpoint_shared_cas.py +++ /dev/null @@ -1,184 +0,0 @@ -from __future__ import annotations - -import asyncio -import multiprocessing -from pathlib import Path -from typing import Any - -from blake3 import blake3 - -import gen_worker.models.cozy_snapshot as snap_mod -import gen_worker.models.cache_paths as cache_paths -from gen_worker.models.cozy_snapshot import CozySnapshotDownloader, ensure_snapshot_async -from gen_worker.models.hub_client import WorkerResolvedRepo, WorkerResolvedRepoFile -from gen_worker.models.refs import TensorhubRef - - -_PAYLOAD = b"endpoint-isolated-model-weights" * 1024 -_BLAKE3 = blake3(_PAYLOAD).hexdigest() -_SNAPSHOT = "a5" * 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(base: Path) -> Path: - return base / "blobs" / "blake3" / _BLAKE3[:2] / _BLAKE3[2:4] / _BLAKE3 - - -def _copy_blob_process( - start: Any, - results: Any, - source: str, - destination: str, - size_bytes: int, - digest: str, -) -> None: - """Run the production CAS publisher in an independently spawned process.""" - if not start.wait(10): - results.put((False, "start barrier timed out")) - return - try: - copied = CozySnapshotDownloader._copy_verified_blob( - Path(source), - Path(destination), - WorkerResolvedRepoFile( - path="model.safetensors", - size_bytes=size_bytes, - blake3=digest, - url="https://tensorhub.invalid/authorized-blob", - ), - ) - results.put((copied, "")) - except BaseException as exc: # pragma: no cover - returned to the parent for diagnosis - results.put((False, repr(exc))) - - -def test_shared_cache_is_enabled_only_for_real_provider_mount(monkeypatch) -> None: - monkeypatch.setattr(cache_paths.os.path, "ismount", lambda _path: False) - assert cache_paths.endpoint_shared_cas_dir() is None - - monkeypatch.setattr(cache_paths.os.path, "ismount", lambda path: path == cache_paths.ENDPOINT_SHARED_CACHE_MOUNT) - assert cache_paths.endpoint_shared_cas_dir() == ( - cache_paths.ENDPOINT_SHARED_CACHE_MOUNT / "tensorhub-cas-v1" - ) - - -def test_second_pod_reuses_verified_endpoint_volume_without_public_get( - tmp_path: Path, monkeypatch, -) -> None: - calls = 0 - - async def _public_get( - _url: str, dst: Path, expected_size: int, expected_blake3: str, - on_bytes=None, - ) -> None: - del expected_size, expected_blake3 - nonlocal calls - calls += 1 - dst.write_bytes(_PAYLOAD) - if on_bytes is not None: - on_bytes(len(_PAYLOAD)) - - monkeypatch.setattr(snap_mod, "_download_one_file", _public_get) - ref = TensorhubRef(owner="org", repo="model") - shared = tmp_path / "endpoint-volume" - - first = asyncio.run(ensure_snapshot_async( - base_dir=tmp_path / "pod-a", shared_base_dir=shared, - ref=ref, resolved=_resolved(), - )) - assert (first / "model.safetensors").read_bytes() == _PAYLOAD - assert _blob(shared).read_bytes() == _PAYLOAD - - second = asyncio.run(ensure_snapshot_async( - base_dir=tmp_path / "pod-b", shared_base_dir=shared, - ref=ref, resolved=_resolved(), - )) - assert (second / "model.safetensors").read_bytes() == _PAYLOAD - assert calls == 1 - - -def test_corrupt_endpoint_volume_blob_falls_back_and_atomically_repairs( - tmp_path: Path, monkeypatch, -) -> None: - calls = 0 - - async def _public_get( - _url: str, dst: Path, expected_size: int, expected_blake3: str, - on_bytes=None, - ) -> None: - del expected_size, expected_blake3 - nonlocal calls - calls += 1 - dst.write_bytes(_PAYLOAD) - if on_bytes is not None: - on_bytes(len(_PAYLOAD)) - - monkeypatch.setattr(snap_mod, "_download_one_file", _public_get) - shared = tmp_path / "endpoint-volume" - corrupt = _blob(shared) - corrupt.parent.mkdir(parents=True) - corrupt.write_bytes(b"x" * len(_PAYLOAD)) - - out = asyncio.run(ensure_snapshot_async( - base_dir=tmp_path / "pod", shared_base_dir=shared, - ref=TensorhubRef(owner="org", repo="model"), resolved=_resolved(), - )) - - assert calls == 1 - assert (out / "model.safetensors").read_bytes() == _PAYLOAD - assert corrupt.read_bytes() == _PAYLOAD - assert not list(corrupt.parent.glob(f".{_BLAKE3}.writer-*")) - - -def test_spawned_process_writers_publish_only_complete_verified_blob(tmp_path: Path) -> None: - payload = _PAYLOAD * 128 - digest = blake3(payload).hexdigest() - shared = ( - tmp_path / "endpoint-volume" / "blobs" / "blake3" / digest[:2] / digest[2:4] - / digest - ) - sources = [tmp_path / f"pod-{i}.blob" for i in range(4)] - for source in sources: - source.write_bytes(payload) - - ctx = multiprocessing.get_context("spawn") - start = ctx.Event() - results = ctx.Queue() - processes = [ - ctx.Process( - target=_copy_blob_process, - args=(start, results, str(source), str(shared), len(payload), digest), - ) - for source in sources - ] - try: - for process in processes: - process.start() - start.set() - for process in processes: - process.join(20) - assert [process.exitcode for process in processes] == [0] * len(processes) - outcomes = [results.get(timeout=5) for _ in processes] - finally: - for process in processes: - if process.is_alive(): - process.terminate() - process.join(5) - results.close() - - assert outcomes == [(True, "")] * len(processes) - assert shared.read_bytes() == payload - assert not list(shared.parent.glob(f".{digest}.writer-*")) diff --git a/tests/test_shared_cas_root_multiwriter.py b/tests/test_shared_cas_root_multiwriter.py new file mode 100644 index 00000000..cf6aee54 --- /dev/null +++ b/tests/test_shared_cas_root_multiwriter.py @@ -0,0 +1,224 @@ +"""th#850: the worker's CAS root can now be a RunPod network volume shared by +several same-endpoint pods concurrently, instead of always being pod-local +disk. There is no separate shared/read-through tier (gw PR #277's design was +superseded) — pointing ``TENSORHUB_CACHE_DIR`` at the mount IS the whole +mechanism, so the ordinary CAS write path must itself be safe when several +independent OS processes race on it. These tests prove that with REAL +multiprocessing (spawn), not mocked concurrency, because the bug class this +guards against (a shared, non-writer-unique temp filename) only manifests +across process boundaries. +""" + +from __future__ import annotations + +import asyncio +import http.server +import multiprocessing +import threading +from pathlib import Path +from typing import Any + +from blake3 import blake3 + +import gen_worker.models.cozy_snapshot as snap_mod +from gen_worker.models.cozy_cas import _download_one_file +from gen_worker.models.cozy_snapshot import ensure_snapshot_async +from gen_worker.models.hub_client import WorkerResolvedRepo, WorkerResolvedRepoFile +from gen_worker.models.refs import TensorhubRef + +_PAYLOAD = b"endpoint-volume-cas-root-payload" * 4096 # ~128KB +_BLAKE3 = blake3(_PAYLOAD).hexdigest() +_SNAPSHOT = "b6" * 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(base: Path) -> Path: + return base / "blobs" / "blake3" / _BLAKE3[:2] / _BLAKE3[2:4] / _BLAKE3 + + +# --------------------------------------------------------------------------- +# Design proof: CAS-root-on-volume needs no separate tier — a second "pod" +# pointed at the SAME base_dir just finds the blob already there. +# --------------------------------------------------------------------------- + +def test_second_pod_on_shared_cas_root_makes_no_network_call( + tmp_path: Path, monkeypatch, +) -> None: + calls = 0 + + async def _public_get( + _url: str, dst: Path, expected_size: int, expected_blake3: str, + on_bytes=None, + ) -> None: + del expected_size, expected_blake3 + nonlocal calls + calls += 1 + dst.write_bytes(_PAYLOAD) + if on_bytes is not None: + on_bytes(len(_PAYLOAD)) + + monkeypatch.setattr(snap_mod, "_download_one_file", _public_get) + ref = TensorhubRef(owner="org", repo="model") + shared_root = tmp_path / "volume" # what TENSORHUB_CACHE_DIR/cas resolves to + + first = asyncio.run(ensure_snapshot_async( + base_dir=shared_root, ref=ref, resolved=_resolved(), + )) + assert (first / "model.safetensors").read_bytes() == _PAYLOAD + assert calls == 1 + + # Second "pod" boots against the exact same CAS root (the volume): a + # fresh ref with a different snapshot key still hits the warmed blob. + ref2 = TensorhubRef(owner="org", repo="model2") + second = asyncio.run(ensure_snapshot_async( + base_dir=shared_root, ref=ref2, resolved=_resolved(), + )) + assert (second / "model.safetensors").read_bytes() == _PAYLOAD + assert calls == 1 # no second network fetch — the blob was already warm + + +# --------------------------------------------------------------------------- +# Multi-writer safety: real OS processes racing the PRIMARY download path +# against one shared destination (simulating several pods on one volume). +# --------------------------------------------------------------------------- + +class _PayloadHandler(http.server.BaseHTTPRequestHandler): + payload = _PAYLOAD + + def do_GET(self) -> None: # noqa: N802 - stdlib API + self.send_response(200) + self.send_header("Content-Length", str(len(self.payload))) + self.end_headers() + self.wfile.write(self.payload) + + def log_message(self, *_a: Any) -> None: # silence + pass + + +def _serve() -> tuple[http.server.ThreadingHTTPServer, str]: + httpd = http.server.ThreadingHTTPServer(("127.0.0.1", 0), _PayloadHandler) + threading.Thread(target=httpd.serve_forever, daemon=True).start() + return httpd, f"http://127.0.0.1:{httpd.server_address[1]}" + + +def _download_worker(start: Any, results: Any, url: str, dst: str) -> None: + if not start.wait(10): + results.put((False, "start barrier timed out")) + return + try: + asyncio.run(_download_one_file( + url, Path(dst), expected_size=len(_PAYLOAD), expected_blake3=_BLAKE3, + )) + results.put((True, "")) + except BaseException as exc: # pragma: no cover - surfaced to parent + results.put((False, repr(exc))) + + +def test_concurrent_processes_racing_same_missing_blob_do_not_corrupt( + tmp_path: Path, +) -> None: + httpd, base_url = _serve() + try: + dst = tmp_path / "blobs" / "blake3" / _BLAKE3[:2] / _BLAKE3[2:4] / _BLAKE3 + dst.parent.mkdir(parents=True) + url = f"{base_url}/blob" + + ctx = multiprocessing.get_context("spawn") + start = ctx.Event() + results = ctx.Queue() + n = 4 + processes = [ + ctx.Process(target=_download_worker, args=(start, results, url, str(dst))) + for _ in range(n) + ] + try: + for p in processes: + p.start() + start.set() + for p in processes: + p.join(30) + assert [p.exitcode for p in processes] == [0] * n + outcomes = [results.get(timeout=5) for _ in processes] + finally: + for p in processes: + if p.is_alive(): + p.terminate() + p.join(5) + results.close() + + assert outcomes == [(True, "")] * n + assert dst.read_bytes() == _PAYLOAD + # No writer-unique temp artifacts left behind by any racing writer. + assert not list(dst.parent.glob(f".{dst.name}.part-*")) + finally: + httpd.shutdown() + + +# --------------------------------------------------------------------------- +# Multi-writer safety: real OS processes racing snapshot MATERIALIZATION +# (blob already present; the race is purely on the .building tmp dir). +# --------------------------------------------------------------------------- + +def _materialize_worker(start: Any, results: Any, base_dir: str) -> None: + if not start.wait(10): + results.put((False, "start barrier timed out")) + return + try: + snap_dir = asyncio.run(ensure_snapshot_async( + base_dir=Path(base_dir), + ref=TensorhubRef(owner="org", repo="model"), + resolved=_resolved(), + )) + content = (snap_dir / "model.safetensors").read_bytes() + results.put((content == _PAYLOAD, "")) + except BaseException as exc: # pragma: no cover - surfaced to parent + results.put((False, repr(exc))) + + +def test_concurrent_processes_racing_same_snapshot_materialization( + tmp_path: Path, +) -> None: + base_dir = tmp_path / "volume" + blob = _blob(base_dir) + blob.parent.mkdir(parents=True) + blob.write_bytes(_PAYLOAD) # pre-warmed: this race is materialization-only + + ctx = multiprocessing.get_context("spawn") + start = ctx.Event() + results = ctx.Queue() + n = 4 + processes = [ + ctx.Process(target=_materialize_worker, args=(start, results, str(base_dir))) + for _ in range(n) + ] + try: + for p in processes: + p.start() + start.set() + for p in processes: + p.join(30) + assert [p.exitcode for p in processes] == [0] * n + outcomes = [results.get(timeout=5) for _ in processes] + finally: + for p in processes: + if p.is_alive(): + p.terminate() + p.join(5) + results.close() + + assert outcomes == [(True, "")] * n + snaps_root = base_dir / "snapshots" + assert not list(snaps_root.glob(f"{_SNAPSHOT}.building-*")) From 1e1149b21225a7f7c469016d01ae213ac88397c5 Mon Sep 17 00:00:00 2001 From: Paul Fidika Date: Sun, 19 Jul 2026 23:04:47 -0600 Subject: [PATCH 2/2] gw#598: digest-verify reused CAS state at the CAS layer, once per root per boot The fresh-materialization path trusted bytes it did not download: an already-materialized snap_dir short-circuited with no check at all, and _blob_usable reused existing blobs on a size-only check. Safe when the CAS root's only writer was this pod; on a th#850 shared persistent volume those are another pod's bytes of any age. Fix lives in the CAS layer (cozy_snapshot), not the executor -- the reuse decision and its verification belong together, stubbed-download tests are untouched, and genuinely fresh downloads are never hashed twice. Mechanics: process-local trust sets keyed by (resolved root, digest/key). Reused blobs get a full BLAKE3 once per root+digest per process (_blob_trusted; corrupt -> unlink -> re-download); a reused materialized tree is verified once per root+key per process (_verify_materialized_tree: manifest size+blake3 + structural safetensors for reassembled originals; corrupt -> quarantine tree + bad blobs -> rebuild). A freshly built tree is trusted without a second pass since every input was verified this process. The executor's existing first-use-per-boot verify for residency-tracked refs is unchanged and now backed by the same guarantee underneath. Chosen tradeoff (recorded in gw#598): verify-on-first-use-per-process, NOT re-hash-every-boot -- warm-volume boot pays one hash pass over the model (the intended defense), cold boots pay nothing extra. Test: test_preseeded_unverified_cas_bytes_are_verified_before_trust -- corrupt right-size blob hardlinked into a materialized tree for a ref residency never saw; ensure_local must detect, quarantine, re-download, and repair the shared blob. test_disk_gc fixtures now declare truthful digests of their fake bytes. Full suite 1519 passed / 0 failed. --- src/gen_worker/models/cozy_snapshot.py | 138 +++++++++++++++++++++++-- tests/test_disk_gc.py | 8 +- tests/test_snapshot_corruption.py | 32 ++++++ 3 files changed, 171 insertions(+), 7 deletions(-) diff --git a/src/gen_worker/models/cozy_snapshot.py b/src/gen_worker/models/cozy_snapshot.py index 2bb8d363..fdc24166 100644 --- a/src/gen_worker/models/cozy_snapshot.py +++ b/src/gen_worker/models/cozy_snapshot.py @@ -12,7 +12,7 @@ from pathlib import Path from typing import Any, Callable, Dict, List, Optional, Sequence, Set -from .cozy_cas import _download_one_file as _download_one_file +from .cozy_cas import _blake3_file, _download_one_file as _download_one_file from .cozy_cas import _norm_rel_path, fsync_dir from .download import components_present, select_component_paths from .hub_client import WorkerResolvedRepo, WorkerResolvedRepoFile @@ -39,6 +39,24 @@ def _inflight_blob_lock(digest: str) -> "asyncio.Lock": _INFLIGHT_BLOB_LOCKS.clear() # bounded memory; races just re-lock return _INFLIGHT_BLOB_LOCKS.setdefault(digest, 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. +_TRUST_CAP = 65536 +_VERIFIED_BLOBS: Set[tuple] = set() +_TRUSTED_SNAPSHOTS: Set[tuple] = set() + + +def _mark_trusted(s: Set[tuple], item: tuple) -> None: + if len(s) > _TRUST_CAP: + s.clear() # bounded memory; worst case is a re-verify + s.add(item) + ProgressFn = Callable[[int, Optional[int]], None] # Free space that must remain after downloading the missing blobs. @@ -221,14 +239,15 @@ async def ensure_snapshot( # snapshot). key = snapshot_dir_key(res.snapshot_digest, components) snap_dir = snaps_root / key - if snap_dir.exists(): + trust_key = (str(snaps_root.resolve()), key) + if snap_dir.exists() and trust_key in _TRUSTED_SNAPSHOTS: _log.info("snapshot_cached key=%s", key[:24]) return snap_dir # Coordinate concurrent builders via threading (works across event loops). loop = asyncio.get_running_loop() with _SNAP_LOCK: - if snap_dir.exists(): + if snap_dir.exists() and trust_key in _TRUSTED_SNAPSHOTS: return snap_dir entry = _SNAP_ENTRIES.get(key) if entry is None: @@ -246,6 +265,25 @@ async def ensure_snapshot( return snap_dir try: + if snap_dir.exists(): + # gw#598: a materialized tree this process has not produced or + # checked (another pod's writes on a shared volume root, or a + # previous boot's) is verified ONCE before reuse; corruption + # quarantines tree + bad blobs and falls through to a rebuild. + ok, bad = await asyncio.to_thread( + _verify_materialized_tree, snap_dir, res.files + ) + if ok: + _mark_trusted(_TRUSTED_SNAPSHOTS, trust_key) + _log.info("snapshot_cached key=%s (verified first use)", key[:24]) + return snap_dir + _log.warning( + "snapshot_reuse_corrupt key=%s bad_files=%d; quarantining and rebuilding", + key[:24], len(bad), + ) + await asyncio.to_thread( + _quarantine_materialized, snap_dir, blobs_root, bad + ) _log.info("snapshot_build_start key=%s files=%d", key[:24], len(res.files)) await self._ensure_blobs( blobs_root, @@ -259,6 +297,10 @@ async def ensure_snapshot( await asyncio.to_thread( self._materialize_snapshot, blobs_root, snaps_root, snap_dir, res ) + # Every input was verified this process (fresh downloads by the + # downloader, reused blobs by _blob_trusted), so the freshly + # materialized tree is trustworthy without a second hash pass. + _mark_trusted(_TRUSTED_SNAPSHOTS, trust_key) _log.info("snapshot_build_done key=%s", key[:24]) return snap_dir except BaseException as exc: @@ -377,16 +419,38 @@ def _on_bytes(n: int, *, network: bool = False) -> None: max_conc = 4 sem = asyncio.Semaphore(max_conc) + blobs_root_id = str(blobs_root.resolve()) + + async def _blob_trusted(dst: Path, f: WorkerResolvedRepoFile, digest: str) -> bool: + """Reusable AND digest-trusted (gw#598): size gate as before, plus + a full BLAKE3 check once per (root, digest) per process — reused + bytes on a shared volume root are another pod's writes of any age + and must earn the same trust as a fresh verified download.""" + if not self._blob_usable(dst, f): + return False + vkey = (blobs_root_id, digest) + if vkey in _VERIFIED_BLOBS: + return True + got = await asyncio.to_thread(_blake3_file, dst) + if got.lower() == digest: + _mark_trusted(_VERIFIED_BLOBS, vkey) + return True + _log.warning( + "blob_corrupt path=%s digest=%s blake3 mismatch on reuse; re-downloading", + f.path, digest[:16], + ) + dst.unlink(missing_ok=True) + return False async def _dl(f: WorkerResolvedRepoFile) -> None: digest = f.blake3.strip().lower() dst = _blob_path(blobs_root, digest) dst.parent.mkdir(parents=True, exist_ok=True) - if self._blob_usable(dst, f): + if await _blob_trusted(dst, f, digest): _log.info("blob_cached path=%s digest=%s", f.path, digest[:16]) return async with _inflight_blob_lock(digest): - if self._blob_usable(dst, f): + if await _blob_trusted(dst, f, digest): _log.info("blob_shared_inflight path=%s digest=%s (sibling ref downloaded it)", f.path, digest[:16]) return @@ -394,7 +458,7 @@ async def _dl(f: WorkerResolvedRepoFile) -> None: async def _dl_locked(f: WorkerResolvedRepoFile, digest: str, dst: Path) -> None: async with sem: - if self._blob_usable(dst, f): + if await _blob_trusted(dst, f, digest): return _log.info("blob_download_start path=%s size=%s digest=%s", f.path, f.size_bytes, digest[:16]) @@ -419,6 +483,8 @@ async def _dl_locked(f: WorkerResolvedRepoFile, digest: str, dst: Path) -> None: expected_blake3=digest, on_bytes=lambda n: _on_bytes(n, network=True), ) + # Both download paths verified size+BLAKE3 before publishing. + _mark_trusted(_VERIFIED_BLOBS, (blobs_root_id, digest)) _log.info("blob_download_done path=%s digest=%s", f.path, digest[:16]) await asyncio.gather(*(_dl(f) for f in unique)) @@ -536,6 +602,66 @@ def _materialize_regular( +def _verify_materialized_tree( + snap_dir: Path, files: List[WorkerResolvedRepoFile], +) -> tuple: + """Integrity of a reused materialized snapshot (worker thread, blocking). + + Manifest-covered regular files are checked against declared size AND + BLAKE3; reassembled chunked originals (which the manifest digests only + part-wise) get the structural safetensors check. Returns + ``(ok, bad)`` — hex digests in ``bad`` name blobs to quarantine.""" + from .loading import safetensors_file_valid + + bad: List[str] = [] + covered: Set[Path] = set() + for f in files: + if _is_parts_manifest(f.path) or _is_part_file(f.path): + continue # not materialized: parts live only in blobs/ + try: + dst = snap_dir / _norm_rel_path(f.path) + except ValueError: + continue + covered.add(dst) + digest = _strip_blake3_prefix((f.blake3 or "").strip().lower()) + try: + if not dst.exists(): + raise ValueError("missing") + if f.size_bytes and dst.stat().st_size != int(f.size_bytes): + raise ValueError("size mismatch") + if digest and _blake3_file(dst).lower() != digest: + raise ValueError("blake3 mismatch") + except (OSError, ValueError) as exc: + _log.warning("snapshot reuse file %s/%s corrupt: %s", snap_dir.name, f.path, exc) + bad.append(digest or f.path) + try: + candidates = sorted(snap_dir.rglob("*.safetensors")) + except OSError: + candidates = [] + for st in candidates: + if st in covered: + continue + if not safetensors_file_valid(st): + _log.warning("snapshot reuse file %s structurally invalid (truncated?)", st) + bad.append(str(st.relative_to(snap_dir))) + return (not bad, bad) + + +def _quarantine_materialized(snap_dir: Path, blobs_root: Path, bad: Any) -> None: + """Delete a corrupt reused tree AND the corrupt blobs it links, so the + rebuild re-downloads clean bytes instead of re-linking the same rot.""" + shutil.rmtree(snap_dir, ignore_errors=True) + for raw in bad or (): + digest = _strip_blake3_prefix(str(raw or "")).strip().lower() + if "/" in digest or "." in digest or len(digest) < 4: + continue # path-shaped entry (structural failure), not a digest + try: + _blob_path(blobs_root, digest).unlink(missing_ok=True) + except OSError: + continue + fsync_dir(snap_dir.parent) + + def delete_blobs(base_dir: Path, digests: Any) -> None: """Remove specific CAS blobs (quarantine of digest-mismatched content, gw#408) so a re-materialization re-downloads them instead of re-linking diff --git a/tests/test_disk_gc.py b/tests/test_disk_gc.py index 764801f3..08b5d81e 100644 --- a/tests/test_disk_gc.py +++ b/tests/test_disk_gc.py @@ -24,8 +24,14 @@ def _snapshot(digest: str, size: int) -> pb.Snapshot: + # Truthful digest of what _fake_download actually writes: since gw#598 + # the executor tree-verifies every fresh materialization before trusting + # it, so fixtures must declare the real hash of their fake bytes. + from blake3 import blake3 + return pb.Snapshot(digest=digest, files=[ - pb.SnapshotFile(path="w.bin", size_bytes=size, blake3="ab" * 16, + pb.SnapshotFile(path="w.bin", size_bytes=size, + blake3=blake3(b"\0" * size).hexdigest(), url="http://example.invalid/w.bin"), ]) diff --git a/tests/test_snapshot_corruption.py b/tests/test_snapshot_corruption.py index 3150b78e..927d2a6d 100644 --- a/tests/test_snapshot_corruption.py +++ b/tests/test_snapshot_corruption.py @@ -118,6 +118,38 @@ def _counting(self, path, snapshot): assert verifies["n"] == 1 and calls["n"] == 1 +def test_preseeded_unverified_cas_bytes_are_verified_before_trust( + tmp_path: Path, monkeypatch, +) -> None: + """gw#598: a ref the residency layer has NEVER tracked whose bytes already + exist under the CAS root (another pod's writes on a shared th#850 volume) + must be digest-verified before being trusted — the fresh-materialization + path used to mark such reuse verified without any hash.""" + import os + + content = _tiny_safetensors() + snap = _snapshot("44" * 32, content) + calls = _wire_download(monkeypatch, content) + + # Simulate another pod's CAS state: a corrupt blob (right size, wrong + # bytes — every size-only check passes) hardlinked into a fully + # materialized snapshot dir; no residency/ref-index record anywhere. + digest = blake3(content).hexdigest() + blob = tmp_path / "blobs" / "blake3" / digest[:2] / digest[2:4] / digest + blob.parent.mkdir(parents=True) + blob.write_bytes(b"\0" * len(content)) + snap_dir = tmp_path / "snapshots" / ("44" * 32) + snap_dir.mkdir(parents=True) + os.link(blob, snap_dir / "model.safetensors") + + store = ModelStore(_noop_emit, cache_dir=tmp_path) + path = asyncio.run(store.ensure_local("e2e/sdxl-c", snap)) + + assert calls["n"] == 1 # corrupt reuse detected -> quarantine -> re-download + assert (path / "model.safetensors").read_bytes() == content + assert blob.read_bytes() == content # bad blob replaced, not left behind + + # --------------------------------------------------------------------------- # # Load-failure path: digest verify -> quarantine -> re-download -> retry once # --------------------------------------------------------------------------- #