docs(examples): VolumeCache warm-cache example + transport benchmarks

05_data_workflows/05_volume_warm_cache/README.md

@@ -0,0 +1,75 @@
+# 05_volume_warm_cache
+
+Warm-cache a model on a Runpod network volume with `VolumeCache`, so cold starts
+restore the model from the volume instead of re-downloading it.
+
+## Status
+
+`VolumeCache` is a `runpod-python` serverless primitive added in
+**SLS-367 / [runpod-python PR #538](https://github.com/runpod/runpod-python/pull/538)**,
+which is merged and ships in `runpod` **1.12.0**.
+
+One gate remains before this example runs end-to-end: the deployed
+**flash-worker** image must be built against `runpod>=1.12.0`. The model load
+runs inside the remote GPU worker, so it is that image's `runpod`, not this
+example's local env, that must have `VolumeCache`. Until flash-worker ships it,
+this example is illustrative.
+
+## Overview
+
+`VolumeCache` keeps a browsable mirror of local cache directories on a mounted
+network volume and reconciles them on each use:
+
+- **On enter** (`hydrate`): copy files that are missing or newer on the volume
+  into the container.
+- **On exit** (`sync`): copy files that are missing or newer in the container
+  back to the volume (in the background).
+
+The GPU worker wraps its Stable Diffusion load in `with VolumeCache(...)`. The
+first cold worker downloads the weights and syncs them to the volume; every later
+cold worker hydrates from the volume and skips the download.
+
+## How this differs from `01_network_volumes`
+
+`01_network_volumes` sets `HF_HUB_CACHE` to a path **on** the volume, so every
+model read hits the network mount. This example keeps the Hugging Face cache on
+**local disk** and uses `VolumeCache` to mirror it to the volume — inference
+reads stay local (fast) while cold starts stay warm (persisted).
+
+| | 01_network_volumes | 02_volume_warm_cache |
+| --- | --- | --- |
+| Model cache location | On the network volume (`HF_HUB_CACHE` → `/runpod-volume/...`) | Local disk (default `~/.cache/huggingface`) |
+| Volume role | Live cache (read/written directly) | Warm mirror (hydrated in, synced out) |
+| Inference read speed | Network mount | Local disk |
+| Cold-start download | Skipped (reads from volume) | Skipped (restored to local disk) |
+
+## Quick Start
+
+```bash
+cd 05_data_workflows/05_volume_warm_cache
+uv pip install -r requirements.txt
+uv run flash login          # or set RUNPOD_API_KEY in .env
+uv run flash dev            # serves at http://localhost:8888
+```
+
+Generate an image (provisions a real GPU worker with the volume attached):
+
+```bash
+curl -X POST http://localhost:8888/gpu_worker/runsync \
+  -H "Content-Type: application/json" \
+  -d '{"input": {"prompt": "a sunset over mountains"}}'
+```
+
+## Benchmark: direct vs VolumeCache
+
+`benchmark/` contains an A/B that measures cold-start model-load time for the two
+caching strategies (direct HF-on-volume vs local-disk-mirrored-to-volume),
+reporting first-run vs warm load time and volume/local storage. See
+[benchmark/README.md](./benchmark/README.md). The direct arm is runnable today;
+the VolumeCache arm needs a flash-worker image built against `runpod>=1.12.0`.
+
+## What You'll Learn
+
+- How to warm-cache model weights across cold starts with `runpod.serverless.VolumeCache`
+- Why mirroring a local cache to a volume differs from mounting the cache on the volume
+- The `with VolumeCache(dirs=[...]):` closure pattern inside a Flash `@Endpoint`

05_data_workflows/05_volume_warm_cache/benchmark/.gitignore

@@ -0,0 +1,2 @@
+.flash/
+results.json

05_data_workflows/05_volume_warm_cache/benchmark/README.md

@@ -0,0 +1,60 @@
+# Benchmark: direct volume cache vs VolumeCache
+
+An A/B that measures **cold-start model-load time** for two ways of caching a
+model on a Runpod network volume:
+
+- **direct** (`bench_direct.py`) — `HF_HOME`/`HF_HUB_CACHE` point at a directory
+  **on** the volume; the model is read from the network mount on every cold start.
+- **volumecache** (`bench_volumecache.py`) — the HF cache stays on **local disk**
+  and `runpod.serverless.VolumeCache` hydrates it from the volume before load and
+  syncs new files back after; reads are local, the volume is a warm mirror.
+
+Both load the same model (`stable-diffusion-v1-5`) and self-report metrics;
+`benchmark.py` drives them and prints a comparison.
+
+## What it measures
+
+| Metric | Meaning |
+| --- | --- |
+| first-run load | Cold start with an empty cache — includes the download (and, for volumecache, the sync back). One-time per endpoint. |
+| warm load (mean/min) | Cold start with the cache already populated — the number that matters at scale. For volumecache this **includes** the hydrate copy. |
+| hydrate | Time volumecache spent copying volume → local (0 for direct). |
+| volume / local | Bytes stored on the volume and on local disk (volumecache double-stores). |
+
+## Prerequisites and status
+
+- Both endpoints need a network volume attached (the workers declare one).
+- **Direct arm: runnable on Flash today.**
+- **VolumeCache arm requires a flash-worker image built against `runpod>=1.12.0`,
+  which is where `VolumeCache` ships (SLS-367 /
+  [runpod-python PR #538](https://github.com/runpod/runpod-python/pull/538)).**
+  It cannot be added at runtime via `dependencies` because the worker's `runpod`
+  is already imported before the handler runs. Until flash-worker ships it, run
+  the direct arm alone with `--skip-volumecache`.
+
+This harness has **not been run yet** — no results are published here. Numbers
+should be filled in from a real run.
+
+## Running
+
+```bash
+# From this benchmark/ directory (so flash dev serves only the two bench workers):
+uv run flash dev            # provisions real GPU endpoints; serves at :8888
+
+# In another shell — direct arm only (runnable today):
+python benchmark.py --skip-volumecache --warm-trials 3
+
+# Full A/B (once the flash-worker image includes VolumeCache):
+python benchmark.py --warm-trials 3
+```
+
+Confirm the exact routes at `http://localhost:8888/docs` and pass
+`--direct-route` / `--vc-route` if they differ. The scale-to-0 wait must exceed
+each arm's worker `idle_timeout`, or the "warm" trials reuse a still-alive worker
+and report warm reuse as a cold start. The two arms differ (direct `idle_timeout=30`,
+volumecache `idle_timeout=180`), so `--cold-wait` is derived per arm by default
+(direct 45s, volumecache 195s); pass `--cold-wait` to override both.
+
+> Running provisions real GPU workers and a network volume — it costs money and
+> cold starts are slow. Tear down with Ctrl+C on `flash dev` (or `flash undeploy`)
+> and verify the endpoints are gone afterward.

05_data_workflows/05_volume_warm_cache/benchmark/bench_direct.py

@@ -0,0 +1,82 @@
+# Benchmark worker A: DIRECT caching.
+#
+# HF_HOME / HF_HUB_CACHE point at a directory ON the network volume, so the model
+# is read directly from the (network) mount on every cold start. This is the
+# 01_network_volumes approach. Paired with bench_volumecache.py; driven by
+# benchmark.py.
+#
+# Runnable on Flash today (no VolumeCache dependency).
+import logging
+
+from runpod_flash import Endpoint, GpuType, NetworkVolume
+
+logger = logging.getLogger(__name__)
+
+# Shared volume so both strategies compete on the same storage medium.
+volume = NetworkVolume(name="flash-05-bench-volume", size=50)
+
+HF_ON_VOLUME = "/runpod-volume/hf"
+
+
+@Endpoint(
+    name="bench_direct",
+    gpu=GpuType.NVIDIA_GEFORCE_RTX_5090,
+    workers=(0, 1),  # single worker so each call is one clean measurement
+    idle_timeout=30,  # scale to 0 between trials -> next call is a cold start
+    volume=volume,
+    env={"HF_HOME": HF_ON_VOLUME, "HF_HUB_CACHE": HF_ON_VOLUME},
+    dependencies=["torch", "diffusers", "transformers", "accelerate"],
+)
+class BenchDirect:
+    def __init__(self):
+        # Imports/paths live inside the body: @Endpoint ships only the function
+        # body to the worker.
+        import logging
+        import os
+        import time
+
+        import torch
+        from diffusers import StableDiffusionPipeline
+
+        self.logger = logging.getLogger(__name__)
+
+        hf = os.environ.get("HF_HOME", "/runpod-volume/hf")
+        # First run = the on-volume cache was empty before this cold start.
+        self._first_run = not (os.path.isdir(hf) and any(os.scandir(hf)))
+
+        t0 = time.perf_counter()
+        self.pipe = StableDiffusionPipeline.from_pretrained(
+            "runwayml/stable-diffusion-v1-5",
+            torch_dtype=torch.float16,
+            safety_checker=None,
+            requires_safety_checker=False,
+            use_safetensors=True,
+            low_cpu_mem_usage=True,
+        ).to("cuda")
+        self._load_seconds = time.perf_counter() - t0
+        self.logger.info(
+            f"[direct] load={self._load_seconds:.2f}s first_run={self._first_run}"
+        )
+
+    async def benchmark(self) -> dict:
+        """Report the cold-start metrics captured during __init__."""
+        import os
+
+        def dir_bytes(path):
+            total = 0
+            for root, _dirs, files in os.walk(path):
+                for name in files:
+                    try:
+                        total += os.path.getsize(os.path.join(root, name))
+                    except OSError:
+                        pass
+            return total
+
+        return {
+            "mode": "direct",
+            "first_run": self._first_run,
+            "load_seconds": round(self._load_seconds, 2),
+            "hydrate_seconds": 0.0,
+            "volume_bytes": dir_bytes("/runpod-volume/hf"),
+            "local_bytes": 0,
+        }

05_data_workflows/05_volume_warm_cache/benchmark/bench_volumecache.py

@@ -0,0 +1,105 @@
+# Benchmark worker B: VOLUMECACHE strategy.
+#
+# HF cache stays on local disk; runpod.serverless.VolumeCache hydrates it from
+# the network volume before the load and syncs new files back after. Reads are
+# local; the volume is a warm mirror. Paired with bench_direct.py; driven by
+# benchmark.py.
+#
+# PREREQUISITE: the deployed flash-worker image must be built against
+# runpod>=1.12.0, where VolumeCache ships (SLS-367 / runpod-python PR #538).
+# It cannot be added at runtime via the
+# dependencies list because the worker's runpod is already imported before the
+# handler runs. Until flash-worker ships it, this arm will fail to import
+# VolumeCache — see benchmark/README.md.
+import logging
+
+from runpod_flash import Endpoint, GpuType, NetworkVolume
+
+logger = logging.getLogger(__name__)
+
+# Same volume as bench_direct so both strategies compete on the same medium.
+volume = NetworkVolume(name="flash-05-bench-volume", size=50)
+
+
+@Endpoint(
+    name="bench_volumecache",
+    gpu=GpuType.NVIDIA_GEFORCE_RTX_5090,
+    workers=(0, 1),
+    # Keep the worker alive long enough for the background mirror-back sync to
+    # finish before scale-to-0 (otherwise the volume mirror never populates and
+    # the next cold worker can't hydrate).
+    idle_timeout=180,
+    volume=volume,
+    dependencies=["torch", "diffusers", "transformers", "accelerate"],
+)
+class BenchVolumeCache:
+    def __init__(self):
+        import logging
+        import os
+        import time
+
+        import torch
+        from diffusers import StableDiffusionPipeline
+        from runpod.serverless import VolumeCache
+
+        self.logger = logging.getLogger(__name__)
+
+        # HF cache on local disk. The flash-worker image sets HF_HOME=/hf-cache,
+        # so honor it (falling back to the HF default) — VolumeCache mirrors that
+        # local dir to the volume.
+        hf = os.environ.get("HF_HOME") or os.path.expanduser("~/.cache/huggingface")
+        self._hf = hf
+        self._vc = VolumeCache(dirs=[hf])
+
+        # Mirror location is {volume}/.cache/{RUNPOD_ENDPOINT_ID}; empty = first run.
+        self._mirror = os.path.join(
+            "/runpod-volume", ".cache", os.environ.get("RUNPOD_ENDPOINT_ID", "")
+        )
+        self._first_run = not (
+            os.path.isdir(self._mirror) and any(os.scandir(self._mirror))
+        )
+
+        t0 = time.perf_counter()
+        self._files_hydrated = self._vc.hydrate()  # volume -> local
+        self._hydrate_seconds = time.perf_counter() - t0
+
+        self.pipe = StableDiffusionPipeline.from_pretrained(
+            "runwayml/stable-diffusion-v1-5",
+            torch_dtype=torch.float16,
+            safety_checker=None,
+            requires_safety_checker=False,
+            use_safetensors=True,
+            low_cpu_mem_usage=True,
+        ).to("cuda")
+        # Total cold-start-to-ready time, including the hydrate copy.
+        self._load_seconds = time.perf_counter() - t0
+
+        self._vc.sync()  # local -> volume, in the background
+        self.logger.info(
+            f"[volumecache] load={self._load_seconds:.2f}s "
+            f"hydrate={self._hydrate_seconds:.2f}s first_run={self._first_run}"
+        )
+
+    async def benchmark(self) -> dict:
+        """Report the cold-start metrics captured during __init__."""
+        import os
+
+        def dir_bytes(path):
+            total = 0
+            for root, _dirs, files in os.walk(path):
+                for name in files:
+                    try:
+                        total += os.path.getsize(os.path.join(root, name))
+                    except OSError:
+                        pass
+            return total
+
+        return {
+            "mode": "volumecache",
+            "first_run": self._first_run,
+            "files_hydrated": self._files_hydrated,
+            "load_seconds": round(self._load_seconds, 2),
+            "hydrate_seconds": round(self._hydrate_seconds, 2),
+            "local_bytes": dir_bytes(self._hf),
+            "volume_bytes": dir_bytes(self._mirror),
+        }