feat(serverless): network-volume warm cache (VolumeCache)#531
Conversation
…test state Wrap build_default_cache() parse+construct in try/except to prevent malformed RUNPOD_VOLUME_CACHE_MAX_GB from crashing worker startup; degrades to cache-disabled instead. Add sync() no-op to fake cache in test_run_worker_hydrates_registered_cache and wrap test body with reset_builtin_state_for_test() to prevent module-state leakage to subsequent tests (sync_after_job could fire and hit AttributeError).
|
Capy auto-review is paused for this organization because the usage-cycle auto-review limit has been reached. Increase the limit or turn it off in billing settings to resume automatic reviews. |
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
![github-code-quality[bot]](https://avatars.githubusercontent.com/u/9919?s=60&v=4){kind=small}
|
Promptless prepared a documentation update related to this change. Triggered by runpod-python PR #531 Documents the new Serverless network-volume warm cache: a new page covering the Review: Document Serverless network-volume warm cache (VolumeCache) |
…che-for-serverless-volumecache
There was a problem hiding this comment.
Pull request overview
This PR introduces a reusable serverless SDK primitive (VolumeCache) for warming local model/cache directories from a mounted network volume and syncing back deltas as per-worker tar “shards”, then wires it into the serverless worker lifecycle to persist model weights across worker recycling (SLS-367).
Changes:
- Added
runpod.serverless.utils.rp_volume_cache.VolumeCachewithhydrate(),sync(), andwarm()plus built-in helpers (build_default_cache,sync_after_job). - Auto-integrated the cache into the worker loop (
run_workerhydrates/registers) and job execution (run_jobdispatches a one-time async sync after a successful handler run). - Expanded public exports (
runpod.serverless.VolumeCache,runpod.serverless.utils.VolumeCache) and added a comprehensive unit test suite.
Reviewed changes
Copilot reviewed 8 out of 8 changed files in this pull request and generated 3 comments.
Show a summary per file
| File | Description |
|---|---|
| tests/test_serverless/test_utils/test_rp_volume_cache.py | New unit test suite covering VolumeCache behavior, safety, retention, and worker/job wiring. |
| tests/test_serverless/test_utils_init.py | Updates runpod.serverless.utils.__all__ expectations to include VolumeCache. |
| tests/test_serverless/test_init.py | Updates runpod.serverless.__all__ expectations to include VolumeCache. |
| runpod/serverless/worker.py | Hydrates and registers the default cache before starting job scaling. |
| runpod/serverless/utils/rp_volume_cache.py | New implementation of VolumeCache plus built-in env-driven wiring helpers. |
| runpod/serverless/utils/init.py | Re-exports VolumeCache from runpod.serverless.utils. |
| runpod/serverless/modules/rp_job.py | Dispatches a one-time background cache sync after non-streaming jobs. |
| runpod/serverless/init.py | Re-exports VolumeCache from runpod.serverless. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
…e, namespace validation, lint
|
Addressed review feedback in ede2d5c: Copilot (behavioral):
CodeQL / code-quality:
Tests added for streaming sync, retention resilience, and namespace validation. Full suite green (566 tests, 94% coverage). Still open for a maintainer decision before merge: default-on vs opt-in for the serverless built-in when a volume is mounted (currently on, disabled via |
| 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) | ||
|
|
| 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 |
Closes SLS-367.
Summary
Adds a directory-agnostic, bidirectional network-volume warm cache to the serverless SDK, and auto-wires it into the worker loop to cache model weights across worker recycling. This generalizes the mechanism Flash implements in flash-worker (
CacheSyncManager) into a reusable runpod-python primitive that any consumer can point at its own directories.What's in it
VolumeCache(runpod/serverless/utils/rp_volume_cache.py, public asrunpod.serverless.VolumeCache) — stdlib-only, best-effort (never raises into the worker loop):hydrate()— extract per-endpoint tar shards from a mounted volume into configured dirs on cold start; idempotent via a container-local marker.sync()— pack the post-baseline delta into a new worker-owned shard ({worker_id}-{seq}.tar) with an atomic publish. Per-worker shards are collision-free under concurrent workers (no shared read-modify-write).warm()context manager — hydrate on enter, sync on exit; places hydration before the model download by construction..tmpsweep.Serverless built-in:
run_workerhydrates + registers a default cache before jobs start;run_jobfires a guarded, threaded sync once after the first successful job. Gated on a mounted volume andRUNPOD_VOLUME_CACHE. Model dirs discovered fromHF_HOME/HF_HUB_CACHE/TORCH_HOME, extensible viaRUNPOD_CACHE_DIRS; cap viaRUNPOD_VOLUME_CACHE_MAX_GB.Notes
>=3.10floor (tar data filter applied only when the runtime provides it).Test plan
tests/test_serverless/test_utils/test_rp_volume_cache.py(34 tests): availability gating, delta detection + exclusions, coarse-mtime tolerance, multi-shard last-writer-wins, cross-worker distinct-shard hydration, retention pruning, extract-safety (traversal / symlink / hardlink / sibling-prefix), best-effort swallow vs re-raise,warm()ordering incl. on-exception, worker-loop wiring + env gating,build_default_cachedisabled / no-volume / bad-env / available branches.make quality-checkpasses (557 tests, 94% coverage, above the 90% gate).Follow-ups (separate)
CacheSyncManagerto this primitive (dirs=["/root/.cache"]), then delete Flash's bespoke tar/find logic.