milvus-io
/
milvus
mirror of https://github.com/milvus-io/milvus.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity
milvus/tests/milvus_python_test/utils.py

697 lines
20 KiB
Python
Raw Blame History

import os
import sys
import random
import string
import struct
import logging
import time, datetime
import copy
import numpy as np
from milvus import Milvus, IndexType, MetricType
port = 19530
epsilon = 0.000001
default_flush_interval = 1
big_flush_interval = 1000
all_index_types = [
IndexType.FLAT,
IndexType.IVFLAT,
IndexType.IVF_SQ8,
IndexType.IVF_SQ8H,
IndexType.IVF_PQ,
IndexType.HNSW,
IndexType.RNSG,
IndexType.ANNOY
]
def get_milvus(host, port, uri=None, handler=None, **kwargs):
if handler is None:
handler = "GRPC"
try_connect = kwargs.get("try_connect", True)
if uri is not None:
milvus = Milvus(uri=uri, handler=handler, try_connect=try_connect)
else:
milvus = Milvus(host=host, port=port, handler=handler, try_connect=try_connect)
return milvus
def disable_flush(connect):
status, reply = connect.set_config("storage", "auto_flush_interval", big_flush_interval)
assert status.OK()
def enable_flush(connect):
# reset auto_flush_interval=1
status, reply = connect.set_config("storage", "auto_flush_interval", default_flush_interval)
assert status.OK()
status, config_value = connect.get_config("storage", "auto_flush_interval")
assert status.OK()
assert config_value == str(default_flush_interval)
def gen_inaccuracy(num):
return num / 255.0
def gen_vectors(num, dim):
return [[random.random() for _ in range(dim)] for _ in range(num)]
def gen_binary_vectors(num, dim):
raw_vectors = []
binary_vectors = []
for i in range(num):
raw_vector = [random.randint(0, 1) for i in range(dim)]
raw_vectors.append(raw_vector)
binary_vectors.append(bytes(np.packbits(raw_vector, axis=-1).tolist()))
return raw_vectors, binary_vectors
def jaccard(x, y):
x = np.asarray(x, np.bool)
y = np.asarray(y, np.bool)
return 1 - np.double(np.bitwise_and(x, y).sum()) / np.double(np.bitwise_or(x, y).sum())
def hamming(x, y):
x = np.asarray(x, np.bool)
y = np.asarray(y, np.bool)
return np.bitwise_xor(x, y).sum()
def tanimoto(x, y):
x = np.asarray(x, np.bool)
y = np.asarray(y, np.bool)
return -np.log2(np.double(np.bitwise_and(x, y).sum()) / np.double(np.bitwise_or(x, y).sum()))
def substructure(x, y):
x = np.asarray(x, np.bool)
y = np.asarray(y, np.bool)
return 1 - np.double(np.bitwise_and(x, y).sum()) / np.count_nonzero(y)
def superstructure(x, y):
x = np.asarray(x, np.bool)
y = np.asarray(y, np.bool)
return 1 - np.double(np.bitwise_and(x, y).sum()) / np.count_nonzero(x)
def gen_binary_sub_vectors(vectors, length):
raw_vectors = []
binary_vectors = []
dim = len(vectors[0])
for i in range(length):
raw_vector = [0 for i in range(dim)]
vector = vectors[i]
for index, j in enumerate(vector):
if j == 1:
raw_vector[index] = 1
raw_vectors.append(raw_vector)
binary_vectors.append(bytes(np.packbits(raw_vector, axis=-1).tolist()))
return raw_vectors, binary_vectors
def gen_binary_super_vectors(vectors, length):
raw_vectors = []
binary_vectors = []
dim = len(vectors[0])
for i in range(length):
cnt_1 = np.count_nonzero(vectors[i])
raw_vector = [1 for i in range(dim)]
raw_vectors.append(raw_vector)
binary_vectors.append(bytes(np.packbits(raw_vector, axis=-1).tolist()))
return raw_vectors, binary_vectors
def gen_single_vector(dim):
return [[random.random() for _ in range(dim)]]
def gen_vector(nb, d, seed=np.random.RandomState(1234)):
xb = seed.rand(nb, d).astype("float32")
return xb.tolist()
def gen_unique_str(str_value=None):
prefix = "".join(random.choice(string.ascii_letters + string.digits) for _ in range(8))
return "test_" + prefix if str_value is None else str_value + "_" + prefix
def gen_long_str(num):
string = ''
for _ in range(num):
char = random.choice('tomorrow')
string += char
def gen_invalid_ips():
ips = [
# "255.0.0.0",
# "255.255.0.0",
# "255.255.255.0",
# "255.255.255.255",
"127.0.0",
# "123.0.0.2",
"12-s",
" ",
"12 s",
"BB。A",
" siede ",
"(mn)",
"中文",
"a".join("a" for _ in range(256))
]
return ips
def gen_invalid_ports():
ports = [
# empty
" ",
-1,
# too big port
100000,
# not correct port
39540,
"BB。A",
" siede ",
"(mn)",
"中文"
]
return ports
def gen_invalid_uris():
ip = None
uris = [
" ",
"中文",
# invalid protocol
# "tc://%s:%s" % (ip, port),
# "tcp%s:%s" % (ip, port),
# # invalid port
# "tcp://%s:100000" % ip,
# "tcp://%s: " % ip,
# "tcp://%s:19540" % ip,
# "tcp://%s:-1" % ip,
# "tcp://%s:string" % ip,
# invalid ip
"tcp:// :19530",
# "tcp://123.0.0.1:%s" % port,
"tcp://127.0.0:19530",
# "tcp://255.0.0.0:%s" % port,
# "tcp://255.255.0.0:%s" % port,
# "tcp://255.255.255.0:%s" % port,
# "tcp://255.255.255.255:%s" % port,
"tcp://\n:19530",
]
return uris
def gen_invalid_collection_names():
collection_names = [
"12-s",
" ",
# "",
# None,
"12 s",
"BB。A",
"c|c",
" siede ",
"(mn)",
"pip+",
"=c",
"中文",
"a".join("a" for i in range(256))
]
return collection_names
def gen_invalid_top_ks():
top_ks = [
0,
-1,
None,
[1,2,3],
(1,2),
{"a": 1},
" ",
"",
"String",
"12-s",
"BB。A",
" siede ",
"(mn)",
"pip+",
"=c",
"中文",
"a".join("a" for i in range(256))
]
return top_ks
def gen_invalid_dims():
dims = [
0,
-1,
100001,
1000000000000001,
None,
False,
[1,2,3],
(1,2),
{"a": 1},
" ",
"",
"String",
"12-s",
"BB。A",
" siede ",
"(mn)",
"pip+",
"=c",
"中文",
"a".join("a" for i in range(256))
]
return dims
def gen_invalid_file_sizes():
file_sizes = [
0,
-1,
1000000000000001,
None,
False,
[1,2,3],
(1,2),
{"a": 1},
" ",
"",
"String",
"12-s",
"BB。A",
" siede ",
"(mn)",
"pip+",
"=c",
"中文",
"a".join("a" for i in range(256))
]
return file_sizes
def gen_invalid_index_types():
invalid_types = [
0,
-1,
100,
1000000000000001,
# None,
False,
[1,2,3],
(1,2),
{"a": 1},
" ",
"",
"String",
"12-s",
"BB。A",
" siede ",
"(mn)",
"pip+",
"=c",
"中文",
"a".join("a" for i in range(256))
]
return invalid_types
def gen_invalid_params():
params = [
9999999999,
-1,
# None,
[1,2,3],
(1,2),
{"a": 1},
" ",
"",
"String",
"12-s",
"BB。A",
" siede ",
"(mn)",
"pip+",
"=c",
"中文"
]
return params
def gen_invalid_nprobes():
nprobes = [
0,
-1,
1000000000000001,
None,
[1,2,3],
(1,2),
{"a": 1},
" ",
"",
"String",
"12-s",
"BB。A",
" siede ",
"(mn)",
"pip+",
"=c",
"中文"
]
return nprobes
def gen_invalid_metric_types():
metric_types = [
0,
-1,
1000000000000001,
# None,
[1,2,3],
(1,2),
{"a": 1},
" ",
"",
"String",
"12-s",
"BB。A",
" siede ",
"(mn)",
"pip+",
"=c",
"中文"
]
return metric_types
def gen_invalid_vectors():
invalid_vectors = [
"1*2",
[],
[1],
[1,2],
[" "],
['a'],
[None],
None,
(1,2),
{"a": 1},
" ",
"",
"String",
"12-s",
"BB。A",
" siede ",
"(mn)",
"pip+",
"=c",
"中文",
"a".join("a" for i in range(256))
]
return invalid_vectors
def gen_invalid_vector_ids():
invalid_vector_ids = [
1.0,
-1.0,
None,
# int 64
10000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000000,
" ",
"",
"String",
"BB。A",
" siede ",
"(mn)",
"=c",
"中文",
]
return invalid_vector_ids
def gen_invalid_cache_config():
invalid_configs = [
0,
-1,
9223372036854775808,
[1,2,3],
(1,2),
{"a": 1},
" ",
"",
"String",
"12-s",
"BB。A",
" siede ",
"(mn)",
"pip+",
"=c",
"中文",
"'123'",
"さようなら"
]
return invalid_configs
def gen_invalid_gpu_config():
invalid_configs = [
-1,
[1,2,3],
(1,2),
{"a": 1},
" ",
"",
"String",
"12-s",
"BB。A",
" siede ",
"(mn)",
"pip+",
"=c",
"中文",
"'123'",
]
return invalid_configs
def gen_invaild_search_params():
invalid_search_key = 100
search_params = []
for index_type in all_index_types:
if index_type == IndexType.FLAT:
continue
search_params.append({"index_type": index_type, "search_param": {"invalid_key": invalid_search_key}})
if index_type in [IndexType.IVFLAT, IndexType.IVF_SQ8, IndexType.IVF_SQ8H, IndexType.IVF_PQ]:
for nprobe in gen_invalid_params():
ivf_search_params = {"index_type": index_type, "search_param": {"nprobe": nprobe}}
search_params.append(ivf_search_params)
elif index_type == IndexType.HNSW:
for ef in gen_invalid_params():
hnsw_search_param = {"index_type": index_type, "search_param": {"ef": ef}}
search_params.append(hnsw_search_param)
elif index_type == IndexType.RNSG:
for search_length in gen_invalid_params():
nsg_search_param = {"index_type": index_type, "search_param": {"search_length": search_length}}
search_params.append(nsg_search_param)
search_params.append({"index_type": index_type, "search_param": {"invalid_key": 100}})
elif index_type == IndexType.ANNOY:
for search_k in gen_invalid_params():
if isinstance(search_k, int):
continue
annoy_search_param = {"index_type": index_type, "search_param": {"search_k": search_k}}
search_params.append(annoy_search_param)
return search_params
def gen_invalid_index():
index_params = []
for index_type in gen_invalid_index_types():
index_param = {"index_type": index_type, "index_param": {"nlist": 1024}}
index_params.append(index_param)
for nlist in gen_invalid_params():
index_param = {"index_type": IndexType.IVFLAT, "index_param": {"nlist": nlist}}
index_params.append(index_param)
for M in gen_invalid_params():
index_param = {"index_type": IndexType.HNSW, "index_param": {"M": M, "efConstruction": 100}}
index_params.append(index_param)
for efConstruction in gen_invalid_params():
index_param = {"index_type": IndexType.HNSW, "index_param": {"M": 16, "efConstruction": efConstruction}}
index_params.append(index_param)
for search_length in gen_invalid_params():
index_param = {"index_type": IndexType.RNSG,
"index_param": {"search_length": search_length, "out_degree": 40, "candidate_pool_size": 50,
"knng": 100}}
index_params.append(index_param)
for out_degree in gen_invalid_params():
index_param = {"index_type": IndexType.RNSG,
"index_param": {"search_length": 100, "out_degree": out_degree, "candidate_pool_size": 50,
"knng": 100}}
index_params.append(index_param)
for candidate_pool_size in gen_invalid_params():
index_param = {"index_type": IndexType.RNSG, "index_param": {"search_length": 100, "out_degree": 40,
"candidate_pool_size": candidate_pool_size,
"knng": 100}}
index_params.append(index_param)
index_params.append({"index_type": IndexType.IVF_FLAT, "index_param": {"invalid_key": 1024}})
index_params.append({"index_type": IndexType.HNSW, "index_param": {"invalid_key": 16, "efConstruction": 100}})
index_params.append({"index_type": IndexType.RNSG,
"index_param": {"invalid_key": 100, "out_degree": 40, "candidate_pool_size": 300,
"knng": 100}})
for invalid_n_trees in gen_invalid_params():
index_params.append({"index_type": IndexType.ANNOY, "index_param": {"n_trees": invalid_n_trees}})
return index_params
def gen_index():
nlists = [1, 1024, 16384]
pq_ms = [128, 64, 32, 16, 8, 4]
Ms = [5, 24, 48]
efConstructions = [100, 300, 500]
search_lengths = [10, 100, 300]
out_degrees = [5, 40, 300]
candidate_pool_sizes = [50, 100, 300]
knngs = [5, 100, 300]
index_params = []
for index_type in all_index_types:
if index_type == IndexType.FLAT:
index_params.append({"index_type": index_type, "index_param": {"nlist": 1024}})
elif index_type in [IndexType.IVFLAT, IndexType.IVF_SQ8, IndexType.IVF_SQ8H]:
ivf_params = [{"index_type": index_type, "index_param": {"nlist": nlist}} \
for nlist in nlists]
index_params.extend(ivf_params)
elif index_type == IndexType.IVF_PQ:
ivf_pq_params = [{"index_type": index_type, "index_param": {"nlist": nlist, "m": m}} \
for nlist in nlists \
for m in pq_ms]
index_params.extend(ivf_pq_params)
elif index_type == IndexType.HNSW:
hnsw_params = [{"index_type": index_type, "index_param": {"M": M, "efConstruction": efConstruction}} \
for M in Ms \
for efConstruction in efConstructions]
index_params.extend(hnsw_params)
elif index_type == IndexType.RNSG:
nsg_params = [{"index_type": index_type,
"index_param": {"search_length": search_length, "out_degree": out_degree,
"candidate_pool_size": candidate_pool_size, "knng": knng}} \
for search_length in search_lengths \
for out_degree in out_degrees \
for candidate_pool_size in candidate_pool_sizes \
for knng in knngs]
index_params.extend(nsg_params)
return index_params
def gen_simple_index():
params = [
{"nlist": 1024},
{"nlist": 1024},
{"nlist": 1024},
{"nlist": 1024},
{"nlist": 1024, "m": 16},
{"M": 48, "efConstruction": 500},
{"search_length": 50, "out_degree": 40, "candidate_pool_size": 100, "knng": 50},
{"n_trees": 4}
]
index_params = []
for i in range(len(all_index_types)):
index_params.append({"index_type": all_index_types[i], "index_param": params[i]})
return index_params
def get_search_param(index_type):
if index_type in [IndexType.FLAT, IndexType.IVFLAT, IndexType.IVF_SQ8, IndexType.IVF_SQ8H, IndexType.IVF_PQ]:
return {"nprobe": 32}
elif index_type == IndexType.HNSW:
return {"ef": 64}
elif index_type == IndexType.RNSG:
return {"search_length": 100}
elif index_type == IndexType.ANNOY:
return {"search_k": 100}
else:
logging.getLogger().info("Invalid index_type.")
def assert_has_collection(conn, collection_name):
status, ok = conn.has_collection(collection_name)
return status.OK() and ok
def assert_equal_vector(v1, v2):
if len(v1) != len(v2):
assert False
for i in range(len(v1)):
assert abs(v1[i] - v2[i]) < epsilon
def restart_server(helm_release_name):
res = True
timeout = 120
from kubernetes import client, config
client.rest.logger.setLevel(logging.WARNING)
namespace = "milvus"
# service_name = "%s.%s.svc.cluster.local" % (helm_release_name, namespace)
config.load_kube_config()
v1 = client.CoreV1Api()
pod_name = None
# config_map_names = v1.list_namespaced_config_map(namespace, pretty='true')
# body = {"replicas": 0}
pods = v1.list_namespaced_pod(namespace)
for i in pods.items:
if i.metadata.name.find(helm_release_name) != -1 and i.metadata.name.find("mysql") == -1:
pod_name = i.metadata.name
break
# v1.patch_namespaced_config_map(config_map_name, namespace, body, pretty='true')
# status_res = v1.read_namespaced_service_status(helm_release_name, namespace, pretty='true')
# print(status_res)
if pod_name is not None:
try:
v1.delete_namespaced_pod(pod_name, namespace)
except Exception as e:
logging.error(str(e))
logging.error("Exception when calling CoreV1Api->delete_namespaced_pod")
res = False
return res
time.sleep(5)
# check if restart successfully
pods = v1.list_namespaced_pod(namespace)
for i in pods.items:
pod_name_tmp = i.metadata.name
if pod_name_tmp.find(helm_release_name) != -1:
logging.debug(pod_name_tmp)
start_time = time.time()
while time.time() - start_time > timeout:
status_res = v1.read_namespaced_pod_status(pod_name_tmp, namespace, pretty='true')
if status_res.status.phase == "Running":
break
time.sleep(1)
if time.time() - start_time > timeout:
logging.error("Restart pod: %s timeout" % pod_name_tmp)
res = False
return res
else:
logging.error("Pod: %s not found" % helm_release_name)
res = False
return res
Reference in New Issue View Git Blame Copy Permalink