Update vectordb_bench/backend/clients/memorydb/cli.py

Co-authored-by: Jonathan S. Katz <jkatz@users.noreply.github.com>
zilliztech · Jul 9, 2024 · 599b8e1 · 599b8e1
1 parent 99442bf
commit 599b8e1
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 3 deletions.
diff --git a/vectordb_bench/backend/clients/memorydb/cli.py b/vectordb_bench/backend/clients/memorydb/cli.py
@@ -44,7 +44,16 @@ class MemoryDBTypedDict(TypedDict):
             is_flag=True,
             show_default=True,
             default=False,
-            help="Cluster Mode Disabled (CMD), use this flag when testing locally on a single node instance. In production, MemoryDB only supports CME mode",
+            help="Cluster Mode Disabled (CMD), use this flag when testing locally on a single node instance. In production, MemoryDB only supports cluster mode (CME)",
+        ),
+    ]
+    insert_batch_size: Annotated[
+        int,
+        click.option(
+            "--insert-batch-size",
+            type=int,
+            default=10,
+            help="Batch size for inserting data. Adjust this as needed, but don't make it too big",
         ),
     ]
 
@@ -73,6 +82,7 @@ def MemoryDB(**parameters: Unpack[MemoryDBHNSWTypedDict]):
             M=parameters["m"],
             ef_construction=parameters["ef_construction"],
             ef_runtime=parameters["ef_runtime"],
+            insert_batch_size=parameters["insert_batch_size"],
         ),
         **parameters,
     )
diff --git a/vectordb_bench/backend/clients/memorydb/config.py b/vectordb_bench/backend/clients/memorydb/config.py
@@ -24,7 +24,7 @@ def to_dict(self) -> dict:
 
 class MemoryDBIndexConfig(BaseModel, DBCaseConfig):
     metric_type: MetricType | None = None
-    insert_batch_size: int | None = 10  # Adjust this as needed, but don't make too big
+    insert_batch_size: int
 
     def parse_metric(self) -> str:
         if self.metric_type == MetricType.L2:

diff --git a/vectordb_bench/backend/clients/memorydb/memorydb.py b/vectordb_bench/backend/clients/memorydb/memorydb.py
@@ -29,7 +29,7 @@ def __init__(
         self.case_config = db_case_config
         self.collection_name = INDEX_NAME
         self.target_nodes = RedisCluster.RANDOM if not self.db_config["cmd"] else None
-        self.insert_batch_size = db_case_config.insert_batch_size or 10
+        self.insert_batch_size = db_case_config.insert_batch_size
         self.dbsize = kwargs.get("num_rows")
 
         # Create a MemoryDB connection, if db has password configured, add it to the connection here and in init():