diff --git a/deeptrack/features.py b/deeptrack/features.py index feb29ad7f..54174bc56 100644 --- a/deeptrack/features.py +++ b/deeptrack/features.py @@ -154,11 +154,17 @@ from __future__ import annotations +import contextlib +import contextvars import itertools import operator import random +import threading import warnings -from typing import Any, Callable, Iterable, Literal, TYPE_CHECKING +import weakref +from typing import ( + Any, Callable, Iterable, Iterator, Literal, TYPE_CHECKING, +) import array_api_compat as apc import numpy as np @@ -167,7 +173,7 @@ from pint import Quantity from deeptrack.backend import config, TORCH_AVAILABLE, xp -from deeptrack.backend.core import DeepTrackNode +from deeptrack.backend.core import DeepTrackDataObject, DeepTrackNode from deeptrack.backend.units import ConversionTable from deeptrack.properties import Property, PropertyDict, SequentialProperty from deeptrack.sources import SourceItem @@ -240,6 +246,68 @@ MERGE_STRATEGY_APPEND: int = 1 +# Resolution bookkeeping. +# +# `_RESOLUTION_DEPTH` counts the nested `Feature.__call__()` frames currently +# executing, and `_RESOLUTION_EPOCH` holds the epoch of the current top-level +# (non-nested) call. Together they let a feature tell whether an input it +# received explicitly was set during the current resolution or during a +# previous one, so that a stale input is not silently reused when the feature +# is later called without input (e.g. when evaluated as a property). +# +# Both are context variables, so independent threads (or asyncio tasks) each +# have their own depth and epoch and do not see each other's calls as nested. +# Epoch numbers come from one process-wide counter, so they never collide +# across contexts. +_RESOLUTION_DEPTH: contextvars.ContextVar[int] = contextvars.ContextVar( + "_RESOLUTION_DEPTH", default=0 +) +_RESOLUTION_EPOCH: contextvars.ContextVar[int] = contextvars.ContextVar( + "_RESOLUTION_EPOCH", default=0 +) +_EPOCH_COUNTER = itertools.count(1) +_EPOCH_LOCK = threading.Lock() + + +@contextlib.contextmanager +def _resolution_scope() -> Iterator[None]: + """Track nested `Feature.__call__()` frames and resolution epochs.""" + + depth = _RESOLUTION_DEPTH.get() + epoch_token = None + if depth == 0: + with _EPOCH_LOCK: + epoch = next(_EPOCH_COUNTER) + epoch_token = _RESOLUTION_EPOCH.set(epoch) + depth_token = _RESOLUTION_DEPTH.set(depth + 1) + try: + yield + finally: + _RESOLUTION_DEPTH.reset(depth_token) + if epoch_token is not None: + _RESOLUTION_EPOCH.reset(epoch_token) + + +def _cache_ID( + node: DeepTrackNode, + _ID: tuple[int, ...], +) -> tuple[int, ...]: + """Return the `_ID` under which `node`'s cache holds the entry for `_ID`. + + A cache resolves a read at an `_ID` longer than its key length by + trimming it, but rejects a write there. Writing at the trimmed `_ID` + instead lands in the entry reads come from. An `_ID` that fits the key + length, or a cache without key length yet, is returned unchanged. + + """ + + keylength = node.data._keylength + if keylength is not None and len(_ID) > keylength: + return _ID[:keylength] + + return _ID + + class Feature(DeepTrackNode): """Base feature class. @@ -683,6 +751,15 @@ def __init__( self._input = DeepTrackNode(node_name="_input", action=_input) self._input.add_child(self) + # Maps each cache cell of `_input` whose value was set explicitly by + # `.__call__()` to the resolution epoch in which it was set (see + # `_RESOLUTION_EPOCH`). Keyed by cell rather than by `_ID`, so that + # `_ID`s the cache resolves to the same cell share one record, and so + # that records disappear with their cells when the cache is reset. + self._explicit_input_epochs: weakref.WeakKeyDictionary[ + DeepTrackDataObject, int + ] = weakref.WeakKeyDictionary() + # Random seed node (for deterministic behavior if desired). self._random_seed = DeepTrackNode( node_name="_random_seed", @@ -750,8 +827,14 @@ def __call__( data_list: Any, optional The input data to the feature or pipeline. It is most commonly a list of NumPy arrays or PyTorch tensors, but it can be anything. - Defaults to `None`, in which case the feature uses the previous set - of input values or propagates properties. + Defaults to `None`, in which case the feature is evaluated without + input. An input set explicitly by a previous call is not reused: + it is reset so that the feature falls back to its default input + (an empty list, or the output of the feature wired as `_input`). + Within a single resolution, however, a feature that received an + input and is then evaluated again without input (for example as + the property of another feature) keeps that input, so that it + yields one consistent value throughout the pipeline. **kwargs: Any Additional parameters passed to the pipeline. These override properties with matching names. For example, calling @@ -781,10 +864,11 @@ def __call__( >>> feature(np.array([1, 2, 3])) array([3, 4, 5]) - Execute the feature with previously set input: + Execute the feature without input. The previously set input is + reset, so the feature is applied to an empty input list: - >>> feature() # Uses stored input - array([3, 4, 5]) + >>> feature() + [] Execute the feature with new input: @@ -815,13 +899,46 @@ def _should_set_input(value: Any) -> bool: return True - with config.with_backend(self._backend): + def _is_empty_input(value: Any) -> bool: + return value is None or ( + isinstance(value, list) and len(value) == 0 + ) + + with _resolution_scope(), config.with_backend(self._backend): + epoch = _RESOLUTION_EPOCH.get() + # If data_list is a Source, activate it. self._activate_sources(data_list) - # If the input is not empty, set the value of the input. + # `_ID` is the logical identifier, used for properties and + # arguments. The `_input` cache may resolve it to a shorter + # entry; see `_cache_ID()`. + input_ID = _cache_ID(self._input, _ID) + + # If the input is not empty, set the value of the input and + # remember in which epoch it was set. if _should_set_input(data_list): - self._input.set_value(data_list, _ID=_ID) + self._input.set_value(data_list, _ID=input_ID) + for cell in self._input_cells(input_ID): + self._explicit_input_epochs[cell] = epoch + + # If no input is given but an input was set explicitly during a + # previous resolution, that input is stale: reset it so that the + # feature falls back to its default input (an empty list, or the + # output of the feature wired as `_input`). An input set during + # the current resolution is kept, so that a feature evaluated + # both with input and as a property in the same pipeline yields + # a single consistent value. + elif _is_empty_input(data_list): + stale_cells = [ + cell + for cell in self._input_cells(input_ID) + if self._explicit_input_epochs.get(cell, epoch) != epoch + ] + if stale_cells: + for cell in stale_cells: + del self._explicit_input_epochs[cell] + self._input.invalidate(_ID=input_ID) # A dict to store values of self.arguments before updating them. original_values: dict[str, Any] = {} @@ -847,7 +964,7 @@ def _should_set_input(value: Any) -> bool: # This executes the feature. # DeepTrackNode will determine if it needs to be recalculated. # If it does, it will call the `.action()` method. - output = super().__call__(_ID=_ID) + output = super().__call__(_ID=_cache_ID(self, _ID)) finally: # If there are self.arguments, reset the values @@ -862,6 +979,27 @@ def _should_set_input(value: Any) -> bool: resolve = __call__ + def _input_cells( + self: Feature, + _ID: tuple[int, ...] = (), + ) -> list[DeepTrackDataObject]: + """Return the cache cells of `_input` that `_ID` refers to. + + An `_ID` shorter than the cache key length matches several cells by + prefix; an empty cache yields no cell. + + """ + + try: + cells = self._input.data[_ID] + except KeyError: + return [] + + if isinstance(cells, dict): + return list(cells.values()) + + return [cells] + def to_sequential( self: Feature, **kwargs: Any, @@ -1505,7 +1643,7 @@ def _action( """ # Retrieve the input images. - inputs = self._input(_ID=_ID) + inputs = self._input(_ID=_cache_ID(self._input, _ID)) # Get the current property values. properties_copy = self.properties(_ID=_ID).copy() @@ -5552,20 +5690,12 @@ class Stack(Feature): Note ---- - If a feature is called directly, its result is cached internally. This can - affect how it behaves when reused in chained pipelines. For example: - - >>> stack_feature = dt.Stack(value=2) - >>> _ = stack_feature(1) # Evaluate the feature and cache the output - >>> (1 & stack_feature)() - [1, 1, 2] - - To ensure consistent behavior when reusing a feature after calling it, - reset its state using instead: + An input given to a feature by a direct call is not reused when the + feature is later evaluated without input, for example inside a chained + pipeline: >>> stack_feature = dt.Stack(value=2) - >>> _ = stack_feature(1) - >>> stack_feature.update() # clear cached state + >>> _ = stack_feature(1) # Evaluate the feature with input 1 >>> (1 & stack_feature)() [1, 2] diff --git a/tests/test_features.py b/tests/test_features.py index f4fb48b48..59320b3bb 100644 --- a/tests/test_features.py +++ b/tests/test_features.py @@ -8,6 +8,8 @@ import itertools import operator +import gc +import threading import unittest import warnings @@ -1389,6 +1391,210 @@ def test_Feature_arithmetic(self): input_2 = [10, 20] self.assertListEqual(pipeline(input_2), [-input_2[0], -input_2[1]]) + def test_Feature_input_reset_after_direct_call(self): + # Issue 352: an input set by a previous direct call must not leak + # into a later evaluation without input. + + stack_feature = features.Stack(value=2) + stack_feature(3) + self.assertEqual(stack_feature(), 2) + + stack_feature = features.Stack(value=2) + stack_feature(1) + self.assertListEqual((1 & stack_feature)(), [1, 2]) + + # Also when the feature is nested in a pipeline. + stack_feature = features.Stack(value=2) + pipeline = stack_feature >> features.Add(b=1) + pipeline(5) + self.assertEqual(pipeline(), 3) + + # A new input still invalidates as before. + stack_feature = features.Stack(value=2) + stack_feature(3) + self.assertListEqual(stack_feature(2), [2, 2]) + + # A feature evaluated with input and as a property within the same + # resolution keeps a single consistent value. + inp = features.DummyFeature() + pipeline = inp - inp * 2 + self.assertEqual(pipeline(10), -10) + self.assertEqual(pipeline(10), -10) + + def test_Feature_input_reset_recomputes_once(self): + # Sequence: explicit input, empty call (recompute on the default + # input), empty call (cached). `get` is counted directly, so the + # test does not rely on property caching. + + calls = [] + + class Counting(features.Feature): + __distributed__ = False + + def get(self, inputs, **kwargs): + calls.append(list(inputs)) + return inputs + + feature = Counting() + self.assertListEqual(feature([1, 2]), [1, 2]) + self.assertListEqual(calls, [[1, 2]]) + + self.assertListEqual(feature(), []) + self.assertListEqual(calls, [[1, 2], []]) + + self.assertListEqual(feature(), []) + self.assertListEqual(calls, [[1, 2], []]) + + # A reducer wired to a feature falls back to that feature's output + # after an explicit input. + from deeptrack import statistics + + inp = features.Value([1, 0]) + mean = statistics.Mean(inp) + self.assertEqual(mean(5), 5.0) + self.assertListEqual(list(mean()), [1.0, 0.0]) + + def test_Feature_input_reset_respects_ID_aliases(self): + # `_ID`s longer than the cache key length are trimmed to the same + # cell: the stale check must follow the cache, not the raw tuple. + + stack_feature = features.Stack(value=2) + stack_feature(3) + self.assertEqual(stack_feature(_ID=(0,)), 2) + # The feature stays usable at its original `_ID` after the reset. + self.assertEqual(stack_feature(), 2) + + stack_feature = features.Stack(value=2) + stack_feature(3) + self.assertEqual(features.Repeat(stack_feature, N=1)(), 2) + + # Public composition mixing a nested use (longer `_ID`) and a direct + # use (original `_ID`) of a feature called before. + stack_feature = features.Stack(value=2) + stack_feature(3) + pipeline = features.Repeat(stack_feature, N=1) & stack_feature + self.assertListEqual(pipeline(), [2, 2]) + + # Round trips between `_ID` lengths: a longer `_ID` writes to the + # trimmed entry, so it neither fails nor changes the key length. + stack_feature = features.Stack(value=2) + stack_feature(3) + self.assertListEqual(stack_feature(5, _ID=(0,)), [5, 2]) + self.assertEqual(stack_feature(), 2) + self.assertListEqual(stack_feature(6), [6, 2]) + self.assertEqual(stack_feature(_ID=(0, 1)), 2) + self.assertEqual(stack_feature(), 2) + + # Two aliases of one cell within the same resolution share the + # input set during that resolution. + stack_feature = features.Stack(value=2) + stack_feature(3) + probe = features.Value( + lambda: (stack_feature(3, _ID=(0,)), stack_feature()) + ) + self.assertEqual(probe(), ([3, 2], [3, 2])) + + # Distinct cells are reset independently: resetting (0,) leaves the + # input stored at (1,) untouched. + stack_feature = features.Stack(value=2) + stack_feature(3, _ID=(0,)) + stack_feature(4, _ID=(1,)) + self.assertEqual(stack_feature(_ID=(0,)), 2) + self.assertTrue(stack_feature._input.is_valid(_ID=(1,))) + self.assertEqual(stack_feature._input(_ID=(1,)), 4) + + # A stale sibling reached through an alias does not take an input + # received during the current resolution down with it. + stack_feature = features.Stack(value=2) + stack_feature(3, _ID=(0,)) + probe = features.Value( + lambda: ( + stack_feature(4, _ID=(1,)), + stack_feature(_ID=(0, 99)), + stack_feature(_ID=(1, 99)), + ) + ) + self.assertEqual(probe(), ([4, 2], 2, [4, 2])) + + def test_Feature_cache_ID_trimming_keeps_logical_ID(self): + # Only the caches that still hold a shorter key length see a trimmed + # `_ID`; properties keep the logical one, so distinct `_ID`s sample + # independently even when `_input` kept key length zero. + + counter = iter(range(10)) + feature = features.Value(lambda: next(counter)) + self.assertEqual(feature(), 0) + feature.value.update() + self.assertEqual((feature(_ID=(0,)), feature(_ID=(1,))), (1, 2)) + + feature = features.Value(lambda _ID: _ID) + feature() + feature.value.update() + self.assertEqual(feature(_ID=(7,)), (7,)) + + def test_Feature_input_reset_metadata_purged_on_update(self): + # Records of explicit inputs live with the cache cells, so resetting + # the cache through `update()` drops them for every component. + + stack_feature = features.Stack(value=2) + pipeline = features.Value(1) >> stack_feature + for i in range(4): + pipeline(_ID=(i,)) + self.assertEqual(len(stack_feature._explicit_input_epochs), 4) + + pipeline.update() + gc.collect() + self.assertEqual(len(stack_feature._explicit_input_epochs), 0) + self.assertEqual(len(pipeline._explicit_input_epochs), 0) + + def test_Feature_input_reset_is_thread_isolated(self): + # A call blocked in another thread must not make this thread's + # top-level calls look nested (which would share one epoch and keep + # the stale input). + + entered = threading.Event() + release = threading.Event() + + class Blocking(features.Feature): + def get(self, inputs, **kwargs): + entered.set() + release.wait(timeout=10) + return inputs + + result = {} + errors = [] + + def guarded(target, *args): + def run(): + try: + target(*args) + except Exception as error: # noqa: BLE001 + errors.append(error) + return run + + def worker(): + stack_feature = features.Stack(value=2) + stack_feature(3) + result["value"] = stack_feature() + + blocking = Blocking() + blocker = threading.Thread(target=guarded(blocking, 1)) + thread = threading.Thread(target=guarded(worker)) + try: + blocker.start() + self.assertTrue(entered.wait(timeout=10)) + thread.start() + thread.join(timeout=10) + finally: + release.set() + blocker.join(timeout=10) + + self.assertFalse(thread.is_alive()) + self.assertFalse(blocker.is_alive()) + if errors: + raise errors[0] + self.assertEqual(result["value"], 2) + def test_Features_chain_lambda(self): value = features.Value(value=1)