Add analysis and caching

Author: JBorrow

tilemaker/analysis/core.py

@@ -0,0 +1,84 @@
+"""
+Analysis caching and producing services.
+"""
+
+import uuid
+from abc import ABC, abstractmethod
+
+import structlog
+from structlog.types import FilteringBoundLogger
+
+from tilemaker.metadata.definitions import MapGroup
+from tilemaker.providers.core import Tiles
+
+from .types import SLUG_TO_TYPE
+
+
+class ProductNotFoundError(Exception):
+    pass
+
+
+class AnalysisProvider(ABC):
+    internal_provider_id: str
+    logger: FilteringBoundLogger
+
+    def __init__(self, internal_provider_id: str | None):
+        self.internal_provider_id = internal_provider_id or str(uuid.uuid4())
+        self.logger = structlog.get_logger()
+
+    @abstractmethod
+    def pull(self, analysis_id: str, grants: set[str]):
+        return
+
+    @abstractmethod
+    def push(self, product: "AnalysisProduct"):
+        return
+
+
+class Analyses:
+    pullable: list[AnalysisProvider]
+    pushable: list[AnalysisProvider]
+    tiles: Tiles
+    metadata: list[MapGroup]
+
+    def __init__(
+        self,
+        pullable: list[AnalysisProvider],
+        pushable: list[AnalysisProvider],
+        tiles: Tiles,
+        metadata: list[MapGroup],
+    ):
+        self.pullable = pullable
+        self.pushable = pushable
+        self.tiles = tiles
+        self.metadata = metadata
+
+    def pull(self, analysis_id: str, grants: set[str]) -> "AnalysisProduct":
+        for provider in self.pullable:
+            try:
+                product = provider.pull(analysis_id=analysis_id, grants=grants)
+
+                if product.grant and product.grant not in grants:
+                    raise ProductNotFoundError(f"Product {analysis_id} not found")
+            except ProductNotFoundError:
+                continue
+
+        # We couldn't find it. Try building it?
+        for slug, analysis_type in SLUG_TO_TYPE.items():
+            if slug in analysis_id:
+                product = analysis_type.build(
+                    tiles=self.tiles, metadata=self.metadata, analysis_id=analysis_id
+                )
+
+                self.push(product)
+
+            if product.grant and product.grant not in grants:
+                raise ProductNotFoundError(f"Product {analysis_id} not found")
+
+        return product
+
+    def push(self, product: "AnalysisProduct"):
+        for provider in self.pushable:
+            provider.push(product)
+
+        return

tilemaker/analysis/histogram.py

@@ -0,0 +1,82 @@
+"""
+Histogram generation and storage.
+"""
+
+from time import perf_counter
+
+import numpy as np
+import structlog
+
+from tilemaker.metadata.definitions import MapGroup
+from tilemaker.providers.core import PullableTile, TileNotFoundError, Tiles
+
+from .products import AnalysisProduct
+
+
+class HistogramProduct(AnalysisProduct):
+    layer_id: str
+
+    counts: list[int]
+    edges: list[float]
+
+    @property
+    def hash(self):
+        return f"hist-{self.layer_id}"
+
+    @classmethod
+    def build(
+        cls, tiles: Tiles, metadata: list[MapGroup], analysis_id: str
+    ) -> "HistogramProduct":
+        log = structlog.get_logger()
+
+        layer_id = analysis_id.replace("hist-", "")
+
+        log = log.bind(layer_id=layer_id)
+
+        for map_group in metadata:
+            layer = map_group.get_layer(layer_id)
+            if layer is not None:
+                break
+
+        if layer is None:
+            log.info("histogram.layer_not_found")
+            raise TileNotFoundError(f"Layer {layer_id} not found")
+
+        timing_start = perf_counter()
+
+        start = layer.vmin * 4
+        end = layer.vmax * 4
+        bins = 128
+
+        log = log.bind(start=start, end=end, bins=bins)
+
+        edges = np.linspace(start, end, bins + 1)
+        counts = np.zeros(bins)
+
+        for tile_x in [0, 1]:
+            tile, pushable = tiles.pull(
+                PullableTile(
+                    layer_id=layer_id,
+                    x=tile_x,
+                    y=0,
+                    level=0,
+                    # Bypass auth for this generation process
+                    grants=set(layer.grant) if layer.grant is not None else None,
+                )
+            )
+
+            tiles.push(pushable)
+
+            if tile.data is not None:
+                counts += np.histogram(tile.data, bins=edges)[0]
+
+        timing_end = perf_counter()
+        log = log.bind(dt=timing_end - timing_start)
+        log.info("histogram.built")
+
+        return cls(
+            layer_id=layer_id,
+            counts=counts,
+            edges=edges,
+            grant=layer.grant,
+        )

tilemaker/analysis/products.py

@@ -0,0 +1,19 @@
+from abc import ABC, abstractmethod
+
+from pydantic import BaseModel
+
+
+class AnalysisProduct(BaseModel, ABC):
+    layer_id: str
+    grant: str | None
+    source: str | None = None
+
+    @property
+    @abstractmethod
+    def hash(self):
+        return
+
+    @classmethod
+    @abstractmethod
+    def build(cls, tiles: "Tiles", metadata: list["MapGroup"], layer_id: str):
+        return

tilemaker/analysis/providers.py

@@ -0,0 +1,89 @@
+from cachetools import LFUCache
+from pymemcache.client.base import Client
+
+from .core import AnalysisProvider, ProductNotFoundError
+from .products import AnalysisProduct
+from .types import AnalysisType
+
+
+class InMemoryAnalysisCache(AnalysisProvider):
+    """
+    A simple in-memory cache for tiles.
+    """
+
+    cache: LFUCache
+
+    def __init__(self, cache_size: int = 8192, internal_provider_id: str | None = None):
+        self.cache = LFUCache(maxsize=cache_size)
+        super().__init__(internal_provider_id=internal_provider_id)
+
+    def pull(self, analysis_id: str, grants: set[str]):
+        log = self.logger.bind(analysis_id=analysis_id)
+
+        cached = self.cache.get(analysis_id, None)
+
+        if cached is None:
+            log.debug("inmemory.miss")
+            raise ProductNotFoundError(f"Product {analysis_id} not found in cache")
+
+        if cached.grant and cached.grant not in grants:
+            log = log.bind(product_grant=cached.grant, user_grants=grants)
+            log.debug("inmemory.proprietary_hidden")
+            raise ProductNotFoundError(f"Product {analysis_id} not found in cache")
+
+        log.debug("inmemory.pulled")
+        return cached
+
+    def push(self, product: AnalysisType):
+        log = self.logger.bind(analysis_id=product.hash)
+
+        if product.source == self.internal_provider_id:
+            log.debug("inmemory.present")
+
+        product.source = self.internal_provider_id
+        self.cache[product.hash] = product
+        log.debug("inmemory.pushed")
+
+
+class MemcachedAnalysisCache(AnalysisProvider):
+    """
+    A cache that uses Memcached for storing tiles.
+    """
+
+    client: Client
+
+    def __init__(self, client: Client, internal_provider_id: str | None = None):
+        self.client = client
+        super().__init__(
+            internal_provider_id=internal_provider_id or "memcached-analysis"
+        )
+
+    def pull(self, analysis_id: str, grants: set[str]):
+        log = self.logger.bind(analysis_id=analysis_id)
+
+        res = self.client.get(analysis_id, None)
+
+        if res is None:
+            log.debug("memcached.miss")
+            raise ProductNotFoundError(f"Product {analysis_id} not found in cache")
+
+        res = AnalysisType.model_validate_json(res)
+
+        if res.grant and res.grant not in grants:
+            log = log.bind(product_grant=res.grant, user_grants=grants)
+            log.debug("memcached.proprietary_hidden")
+            raise ProductNotFoundError(f"Product {analysis_id} not found in cache")
+
+        log.debug("memcached.pulled")
+
+        return res
+
+    def push(self, product: AnalysisProduct):
+        log = self.logger.bind(analysis_id=product.hash)
+
+        if product.source == self.internal_provider_id:
+            log.debug("memcached.present")
+
+        product.source = self.internal_provider_id
+        self.client.set(product.hash, product.model_dump_json(), noreply=True)
+        log.debug("memcached.pushed")

tilemaker/analysis/types.py

@@ -0,0 +1,5 @@
+from .histogram import HistogramProduct
+
+SLUG_TO_TYPE = {"hist": HistogramProduct}
+
+AnalysisType = HistogramProduct

tilemaker/metadata/definitions.py

@@ -71,18 +71,25 @@
 from pathlib import Path
 from typing import Literal
 
+import structlog
 from astropy import units
 from astropy.io import fits
 from astropy.wcs import WCS
 from pydantic import BaseModel, RootModel
 
 
 def parse_config(config: Path) -> list["MapGroup"]:
+    log = structlog.get_logger()
+    log = log.bind(config_path=str(config))
+
     MapGroupList = RootModel[list[MapGroup]]
 
     with open(config, "r") as handle:
         mgl = MapGroupList.model_validate_json(handle.read()).root
 
+    log = log.bind(number_of_groups=len(mgl))
+    log.info("config.parsed")
+
     return mgl
 
 
@@ -235,3 +242,12 @@ class MapGroup(AuthenticatedModel):
     description: str
 
     maps: list[Map]
+
+    def get_layer(self, layer_id: str) -> Layer | None:
+        for map in self.maps:
+            for band in map.bands:
+                for layer in band.layers:
+                    if layer.layer_id == layer_id:
+                        return layer
+
+        return None

tilemaker/metadata/generation.py

@@ -41,8 +41,6 @@ def map_group_from_fits(
             )
         )
 
-    print(maps)
-
     return MapGroup(
         name="Auto-Populated", description="No description provided", maps=maps
     )
@@ -54,7 +52,7 @@ def layers_from_fits(
     unit_override: str | None = None,
 ) -> list[Layer]:
     log = structlog.get_logger()
-    log = log.bind(filename=filename)
+    log = log.bind(filename=str(filename))
     data = fits.open(filename)
 
     if force:

tilemaker/providers/core.py

@@ -20,7 +20,7 @@ class PullableTile(BaseModel):
     x: int
     y: int
     level: int
-    grants: set[str]
+    grants: set[str] | None
 
     @property
     def hash(self) -> str: