diff --git a/README.md b/README.md index b34014b19..6ad9c6669 100644 --- a/README.md +++ b/README.md @@ -149,6 +149,19 @@ runpod.serverless.start({"handler": handler}) See [Worker Fitness Checks](https://github.com/runpod/runpod-python/blob/main/docs/serverless/worker_fitness_checks.md) documentation for more examples and best practices. +### Network-Volume Warm Cache + +When a network volume is attached, `VolumeCache` warms local directories (such as a model cache) across cold starts — hydrating them on startup and syncing new files back on exit — so a repeated multi-GB model download becomes a one-time cost per endpoint. It is stdlib-only and best-effort. + +```python +from runpod.serverless import VolumeCache + +with VolumeCache(dirs=["/root/.cache/huggingface"]): + model = load_model() +``` + +See [Network-Volume Warm Cache](https://github.com/runpod/runpod-python/blob/main/docs/serverless/volume_cache.md) documentation for configuration and details. + ## 📚 | API Language Library (GraphQL Wrapper) When interacting with the Runpod API you can use this library to make requests to the API. diff --git a/docs/serverless/volume_cache.md b/docs/serverless/volume_cache.md new file mode 100644 index 000000000..d3ebebda9 --- /dev/null +++ b/docs/serverless/volume_cache.md @@ -0,0 +1,110 @@ +# Network-Volume Warm Cache (VolumeCache) + +`VolumeCache` warms local directories across serverless workers using a mounted +network volume. It keeps a browsable mirror of your cache directories on the +volume and reconciles it against the container on each use: on cold start it +restores previously-cached files into place, and after use it copies newly +written files back to the volume. This turns a repeated multi-GB model +download on every cold start into a one-time cost per endpoint. + +It is stdlib-only and best-effort: any failure degrades to a cold worker and +never raises into your handler or worker loop. + +## Requirements + +- A **network volume** attached to the endpoint (mounted at `/runpod-volume`). +- Set on serverless automatically: `RUNPOD_ENDPOINT_ID` (used to scope the + mirror per endpoint/namespace). + +If no volume is mounted, or the namespace is empty, every operation is a safe +no-op. + +## Usage + +`VolumeCache` is a context-manager closure — using it around a model load +hydrates the cache before the block runs and syncs any changes back after: + +```python +from runpod.serverless import VolumeCache + +with VolumeCache(dirs=["/root/.cache/huggingface"]): + model = load_model() # downloads land in the cached directory +``` + +- **On enter**, `hydrate()` copies files that are missing or newer on the + volume mirror into the container. +- **On exit**, `sync()` copies files that are missing or newer in the + container onto the volume mirror. By default this runs on a background + daemon thread and returns immediately, so the `with` block doesn't block on + the sync; a process-exit hook joins any outstanding syncs so short-lived + processes (including local test runs) still complete the sync before + exiting. + +You can also call the phases directly when they happen at different points in +your worker's lifecycle: + +```python +vc = VolumeCache(dirs=["/data/models"], namespace="my-model-cache") + +vc.hydrate() # restore cached files (e.g. at startup) +model = load_model() # populate the cache +vc.sync(background=False) # persist new files back to the volume, inline +``` + +## Constructor + +| Argument | Default | Purpose | +| --- | --- | --- | +| `dirs` | required | Local directories to cache. | +| `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 + +- **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. 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 + +- **Cold-scale write amplification.** If N workers cold-start at the same time, + each may miss the still-empty mirror, download the model, and sync a full + copy back. There's no coordination between concurrent syncs. +- **Background sync on short-lived processes.** `sync()` schedules the copy on + 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/examples/serverless/volume_cache_handler.py b/examples/serverless/volume_cache_handler.py new file mode 100644 index 000000000..a0b4e08c2 --- /dev/null +++ b/examples/serverless/volume_cache_handler.py @@ -0,0 +1,48 @@ +"""Warm-cache model weights across cold starts with VolumeCache. + +`VolumeCache` keeps a browsable mirror of local cache directories on a mounted +network volume and reconciles the two on each use: it restores cached files on +cold start (`hydrate`) and syncs new downloads back afterward (`sync`). Wrapping +a model load in `with VolumeCache(...)` turns a repeated multi-GB download into a +one-time cost per endpoint. + +Why not just point HF_HOME at the volume? Because then every read hits the +network mount. Here the Hugging Face cache stays on fast local disk and +VolumeCache mirrors it to the volume, so inference reads are local while cold +starts stay warm. + +Requirements to run: +- Attach a network volume to the endpoint (mounted at /runpod-volume). +- RUNPOD_ENDPOINT_ID is set automatically on Runpod serverless (it scopes the + cache per endpoint). Without a mounted volume or endpoint id, VolumeCache is a + safe no-op and the handler still works. +- pip install "transformers" "torch" (the model library used below). + +Local test: + python volume_cache_handler.py --rp_serve_api +""" + +import os + +import runpod +from runpod.serverless import VolumeCache + +# Hugging Face caches models here by default. Keep it on local disk and let +# VolumeCache mirror it to the network volume. +HF_CACHE = os.environ.get("HF_HOME", os.path.expanduser("~/.cache/huggingface")) + +# Hydrate the cache from the volume before loading, then sync new files back on +# exit. On a warm endpoint the model is already local and nothing is downloaded. +with VolumeCache(dirs=[HF_CACHE]): + from transformers import pipeline + + classifier = pipeline("sentiment-analysis") + + +def handler(job): + """Classify the sentiment of an input string.""" + text = job["input"].get("text", "") + return {"input": text, "predictions": classifier(text)} + + +runpod.serverless.start({"handler": handler}) diff --git a/runpod/serverless/__init__.py b/runpod/serverless/__init__.py index 7905731f3..052452073 100644 --- a/runpod/serverless/__init__.py +++ b/runpod/serverless/__init__.py @@ -17,12 +17,14 @@ from .modules.rp_logger import RunPodLogger from .modules.rp_progress import progress_update from .modules.rp_fitness import register_fitness_check +from .utils.rp_volume_cache import VolumeCache __all__ = [ "start", "progress_update", "register_fitness_check", - "runpod_version" + "runpod_version", + "VolumeCache", ] log = RunPodLogger() diff --git a/runpod/serverless/utils/__init__.py b/runpod/serverless/utils/__init__.py index 9036a4ddd..f0c44cc8d 100644 --- a/runpod/serverless/utils/__init__.py +++ b/runpod/serverless/utils/__init__.py @@ -2,9 +2,11 @@ from .rp_download import download_files_from_urls from .rp_upload import upload_file_to_bucket, upload_in_memory_object +from .rp_volume_cache import VolumeCache __all__ = [ "download_files_from_urls", "upload_file_to_bucket", - "upload_in_memory_object" + "upload_in_memory_object", + "VolumeCache", ] diff --git a/runpod/serverless/utils/rp_volume_cache.py b/runpod/serverless/utils/rp_volume_cache.py new file mode 100644 index 000000000..a4c75449a --- /dev/null +++ b/runpod/serverless/utils/rp_volume_cache.py @@ -0,0 +1,491 @@ +"""Directory-mirror warm cache between local directories and a network volume. + +``VolumeCache`` keeps a browsable mirror of one or more local directories on a +mounted network volume, and reconciles the two directions on demand: + +- ``hydrate()``: copy files that are missing or newer on the volume mirror + into the container (used on cold start, before the cache is populated). +- ``sync()``: copy files that are missing or newer in the container onto the + volume mirror (used after the cache has been populated/updated). + +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 + +log = RunPodLogger() + +_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: + """Warm-cache local directories across serverless workers via a network volume. + + Maintains a mirror of ``dirs`` at ``{volume_path}/.cache/{namespace}``. + ``hydrate()`` reconciles volume -> container; ``sync()`` reconciles + container -> volume. Used as a context manager, the object is itself the + "warm cache" closure: hydrate on enter, sync (in the background by + default) on exit. + + Requires a network volume mounted at ``volume_path`` (default + ``/runpod-volume``) and a non-empty ``namespace`` (default + ``RUNPOD_ENDPOINT_ID``); otherwise ``available`` is False and all + operations are no-ops. + + Args: + dirs: Local directories to cache (e.g. a model cache like ``HF_HOME``). + namespace: Isolation key for the on-volume mirror. Must be a single + safe path component. Defaults to ``RUNPOD_ENDPOINT_ID``. + 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"]): + ... model = load_model() # downloads land in the cached dir + """ + + _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, max_workers=None + ): + if isinstance(dirs, (str, bytes, os.PathLike)): + # A bare path is one directory, not an iterable of path characters: + # `dirs="/root/.cache"` must not be split into "/", "r", "o", ... + dirs = [dirs] + 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 ( + os.path.isabs(self._namespace) + or os.sep in self._namespace + or "/" in self._namespace + or "\\" in self._namespace + or self._namespace in (".", "..") + ): + raise ValueError( + f"namespace must be a single safe path component, got {self._namespace!r}" + ) + 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.""" + return bool(self._namespace) and os.path.isdir(self._volume_path) + + # ----------------------------------------------------------------- # + # reconcile helpers + # ----------------------------------------------------------------- # + + def _iter_files(self, root): + for dirpath, _dirs, files in os.walk(root): + for name in files: + path = os.path.join(dirpath, name) + if name.endswith(".lock") or name.endswith(".rpvc.tmp"): + continue + if any(sub in path for sub in self._EXCLUDE_SUBSTRINGS): + continue + try: + if os.path.islink(path): + continue + except OSError: + 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: + s = os.stat(src_path) + except OSError: + return False + try: + d = os.stat(dst_path) + except OSError: + return True + return s.st_size != d.st_size or s.st_mtime > d.st_mtime + _MTIME_TOLERANCE + + def _is_safe_dest(self, dst_abs): + target = os.path.realpath(dst_abs) + return any(target == d or target.startswith(d + os.sep) for d in self._dirs) + + def _copy_file(self, src, dst): + tmp = f"{dst}.{os.getpid()}.{threading.get_ident()}.rpvc.tmp" + try: + os.makedirs(os.path.dirname(dst), exist_ok=True) + shutil.copy2(src, tmp) # preserves mtime so future diffs converge + os.replace(tmp, dst) # atomic; last-writer-wins under concurrency + return True + except OSError as exc: + log.debug(f"VolumeCache: skip {src} -> {dst}: {exc}") + try: + if os.path.exists(tmp): + os.remove(tmp) + except OSError: + # best-effort cleanup; a leftover temp file is overwritten next sync + 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: + # Stream the success count rather than materializing a results list, + # which would hold one entry per file for large (40k+) trees. + return sum( + 1 for ok in pool.map(lambda sd: self._copy_file(sd[0], sd[1]), todo) if ok + ) + + @staticmethod + def _within(base, path): + """True iff ``path`` resolves inside ``base`` (symlinks resolved).""" + base_real = os.path.realpath(base) + path_real = os.path.realpath(path) + return path_real == base_real or path_real.startswith(base_real + os.sep) + + def _guard(self, fn, default): + try: + return fn() + except Exception as exc: # best-effort: never break the worker + if not self._best_effort: + raise + 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 + # ----------------------------------------------------------------- # + + def hydrate(self): + """Reconcile volume mirror -> container. + + Copy files missing or newer in the container. Returns the number of + files copied. No-op (0) if unavailable. + """ + if not self.available: + return 0 + return self._guard(self._do_hydrate, 0) + + def _do_hydrate(self): + manifest = self._read_manifest() + if manifest is None: + return 0 + 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) + src = os.path.join(self._big_root, rel) + # Guard both ends against a crafted manifest: the destination must + # stay inside a configured dir, and the source must resolve inside + # big/ so a malicious entry can't read outside the mirror. + if self._is_safe_dest(dst) and self._within(self._big_root, src): + big_pairs.append((src, 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. + + Copy files missing or newer on the volume. When ``background=True`` + (default), run on a daemon thread and return immediately; a + process-exit hook joins outstanding syncs so short-lived processes + still complete. When False, run inline and return when done. + """ + if not self.available: + return + if background: + t = threading.Thread(target=lambda: self._guard(self._do_sync, 0), daemon=True) + _register_pending(t) + t.start() + else: + self._guard(self._do_sync, 0) + + def _do_sync(self): + all_files = [] + for root in self._dirs: + 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") + prior_small = {e["path"] for e in (manifest.get("small", []) if manifest else [])} + deleted = bool(prior_small - {m["path"] for m in small_meta}) + if changed or deleted: + # Repack on any small-set change, including deletions. small.tar + # is extracted member-by-member on hydrate, so a stale archive + # that still holds a locally-deleted file would resurrect it; + # _changed_vs_manifest only reports additions/modifications, so + # deletions are detected separately here. + 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 + # ----------------------------------------------------------------- # + + def __enter__(self): + self.hydrate() + return self + + def __exit__(self, *exc): + self.sync(background=True) + 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 +# ----------------------------------------------------------------------- # + +_pending_syncs = [] + + +def _register_pending(thread): + _pending_syncs[:] = [t for t in _pending_syncs if t.is_alive()] + _pending_syncs.append(thread) + + +def _join_pending_syncs(): + for t in list(_pending_syncs): + try: + t.join(timeout=_JOIN_TIMEOUT_SECONDS) + if t.is_alive(): + log.warn( + "VolumeCache: background sync did not finish within " + f"{_JOIN_TIMEOUT_SECONDS:.0f}s at exit; cache may be incomplete" + ) + except Exception: + # best-effort shutdown: a join failure must not block interpreter exit + pass + _pending_syncs.clear() + + +def _reset_pending_for_test(): + _pending_syncs.clear() + + +atexit.register(_join_pending_syncs) diff --git a/tests/test_serverless/test_init.py b/tests/test_serverless/test_init.py index de2120415..a41c05a25 100644 --- a/tests/test_serverless/test_init.py +++ b/tests/test_serverless/test_init.py @@ -24,7 +24,8 @@ def test_expected_public_symbols(self): 'start', 'progress_update', 'register_fitness_check', - 'runpod_version' + 'runpod_version', + 'VolumeCache' } actual_symbols = set(runpod.serverless.__all__) assert expected_symbols == actual_symbols, f"Expected {expected_symbols}, got {actual_symbols}" @@ -79,24 +80,24 @@ def test_private_symbols_not_exported(self): def test_all_covers_public_api_only(self): """Test that __all__ contains only the intended public API.""" # Get all non-private attributes from the module - module_attrs = {name for name in dir(runpod.serverless) + module_attrs = {name for name in dir(runpod.serverless) if not name.startswith('_')} - + # Filter out imported modules and types that shouldn't be public expected_private_attrs = { - 'argparse', 'json', 'os', 'signal', 'sys', 'time', + 'argparse', 'json', 'os', 'signal', 'sys', 'time', 'worker', 'rp_fastapi', 'log', 'parser', 'Any', 'Dict', # Type hints 'modules', 'utils', # Sub-modules 'RunPodLogger' # Internal logger class } - + public_attrs = module_attrs - expected_private_attrs all_symbols = set(runpod.serverless.__all__) - + # All symbols in __all__ should be actual public API assert all_symbols.issubset(public_attrs), f"__all__ contains non-public symbols: {all_symbols - public_attrs}" - + # Expected public API should be exactly what's in __all__ - expected_public_api = {'start', 'progress_update', 'register_fitness_check', 'runpod_version'} + expected_public_api = {'start', 'progress_update', 'register_fitness_check', 'runpod_version', 'VolumeCache'} assert all_symbols == expected_public_api, f"Expected {expected_public_api}, got {all_symbols}" diff --git a/tests/test_serverless/test_modules/test_job.py b/tests/test_serverless/test_modules/test_job.py index 96590c1e6..1ec5ce353 100644 --- a/tests/test_serverless/test_modules/test_job.py +++ b/tests/test_serverless/test_modules/test_job.py @@ -439,3 +439,4 @@ async def test_run_job_generator_exception(self): assert mock_log.error.call_count == 1 assert mock_log.info.call_count == 1 mock_log.info.assert_called_with("Finished running generator.", "123") + diff --git a/tests/test_serverless/test_utils/test_rp_volume_cache.py b/tests/test_serverless/test_utils/test_rp_volume_cache.py new file mode 100644 index 000000000..aeeefaaaf --- /dev/null +++ b/tests/test_serverless/test_utils/test_rp_volume_cache.py @@ -0,0 +1,813 @@ +import json +import os +import shutil +import subprocess +import tarfile +import threading +import time + +import pytest + +import runpod.serverless.utils.rp_volume_cache as vcmod + +VolumeCache = vcmod.VolumeCache + + +@pytest.fixture(autouse=True) +def _reset_pending(): + vcmod._reset_pending_for_test() + yield + vcmod._join_pending_syncs() + vcmod._reset_pending_for_test() + + +def _mk_cache_with_volume(tmp_path, namespace="ep1"): + cache = tmp_path / "cache" + cache.mkdir() + vol = tmp_path / "volume" + vol.mkdir() + vc = VolumeCache([str(cache)], namespace=namespace, volume_path=str(vol)) + return vc, cache, vol + + +# --------------------------------------------------------------------------- # +# available +# --------------------------------------------------------------------------- # + + +def test_unavailable_when_volume_dir_missing(tmp_path): + vc = VolumeCache( + [str(tmp_path / "cache")], namespace="ep1", volume_path=str(tmp_path / "no-volume") + ) + assert vc.available is False + + +def test_unavailable_without_namespace(tmp_path, monkeypatch): + monkeypatch.delenv("RUNPOD_ENDPOINT_ID", raising=False) + vol = tmp_path / "volume" + vol.mkdir() + vc = VolumeCache([str(tmp_path / "cache")], volume_path=str(vol)) + assert vc.available is False + + +def test_available_when_volume_present_and_namespace_set(tmp_path): + vc, _cache, _vol = _mk_cache_with_volume(tmp_path) + assert vc.available is True + + +def test_namespace_defaults_to_endpoint_id(tmp_path, monkeypatch): + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "endpoint-xyz") + vol = tmp_path / "volume" + vol.mkdir() + vc = VolumeCache([str(tmp_path / "cache")], volume_path=str(vol)) + assert vc._mirror_root == os.path.join(str(vol), ".cache", "endpoint-xyz") + + +# --------------------------------------------------------------------------- # +# namespace validation +# --------------------------------------------------------------------------- # + + +@pytest.mark.parametrize("bad_namespace", ["../evil", "a/b", "/etc", "..", "a\\b", "."]) +def test_namespace_rejects_unsafe_values(tmp_path, bad_namespace): + vol = tmp_path / "volume" + vol.mkdir() + with pytest.raises(ValueError): + VolumeCache([str(tmp_path / "cache")], namespace=bad_namespace, volume_path=str(vol)) + + +def test_namespace_accepts_normal_value(tmp_path): + vc, _cache, _vol = _mk_cache_with_volume(tmp_path, namespace="ep1") + assert vc._namespace == "ep1" + + +def test_scalar_dirs_path_is_treated_as_single_dir(tmp_path): + # A bare string/Path must be one directory, not iterated character-by-character + # (which would put "/" and other roots into _dirs). + vol = tmp_path / "volume" + vol.mkdir() + cache = tmp_path / "cache" + expected = [os.path.realpath(str(cache))] + + vc_str = VolumeCache(str(cache), namespace="ep1", volume_path=str(vol)) + assert vc_str._dirs == expected + + vc_path = VolumeCache(cache, namespace="ep1", volume_path=str(vol)) # os.PathLike + assert vc_path._dirs == expected + + vc_list = VolumeCache([str(cache)], namespace="ep1", volume_path=str(vol)) + assert vc_list._dirs == expected # list input still works + + +# --------------------------------------------------------------------------- # +# sync -> hydrate round trip +# --------------------------------------------------------------------------- # + + +def test_sync_then_hydrate_round_trip(tmp_path): + vc, cache, vol = _mk_cache_with_volume(tmp_path) + (cache / "model.bin").write_text("weights") + + vc.sync(background=False) + (cache / "model.bin").unlink() + + fresh = VolumeCache([str(cache)], namespace="ep1", volume_path=str(vol)) + copied = fresh.hydrate() + + assert copied == 1 + assert (cache / "model.bin").read_text() == "weights" + + +def test_sync_is_idempotent_on_unchanged_file(tmp_path): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + (cache / "model.bin").write_text("weights") + + assert vc.sync(background=False) is None + first_copied = vc._do_sync() + assert first_copied == 0 # already synced above; nothing new to copy + + +def test_sync_recopies_modified_file(tmp_path): + vc, cache, vol = _mk_cache_with_volume(tmp_path) + f = cache / "model.bin" + f.write_text("v1") + vc.sync(background=False) + + 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 + + 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): + vc, cache, vol = _mk_cache_with_volume(tmp_path) + (cache / "model.bin").write_text("weights") + vc.sync(background=False) + (cache / "model.bin").unlink() + + fresh = VolumeCache([str(cache)], namespace="ep1", volume_path=str(vol)) + assert fresh.hydrate() == 1 + assert fresh.hydrate() == 0 # already up to date + + +def test_hydrate_does_not_overwrite_newer_container_file(tmp_path): + # Same size, later mtime: _needs_copy must key off mtime here since size + # alone can't distinguish the versions. + vc, cache, vol = _mk_cache_with_volume(tmp_path) + f = cache / "model.bin" + f.write_text("weightsA") + vc.sync(background=False) + + # Container file is newer than the mirror copy (same size). + f.write_text("weightsB") + os.utime(f, (time.time() + 100, time.time() + 100)) + + fresh = VolumeCache([str(cache)], namespace="ep1", volume_path=str(vol)) + fresh.hydrate() + assert f.read_text() == "weightsB" + + +# --------------------------------------------------------------------------- # +# exclusions +# --------------------------------------------------------------------------- # + + +def test_sync_skips_excluded_paths(tmp_path): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + (cache / "refs").mkdir() + (cache / "refs" / "main").write_text("ref") + (cache / "locked.lock").write_text("lock") + + copied = vc._do_sync() + assert copied == 0 + + +def test_sync_skips_symlink_source(tmp_path): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + real = cache / "real.bin" + real.write_text("real-data") + link = cache / "link.bin" + try: + link.symlink_to(real) + except (OSError, NotImplementedError): + pytest.skip("symlinks not supported on this platform") + + copied = vc._do_sync() + assert copied == 1 # only real.bin, not the symlink + + manifest = vc._read_manifest() + manifest_paths = [e["path"] for e in manifest["small"] + manifest["big"]] + assert str(link) not in manifest_paths + + +# --------------------------------------------------------------------------- # +# hydrate destination safety +# --------------------------------------------------------------------------- # + + +def test_hydrate_skips_unsafe_destination(tmp_path): + 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) + + +def test_within_confines_source_to_base(tmp_path): + # Defense-in-depth: the big-bucket source must resolve inside big/. + vc, _cache, _vol = _mk_cache_with_volume(tmp_path) + base = vc._big_root + assert vc._within(base, os.path.join(base, "root/.cache/x")) is True + assert vc._within(base, base) is True + assert vc._within(base, os.path.join(base, "..", "escape.txt")) is False + assert vc._within(base, "/etc/passwd") is False + + +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) + + +# --------------------------------------------------------------------------- # +# context manager +# --------------------------------------------------------------------------- # + + +def test_context_manager_hydrates_on_enter_and_syncs_on_exit(tmp_path): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + (cache / "model.bin").write_text("weights") + + with vc as ctx: + assert ctx is vc + + vcmod._join_pending_syncs() + + 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): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + (cache / "model.bin").write_text("weights") + + raised = False + try: + with vc: + raise ValueError("boom") + except ValueError: + raised = True + assert raised + + vcmod._join_pending_syncs() + m = vc._read_manifest() + assert [e["path"] for e in m["small"]] == [str(cache / "model.bin")] + + +# --------------------------------------------------------------------------- # +# best-effort +# --------------------------------------------------------------------------- # + + +def test_best_effort_swallows_sync_failure_inline(tmp_path, monkeypatch): + vc, _cache, _vol = _mk_cache_with_volume(tmp_path) + + def boom(): + raise RuntimeError("disk exploded") + + monkeypatch.setattr(vc, "_do_sync", boom) + vc.sync(background=False) # must not raise + + +def test_best_effort_false_reraises_inline(tmp_path, monkeypatch): + vc, _cache, _vol = _mk_cache_with_volume(tmp_path) + vc._best_effort = False + + def boom(): + raise RuntimeError("disk exploded") + + monkeypatch.setattr(vc, "_do_sync", boom) + with pytest.raises(RuntimeError): + vc.sync(background=False) + + +def test_best_effort_swallows_sync_failure_background(tmp_path, monkeypatch): + vc, _cache, _vol = _mk_cache_with_volume(tmp_path) + + def boom(): + raise RuntimeError("disk exploded") + + monkeypatch.setattr(vc, "_do_sync", boom) + vc.sync(background=True) + vcmod._join_pending_syncs() # must not raise / must not crash the thread + + +# --------------------------------------------------------------------------- # +# unavailable no-ops +# --------------------------------------------------------------------------- # + + +# --------------------------------------------------------------------------- # +# low-level helper edge cases (coverage of defensive branches) +# --------------------------------------------------------------------------- # + + +def test_iter_files_swallows_islink_oserror(tmp_path, monkeypatch): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + (cache / "f.bin").write_text("data") + + def flaky_islink(path): + raise OSError("boom") + + monkeypatch.setattr(os.path, "islink", flaky_islink) + assert list(vc._iter_files(str(cache))) == [] + + +def test_needs_copy_false_when_src_missing(tmp_path): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + assert vc._needs_copy(str(cache / "missing.bin"), str(cache / "dst.bin")) is False + + +def test_do_sync_skips_missing_dirs(tmp_path): + vol = tmp_path / "volume" + vol.mkdir() + vc = VolumeCache([str(tmp_path / "does-not-exist")], namespace="ep1", volume_path=str(vol)) + assert vc._do_sync() == 0 + + +def test_copy_file_returns_false_and_cleans_tmp_on_failure(tmp_path, monkeypatch): + vc, cache, _vol = _mk_cache_with_volume(tmp_path) + src = cache / "f.bin" + src.write_text("data") + dst = cache / "dst.bin" + + def flaky_copy2(_src, _dst): + raise OSError("disk full") + + monkeypatch.setattr(shutil, "copy2", flaky_copy2) + assert vc._copy_file(str(src), str(dst)) is False + assert not os.path.exists(str(dst) + ".rpvc.tmp") + + +def test_join_pending_syncs_swallows_join_exception(): + class FlakyThread: + def join(self): + raise RuntimeError("join blew up") + + vcmod._pending_syncs.append(FlakyThread()) + vcmod._join_pending_syncs() # must not raise + assert vcmod._pending_syncs == [] + + +def test_hydrate_noop_when_unavailable(tmp_path): + vc = VolumeCache( + [str(tmp_path / "cache")], namespace="ep1", volume_path=str(tmp_path / "missing") + ) + assert vc.hydrate() == 0 + + +def test_sync_noop_when_unavailable(tmp_path): + vc = VolumeCache( + [str(tmp_path / "cache")], namespace="ep1", volume_path=str(tmp_path / "missing") + ) + 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 +# --------------------------------------------------------------------------- # + + +def test_volumecache_exported_from_serverless(): + from runpod import serverless as sls + from runpod.serverless import utils as sls_utils + + assert sls.VolumeCache is sls_utils.VolumeCache + 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 +# --------------------------------------------------------------------------- # + + +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") + 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_hydrate_does_not_resurrect_deleted_small_file(tmp_path): + # A small file deleted locally while other small files remain must be + # repacked out of small.tar; _changed_vs_manifest reports only + # additions/modifications, so without a deletion check the stale archive + # would resurrect the file on a later hydrate. + vc, cache, vol = _mk_cache_with_volume(tmp_path) + a = cache / "a.bin" + b = cache / "b.bin" + a.write_bytes(b"aaa") + b.write_bytes(b"bbb") + vc.sync(background=False) + assert sorted(e["path"] for e in vc._read_manifest()["small"]) == sorted( + [str(a), str(b)] + ) + + # Delete b; re-sync. a is unchanged, but the small set shrank -> repack. + b.unlink() + vc._do_sync() + assert [e["path"] for e in vc._read_manifest()["small"]] == [str(a)] + with tarfile.open(vc._small_archive_path) as tf: + assert os.path.relpath(str(b), "/") not in tf.getnames() + + # Fresh cold worker: both gone locally. Hydrate restores a, never b. + a.unlink() + fresh = VolumeCache([str(cache)], namespace="ep1", volume_path=str(vol)) + fresh.hydrate() + assert a.read_bytes() == b"aaa" + assert not b.exists() + + +def test_join_pending_syncs_bounded_by_timeout(monkeypatch): + vcmod._reset_pending_for_test() + monkeypatch.setattr(vcmod, "_JOIN_TIMEOUT_SECONDS", 0.05) + started = threading.Event() + release = threading.Event() + + def slow(): + started.set() + release.wait(5) + + t = threading.Thread(target=slow, daemon=True) + vcmod._register_pending(t) + t.start() + started.wait(1) + # Must return promptly despite the thread still running (bounded join). + 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 diff --git a/tests/test_serverless/test_utils_init.py b/tests/test_serverless/test_utils_init.py index 296f4b9aa..bff2bb337 100644 --- a/tests/test_serverless/test_utils_init.py +++ b/tests/test_serverless/test_utils_init.py @@ -23,7 +23,8 @@ def test_expected_public_symbols(self): expected_symbols = { 'download_files_from_urls', 'upload_file_to_bucket', - 'upload_in_memory_object' + 'upload_in_memory_object', + 'VolumeCache' } actual_symbols = set(runpod.serverless.utils.__all__) assert expected_symbols == actual_symbols, f"Expected {expected_symbols}, got {actual_symbols}" @@ -69,23 +70,24 @@ def test_no_duplicate_symbols_in_all(self): def test_all_covers_public_api_only(self): """Test that __all__ contains only the intended public API.""" # Get all non-private attributes from the module - module_attrs = {name for name in dir(runpod.serverless.utils) + module_attrs = {name for name in dir(runpod.serverless.utils) if not name.startswith('_')} - + # Filter out any imported modules that shouldn't be public expected_private_attrs = set() # No private imports in this module - + public_attrs = module_attrs - expected_private_attrs all_symbols = set(runpod.serverless.utils.__all__) - + # All symbols in __all__ should be actual public API assert all_symbols.issubset(public_attrs), f"__all__ contains non-public symbols: {all_symbols - public_attrs}" - + # Expected public API should be exactly what's in __all__ expected_public_api = { - 'download_files_from_urls', - 'upload_file_to_bucket', - 'upload_in_memory_object' + 'download_files_from_urls', + 'upload_file_to_bucket', + 'upload_in_memory_object', + 'VolumeCache' } assert all_symbols == expected_public_api, f"Expected {expected_public_api}, got {all_symbols}"