mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-05 05:18:52 +08:00
65 lines
2.3 KiB
Python
65 lines
2.3 KiB
Python
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 |