Improvements to the matrix cache (#1340)

* Improvements to the matrix cache

- Add a timelimit for the matrix in the cache.
Once the timelimit is reached, the matrix can be removed.

- If a DatasetAccessError occurs, then remove the dataset
from the matrix cache.

Fixes #1322

server/app/app.py

@@ -239,8 +239,7 @@ def get_api_resources(bp_api):
 
 
 class Server:
-    def __init__(self, matrix_data_cache_manager, annotations, app_config):
-
+    def __init__(self, app_config):
         self.app = Flask(__name__, static_folder="../common/web/static")
         self._before_adding_routes(app_config)
         self.app.json_encoder = Float32JSONEncoder
@@ -268,9 +267,8 @@ class Server:
             resources = get_api_resources(bp_api)
             self.app.register_blueprint(resources.blueprint)
             self.app.add_url_rule("/d/<dataset>/", "dataset_index", dataset_index, methods=["GET"])
-
-        self.app.matrix_data_cache_manager = matrix_data_cache_manager
-        self.app.annotations = annotations
+        self.app.matrix_data_cache_manager = app_config.matrix_data_cache_manager
+        self.app.annotations = app_config.user_annotations
         self.app.app_config = app_config
 
     def _before_adding_routes(self, app_config):

server/cli/launch.py

@@ -11,7 +11,6 @@ from flask_cors import CORS
 
 from server.common.utils import sort_options
 from server.common.errors import DatasetAccessError, ConfigurationError
-from server.data_common.matrix_loader import MatrixDataCacheManager
 from server.common.app_config import AppConfig
 from server.common.default_config import default_config
 from server.app.app import Server
@@ -276,8 +275,8 @@ class CliLaunchServer(Server):
     the CLI runs a local web server, and needs to enable a few more features.
     """
 
-    def __init__(self, matrix_data_cache_manager, annotations, app_config):
-        super().__init__(matrix_data_cache_manager, annotations, app_config)
+    def __init__(self, app_config):
+        super().__init__(app_config)
 
     def _before_adding_routes(self, app_config):
         self.app.config["COMPRESS_MIMETYPES"] = [
@@ -388,12 +387,11 @@ def launch(
         # process the configuration
         #  any errors will be thrown as an exception.
         #  any info messages will be passed to the messagefn function.
-        matrix_data_cache_manager = MatrixDataCacheManager()
 
         def messagefn(message):
             click.echo("[cellxgene] " + message)
 
-        app_config.complete_config(matrix_data_cache_manager, messagefn)
+        app_config.complete_config(messagefn)
 
         # Use a default secret if one is not provided
         if not app_config.server__flask_secret_key:
@@ -403,10 +401,9 @@ def launch(
         raise click.ClickException(e)
 
     handle_scripts(scripts)
-    user_annotations = app_config.user_annotations
 
     # create the server
-    server = CliLaunchServer(matrix_data_cache_manager, user_annotations, app_config)
+    server = CliLaunchServer(app_config)
 
     if not app_config.server__verbose:
         log = logging.getLogger("werkzeug")

server/common/app_config.py

@@ -72,6 +72,7 @@ class AppConfig(object):
             self.multi_dataset__index = dc["multi_dataset"]["index"]
             self.multi_dataset__allowed_matrix_types = dc["multi_dataset"]["allowed_matrix_types"]
             self.multi_dataset__matrix_cache__max_datasets = dc["multi_dataset"]["matrix_cache"]["max_datasets"]
+            self.multi_dataset__matrix_cache__timelimit_s = dc["multi_dataset"]["matrix_cache"]["timelimit_s"]
 
             self.single_dataset__datapath = dc["single_dataset"]["datapath"]
             self.single_dataset__obs_names = dc["single_dataset"]["obs_names"]
@@ -108,6 +109,9 @@ class AppConfig(object):
         # The annotation object is created during complete_config and stored here.
         self.user_annotations = None
 
+        # The matrix data cache manager is created during the complete_config and stored here.
+        self.matrix_data_cache_manager = None
+
         # Set to true when config_completed is called
         self.is_completed = False
 
@@ -146,12 +150,10 @@ class AppConfig(object):
 
         self.is_completed = False
 
-    def complete_config(self, matrix_data_cache_manager=None, messagefn=None):
+    def complete_config(self, messagefn=None):
         """The configure options are checked, and any additional setup based on the config
         parameters is done"""
 
-        if matrix_data_cache_manager is None:
-            matrix_data_cache_manager = MatrixDataCacheManager()
         if messagefn is None:
 
             def noop(message):
@@ -162,7 +164,7 @@ class AppConfig(object):
         # TODO: to give better error messages we can add a mapping between where each config
         # attribute originated (e.g. command line argument or config file), then in the error
         # messages we can give correct context for attributes with bad value.
-        context = dict(matrix_cache=matrix_data_cache_manager, messagefn=messagefn)
+        context = dict(messagefn=messagefn)
 
         self.handle_server(context)
         self.handle_single_dataset(context)
@@ -242,6 +244,10 @@ class AppConfig(object):
             if self.multi_dataset__dataroot is not None:
                 raise ConfigurationError("must supply only one of datapath or dataroot")
 
+        # create the matrix data cache manager:
+        if self.matrix_data_cache_manager is None:
+            self.matrix_data_cache_manager = MatrixDataCacheManager(max_cached=1, timelimit_s=None)
+
         # preload this data set
         matrix_data_loader = MatrixDataLoader(self.single_dataset__datapath)
         try:
@@ -278,6 +284,7 @@ class AppConfig(object):
         self.__check_attr("multi_dataset__index", (type(None), bool, str))
         self.__check_attr("multi_dataset__allowed_matrix_types", (tuple, list))
         self.__check_attr("multi_dataset__matrix_cache__max_datasets", int)
+        self.__check_attr("multi_dataset__matrix_cache__timelimit_s", (type(None), int, float))
 
         if self.multi_dataset__dataroot is None:
             return
@@ -289,8 +296,12 @@ class AppConfig(object):
             except ValueError:
                 raise ConfigurationError(f'Invalid matrix type in "allowed_matrix_types": {mtype}')
 
-        # matrix cache
-        MatrixDataCacheManager.set_max_datasets(self.multi_dataset__matrix_cache__max_datasets)
+        # create the matrix data cache manager:
+        if self.matrix_data_cache_manager is None:
+            self.matrix_data_cache_manager = MatrixDataCacheManager(
+                max_cached=self.multi_dataset__matrix_cache__max_datasets,
+                timelimit_s=self.multi_dataset__matrix_cache__timelimit_s,
+            )
 
     def handle_user_annotations(self, context):
         self.__check_attr("user_annotations__enable", bool)
@@ -328,7 +339,7 @@ class AppConfig(object):
             # if the user has specified a fixed label file, go ahead and validate it
             # so that we can remove errors early in the process.
             if self.single_dataset__datapath and self.user_annotations__local_file_csv__file:
-                with context["matrix_cache"].data_adaptor(self.single_dataset__datapath, self) as data_adaptor:
+                with self.matrix_data_cache_manager.data_adaptor(self.single_dataset__datapath, self) as data_adaptor:
                     data_adaptor.check_new_labels(self.user_annotations.read_labels(data_adaptor))
 
             if self.user_annotations__ontology__enable or self.user_annotations__ontology__obo_location:
@@ -372,7 +383,7 @@ class AppConfig(object):
         self.__check_attr("diffexp__lfc_cutoff", float)
 
         if self.single_dataset__datapath:
-            with context["matrix_cache"].data_adaptor(self.single_dataset__datapath, self) as data_adaptor:
+            with self.matrix_data_cache_manager.data_adaptor(self.single_dataset__datapath, self) as data_adaptor:
                 if self.diffexp__enable and data_adaptor.parameters.get("diffexp_may_be_slow", False):
                     context["messagefn"](
                         f"CAUTION: due to the size of your dataset, "

server/common/default_config.py

@@ -36,6 +36,10 @@ multi_dataset:
     # is evicted from the cache first.
     max_datasets: 5
 
+    # A matrix is automatically removed from the cache after timelimit_s number of seconds.
+    # If timelimit_s is set to None, then there is no time limit.
+    timelimit_s: 30
+
 single_dataset:
   datapath: null
   obs_names: null

server/data_common/matrix_loader.py

@@ -47,7 +47,7 @@ class MatrixDataCacheItem(object):
                 # necessary to hold the reader lock after an exception, since
                 # the release will occur when the context exits.
                 self.data_lock.w_demote()
-                raise e
+                raise DatasetAccessError(str(e))
 
         # demote the write lock to a read lock.
         self.data_lock.w_demote()
@@ -64,90 +64,145 @@ class MatrixDataCacheItem(object):
                 self.data_adaptor.cleanup()
                 self.data_adaptor = None
 
+    def attempt_delete(self):
+        """Delete, but only if the write lock can be immediately locked.  Return True if the delete happened"""
+        if self.data_lock.w_acquire_non_blocking():
+            if self.data_adaptor:
+                try:
+                    self.data_adaptor.cleanup()
+                    self.data_adaptor = None
+                except Exception:
+                    # catch all exceptions to ensure the lock is released
+                    pass
+
+                self.data_lock.w_release()
+                return True
+        else:
+            return False
+
+
+class MatrixDataCacheInfo(object):
+    def __init__(self, cache_item, timestamp):
+        # The MatrixDataCacheItem in the cache
+        self.cache_item = cache_item
+        # The last time the cache_item was accessed
+        self.last_access = timestamp
+        # The number of times the cache_item was accessed (used for testing)
+        self.num_access = 1
+
 
 class MatrixDataCacheManager(object):
     """A class to manage the cached datasets.   This is intended to be used as a context manager
     for handling api requests.  When the context is created, the data_adator is either loaded or
     retrieved from a cache.  In either case, the reader lock is taken during this time, and release
     when the context ends.  This class currently implements a simple least recently used cache,
-    which can delete a dataset from the cache to make room for a new oneo
+    which can delete a dataset from the cache to make room for a new one.
 
     This is the intended usage pattern:
 
-           m = MatrixDataCacheManager()
+           m = MatrixDataCacheManager(max_cached=..., timelimmit_s = ...)
            with m.data_adaptor(location, app_config) as data_adaptor:
                # use the data_adaptor for some operation
     """
 
-    #  The number of datasets to cache.  When MAX_CACHED is reached, the least recently used
-    #  cache is replaced with the newly requested one.
-    #  TODO:  This is very simple.  This can be improved by taking into account how much space is actually
-    #         taken by each dataset, instead of arbitrarily picking a max datasets to cache.
-    #         Also, this should be controlled by a configuration parameter.
-    MAX_CACHED = 5
-
-    @staticmethod
-    def set_max_datasets(max_cached):
-        MatrixDataCacheManager.MAX_CACHED = max_cached
-
-    # FIXME:   If the number of active datasets exceeds the MAX_CACHED, then each request could
+    # FIXME:   If the number of active datasets exceeds the max_cached, then each request could
     # lead to a dataset being deleted and a new only being opened: the cache will get thrashed.
     # In this case, we may need to send back a 503 (Server Unavailable), or some other error message.
 
-    # FIXME:  If the actual dataset is changed.  E.g. a new set of datafiles replaces an existing set,
-    # then the cache will not react to this.   Ideally this would invalidate the cache.  One solution is
-    # to keep a small metadata file associated with each dataset, which contains versioning information.
-    # When the dataset is accessed, the current version can be compared with the cached version, and if
-    # there is a mismatch, then the cache can be refreshed.
+    # NOTE:  If the actual dataset is changed.  E.g. a new set of datafiles replaces an existing set,
+    # then the cache will not react to this, however once the cache time limit is reached, the dataset
+    # will automatically be refreshed.
 
-    def __init__(self):
-        # key is location, value is tuple of (MatrixDataCacheItem, last_accessed)
+    def __init__(self, max_cached, timelimit_s=None):
+        # key is location, value is a MatrixDataCacheInfo
         self.datasets = {}
+
+        # lock to protect the datasets
         self.lock = threading.Lock()
 
+        #  The number of datasets to cache.  When max_cached is reached, the least recently used
+        #  cache is replaced with the newly requested one.
+        #  TODO:  This is very simple.  This can be improved by taking into account how much space is actually
+        #         taken by each dataset, instead of arbitrarily picking a max datasets to cache.
+        self.max_cached = max_cached
+
+        # items are automatically removed from the cache once this time limit is reached
+        self.timelimit_s = timelimit_s
+
     @contextmanager
     def data_adaptor(self, location, app_config):
         # create a loader for to this location if it does not already exist
 
         delete_adaptor = None
         data_adaptor = None
+        cache_item = None
 
         with self.lock:
-            value = self.datasets.get(location)
-            if value is not None:
-                cache_item = value[0]
-                last_accessed = time.time()
-                self.datasets[location] = (cache_item, last_accessed)
-                data_adaptor = cache_item.acquire_existing()
+            self.evict_old_datasets()
+            info = self.datasets.get(location)
+            if info is not None:
+                info.last_access = time.time()
+                info.num_access += 1
+                self.datasets[location] = info
+                data_adaptor = info.cache_item.acquire_existing()
+                cache_item = info.cache_item
 
             if data_adaptor is None:
                 while True:
-                    # find the last access times for each loader
-                    if len(self.datasets) < self.MAX_CACHED:
+                    if len(self.datasets) < self.max_cached:
                         break
 
                     items = list(self.datasets.items())
-                    sorted(items, key=lambda x: x[1][1])
+                    items = sorted(items, key=lambda x: x[1].last_access)
                     # close the least recently used loader
                     oldest = items[0]
-                    oldest_cache = oldest[1][0]
+                    oldest_cache = oldest[1].cache_item
                     oldest_key = oldest[0]
                     del self.datasets[oldest_key]
                     delete_adaptor = oldest_cache
 
-                last_accessed = time.time()
                 loader = MatrixDataLoader(location, app_config=app_config)
                 cache_item = MatrixDataCacheItem(loader)
-                self.datasets[location] = (cache_item, last_accessed)
+                item = MatrixDataCacheInfo(cache_item, time.time())
+                self.datasets[location] = item
 
         try:
+            assert(cache_item)
             if delete_adaptor:
                 delete_adaptor.delete()
             if data_adaptor is None:
                 data_adaptor = cache_item.acquire_and_open(app_config)
             yield data_adaptor
-        finally:
+        except DatasetAccessError:
             cache_item.release()
+            with self.lock:
+                del self.datasets[location]
+                cache_item.delete()
+            cache_item = None
+            raise
+
+        finally:
+            if cache_item:
+                cache_item.release()
+
+    def evict_old_datasets(self):
+        # must be called with the lock held
+        if self.timelimit_s is None:
+            return
+
+        now = time.time()
+        to_del = []
+        for key, info in self.datasets.items():
+            if (now - info.last_access) > self.timelimit_s:
+                # remove the data_cache when if it has been in the cache too long
+                to_del.append((key, info))
+
+        for key, info in to_del:
+            # try and get the write_lock for the dataset.
+            # if this returns false, it means the dataset is being used, and should
+            # not be removed.
+            if info.cache_item.attempt_delete():
+                del self.datasets[key]
 
 
 class MatrixDataType(Enum):

server/data_common/rwlock.py

@@ -100,6 +100,16 @@ class RWLock(object):
         self.d_lock.acquire()
         self.w_lock.acquire()
 
+    def w_acquire_non_blocking(self):
+        # if d_lock and w_lock can be acquired without blocking, acquire and return True,
+        # else immediately return False.
+        if self.d_lock.acquire(blocking=False):
+            if self.w_lock.acquire(blocking=False):
+                return True
+            else:
+                self.d_lock.release()
+        return False
+
     def w_release(self):
         self.w_lock.release()
         self.d_lock.release()

server/data_cxg/cxg_adaptor.py

@@ -250,6 +250,8 @@ class CxgAdaptor(DataAdaptor):
         with ServerTiming.time(f"layout.lsuri"):
             pemb = self.get_path("emb")
             embeddings = [os.path.basename(p) for (p, t) in self.lsuri(pemb) if t == "array"]
+        if len(embeddings) == 0:
+            raise DatasetAccessError("cxg matrix missing embeddings")
         return embeddings
 
     @staticmethod

server/eb/app.py

@@ -23,7 +23,6 @@ sys.path.append(SERVERDIR)
 try:
     from server.common.app_config import AppConfig
     from server.app.app import Server
-    from server.data_common.matrix_loader import MatrixDataCacheManager
     from server.common.data_locator import DataLocator
 except Exception:
     logging.critical("Exception importing server modules", exc_info=True)
@@ -31,8 +30,8 @@ except Exception:
 
 
 class WSGIServer(Server):
-    def __init__(self, matrix_data_cache_manager, annotations, app_config):
-        super().__init__(matrix_data_cache_manager, annotations, app_config)
+    def __init__(self, app_config):
+        super().__init__(app_config)
 
     def _before_adding_routes(self, app_config):
         csp = {"default-src": "'self' 'unsafe-inline' 'unsafe-eval'", "img-src": ["'self'", "data:"]}
@@ -75,8 +74,7 @@ try:
         multi_dataset__allowed_matrix_types=["cxg"],
     )
 
-    matrix_data_cache_manager = MatrixDataCacheManager()
-    app_config.complete_config(matrix_data_cache_manager, logging.info)
+    app_config.complete_config(logging.info)
 
     if not app_config.server__flask_secret_key:
         logging.critical(
@@ -86,7 +84,7 @@ try:
 
     user_annotations = app_config.user_annotations
 
-    server = WSGIServer(matrix_data_cache_manager, user_annotations, app_config)
+    server = WSGIServer(app_config)
 
     debug = False
     application = server.app

server/test/test_matrixcache.py

@@ -0,0 +1,124 @@
+import unittest
+from server.data_common.matrix_loader import MatrixDataCacheManager
+from server.common.app_config import AppConfig
+from server.common.errors import DatasetAccessError
+import tempfile
+import shutil
+import os
+import time
+
+
+class MatrixCacheTest(unittest.TestCase):
+    def setup(self):
+        pass
+
+    def make_temporay_datasets(self, dirname, num):
+        source = "test/test_datasets/pbmc3k.cxg"
+        for i in range(num):
+            target = os.path.join(dirname, str(i) + ".cxg")
+            shutil.copytree(source, target)
+
+    def use_dataset(self, matrix_cache, dirname, app_config, dataset_index):
+        with matrix_cache.data_adaptor(os.path.join(dirname, str(dataset_index) + ".cxg"), app_config) as adaptor:
+            pass
+        return adaptor
+
+    def use_dataset_with_error(self, matrix_cache, dirname, app_config, dataset_index):
+        try:
+            with matrix_cache.data_adaptor(os.path.join(dirname, str(dataset_index) + ".cxg"), app_config):
+                raise DatasetAccessError("something bad happened")
+        except DatasetAccessError:
+            # the MatrixDataCacheManager rethrows the exception, so catch and ignore
+            pass
+
+    def get_datasets(self, matrix_cache, dirname):
+        datasets = matrix_cache.datasets
+        result = {}
+        for k, v in datasets.items():
+            # filter out the dirname and the .cxg from the name
+            newk = int(k[len(dirname) + 1 : -4])
+            result[newk] = v
+
+        return result
+
+    def check_datasets(self, matrix_cache, dirname, expected):
+        res = self.get_datasets(matrix_cache, dirname)
+        actual = res.keys()
+        self.assertSetEqual(set(actual), set(expected))
+
+    def test_basic(self):
+        with tempfile.TemporaryDirectory() as dirname:
+            self.make_temporay_datasets(dirname, 5)
+            app_config = AppConfig()
+            m = MatrixDataCacheManager(max_cached=3, timelimit_s=None)
+
+            # should have only dataset 0
+            self.use_dataset(m, dirname, app_config, 0)
+            self.check_datasets(m, dirname, [0])
+
+            # should have datasets 0, 1
+            self.use_dataset(m, dirname, app_config, 1)
+            self.check_datasets(m, dirname, [0, 1])
+
+            # should have datasets 0, 1, 2
+            self.use_dataset(m, dirname, app_config, 2)
+            self.check_datasets(m, dirname, [0, 1, 2])
+
+            # should have datasets 1, 2, 3
+            self.use_dataset(m, dirname, app_config, 3)
+            self.check_datasets(m, dirname, [1, 2, 3])
+
+            # use dataset 1, making is more recent than dataset 2
+            self.use_dataset(m, dirname, app_config, 1)
+            self.check_datasets(m, dirname, [1, 2, 3])
+
+            # use dataset 4, should have 1,3,4
+            self.use_dataset(m, dirname, app_config, 4)
+            self.check_datasets(m, dirname, [1, 3, 4])
+
+            # use dataset 4 a few more times, get the count to 3
+            self.use_dataset(m, dirname, app_config, 4)
+            self.use_dataset(m, dirname, app_config, 4)
+
+            datasets = self.get_datasets(m, dirname)
+            self.assertEqual(datasets[1].num_access, 2)
+            self.assertEqual(datasets[3].num_access, 1)
+            self.assertEqual(datasets[4].num_access, 3)
+
+    def test_timelimit(self):
+        with tempfile.TemporaryDirectory() as dirname:
+            self.make_temporay_datasets(dirname, 2)
+
+            app_config = AppConfig()
+            m = MatrixDataCacheManager(max_cached=3, timelimit_s=1)
+
+            adaptor = self.use_dataset(m, dirname, app_config, 0)
+            adaptor1 = self.use_dataset(m, dirname, app_config, 0)
+            self.assertTrue(adaptor is adaptor1)
+
+            # wait until the timelimit expires and check that there is a new adaptor
+            time.sleep(1.1)
+            adaptor2 = self.use_dataset(m, dirname, app_config, 0)
+            self.assertTrue(adaptor is not adaptor2)
+            self.check_datasets(m, dirname, [0])
+
+            # now load a different dataset and see if dataset 0 gets evicted
+            time.sleep(1.1)
+            self.use_dataset(m, dirname, app_config, 1)
+            self.check_datasets(m, dirname, [1])
+
+    def test_access_error(self):
+        with tempfile.TemporaryDirectory() as dirname:
+            self.make_temporay_datasets(dirname, 1)
+
+            app_config = AppConfig()
+            m = MatrixDataCacheManager(max_cached=3, timelimit_s=1)
+
+            # use the 0 datasets
+            self.use_dataset(m, dirname, app_config, 0)
+            self.check_datasets(m, dirname, [0])
+
+            # use the 0 datasets, but this time a DatasetAccessError is raised.
+            # verify that dataset is removed from the cache.
+            self.use_dataset_with_error(m, dirname, app_config, 0)
+            self.check_datasets(m, dirname, [])