From 4044682a48f8dec760d84445e70277713adc893f Mon Sep 17 00:00:00 2001 From: Janick Martinez Esturo Date: Sun, 11 Oct 2026 14:04:00 +0200 Subject: [PATCH] perf(data): node cache for component readers Opening a group / array constructs its zarr-python object, which decodes the node's metadata on every access (~0.5 ms per frame lookup with zarr-python 2, dominating reads of small frame arrays). Component readers now cache the nodes they open (with their decoded metadata, but not their data) in a thread-safe nodes.NodeCache shared by all member groups of the component. The cache size is configurable via SequenceComponentGroupsReader(node_cache_size=...): None (default) caches all accessed nodes, N the N most recently used nodes, and 0 disables caching. With the cache, repeated frame lookups take ~2 us instead of ~470 us. Component stores are opened concurrently (hiding the latency of remote archive index reads), their hierarchies are now opened from their consolidated metadata sequentially on the calling thread, each as soon as its store is opened (while the remaining stores are still being opened): opening hierarchies is pure-Python CPU work, which concurrent threads only slow down. --- docs/data/storage_and_access.rst | 8 ++ ncore/impl/data/nodes.py | 99 +++++++++++++++--- ncore/impl/data/stores.py | 4 +- ncore/impl/data/stores_test.py | 72 +++++++++++++ ncore/impl/data/v4/components.py | 40 +++++--- ncore/impl/data/v4/components_test.py | 140 ++++++++++++++++++++++++-- 6 files changed, 327 insertions(+), 36 deletions(-) diff --git a/docs/data/storage_and_access.rst b/docs/data/storage_and_access.rst index 38d9911d..202799ca 100644 --- a/docs/data/storage_and_access.rst +++ b/docs/data/storage_and_access.rst @@ -361,6 +361,14 @@ Performance Recommendations This is especially important for remote stores where each metadata lookup would otherwise be a separate round-trip. +* **Keep the node cache enabled** (the default). Component readers cache the + groups / arrays they open (including their decoded metadata, but not their + data), so repeated accesses (e.g., of frames) skip re-opening them. The + ``node_cache_size`` parameter on + :class:`~ncore.data.v4.SequenceComponentGroupsReader` bounds the number of + cached nodes per component reader for very long sequences (``None`` caches + all accessed nodes, ``0`` disables caching). + * **Tune the block size and cache type** for your workload. The ``default_block_size`` and ``default_cache_type`` parameters on ``UPath`` control how much data is fetched per HTTP request and how it is cached diff --git a/ncore/impl/data/nodes.py b/ncore/impl/data/nodes.py index 7d526a1a..b26c8f6b 100644 --- a/ncore/impl/data/nodes.py +++ b/ncore/impl/data/nodes.py @@ -22,6 +22,9 @@ from __future__ import annotations +import collections +import threading + from typing import TYPE_CHECKING, Any, Iterator, List, Literal, Optional, Tuple, Type, TypeVar, Union, cast import numpy as np @@ -48,6 +51,53 @@ T = TypeVar("T") +class NodeCache: + """Cache of the opened nodes of a hierarchy opened for reading, by node path. + + Opening a node constructs its zarr-python object (decoding its metadata), which dominates repeated reads of small + arrays (e.g., of frames). Cached nodes reference their metadata only, never array data. ``max_entries=None`` caches + all opened nodes, ``max_entries=N`` the N most recently used nodes (evicted nodes are re-opened on access). + + All methods are thread-safe. + """ + + def __init__(self, max_entries: Optional[int] = None) -> None: + if max_entries is not None and max_entries < 1: + raise ValueError(f"max_entries needs to be positive or None (unbounded), got {max_entries}") + self._max_entries = max_entries + self._nodes: collections.OrderedDict[str, ZarrNode] = collections.OrderedDict() + self._lock = threading.Lock() + + @property + def max_entries(self) -> Optional[int]: + return self._max_entries + + def __len__(self) -> int: + return len(self._nodes) + + def keys(self) -> List[str]: + """The paths of the cached nodes, from least to most recently used""" + with self._lock: + return list(self._nodes) + + def _open(self, group: ZarrGroup, path: str) -> ZarrNode: + key = _Node._join(group.path, path) + with self._lock: + if (node := self._nodes.get(key)) is not None: + self._nodes.move_to_end(key) + return node + + # (concurrent first accesses open nodes redundantly, which is cheaper than serializing all accesses) + node = _backend.open_member(group, path) + + with self._lock: + self._nodes[key] = node + # (at most one entry is added per call, so at most the least recently used node is evicted) + if self._max_entries is not None and len(self._nodes) > self._max_entries: + self._nodes.popitem(last=False) + return node + + class _Node: """Common functionality of groups and arrays""" @@ -100,9 +150,10 @@ def _typed_attr(self, name: str, expected: Type[T]) -> T: raise TypeError(f"Attribute {name!r} of {self.path!r} is not of type {expected.__name__}: {value!r}") return cast(T, value) - def _member_path(self, path: str) -> str: - """The path of a (nested) member of the node within its store""" - return f"{self.path}/{path}" if self.path else path + @staticmethod + def _join(parent: str, path: str) -> str: + """The path of a (nested) member of a node within its store""" + return f"{parent}/{path}" if parent else path def __repr__(self) -> str: return f"{type(self).__name__}({self.path!r}, zarr_format={self.zarr_format})" @@ -159,11 +210,15 @@ def read_bytes(self) -> bytes: class Group(_Node): - """A group of an NCore data store""" + """A group of an NCore data store. - def __init__(self, group: ZarrGroup, store: Optional[Store] = None) -> None: + Members of groups using a :class:`NodeCache` are opened via the cache (also members of member groups). + """ + + def __init__(self, group: ZarrGroup, cache: Optional[NodeCache] = None, store: Optional[Store] = None) -> None: super().__init__(group) self._group = group + self._cache = cache self._store = store @property @@ -176,23 +231,37 @@ def store(self) -> Optional[Store]: """The store of the group (None for groups not opened / created via a :class:`~ncore.impl.data.stores.Store`)""" return self._store + @property + def cache(self) -> Optional[NodeCache]: + """The node cache used to open members (None if not cached)""" + return self._cache + + def with_cache(self, cache: Optional[NodeCache]) -> Group: + """Returns the group, using the given node cache to open its members""" + return Group(self._group, cache, self._store) + # -- reading --------------------------------------------------------------- + def _open(self, path: str) -> ZarrNode: + if self._cache is not None: + return self._cache._open(self._group, path) + return _backend.open_member(self._group, path) + def node(self, path: str) -> Union[Group, Array]: """Returns a (nested) member group or array (raises KeyError if it doesn't exist)""" - node = _backend.open_member(self._group, path) - return Group(node, self._store) if isinstance(node, ZarrGroup) else Array(node) + node = self._open(path) + return Group(node, self._cache, self._store) if isinstance(node, ZarrGroup) else Array(node) def group(self, path: str) -> Group: """Returns a (nested) member group (raises KeyError if it doesn't exist or isn't a group)""" - if not isinstance(node := _backend.open_member(self._group, path), ZarrGroup): - raise KeyError(f"{self._member_path(path)} is not a group") - return Group(node, self._store) + if not isinstance(node := self._open(path), ZarrGroup): + raise KeyError(f"{_Node._join(self.path, path)} is not a group") + return Group(node, self._cache, self._store) def array(self, path: str) -> Array: """Returns a (nested) member array (raises KeyError if it doesn't exist or isn't an array)""" - if not isinstance(node := _backend.open_member(self._group, path), ZarrArray): - raise KeyError(f"{self._member_path(path)} is not an array") + if not isinstance(node := self._open(path), ZarrArray): + raise KeyError(f"{_Node._join(self.path, path)} is not an array") return Array(node) def members(self) -> List[str]: @@ -206,7 +275,7 @@ def __contains__(self, name: str) -> bool: def groups(self) -> Iterator[Tuple[str, Group]]: """Iterates over the direct member groups (ordered by name)""" for name, group in sorted(_backend.member_groups(self._group), key=lambda member: member[0]): - yield name, Group(group, self._store) + yield name, Group(group, self._cache, self._store) # (zarr-python style read access, without the equality semantics of mappings: groups compare by identity) @@ -242,11 +311,11 @@ def __len__(self) -> int: def create_group(self, name: str, *, attributes: Optional[Attributes] = None) -> Group: """Creates a member group (raises if a member with the name exists already)""" - return Group(_backend.create_group(self._group, name, attributes), self._store) + return Group(_backend.create_group(self._group, name, attributes), store=self._store) def require_group(self, name: str) -> Group: """Returns a member group, creating it if it doesn't exist yet""" - return Group(_backend.require_group(self._group, name), self._store) + return Group(_backend.require_group(self._group, name), store=self._store) def create_array( self, diff --git a/ncore/impl/data/stores.py b/ncore/impl/data/stores.py index bea697db..c4957dda 100644 --- a/ncore/impl/data/stores.py +++ b/ncore/impl/data/stores.py @@ -54,12 +54,12 @@ def zarr(self) -> ZarrStore: def create_root_group(self, attributes: Attributes) -> Group: """Creates the root group of the hierarchy of the store""" - return Group(_backend.create_root_group(self._zarr_store, 2, attributes), self) + return Group(_backend.create_root_group(self._zarr_store, 2, attributes), store=self) def open_root_group(self, consolidated: bool = True) -> Group: """Opens the root group of the hierarchy of the store for reading, via its compressed consolidated metadata if consolidated (raises KeyError if the store has no compressed consolidated metadata)""" - return Group(_backend.open_root_group(self._zarr_store, consolidated), self) + return Group(_backend.open_root_group(self._zarr_store, consolidated), store=self) def consolidate(self) -> None: """Stores the compressed consolidated metadata of the hierarchy of the store (after writing all nodes)""" diff --git a/ncore/impl/data/stores_test.py b/ncore/impl/data/stores_test.py index dabf75a5..98ef70db 100644 --- a/ncore/impl/data/stores_test.py +++ b/ncore/impl/data/stores_test.py @@ -546,6 +546,78 @@ def test_names(self): self.assertEqual(self.root["a/b/c"].name, "/a/b/c") +class TestNodeCache(unittest.TestCase): + """Node caches of groups (disabled / unbounded / least recently used)""" + + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + path = os.path.join(self._tmp.name, "store.itar") + with IndexedTarZarrStore(path, mode="w") as store: + sub = store.create_root_group({}).create_group("g") + for i in range(8): + sub.create_array(f"a{i}", np.full(3, i)) + store.consolidate() + self.root = IndexedTarZarrStore(path).open_root_group() + + def tearDown(self): + self._tmp.cleanup() + + def test_uncached(self): + self.assertIsNone(self.root.cache) + self.assertIsNot(self.root.array("g/a1").zarr, self.root.array("g/a1").zarr) + + def test_unbounded(self): + root = self.root.with_cache(cache := nodes.NodeCache()) + for i in range(8): + np.testing.assert_array_equal(root.array(f"g/a{i}").read(), [i] * 3) + self.assertEqual(len(cache), 8) + self.assertIs(root.array("g/a3").zarr, root.array("g/a3").zarr) + + def test_member_groups_share_the_cache(self): + root = self.root.with_cache(cache := nodes.NodeCache()) + group = root.group("g") + self.assertIs(group.cache, cache) + node = root.node("g") + assert isinstance(node, nodes.Group) + self.assertIs(node.cache, cache) + self.assertIs(cast(nodes.Array, root.node("g/a2")).zarr, root.array("g/a2").zarr) + self.assertIs(group.array("a3").zarr, root.array("g/a3").zarr) + self.assertIs(next(root.groups())[1].cache, cache) + self.assertEqual(cache.keys(), ["g", "g/a2", "g/a3"]) + + def test_zarr_python_style_access_uses_the_cache(self): + root = self.root.with_cache(cache := nodes.NodeCache()) + assert isinstance(group := root["g"], nodes.Group) + self.assertIs(group.cache, cache) + self.assertIs(cast(nodes.Array, root["g/a1"]).zarr, cast(nodes.Array, group.get("a1")).zarr) + self.assertIsNone(group.get("missing")) + members = dict(group.items()) + self.assertIs(cast(nodes.Array, members["a2"]).zarr, root.array("g/a2").zarr) + self.assertEqual(sorted(cache.keys()), ["g"] + [f"g/a{i}" for i in range(8)]) + + def test_least_recently_used(self): + root = self.root.with_cache(cache := nodes.NodeCache(max_entries=3)) + a0 = root.array("g/a0").zarr + root.array("g/a1") + root.array("g/a2") + self.assertIs(root.array("g/a0").zarr, a0) # (a1 is the least recently used node now) + root.array("g/a3") + self.assertEqual(cache.keys(), ["g/a2", "g/a0", "g/a3"]) + for i in range(8): # evicted nodes are re-opened + np.testing.assert_array_equal(root.array(f"g/a{i}").read(), [i] * 3) + self.assertEqual(len(cache), 3) + + def test_missing_nodes_are_not_cached(self): + root = self.root.with_cache(cache := nodes.NodeCache()) + with self.assertRaises(KeyError): + root.array("g/missing") + self.assertEqual(len(cache), 0) + + def test_invalid_sizes(self): + with self.assertRaises(ValueError): + nodes.NodeCache(max_entries=0) + + class TestZeroCopyReads(unittest.TestCase): """Read results are read-only, arrays stored as a single uncompressed chunk are read as views of their records""" diff --git a/ncore/impl/data/v4/components.py b/ncore/impl/data/v4/components.py index 770feb69..8c2be1a0 100644 --- a/ncore/impl/data/v4/components.py +++ b/ncore/impl/data/v4/components.py @@ -346,6 +346,7 @@ def __init__( open_consolidated: bool = True, itar_index_tail_read_size: int = 1 << 20, # 1 MiB default max_threads: int | None = None, + node_cache_size: int | None = None, ): """Initialize a SequenceComponentReader for a virtual sequence represented by a list of components. @@ -361,10 +362,18 @@ def __init__( single I/O call to initialize the tar file index to minimize reads. Default 1 MiB size fits typical use cases, but may be increased if the compressed index size is expected to be larger or decreased if remote access chunk sizes are smaller. - max_threads: The maximum number of threads used to load the different components (if None, - use interpreter-default number of threads for a ThreadPoolExecutor) + max_threads: The maximum number of threads used to open the component stores concurrently (if + None, use interpreter-default number of threads for a ThreadPoolExecutor) + node_cache_size: Number of opened groups / arrays (with their decoded meta-data, but not their data) + each component reader caches for repeated accesses (e.g., of frames). None caches + all accessed nodes for the lifetime of the reader, a positive number the most recently + used nodes only, and 0 disables caching. """ + if node_cache_size is not None and node_cache_size < 0: + raise ValueError(f"node_cache_size needs to be None or non-negative, got {node_cache_size}") + self._node_cache_size = node_cache_size + component_group_upaths: List[UPath] = self.expand_component_group_paths(component_group_paths) assert len(component_group_upaths), "No component inputs provided" @@ -372,10 +381,13 @@ def __init__( # Load component stores concurrently (to hide latency) and check for sequence consistency self._component_stores: Dict[str, Tuple[nodes.Group, UPath]] = {} # use str as the generic path / URL type + # Component stores are opened concurrently (hiding the latency of remote archive index reads), their + # hierarchies are opened from their consolidated metadata sequentially on the calling thread (pure-Python CPU + # work that threads only slow down), each as soon as its store is opened with concurrent.futures.ThreadPoolExecutor(max_workers=max_threads) as executor: - def thread_load_component_store(component_store_upath: UPath) -> Tuple[nodes.Group, UPath]: - """Thread-executed shard opening""" + def thread_open_component_store(component_store_upath: UPath) -> Tuple[stores.Store, UPath]: + """Thread-executed component store opening""" # Make sure paths are absolute at this point component_store_upath = component_store_upath.absolute() @@ -396,16 +408,14 @@ def thread_load_component_store(component_store_upath: UPath) -> Tuple[nodes.Gro else: component_store = stores.ZarrStoreAdapter.directory(component_store_upath) - return component_store.open_root_group(open_consolidated), component_store_upath + return component_store, component_store_upath - for future in concurrent.futures.as_completed( - [ - executor.submit(thread_load_component_store, component_store_upath) - for component_store_upath in component_group_upaths - ] + # (hierarchies are opened as soon as their stores are opened, in order, while the remaining stores + # are still being opened) + for component_store, component_store_path in executor.map( + thread_open_component_store, component_group_upaths ): - # Note: thread completion order is not relevant here - component_root, component_store_path = future.result() + component_root = component_store.open_root_group(open_consolidated) # Check sequence compatibility component_root_version = component_root.attr_str("version") @@ -491,7 +501,11 @@ def open_component_readers( if not component_reader_type.supports_component_version(component_group.attr_str("component_version")): continue - ret[component_instance_name] = component_reader_type(component_instance_name, component_group) + # each reader caches the nodes of its component + cache = None if self._node_cache_size == 0 else nodes.NodeCache(self._node_cache_size) + ret[component_instance_name] = component_reader_type( + component_instance_name, component_group.with_cache(cache) + ) return ret diff --git a/ncore/impl/data/v4/components_test.py b/ncore/impl/data/v4/components_test.py index 94247cbe..3ea3da56 100644 --- a/ncore/impl/data/v4/components_test.py +++ b/ncore/impl/data/v4/components_test.py @@ -16,10 +16,11 @@ import dataclasses import io import tempfile +import threading import types import unittest -from typing import Any, Dict, List, Literal, Tuple, cast +from typing import Any, Dict, List, Literal, Optional, Tuple, cast from unittest import mock import numpy as np @@ -31,7 +32,7 @@ from ncore.impl.common.transformations import HalfClosedInterval from ncore.impl.common.util import unpack_optional -from ncore.impl.data import stores +from ncore.impl.data import nodes, stores from ncore.impl.data.itar import IndexedTarArchive from ncore.impl.data.json_types import JsonDocument from ncore.impl.data.nodes import BloscCompression @@ -120,11 +121,13 @@ def setUp(self): @parameterized.expand( [ - (False,), - (True,), + (False, None), + (True, None), + (True, 0), # uncached nodes + (True, 1), # nodes evicted on most accesses ] ) - def test_reload(self, open_consolidated: bool): + def test_reload(self, open_consolidated: bool, node_cache_size: Optional[int]): """Test to make sure serialized data is faithfully reloaded""" tempdir = tempfile.TemporaryDirectory() @@ -606,7 +609,9 @@ def normalize_points(vectors: np.ndarray) -> Tuple[np.ndarray, np.ndarray]: store_paths += store_writer.finalize() ## Reload sequence and verify consistency - store_reader = SequenceComponentGroupsReader(store_paths, open_consolidated=open_consolidated) + store_reader = SequenceComponentGroupsReader( + store_paths, open_consolidated=open_consolidated, node_cache_size=node_cache_size + ) # check sequence data self.assertEqual(store_reader.sequence_id, ref_sequence_id) @@ -2929,6 +2934,129 @@ def test_schema_json_roundtrip(self) -> None: self.assertIsNone(rt.quantization) +class TestReaderOpening(unittest.TestCase): + """Component stores are opened concurrently, their hierarchies sequentially on the calling thread""" + + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + writer = SequenceComponentGroupsWriter( + output_dir_path=UPath(self._tmp.name), + store_base_name="seq", + sequence_id="seq", + sequence_timestamp_interval_us=HalfClosedInterval(0, 10_000_001), + generic_meta_data={}, + ) + for name in ("a", "b", "c"): + writer.register_component_writer(PosesComponent.Writer, name, f"poses_{name}").store_static_pose( + "x", "y", np.eye(4, dtype=np.float32) + ) + self.store_paths = writer.finalize() + + def tearDown(self): + self._tmp.cleanup() + + def test_hierarchies_are_opened_on_the_calling_thread(self): + threads = [] + open_root_group = stores.Store.open_root_group + + def recording_open_root_group(store: stores.Store, consolidated: bool = True) -> nodes.Group: + threads.append(threading.current_thread()) + return open_root_group(store, consolidated) + + with mock.patch.object(stores.Store, "open_root_group", autospec=True, side_effect=recording_open_root_group): + reader = SequenceComponentGroupsReader(self.store_paths, max_threads=4) + self.assertEqual(len(threads), len(self.store_paths)) + self.assertEqual(set(threads), {threading.current_thread()}) + self.assertEqual(sorted(reader.open_component_readers(PosesComponent.Reader)), ["a", "b", "c"]) + + def test_hierarchies_are_opened_while_stores_are_opened(self): + """The hierarchy of an opened store is opened while the remaining stores are still being opened""" + events: List[str] = [] + opened = threading.Event() + open_root_group = stores.Store.open_root_group + init = stores.IndexedTarZarrStore.__init__ + + def recording_init(store: stores.IndexedTarZarrStore, path: Any, *args: Any, **kwargs: Any) -> None: + if str(path).endswith(self.store_paths[-1].name): + # (the last store waits for the hierarchy of the first store to be opened) + self.assertTrue(opened.wait(timeout=10)) + init(store, path, *args, **kwargs) + events.append(f"store {UPath(path).name}") + + def recording_open_root_group(store: stores.Store, consolidated: bool = True) -> nodes.Group: + events.append("hierarchy") + opened.set() + return open_root_group(store, consolidated) + + with mock.patch.object(stores.IndexedTarZarrStore, "__init__", autospec=True, side_effect=recording_init): + with mock.patch.object( + stores.Store, "open_root_group", autospec=True, side_effect=recording_open_root_group + ): + SequenceComponentGroupsReader(self.store_paths, max_threads=len(self.store_paths)) + last = events.index(f"store {self.store_paths[-1].name}") + self.assertIn("hierarchy", events[:last]) + + +class TestReaderNodeCaches(unittest.TestCase): + """Component readers cache the opened nodes of their component (sized by node_cache_size)""" + + def setUp(self): + self._tmp = tempfile.TemporaryDirectory() + writer = SequenceComponentGroupsWriter( + output_dir_path=UPath(self._tmp.name), + store_base_name="seq", + sequence_id="seq", + sequence_timestamp_interval_us=HalfClosedInterval(0, 10_000_001), + generic_meta_data={}, + ) + lidar_writer = writer.register_component_writer(LidarSensorComponent.Writer, "lidar", "lidar") + for i in range(4): + direction = np.tile(np.array([1.0, 0.0, 0.0], dtype=np.float32), (5, 1)) + lidar_writer.store_frame( + direction=direction, + timestamp_us=np.full(5, 10 + i * 1_000_000, dtype=np.uint64), + model_element=None, + distance_m=np.full((1, 5), i, dtype=np.float32), + intensity=np.ones((1, 5), dtype=np.float32), + frame_timestamps_us=np.array([10 + i * 1_000_000, 20 + i * 1_000_000], dtype=np.uint64), + generic_data={}, + generic_meta_data={}, + ) + self.store_paths = writer.finalize() + + def tearDown(self): + self._tmp.cleanup() + + @parameterized.expand([(None,), (0,), (1,), (3,)]) + def test_node_cache_sizes(self, node_cache_size: Optional[int]): + reader = SequenceComponentGroupsReader(self.store_paths, node_cache_size=node_cache_size) + lidar = reader.open_component_readers(LidarSensorComponent.Reader)["lidar"] + frames = [int(ts) for ts in lidar.frames_timestamps_us[:, 1]] + for _ in range(2): + for i, ts in enumerate(frames): + np.testing.assert_array_equal(lidar.get_frame_ray_bundle_return_data(ts, "distance_m", 0), [i] * 5) + + cache = lidar._group.cache + if node_cache_size == 0: + self.assertIsNone(cache) + return + assert cache is not None + self.assertEqual(cache.max_entries, node_cache_size) + # every frame read opens the frame's returns group and its distance_m array + accessed = ["lidars/lidar/frames"] + [ + f"lidars/lidar/frames/{ts}/ray_bundle_returns{suffix}" for ts in frames for suffix in ("", "/distance_m") + ] + if node_cache_size is None: + self.assertEqual(sorted(cache.keys()), sorted(accessed)) + else: + # the most recently used nodes only + self.assertEqual(cache.keys(), accessed[-node_cache_size:]) + + def test_invalid_node_cache_size(self): + with self.assertRaises(ValueError): + SequenceComponentGroupsReader(self.store_paths, node_cache_size=-1) + + # ============================================================================== # Tests for zero-dimension array support # ==============================================================================