Keydb support

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

@@ -66,8 +66,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 +94,19 @@ 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 self.verbose:
                 logger.info("Version Information:")
                 vers = self.versions.as_dict()
                 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 +126,19 @@ 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}"
-            )
-
-            redis_builder.build_from_git(
-                self.versions.REDIS_URL, self.versions.REDIS_BRANCH
+                f"Building {database_name} version {self.versions.REDIS} from {self.versions.REDIS_URL}"
             )
-            redis_builder.cleanup()
-        logger.info("Redis build complete!")
+            database_builder.build_from_git(self.versions.REDIS_URL, self.versions.REDIS_BRANCH)
+            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

@@ -59,12 +59,12 @@ def site(self):
 
     def dbcli(self):
         install_path = get_install_path()
-        script_path = install_path.joinpath("_core/bin/redis-cli").resolve()
+        script_path = next((install_path / "_core" / "bin").glob("*-cli")).resolve()
         if script_path.is_file():
             print(script_path)
             exit(0)
         else:
-            print("Redis dependencies not found")
+            print("Database (Redis or KeyDB) dependencies not found")
             exit(1)
 
 

smartsim/_core/_install/buildenv.py

@@ -211,10 +211,11 @@ class Versioner:
     ONNX = Version_(REDISAI.onnx)
 
     def as_dict(self):
+        database_name = "REDIS" if "redis" in self.REDIS_URL else "KEYDB"
         packages = [
             "SMARTSIM",
             "SMARTREDIS",
-            "REDIS",
+            database_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.
     """
 
@@ -131,8 +131,8 @@ 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()
+        server = next(self.bin_path.glob("*-server"), None) is not None
+        cli = next(self.bin_path.glob("*-cli"), None) is not None
         return server and cli
 
     def build_from_git(self, git_url, branch):
@@ -142,16 +142,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 = "redis" if "redis" in git_url else "KeyDB"
+        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 +167,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 +178,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 +246,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 +280,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

tests/test_config.py

@@ -10,9 +10,9 @@
 def test_all_config_defaults():
     config = Config()
     assert Path(config.redisai).is_file()
-    assert Path(config.redis_exe).is_file()
-    assert Path(config.redis_cli).is_file()
-    assert Path(config.redis_conf).is_file()
+    assert Path(config.database_exe).is_file()
+    assert Path(config.database_cli).is_file()
+    assert Path(config.database_conf).is_file()
 
     # these will be changed so we will just run them
     assert config.log_level
@@ -38,25 +38,25 @@ def test_redisai():
 
 def test_redis_conf():
     config = Config()
-    assert Path(config.redis_conf).is_file()
-    assert isinstance(config.redis_conf, str)
+    assert Path(config.database_conf).is_file()
+    assert isinstance(config.database_conf, str)
 
     os.environ["REDIS_CONF"] = "not/a/path"
     config = Config()
     with pytest.raises(SSConfigError):
-        config.redis_conf
+        config.database_conf
     os.environ.pop("REDIS_CONF")
 
 
 def test_redis_exe():
     config = Config()
-    assert Path(config.redis_exe).is_file()
-    assert isinstance(config.redis_exe, str)
+    assert Path(config.database_exe).is_file()
+    assert isinstance(config.database_exe, str)
 
     os.environ["REDIS_PATH"] = "not/a/path"
     config = Config()
     with pytest.raises(SSConfigError):
-        config.redis_exe
+        config.database_exe
     os.environ.pop("REDIS_PATH")
 
 
@@ -68,5 +68,5 @@ def test_redis_cli():
     os.environ["REDIS_CLI_PATH"] = "not/a/path"
     config = Config()
     with pytest.raises(SSConfigError):
-        config.redis_cli
+        config.database_cli
     os.environ.pop("REDIS_CLI_PATH")