chore: import upstream snapshot with attribution

2026-07-13 12:31:17 +08:00
commit 498b235461
5446 changed files with 2748612 additions and 0 deletions
@@ -0,0 +1,65 @@
+import json
+from utils.util_log import test_log as log
+
+all_index_types = ["FLAT", "IVF_FLAT", "IVF_SQ8", "IVF_PQ", "HNSW", "BIN_FLAT", "BIN_IVF_FLAT"]
+
+default_index_params = [{"nlist": 128}, {"nlist": 128}, {"nlist": 128}, {"nlist": 128, "m": 16, "nbits": 8},
+                        {"M": 48, "efConstruction": 500}, {"nlist": 128}, {"nlist": 128}]
+
+index_params_map = dict(zip(all_index_types, default_index_params))
+
+def gen_index_param(index_type):
+    metric_type = "L2"
+    if "BIN" in index_type:
+        metric_type = "HAMMING"
+    index_param = {
+        "index_type": index_type,
+        "params": index_params_map[index_type],
+        "metric_type": metric_type
+    }
+    return index_param
+
+
+def gen_search_param(index_type, metric_type="L2"):
+    search_params = []
+    if index_type in ["FLAT", "IVF_FLAT", "IVF_SQ8", "IVF_PQ"]:
+        for nprobe in [10]:
+            ivf_search_params = {"metric_type": metric_type, "params": {"nprobe": nprobe}}
+            search_params.append(ivf_search_params)
+    elif index_type in ["BIN_FLAT", "BIN_IVF_FLAT"]:
+        for nprobe in [10]:
+            bin_search_params = {"metric_type": "HAMMING", "params": {"nprobe": nprobe}}
+            search_params.append(bin_search_params)
+    elif index_type in ["HNSW"]:
+        for ef in [64]:
+            hnsw_search_param = {"metric_type": metric_type, "params": {"ef": ef}}
+            search_params.append(hnsw_search_param)
+    elif index_type == "ANNOY":
+        for search_k in [1000]:
+            annoy_search_param = {"metric_type": metric_type, "params": {"search_k": search_k}}
+            search_params.append(annoy_search_param)
+    else:
+        print("Invalid index_type.")
+        raise Exception("Invalid index_type.")
+    return search_params
+
+
+def get_deploy_test_collections():
+    try:
+        with open("/tmp/ci_logs/deploy_test_all_collections.json", "r") as f:
+            data = json.load(f)
+            collections = data["all"]
+    except Exception as e:
+        log.error(f"get_all_collections error: {e}")
+        return []
+    return collections
+
+def get_chaos_test_collections():
+    try:
+        with open("/tmp/ci_logs/chaos_test_all_collections.json", "r") as f:
+            data = json.load(f)
+            collections = data["all"]
+    except Exception as e:
+        log.error(f"get_all_collections error: {e}")
+        return []
+    return collections