add a basic caching implementation in general and for Dataset object to start with

Author: aalok-sathe

langbrainscore/interface/cacheable.py

@@ -1,10 +1,13 @@
 
 
-from abc import ABC, abstractmethod, abstractclassmethod
+import pickle
 import typing
+from abc import ABC, abstractclassmethod, abstractmethod
 from pathlib import Path
-from diskcache import Cache
-
+import yaml
+import xarray as xr
+from langbrainscore.utils.cache import get_cache_root
+from langbrainscore.utils.logging import log
 
 T = typing.TypeVar('T')
 
@@ -13,24 +16,57 @@ class _Cacheable(ABC):
     A class used to define a common interface for Object caching in LangBrainscore
     '''
 
-    @abstractclassmethod
-    def _get_netcdf_cacheable_objects(cls) -> typing.Iterable[str]:
-        return ()
+    # @abstractclassmethod
+    # @classmethod
+    def _get_xarray_objects(self) -> typing.Iterable[str]:
+        '''
+        returns all (visible) attributes of self that are instances of xarray
+        '''
+        keys = []
+        for key, ob in vars(self).items():
+            if isinstance(ob, xr.DataArray):
+                keys += [key]
+        return keys
 
-    @abstractclassmethod
-    def _get_meta_attributes(cls) -> typing.Iterable[str]:
-        return ()
+    # @abstractclassmethod
+    # def _get_meta_attributes(cls) -> typing.Iterable[str]:
+    #     return ()
 
-    def to_cache(self, filename) -> None:
+    def to_cache(self, identifier_string: str, overwrite = True,
+                 xarray_serialization_backend='to_zarr', cache_dir = None) -> None:
         '''
         dump this object to cache. this method implementation will serve
         as the default implementation. it is recommended that this be left
         as-is for compatibility with caching across the library.
         '''
-        NotImplemented
+        root = Path(cache_dir or get_cache_root()).expanduser().resolve()
+        root /= identifier_string
+        root /= self.__class__.__name__
+        root.mkdir(exist_ok=True, parents=True)
+        log(f'caching {self} to {root}')
+
+        with (root / 'xarray_object_names.yml').open('w') as f:
+            yaml.dump(self._get_xarray_objects(), f, yaml.SafeDumper)
+
+        kwargs = {}
+        if overwrite and 'zarr' in xarray_serialization_backend: 
+            kwargs.update({'mode':'w'})
+        for ob_name in self._get_xarray_objects():
+            ob = getattr(self, ob_name)
+            tgt_dir = root / (ob_name + '.xr')
+            
+            dump_object = getattr(ob.to_dataset(name='data'), xarray_serialization_backend)
+            dump_object(tgt_dir, **kwargs)
+
+        with (root / f'{self.__class__.__name__}.pkl').open('wb') as f:
+            pickle.dump(self, f)
+
 
     # NB comment from Guido: https://github.com/python/typing/issues/58#issuecomment-194569410
-    def from_cache(cls: T, filename) -> typing.Callable[..., T]:
+    @classmethod
+    def from_cache(cls, identifier_string: str, 
+                   xarray_deserialization_backend='open_zarr',
+                   cache_dir = None) -> T:
         '''
         construct an object from cache. subclasses must start with the
         object returned by a call to this method like so:
@@ -41,12 +77,22 @@ def from_cache(cls: T, filename) -> typing.Callable[..., T]:
             return ob
             
         '''
-        return NotImplemented
+        root = Path(cache_dir or get_cache_root()).expanduser().resolve()
+        root /= identifier_string
+        root /= cls.__name__
+        root.mkdir(exist_ok=True, parents=True)
+        log(f'loading cache for {cls} from {root} {identifier_string}')
+
+        with (root/f'{cls.__name__}.pkl').open('rb') as f:
+            ob = pickle.load(f)
+
+        with (root / 'xarray_object_names.yml').open('r') as f:
+            xarray_object_names = yaml.load(f, yaml.SafeLoader)
 
-        C = Cache()
-        ob = cls()
-        for attr in cls._get_meta_attributes():
-            thing = None # retrieve from cache
-            setattr(ob, attr, thing)
+        for attr in xarray_object_names:
+            tgt_dir = root / (attr + '.xr')
+            load_object = getattr(xr, xarray_deserialization_backend)
+            xarray_instance = load_object(tgt_dir)
+            setattr(ob, attr, xarray_instance.data)
 
-        return ob
+        return ob

langbrainscore/interface/dataset.py

@@ -1,13 +1,13 @@
 import xarray as xr
 from langbrainscore.interface.cacheable import _Cacheable
+from abc import ABC
 
-
-# TODO: class _Dataset(_Cacheable, ABC)
-class _Dataset:
+class _Dataset(_Cacheable, ABC):
     """
     wrapper class for xarray DataArray that confirms format adheres to interface.
     """
 
+
     def __init__(self, xr_obj: xr.DataArray) -> "_Dataset":
         """
         accepts an xarray with the following core

langbrainscore/utils/cache.py

@@ -5,17 +5,21 @@
 
 import typing
 from pathlib import Path
+import os
 
 def get_cache_root(prefix: typing.Union[str, Path] = '~/.cache') -> Path:
     '''
 
     '''
+    if 'LBS_CACHE' in os.environ:
+        prefix = os.environ['LBS_CACHE']
+
     prefix = Path(prefix).expanduser().resolve()
     root = prefix / 'langbrainscore'
 
-    act = root / 'encoder_activations'
-    data = root / 'datasets' 
-    results = root / 'results'
+    # act = root / 'encoder_activations'
+    # data = root / 'datasets' 
+    # results = root / 'results'
 
     root.mkdir(parents=True, exist_ok=True)
-    return prefix
+    return root