vesuvius.predict: --chunk_cache_mb exposes the in-memory chunk cache (#1325)

[CVasilopoulos]

In one sentence: vesuvius.predict --chunk_cache_mb 512 turns on the in-memory LRU chunk cache that #1373 added to Volume, so overlapping patches reuse chunks already downloaded instead of fetching them again.

One real example: Starting with the published Scroll 1 volume https://dl.ash2txt.org/full-scrolls/Scroll1/PHercParis4.volpkg/volumes_zarr_standardized/54keV_7.91um_Scroll1A.zarr, I ran vesuvius.predict on a 192³ region with 192³ patches and 4 DataLoader workers. Without the flag it made 1000 chunk requests (796 MB). With --chunk_cache_mb 512 it made 250 (199 MB). With --chunk_cache_mb 512 --num_workers 0 it made 125, one per distinct chunk. The logits were byte-identical in all three runs.

Before: #1373 made open_zarr(cache=True) and Volume(cache=True, cache_size_mb=...) available, but nothing in the inference path passes them. VCDataset has no cache argument, and vesuvius.predict has no flag. Patches overlap by 50% by default and are not aligned to chunks, so every patch fetches all the chunks it touches again. That is the amplification @aistae measured in #1325 and @TAUIL-Abd-Elilah reproduced by counting HTTP requests. On main, --chunk_cache_mb is error: unrecognized arguments.

After this PR: --chunk_cache_mb N passes cache=True, cache_size_mb=N from Inferer through VCDataset to Volume. Each DataLoader worker keeps its own cache of up to N MB of stored chunk bytes, evicting least-recently-used chunks. The default is 0, which leaves predict exactly as it is today: the same requests and the same output. Negative values are rejected. Like #1373, the cache needs zarr 3. With zarr 2, open_zarr raises its existing NotImplementedError.

Proof: Upstream main 757f70c01 compared with this branch on the same base, same volume, same --bbox 6400:6592,5120:5312,5120:5312 (level 0, inside the scroll), overlap 0.5, batch 1, --disable_tta, CPU. Requests are counted, not timed. A sitecustomize.py on PYTHONPATH wraps fsspec's HTTPFileSystem._cat_file in every process, including the DataLoader workers, and appends one line per request. chunk GETs counts successful requests for 0/z/y/x keys. logits sha256 hashes the written logits_part_0.zarr and coordinates_part_0.zarr with patches sorted by coordinate.

Patch 192³ (the patch size of hf://scrollprize/surface_recto), 64 patches over 125 distinct chunks:

run	chunk GETs	MB downloaded	GETs per distinct chunk	logits sha256
main, default (4 workers)	1000	796.2	8.00	69eaa10571db94f6
this PR, same command	1000	796.2	8.00	69eaa10571db94f6
this PR, --chunk_cache_mb 512 (4 workers)	250	199.0	2.00	69eaa10571db94f6
this PR, --chunk_cache_mb 512 --num_workers 0	125	99.4	1.00	69eaa10571db94f6
main, --chunk_cache_mb 512	error: unrecognized arguments: --chunk_cache_mb 512, exit 2

Patch 128³ (the configuration in #1325), 125 patches over 64 distinct chunks:

run	chunk GETs	MB downloaded	GETs per distinct chunk	logits sha256
main, default (4 workers)	1000	801.6	15.62	1f2f92447af1dfc4
this PR, same command	1000	801.6	15.62	1f2f92447af1dfc4
this PR, --chunk_cache_mb 512 (4 workers)	241	195.1	3.77	1f2f92447af1dfc4
this PR, --chunk_cache_mb 512 --num_workers 0	64	52.0	1.00	1f2f92447af1dfc4

Excerpt of the terminal output (192³, main, then this PR with the flag):

$ vesuvius.predict --model_path tiny_train_py.pth --input_dir https://dl.ash2txt.org/full-scrolls/Scroll1/PHercParis4.volpkg/volumes_zarr_standardized/54keV_7.91um_Scroll1A.zarr --output_dir out-p192-before-default --device cpu --disable_tta --patch_size 192,192,192 --bbox 6400:6592,5120:5312,5120:5312
Output shape: (64, 1, 192, 192, 192)
EXIT=0
chunk GETs: 1000 | distinct chunks: 125 | amplification: 8.00x | chunk MB downloaded: 796.2 | processes that fetched chunks: 4 | metadata GETs: 21 | failed requests (404 etc.): 40

$ vesuvius.predict ... --patch_size 192,192,192 --bbox 6400:6592,5120:5312,5120:5312 --chunk_cache_mb 512
Output shape: (64, 1, 192, 192, 192)
EXIT=0
chunk GETs: 250 | distinct chunks: 125 | amplification: 2.00x | chunk MB downloaded: 199.0 | processes that fetched chunks: 4 | metadata GETs: 8 | failed requests (404 etc.): 40

All eight counts match a simulation of the sliding-window grid over the 128³ chunk grid, with PyTorch's round-robin dispatch of batches to workers and one LRU per worker. Main and this PR print the same counts for the same command.

Tests: the new tests/models/run/test_inference_chunk_cache.py has 6 passed with this PR and 6 failed against main's source. tests/models/run + tests/data/test_chunk_cache.py + tests/data/test_vc_dataset_bbox.py go from 73 passed, 1 skipped on main to 79 passed, 1 skipped.

Why / where this is useful:

I picked this up because #1373 had already put a chunk cache into Volume, but vesuvius.predict, the tool people actually stream scrolls with, still had no way to turn it on. I ran the same Scroll 1 region on main and on this branch, counted the chunk requests in every worker, and checked that the logits come out byte-identical.

Anyone running vesuvius.predict against a remote volume, for a --bbox region or a whole scroll, pays for the same chunks several times over. As @TAUIL-Abd-Elilah noted in #1325, their link caps near 11 MB/s, so request count decides whether a run fits in a night. With this flag the requests for one region drop about 4 times with the default 4 workers. With --num_workers 0 they drop to one per chunk, 8 to 16 times fewer here. It uses cache code that is already merged and tested. It needs no local disk and cannot serve stale data across runs.

• I personally verified that the example and proof above were produced by this PR on the stated data.

Details

What changed

• vesuvius/src/vesuvius/models/run/inference.py:
  • a --chunk_cache_mb parser flag (int, default 0; --chunk-cache-mb also works);
  • an Inferer(chunk_cache_mb=0) argument with a >= 0 check;
  • _create_dataset_and_loader passes cache=chunk_cache_mb > 0, cache_size_mb=chunk_cache_mb to VCDataset.
• vesuvius/src/vesuvius/data/vc_dataset.py: VCDataset(cache=False, cache_size_mb=256) is documented and forwarded to Volume, with the same names and defaults as Volume.
• vesuvius/docs/inference.md: one row in the argument table.
• vesuvius/tests/models/run/test_inference_chunk_cache.py covers:
  • the parser default;
  • main passing the value to Inferer;
  • Inferer forwarding cache/cache_size_mb for 0 and 512;
  • the negative-value check;
  • a local VCDataset with the cache whose store is a CacheStore, survives pickle as DataLoader workers receive it, reads the same patches as an uncached dataset, and records hits within its size bound.
• vesuvius/tests/models/run/test_inference_nnunet_normalization.py: the hand-built stub Inferer gains chunk_cache_mb = 0.

Design notes

• Off by default: vesuvius: opt-in zarr 3 chunk cache via zarr's built-in CacheStore #1373 kept the cache opt-in, and turning it on would change memory use for every existing run. If you would rather have it on by default, the default can change, but zarr 2 installs and local inputs would then need a guard.
• One cache per worker: worker processes do not share memory, and batches go to workers round-robin. A worker only sees every fourth batch, so it can reuse only part of the overlap. That is why 4 workers give 250 requests rather than 125. --num_workers 0 reaches one request per chunk, at the cost of reading in the main process.
• Memory: up to N × max(1, num_workers) MB of stored chunk bytes. For this compressed volume a chunk is about 0.8 MB, so 512 MB holds about 640 chunks per worker.
• Traversal order: feat(data): visit patch positions in chunk-local order — makes any chunk cache far more effective #1331 (@aistae, Morton ordering, closed by the time-limit bot) would raise the hit rate further. I kept this PR to the flag only.

Tested

• Linux x86_64, CPU only, in python:3.14-slim with Python 3.14.7, zarr 3.3.0, fsspec 2026.7.0 and torch 2.14.0+cpu. Each run was a container capped at 12 GB and 8 CPUs. Not tested on GPU, macOS or Windows.
• Volume: https://dl.ash2txt.org/full-scrolls/Scroll1/PHercParis4.volpkg/volumes_zarr_standardized/54keV_7.91um_Scroll1A.zarr. Level 0, shape 14376×7888×8096, chunks 128³, blosc/zstd, about 0.8 MB per chunk.
• Command: vesuvius.predict --model_path tiny_train_py.pth --input_dir <volume> --output_dir <out> --device cpu --disable_tta [--patch_size 192,192,192] --bbox 6400:6592,5120:5312,5120:5312 [--chunk_cache_mb 512] [--num_workers 0]
• Model: a randomly initialised 2-stage train_py UNet (4,837 parameters, seed 0), built with NetworkFromConfig as in test_inference_legacy_checkpoint.py. It only drives the read path, which does not depend on the model. Its logits are compared byte for byte and not interpreted. I did not use surface_recto (819 MB), because 125 patches of it on CPU per run was not practical here.

Limitations

• The cache lasts for one run. Repeated runs over the same region fetch again. vesuvius: persistent on-disk chunk cache for remote volumes #1543 (@blackemberlabs) proposed a persistent on-disk cache with --chunk_cache_dir, and the two can coexist.
• CacheStore does not cache misses. Absent chunks in sparse or masked volumes are still requested each time. In these runs, the 40 failed requests in every row are zarr v2 metadata probes (0/.zgroup, 0/.zattrs, zarr.json, per-level .zattrs/.zgroup), not chunks.
• Wall-clock time is not reported as a result. It is dominated by the link: two identical uncached runs took 168 s and 124 s.

Prior work and credit

• @aistae reported vesuvius.predict has no chunk cache: inference re-downloads the same chunks, ~17x more traffic than needed #1325 with measurements and proposed wrapping the store in an LRU cache.
• @TAUIL-Abd-Elilah worked out the counting method (wrap HTTPFileSystem._cat_file) and its DataLoader-worker caveat, and the simulation that the counts here match.
• @robertlangdonn wrote the cache itself in vesuvius: fix zarr 3 Volume multiscale reads #1177/vesuvius: opt-in zarr 3 chunk cache via zarr's built-in CacheStore #1373, using zarr's CacheStore at @bruniss's suggestion. This PR only exposes it in predict.
• fix(data): optional on-disk chunk cache for remote reads (refs #1325) #1327 (@aistae, on-disk simplecache, closed by its author in favour of vesuvius: fix zarr 3 Volume multiscale reads #1177) and vesuvius: persistent on-disk chunk cache for remote volumes #1543 (@blackemberlabs, persistent on-disk cache wired into predict, closed unreviewed by the time-limit bot on 2026-09-16) took the disk route. This PR copies no code from either. If vesuvius: persistent on-disk chunk cache for remote volumes #1543 is reopened, the two touch neighbouring lines in inference.py and vc_dataset.py, and I am happy to rebase.
• I found no other PR or branch that wires the vesuvius: opt-in zarr 3 chunk cache via zarr's built-in CacheStore #1373 cache into vesuvius.predict.

AI-assisted (Claude Code), human-directed.

[vercel]

@CVasilopoulos is attempting to deploy a commit to the scroll Team on Vercel.

A member of the Team first needs to authorize it.