From d50a5376193f2e9ccfd17cc827254ed96a50b8be Mon Sep 17 00:00:00 2001 From: Ben Dichter Date: Sat, 14 Mar 2026 09:24:41 +0000 Subject: [PATCH 1/3] Add concurrent chunk fetching for external array links Route remote external array links through zarr + LindiH5ZarrStore instead of h5py + LindiRemfile. This enables concurrent HTTP range requests via LindiH5ZarrStore.getitems(), which zarr calls when reading multiple chunks. The getitems() method separates serial metadata lookup (fast, uses h5py's B-tree cache) from parallel data fetches (N concurrent HTTP requests via ThreadPoolExecutor instead of N serial ones). Local external array links still use h5py directly since there's no concurrency benefit for local I/O. Co-Authored-By: Claude Opus 4.6 (1M context) --- lindi/LindiH5ZarrStore/LindiH5ZarrStore.py | 118 ++++++++++++++++- lindi/LindiH5pyFile/LindiH5pyDataset.py | 36 +++++- tests/test_concurrent_external_link.py | 140 +++++++++++++++++++++ 3 files changed, 289 insertions(+), 5 deletions(-) create mode 100644 tests/test_concurrent_external_link.py diff --git a/lindi/LindiH5ZarrStore/LindiH5ZarrStore.py b/lindi/LindiH5ZarrStore/LindiH5ZarrStore.py index bb532c2..6527f64 100644 --- a/lindi/LindiH5ZarrStore/LindiH5ZarrStore.py +++ b/lindi/LindiH5ZarrStore/LindiH5ZarrStore.py @@ -1,9 +1,12 @@ import json import base64 +import time from typing import Tuple, Union, List, IO, Any, Dict, Callable +from concurrent.futures import ThreadPoolExecutor, as_completed import numpy as np import zarr from zarr.storage import Store, MemoryStore +import requests import h5py from tqdm import tqdm from ._util import ( @@ -150,7 +153,8 @@ def __init__( _opts: LindiH5ZarrStoreOpts, _url: Union[str, None] = None, _entities_to_close: List[Any], - _local_cache: Union[LocalCache, None] = None + _local_cache: Union[LocalCache, None] = None, + _concurrent_max_workers: int = 8 ): """ Do not call the constructor directly. Instead, use the from_file class @@ -161,6 +165,7 @@ def __init__( self._url = _url self._opts = _opts self._local_cache = _local_cache + self._concurrent_max_workers = _concurrent_max_workers self._entities_to_close = _entities_to_close + [self._h5f] # Some datasets do not correspond to traditional chunked datasets. For @@ -325,6 +330,97 @@ def __contains__(self, key): return False return True + def getitems(self, keys, *, contexts=None): + """Fetch multiple keys, with concurrent HTTP fetches for remote chunks.""" + results = {} + remote_chunks = [] # (key, byte_offset, byte_count) + + for key in keys: + parts = [p for p in key.split("/") if p] + if not parts: + continue + key_name = parts[-1] + + # Metadata keys — resolve synchronously + if key_name in ('.zattrs', '.zgroup', '.zarray'): + try: + results[key] = self[key] + except KeyError: + pass + continue + + # Chunk keys — get byte range from h5py metadata + key_parent = "/".join(parts[:-1]) + try: + byte_offset, byte_count, inline_data = self._get_chunk_file_bytes_data(key_parent, key_name) + except Exception: + continue + + if inline_data is not None: + results[key] = inline_data + continue + + # Check local cache + if self._local_cache is not None and self._url is not None: + cached = self._local_cache.get_remote_chunk(url=self._url, offset=byte_offset, size=byte_count) + if cached is not None: + results[key] = cached + continue + + if self._url is not None and (self._url.startswith('http://') or self._url.startswith('https://')): + remote_chunks.append((key, byte_offset, byte_count)) + else: + # Local file — read synchronously (byte range already known) + buf = _read_bytes(self._file, byte_offset, byte_count) + self._try_cache_put(byte_offset, byte_count, buf) + results[key] = buf + + if not remote_chunks: + return self._apply_padding_to_results(results) + + # Pre-resolve URL for DANDI auth + from ..LindiRemfile.LindiRemfile import _resolve_url + resolved_url = _resolve_url(self._url) + + # Single chunk — skip thread pool overhead + if len(remote_chunks) == 1: + key, offset, count = remote_chunks[0] + val = _fetch_bytes_direct(resolved_url, offset, count) + self._try_cache_put(offset, count, val) + results[key] = val + return self._apply_padding_to_results(results) + + # Concurrent fetch + max_workers = min(len(remote_chunks), self._concurrent_max_workers) + with ThreadPoolExecutor(max_workers=max_workers) as pool: + futures = { + pool.submit(_fetch_bytes_direct, resolved_url, offset, count): (key, offset, count) + for key, offset, count in remote_chunks + } + for future in as_completed(futures): + key, offset, count = futures[future] + val = future.result() + self._try_cache_put(offset, count, val) + results[key] = val + + return self._apply_padding_to_results(results) + + def _try_cache_put(self, byte_offset, byte_count, data): + """Write data to the local cache if available.""" + if self._local_cache is not None and self._url is not None: + try: + self._local_cache.put_remote_chunk(url=self._url, offset=byte_offset, size=byte_count, data=data) + except ChunkTooLargeError: + pass + + def _apply_padding_to_results(self, results): + for key in list(results.keys()): + val = results[key] + padded_size = _get_padded_size(self, key, val) + if padded_size is not None: + results[key] = _pad_chunk(val, padded_size) + return results + def __delitem__(self, key): raise Exception("Deleting items is not allowed") @@ -889,3 +985,23 @@ def chunk_fname(self): @property def chunk_bytes(self): return self._chunk_bytes + + +def _fetch_bytes_direct(resolved_url: str, offset: int, length: int) -> bytes: + """Fetch bytes from a resolved URL via HTTP range request. Thread-safe.""" + num_retries = 8 + for try_num in range(num_retries): + try: + range_header = f"bytes={offset}-{offset + length - 1}" + headers = { + "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3", + "Range": range_header + } + response = requests.get(resolved_url, headers=headers) + response.raise_for_status() + return response.content + except Exception as e: + if try_num == num_retries - 1: + raise + time.sleep(0.1 * 2 ** try_num) + assert False, "unreachable" # loop always returns or raises diff --git a/lindi/LindiH5pyFile/LindiH5pyDataset.py b/lindi/LindiH5pyFile/LindiH5pyDataset.py index 8e12b93..bee9bab 100644 --- a/lindi/LindiH5pyFile/LindiH5pyDataset.py +++ b/lindi/LindiH5pyFile/LindiH5pyDataset.py @@ -12,6 +12,7 @@ if TYPE_CHECKING: from .LindiH5pyFile import LindiH5pyFile # pragma: no cover + from ..LindiH5ZarrStore.LindiH5ZarrStore import LindiH5ZarrStore # pragma: no cover # This is a global list of external hdf5 clients, which are used by @@ -20,6 +21,11 @@ # TODO: figure out how to close these clients _external_hdf5_clients: Dict[str, h5py.File] = {} +# Cache of LindiH5ZarrStore instances for remote external array links, +# keyed by URL. Similar to _external_hdf5_clients. +# TODO: figure out how to close these stores (same issue as _external_hdf5_clients) +_external_zarr_stores: Dict[str, "LindiH5ZarrStore"] = {} + class LindiH5pyDataset(h5py.Dataset): def __init__(self, _zarr_array: zarr.Array, _file: "LindiH5pyFile"): @@ -203,10 +209,17 @@ def _get_item_for_zarr(self, zarr_array: zarr.Array, selection: Any): url = external_array_link.get("url", None) name = external_array_link.get("name", None) if url is not None and name is not None: - client = self._get_external_hdf5_client(url) - dataset = client[name] - assert isinstance(dataset, h5py.Dataset) - return dataset[selection] + is_remote = url.startswith("http://") or url.startswith("https://") + if is_remote: + # Use zarr + LindiH5ZarrStore for concurrent chunk fetching + ext_zarr_array = self._get_external_zarr_array(url, name) + return ext_zarr_array[selection] + else: + # Local files — use h5py directly (no concurrency benefit) + client = self._get_external_hdf5_client(url) + dataset = client[name] + assert isinstance(dataset, h5py.Dataset) + return dataset[selection] if self._compound_dtype is not None: # Compound dtype # In this case we index into the compound dtype using the name of the field @@ -252,6 +265,21 @@ def _get_external_hdf5_client(self, url: str) -> h5py.File: _external_hdf5_clients[url] = h5py.File(ff, "r") return _external_hdf5_clients[url] + def _get_external_zarr_array(self, url: str, name: str) -> zarr.Array: + """Get a zarr array for concurrent reading of a remote external array link.""" + from ..LindiH5ZarrStore.LindiH5ZarrStore import LindiH5ZarrStore + from ..LindiH5ZarrStore.LindiH5ZarrStoreOpts import LindiH5ZarrStoreOpts + + if url not in _external_zarr_stores: + # Disable external array links (num_dataset_chunks_threshold=None) + # so all chunks are served through the zarr store + opts = LindiH5ZarrStoreOpts(num_dataset_chunks_threshold=None) + _external_zarr_stores[url] = LindiH5ZarrStore.from_file( + url, opts=opts, local_cache=self._file._local_cache + ) + store = _external_zarr_stores[url] + return zarr.open_array(store=store, path=name, mode='r') + @property def ref(self): if self._readonly: diff --git a/tests/test_concurrent_external_link.py b/tests/test_concurrent_external_link.py new file mode 100644 index 0000000..41c0725 --- /dev/null +++ b/tests/test_concurrent_external_link.py @@ -0,0 +1,140 @@ +import tempfile +import numpy as np +import h5py +import zarr +import lindi +from lindi.LindiH5ZarrStore.LindiH5ZarrStore import LindiH5ZarrStore +from lindi.LindiH5ZarrStore.LindiH5ZarrStoreOpts import LindiH5ZarrStoreOpts + + +def test_getitems_local_chunks(): + """Test getitems on LindiH5ZarrStore with a local chunked dataset.""" + with tempfile.TemporaryDirectory() as tmpdir: + filename = f"{tmpdir}/test.h5" + X = np.random.randn(100, 10) + with h5py.File(filename, "w") as f: + f.create_dataset("dataset1", data=X, chunks=(20, 10)) + + # Use num_dataset_chunks_threshold=None so chunks are served through store + opts = LindiH5ZarrStoreOpts(num_dataset_chunks_threshold=None) + with LindiH5ZarrStore.from_file(filename, url=filename, opts=opts) as store: + # Read via zarr to verify basic functionality + arr = zarr.open_array(store=store, path="dataset1", mode="r") + np.testing.assert_array_equal(arr[:], X) + + # Test getitems with chunk keys + keys = ["dataset1/0.0", "dataset1/1.0", "dataset1/2.0"] + results = store.getitems(keys) + assert len(results) == 3 + for key in keys: + assert key in results + + # Test getitems with metadata keys + meta_keys = ["dataset1/.zarray", "dataset1/.zattrs"] + meta_results = store.getitems(meta_keys) + assert len(meta_results) == 2 + for key in meta_keys: + assert key in meta_results + + # Test getitems with non-existent keys (should be skipped) + mixed_keys = ["dataset1/0.0", "nonexistent/0.0"] + mixed_results = store.getitems(mixed_keys) + assert "dataset1/0.0" in mixed_results + assert "nonexistent/0.0" not in mixed_results + + +def test_getitems_inline_data(): + """Test getitems with a small dataset that is stored inline.""" + with tempfile.TemporaryDirectory() as tmpdir: + filename = f"{tmpdir}/test.h5" + X = np.array([1, 2, 3], dtype=np.float64) + with h5py.File(filename, "w") as f: + f.create_dataset("small", data=X) + + opts = LindiH5ZarrStoreOpts(num_dataset_chunks_threshold=None) + with LindiH5ZarrStore.from_file(filename, url=filename, opts=opts) as store: + # Small arrays should be inline + keys = ["small/0"] + results = store.getitems(keys) + assert len(results) == 1 + + +def test_getitems_single_chunk_shortcut(): + """Test that a single remote chunk skips the thread pool.""" + with tempfile.TemporaryDirectory() as tmpdir: + filename = f"{tmpdir}/test.h5" + X = np.random.randn(1000) + with h5py.File(filename, "w") as f: + f.create_dataset("data", data=X, chunks=(1000,)) + + opts = LindiH5ZarrStoreOpts(num_dataset_chunks_threshold=None) + with LindiH5ZarrStore.from_file(filename, url=filename, opts=opts) as store: + keys = ["data/0"] + results = store.getitems(keys) + assert "data/0" in results + + +def test_external_array_link_via_zarr_store(): + """Test that external array links for local files still work correctly.""" + with tempfile.TemporaryDirectory() as tmpdir: + filename = f"{tmpdir}/test.h5" + X = np.random.randn(50, 12) + with h5py.File(filename, "w") as f: + f.create_dataset("dataset1", data=X, chunks=(10, 6)) + + # Create a LINDI reference with a low threshold so external array link is used + with LindiH5ZarrStore.from_file( + filename, + url=filename, + opts=LindiH5ZarrStoreOpts(num_dataset_chunks_threshold=4), + ) as store: + rfs = store.to_reference_file_system() + + # Read back through LindiH5pyFile — local external links use h5py directly + client = lindi.LindiH5pyFile.from_reference_file_system(rfs) + X2 = client["dataset1"][:] + np.testing.assert_array_equal(X, X2) + + +def test_zarr_store_for_external_array(): + """Test creating a LindiH5ZarrStore with num_dataset_chunks_threshold=None + to serve all chunks (the pattern used for remote external array links).""" + with tempfile.TemporaryDirectory() as tmpdir: + filename = f"{tmpdir}/test.h5" + X = np.random.randn(200, 10) + with h5py.File(filename, "w") as f: + f.create_dataset("dataset1", data=X, chunks=(20, 10)) + + # This is the same pattern used in _get_external_zarr_array + opts = LindiH5ZarrStoreOpts(num_dataset_chunks_threshold=None) + with LindiH5ZarrStore.from_file(filename, opts=opts, url=filename) as store: + arr = zarr.open_array(store=store, path="dataset1", mode="r") + result = arr[:] + np.testing.assert_array_equal(result, X) + + # Test slicing + result_slice = arr[10:30, 3:7] + np.testing.assert_array_equal(result_slice, X[10:30, 3:7]) + + +def test_getitems_empty_keys(): + """Test getitems with empty key list.""" + with tempfile.TemporaryDirectory() as tmpdir: + filename = f"{tmpdir}/test.h5" + with h5py.File(filename, "w") as f: + f.create_dataset("data", data=np.array([1, 2, 3])) + + opts = LindiH5ZarrStoreOpts(num_dataset_chunks_threshold=None) + with LindiH5ZarrStore.from_file(filename, url=filename, opts=opts) as store: + results = store.getitems([]) + assert results == {} + + +if __name__ == "__main__": + test_getitems_local_chunks() + test_getitems_inline_data() + test_getitems_single_chunk_shortcut() + test_external_array_link_via_zarr_store() + test_zarr_store_for_external_array() + test_getitems_empty_keys() + print("All tests passed!") From 797ae26dd17368e8c37b495926d3708b5e1dffbc Mon Sep 17 00:00:00 2001 From: Ben Dichter Date: Sat, 14 Mar 2026 09:29:50 +0000 Subject: [PATCH 2/3] Add byte range coalescing for fewer HTTP requests Merge adjacent/nearby chunk byte ranges into single HTTP requests before concurrent fetching. Two configurable parameters control the behavior: - coalesce_merge_gap (default 256KB): max gap between ranges to merge - coalesce_max_size (default 20MB): max size of a coalesced request Co-Authored-By: Claude Opus 4.6 (1M context) --- lindi/LindiH5ZarrStore/LindiH5ZarrStore.py | 88 ++++++++++++++++++---- tests/test_concurrent_external_link.py | 75 ++++++++++++++++++ 2 files changed, 148 insertions(+), 15 deletions(-) diff --git a/lindi/LindiH5ZarrStore/LindiH5ZarrStore.py b/lindi/LindiH5ZarrStore/LindiH5ZarrStore.py index 6527f64..c4ea245 100644 --- a/lindi/LindiH5ZarrStore/LindiH5ZarrStore.py +++ b/lindi/LindiH5ZarrStore/LindiH5ZarrStore.py @@ -154,7 +154,9 @@ def __init__( _url: Union[str, None] = None, _entities_to_close: List[Any], _local_cache: Union[LocalCache, None] = None, - _concurrent_max_workers: int = 8 + _concurrent_max_workers: int = 8, + _coalesce_merge_gap: int = 256 * 1024, + _coalesce_max_size: int = 20 * 1024 * 1024 ): """ Do not call the constructor directly. Instead, use the from_file class @@ -166,6 +168,8 @@ def __init__( self._opts = _opts self._local_cache = _local_cache self._concurrent_max_workers = _concurrent_max_workers + self._coalesce_merge_gap = _coalesce_merge_gap + self._coalesce_max_size = _coalesce_max_size self._entities_to_close = _entities_to_close + [self._h5f] # Some datasets do not correspond to traditional chunked datasets. For @@ -382,26 +386,33 @@ def getitems(self, keys, *, contexts=None): from ..LindiRemfile.LindiRemfile import _resolve_url resolved_url = _resolve_url(self._url) - # Single chunk — skip thread pool overhead - if len(remote_chunks) == 1: - key, offset, count = remote_chunks[0] - val = _fetch_bytes_direct(resolved_url, offset, count) - self._try_cache_put(offset, count, val) - results[key] = val + # Coalesce adjacent/nearby byte ranges to reduce HTTP round-trips + coalesced = _coalesce_byte_ranges(remote_chunks, merge_gap=self._coalesce_merge_gap, max_size=self._coalesce_max_size) + + # Single request — skip thread pool overhead + if len(coalesced) == 1: + group_offset, group_length, members = coalesced[0] + buf = _fetch_bytes_direct(resolved_url, group_offset, group_length) + for key, offset, count in members: + val = buf[offset - group_offset: offset - group_offset + count] + self._try_cache_put(offset, count, val) + results[key] = val return self._apply_padding_to_results(results) - # Concurrent fetch - max_workers = min(len(remote_chunks), self._concurrent_max_workers) + # Concurrent fetch of coalesced groups + max_workers = min(len(coalesced), self._concurrent_max_workers) with ThreadPoolExecutor(max_workers=max_workers) as pool: futures = { - pool.submit(_fetch_bytes_direct, resolved_url, offset, count): (key, offset, count) - for key, offset, count in remote_chunks + pool.submit(_fetch_bytes_direct, resolved_url, group_offset, group_length): (group_offset, members) + for group_offset, group_length, members in coalesced } for future in as_completed(futures): - key, offset, count = futures[future] - val = future.result() - self._try_cache_put(offset, count, val) - results[key] = val + group_offset, members = futures[future] + buf = future.result() + for key, offset, count in members: + val = buf[offset - group_offset: offset - group_offset + count] + self._try_cache_put(offset, count, val) + results[key] = val return self._apply_padding_to_results(results) @@ -987,6 +998,53 @@ def chunk_bytes(self): return self._chunk_bytes +def _coalesce_byte_ranges(chunks, *, merge_gap, max_size=None): + """Merge byte ranges that are adjacent or within merge_gap bytes of each other. + + Args: + chunks: list of (key, byte_offset, byte_count) tuples + merge_gap: maximum gap in bytes between ranges to merge + max_size: maximum total size in bytes for a coalesced group. If merging + a chunk would exceed this, start a new group instead. + + Returns: + list of (group_offset, group_length, members) where members is a list + of (key, byte_offset, byte_count) tuples from the original input. + """ + if not chunks: + return [] + + # Sort by byte offset + sorted_chunks = sorted(chunks, key=lambda x: x[1]) + + groups = [] + # Start first group + key, offset, count = sorted_chunks[0] + group_start = offset + group_end = offset + count + group_members = [(key, offset, count)] + + for key, offset, count in sorted_chunks[1:]: + chunk_end = offset + count + merged_size = max(group_end, chunk_end) - group_start + within_gap = offset <= group_end + merge_gap + within_size = max_size is None or merged_size <= max_size + if within_gap and within_size: + # Merge into current group + group_end = max(group_end, chunk_end) + group_members.append((key, offset, count)) + else: + # Finalize current group, start new one + groups.append((group_start, group_end - group_start, group_members)) + group_start = offset + group_end = chunk_end + group_members = [(key, offset, count)] + + # Finalize last group + groups.append((group_start, group_end - group_start, group_members)) + return groups + + def _fetch_bytes_direct(resolved_url: str, offset: int, length: int) -> bytes: """Fetch bytes from a resolved URL via HTTP range request. Thread-safe.""" num_retries = 8 diff --git a/tests/test_concurrent_external_link.py b/tests/test_concurrent_external_link.py index 41c0725..74adb90 100644 --- a/tests/test_concurrent_external_link.py +++ b/tests/test_concurrent_external_link.py @@ -130,6 +130,79 @@ def test_getitems_empty_keys(): assert results == {} +def test_coalesce_byte_ranges(): + """Test that _coalesce_byte_ranges merges adjacent/nearby ranges.""" + from lindi.LindiH5ZarrStore.LindiH5ZarrStore import _coalesce_byte_ranges + + # Empty input + assert _coalesce_byte_ranges([], merge_gap=256) == [] + + # Single chunk — no merging + chunks = [("a", 0, 100)] + groups = _coalesce_byte_ranges(chunks, merge_gap=256) + assert len(groups) == 1 + assert groups[0] == (0, 100, [("a", 0, 100)]) + + # Adjacent chunks — should merge + chunks = [("a", 0, 100), ("b", 100, 100)] + groups = _coalesce_byte_ranges(chunks, merge_gap=0) + assert len(groups) == 1 + assert groups[0][0] == 0 # group_start + assert groups[0][1] == 200 # group_length + assert len(groups[0][2]) == 2 # two members + + # Chunks with small gap — should merge + chunks = [("a", 0, 100), ("b", 200, 100)] + groups = _coalesce_byte_ranges(chunks, merge_gap=100) + assert len(groups) == 1 + assert groups[0][0] == 0 + assert groups[0][1] == 300 + + # Chunks with gap larger than threshold — should NOT merge + chunks = [("a", 0, 100), ("b", 500, 100)] + groups = _coalesce_byte_ranges(chunks, merge_gap=100) + assert len(groups) == 2 + + # Unsorted input — should still work + chunks = [("c", 2000, 100), ("a", 0, 100), ("b", 100, 100)] + groups = _coalesce_byte_ranges(chunks, merge_gap=50) + assert len(groups) == 2 + assert groups[0][0] == 0 + assert groups[0][1] == 200 + assert groups[1][0] == 2000 + + # max_size — prevents merging when result would be too large + chunks = [("a", 0, 100), ("b", 100, 100), ("c", 200, 100)] + groups = _coalesce_byte_ranges(chunks, merge_gap=100, max_size=200) + assert len(groups) == 2 + assert groups[0][1] == 200 # first two merged + assert groups[1][1] == 100 # third alone + + # max_size=None — no limit + groups = _coalesce_byte_ranges(chunks, merge_gap=100, max_size=None) + assert len(groups) == 1 + + +def test_coalesce_integration(): + """Test that coalesced fetching returns correct data through zarr.""" + with tempfile.TemporaryDirectory() as tmpdir: + filename = f"{tmpdir}/test.h5" + X = np.random.randn(200, 10) + with h5py.File(filename, "w") as f: + f.create_dataset("dataset1", data=X, chunks=(20, 10)) + + opts = LindiH5ZarrStoreOpts(num_dataset_chunks_threshold=None) + with LindiH5ZarrStore.from_file(filename, opts=opts, url=filename) as store: + # Fetch all 10 chunk keys at once — should coalesce + keys = [f"dataset1/{i}.0" for i in range(10)] + results = store.getitems(keys) + assert len(results) == 10 + + # Verify data through zarr is correct + arr = zarr.open_array(store=store, path="dataset1", mode="r") + np.testing.assert_array_equal(arr[:], X) + + if __name__ == "__main__": test_getitems_local_chunks() test_getitems_inline_data() @@ -137,4 +210,6 @@ def test_getitems_empty_keys(): test_external_array_link_via_zarr_store() test_zarr_store_for_external_array() test_getitems_empty_keys() + test_coalesce_byte_ranges() + test_coalesce_integration() print("All tests passed!") From 9dddcc9e8132236a7985de70f69d68b9591612d2 Mon Sep 17 00:00:00 2001 From: Ben Dichter Date: Sat, 14 Mar 2026 09:39:00 +0000 Subject: [PATCH 3/3] Add benchmark script for concurrent chunk fetching Standalone script in devel/ that compares serial (h5py+LindiRemfile) vs concurrent (zarr+LindiH5ZarrStore) reads from DANDI, asserts data equivalence, and produces a bar chart showing timings and speedup. Usage: python devel/benchmark_concurrent_fetch.py [--dandiset 000473] [-o results.png] Co-Authored-By: Claude Opus 4.6 (1M context) --- devel/benchmark_concurrent_fetch.py | 248 ++++++++++++++++++++++++++++ 1 file changed, 248 insertions(+) create mode 100644 devel/benchmark_concurrent_fetch.py diff --git a/devel/benchmark_concurrent_fetch.py b/devel/benchmark_concurrent_fetch.py new file mode 100644 index 0000000..066e7bd --- /dev/null +++ b/devel/benchmark_concurrent_fetch.py @@ -0,0 +1,248 @@ +"""Benchmark concurrent chunk fetching vs serial reads from DANDI Archive. + +Compares the old path (h5py + LindiRemfile, serial) against the new path +(zarr + LindiH5ZarrStore, concurrent) for reading external array links. + +Produces a bar chart showing timings and speedup for each test case. + +Usage: + python devel/benchmark_concurrent_fetch.py + python devel/benchmark_concurrent_fetch.py --dandiset 000473 + python devel/benchmark_concurrent_fetch.py --dandiset 000409 + python devel/benchmark_concurrent_fetch.py --output benchmark_results.png +""" +import argparse +import time +import tempfile +import numpy as np +import h5py +import zarr +import matplotlib.pyplot as plt +import lindi +from lindi.LindiRemfile.LindiRemfile import LindiRemfile +from lindi.LindiH5ZarrStore.LindiH5ZarrStore import LindiH5ZarrStore +from lindi.LindiH5ZarrStore.LindiH5ZarrStoreOpts import LindiH5ZarrStoreOpts + + +DANDISETS = { + "000473": { + "url": "https://api.dandiarchive.org/api/assets/11f512ba-5bcf-4230-a8cb-dc8d36db38cb/download/", + "dataset": "processing/ecephys/LFP/LFP/data", + "label": "000473 LFP", + "slices": [ + (np.s_[:1000], "[:1000]"), + (np.s_[:5000], "[:5000]"), + (np.s_[:10000], "[:10000]"), + ], + }, + "000409": { + "url": "https://api.dandiarchive.org/api/assets/c04f6b30-82bf-40e1-9210-34f0bcd8be24/download/", + "dataset": "acquisition/ElectricalSeriesAp/data", + "label": "000409 Neuropixels", + "slices": [ + (np.s_[:500], "[:500]"), + (np.s_[:1000], "[:1000]"), + (np.s_[:2000], "[:2000]"), + ], + }, +} + + +def read_serial(url, dataset_name, selection): + """Old path: serial reads through h5py + LindiRemfile.""" + remf = LindiRemfile(url, verbose=False, local_cache=None) + with h5py.File(remf, "r") as f: + return f[dataset_name][selection] + + +def read_concurrent(url, dataset_name, selection): + """New path: concurrent reads through zarr + LindiH5ZarrStore.""" + opts = LindiH5ZarrStoreOpts(num_dataset_chunks_threshold=None) + with LindiH5ZarrStore.from_file(url, opts=opts) as store: + arr = zarr.open_array(store=store, path=dataset_name, mode="r") + return arr[selection] + + +def read_lindi_file(lindi_json_path, dataset_name, selection): + """Full pipeline: LindiH5pyFile (concurrent path for remote external links).""" + f = lindi.LindiH5pyFile.from_lindi_file(lindi_json_path, mode="r") + return f[dataset_name][selection] + + +def benchmark_one(func, *args, **kwargs): + """Time a single call. Returns (result, elapsed_seconds).""" + t0 = time.perf_counter() + result = func(*args, **kwargs) + elapsed = time.perf_counter() - t0 + return result, elapsed + + +def run_benchmark(dandiset_key): + """Run benchmarks for one dandiset. Returns list of result dicts.""" + info = DANDISETS[dandiset_key] + url = info["url"] + dataset = info["dataset"] + label = info["label"] + + print(f"\n{'=' * 60}") + print(f"Benchmarking: {label}") + print(f" URL: {url}") + print(f" Dataset: {dataset}") + print(f"{'=' * 60}") + + results = [] + for selection, sel_label in info["slices"]: + print(f"\n Slice: {dataset}{sel_label}") + print(f" {'-' * 40}") + + # Serial (old path) + data_serial, t_serial = benchmark_one(read_serial, url, dataset, selection) + print(f" Serial (h5py+LindiRemfile): {t_serial:.2f}s") + + # Concurrent (new path) + data_concurrent, t_concurrent = benchmark_one(read_concurrent, url, dataset, selection) + print(f" Concurrent (zarr+LindiH5ZarrStore): {t_concurrent:.2f}s") + + # Equivalence check + np.testing.assert_array_equal(data_serial, data_concurrent) + print(f" Data equivalent: shape={data_serial.shape}, dtype={data_serial.dtype}") + + speedup = t_serial / t_concurrent if t_concurrent > 0 else float("inf") + print(f" Speedup: {speedup:.1f}x") + + results.append({ + "dandiset": dandiset_key, + "label": f"{label}\n{sel_label}", + "short_label": sel_label, + "serial": t_serial, + "concurrent": t_concurrent, + "speedup": speedup, + "shape": data_serial.shape, + }) + + return results + + +def run_lindi_file_benchmark(dandiset_key): + """Run the full LindiH5pyFile pipeline benchmark.""" + info = DANDISETS[dandiset_key] + url = info["url"] + dataset = info["dataset"] + selection, sel_label = info["slices"][1] # use middle slice + + print(f"\n{'=' * 60}") + print(f"LindiH5pyFile Pipeline: {info['label']}") + print(f"{'=' * 60}") + + with tempfile.TemporaryDirectory() as tmpdir: + print(" Creating .lindi.json...") + t0 = time.perf_counter() + fname = f"{tmpdir}/test.nwb.lindi.json" + with lindi.LindiH5pyFile.from_hdf5_file(url) as f: + f.write_lindi_file(fname) + t_create = time.perf_counter() - t0 + print(f" Created in {t_create:.2f}s") + + # Serial baseline + data_serial, t_serial = benchmark_one(read_serial, url, dataset, selection) + print(f" Serial (h5py+LindiRemfile): {t_serial:.2f}s") + + # LindiH5pyFile (uses concurrent path) + data_lindi, t_lindi = benchmark_one(read_lindi_file, fname, dataset, selection) + print(f" LindiH5pyFile (concurrent): {t_lindi:.2f}s") + + np.testing.assert_array_equal(data_serial, data_lindi) + speedup = t_serial / t_lindi if t_lindi > 0 else float("inf") + print(f" Data equivalent: shape={data_serial.shape}") + print(f" Speedup: {speedup:.1f}x") + + return { + "label": f"LindiH5pyFile\n{info['label']} {sel_label}", + "serial": t_serial, + "concurrent": t_lindi, + "speedup": speedup, + } + + +def plot_results(all_results, output_path=None): + """Create a bar chart comparing serial vs concurrent timings.""" + labels = [r["label"] for r in all_results] + serial_times = [r["serial"] for r in all_results] + concurrent_times = [r["concurrent"] for r in all_results] + speedups = [r["speedup"] for r in all_results] + + x = np.arange(len(labels)) + width = 0.35 + + fig, ax = plt.subplots(figsize=(max(10, len(labels) * 2), 6)) + bars_serial = ax.bar(x - width / 2, serial_times, width, label="Serial (h5py + LindiRemfile)", color="#d35f5f") + bars_concurrent = ax.bar(x + width / 2, concurrent_times, width, label="Concurrent (zarr + LindiH5ZarrStore)", color="#5f9ed3") + + # Add speedup annotations + for i, (s_time, c_time, speedup) in enumerate(zip(serial_times, concurrent_times, speedups)): + y = max(s_time, c_time) + ax.annotate( + f"{speedup:.1f}x", + xy=(i, y), + xytext=(0, 8), + textcoords="offset points", + ha="center", + fontweight="bold", + fontsize=11, + ) + + ax.set_ylabel("Time (seconds)") + ax.set_title("LINDI Chunk Fetching: Serial vs Concurrent") + ax.set_xticks(x) + ax.set_xticklabels(labels, fontsize=9) + ax.legend() + ax.set_ylim(0, max(serial_times) * 1.3) + + fig.tight_layout() + + if output_path: + fig.savefig(output_path, dpi=150) + print(f"\nPlot saved to: {output_path}") + else: + plt.show() + + +def main(): + parser = argparse.ArgumentParser(description="Benchmark concurrent chunk fetching from DANDI") + parser.add_argument( + "--dandiset", + choices=list(DANDISETS.keys()), + default=None, + help="Run benchmarks for a specific dandiset only (default: run all)", + ) + parser.add_argument( + "--output", "-o", + default=None, + help="Save plot to file instead of displaying (e.g. benchmark.png)", + ) + parser.add_argument( + "--skip-lindi-file", + action="store_true", + help="Skip the full LindiH5pyFile pipeline benchmark (slow due to .lindi.json creation)", + ) + args = parser.parse_args() + + dandisets = [args.dandiset] if args.dandiset else list(DANDISETS.keys()) + + all_results = [] + for key in dandisets: + all_results.extend(run_benchmark(key)) + + if not args.skip_lindi_file: + for key in dandisets: + all_results.append(run_lindi_file_benchmark(key)) + + print(f"\n{'=' * 60}") + print("All data equivalence checks passed!") + print(f"{'=' * 60}") + + plot_results(all_results, output_path=args.output) + + +if __name__ == "__main__": + main()