Keydb support (#180)

keydb flag for smart build added. Users can now
use KeyDB in place of Redis more easily.

[ committed by @EricGustin ]
[ reviewed by @MattToast @ashao @Spartee ]

.gitignore

@@ -17,8 +17,8 @@ smartredis
 # written upon install
 smartsim/version.py
 
-smartsim/_core/bin/redis-server
-smartsim/_core/bin/redis-cli
+smartsim/_core/bin/*-server
+smartsim/_core/bin/*-cli
 
 # created upon install
 smartsim/_core/lib

setup.py

@@ -108,14 +108,14 @@ def finalize_options(self):
 class SmartSimBuild(build_py):
 
     def run(self):
-        redis_builder = builder.RedisBuilder(build_env(),
+        database_builder = builder.DatabaseBuilder(build_env(),
                                              build_env.MALLOC,
                                              build_env.JOBS)
-        if not redis_builder.is_built:
-            redis_builder.build_from_git(versions.REDIS_URL,
+        if not database_builder.is_built:
+            database_builder.build_from_git(versions.REDIS_URL,
                                          versions.REDIS)
 
-            redis_builder.cleanup()
+            database_builder.cleanup()
 
         # run original build_py command
         build_py.run(self)

smartsim/_core/_cli/build.py

@@ -12,6 +12,7 @@
 from smartsim._core._install.builder import BuildError
 from smartsim._core.config import CONFIG
 from smartsim._core.utils.helpers import installed_redisai_backends
+from smartsim.error import SSConfigError
 from smartsim.log import get_logger
 
 smart_logger_format = "[%(name)s] %(levelname)s %(message)s"
@@ -66,8 +67,15 @@ def __init__(self):
             type=str,
             help="Path to custom libtensorflow directory (ONLY USED IF NEEDED)",
         )
+        parser.add_argument(
+            "--keydb",
+            action="store_true",
+            default=False,
+            help="Build KeyDB instead of Redis",
+        )
         args = parser.parse_args(sys.argv[2:])
         self.verbose = args.v
+        self.keydb = args.keydb
 
         # torch and tf build by default
         pt = not args.no_pt
@@ -87,13 +95,24 @@ def __init__(self):
             logger.info("Checking requested versions...")
             self.versions = Versioner()
 
+            if self.keydb:
+                self.versions.REDIS = Version_("6.2.0")
+                self.versions.REDIS_URL = "https://github.com/EQ-Alpha/KeyDB"
+                self.versions.REDIS_BRANCH = "v6.2.0"
+                CONFIG.conf_path = Path(CONFIG.core_path, "config", "keydb.conf")
+                if not CONFIG.conf_path.resolve().is_file():
+                    raise SSConfigError(
+                        "Database configuration file at REDIS_CONF could not be found"
+                    )
+
             if self.verbose:
+                db_name = "KEYDB" if self.keydb else "REDIS"
                 logger.info("Version Information:")
-                vers = self.versions.as_dict()
+                vers = self.versions.as_dict(db_name=db_name)
                 print(tabulate(vers, headers=vers.keys(), tablefmt="github"), "\n")
 
-            # REDIS
-            self.build_redis()
+            # REDIS/KeyDB
+            self.build_database()
 
             # REDISAI
             self.build_redis_ai(
@@ -113,22 +132,21 @@ def __init__(self):
 
         logger.info("SmartSim build complete!")
 
-    def build_redis(self):
-        # check redis installation
-        redis_builder = builder.RedisBuilder(
+    def build_database(self):
+        # check database installation
+        database_name = "KeyDB" if self.keydb else "Redis"
+        database_builder = builder.DatabaseBuilder(
             self.build_env(), self.build_env.MALLOC, self.build_env.JOBS, self.verbose
         )
-
-        if not redis_builder.is_built:
+        if not database_builder.is_built:
             logger.info(
-                f"Building Redis version {self.versions.REDIS} from {self.versions.REDIS_URL}"
+                f"Building {database_name} version {self.versions.REDIS} from {self.versions.REDIS_URL}"
             )
-
-            redis_builder.build_from_git(
+            database_builder.build_from_git(
                 self.versions.REDIS_URL, self.versions.REDIS_BRANCH
             )
-            redis_builder.cleanup()
-        logger.info("Redis build complete!")
+            database_builder.cleanup()
+        logger.info(f"{database_name} build complete!")
 
     def build_redis_ai(
         self, device, torch=True, tf=True, onnx=False, torch_dir=None, libtf_dir=None

smartsim/_core/_cli/clean.py

@@ -53,7 +53,7 @@ def clean(self, _all=False):
 
         bin_path = self._core_path / "bin"
         if bin_path.is_dir() and _all:
-            files_to_remove = ["redis-server", "redis-cli"]
+            files_to_remove = ["redis-server", "redis-cli", "keydb-server", "keydb-cli"]
             removed = False
             for _file in files_to_remove:
                 file_path = bin_path.joinpath(_file)
@@ -62,4 +62,4 @@ def clean(self, _all=False):
                     removed = True
                     file_path.unlink()
             if removed:
-                logger.info("Successfully removed SmartSim Redis installation")
+                logger.info("Successfully removed SmartSim database installation")

smartsim/_core/_cli/cli.py

@@ -58,14 +58,13 @@ def site(self):
         exit(0)
 
     def dbcli(self):
-        install_path = get_install_path()
-        script_path = install_path.joinpath("_core/bin/redis-cli").resolve()
-        if script_path.is_file():
-            print(script_path)
-            exit(0)
-        else:
-            print("Redis dependencies not found")
-            exit(1)
+        bin_path = get_install_path() / "_core" / "bin"
+        for option in bin_path.iterdir():
+            if option.name in ("redis-cli", "keydb-cli"):
+                print(option)
+                exit(0)
+        print("Database (Redis or KeyDB) dependencies not found")
+        exit(1)
 
 
 def main():

smartsim/_core/_install/buildenv.py

@@ -210,11 +210,11 @@ class Versioner:
     TENSORFLOW = Version_(REDISAI.tensorflow)
     ONNX = Version_(REDISAI.onnx)
 
-    def as_dict(self):
+    def as_dict(self, db_name="REDIS"):
         packages = [
             "SMARTSIM",
             "SMARTREDIS",
-            "REDIS",
+            db_name,
             "REDISAI",
             "TORCH",
             "TENSORFLOW",

smartsim/_core/_install/builder.py

@@ -116,11 +116,11 @@ def run_command(self, cmd, shell=False, out=None, cwd=None):
             raise BuildError(e)
 
 
-class RedisBuilder(Builder):
-    """Class to build Redis from Source
+class DatabaseBuilder(Builder):
+    """Class to build Redis or KeyDB from Source
     Supported build methods:
      - from git
-    See buildenv.py for buildtime configuration of Redis
+    See buildenv.py for buildtime configuration of Redis/KeyDB
     version and url.
     """
 
@@ -130,10 +130,11 @@ def __init__(self, build_env={}, malloc="libc", jobs=None, verbose=False):
 
     @property
     def is_built(self):
-        """Check if Redis is built"""
-        server = self.bin_path.joinpath("redis-server").is_file()
-        cli = self.bin_path.joinpath("redis-cli").is_file()
-        return server and cli
+        """Check if Redis or KeyDB is built"""
+        bin_files = {file.name for file in self.bin_path.iterdir()}
+        redis_files = {"redis-server", "redis-cli"}
+        keydb_files = {"keydb-server", "keydb-cli"}
+        return redis_files.issubset(bin_files) or keydb_files.issubset(bin_files)
 
     def build_from_git(self, git_url, branch):
         """Build Redis from git
@@ -142,16 +143,21 @@ def build_from_git(self, git_url, branch):
         :param branch: branch to checkout
         :type branch: str
         """
-        redis_build_path = Path(self.build_dir, "redis")
+        database_name = "KeyDB" if "KeyDB" in git_url else "redis"
+        database_build_path = Path(self.build_dir, database_name.lower())
 
         # remove git directory if it exists as it should
         # really never exist as we delete after build
+        redis_build_path = Path(self.build_dir, "redis")
+        keydb_build_path = Path(self.build_dir, "keydb")
         if redis_build_path.is_dir():
             shutil.rmtree(str(redis_build_path))
+        if keydb_build_path.is_dir():
+            shutil.rmtree(str(keydb_build_path))
 
-        # Check Redis URL
+        # Check database URL
         if not self.is_valid_url(git_url):
-            raise BuildError(f"Malformed Redis URL: {git_url}")
+            raise BuildError(f"Malformed {database_name} URL: {git_url}")
 
         # clone Redis
         clone_cmd = [
@@ -162,7 +168,7 @@ def build_from_git(self, git_url, branch):
             branch,
             "--depth",
             "1",
-            "redis",
+            database_name,
         ]
         self.run_command(clone_cmd, cwd=self.build_dir)
 
@@ -173,16 +179,16 @@ def build_from_git(self, git_url, branch):
             str(self.jobs),
             f"MALLOC={self.malloc}",
         ]
-        self.run_command(build_cmd, cwd=str(redis_build_path))
+        self.run_command(build_cmd, cwd=str(database_build_path))
 
         # move redis binaries to smartsim/smartsim/_core/bin
-        redis_src_dir = redis_build_path / "src"
-        self.copy_file(
-            redis_src_dir / "redis-server", self.bin_path / "redis-server", set_exe=True
-        )
-        self.copy_file(
-            redis_src_dir / "redis-cli", self.bin_path / "redis-cli", set_exe=True
-        )
+        database_src_dir = database_build_path / "src"
+        server_source = database_src_dir / (database_name.lower() + "-server")
+        server_destination = self.bin_path / (database_name.lower() + "-server")
+        cli_source = database_src_dir / (database_name.lower() + "-cli")
+        cli_destination = self.bin_path / (database_name.lower() + "-cli")
+        self.copy_file(server_source, server_destination, set_exe=True)
+        self.copy_file(cli_source, cli_destination, set_exe=True)
 
 
 class RedisAIBuilder(Builder):
@@ -241,12 +247,12 @@ def symlink_libtf(self, device):
         :param device: cpu or gpu
         :type device: str
         """
-        rai_deps_path = sorted(self.rai_build_path.glob(
-            os.path.join("deps", f"*{device}*")
-        ))
+        rai_deps_path = sorted(
+            self.rai_build_path.glob(os.path.join("deps", f"*{device}*"))
+        )
         if not rai_deps_path:
             raise FileNotFoundError("Could not find RedisAI 'deps' directory")
-        
+
         # There should only be one path for a given device,
         # and this should hold even if in the future we use
         # an external build of RedisAI
@@ -275,11 +281,15 @@ def symlink_libtf(self, device):
         if src_libtf_lib_dir.is_dir():
             library_files = sorted(src_libtf_lib_dir.glob("*"))
             if not library_files:
-                raise FileNotFoundError(f"Could not find libtensorflow library files in {src_libtf_lib_dir}")  
+                raise FileNotFoundError(
+                    f"Could not find libtensorflow library files in {src_libtf_lib_dir}"
+                )
         else:
             library_files = sorted(libtf_path.glob("lib*.so*"))
             if not library_files:
-                raise FileNotFoundError(f"Could not find libtensorflow library files in {libtf_path}")        
+                raise FileNotFoundError(
+                    f"Could not find libtensorflow library files in {libtf_path}"
+                )
 
         for src_file in library_files:
             dst_file = rai_libtf_lib_dir / src_file.name

smartsim/_core/config/config.py

@@ -102,30 +102,30 @@ def redisai(self) -> str:
         return str(redisai)
 
     @property
-    def redis_conf(self) -> str:
+    def database_conf(self) -> str:
         conf = Path(os.environ.get("REDIS_CONF", self.conf_path)).resolve()
         if not conf.is_file():
             raise SSConfigError(
-                "Redis configuration file at REDIS_CONF could not be found"
+                "Database configuration file at REDIS_CONF could not be found"
             )
         return str(conf)
 
     @property
-    def redis_exe(self) -> str:
+    def database_exe(self) -> str:
         try:
-            redis_exe = self.bin_path / "redis-server"
-            redis = Path(os.environ.get("REDIS_PATH", redis_exe)).resolve()
-            exe = expand_exe_path(str(redis))
+            database_exe = next(self.bin_path.glob("*-server"))
+            database = Path(os.environ.get("REDIS_PATH", database_exe)).resolve()
+            exe = expand_exe_path(str(database))
             return exe
         except (TypeError, FileNotFoundError) as e:
             raise SSConfigError(
-                "Specified Redis binary at REDIS_PATH could not be used"
+                "Specified database binary at REDIS_PATH could not be used"
             ) from e
 
     @property
-    def redis_cli(self) -> str:
+    def database_cli(self) -> str:
         try:
-            redis_cli_exe = self.bin_path / "redis-cli"
+            redis_cli_exe = next(self.bin_path.glob("*-cli"))
             redis_cli = Path(os.environ.get("REDIS_CLI_PATH", redis_cli_exe)).resolve()
             exe = expand_exe_path(str(redis_cli))
             return exe

smartsim/_core/launcher/colocated.py

@@ -121,8 +121,8 @@ def _build_colocated_wrapper_cmd(port=6780,
 
     # collect DB binaries and libraries from the config
     db_cmd = [
-        CONFIG.redis_exe,
-        CONFIG.redis_conf,
+        CONFIG.database_exe,
+        CONFIG.database_conf,
         "--loadmodule",
         CONFIG.redisai
     ]

smartsim/_core/utils/redis.py

@@ -62,7 +62,7 @@ def create_cluster(hosts, ports):  # cov-wlm
             ip_list.append(address)
 
     # call cluster command
-    redis_cli = CONFIG.redis_cli
+    redis_cli = CONFIG.database_cli
     cmd = [redis_cli, "--cluster", "create"]
     cmd += ip_list
     cmd += ["--cluster-replicas", "0"]

smartsim/database/orchestrator.py

@@ -179,7 +179,7 @@ def _detect_command(launcher):
             # will raise SSConfigError if not found
             self._redis_exe
             self._redis_conf
-            CONFIG.redis_cli
+            CONFIG.database_cli
         except SSConfigError as e:
             msg = "SmartSim not installed with pre-built extensions (Redis)\n"
             msg += "Use the `smart` cli tool to install needed extensions\n"
@@ -299,11 +299,11 @@ def _rai_module(self):
 
     @property
     def _redis_exe(self):
-        return CONFIG.redis_exe
+        return CONFIG.database_exe
 
     @property
     def _redis_conf(self):
-        return CONFIG.redis_conf
+        return CONFIG.database_conf
 
     def set_cpus(self, num_cpus):
         """Set the number of CPUs available to each database shard