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 # ==============================================================================