Configurable fsspec reads and coalescing block cache

datalad_fuse/caching.py

@@ -0,0 +1,94 @@
+"""fsspec caches suited to the random access patterns of, e.g., HDF5 readers
+
+.. note::
+   `CoalescingBlockCache` only differs from fsspec's ``"blockcache"`` in
+   fetching runs of consecutive missing blocks with one request.  Once
+   fsspec's ``BlockCache`` does so too (proposed in
+   https://github.com/CodyCBakerPhD/filesystem_spec/pull/1), this module can be
+   dropped in favor of ``cache_type="blockcache"``.
+"""
+
+from __future__ import annotations
+
+from collections import OrderedDict
+from typing import Callable, Optional
+
+from fsspec.caching import BaseCache, register_cache
+
+#: Name under which `CoalescingBlockCache` is registered with fsspec, to be
+#: passed as ``cache_type``
+COALESCING_BLOCK_CACHE = "coalescing-blockcache"
+
+
+class CoalescingBlockCache(BaseCache):
+    """Keep the least recently used blocks of a file, fetching runs of
+    consecutive missing blocks with one request each
+
+    Like fsspec's ``"blockcache"``, which suits readers jumping over many
+    small pieces scattered over a file (where the default read-ahead cache
+    would discard what it read ahead at each jump), but a read spanning many
+    missing blocks is fetched with one request per run of consecutive missing
+    blocks rather than with one request per block.
+
+    Parameters
+    ----------
+    blocksize : int
+        Size of the blocks in bytes.
+    fetcher : callable
+        ``fetcher(start, end)`` returns the bytes of the file in that range.
+    size : int
+        Size of the file.
+    maxblocks : int
+        Maximum number of blocks kept (more are kept while a single read
+        needs them).
+    """
+
+    name = COALESCING_BLOCK_CACHE
+
+    def __init__(
+        self,
+        blocksize: int,
+        fetcher: Callable[[int, int], bytes],
+        size: int,
+        maxblocks: int = 32,
+    ) -> None:
+        super().__init__(blocksize, fetcher, size)
+        self.maxblocks = maxblocks
+        self.blocks: OrderedDict[int, bytes] = OrderedDict()
+
+    def _fetch(self, start: Optional[int], end: Optional[int]) -> bytes:
+        if start is None:
+            start = 0
+        if end is None or end > self.size:
+            end = self.size
+        if start >= end:
+            return b""
+        first = start // self.blocksize
+        last = (end - 1) // self.blocksize
+        numbers = range(first, last + 1)
+        missing = [n for n in numbers if n not in self.blocks]
+        i = 0
+        while i < len(missing):
+            j = i
+            while j + 1 < len(missing) and missing[j + 1] == missing[j] + 1:
+                j += 1
+            run_start = missing[i] * self.blocksize
+            run_end = min((missing[j] + 1) * self.blocksize, self.size)
+            data = self.fetcher(run_start, run_end)
+            for n in range(missing[i], missing[j] + 1):
+                offset = (n - missing[i]) * self.blocksize
+                self.blocks[n] = data[offset : offset + self.blocksize]
+            i = j + 1
+        out = []
+        for n in numbers:
+            self.blocks.move_to_end(n)
+            out.append(self.blocks[n])
+        # Evict the least recently used blocks, but not those just read
+        while len(self.blocks) > max(self.maxblocks, len(numbers)):
+            self.blocks.popitem(last=False)
+        data = b"".join(out)
+        offset = first * self.blocksize
+        return data[start - offset : end - offset]
+
+
+register_cache(CoalescingBlockCache, clobber=True)

datalad_fuse/fsspec.py

@@ -27,6 +27,7 @@
 from fsspec.implementations.http import HTTPFileSystem
 import methodtools
 
+from . import caching as _caching  # noqa: F401  (registers its cache)
 from .consts import CACHE_SIZE
 from .utils import AnnexKey, is_annex_dir_or_key
 
@@ -68,6 +69,15 @@ class DatasetAdapter:
         If true, paths of key files under ``.git/annex/objects/`` (the targets
         of annexed symlinks) are opened as annexed content, fetched from a
         remote URL if not present locally.
+    fsspec_options : dict, optional
+        Options for fsspec's ``HTTPFileSystem`` reading remote URLs:
+        ``block_size`` (the size in bytes of its requests), ``cache_type``
+        (the name of the fsspec cache keeping what was read) and
+        ``cache_options`` (e.g. ``{"maxblocks": 512}``); fsspec's defaults if
+        not given.  Readers jumping over many small pieces scattered over a
+        file (e.g., HDF5) read much less with ``"coalescing-blockcache"`` (see
+        :mod:`datalad_fuse.caching`) and a small ``block_size`` than with
+        fsspec's default read-ahead cache.
 
     Notes
     -----
@@ -76,7 +86,11 @@ class DatasetAdapter:
     """
 
     def __init__(
-        self, path: str | Path, caching: bool, mode_transparent: bool = False
+        self,
+        path: str | Path,
+        caching: bool,
+        mode_transparent: bool = False,
+        fsspec_options: Optional[dict[str, Any]] = None,
     ) -> None:
         self.path = Path(path)
         self.mode_transparent = mode_transparent
@@ -90,7 +104,7 @@ def __init__(
             ds.repo.get_commit_date(), tz=timezone.utc
         )
         self.caching = caching
-        fs = HTTPFileSystem(get_client=get_client)
+        fs = HTTPFileSystem(get_client=get_client, **(fsspec_options or {}))
         if self.caching:
             self.fs = CachingFileSystem(
                 fs=fs,
@@ -596,6 +610,8 @@ class FsspecAdapter:
         Passed to each :class:`DatasetAdapter`.
     mode_transparent : bool
         Passed to each :class:`DatasetAdapter`.
+    **adapter_options
+        Further options passed to each :class:`DatasetAdapter`.
 
     Notes
     -----
@@ -604,11 +620,16 @@ class FsspecAdapter:
     """
 
     def __init__(
-        self, root: str | Path, caching: bool, mode_transparent: bool = False
+        self,
+        root: str | Path,
+        caching: bool,
+        mode_transparent: bool = False,
+        **adapter_options: Any,
     ) -> None:
         self.root = Path(root)
         self.mode_transparent = mode_transparent
         self.caching = caching
+        self.adapter_options = adapter_options
         self.datasets: dict[Path, DatasetAdapter] = {}
 
     def __enter__(self) -> FsspecAdapter:
@@ -658,6 +679,7 @@ def resolve_dataset(self, filepath: str | Path) -> tuple[DatasetAdapter, str]:
                 dspath,
                 mode_transparent=self.mode_transparent,
                 caching=self.caching,
+                **self.adapter_options,
             )
         relpath = str(Path(filepath).relative_to(dspath))
         return dsap, relpath

datalad_fuse/tests/test_caching.py

@@ -0,0 +1,61 @@
+from __future__ import annotations
+
+from pathlib import Path
+
+from datalad.api import Dataset
+from fsspec.caching import caches
+import pytest
+
+from datalad_fuse.caching import COALESCING_BLOCK_CACHE, CoalescingBlockCache
+from datalad_fuse.fsspec import FsspecAdapter
+
+
+def test_coalescing_block_cache() -> None:
+    assert caches[COALESCING_BLOCK_CACHE] is CoalescingBlockCache
+    content = bytes(range(256)) * 40
+    requests: list[tuple[int, int]] = []
+
+    def fetcher(start: int, end: int) -> bytes:
+        requests.append((start, end))
+        return content[start:end]
+
+    c = CoalescingBlockCache(100, fetcher, len(content), maxblocks=5)
+    # The blocks missing for a read are fetched with one request
+    assert c._fetch(150, 420) == content[150:420]
+    assert requests == [(100, 500)]
+    # Cached blocks are not fetched again, and runs of missing blocks around
+    # them are fetched with one request each
+    assert c._fetch(30, 720) == content[30:720]
+    assert requests[1:] == [(0, 100), (500, 800)]
+    # Reads at the end of the content, and beyond
+    assert c._fetch(10_190, 20_000) == content[10_190:]
+    assert requests[3:] == [(10_100, 10_240)]
+    assert c._fetch(20_000, 20_100) == b""
+    # Only the 5 most recently used blocks are kept (but at least the blocks
+    # of the last read)
+    assert sorted(c.blocks) == [5, 6, 7, 101, 102]
+    requests.clear()
+    assert c._fetch(0, 10_240) == content
+    assert requests == [(0, 500), (800, 10_100)]
+    assert len(c.blocks) == 103
+
+
+@pytest.mark.usefixtures("tmp_home")
+def test_adapter_fsspec_options(tmp_path: Path) -> None:
+    ds = Dataset(tmp_path / "ds").create()
+    (ds.pathobj / "data.bin").write_bytes(b"0123456789" * 1000)
+    ds.save(message="Add data")
+    with FsspecAdapter(
+        ds.pathobj,
+        caching=False,
+        fsspec_options={
+            "block_size": 1024,
+            "cache_type": COALESCING_BLOCK_CACHE,
+            "cache_options": {"maxblocks": 7},
+        },
+    ) as adapter:
+        dsap, relpath = adapter.resolve_dataset(ds.pathobj / "data.bin")
+        assert relpath == "data.bin"
+        assert dsap.fs.block_size == 1024
+        assert dsap.fs.cache_type == COALESCING_BLOCK_CACHE
+        assert dsap.fs.cache_options == {"maxblocks": 7}

docs/source/api.rst

@@ -35,6 +35,12 @@ Opening files: ``datalad_fuse.fsspec``
    :undoc-members:
    :show-inheritance:
 
+.. currentmodule:: datalad_fuse.caching
+
+.. autoclass:: CoalescingBlockCache
+
+.. autodata:: COALESCING_BLOCK_CACHE
+
 
 git-annex helpers: ``datalad_fuse.utils``
 =========================================