rename "metadata" to "config" to match benchmark wordings, add chromadb volume, fix CHROMA_WORKERS, allow more chromadb versions in pyproject.toml

LukasWestholt · LukasWestholt · commit 6eb164791aa1 · 2024-09-30T21:17:57.000+02:00
diff --git a/benchmark/convert.py b/benchmark/convert.py
@@ -60,7 +60,7 @@ def main():
         data = json.load(open(search_file))
         experiment_name = data["params"]["experiment"]
         dataset_name = data["params"]["dataset"]
-        engine_params = data["params"].get("config", {})
+        engine_params = data["params"]["config"]
         parallel = data["params"]["parallel"]
         engine_name = data["params"]["engine"]
 
diff --git a/engine/clients/chroma/configure.py b/engine/clients/chroma/configure.py
@@ -37,7 +37,7 @@ def recreate(self, dataset: Dataset, collection_params):
         params = self.collection_params
         params["metadata"] = dict(
             {"hnsw:space": self.DISTANCE_MAPPING.get(dataset.config.distance)},
-            **params.get("metadata", {}),
+            **params.get("config", {}),
         )
         self.client.create_collection(
             name=CHROMA_COLLECTION_NAME,
diff --git a/engine/clients/chroma/parser.py b/engine/clients/chroma/parser.py
@@ -27,7 +27,6 @@ def build_condition(
                 condition = {**condition, **or_subfilters[0]}
 
         return condition
-        # return {k: v for d in [flt for xs in [and_subfilters, or_subfilters] for flt in xs] for k, v in d.items()}
 
     def build_exact_match_filter(self, field_name: str, value: FieldValue) -> Where:
         return {field_name: value}
diff --git a/engine/clients/chroma/search.py b/engine/clients/chroma/search.py
@@ -40,5 +40,5 @@ def search_one(cls, query: Query, top: int) -> List[Tuple[int, float]]:
     def setup_search(self):
         metadata = self.collection.metadata.copy()
         metadata.pop("hnsw:space", None)  # Not allowed in the collection.modify method
-        metadata.update(self.search_params.get("metadata", {}))
+        metadata.update(self.search_params.get("config", {}))
         self.collection.modify(metadata=metadata)
diff --git a/engine/clients/chroma/upload.py b/engine/clients/chroma/upload.py
@@ -22,6 +22,7 @@ def init_client(cls, host, distance, connection_params, upload_params):
 
     @classmethod
     def upload_batch(cls, batch: List[Record]):
+        # assert len(batch) <= cls.client.get_max_batch_size() # commented for performance reasons
         ids, vectors, payloads = [], [], []
         for point in batch:
             ids.append(str(point.id))
diff --git a/engine/servers/chroma-single-node/docker-compose.yaml b/engine/servers/chroma-single-node/docker-compose.yaml
@@ -1,8 +1,8 @@
 services:
   chromadb_bench:
     image: ${CONTAINER_REGISTRY:-docker.io}/chromadb/chroma:0.5.7
-    #volumes:
-    #  - ./chromadb:/chroma/chroma
+    volumes:
+      - ./chromadb:/chroma/chroma
     ports:
       - "8000:8000"
     logging:
@@ -13,7 +13,7 @@ services:
     environment:
       IS_PERSISTENT: TRUE
       ANONYMIZED_TELEMETRY: False
-      CHROMA_WORKERS: 2
+      CHROMA_WORKERS: 1
     deploy:
       resources:
         limits:
diff --git a/experiments/configurations/chroma-single-node.json b/experiments/configurations/chroma-single-node.json
@@ -5,9 +5,7 @@
     "connection_params": {},
     "collection_params": {},
     "search_params": [
-      {
-        "parallel": 8
-      }
+      { "parallel": 8, "config": {} }
     ],
     "upload_params": {
       "parallel": 16,
@@ -19,14 +17,14 @@
     "engine": "chroma",
     "connection_params": {},
     "collection_params": {
-      "metadata": {
+      "config": {
         "hnsw:M": 16,
         "hnsw:construction_ef": 128
       }
     },
     "search_params": [
-      { "parallel": 1, "metadata": {"hnsw:search_ef": 128 } }, { "parallel": 1, "metadata": {"hnsw:search_ef": 256 } }, { "parallel": 1, "metadata": {"hnsw:search_ef": 512 } },
-      { "parallel": 100, "metadata": {"hnsw:search_ef": 128 } }, { "parallel": 100, "metadata": {"hnsw:search_ef": 256 } }, { "parallel": 100, "metadata": {"hnsw:search_ef": 512 } }
+      { "parallel": 1, "config": {"hnsw:search_ef": 128 } }, { "parallel": 1, "config": {"hnsw:search_ef": 256 } }, { "parallel": 1, "config": {"hnsw:search_ef": 512 } },
+      { "parallel": 100, "config": {"hnsw:search_ef": 128 } }, { "parallel": 100, "config": {"hnsw:search_ef": 256 } }, { "parallel": 100, "config": {"hnsw:search_ef": 512 } }
     ],
     "upload_params": { "parallel": 16 }
   },
@@ -35,14 +33,14 @@
     "engine": "chroma",
     "connection_params": {},
     "collection_params": {
-      "metadata": {
+      "config": {
         "hnsw:M": 32,
         "hnsw:construction_ef": 128
       }
     },
     "search_params": [
-      { "parallel": 1, "metadata": {"hnsw:search_ef": 128 } }, { "parallel": 1, "metadata": {"hnsw:search_ef": 256 } }, { "parallel": 1, "metadata": {"hnsw:search_ef": 512 } },
-      { "parallel": 100, "metadata": {"hnsw:search_ef": 128 } }, { "parallel": 100, "metadata": {"hnsw:search_ef": 256 } }, { "parallel": 100, "metadata": {"hnsw:search_ef": 512 } }
+      { "parallel": 1, "config": {"hnsw:search_ef": 128 } }, { "parallel": 1, "config": {"hnsw:search_ef": 256 } }, { "parallel": 1, "config": {"hnsw:search_ef": 512 } },
+      { "parallel": 100, "config": {"hnsw:search_ef": 128 } }, { "parallel": 100, "config": {"hnsw:search_ef": 256 } }, { "parallel": 100, "config": {"hnsw:search_ef": 512 } }
     ],
     "upload_params": { "parallel": 16 }
   },
@@ -51,14 +49,14 @@
     "engine": "chroma",
     "connection_params": {},
     "collection_params": {
-      "metadata": {
+      "config": {
         "hnsw:M": 32,
         "hnsw:construction_ef": 256
       }
     },
     "search_params": [
-      { "parallel": 1, "metadata": {"hnsw:search_ef": 128 } }, { "parallel": 1, "metadata": {"hnsw:search_ef": 256 } }, { "parallel": 1, "metadata": {"hnsw:search_ef": 512 } },
-      { "parallel": 100, "metadata": {"hnsw:search_ef": 128 } }, { "parallel": 100, "metadata": {"hnsw:search_ef": 256 } }, { "parallel": 100, "metadata": {"hnsw:search_ef": 512 } }
+      { "parallel": 1, "config": {"hnsw:search_ef": 128 } }, { "parallel": 1, "config": {"hnsw:search_ef": 256 } }, { "parallel": 1, "config": {"hnsw:search_ef": 512 } },
+      { "parallel": 100, "config": {"hnsw:search_ef": 128 } }, { "parallel": 100, "config": {"hnsw:search_ef": 256 } }, { "parallel": 100, "config": {"hnsw:search_ef": 512 } }
     ],
     "upload_params": { "parallel": 16 }
   },
@@ -67,14 +65,14 @@
     "engine": "chroma",
     "connection_params": {},
     "collection_params": {
-      "metadata": {
+      "config": {
         "hnsw:M": 32,
         "hnsw:construction_ef": 512
       }
     },
     "search_params": [
-      { "parallel": 1, "metadata": {"hnsw:search_ef": 128 } }, { "parallel": 1, "metadata": {"hnsw:search_ef": 256 } }, { "parallel": 1, "metadata": {"hnsw:search_ef": 512 } },
-      { "parallel": 100, "metadata": {"hnsw:search_ef": 128 } }, { "parallel": 100, "metadata": {"hnsw:search_ef": 256 } }, { "parallel": 100, "metadata": {"hnsw:search_ef": 512 } }
+      { "parallel": 1, "config": {"hnsw:search_ef": 128 } }, { "parallel": 1, "config": {"hnsw:search_ef": 256 } }, { "parallel": 1, "config": {"hnsw:search_ef": 512 } },
+      { "parallel": 100, "config": {"hnsw:search_ef": 128 } }, { "parallel": 100, "config": {"hnsw:search_ef": 256 } }, { "parallel": 100, "config": {"hnsw:search_ef": 512 } }
     ],
     "upload_params": { "parallel": 16 }
   },
@@ -83,14 +81,14 @@
     "engine": "chroma",
     "connection_params": {},
     "collection_params": {
-      "metadata": {
+      "config": {
         "hnsw:M": 64,
         "hnsw:construction_ef": 256
       }
     },
     "search_params": [
-      { "parallel": 1, "metadata": {"hnsw:search_ef": 128 } }, { "parallel": 1, "metadata": {"hnsw:search_ef": 256 } }, { "parallel": 1, "metadata": {"hnsw:search_ef": 512 } },
-      { "parallel": 100, "metadata": {"hnsw:search_ef": 128 } }, { "parallel": 100, "metadata": {"hnsw:search_ef": 256 } }, { "parallel": 100, "metadata": {"hnsw:search_ef": 512 } }
+      { "parallel": 1, "config": {"hnsw:search_ef": 128 } }, { "parallel": 1, "config": {"hnsw:search_ef": 256 } }, { "parallel": 1, "config": {"hnsw:search_ef": 512 } },
+      { "parallel": 100, "config": {"hnsw:search_ef": 128 } }, { "parallel": 100, "config": {"hnsw:search_ef": 256 } }, { "parallel": 100, "config": {"hnsw:search_ef": 512 } }
     ],
     "upload_params": { "parallel": 16 }
   },
@@ -99,14 +97,14 @@
     "engine": "chroma",
     "connection_params": {},
     "collection_params": {
-      "metadata": {
+      "config": {
         "hnsw:M": 64,
         "hnsw:construction_ef": 512
       }
     },
     "search_params": [
-      { "parallel": 1, "metadata": {"hnsw:search_ef": 128 } }, { "parallel": 1, "metadata": {"hnsw:search_ef": 256 } }, { "parallel": 1, "metadata": {"hnsw:search_ef": 512 } },
-      { "parallel": 100, "metadata": {"hnsw:search_ef": 128 } }, { "parallel": 100, "metadata": {"hnsw:search_ef": 256 } }, { "parallel": 100, "metadata": {"hnsw:search_ef": 512 } }
+      { "parallel": 1, "config": {"hnsw:search_ef": 128 } }, { "parallel": 1, "config": {"hnsw:search_ef": 256 } }, { "parallel": 1, "config": {"hnsw:search_ef": 512 } },
+      { "parallel": 100, "config": {"hnsw:search_ef": 128 } }, { "parallel": 100, "config": {"hnsw:search_ef": 256 } }, { "parallel": 100, "config": {"hnsw:search_ef": 512 } }
     ],
     "upload_params": { "parallel": 16 }
   }
diff --git a/pyproject.toml b/pyproject.toml
@@ -20,7 +20,7 @@ opensearch-py = "^2.3.2"
 tqdm = "^4.66.1"
 psycopg = {extras = ["binary"], version = "^3.1.17"}
 pgvector = "^0.2.4"
-chromadb = "0.5.7"
+chromadb = { version = ">=0.5.5,<0.6.0" }
 
 [tool.poetry.dev-dependencies]
 pre-commit = "^2.20.0"

Original file line number	Diff line number	Diff line change
`@@ -37,7 +37,7 @@ def recreate(self, dataset: Dataset, collection_params):`
`37`	`37`	`params = self.collection_params`
`38`	`38`	`params["metadata"] = dict(`
`39`	`39`	`{"hnsw:space": self.DISTANCE_MAPPING.get(dataset.config.distance)},`
`40`		`- **params.get("metadata", {}),`
	`40`	`+ **params.get("config", {}),`
`41`	`41`	`)`
`42`	`42`	`self.client.create_collection(`
`43`	`43`	`name=CHROMA_COLLECTION_NAME,`