refactor: warm cache

.flake8

@@ -1,3 +1,13 @@
 [flake8]
+# Ignoring
+# F401: module imported but not used https://www.flake8rules.com/rules/F401.html
+# F402: Import module from line n shadowed by loop variable https://www.flake8rules.com/rules/F402.html
 per-file-ignores = __init__.py:F401,E402
-ignore = E501,W503,F722, E731
+# Ignoring:
+# E501: line too long
+# W503: Line break occurred before a binary operator https://www.flake8rules.com/rules/W503.html
+# n.b. PEP now recommends linebreak BEFORE operator. see https://peps.python.org/pep-0008/#should-a-line-break-before-or-after-a-binary-operator
+# E731: do not assign a lambda expression, use a def
+ignore = E501,W503,E731
+
+exclude = venv,examples,docs,config_schema,.ebrains

.github/workflows/siibra-testing.yml

@@ -2,6 +2,17 @@ name: '[test] unit test'
 
 on: [push]
 jobs:
+  lint:
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v4
+    - name: Set up Python 3.10
+      uses: actions/setup-python@v4
+      with:
+        python-version: '3.10'
+    - run: pip install flake8
+    - run: flake8 .
   unit-tests:
     runs-on: ubuntu-latest
     env:
@@ -27,12 +38,6 @@ jobs:
         python -m pip install -U .
         pip install -r requirements-test.txt
         if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
-    - name: Lint with flake8
-      run: |
-        # stop the build if there are Python syntax errors or undefined names
-        flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
-        # exit-zero treats all errors as warnings. The GitHub editor is 127 chars wide
-        flake8 . --count --exit-zero --max-complexity=10 --max-line-length=127 --statistics
     - name: Install test dependencies
       run:  pip install pytest pytest-cov coverage
     - name: Run test with pytest

e2e/volumes/test_sparsemap_cache_uniqueness.py

@@ -4,4 +4,4 @@
 def test_sparsemap_cache_uniqueness():
     mp157 = siibra.get_map("julich 3.0", "colin 27", "statistical", spec="157")
     mp175 = siibra.get_map("julich 3.0", "colin 27", "statistical", spec="175")
-    assert mp157.sparse_index.probs[0] != mp175.sparse_index.probs[0]
+    assert mp157.sparse_index.probs[0] != mp175.sparse_index.probs[0]

siibra/__init__.py

@@ -33,6 +33,8 @@
     EbrainsRequest as _EbrainsRequest,
     CACHE as cache
 )
+from .retrieval.cache import Warmup, WarmupLevel
+
 from . import configuration
 from .configuration import factory
 from . import features, livequeries
@@ -113,7 +115,7 @@ def set_cache_size(maxsize_gbyte: int):
     set_cache_size(float(_os.environ.get("SIIBRA_CACHE_SIZE_GIB")))
 
 
-def warm_cache():
+def warm_cache(level=WarmupLevel.INSTANCE):
     """
     Preload preconfigured siibra concepts.
 
@@ -122,12 +124,7 @@ def warm_cache():
     features. By preloading the instances, siibra commits all preconfigurations
     to the memory at once instead of commiting them when required.
     """
-    _ = _atlas.Atlas.registry()
-    _ = _space.Space.registry()
-    _ = _parcellation.Parcellation.registry()
-    _ = _parcellationmap.Map.registry()
-    features.warm_cache()
-    livequeries.warm_cache()
+    Warmup.warmup(level)
 
 
 def __dir__():

siibra/core/__init__.py

@@ -14,10 +14,3 @@
 # limitations under the License.
 """:ref:`Main siibra concepts<mainconcepts>`"""
 from . import atlas, parcellation, space
-
-
-def warm_cache():
-    """Preload preconfigured siibra core concepts."""
-    _ = atlas.Atlas.registry()
-    _ = space.Space.registry()
-    _ = parcellation.Parcellation.registry()

siibra/core/concept.py

@@ -21,12 +21,19 @@
     Species,
     TypePublication
 )
+from ..retrieval import cache
 
 import re
-from typing import TypeVar, Type, Union, List, TYPE_CHECKING
+from typing import TypeVar, Type, Union, List, TYPE_CHECKING, Dict
 
 T = TypeVar("T", bound="AtlasConcept")
-_REGISTRIES = {}
+_REGISTRIES: Dict[Type[T], InstanceTable[T]] = {}
+
+
+@cache.Warmup.register_warmup_fn(is_factory=True)
+def _atlas_concept_warmup():
+    return [cls.registry for cls in _REGISTRIES]
+
 
 if TYPE_CHECKING:
     from ..retrieval.datasets import EbrainsDataset

siibra/exceptions.py

@@ -41,3 +41,7 @@ class SpaceWarpingFailedError(RuntimeError):
 
 class NoVolumeFound(RuntimeError):
     pass
+
+
+class WarmupRegException(Exception):
+    pass

siibra/features/__init__.py

@@ -23,6 +23,8 @@
 
 from typing import Union
 from .feature import Feature
+from ..retrieval import cache
+from ..commons import siibra_tqdm
 
 get = Feature._match
 
@@ -46,12 +48,31 @@ def __getattr__(attr: str):
         raise AttributeError(f"No such attribute: {__name__}.{attr} " + hint)
 
 
-def warm_cache():
+@cache.Warmup.register_warmup_fn()
+def _warm_feature_cache_insntaces():
     """Preload preconfigured multimodal data features."""
     for ftype in TYPES.values():
         _ = ftype._get_instances()
-    from ..livequeries.bigbrain import WagstylProfileLoader
-    WagstylProfileLoader._load()
+
+
+@cache.Warmup.register_warmup_fn(cache.WarmupLevel.DATA, is_factory=True)
+def _warm_feature_cache_data():
+    return_callables = []
+    for ftype in TYPES.values():
+        instances = ftype._get_instances()
+        tally = siibra_tqdm(desc=f"Warming data {ftype.__name__}", total=len(instances))
+        for f in instances:
+            def get_data():
+                # TODO
+                # the try catch is as a result of https://github.com/FZJ-INM1-BDA/siibra-python/issues/509
+                # sometimes f.data can fail
+                try:
+                    _ = f.data
+                except Exception:
+                    ...
+                tally.update(1)
+            return_callables.append(get_data)
+    return return_callables
 
 
 def render_ascii_tree(class_or_classname: Union[type, str]):

siibra/features/tabular/receptor_density_fingerprint.py

@@ -170,10 +170,11 @@ def polar_plot(self, *args, backend='matplotlib', **kwargs):
                             self.data["mean"] + self.data["std"]
                         ]
                     ),
-                    "cat":
-                        len(self.data) * ["mean"] +\
-                        len(self.data) * ["mean - std"] +\
-                        len(self.data) * ["mean + std"]
+                    "cat": (
+                        len(self.data) * ["mean"]
+                        + len(self.data) * ["mean - std"]
+                        + len(self.data) * ["mean + std"]
+                    )
                 }
             )
             return line_polar(

siibra/livequeries/__init__.py

@@ -17,9 +17,3 @@
 from .allen import AllenBrainAtlasQuery
 from .bigbrain import LayerwiseBigBrainIntensityQuery, BigBrainProfileQuery
 from .ebrains import EbrainsFeatureQuery
-
-
-def warm_cache():
-    """Preload downloadable livequeries."""
-    from .bigbrain import WagstylProfileLoader
-    WagstylProfileLoader()

siibra/livequeries/bigbrain.py

@@ -77,6 +77,9 @@ def __len__(self):
         return self._vertices.shape[0]
 
 
+cache.Warmup.register_warmup_fn()(lambda: WagstylProfileLoader._load())
+
+
 class BigBrainProfileQuery(query.LiveQuery, args=[], FeatureType=bigbrain_intensity_profile.BigBrainIntensityProfile):
 
     def __init__(self):

siibra/livequeries/ebrains.py

@@ -19,7 +19,7 @@
 
 from ..commons import logger, siibra_tqdm
 from ..features import anchor as _anchor
-from ..retrieval import requests, datasets
+from ..retrieval import requests, datasets, cache
 from ..core import parcellation, region
 
 from collections import defaultdict
@@ -140,3 +140,6 @@ def query(self, region: region.Region):
                 f"Its version history is: {curr.version_history}"
             )
             yield curr
+
+
+cache.Warmup.register_warmup_fn(cache.WarmupLevel.DATA)(lambda: EbrainsFeatureQuery.loader.data)

siibra/retrieval/cache.py

@@ -12,14 +12,19 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-"""Maintaining and hadnling caching files on disk."""
+"""Maintaining and handling caching files on disk."""
 
 import hashlib
 import os
 from appdirs import user_cache_dir
 import tempfile
+from functools import wraps
+from enum import Enum
+from typing import Callable, List, NamedTuple, Union
+from concurrent.futures import ThreadPoolExecutor
 
-from ..commons import logger, SIIBRA_CACHEDIR, SKIP_CACHEINIT_MAINTENANCE
+from ..commons import logger, SIIBRA_CACHEDIR, SKIP_CACHEINIT_MAINTENANCE, siibra_tqdm
+from ..exceptions import WarmupRegException
 
 
 def assert_folder(folder):
@@ -137,3 +142,72 @@ def build_filename(self, str_rep: str, suffix=None):
 
 
 CACHE = Cache.instance()
+
+
+class WarmupLevel(int, Enum):
+    INSTANCE = 1
+    DATA = 5
+
+
+class WarmupParam(NamedTuple):
+    level: Union[int, WarmupLevel]
+    fn: Callable
+    is_factory: bool = False
+
+
+class Warmup:
+
+    _warmup_fns: List[WarmupParam] = []
+
+    @staticmethod
+    def fn_eql(wrapped_fn, original_fn):
+        return wrapped_fn is original_fn or wrapped_fn.__wrapped__ is original_fn
+
+    @classmethod
+    def is_registered(cls, fn):
+        return len([warmup_fn.fn
+                    for warmup_fn in cls._warmup_fns
+                    if cls.fn_eql(warmup_fn.fn, fn)]) > 0
+
+    @classmethod
+    def register_warmup_fn(cls, warmup_level: WarmupLevel = WarmupLevel.INSTANCE, *, is_factory=False):
+        def outer(fn):
+            if cls.is_registered(fn):
+                raise WarmupRegException
+
+            @wraps(fn)
+            def inner(*args, **kwargs):
+                return fn(*args, **kwargs)
+
+            cls._warmup_fns.append(WarmupParam(warmup_level, inner, is_factory))
+            return inner
+        return outer
+
+    @classmethod
+    def deregister_warmup_fn(cls, original_fn):
+        cls._warmup_fns = [
+            warmup_fn for warmup_fn in cls._warmup_fns
+            if not cls.fn_eql(warmup_fn.fn, original_fn)
+        ]
+
+    @classmethod
+    def warmup(cls, warmup_level: WarmupLevel = WarmupLevel.INSTANCE, *, max_workers=4):
+        all_fns = [warmup for warmup in cls._warmup_fns if warmup.level <= warmup_level]
+
+        def call_fn(fn: WarmupParam):
+            return_val = fn.fn()
+            if not fn.is_factory:
+                return
+            for f in return_val:
+                f()
+
+        with ThreadPoolExecutor(max_workers=max_workers) as ex:
+            for _ in siibra_tqdm(
+                ex.map(
+                    call_fn,
+                    all_fns
+                ),
+                desc="Warming cache",
+                total=len(all_fns),
+            ):
+                ...

siibra/retrieval/exceptions/__init__.py

@@ -14,6 +14,7 @@
 # limitations under the License.
 """Exceptions concerning file retrieval processes."""
 
+
 class NoSiibraConfigMirrorsAvailableException(Exception):
     pass
 

siibra/vocabularies/__init__.py

@@ -29,7 +29,7 @@ def runtime_path(fname: str):
 
 with open(runtime_path('gene_names.json'), 'r') as f:
     _gene_names = json.load(f)
-    GENE_NAMES = InstanceTable[str](
+    GENE_NAMES = InstanceTable(
         elements={
             k: {'symbol': k, 'description': v}
             for k, v in _gene_names.items()

siibra/volumes/__init__.py

@@ -21,8 +21,3 @@
 from ..commons import logger
 from typing import List, Union
 import numpy as np
-
-
-def warm_cache():
-    """Preload preconfigured parcellation maps."""
-    _ = Map.registry()