Add Redis cache option

Author: blythed

CHANGELOG.md

@@ -25,6 +25,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Deprecate "free" queries
 - Create simpler developer contract for databackend
 - Snowflake native data-backend implementation
+- Add redis cache inside `db.metadata` for quick multi-process loading
 
 #### Bug Fixes
 

plugins/redis/pyproject.toml

@@ -0,0 +1,82 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "superduper_redis"
+readme = "README.md"
+description = "superduper allows users to work with arbitrary sklearn estimators, with additional support for pre-, post-processing and input/ output data-types."
+license = {file = "LICENSE"}
+maintainers = [{name = "superduper.io, Inc.", email = "[email protected]"}]
+keywords = [
+    "databases",
+    "mongodb",
+    "data-science",
+    "machine-learning",
+    "mlops",
+    "vector-database",
+    "ai",
+]
+requires-python = ">=3.10"
+dynamic = ["version"]
+dependencies = [
+    "redis"
+]
+
+[project.optional-dependencies]
+test = [
+    # Annotation plugin dependencies will be installed in CI
+    # :CI: plugins/mongodb
+]
+
+[project.urls]
+homepage = "https://superduper.io"
+documentation = "https://docs.superduper.io/docs/intro"
+source = "https://github.com/superduper-io/superduper"
+
+[tool.setuptools.packages.find]
+include = ["superduper_redis*"]
+
+[tool.setuptools.dynamic]
+version = {attr = "superduper_redis.__version__"}
+
+[tool.black]
+skip-string-normalization = true
+target-version = ["py38"]
+
+[tool.mypy]
+ignore_missing_imports = true
+no_implicit_optional = true
+warn_unused_ignores = true
+disable_error_code = ["has-type", "attr-defined", "assignment", "misc", "override", "call-arg"]
+
+[tool.pytest.ini_options]
+addopts = "-W ignore"
+
+[tool.ruff.lint]
+extend-select = [
+    "I", # Missing required import (auto-fixable)
+    "F", # PyFlakes
+    #"W", # PyCode Warning
+    "E", # PyCode Error
+    #"N", # pep8-naming
+    "D", # pydocstyle
+]
+ignore = [
+  "D100", # Missing docstring in public module
+  "D104", # Missing docstring in public package
+  "D107", # Missing docstring in __init__
+  "D105", # Missing docstring in magic method
+  "D203", # 1 blank line required before class docstring
+  "D212", # Multi-line docstring summary should start at the first line
+  "D213", # Multi-line docstring summary should start at the second line
+  "D401",
+  "E402",
+]
+
+[tool.ruff.lint.isort]
+combine-as-imports = true
+
+[tool.ruff.lint.per-file-ignores]
+"test/**" = ["D"]
+"plugin_test/**" = ["D"]

plugins/redis/superduper_redis/__init__.py

@@ -0,0 +1,5 @@
+from .cache import RedisCache as Cache
+
+__version__ = '0.6.0'
+
+__all__ = ['Cache']

plugins/redis/superduper_redis/cache.py

@@ -0,0 +1,65 @@
+import json
+import re
+
+import redis
+from superduper import logging
+from superduper.backends.base.cache import Cache
+
+
+class RedisCache(Cache):
+    """Local cache for caching components.
+
+    :param init_cache: Initialize cache
+    """
+
+    def __init__(self, uri: str = 'redis://localhost:6379/0'):
+        logging.info('Using Redis cache')
+        logging.info(f'Connecting to Redis cache at {uri}')
+        self.redis = redis.Redis.from_url(uri, decode_responses=True)
+        logging.info(f'Connecting to Redis cache at {uri}... DONE')
+
+    def __delitem__(self, item):
+        self.redis.delete(':'.join(item))
+
+    def __setitem__(self, key, value):
+        key = ':'.join(key)
+        self.redis.set(key, json.dumps(value))
+
+    def keys(self, *pattern):
+        """Get keys from the cache.
+
+        :param pattern: The pattern to search for.
+        """
+        pattern = ':'.join(pattern)
+        strings = list(self.redis.keys(pattern))
+        return [tuple(re.split(':', string)) for string in strings]
+
+    def __getitem__(self, item):
+        out = self.redis.get(':'.join(item))
+        if out is None:
+            raise KeyError(item)
+        return json.loads(out)
+
+    def initialize(self):
+        """Initialize the cache."""
+        pass
+
+    def drop(self, force: bool = False):
+        """Drop component from the cache.
+
+        :param uuid: Component uuid.
+        """
+        self.redis.flushdb()
+
+    @property
+    def db(self):
+        """Get the ``db``."""
+        return self._db
+
+    @db.setter
+    def db(self, value):
+        """Set the ``db``.
+
+        :param value: The value to set the ``db`` to.
+        """
+        self._db = value

superduper/backends/base/cache.py

@@ -1,16 +1,101 @@
+import typing as t
 from abc import abstractmethod
 
-from superduper.backends.base.backends import BaseBackend
 from superduper.components.component import Component
 
 
-class Cache(BaseBackend):
+class Cache:
     """Cache object for caching components.
 
     # noqa
     """
 
     @abstractmethod
-    def __getitem__(self, *item) -> Component:
+    def __getitem__(self, *item) -> t.Dict | t.List:
         """Get a component from the cache."""
         pass
+
+    @abstractmethod
+    def keys(self, *pattern) -> t.List[str]:
+        """Get the keys from the cache.
+
+        :param pattern: The pattern to match.
+
+        >>> cache.keys('*', '*', '*')
+        >>> cache.keys('Model', '*', '*')
+        >>> cache.keys('Model', 'my_model', '*')
+        >>> cache.keys('*', '*', '1234567890')
+        """
+
+    def get_with_uuid(self, uuid: str):
+        """Get a component from the cache with a specific uuid.
+
+        :param uuid: The uuid of the component to get.
+        """
+        key = self.keys('*', '*', uuid)
+        if not key:
+            return None
+        else:
+            key = key[0]
+
+        try:
+            return self[key]
+        except KeyError:
+            return
+
+    def get_with_component(self, component: str):
+        """Get all components from the cache of a certain type.
+
+        :param component: The component to get.
+        """
+        keys = self.keys(component, '*', '*')
+        return [self[k] for k in keys]
+
+    def get_with_component_identifier(self, component: str, identifier: str):
+        """Get a component from the cache with a specific identifier.
+
+        :param component: The component to get.
+        :param identifier: The identifier of the component to
+        """
+        keys = self.keys(component, identifier, '*')
+        out = [self[k] for k in keys]
+        if not out:
+            return None
+        out = max(out, key=lambda x: x['version'])  # type: ignore[arg-type, call-overload]
+        return out
+
+    def get_with_component_identifier_version(
+        self, component: str, identifier: str, version: int
+    ):
+        """Get a component from the cache with a specific version.
+
+        :param component: The component to get.
+        :param identifier: The identifier of the component to get.
+        :param version: The version of the component to get.
+        """
+        keys = self.keys(component, identifier, '*')
+        out = [self[k] for k in keys]
+        try:
+            return next(r for r in out if r['version'] == version)  # type: ignore[call-overload]
+        except StopIteration:
+            return
+
+    def __contains__(self, key: str) -> bool:
+        return key in self.keys()
+
+    @abstractmethod
+    def __setitem__(self, key: t.Tuple[str, str, str], value: t.Dict) -> None:
+        pass
+
+    def delete_uuid(self, uuid: str):
+        """Delete a component from the cache.
+
+        :param uuid: The uuid of the component to delete.
+        """
+        keys = self.keys('*', '*', uuid)
+        for key in keys:
+            del self[key]  # type: ignore[arg-type]
+
+    @abstractmethod
+    def __delitem__(self, key: t.Tuple[str, str, str]):
+        pass

superduper/backends/base/cluster.py

@@ -111,13 +111,3 @@ def initialize(self, with_compute: bool = False):
         self.cdc.initialize()
 
         logging.info(f"Cluster initialized in {time.time() - start:.2f} seconds.")
-
-    def drop_component(self, uuid: str):
-        """Drop component and its services rom the cluster.
-
-        :param uuid: Component uuid.
-        """
-        try:
-            del self.cache[uuid]
-        except KeyError:
-            pass