diff --git a/README.md b/README.md index b34014b1..6ad9c666 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 00000000..be780458 --- /dev/null +++ b/docs/serverless/volume_cache.md @@ -0,0 +1,89 @@ +# 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. | + +## 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. +- **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. + +## 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. diff --git a/examples/serverless/volume_cache_handler.py b/examples/serverless/volume_cache_handler.py new file mode 100644 index 00000000..a0b4e08c --- /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 7905731f..05245207 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 9036a4dd..f0c44cc8 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 00000000..6c6d1342 --- /dev/null +++ b/runpod/serverless/utils/rp_volume_cache.py @@ -0,0 +1,249 @@ +"""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). + +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``). +""" + +import atexit +import os +import shutil +import threading + +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 + + +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. + + 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): + 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 + + @property + def _mirror_root(self): + return os.path.join(self._volume_path, ".cache", self._namespace) + + @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 + + @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 _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 + + # ----------------------------------------------------------------- # + # 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): + root = self._mirror_root + if not os.path.isdir(root): + 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 + + 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): + os.makedirs(self._mirror_root, exist_ok=True) + copied = 0 + 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 + + # ----------------------------------------------------------------- # + # context manager + # ----------------------------------------------------------------- # + + def __enter__(self): + self.hydrate() + return self + + def __exit__(self, *exc): + self.sync(background=True) + return None + + +# ----------------------------------------------------------------------- # +# 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 de212041..a41c05a2 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 96590c1e..1ec5ce35 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 00000000..4c85bd74 --- /dev/null +++ b/tests/test_serverless/test_utils/test_rp_volume_cache.py @@ -0,0 +1,399 @@ +import os +import shutil +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" + + +# --------------------------------------------------------------------------- # +# 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)) + + 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" + + +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 + + mirror_link = os.path.join(vc._mirror_root, os.path.relpath(str(link), "/")) + assert not os.path.exists(mirror_link) + + +# --------------------------------------------------------------------------- # +# hydrate destination safety +# --------------------------------------------------------------------------- # + + +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: + fh.write("malicious") + + copied = vc.hydrate() + assert copied == 0 + assert not (tmp_path / "outside" / "evil.txt").exists() + + +# --------------------------------------------------------------------------- # +# 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() + + mirror_file = os.path.join(vc._mirror_root, os.path.relpath(str(cache / "model.bin"), "/")) + assert os.path.exists(mirror_file) + + +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() + mirror_file = os.path.join(vc._mirror_root, os.path.relpath(str(cache / "model.bin"), "/")) + assert os.path.exists(mirror_file) + + +# --------------------------------------------------------------------------- # +# 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 + + +# --------------------------------------------------------------------------- # +# 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__ + + +# --------------------------------------------------------------------------- # +# 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") + 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) + + +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() diff --git a/tests/test_serverless/test_utils_init.py b/tests/test_serverless/test_utils_init.py index 296f4b9a..bff2bb33 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}"