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
 
@@ -198,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.
@@ -211,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)
@@ -233,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
@@ -272,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(
@@ -293,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,
+        ),
     )
 
 
@@ -309,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}
@@ -430,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)
@@ -439,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=}"
@@ -456,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/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()

tests/_backend/test_post_dry_run.py

@@ -1,3 +1,6 @@
+import collections
+import dirty_equals
+
 from typing import Callable, List, Optional
 
 import openeo_driver.ProcessGraphDeserializer
@@ -19,7 +22,7 @@
     determine_global_extent,
 )
 from openeogeotrellis.load_stac import _ProjectionMetadata
-
+from openeogeotrellis.util.caching import AlwaysCallWithoutCache
 
 class TestGridInfo:
     def test_minimal(self):
@@ -979,6 +982,53 @@ def test_determine_global_extent_load_stac_minimal(
             },
         }
 
+    @pytest.mark.parametrize(
+        ["cache", "expected_searches"],
+        [
+            (None, 1),  # cache=None: use default caching
+            (AlwaysCallWithoutCache(), 3),
+        ],
+    )
+    def test_determine_global_extent_load_stac_caching(
+        self,
+        dummy_catalog,
+        extract_source_constraints,
+        dummy_stac_api_server,
+        dummy_stac_api,
+        cache,
+        expected_searches,
+    ):
+        # Multiple source constraints with same essentials,
+        # and some differences in non-essential details
+        source_constraints = [
+            (
+                ("load_stac", (f"{dummy_stac_api}/collections/collection-123", (), ("B02",))),
+                {"bands": ["B02"], "breakfast": "cereal"},
+            ),
+            (
+                ("load_stac", (f"{dummy_stac_api}/collections/collection-123", (), ("B02",))),
+                {"bands": ["B02"], "shoe:size": 42},
+            ),
+            (
+                ("load_stac", (f"{dummy_stac_api}/collections/collection-123", (), ("B02",))),
+                {"bands": ["B02"], "pets": ["dog"]},
+            ),
+        ]
+        global_extent = determine_global_extent(
+            source_constraints=source_constraints, catalog=dummy_catalog, cache=cache
+        )
+        expected = BoundingBox(west=2, south=49, east=7, north=52, crs="EPSG:4326").approx(abs=1e-6)
+        assert global_extent == {
+            "global_extent": expected,
+            "global_extent_variants": {
+                "assets_full_bbox": expected,
+            },
+        }
+        # Check request history histogram by path for search requests
+        assert collections.Counter(
+            (r["method"], r["path"]) for r in dummy_stac_api_server.request_history
+        ) == dirty_equals.IsPartialDict({("GET", "/search"): expected_searches})
+
     def test_extract_spatial_extent_from_constraint_load_collection_type_stac_minimal(
         self, dummy_catalog, extract_source_constraints, dummy_stac_api
     ):