post_dry_run: add caching to determine_global_extent workflow #1565

introduces caching on _extract_spatial_extent_from_load_stac_item_collection calls

involves introducing new caching utilities (GetOrCallCacheInterface)
and making some data structs cache key compatible (__hash__ + __eq__)

Author: soxofaan

openeogeotrellis/_backend/post_dry_run.py

@@ -17,6 +17,7 @@
 from openeogeotrellis.constants import EVAL_ENV_KEY
 
 import openeogeotrellis.load_stac
+from openeogeotrellis.util.caching import GetOrCallCache, GetOrCallCacheInterface, AlwaysCallWithoutCache
 from openeogeotrellis.util.geometry import BoundingBoxMerger
 from openeogeotrellis.util.math import logarithmic_round
 
@@ -57,14 +58,15 @@ def __init__(
     def __repr__(self) -> str:
         return f"_GridInfo(crs={self.crs_raw!r}, resolution={self.resolution!r}, extent_x={self.extent_x!r}, extent_y={self.extent_y!r})"
 
+    def _key(self) -> tuple:
+        return (self.crs_raw, self.resolution, self.extent_x, self.extent_y)
+
+    def __hash__(self):
+        return hash(self._key())
     def __eq__(self, other) -> bool:
-        return (
-            isinstance(other, _GridInfo)
-            and self.crs_raw == other.crs_raw
-            and self.resolution == other.resolution
-            and self.extent_x == other.extent_x
-            and self.extent_y == other.extent_y
-        )
+        if isinstance(other, _GridInfo):
+            return self._key() == other._key()
+        return NotImplemented
 
     @classmethod
     def from_datacube_metadata(cls, metadata: dict) -> _GridInfo:
@@ -197,7 +199,10 @@ class AlignedExtentResult:
 
 
 def _extract_spatial_extent_from_constraint(
-    source_constraint: SourceConstraint, *, catalog: AbstractCollectionCatalog
+    source_constraint: SourceConstraint,
+    *,
+    catalog: AbstractCollectionCatalog,
+    cache: Optional[GetOrCallCacheInterface] = None,
 ) -> Union[None, AlignedExtentResult]:
     """
     Extract spatial extent from given source constraint (if any), and align it to target grid.
@@ -210,18 +215,22 @@ def _extract_spatial_extent_from_constraint(
     if source_process == "load_collection":
         collection_id = source_id[1][0]
         return _extract_spatial_extent_from_constraint_load_collection(
-            collection_id=collection_id, constraint=constraint, catalog=catalog
+            collection_id=collection_id, constraint=constraint, catalog=catalog, cache=cache
         )
     elif source_process == "load_stac":
         url = source_id[1][0]
-        return _extract_spatial_extent_from_constraint_load_stac(stac_url=url, constraint=constraint)
+        return _extract_spatial_extent_from_constraint_load_stac(stac_url=url, constraint=constraint, cache=cache)
     else:
         # TODO?
         return None
 
 
 def _extract_spatial_extent_from_constraint_load_collection(
-    collection_id: str, *, constraint: dict, catalog: AbstractCollectionCatalog
+    collection_id: str,
+    *,
+    constraint: dict,
+    catalog: AbstractCollectionCatalog,
+    cache: Optional[GetOrCallCacheInterface] = None,
 ) -> Union[None, AlignedExtentResult]:
     try:
         metadata = catalog.get_collection_metadata(collection_id)
@@ -232,7 +241,7 @@ def _extract_spatial_extent_from_constraint_load_collection(
         stac_url = deep_get(metadata, "_vito", "data_source", "url")
         load_stac_feature_flags = deep_get(metadata, "_vito", "data_source", "load_stac_feature_flags", default={})
         return _extract_spatial_extent_from_constraint_load_stac(
-            stac_url=stac_url, constraint=constraint, feature_flags=load_stac_feature_flags
+            stac_url=stac_url, constraint=constraint, feature_flags=load_stac_feature_flags, cache=cache
         )
 
     # TODO Extracting pixel grid info from collection metadata might might be unreliable
@@ -271,7 +280,11 @@ def _extract_spatial_extent_from_constraint_load_collection(
 
 
 def _extract_spatial_extent_from_constraint_load_stac(
-    stac_url: str, *, constraint: dict, feature_flags: Optional[dict] = None
+    stac_url: str,
+    *,
+    constraint: dict,
+    feature_flags: Optional[dict] = None,
+    cache: Optional[GetOrCallCacheInterface] = None,
 ) -> Union[None, AlignedExtentResult]:
     spatial_extent_from_pg = constraint.get("spatial_extent") or constraint.get("weak_spatial_extent")
     spatiotemporal_extent = openeogeotrellis.load_stac._spatiotemporal_extent_from_load_params(
@@ -292,12 +305,23 @@ def _extract_spatial_extent_from_constraint_load_stac(
 
     pixel_buffer_size = deep_get(constraint, "pixel_buffer", "buffer_size", default=None)
 
-    return _extract_spatial_extent_from_load_stac_item_collection(
-        stac_url=stac_url,
-        spatiotemporal_extent=spatiotemporal_extent,
-        property_filter_pg_map=property_filter_pg_map,
-        resample_grid=resample_grid,
-        pixel_buffer_size=pixel_buffer_size,
+    return (cache or AlwaysCallWithoutCache()).get_or_call(
+        key=(
+            stac_url,
+            spatiotemporal_extent,
+            str(property_filter_pg_map),
+            resample_grid,
+            pixel_buffer_size,
+            str(feature_flags),
+        ),
+        callback=lambda: _extract_spatial_extent_from_load_stac_item_collection(
+            stac_url=stac_url,
+            spatiotemporal_extent=spatiotemporal_extent,
+            property_filter_pg_map=property_filter_pg_map,
+            resample_grid=resample_grid,
+            pixel_buffer_size=pixel_buffer_size,
+            feature_flags=feature_flags,
+        ),
     )
 
 
@@ -308,6 +332,7 @@ def _extract_spatial_extent_from_load_stac_item_collection(
     property_filter_pg_map: Optional[openeogeotrellis.load_stac.PropertyFilterPGMap] = None,
     resample_grid: Optional[_GridInfo] = None,
     pixel_buffer_size: Union[Tuple[float, float], int, float, None] = None,
+    feature_flags: Optional[dict] = None,
 ) -> Union[None, AlignedExtentResult]:
     extent_orig: Union[BoundingBox, None] = spatiotemporal_extent.spatial_extent.as_bbox()
     extent_variants = {"original": extent_orig}
@@ -429,6 +454,7 @@ def determine_global_extent(
     *,
     source_constraints: List[SourceConstraint],
     catalog: AbstractCollectionCatalog,
+    cache: Optional[GetOrCallCacheInterface] = None,
 ) -> dict:
     """
     Go through all source constraints, extract the aligned extent from each (possibly with variations)
@@ -438,9 +464,16 @@ def determine_global_extent(
     #       e.g. add stats to AlignedExtentResult for better informed decision?
     aligned_merger = BoundingBoxMerger()
     variant_mergers: Dict[str, BoundingBoxMerger] = collections.defaultdict(BoundingBoxMerger)
+    # Enable the implementation to use caching
+    # when going through all these source constraints
+    # as there might be duplication that's hidden
+    # by differences from non-relevant details
+    cache = cache or GetOrCallCache(max_size=100)
     for source_id, constraint in source_constraints:
         try:
-            aligned_extent_result = _extract_spatial_extent_from_constraint((source_id, constraint), catalog=catalog)
+            aligned_extent_result = _extract_spatial_extent_from_constraint(
+                (source_id, constraint), catalog=catalog, cache=cache
+            )
         except Exception as e:
             raise SpatialExtentExtractionError(
                 f"Failed to extract spatial extent from {source_id=} with {constraint=}: {e=}"
@@ -455,6 +488,7 @@ def determine_global_extent(
                     _log.info(
                         f"determine_global_extent: skipping {name=} from {source_id=} with {aligned_extent_result.variants=}"
                     )
+    _log.info(f"_extract_spatial_extent_from_constraint cache stats: {cache!s}")
 
     global_extent: BoundingBox = aligned_merger.get()
     global_extent_variants: Dict[str, BoundingBox] = {name: merger.get() for name, merger in variant_mergers.items()}

openeogeotrellis/load_stac.py

@@ -931,13 +931,26 @@ class _TemporalExtent:
     """
 
     # TODO: move this to a more generic location for better reuse
+    # TODO: enforce/ensure immutability
+    # TODO: re-implement in dataclasses/attrs to better enforce immutability and simplify equality/hash implementation
 
     __slots__ = ("from_date", "to_date")
 
     def __init__(self, from_date: DateTimeLikeOrNone, to_date: DateTimeLikeOrNone):
         self.from_date: Union[datetime.datetime, None] = to_datetime_utc_unless_none(from_date)
         self.to_date: Union[datetime.datetime, None] = to_datetime_utc_unless_none(to_date)
 
+    def _key(self) -> tuple:
+        return (self.from_date, self.to_date)
+
+    def __hash__(self):
+        return hash(self._key())
+
+    def __eq__(self, other):
+        if isinstance(other, _TemporalExtent):
+            return self._key() == other._key()
+        return NotImplemented
+
     def as_tuple(self) -> Tuple[Union[datetime.datetime, None], Union[datetime.datetime, None]]:
         return self.from_date, self.to_date
 
@@ -994,6 +1007,8 @@ class _SpatialExtent:
     """
 
     # TODO: move this to a more generic location for better reuse
+    # TODO: enforce/ensure immutability
+    # TODO: re-implement in dataclasses/attrs to better enforce immutability and simplify equality/hash implementation
 
     __slots__ = ("_bbox", "_bbox_lonlat_shape")
 
@@ -1003,6 +1018,17 @@ def __init__(self, *, bbox: Union[BoundingBox, None]):
         # cache for shapely polygon in lon/lat
         self._bbox_lonlat_shape = self._bbox.reproject("EPSG:4326").as_polygon() if self._bbox else None
 
+    def _key(self) -> tuple:
+        return (self._bbox,)
+
+    def __hash__(self):
+        return hash(self._key())
+
+    def __eq__(self, other):
+        if isinstance(other, _SpatialExtent):
+            return self._key() == other._key()
+        return NotImplemented
+
     def as_bbox(self, crs: Optional[str] = None) -> Union[BoundingBox, None]:
         bbox = self._bbox
         if bbox and crs:
@@ -1018,6 +1044,9 @@ def intersects(self, bbox: Union[List[float], Tuple[float, float, float, float],
 
 class _SpatioTemporalExtent:
     # TODO: move this to a more generic location for better reuse
+    # TODO: enforce/ensure immutability
+    # TODO: re-implement in dataclasses/attrs to better enforce immutability and simplify equality/hash implementation
+
     def __init__(
         self,
         *,
@@ -1028,6 +1057,17 @@ def __init__(
         self._spatial_extent = _SpatialExtent(bbox=bbox)
         self._temporal_extent = _TemporalExtent(from_date=from_date, to_date=to_date)
 
+    def _key(self) -> tuple:
+        return (self._spatial_extent, self._temporal_extent)
+
+    def __hash__(self):
+        return hash(self._key())
+
+    def __eq__(self, other):
+        if isinstance(other, _SpatioTemporalExtent):
+            return self._key() == other._key()
+        return NotImplemented
+
     @property
     def spatial_extent(self) -> _SpatialExtent:
         return self._spatial_extent

openeogeotrellis/util/caching.py

@@ -0,0 +1,100 @@
+import collections
+
+from typing import Dict, Any, Optional, Callable
+
+
+class GetOrCallCacheInterface:
+    """
+    Interface for "memoizing" function calls with an explicit cache key.
+
+    Unlike the well known `functools.lru` cache,
+    it requires to specify an explicit cache key,
+    which gives more control over how to handle arguments, e.g.:
+    - serialize or transform arguments that are not hashable by default
+    - ignore cetain arguments
+    - add cache key components that are not arguments to the function
+
+    Usage example:
+
+         result = cache.get_or_call(
+            key=(name, str(sorted(options.items())),
+            callable=lambda : function_to_cache(name, options)
+        )
+
+    """
+
+    def get_or_call(self, key, callback: Callable[[], Any]) -> Any:
+        """
+        Try to get item from cache.
+        If not available: call callable to build it and store result in cache.
+
+        :param key: key to store item at (can be a simple string,
+            or something more complex like a tuple of strings/ints)
+        :param callback: item builder to call when item is not in cache
+        :return: item (from cache or freshly built)
+        """
+        raise NotImplementedError()
+
+
+class AlwaysCallWithoutCache(GetOrCallCacheInterface):
+    """Don't cache, just call"""
+
+    def get_or_call(self, key, callback: Callable[[], Any]) -> Any:
+        return callback()
+
+
+class GetOrCallCache(GetOrCallCacheInterface):
+    """
+    In-memory dictionary based cache for "memoizing" function calls.
+
+    Supports a maximum size by pruning least frequently used items.
+    """
+
+    # TODO: support least recently used (LRU) pruning strategy as well?
+
+    def __init__(self, max_size: Optional[int] = None):
+        self._cache = {}
+        self._usage_stats = {}
+        self._max_size = max_size
+        self._stats = collections.defaultdict(int)
+
+    def get_or_call(self, key, callback: Callable[[], Any]) -> Any:
+        self._stats["get_or_call"] += 1
+        if key in self._cache:
+            self._usage_stats[key] += 1
+            self._stats["hit"] += 1
+            return self._cache[key]
+        else:
+            value = callback()
+            if self._ensure_capacity_for(count=1):
+                self._cache[key] = value
+                self._usage_stats[key] = 1
+            self._stats["miss"] += 1
+            return value
+
+    def _ensure_capacity_for(self, count: int = 1) -> bool:
+        """
+        Make sure there is room (according to max_size setting)
+        for adding the given number of new items
+        by pruning least used items from the cache if necessary.
+        """
+        if self._max_size is not None:
+            to_prune = len(self._cache) + count - self._max_size
+            if to_prune > 0:
+                least_used_keys = sorted(self._usage_stats, key=self._usage_stats.get)[:to_prune]
+                for key in least_used_keys:
+                    del self._cache[key]
+                    del self._usage_stats[key]
+                    self._stats["prune"] += 1
+            return len(self._cache) + count <= self._max_size
+        else:
+            return True
+
+    def stats(self) -> dict:
+        return dict(self._stats)
+
+    def __str__(self):
+        return f"<GetOrCallCachesize size={len(self._cache)} stats={dict(self._stats)}>"
+
+    def clear(self):
+        self._cache.clear()