diff --git a/docs/serverless/volume_cache.md b/docs/serverless/volume_cache.md index be780458..d3ebebda 100644 --- a/docs/serverless/volume_cache.md +++ b/docs/serverless/volume_cache.md @@ -59,24 +59,41 @@ vc.sync(background=False) # persist new files back to the volume, inli | `namespace` | `RUNPOD_ENDPOINT_ID` | Isolation key for the on-volume mirror. Must be a single safe path component. | | `volume_path` | `/runpod-volume` | Network-volume mount point. | | `best_effort` | `True` | Swallow and log errors instead of raising. Set `False` while debugging. | +| `max_workers` | `min(32, (os.cpu_count() or 4) * 4)` | Thread count for parallel copy of large files (I/O-bound). | ## How it works -- **Directory mirror.** Cached files live at `{volume_path}/.cache/{namespace}`, - laid out at the same absolute path they occupy in the container (so the - mirror is directly browsable — no archive format to unpack). -- **Per-file, atomic reconcile.** Each direction (`hydrate`/`sync`) walks the - source tree and copies any file whose size differs or whose mtime is newer - than the destination's (beyond a small tolerance for coarse network-filesystem - mtimes). Copies are written to a temp file and atomically renamed into place - (`os.replace`), so a crash mid-copy never leaves a half-written file visible. -- **Idempotent.** Copies preserve mtime (`shutil.copy2`), so re-running - `hydrate`/`sync` after nothing has changed copies zero files. +- **Size-bucketed mirror.** Cached files live at + `{volume_path}/.cache/{namespace}`, split by size: files below 256 KiB are + packed into a single `small.tar` archive (collapsing per-file metadata + round-trips on the network volume), and larger files are copied unpacked + into a `big/` subdirectory, preserving their original relative path. A + versioned `manifest.json` — written last — records file metadata (size, + mtime) for every cached file and is the commit marker for a complete + mirror; a mirror without a valid, current-version manifest is treated as + absent. +- **Incremental large files, whole-archive small files.** `big/` transfers + are diffed per file by size/mtime against the manifest, so unchanged large + files are skipped. The `small.tar` archive is re-packed as a whole whenever + any small file has changed, since unpacking and re-diffing many tiny files + individually is slower than the network volume's per-file overhead. +- **Parallel copy.** Large-file transfers run across a thread pool sized by + `max_workers` (default `min(32, (os.cpu_count() or 4) * 4)`), since the + work is I/O-bound. +- **Packing/extraction.** Packing the small-file archive uses the `tar` + binary when available (falling back to the stdlib `tarfile` module + otherwise). Extraction always goes through `tarfile`, validating every + member's resolved path against the configured `dirs` before writing. +- **Idempotent.** Re-running `hydrate`/`sync` after nothing has changed + copies zero files and does not repack `small.tar`. `manifest.json` is still + refreshed on every `sync` call — a small atomic write that also drops + entries for files deleted locally since the last sync. - **Last-writer-wins.** Under concurrent workers, the mirror simply reflects whichever worker synced most recently — there's no locking or merge. -- **Safety.** Symlinked sources are never followed/copied. Hydration destinations - are checked to resolve inside one of the configured `dirs` before any write, - so a mirror entry can't be used to write outside the cached directories. +- **Safety.** Symlinked sources are never followed/copied. Every archive + member and every big-file destination is checked to resolve inside one of + the configured `dirs` before any write, so a mirror entry can't be used to + write outside the cached directories. ## Limitations @@ -87,3 +104,7 @@ vc.sync(background=False) # persist new files back to the volume, inli a daemon thread; if the process exits without going through normal interpreter shutdown (e.g. `os._exit`, `SIGKILL`), the atexit hook never runs and the sync may not complete. +- **Orphaned big files are never pruned.** If a large file is deleted or + renamed locally, its `big/` copy stays on the volume — hydrate + is manifest-driven and simply ignores it, but volume space grows across + model-version swaps. Same behavior as the prior flat-mirror design. diff --git a/runpod/serverless/utils/rp_volume_cache.py b/runpod/serverless/utils/rp_volume_cache.py index 6c6d1342..6c410019 100644 --- a/runpod/serverless/utils/rp_volume_cache.py +++ b/runpod/serverless/utils/rp_volume_cache.py @@ -8,16 +8,25 @@ - ``sync()``: copy files that are missing or newer in the container onto the volume mirror (used after the cache has been populated/updated). -Both directions are per-file, atomic (write-to-temp + ``os.replace``), and -idempotent: unchanged files are skipped on repeat calls because ``shutil.copy2`` -preserves mtimes. Stdlib-only and best-effort: any failure degrades to a cold -worker and never raises into the caller (unless ``best_effort=False``). +Transport adapts to the tree's shape: files below 256 KiB are packed into a +single ``small.tar`` on the volume (collapsing per-file metadata round-trips), +larger files are copied unpacked into ``big/`` across a thread pool, and a +versioned ``manifest.json`` (written last) records the layout. Large-file +transfers stay incremental (size/mtime diff); the small-file archive is +re-packed whole when any small file changes. Packing uses the ``tar`` binary +when present (else stdlib ``tarfile``); extraction always uses ``tarfile``, +validating every member against the configured dirs. Best-effort: any failure +degrades to a cold worker and never raises unless ``best_effort=False``. """ import atexit +import json import os import shutil +import subprocess +import tarfile import threading +from concurrent.futures import ThreadPoolExecutor from runpod.serverless.modules.rp_logger import RunPodLogger @@ -25,6 +34,11 @@ _MTIME_TOLERANCE = 2.0 # seconds; tolerate coarse (NFS) mtime granularity _JOIN_TIMEOUT_SECONDS = 30.0 # bound the atexit wait so stalled volume I/O can't hang shutdown +_SMALL_FILE_THRESHOLD = 256 * 1024 # bytes; <-threshold packs into one tar (metadata collapse) +_FORMAT_VERSION = 1 +_MANIFEST_NAME = "manifest.json" +_SMALL_ARCHIVE_NAME = "small.tar" +_BIG_SUBDIR = "big" class VolumeCache: @@ -48,6 +62,9 @@ class VolumeCache: volume_path: Network-volume mount point. Defaults to ``/runpod-volume``. best_effort: When True (default), swallow and log errors instead of raising. + max_workers: Thread count for parallel copy of large files. Defaults + to ``min(32, (os.cpu_count() or 4) * 4)`` (I/O-bound, so + oversubscribing the CPU count is fine). Example: >>> with VolumeCache(dirs=["/root/.cache/huggingface"]): @@ -56,7 +73,9 @@ class VolumeCache: _EXCLUDE_SUBSTRINGS = (os.sep + "refs" + os.sep, os.sep + ".no_exist" + os.sep) - def __init__(self, dirs, *, namespace=None, volume_path="/runpod-volume", best_effort=True): + def __init__( + self, dirs, *, namespace=None, volume_path="/runpod-volume", best_effort=True, max_workers=None + ): self._dirs = [os.path.realpath(os.fspath(d)) for d in dirs] self._namespace = namespace or os.environ.get("RUNPOD_ENDPOINT_ID") or "" if self._namespace and ( @@ -71,11 +90,24 @@ def __init__(self, dirs, *, namespace=None, volume_path="/runpod-volume", best_e ) self._volume_path = os.fspath(volume_path) self._best_effort = best_effort + self._max_workers = max_workers or min(32, (os.cpu_count() or 4) * 4) @property def _mirror_root(self): return os.path.join(self._volume_path, ".cache", self._namespace) + @property + def _manifest_path(self): + return os.path.join(self._mirror_root, _MANIFEST_NAME) + + @property + def _small_archive_path(self): + return os.path.join(self._mirror_root, _SMALL_ARCHIVE_NAME) + + @property + def _big_root(self): + return os.path.join(self._mirror_root, _BIG_SUBDIR) + @property def available(self): """True iff volume_path is a mounted dir AND namespace is non-empty.""" @@ -100,6 +132,65 @@ def _iter_files(self, root): continue yield path + def _partition(self, paths): + small, big = [], [] + for p in paths: + try: + size = os.stat(p).st_size + except OSError: + continue + (small if size < _SMALL_FILE_THRESHOLD else big).append(p) + return small, big + + def _file_meta(self, path): + try: + st = os.stat(path) + except OSError: + return None + return {"path": path, "size": st.st_size, "mtime": st.st_mtime} + + def _write_manifest(self, small_meta, big_meta): + os.makedirs(self._mirror_root, exist_ok=True) + payload = { + "version": _FORMAT_VERSION, + "threshold": _SMALL_FILE_THRESHOLD, + "small": small_meta, + "big": big_meta, + } + tmp = f"{self._manifest_path}.{os.getpid()}.{threading.get_ident()}.rpvc.tmp" + with open(tmp, "w") as fh: + json.dump(payload, fh) + os.replace(tmp, self._manifest_path) + + def _read_manifest(self): + try: + with open(self._manifest_path) as fh: + obj = json.load(fh) + except (OSError, ValueError): + return None + if not isinstance(obj, dict) or obj.get("version") != _FORMAT_VERSION: + # Not a dict, or an unknown/future format version: treat as absent + # so callers self-heal (sync repacks and rewrites the manifest) + # instead of raising downstream on unexpected shapes. + return None + return obj + + def _meta_satisfied_by_local(self, meta): + try: + st = os.stat(meta["path"]) + except OSError: + return False + return st.st_size == meta["size"] and st.st_mtime >= meta["mtime"] - _MTIME_TOLERANCE + + def _changed_vs_manifest(self, metas, manifest, key): + prior = {e["path"]: e for e in (manifest.get(key, []) if manifest else [])} + changed = [] + for m in metas: + p = prior.get(m["path"]) + if p is None or p["size"] != m["size"] or m["mtime"] > p["mtime"] + _MTIME_TOLERANCE: + changed.append(m) + return changed + @staticmethod def _needs_copy(src_path, dst_path): try: @@ -133,6 +224,14 @@ def _copy_file(self, src, dst): pass return False + def _copy_parallel(self, pairs): + todo = [(s, d) for s, d in pairs if self._needs_copy(s, d)] + if not todo: + return 0 + with ThreadPoolExecutor(max_workers=self._max_workers) as pool: + results = list(pool.map(lambda sd: self._copy_file(sd[0], sd[1]), todo)) + return sum(1 for ok in results if ok) + def _guard(self, fn, default): try: return fn() @@ -142,6 +241,85 @@ def _guard(self, fn, default): log.warn(f"VolumeCache operation failed: {exc}") return default + def _tar_binary(self): + return shutil.which("tar") + + def _pack_small(self, files): + os.makedirs(self._mirror_root, exist_ok=True) + tmp = f"{self._small_archive_path}.{os.getpid()}.{threading.get_ident()}.rpvc.tmp" + rels = [os.path.relpath(f, "/") for f in files] + try: + if self._tar_binary(): + listing = f"{tmp}.list" + try: + with open(listing, "w") as fh: + fh.write("\n".join(rels)) + subprocess.run( + ["tar", "-C", "/", "-c", "--no-recursion", "-f", tmp, "-T", listing], + check=True, + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + # Suppress BSD tar AppleDouble (._*) sidecars; no-op for GNU tar. + env={**os.environ, "COPYFILE_DISABLE": "1"}, + ) + finally: + _silent_remove(listing) + else: + with tarfile.open(tmp, "w") as tf: + for f, rel in zip(files, rels): + tf.add(f, arcname=rel, recursive=False) + os.replace(tmp, self._small_archive_path) + return True + except (OSError, subprocess.SubprocessError, tarfile.TarError) as exc: + log.debug(f"VolumeCache: pack small bucket failed: {exc}") + _silent_remove(tmp) + return False + + def _extract_small(self): + try: + tf = tarfile.open(self._small_archive_path) + except (OSError, tarfile.TarError): + return 0 + extracted = 0 + try: + try: + members = tf.getmembers() + except (OSError, tarfile.TarError) as exc: + log.debug(f"VolumeCache: corrupt archive {self._small_archive_path}: {exc}") + return extracted + for member in members: + # Skip non-files (directories, symlinks, etc.) + if not member.isfile(): + continue + dst = os.path.join("/", member.name) + # Verify the destination is safe (within configured directories) + if not self._is_safe_dest(dst): + continue + # Check if the local file already satisfies this archive entry + meta = {"path": dst, "size": member.size, "mtime": member.mtime} + if self._meta_satisfied_by_local(meta): + continue + # Extract this member + try: + self._extract_member(tf, member) + extracted += 1 + except (OSError, tarfile.TarError) as exc: + log.debug(f"VolumeCache: skip extract {member.name}: {exc}") + finally: + tf.close() + if extracted: + log.info(f"VolumeCache: extracted {extracted} small file(s)") + return extracted + + @staticmethod + def _extract_member(tf, member): + # Extract to "/" using the data filter when available (3.12+ path-traversal + # hardening); the per-member _is_safe_dest check above is the portable guard. + if hasattr(tarfile, "data_filter"): + tf.extract(member, "/", filter="data") + else: + tf.extract(member, "/") + # ----------------------------------------------------------------- # # hydrate / sync # ----------------------------------------------------------------- # @@ -157,19 +335,22 @@ def hydrate(self): return self._guard(self._do_hydrate, 0) def _do_hydrate(self): - root = self._mirror_root - if not os.path.isdir(root): + manifest = self._read_manifest() + if manifest is None: return 0 - copied = 0 - for m in self._iter_files(root): - dst = os.path.join("/", os.path.relpath(m, root)) - if not self._is_safe_dest(dst): - continue # skip anything that would escape configured dirs - if self._needs_copy(m, dst) and self._copy_file(m, dst): - copied += 1 - if copied: - log.info(f"VolumeCache: hydrated {copied} file(s) from {root}") - return copied + restored = 0 + if manifest.get("small"): + restored += self._extract_small() + big_pairs = [] + for entry in manifest.get("big", []): + rel = os.path.relpath(entry["path"], "/") + dst = os.path.join("/", rel) + if self._is_safe_dest(dst): + big_pairs.append((os.path.join(self._big_root, rel), dst)) + restored += self._copy_parallel(big_pairs) + if restored: + log.info(f"VolumeCache: hydrated {restored} file(s) from {self._mirror_root}") + return restored def sync(self, *, background=True): """Reconcile container -> volume mirror. @@ -189,18 +370,41 @@ def sync(self, *, background=True): self._guard(self._do_sync, 0) def _do_sync(self): - os.makedirs(self._mirror_root, exist_ok=True) - copied = 0 + all_files = [] for root in self._dirs: - if not os.path.isdir(root): - continue - for f in self._iter_files(root): - dst = os.path.join(self._mirror_root, os.path.relpath(f, "/")) - if self._needs_copy(f, dst) and self._copy_file(f, dst): - copied += 1 - if copied: - log.info(f"VolumeCache: synced {copied} file(s) to {self._mirror_root}") - return copied + if os.path.isdir(root): + all_files.extend(self._iter_files(root)) + small_files, big_files = self._partition(all_files) + + manifest = self._read_manifest() + small_meta = [m for m in (self._file_meta(f) for f in small_files) if m] + small_transferred = 0 + if small_meta: + changed = self._changed_vs_manifest(small_meta, manifest, "small") + if changed: + if self._pack_small([m["path"] for m in small_meta]): + small_transferred = len(changed) + else: + # pack failed (no tar binary and tarfile unusable): reclassify + # the small files as big and copy them unpacked instead. + big_files = big_files + [m["path"] for m in small_meta] + small_meta = [] + # else: unchanged -> the existing archive is still current + if not small_meta: + # No small files (or pack failed and reclassified them as big): + # remove any stale archive before the manifest is written so a + # crash can't leave manifest.small=[] with a live archive. + _silent_remove(self._small_archive_path) + + big_pairs = [(f, os.path.join(self._big_root, os.path.relpath(f, "/"))) for f in big_files] + big_copied = self._copy_parallel(big_pairs) + big_meta = [m for m in (self._file_meta(f) for f in big_files) if m] + + self._write_manifest(small_meta, big_meta) + total = big_copied + small_transferred + if total: + log.info(f"VolumeCache: synced {total} file(s) to {self._mirror_root}") + return total # ----------------------------------------------------------------- # # context manager @@ -215,6 +419,19 @@ def __exit__(self, *exc): return None +# ----------------------------------------------------------------------- # +# utilities +# ----------------------------------------------------------------------- # + + +def _silent_remove(path): + try: + os.remove(path) + except OSError: + # best-effort cleanup; a leftover temp is overwritten or ignored next run + pass + + # ----------------------------------------------------------------------- # # background sync completion # ----------------------------------------------------------------------- # diff --git a/tests/test_serverless/test_utils/test_rp_volume_cache.py b/tests/test_serverless/test_utils/test_rp_volume_cache.py index 4c85bd74..d3d91e41 100644 --- a/tests/test_serverless/test_utils/test_rp_volume_cache.py +++ b/tests/test_serverless/test_utils/test_rp_volume_cache.py @@ -1,5 +1,8 @@ +import json import os import shutil +import subprocess +import tarfile import threading import time @@ -107,7 +110,7 @@ def test_sync_is_idempotent_on_unchanged_file(tmp_path): def test_sync_recopies_modified_file(tmp_path): - vc, cache, _vol = _mk_cache_with_volume(tmp_path) + vc, cache, vol = _mk_cache_with_volume(tmp_path) f = cache / "model.bin" f.write_text("v1") vc.sync(background=False) @@ -115,14 +118,12 @@ def test_sync_recopies_modified_file(tmp_path): time.sleep(0.01) f.write_text("v2-longer") os.utime(f, (time.time() + 10, time.time() + 10)) + assert vc._do_sync() == 1 # one changed small file repacked - copied = vc._do_sync() - assert copied == 1 - - mirror_file = os.path.join(vc._mirror_root, os.path.relpath(str(f), "/")) - with open(mirror_file) as fh: - content = fh.read() - assert content == "v2-longer" + f.unlink() + fresh = VolumeCache([str(cache)], namespace="ep1", volume_path=str(vol)) + fresh.hydrate() + assert f.read_text() == "v2-longer" def test_hydrate_skips_unchanged_after_first_hydrate(tmp_path): @@ -181,8 +182,9 @@ def test_sync_skips_symlink_source(tmp_path): copied = vc._do_sync() assert copied == 1 # only real.bin, not the symlink - mirror_link = os.path.join(vc._mirror_root, os.path.relpath(str(link), "/")) - assert not os.path.exists(mirror_link) + manifest = vc._read_manifest() + manifest_paths = [e["path"] for e in manifest["small"] + manifest["big"]] + assert str(link) not in manifest_paths # --------------------------------------------------------------------------- # @@ -191,18 +193,37 @@ def test_sync_skips_symlink_source(tmp_path): def test_hydrate_skips_unsafe_destination(tmp_path): - vc, cache, vol = _mk_cache_with_volume(tmp_path) - - # Craft a mirror file whose mapped destination escapes the configured dirs. - escape_rel = os.path.relpath(str(tmp_path / "outside" / "evil.txt"), "/") - mirror_file = os.path.join(vc._mirror_root, escape_rel) - os.makedirs(os.path.dirname(mirror_file), exist_ok=True) - with open(mirror_file, "w") as fh: + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + os.makedirs(vc._big_root, exist_ok=True) + # craft a manifest whose big entry maps outside the configured dirs + escape = str(tmp_path / "outside" / "evil.txt") + big_src = os.path.join(vc._big_root, os.path.relpath(escape, "/")) + os.makedirs(os.path.dirname(big_src), exist_ok=True) + with open(big_src, "w") as fh: fh.write("malicious") + vc._write_manifest([], [{"path": escape, "size": 9, "mtime": 1.0}]) + assert vc.hydrate() == 0 + assert not os.path.exists(escape) - copied = vc.hydrate() - assert copied == 0 - assert not (tmp_path / "outside" / "evil.txt").exists() + +def test_hydrate_no_manifest_is_noop(tmp_path): + vc, _cache, _vol = _mk_cache_with_volume(tmp_path) + os.makedirs(vc._mirror_root, exist_ok=True) + assert vc._do_hydrate() == 0 + + +def test_hydrate_restores_big_and_small(tmp_path): + vc, cache, vol = _mk_cache_with_volume(tmp_path) + (cache / "tiny.bin").write_bytes(b"s" * 100) + (cache / "weights.bin").write_bytes(b"w" * (256 * 1024 + 5)) + vc.sync(background=False) + (cache / "tiny.bin").unlink() + (cache / "weights.bin").unlink() + + fresh = VolumeCache([str(cache)], namespace="ep1", volume_path=str(vol)) + assert fresh.hydrate() == 2 + assert (cache / "tiny.bin").read_bytes() == b"s" * 100 + assert (cache / "weights.bin").read_bytes() == b"w" * (256 * 1024 + 5) # --------------------------------------------------------------------------- # @@ -219,8 +240,8 @@ def test_context_manager_hydrates_on_enter_and_syncs_on_exit(tmp_path): vcmod._join_pending_syncs() - mirror_file = os.path.join(vc._mirror_root, os.path.relpath(str(cache / "model.bin"), "/")) - assert os.path.exists(mirror_file) + m = vc._read_manifest() + assert [e["path"] for e in m["small"]] == [str(cache / "model.bin")] def test_context_manager_does_not_suppress_exceptions(tmp_path): @@ -236,8 +257,8 @@ def test_context_manager_does_not_suppress_exceptions(tmp_path): assert raised vcmod._join_pending_syncs() - mirror_file = os.path.join(vc._mirror_root, os.path.relpath(str(cache / "model.bin"), "/")) - assert os.path.exists(mirror_file) + m = vc._read_manifest() + assert [e["path"] for e in m["small"]] == [str(cache / "model.bin")] # --------------------------------------------------------------------------- # @@ -349,6 +370,76 @@ def test_sync_noop_when_unavailable(tmp_path): assert vc.sync(background=False) is None +# --------------------------------------------------------------------------- # +# manifest and metadata +# --------------------------------------------------------------------------- # + + +def test_manifest_round_trip(tmp_path): + vc, _cache, _vol = _mk_cache_with_volume(tmp_path) + os.makedirs(vc._mirror_root, exist_ok=True) + small = [{"path": "/a", "size": 1, "mtime": 100.0}] + big = [{"path": "/b", "size": 999999, "mtime": 200.0}] + vc._write_manifest(small, big) + m = vc._read_manifest() + assert m["version"] == 1 and m["threshold"] == 256 * 1024 + assert m["small"] == small and m["big"] == big + + +def test_read_manifest_none_when_absent(tmp_path): + vc, _cache, _vol = _mk_cache_with_volume(tmp_path) + assert vc._read_manifest() is None + + +def test_read_manifest_none_when_corrupt(tmp_path): + vc, _cache, _vol = _mk_cache_with_volume(tmp_path) + os.makedirs(vc._mirror_root, exist_ok=True) + with open(vc._manifest_path, "w") as fh: + fh.write("{not json") + assert vc._read_manifest() is None + + +def test_read_manifest_none_when_valid_json_not_object(tmp_path): + vc, _cache, _vol = _mk_cache_with_volume(tmp_path) + os.makedirs(vc._mirror_root, exist_ok=True) + with open(vc._manifest_path, "w") as fh: + fh.write("42") + assert vc._read_manifest() is None + + +def test_read_manifest_none_when_unknown_version(tmp_path): + vc, _cache, _vol = _mk_cache_with_volume(tmp_path) + os.makedirs(vc._mirror_root, exist_ok=True) + with open(vc._manifest_path, "w") as fh: + json.dump({"version": vcmod._FORMAT_VERSION + 1, "small": [], "big": []}, fh) + assert vc._read_manifest() is None + + +def test_meta_satisfied_by_local(tmp_path): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + f = cache / "m.bin" + f.write_bytes(b"1234") + st = os.stat(str(f)) + assert vc._meta_satisfied_by_local({"path": str(f), "size": 4, "mtime": st.st_mtime}) is True + # older manifest mtime -> local (newer/equal) still satisfies + assert vc._meta_satisfied_by_local({"path": str(f), "size": 4, "mtime": st.st_mtime - 100}) is True + # size mismatch -> not satisfied + assert vc._meta_satisfied_by_local({"path": str(f), "size": 5, "mtime": st.st_mtime}) is False + # missing local -> not satisfied + assert vc._meta_satisfied_by_local({"path": str(cache / "gone"), "size": 1, "mtime": 1.0}) is False + + +def test_changed_vs_manifest(tmp_path): + vc, _cache, _vol = _mk_cache_with_volume(tmp_path) + manifest = {"small": [{"path": "/a", "size": 1, "mtime": 100.0}], "big": []} + metas = [ + {"path": "/a", "size": 1, "mtime": 100.0}, # unchanged + {"path": "/c", "size": 2, "mtime": 100.0}, # new + ] + changed = vc._changed_vs_manifest(metas, manifest, "small") + assert [m["path"] for m in changed] == ["/c"] + + # --------------------------------------------------------------------------- # # exports # --------------------------------------------------------------------------- # @@ -362,6 +453,42 @@ def test_volumecache_exported_from_serverless(): assert "VolumeCache" in sls.__all__ +# --------------------------------------------------------------------------- # +# size partition and packed format +# --------------------------------------------------------------------------- # + + +def test_partition_splits_at_threshold(tmp_path): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + small = cache / "small.bin" + small.write_bytes(b"x" * (256 * 1024 - 1)) + exact = cache / "exact.bin" + exact.write_bytes(b"x" * (256 * 1024)) # >= threshold -> big + big = cache / "big.bin" + big.write_bytes(b"x" * (256 * 1024 + 1)) + s, b = vc._partition([str(small), str(exact), str(big)]) + assert s == [str(small)] + assert sorted(b) == sorted([str(exact), str(big)]) + + +def test_partition_drops_unstatable(tmp_path): + vc, _cache, _vol = _mk_cache_with_volume(tmp_path) + s, b = vc._partition([str(tmp_path / "gone.bin")]) + assert s == [] and b == [] + + +def test_default_max_workers_is_positive(tmp_path): + vc, _cache, _vol = _mk_cache_with_volume(tmp_path) + assert isinstance(vc._max_workers, int) and vc._max_workers >= 1 + + +def test_mirror_layout_paths(tmp_path): + vc, _cache, vol = _mk_cache_with_volume(tmp_path) + assert vc._manifest_path == os.path.join(vc._mirror_root, "manifest.json") + assert vc._small_archive_path == os.path.join(vc._mirror_root, "small.tar") + assert vc._big_root == os.path.join(vc._mirror_root, "big") + + # --------------------------------------------------------------------------- # # v2 review fixes # --------------------------------------------------------------------------- # @@ -371,12 +498,71 @@ def test_iter_files_skips_inflight_temp_files(tmp_path): vc, cache, _vol = _mk_cache_with_volume(tmp_path) (cache / "model.bin").write_text("real") (cache / "model.bin.12345.67890.rpvc.tmp").write_text("partial") - copied = vc._do_sync() - assert copied == 1 # only model.bin, not the .rpvc.tmp file - tmp_mirror = os.path.join( - vc._mirror_root, os.path.relpath(str(cache / "model.bin.12345.67890.rpvc.tmp"), "/") - ) - assert not os.path.exists(tmp_mirror) + assert vc._do_sync() == 1 # only model.bin + m = vc._read_manifest() + packed = [e["path"] for e in m["small"]] + [e["path"] for e in m["big"]] + assert str(cache / "model.bin") in packed + assert str(cache / "model.bin.12345.67890.rpvc.tmp") not in packed + + +def test_sync_writes_manifest_and_archive_for_small(tmp_path): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + (cache / "tiny.bin").write_bytes(b"x" * 100) + vc._do_sync() + assert os.path.exists(vc._manifest_path) + assert os.path.exists(vc._small_archive_path) + m = vc._read_manifest() + assert [e["path"] for e in m["small"]] == [str(cache / "tiny.bin")] + + +def test_sync_big_file_is_unpacked_and_incremental(tmp_path): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + big = cache / "weights.bin" + big.write_bytes(b"x" * (256 * 1024 + 10)) + assert vc._do_sync() == 1 + assert os.path.exists(os.path.join(vc._big_root, os.path.relpath(str(big), "/"))) + assert vc._do_sync() == 0 # unchanged -> no recopy, no repack + + +def test_sync_reclassifies_small_as_big_without_tar(tmp_path, monkeypatch): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + (cache / "tiny.bin").write_bytes(b"x" * 100) + monkeypatch.setattr(vc, "_tar_binary", lambda: None) + monkeypatch.setattr(tarfile, "open", lambda *a, **k: (_ for _ in ()).throw(OSError("no tar"))) + assert vc._do_sync() == 1 + m = vc._read_manifest() + assert m["small"] == [] # nothing packed + assert [e["path"] for e in m["big"]] == [str(cache / "tiny.bin")] # unpacked instead + assert not os.path.exists(vc._small_archive_path) + + +def test_hydrate_ignores_stale_archive_after_reclassify(tmp_path, monkeypatch): + vc, cache, vol = _mk_cache_with_volume(tmp_path) + tiny = cache / "tiny.bin" + tiny.write_bytes(b"v1" * 10) + + vc.sync(background=False) + assert os.path.exists(vc._small_archive_path) + manifest = vc._read_manifest() + assert [e["path"] for e in manifest["small"]] == [str(tiny)] + + # Modify the file, then force sync #2's pack to fail so it reclassifies + # the file as big -- the stale small.tar must not survive this. + tiny.write_bytes(b"v2-longer-content") + monkeypatch.setattr(vc, "_tar_binary", lambda: None) + monkeypatch.setattr(tarfile, "open", lambda *a, **k: (_ for _ in ()).throw(OSError("no tar"))) + vc._do_sync() + + manifest = vc._read_manifest() + assert manifest["small"] == [] + assert not os.path.exists(vc._small_archive_path) + + # Read normally (unpatched) from here on. + tiny.unlink() + fresh = VolumeCache([str(cache)], namespace="ep1", volume_path=str(vol)) + fresh.hydrate() + + assert tiny.read_bytes() == b"v2-longer-content" def test_join_pending_syncs_bounded_by_timeout(monkeypatch): @@ -397,3 +583,173 @@ def slow(): vcmod._join_pending_syncs() release.set() vcmod._reset_pending_for_test() + + +# --------------------------------------------------------------------------- # +# pack small +# --------------------------------------------------------------------------- # + + +def _rel(p): + return os.path.relpath(p, "/") + + +def test_pack_small_with_binary_round_trips(tmp_path): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + os.makedirs(vc._mirror_root, exist_ok=True) + (cache / "a.txt").write_text("aaa") + (cache / "sub").mkdir() + (cache / "sub" / "b.txt").write_text("bbb") + files = [str(cache / "a.txt"), str(cache / "sub" / "b.txt")] + assert vc._pack_small(files) is True + with tarfile.open(vc._small_archive_path) as tf: + names = set(tf.getnames()) + assert _rel(files[0]) in names and _rel(files[1]) in names + + +def test_pack_small_falls_back_to_tarfile(tmp_path, monkeypatch): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + os.makedirs(vc._mirror_root, exist_ok=True) + (cache / "a.txt").write_text("aaa") + monkeypatch.setattr(vc, "_tar_binary", lambda: None) + assert vc._pack_small([str(cache / "a.txt")]) is True + with tarfile.open(vc._small_archive_path) as tf: + assert _rel(str(cache / "a.txt")) in tf.getnames() + + +def test_pack_small_atomic_no_partial_on_failure(tmp_path, monkeypatch): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + os.makedirs(vc._mirror_root, exist_ok=True) + (cache / "a.txt").write_text("aaa") + monkeypatch.setattr(vc, "_tar_binary", lambda: None) + + def boom(*a, **k): + raise OSError("no space") + + monkeypatch.setattr(tarfile, "open", boom) + assert vc._pack_small([str(cache / "a.txt")]) is False + assert not os.path.exists(vc._small_archive_path) # atomic: no partial archive + + +def test_pack_small_subprocess_tar_failure_leaves_no_partial(tmp_path, monkeypatch): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + os.makedirs(vc._mirror_root, exist_ok=True) + (cache / "a.txt").write_text("aaa") + # Force the subprocess-tar branch, then fail the subprocess call itself. + monkeypatch.setattr(vc, "_tar_binary", lambda: "/usr/bin/tar") + + def boom(*a, **k): + raise subprocess.CalledProcessError(1, "tar") + + monkeypatch.setattr(subprocess, "run", boom) + assert vc._pack_small([str(cache / "a.txt")]) is False + assert not os.path.exists(vc._small_archive_path) # atomic: no partial archive + # No leftover .list temp file in the mirror root. + assert not any(name.endswith(".list") for name in os.listdir(vc._mirror_root)) + + +# --------------------------------------------------------------------------- # +# extract small +# --------------------------------------------------------------------------- # + + +def test_extract_small_restores_files(tmp_path, monkeypatch): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + monkeypatch.setattr(vc, "_tar_binary", lambda: None) # deterministic pure-Python tarfile pack + os.makedirs(vc._mirror_root, exist_ok=True) + (cache / "a.txt").write_text("hello") + vc._pack_small([str(cache / "a.txt")]) + (cache / "a.txt").unlink() + assert vc._extract_small() == 1 + assert (cache / "a.txt").read_text() == "hello" + + +def test_extract_small_skips_unsafe_member(tmp_path, monkeypatch): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + monkeypatch.setattr(vc, "_tar_binary", lambda: None) # deterministic pure-Python tarfile pack + os.makedirs(vc._mirror_root, exist_ok=True) + outside = tmp_path / "outside" / "evil.txt" + outside.parent.mkdir() + outside.write_text("malicious") + vc._pack_small([str(outside)]) # archive contains an escaping path + outside.unlink() + assert vc._extract_small() == 0 # refused by _is_safe_dest + assert not outside.exists() + + +def test_extract_small_incremental_skips_satisfied(tmp_path, monkeypatch): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + monkeypatch.setattr(vc, "_tar_binary", lambda: None) # deterministic pure-Python tarfile pack + os.makedirs(vc._mirror_root, exist_ok=True) + (cache / "a.txt").write_text("hello") + vc._pack_small([str(cache / "a.txt")]) + # local file already present and current -> nothing to extract + assert vc._extract_small() == 0 + + +def test_extract_small_zero_when_archive_missing(tmp_path): + vc, _cache, _vol = _mk_cache_with_volume(tmp_path) + os.makedirs(vc._mirror_root, exist_ok=True) + assert vc._extract_small() == 0 + + +def test_extract_small_never_raises_when_getmembers_raises(tmp_path, monkeypatch): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + monkeypatch.setattr(vc, "_tar_binary", lambda: None) # deterministic pure-Python tarfile pack + os.makedirs(vc._mirror_root, exist_ok=True) + (cache / "a.txt").write_text("hello") + vc._pack_small([str(cache / "a.txt")]) + + # tarfile.open() succeeds (archive exists and opens fine), but the inner + # tf.getmembers() call raises tarfile.TarError. This exercises the inner + # guard in _extract_small, distinct from the outer open()-time guard. + class FakeTarFile: + def getmembers(self): + raise tarfile.TarError("corrupt member table") + + def close(self): + pass + + monkeypatch.setattr(tarfile, "open", lambda *a, **k: FakeTarFile()) + assert vc._extract_small() == 0 + + +def test_extract_small_skips_non_file_members(tmp_path): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + os.makedirs(vc._mirror_root, exist_ok=True) + (cache / "sub").mkdir() + (cache / "sub" / "a.txt").write_text("hello") + # Build the archive by hand so it contains a real directory member + # alongside the file member; _extract_small must skip the directory + # (member.isfile() is False) without raising. + with tarfile.open(vc._small_archive_path, "w") as tf: + tf.add(str(cache / "sub"), arcname=_rel(str(cache / "sub")), recursive=False) + tf.add(str(cache / "sub" / "a.txt"), arcname=_rel(str(cache / "sub" / "a.txt"))) + (cache / "sub" / "a.txt").unlink() + assert vc._extract_small() == 1 + assert (cache / "sub" / "a.txt").read_text() == "hello" + + +# --------------------------------------------------------------------------- # +# copy_parallel (Task 5) +# --------------------------------------------------------------------------- # + + +def test_copy_parallel_copies_needed_only(tmp_path): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + src1 = cache / "s1" + src1.write_text("one") + src2 = cache / "s2" + src2.write_text("two") + dst1 = tmp_path / "d1" + dst2 = tmp_path / "d2" + # pre-satisfy dst1 so it is skipped + shutil.copy2(str(src1), str(dst1)) + n = vc._copy_parallel([(str(src1), str(dst1)), (str(src2), str(dst2))]) + assert n == 1 + assert dst2.read_text() == "two" + + +def test_copy_parallel_empty_is_zero(tmp_path): + vc, _cache, _vol = _mk_cache_with_volume(tmp_path) + assert vc._copy_parallel([]) == 0