Add blosc2.set_prefetch(): read-ahead for on-disk arrays

bench/optim_tips/tip_16_prefetch.py

@@ -0,0 +1,90 @@
+#######################################################################
+# Copyright (c) 2019-present, Blosc Development Team <[email protected]>
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+#######################################################################
+
+# Tip 16: blosc2.set_prefetch() lets a few background threads read the chunks a
+# slice needs, in large parallel requests and ahead of decompression.  Without
+# it, a slice of an on-disk array reads its blocks one after the other, and
+# every read that misses the page cache waits for the storage.
+#
+# The benefit depends on the storage latency: on a local NVMe disk cold patches
+# get ~1.5-3x faster, on NFS ~3-5x.  Every read here is cold: the file is
+# dropped from the page cache before each one (Linux; elsewhere the reads are
+# warm, and both variants measure the same).
+
+import os
+from pathlib import Path
+
+import numpy as np
+from common import fmt_bytes, measure, save_plot
+
+import blosc2
+
+SHAPE, CHUNKS, BLOCKS, PATCH = (4, 256, 256, 256), (1, 64, 128, 128), (1, 16, 32, 128), (128, 128, 128)
+URLPATH = str(Path(__file__).parent / "tip_16.b2nd")
+N_READS = 40
+
+
+def _build():
+    rng = np.random.default_rng(0)
+    data = blosc2.zeros(SHAPE, dtype=np.float32, chunks=CHUNKS, blocks=BLOCKS, urlpath=URLPATH, mode="w")
+    for c in range(SHAPE[0]):
+        data[c] = rng.normal(size=SHAPE[1:]).astype(np.float32).round(2)
+
+
+_build()
+
+_starts = np.random.default_rng(1).integers(
+    0, [s - p + 1 for s, p in zip(SHAPE[1:], PATCH, strict=True)], size=(N_READS, 3)
+)
+
+
+def _drop_from_page_cache():
+    if hasattr(os, "posix_fadvise"):
+        fd = os.open(URLPATH, os.O_RDONLY)
+        os.posix_fadvise(fd, 0, 0, os.POSIX_FADV_DONTNEED)
+        os.close(fd)
+
+
+def _cold_patch_reads():
+    total = 0.0
+    for start in _starts:
+        _drop_from_page_cache()
+        arr = blosc2.open(URLPATH, mmap_mode="r")
+        total += arr[
+            (slice(None),) + tuple(slice(s, s + p) for s, p in zip(start, PATCH, strict=True))
+        ].sum()
+        del arr  # unmapped, so that the next drop reaches every page
+    return total
+
+
+def naive():
+    return _cold_patch_reads()
+
+
+def tip():
+    blosc2.set_prefetch(4)  # per process; BLOSC_PREFETCH=4 in the environment does the same
+    return _cold_patch_reads()
+
+
+if __name__ == "__main__":
+    naive_t, naive_m = measure(__file__, "naive")
+    tip_t, tip_m = measure(__file__, "tip")
+
+    print(f"naive  without read-ahead   : {naive_t:.3f}s  peak {fmt_bytes(naive_m)}")
+    print(f"tip    set_prefetch(4)      : {tip_t:.3f}s  peak {fmt_bytes(tip_m)}")
+    print(f"speedup: {naive_t / tip_t:.1f}x")
+
+    save_plot(
+        "tip_16_prefetch.png",
+        f"set_prefetch(4) vs no read-ahead — {N_READS} cold patch reads",
+        "no read-ahead",
+        "set_prefetch(4)",
+        naive_t,
+        tip_t,
+        naive_m,
+        tip_m,
+    )

doc/guides/optimization_tips.md

@@ -514,6 +514,28 @@ With 8 concurrent readers doing random slice reads over a 269 MB array, mmap was
 
 *Benchmark for this tip: [`tip_10_mmap_many_readers.py`](https://github.com/Blosc/python-blosc2/blob/main/bench/optim_tips/tip_10_mmap_many_readers.py)*
 
+## Reading cold data from disk? Turn on read-ahead
+
+A slice of an on-disk array reads the blocks it needs one after the other, and every read that misses the page cache waits for the storage. On a network file system (NFS, SMB, ...) each of those waits is a round trip to the server. {func}`blosc2.set_prefetch` starts a few background threads that fetch the chunks a slice needs in large, parallel requests, while the slice decompresses the chunks that have already arrived. It works for `a[...]`, {meth}`NDArray.slice <blosc2.NDArray.slice>`, orthogonal selections and {class}`SChunk <blosc2.SChunk>` slices:
+
+```python
+# Avoid (data not in the page cache yet): one read per block, one after the other
+arr = blosc2.open(path, mmap_mode="r")
+patch = arr[:, 10:138, 20:148, 30:158]
+
+# Prefer: let a few threads fetch the chunks ahead of decompression (once per process;
+# BLOSC_PREFETCH=4 in the environment does the same)
+blosc2.set_prefetch(4)
+arr = blosc2.open(path, mmap_mode="r")
+patch = arr[:, 10:138, 20:148, 30:158]
+```
+
+![set_prefetch vs no read-ahead](optim_tips/tip_16_prefetch.png)
+
+On a local NVMe disk, 40 cold patch reads were **~1.7x faster**; on NFS, cold patches of 3D images get 3-5x faster. Chunks already in the page cache are not read again, so warm reads cost about the same as without read-ahead. The threads mostly wait for the storage, but the kernel spends CPU on every byte they read, more so the more threads read one file at once: 4 threads with the default 1 MiB pieces give most of the speed on network and local storage alike; on local disks, `set_prefetch(2, piece_size=8 << 20)` costs less CPU for some of the speed. Of a chunk that `a[...]` or `slice()` covers only in part, only the blocks the slice needs are read. With many reader processes at once, 1-2 threads each are enough; when the CPU is the limit (tens of processes reading cold data on one node), read-ahead can be slower than none. The read-ahead threads are separate from the decompression threads (`nthreads`), so a program with many reader processes can keep `nthreads=1` and still use them.
+
+*Benchmark for this tip: [`tip_16_prefetch.py`](https://github.com/Blosc/python-blosc2/blob/main/bench/optim_tips/tip_16_prefetch.py)*
+
 ## Use `.b2z` to group related data into one memory-mapped file
 
 `.b2z` packs related data — a {class}`~blosc2.CTable`'s columns and indexes, hierarchies of arrays, or both — into a single-file container that blosc2 reads *in place*: opened read-only, nothing is ever unpacked, and with `mmap_mode="r"` every member is memory-mapped at its offset inside the container. Beyond the convenience of one file to copy or archive, this buys you:

doc/reference/low_level.rst

@@ -43,6 +43,7 @@ Set / get compression params
     register_filter
     set_blocksize
     set_nthreads
+    set_prefetch
     set_releasegil
     set_compressor
     get_compressor

src/blosc2/__init__.py

@@ -491,6 +491,7 @@ def _raise(exc):
     set_blocksize,
     set_compressor,
     set_nthreads,
+    set_prefetch,
     set_releasegil,
     unpack,
     unpack_array,
@@ -1116,6 +1117,7 @@ def _raise(exc):
     "set_blocksize",
     "set_compressor",
     "set_nthreads",
+    "set_prefetch",
     "set_releasegil",
     "sign",
     "signbit",