diff --git a/README.md b/README.md index b34014b1..b84c089f 100644 --- a/README.md +++ b/README.md @@ -149,6 +149,26 @@ 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 — restoring them on startup and syncing new files back — so a repeated multi-GB model download becomes a one-time cost per endpoint. It is opt-in and best-effort. + +```bash +# Enable the built-in warm cache (requires a mounted network volume) +RUNPOD_VOLUME_CACHE=1 +``` + +Or use it explicitly around a model load: + +```python +from runpod.serverless import VolumeCache + +with VolumeCache(dirs=["/root/.cache/huggingface"]).warm(): + 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..20fb4598 --- /dev/null +++ b/docs/serverless/volume_cache.md @@ -0,0 +1,139 @@ +# Network-Volume Warm Cache (VolumeCache) + +`VolumeCache` warms local directories across serverless workers using a mounted +network volume. On cold start it restores previously-synced directories (for +example a model cache) from the volume; after use it syncs newly written files +back so the next cold worker starts warm. 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 the 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 cache + per endpoint). + +If no volume is mounted, every operation is a safe no-op. + +## Option 1: Built-in (opt-in, zero code) + +The worker can hydrate a model cache at startup and sync it after the first job, +with no changes to your handler. It is **off by default** — enable it with an +environment variable: + +```bash +# Enable the built-in warm cache (opt-in) +RUNPOD_VOLUME_CACHE=1 +``` + +When enabled and a volume is mounted, the worker: + +1. Hydrates the model cache from the volume before the first job runs. +2. Syncs new files to the volume once, after the first successful job. + +By default it caches the directories pointed to by `HF_HOME`, `HF_HUB_CACHE`, +and `TORCH_HOME` (whichever are set; `HF_HOME` defaults to +`~/.cache/huggingface`). Add more directories with `RUNPOD_CACHE_DIRS`. + +### Configuration + +| Variable | Default | Purpose | +| --- | --- | --- | +| `RUNPOD_VOLUME_CACHE` | unset (off) | Set to `1`/`true`/`yes`/`on` to enable the built-in. | +| `RUNPOD_CACHE_DIRS` | — | Extra directories to cache, `os.pathsep`-separated (`:` on Linux). | +| `RUNPOD_VOLUME_CACHE_MAX_GB` | `50` | Prune oldest shards once the endpoint's cache exceeds this size. | +| `HF_HOME` / `HF_HUB_CACHE` / `TORCH_HOME` | — | Auto-discovered model-cache locations. | + +### Example + +For most model-serving workers, enabling the built-in and letting the model +download into `HF_HOME` on the first request is all that is needed: + +```python +# my_worker.py +import runpod + +def handler(job): + # First cold worker downloads the model into HF_HOME; the built-in syncs it + # to the volume. Subsequent cold workers hydrate it and skip the download. + ... + +runpod.serverless.start({"handler": handler}) +``` + +```bash +# Template / container env +RUNPOD_VOLUME_CACHE=1 +``` + +> Note: hydration only helps model loads that happen **after** the worker starts +> (lazy or in-handler loading). If your model loads at module import time, use +> Option 2 to place `warm()` around the load. + +## Option 2: Explicit API + +Import `VolumeCache` and control caching yourself. This works for any +directories, and the `warm()` context manager guarantees hydration happens +before your model load regardless of when it runs. + +```python +from runpod.serverless import VolumeCache + +vc = VolumeCache(dirs=["/root/.cache/huggingface"]) + +with vc.warm(): # hydrate on enter, sync the delta on exit + model = load_model() # downloads land in the cached directory +``` + +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", max_size_gb=100) + +vc.hydrate() # restore cached files (e.g. at startup) +model = load_model() # populate the cache +vc.sync() # persist new files back to the volume +``` + +### Constructor + +| Argument | Default | Purpose | +| --- | --- | --- | +| `dirs` | required | Local directories to cache. | +| `namespace` | `RUNPOD_ENDPOINT_ID` | Isolation key for the on-volume shards. Must be a single safe path component. | +| `volume_path` | `/runpod-volume` | Network-volume mount point. | +| `max_size_gb` | `None` | Prune oldest shards past this cap; `None` = no cap. | +| `best_effort` | `True` | Swallow and log errors instead of raising. Set `False` while debugging. | + +## How it works + +- **Per-endpoint, sharded storage.** Each worker writes its delta as its own tar + shard under `//.cache//`, published atomically. Per-worker + shards are collision-free under concurrent workers — no shared file is + rewritten. Hydration extracts all shards for the namespace (newest wins on + overlap). +- **Delta only.** `sync()` packs just the files written since the last hydrate, + not the whole cache. +- **Retention.** With `max_size_gb` set, the oldest shards are pruned once the + namespace exceeds the cap. +- **Safety.** Extraction rejects members that would escape the configured + directories (traversal, symlinks, absolute paths). + +## Limitations + +- **Cold-scale write amplification.** If N workers cold-start at the same time, + each misses the still-empty cache, downloads the model, and writes a full-size + shard. Total volume writes on the first scale-up are ~N×; bounded thereafter by + retention. +- **Lost first warm on aggressive recycle.** `sync()` runs in a background thread + after the job response. A large-model sync may not finish before the worker is + recycled; the partial shard is cleaned up and the warm is retried on the next + worker. + +## Status + +The built-in is opt-in for now (`RUNPOD_VOLUME_CACHE=1`). Whether it should +default on when a volume is present is still under review. 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/modules/rp_job.py b/runpod/serverless/modules/rp_job.py index a45cebc6..94de75b0 100644 --- a/runpod/serverless/modules/rp_job.py +++ b/runpod/serverless/modules/rp_job.py @@ -15,6 +15,7 @@ from ...version import __version__ as runpod_version from ..utils import rp_debugger +from ..utils.rp_volume_cache import sync_after_job from .rp_handler import is_generator from .rp_http import send_result, stream_result from .rp_tips import check_return_size @@ -282,6 +283,8 @@ async def run_job(handler: Callable, job: Dict[str, Any]) -> Dict[str, Any]: if run_result.get("output") == {}: run_result.pop("output") + sync_after_job() # fire-and-forget cache warm on any successful output + check_return_size(run_result) # Checks the size of the return body. except Exception as err: @@ -329,6 +332,8 @@ async def run_job_generator( log.debug(f"Generator output: {output_partial}", job["id"]) yield {"output": output_partial} + sync_after_job() # fire-and-forget cache warm on successful generator completion + except Exception as err: log.error(err, job["id"]) yield {"error": f"handler: {str(err)} \ntraceback: {traceback.format_exc()}"} 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..1523526f --- /dev/null +++ b/runpod/serverless/utils/rp_volume_cache.py @@ -0,0 +1,299 @@ +"""Bidirectional warm-cache sync between local directories and a network volume.""" + +import contextlib +import glob +import os +import tarfile +import tempfile +import time +import uuid +import threading + +from runpod.serverless.modules.rp_logger import RunPodLogger + +log = RunPodLogger() + +_BASELINE_EPSILON_SECONDS = 2.0 # tolerate coarse (1s) network-filesystem mtime granularity + + +class VolumeCache: + """Warm-cache local directories across serverless workers via a network volume. + + On cold start, ``hydrate()`` extracts previously-synced directories from a + mounted network volume into place; after use, ``sync()`` packs the newly + written files into a per-worker tar shard on the volume so the next cold + worker starts warm. Stdlib-only and best-effort: any failure degrades to a + cold worker and never raises into the caller. + + 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 shards. Must be a single safe + path component. Defaults to ``RUNPOD_ENDPOINT_ID``. + volume_path: Network-volume mount point. Defaults to ``/runpod-volume``. + max_size_gb: Prune oldest shards past this cap. ``None`` = no cap. + best_effort: When True (default), swallow and log errors instead of + raising. + + Example: + >>> vc = VolumeCache(dirs=["/root/.cache/huggingface"]) + >>> with vc.warm(): # hydrate on enter, sync delta on exit + ... 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", + max_size_gb=None, + 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._max_size_gb = max_size_gb + self._best_effort = best_effort + self._worker_id = os.environ.get("RUNPOD_POD_ID") or uuid.uuid4().hex[:12] + self._baseline = time.time() - _BASELINE_EPSILON_SECONDS + + @property + def _shard_dir(self): + return os.path.join(self._volume_path, ".cache", self._namespace) + + @property + def available(self): + return bool(self._namespace) and os.path.isdir(self._volume_path) + + def _list_shards(self): + d = self._shard_dir + if not os.path.isdir(d): + return [] + shards = [os.path.join(d, f) for f in os.listdir(d) if f.endswith(".tar")] + return sorted(shards, key=os.path.getmtime) + + def _iter_delta_files(self): + for root in self._dirs: + if not os.path.isdir(root): + continue + for dirpath, _dirs, files in os.walk(root): + for name in files: + path = os.path.join(dirpath, name) + if name.endswith(".lock") or name.startswith(".rp_volume_cache"): + continue + if any(sub in path for sub in self._EXCLUDE_SUBSTRINGS): + continue + try: + if os.path.getmtime(path) > self._baseline: + yield path + except OSError: + continue + + 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 sync(self): + if not self.available: + return False + return self._guard(self._do_sync, False) + + def _do_sync(self): + files = list(self._iter_delta_files()) + if not files: + log.debug("VolumeCache: no delta files to sync") + return False + os.makedirs(self._shard_dir, exist_ok=True) + for stale in glob.glob(os.path.join(self._shard_dir, f"{self._worker_id}-*.tar.tmp")): + try: + os.remove(stale) + except OSError: + # best-effort cleanup: ignore temp files that vanish or can't be removed + pass + final = os.path.join(self._shard_dir, f"{self._worker_id}-{time.time_ns():020d}.tar") + tmp = final + ".tmp" + with tarfile.open(tmp, "w") as tar: + for path in files: + tar.add(path, arcname=os.path.relpath(path, "/")) + os.replace(tmp, final) + log.info(f"VolumeCache: synced {len(files)} files to {final}") + self._baseline = time.time() - _BASELINE_EPSILON_SECONDS + self._enforce_retention() + return True + + def _enforce_retention(self): + if not self._max_size_gb: + return + def _size(p): + try: + return os.path.getsize(p) + except OSError: + return 0 + + cap = self._max_size_gb * (1024 ** 3) + shards = self._list_shards() # oldest first + total = sum(_size(s) for s in shards) + for shard in shards: + if total <= cap: + break + size = _size(shard) + try: + os.remove(shard) + total -= size + log.info(f"VolumeCache: pruned old shard {shard}") + except OSError as exc: + log.warn(f"VolumeCache: failed to prune {shard}: {exc}") + + @property + def _marker_path(self): + base = os.path.join(tempfile.gettempdir(), "rp_volume_cache") + return os.path.join(base, f"{self._namespace}.hydrated") + + def _newest_shard_mtime(self): + shards = self._list_shards() + return os.path.getmtime(shards[-1]) if shards else 0.0 + + def _clear_marker_for_test(self): + if os.path.exists(self._marker_path): + os.remove(self._marker_path) + + def hydrate(self): + if not self.available: + return False + return self._guard(self._do_hydrate, False) + + def _do_hydrate(self): + shards = self._list_shards() + if not shards: + return False + newest = self._newest_shard_mtime() + if os.path.exists(self._marker_path) and os.path.getmtime(self._marker_path) >= newest: + log.debug("VolumeCache: cache already hydrated, skipping") + return False + extracted = False + # Use the tar data filter for defense-in-depth where the runtime provides + # it (Python 3.12+, and 3.10.12+/3.11.4+ backports) without requiring it -- + # the >=3.10 floor may predate the API. _is_safe_member is the primary guard. + extract_kwargs = {"filter": "data"} if hasattr(tarfile, "data_filter") else {} + for shard in shards: # oldest -> newest (last wins) + with tarfile.open(shard) as tar: + safe = [m for m in tar.getmembers() if self._is_safe_member(m)] + tar.extractall(path="/", members=safe, **extract_kwargs) + extracted = extracted or bool(safe) + os.makedirs(os.path.dirname(self._marker_path), exist_ok=True) + with open(self._marker_path, "w") as fh: + fh.write(str(newest)) + os.utime(self._marker_path, (newest, newest)) + self._baseline = time.time() - _BASELINE_EPSILON_SECONDS + if extracted: + log.info(f"VolumeCache: hydrated from {len(shards)} shard(s)") + return extracted + + def _is_safe_member(self, member): + if not (member.isfile() or member.isdir()): + return False # reject symlink/hardlink/device/fifo + target = os.path.realpath(os.path.join("/", member.name)) + return any( + target == d or target.startswith(d + os.sep) + for d in self._dirs + ) + + @contextlib.contextmanager + def warm(self): + self.hydrate() + try: + yield self + finally: + self.sync() + + +_ACTIVE_CACHE = None +_SYNCED = False +_sync_lock = threading.Lock() + +_ENABLED_VALUES = ("1", "true", "yes", "on") + + +def _discover_model_dirs(): + dirs = [os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface")] + for var in ("HF_HUB_CACHE", "TORCH_HOME"): + if os.environ.get(var): + dirs.append(os.environ[var]) + extra = os.environ.get("RUNPOD_CACHE_DIRS") + if extra: + dirs.extend(p for p in extra.split(os.pathsep) if p) + return list(dict.fromkeys(dirs)) # de-dupe, preserve order + + +def build_default_cache(): + """Build the built-in VolumeCache when opt-in is enabled, else None. + + The built-in is OFF by default. It activates only when RUNPOD_VOLUME_CACHE is + set to a truthy value ("1"/"true"/"yes"/"on") AND a network volume is mounted. + + Known operational limitations (accepted, not bugs): + 1. Cold-scale write amplification: when N workers cold-start simultaneously, + each misses the still-empty cache and each writes its own full-size shard + on first sync. Bounded by retention (RUNPOD_VOLUME_CACHE_MAX_GB). + 2. Lost first warm on aggressive recycle: sync() runs in a daemon thread + dispatched after the job response; a large-model sync may not finish + before the worker is recycled, killing the thread mid-write. The partial + shard is cleaned up on the next sync from this worker, and the warm is + simply retried then. + """ + if os.environ.get("RUNPOD_VOLUME_CACHE", "").lower() not in _ENABLED_VALUES: + return None + try: + max_gb = float(os.environ.get("RUNPOD_VOLUME_CACHE_MAX_GB", "50")) + vc = VolumeCache(_discover_model_dirs(), max_size_gb=max_gb) + except Exception as exc: # never let cache setup crash worker startup + log.warn(f"VolumeCache: failed to build default cache, disabling: {exc}") + return None + return vc if vc.available else None + + +def set_active_cache(vc): + global _ACTIVE_CACHE + _ACTIVE_CACHE = vc + + +def sync_after_job(): + global _SYNCED + try: + if _ACTIVE_CACHE is None: + return + with _sync_lock: + if _SYNCED: + return + _SYNCED = True + threading.Thread(target=_ACTIVE_CACHE.sync, daemon=True).start() + except Exception as exc: # cache glue must never affect job outcome + log.warn(f"VolumeCache: sync_after_job failed to dispatch: {exc}") + + +def reset_builtin_state_for_test(): + global _ACTIVE_CACHE, _SYNCED + _ACTIVE_CACHE = None + _SYNCED = False diff --git a/runpod/serverless/worker.py b/runpod/serverless/worker.py index 90053ec7..d50df6ec 100644 --- a/runpod/serverless/worker.py +++ b/runpod/serverless/worker.py @@ -9,6 +9,7 @@ from runpod.serverless.modules import rp_logger, rp_local, rp_ping, rp_scale from runpod.serverless.modules.rp_fitness import run_fitness_checks +from runpod.serverless.utils.rp_volume_cache import build_default_cache, set_active_cache log = rp_logger.RunPodLogger() heartbeat = rp_ping.Heartbeat() @@ -48,6 +49,12 @@ def run_worker(config: Dict[str, Any]) -> None: # Start pinging Runpod to show that the worker is alive. heartbeat.start_ping(mirror) + # Warm the local cache from a mounted network volume before any job runs. + volume_cache = build_default_cache() + if volume_cache is not None: + volume_cache.hydrate() + set_active_cache(volume_cache) + # Create a JobScaler responsible for adjusting the concurrency job_scaler = rp_scale.JobScaler(config) job_scaler.start() 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..410d01fa 100644 --- a/tests/test_serverless/test_modules/test_job.py +++ b/tests/test_serverless/test_modules/test_job.py @@ -439,3 +439,69 @@ 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") + + async def test_run_job_generator_success_syncs_cache(self): + """ + Tests that run_job_generator triggers sync_after_job on the success path. + """ + handler = self.handler_gen_success + job = {"id": "123"} + counter = {"n": 0} + + def fake_sync_after_job(): + counter["n"] += 1 + + with patch( + "runpod.serverless.modules.rp_job.log", new_callable=Mock + ), patch( + "runpod.serverless.modules.rp_job.sync_after_job", + side_effect=fake_sync_after_job, + ) as mock_sync: + result = [i async for i in rp_job.run_job_generator(handler, job)] + + assert result == [ + {"output": "partial_output_1"}, + {"output": "partial_output_2"}, + ] + mock_sync.assert_called_once() + assert counter["n"] == 1 + + async def test_run_job_generator_success_syncs_cache_async(self): + """ + Tests that run_job_generator triggers sync_after_job on the success path + for an async generator handler. + """ + handler = self.handler_async_gen_success + job = {"id": "123"} + + with patch( + "runpod.serverless.modules.rp_job.log", new_callable=Mock + ), patch( + "runpod.serverless.modules.rp_job.sync_after_job" + ) as mock_sync: + result = [i async for i in rp_job.run_job_generator(handler, job)] + + assert result == [ + {"output": "partial_output_1"}, + {"output": "partial_output_2"}, + ] + mock_sync.assert_called_once() + + async def test_run_job_generator_exception_does_not_sync_cache(self): + """ + Tests that run_job_generator does NOT trigger sync_after_job on the + exception path. + """ + handler = self.handler_fail + job = {"id": "123"} + + with patch( + "runpod.serverless.modules.rp_job.log", new_callable=Mock + ), patch( + "runpod.serverless.modules.rp_job.sync_after_job" + ) as mock_sync: + result = [i async for i in rp_job.run_job_generator(handler, job)] + + assert len(result) == 1 + assert "error" in result[0] + mock_sync.assert_not_called() 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..0a216e60 --- /dev/null +++ b/tests/test_serverless/test_utils/test_rp_volume_cache.py @@ -0,0 +1,415 @@ +import os +import tarfile +import time +import pytest +import runpod.serverless.utils.rp_volume_cache as vcmod + +VolumeCache = vcmod.VolumeCache + + +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_available_when_volume_present_and_namespace_set(tmp_path): + vol = tmp_path / "volume" + vol.mkdir() + vc = VolumeCache([str(tmp_path / "cache")], namespace="ep1", volume_path=str(vol)) + assert vc.available is True + + +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 + + +@pytest.mark.parametrize("bad_namespace", ["../evil", "a/b", "/etc", ".."]) +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): + vol = tmp_path / "volume" + vol.mkdir() + vc = VolumeCache([str(tmp_path / "cache")], namespace="ep1", volume_path=str(vol)) + assert vc._namespace == "ep1" + + +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._shard_dir == os.path.join(str(vol), ".cache", "endpoint-xyz") + + +def _mk_cache_with_volume(tmp_path): + cache = tmp_path / "cache" + cache.mkdir() + vol = tmp_path / "volume" + vol.mkdir() + vc = VolumeCache([str(cache)], namespace="ep1", volume_path=str(vol)) + return vc, cache, vol + + +def test_sync_packs_files_created_after_baseline(tmp_path): + vc, cache, vol = _mk_cache_with_volume(tmp_path) + vc._baseline = time.time() - 5 + (cache / "model.bin").write_text("weights") + assert vc.sync() is True + shards = vc._list_shards() + assert len(shards) == 1 + with tarfile.open(shards[0]) as tar: + names = tar.getnames() + assert os.path.relpath(str(cache / "model.bin"), "/") in names + + +def test_sync_excludes_files_older_than_baseline(tmp_path): + vc, cache, vol = _mk_cache_with_volume(tmp_path) + old = cache / "old.bin" + old.write_text("x") + os.utime(old, (time.time() - 100, time.time() - 100)) + vc._baseline = time.time() + assert vc.sync() is False + + +def test_sync_skips_excluded_paths(tmp_path): + vc, cache, vol = _mk_cache_with_volume(tmp_path) + vc._baseline = time.time() - 5 + (cache / "refs").mkdir() + (cache / "refs" / "main").write_text("ref") + assert vc.sync() is False + + +def test_sync_shard_names_are_unique_per_call(tmp_path): + vc, cache, vol = _mk_cache_with_volume(tmp_path) + vc._baseline = time.time() - 5 + (cache / "a.bin").write_text("a") + vc.sync() + (cache / "b.bin").write_text("b") + vc.sync() + assert len(vc._list_shards()) == 2 + + +def test_sync_noop_when_unavailable(tmp_path): + vc = VolumeCache([str(tmp_path / "c")], namespace="ep1", + volume_path=str(tmp_path / "missing")) + assert vc.sync() is False + + +def test_sync_tolerates_coarse_mtime_granularity(tmp_path): + # A fresh instance sets baseline to now - epsilon. A file whose mtime is + # floored to the current integer second (as coarse NFS filesystems report) + # must still be picked up, not silently dropped. + cache = tmp_path / "cache" + cache.mkdir() + vol = tmp_path / "volume" + vol.mkdir() + vc = VolumeCache([str(cache)], namespace="ep1", volume_path=str(vol)) + f = cache / "model.bin" + f.write_text("weights") + now = time.time() + os.utime(f, (float(int(now)), float(int(now)))) # floor mtime to integer second + assert vc.sync() is True + + +def test_hydrate_noop_when_no_shards(tmp_path): + vc, cache, vol = _mk_cache_with_volume(tmp_path) + assert vc.hydrate() is False + + +def test_hydrate_restores_files_to_absolute_paths(tmp_path): + vc, cache, vol = _mk_cache_with_volume(tmp_path) + vc._baseline = time.time() - 5 + (cache / "model.bin").write_text("weights") + vc.sync() + (cache / "model.bin").unlink() # simulate a fresh cold worker + assert vc.hydrate() is True + assert (cache / "model.bin").read_text() == "weights" + + +def test_hydrate_later_shard_overwrites_earlier(tmp_path): + vc, cache, vol = _mk_cache_with_volume(tmp_path) + vc._baseline = time.time() - 5 + f = cache / "model.bin" + f.write_text("v1") + vc.sync() + time.sleep(0.01) + f.write_text("v2") + vc._baseline = time.time() - 5 + vc.sync() + f.unlink() + vc._clear_marker_for_test() + vc.hydrate() + assert f.read_text() == "v2" + + +def test_hydrate_is_idempotent_via_marker(tmp_path): + vc, cache, vol = _mk_cache_with_volume(tmp_path) + vc._baseline = time.time() - 5 + (cache / "model.bin").write_text("weights") + vc.sync() + assert vc.hydrate() is True # first hydrate extracts + assert vc.hydrate() is False # marker current -> no-op + + +def test_rejects_member_outside_configured_dirs(tmp_path): + vc, cache, vol = _mk_cache_with_volume(tmp_path) + m = tarfile.TarInfo(name="etc/passwd") # resolves to /etc/passwd, outside cache + m.type = tarfile.REGTYPE + assert vc._is_safe_member(m) is False + + +def test_rejects_symlink_member(tmp_path): + vc, cache, vol = _mk_cache_with_volume(tmp_path) + rel = os.path.relpath(str(cache / "link"), "/") + m = tarfile.TarInfo(name=rel) + m.type = tarfile.SYMTYPE + m.linkname = "/etc/passwd" + assert vc._is_safe_member(m) is False + + +def test_accepts_regular_member_inside_dirs(tmp_path): + vc, cache, vol = _mk_cache_with_volume(tmp_path) + rel = os.path.relpath(str(cache / "model.bin"), "/") + m = tarfile.TarInfo(name=rel) + m.type = tarfile.REGTYPE + assert vc._is_safe_member(m) is True + + +def test_rejects_sibling_prefix_collision(tmp_path): + # An allowed dir ".../cache" must NOT match a sibling ".../cache-evil"; + # this locks the separator-anchored prefix check against substring regressions. + vc, cache, vol = _mk_cache_with_volume(tmp_path) + evil = tmp_path / "cache-evil" + evil.mkdir() + rel = os.path.relpath(str(evil / "x"), "/") + m = tarfile.TarInfo(name=rel) + m.type = tarfile.REGTYPE + assert vc._is_safe_member(m) is False + + +def test_rejects_hardlink_member(tmp_path): + vc, cache, vol = _mk_cache_with_volume(tmp_path) + rel = os.path.relpath(str(cache / "hl"), "/") + m = tarfile.TarInfo(name=rel) + m.type = tarfile.LNKTYPE + m.linkname = "root/.cache/x" + assert vc._is_safe_member(m) is False + + +def test_retention_prunes_oldest_shards_past_cap(tmp_path): + cache = tmp_path / "cache"; cache.mkdir() + vol = tmp_path / "volume"; vol.mkdir() + cap_bytes = 1500 + vc = VolumeCache([str(cache)], namespace="ep1", volume_path=str(vol), + max_size_gb=cap_bytes / (1024 ** 3)) + for i in range(4): + vc._baseline = time.time() - 5 + (cache / f"f{i}.bin").write_text("x" * 800) + vc.sync() + time.sleep(0.01) + total = sum(os.path.getsize(s) for s in vc._list_shards()) + assert total <= cap_bytes + + +def test_retention_tolerates_shard_removed_concurrently(tmp_path, monkeypatch): + # A concurrent worker prunes/removes a shard between _list_shards() and the + # getsize() lookups; the size lookup must degrade to 0 instead of raising. + cache = tmp_path / "cache"; cache.mkdir() + vol = tmp_path / "volume"; vol.mkdir() + cap_bytes = 1000 + vc = VolumeCache([str(cache)], namespace="ep1", volume_path=str(vol), + max_size_gb=cap_bytes / (1024 ** 3)) + for i in range(3): + vc._baseline = time.time() - 5 + (cache / f"f{i}.bin").write_text("x" * 800) + vc.sync() + time.sleep(0.01) + + real_getsize = os.path.getsize + + def flaky_getsize(path): + if str(path).endswith(".tar") and "f0" not in str(path): + raise FileNotFoundError(path) + return real_getsize(path) + + monkeypatch.setattr(os.path, "getsize", flaky_getsize) + # Should not raise despite getsize() failures on some shards. + vc._enforce_retention() + + +def test_no_retention_when_cap_is_none(tmp_path): + vc, cache, vol = _mk_cache_with_volume(tmp_path) + for i in range(3): + vc._baseline = time.time() - 5 + (cache / f"f{i}.bin").write_text("data") + vc.sync() + assert len(vc._list_shards()) == 3 + + +def test_best_effort_swallows_and_returns_default(tmp_path, monkeypatch): + vc, cache, vol = _mk_cache_with_volume(tmp_path) + def boom(): + raise RuntimeError("disk exploded") + monkeypatch.setattr(vc, "_do_sync", boom) + (cache / "x.bin").write_text("x") + vc._baseline = time.time() - 5 + assert vc.sync() is False # swallowed, no raise + + +def test_best_effort_false_reraises(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() + + +def test_warm_hydrates_on_enter_and_syncs_on_exit(tmp_path): + vc, cache, vol = _mk_cache_with_volume(tmp_path) + calls = [] + vc.hydrate = lambda: calls.append("hydrate") + vc.sync = lambda: calls.append("sync") + with vc.warm(): + calls.append("body") + assert calls == ["hydrate", "body", "sync"] + + +def test_warm_syncs_even_on_exception(tmp_path): + vc, cache, vol = _mk_cache_with_volume(tmp_path) + calls = [] + vc.hydrate = lambda: calls.append("hydrate") + vc.sync = lambda: calls.append("sync") + raised = False + try: + with vc.warm(): + raise ValueError("boom") + except ValueError: + raised = True + assert raised + assert calls == ["hydrate", "sync"] + + +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__ + + +def test_build_default_cache_off_by_default(tmp_path, monkeypatch): + # Opt-in: with RUNPOD_VOLUME_CACHE unset, the built-in is disabled even + # when a volume would be available. + monkeypatch.delenv("RUNPOD_VOLUME_CACHE", raising=False) + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "ep1") + monkeypatch.setattr(vcmod.VolumeCache, "available", property(lambda self: True)) + assert vcmod.build_default_cache() is None + + +def test_build_default_cache_disabled_by_env(tmp_path, monkeypatch): + monkeypatch.setenv("RUNPOD_VOLUME_CACHE", "0") + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "ep1") + assert vcmod.build_default_cache() is None + + +def test_build_default_cache_none_when_no_volume(tmp_path, monkeypatch): + monkeypatch.setenv("RUNPOD_VOLUME_CACHE", "1") + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "ep1") + monkeypatch.setattr(vcmod.VolumeCache, "available", property(lambda self: False)) + assert vcmod.build_default_cache() is None + + +def test_discover_model_dirs_includes_hf_home_and_extras(monkeypatch): + monkeypatch.setenv("HF_HOME", "/models/hf") + monkeypatch.setenv("RUNPOD_CACHE_DIRS", "/a" + os.pathsep + "/b") + dirs = vcmod._discover_model_dirs() + assert "/models/hf" in dirs and "/a" in dirs and "/b" in dirs + + +def test_sync_after_job_runs_once(monkeypatch): + vcmod.reset_builtin_state_for_test() + counter = {"n": 0} + fake = type("F", (), {"sync": lambda self: counter.__setitem__("n", counter["n"] + 1)})() + joined = [] + class FakeThread: + def __init__(self, target, daemon=None): self.target = target + def start(self): self.target(); joined.append(True) + monkeypatch.setattr(vcmod.threading, "Thread", FakeThread) + vcmod.set_active_cache(fake) + vcmod.sync_after_job() + vcmod.sync_after_job() + assert counter["n"] == 1 + + +def test_run_worker_hydrates_registered_cache(monkeypatch): + from runpod.serverless import worker + + vcmod.reset_builtin_state_for_test() + try: + async def _noop_fitness_checks(): + return None + + fake = type("F", (), { + "hydrate": lambda self: fake_calls.append("h"), + "sync": lambda self: None, + })() + fake_calls = [] + monkeypatch.setattr(worker, "build_default_cache", lambda: fake, raising=False) + monkeypatch.setattr(worker.rp_scale, "JobScaler", + lambda config: type("J", (), {"start": lambda self: None})()) + monkeypatch.setattr(worker.heartbeat, "start_ping", lambda mirror: None) + monkeypatch.setattr(worker, "run_fitness_checks", _noop_fitness_checks) + worker.run_worker({"handler": lambda job: job}) + assert "h" in fake_calls + finally: + vcmod.reset_builtin_state_for_test() + + +def test_build_default_cache_survives_bad_max_gb(monkeypatch): + vcmod.reset_builtin_state_for_test() + monkeypatch.setenv("RUNPOD_VOLUME_CACHE", "1") + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "ep1") + monkeypatch.setenv("RUNPOD_VOLUME_CACHE_MAX_GB", "not-a-number") + assert vcmod.build_default_cache() is None # degrades, does not raise + + +def test_two_workers_produce_independently_hydratable_shards(tmp_path): + # Distinct worker_ids must write non-colliding shards that both hydrate (spec acceptance). + cache = tmp_path / "cache"; cache.mkdir() + vol = tmp_path / "volume"; vol.mkdir() + a = VolumeCache([str(cache)], namespace="ep1", volume_path=str(vol)) + a._worker_id = "workerA" + a._baseline = time.time() - 5 + (cache / "a.bin").write_text("aaa") + assert a.sync() is True + b = VolumeCache([str(cache)], namespace="ep1", volume_path=str(vol)) + b._worker_id = "workerB" + b._baseline = time.time() - 5 + (cache / "b.bin").write_text("bbb") + assert b.sync() is True + (cache / "a.bin").unlink(); (cache / "b.bin").unlink() + reader = VolumeCache([str(cache)], namespace="ep1", volume_path=str(vol)) + reader._clear_marker_for_test() + assert reader.hydrate() is True + assert (cache / "a.bin").read_text() == "aaa" + assert (cache / "b.bin").read_text() == "bbb" + + +def test_build_default_cache_returns_instance_when_available(monkeypatch): + vcmod.reset_builtin_state_for_test() + monkeypatch.setenv("RUNPOD_VOLUME_CACHE", "1") + monkeypatch.setenv("RUNPOD_ENDPOINT_ID", "ep1") + monkeypatch.setattr(vcmod.VolumeCache, "available", property(lambda self: True)) + vc = vcmod.build_default_cache() + assert isinstance(vc, vcmod.VolumeCache) 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}"