def _add():
milvus = Milvus()
status = milvus.connect(**server_config)
vectors = _generate_vectors(128, 10000)
print('\n\tPID: {}, insert {} vectors'.format(os.getpid(), 10000))
status, _ = milvus.add_vectors(_table_name, vectors)
milvus.disconnect()
class TestConnect:
milvus = Milvus()
DefaultConfig.THRIFTCLIENT_FRAMED = True
cnn_status = milvus.connect(uri='tcp://localhost:9090')
def test_false_connect1(self):
_, res = self.milvus.server_version()
assert res == '0.3.1'
def __init__(self):
try:
self.client = Milvus(host=MILVUS_HOST, port=MILVUS_PORT)
LOGGER.debug(
"Successfully connect to Milvus with IP:{} and PORT:{}".format(
MILVUS_HOST, MILVUS_PORT))
except Exception as e:
LOGGER.error("Failed to connect Milvus: {}".format(e))
sys.exit(1)
def main():
# connect_milvus_server()
milvus = Milvus(host=SERVER_ADDR, port=SERVER_PORT)
create_milvus_collection(milvus)
build_collection(milvus)
conn = connect_postgres_server()
cur = conn.cursor()
create_pg_table(conn, cur)
count = 0
while count < (VEC_NUM // BASE_LEN):
vectors = load_bvecs_data(FILE_PATH, BASE_LEN, count)
vectors_ids = [
id for id in range(count * BASE_LEN, (count + 1) * BASE_LEN)
]
sex = random.choice(['female', 'male'])
get_time = fake.past_datetime(start_date="-120d", tzinfo=None)
is_glasses = random.choice(['True', 'False'])
line = str(vectors_ids[i]) + "|" + sex + "|'" + str(
get_time) + "'|" + str(is_glasses) + "\n"
hybrid_entities = [{
"name": "sex",
"values": sex,
"type": DataType.INT32
}, {
"name": "is_glasses",
"values": get_time,
"type": DataType.INT32
}, {
"name": "get_time",
"values": is_glasses,
"type": DataType.INT32
}, {
"name": "Vec",
"values": vectors,
"type": DataType.FLOAT_VECTOR
}]
time_start = time.time()
ids = milvus.bulk_insert(MILVUS_collection,
hybrid_entities,
ids=vectors_ids)
time_end = time.time()
print(count, "insert milvue time: ", time_end - time_start)
#print(format(ids))
# print(count)
time_start = time.time()
with open(fname, 'w+') as f:
for i in range(len(ids)):
res = client.get_entity_by_id(collection_name, ids)
line = str(res.id) + "|" + res.get(sex) + "|'" + str(
res.get(get_time)) + "'|" + str(res.get(is_glasses)) + "\n"
f.write(line)
copy_data_to_pg(conn, cur)
time_end = time.time()
print(count, "insert pg time: ", time_end - time_start)
count = count + 1
build_pg_index(conn, cur)
def del_vectors(collection_name, ids):
try:
milvus = Milvus(host=MILVUS_ADDR, port=MILVUS_PORT)
milvus.delete_entity_by_id(collection_name=collection_name,
id_array=ids)
except Exception as e:
err_msg = "There was some error when delete vectors"
logger.error(f"{err_msg} : {str(e)}", exc_info=True)
raise MilvusError(err_msg, e)
def test_connect_correct_ip_port(self, args):
'''
target: test connect with corrent ip and port value
method: set correct ip and port
expected: connected is True
'''
milvus = Milvus()
milvus.connect(host=args["ip"], port=args["port"])
assert milvus.connected()
def _test_connect_ip_localhost(self, args):
'''
target: test connect with ip value: localhost
method: set host localhost
expected: connected is True
'''
milvus = Milvus()
milvus.connect(host='localhost', port=args["port"])
assert milvus.connected()
def connect(self):
from milvus import Milvus
if self.milvus_client is None or not self.milvus_client.server_status(
)[0].OK():
self.logger.info(
f'Setting connection to Milvus Server at {self.host}:{self.port}'
)
self.milvus_client = Milvus(self.host, self.port)
return self
def milvus_test(usr_features, IS_INFER, mov_features=None, ids=None):
_HOST = '127.0.0.1'
_PORT = '19530' # default value
table_name = 'recommender_demo'
milvus = Milvus()
param = {'host': _HOST, 'port': _PORT}
status = milvus.connect(**param)
if status.OK():
print("Server connected.")
else:
print("Server connect fail.")
sys.exit(1)
if IS_INFER:
status = milvus.drop_table(table_name)
time.sleep(3)
status, ok = milvus.has_table(table_name)
if not ok:
if mov_features is None:
print("Insert vectors is none!")
sys.exit(1)
param = {
'table_name': table_name,
'dimension': 200,
'index_file_size': 1024, # optional
'metric_type': MetricType.IP # optional
}
print(milvus.create_table(param))
insert_vectors = normaliz_data(mov_features)
status, ids = milvus.insert(table_name=table_name,
records=insert_vectors,
ids=ids)
time.sleep(1)
status, result = milvus.count_table(table_name)
print("rows in table recommender_demo:", result)
search_vectors = normaliz_data(usr_features)
param = {
'table_name': table_name,
'query_records': search_vectors,
'top_k': 5,
'nprobe': 16
}
time1 = time.time()
status, results = milvus.search_vectors(**param)
time2 = time.time()
print("Top\t", "Ids\t", "Title\t", "Score")
for i, re in enumerate(results[0]):
title = paddle.dataset.movielens.movie_info()[int(re.id)].title
print(i, "\t", re.id, "\t", title, "\t", float(re.distance) * 5)
def get_milvus():
milvus = Milvus()
param = {'host': _HOST, 'port': _PORT}
status = milvus.connect(**param)
if status.OK():
print("Server connected.")
else:
print("Server connect fail.")
sys.exit(1)
return milvus
def validate_insert(_table_name):
milvus = Milvus()
milvus.connect(host="localhost", port="19530")
status, count = milvus.get_table_row_count(_table_name)
assert count == vector_num * process_num, f"Error: validate insert not pass: "******"{vector_num * process_num} expected but {count} instead!"
milvus.disconnect()
def _add():
milvus = Milvus()
status = milvus.connect()
if not status.OK:
print(f'PID: {os.getpid()}, connect failed')
status, _ = milvus.add_vectors(_table_name, vectors)
milvus.disconnect()
def _add():
milvus = Milvus(**server_config)
vectors = _generate_vectors(128, 10000)
print('\n\tPID: {}, insert {} vectors'.format(os.getpid(), 10000))
status, _ = milvus.insert(_collection_name, vectors)
if not status.OK():
print("PID {} insert failed: {}".format(os.getpid(),
status.message))
milvus.close()
def validate_insert(_collection_name):
milvus = Milvus(**server_config)
milvus.flush([_collection_name])
status, count = milvus.count_entities(_collection_name)
assert count == 10 * 10000, "Insert validate fail. Vectors num is not matched."
# drop collcetion
print("Drop collection ...")
milvus.drop_collection(_collection_name)
milvus.close()
def test_hooks(self, gip):
with Milvus(*gip) as client:
class FakeSerchHook(BaseSearchHook):
def pre_search(self, *args, **kwargs):
print("Before search ...")
def aft_search(self, *args, **kwargs):
print("Search done ...")
client.set_hook(search=FakeSerchHook())
def _test_search_during_createIndex(self, args):
loops = 10000
table = gen_unique_str()
query_vecs = [vectors[0], vectors[1]]
uri = "tcp://%s:%s" % (args["ip"], args["port"])
id_0 = 0; id_1 = 0
milvus_instance = Milvus()
milvus_instance.connect(uri=uri)
milvus_instance.create_table({'table_name': table,
'dimension': dim,
'index_file_size': index_file_size,
'metric_type': MetricType.L2})
for i in range(10):
status, ids = milvus_instance.add_vectors(table, vectors)
# logging.getLogger().info(ids)
if i == 0:
id_0 = ids[0]; id_1 = ids[1]
def create_index(milvus_instance):
logging.getLogger().info("In create index")
status = milvus_instance.create_index(table, index_params)
logging.getLogger().info(status)
status, result = milvus_instance.describe_index(table)
logging.getLogger().info(result)
def add_vectors(milvus_instance):
logging.getLogger().info("In add vectors")
status, ids = milvus_instance.add_vectors(table, vectors)
logging.getLogger().info(status)
def search(milvus_instance):
logging.getLogger().info("In search vectors")
for i in range(loops):
status, result = milvus_instance.search_vectors(table, top_k, nprobe, query_vecs)
logging.getLogger().info(status)
assert result[0][0].id == id_0
assert result[1][0].id == id_1
milvus_instance = Milvus()
milvus_instance.connect(uri=uri)
p_search = Process(target=search, args=(milvus_instance, ))
p_search.start()
milvus_instance = Milvus()
milvus_instance.connect(uri=uri)
p_create = Process(target=add_vectors, args=(milvus_instance, ))
p_create.start()
p_create.join()
def test_connect_param_priority_no_port(self, args):
'''
target: both host_ip_port / uri are both given, if port is null, use the uri params
method: port set "", check if wrong uri connection is ok
expected: connect raise an exception and connected is false
'''
milvus = Milvus()
uri_value = "tcp://%s:19540" % args["ip"]
milvus.connect(host=args["ip"], port="", uri=uri_value)
assert milvus.connected()
def test_connect_uri(self, args):
'''
target: test connect with correct uri
method: uri format and value are both correct
expected: connected is True
'''
milvus = Milvus()
uri_value = "tcp://%s:%s" % (args["ip"], args["port"])
milvus.connect(uri=uri_value)
assert milvus.connected()
def test_connect_connected(self, args):
'''
target: test connect and disconnect with corrent ip and port value, assert connected value
method: set correct ip and port
expected: connected is False
'''
milvus = Milvus()
milvus.connect(host=args["ip"], port=args["port"])
milvus.disconnect()
assert not milvus.connected()
def test_connect_wrong_uri_wrong_port_null(self, args):
'''
target: test uri connect with port value wouldn't connected
method: set uri port null
expected: connected is True
'''
milvus = Milvus()
uri_value = "tcp://%s:" % args["ip"]
with pytest.raises(Exception) as e:
milvus.connect(uri=uri_value, timeout=1)
def create():
_HOST = 'localhost'
_PORT = '19530'
_collection_name = 'chs_stars_faces_512'
_DIM = 512 # dimension of vector
_INDEX_FILE_SIZE = 256 # max file size of stored index
milvus = Milvus(_HOST, _PORT)
param = {
'collection_name': _collection_name,
'dimension': _DIM,
'index_file_size': _INDEX_FILE_SIZE, # optional
'metric_type': MetricType.IP # optional
}
milvus.create_collection(param)
index_param = {
'nlist': 2048 # 推荐 4 * sqrt(n)
}
status = milvus.create_index(_collection_name, IndexType.IVF_SQ8, index_param)
# with open("chs_stars_features_pca.pickle", "rb") as f:
# pca = pickle.load(f)
#
# with open("../chs_stars_features_pca.csv", "w") as fw, open("../chs_stars_features.csv", "r") as fr:
# reader = csv.reader(fr)
# writer = csv.writer(fw)
# for index, line in enumerate(tqdm(reader)):
# star, fname, features = line
# features = np.array(json.loads(features))
# features = np.resize(features, (1, 512))
# features = normalize(features)
# features = pca.transform(features).squeeze()
# status, ids = milvus.insert(collection_name=_collection_name, records=[features.tolist()], ids=[index])
# if not status.OK():
# print(status)
# continue
# writer.writerow([index, star, fname, features])
with open("../chs_stars_labels.csv", "w") as fw, open("../chs_stars_features.csv", "r") as fr:
reader = csv.reader(fr)
writer = csv.writer(fw)
for index, line in enumerate(tqdm(reader)):
star, fname, features = line
# features = np.array(json.loads(features))
# features = np.resize(features, (1, 512))
#features = normalize(features)
features = json.loads(features)
status, ids = milvus.insert(collection_name=_collection_name, records=[features], ids=[index])
if not status.OK():
print(status)
continue
writer.writerow([index, star, fname])
def _test_add_vector_with_multiprocessing(self, args):
'''
target: test add vectors, with multi processes
method: 10 processed add vectors concurrently
expected: status ok and result length is equal to the length off added vectors
'''
table = gen_unique_str("test_add_vector_with_multiprocessing")
uri = "tcp://%s:%s" % (args["ip"], args["port"])
param = {
'table_name': table,
'dimension': dim,
'index_file_size': index_file_size
}
# create table
milvus = Milvus()
milvus.connect(uri=uri)
milvus.create_table(param)
vector = gen_single_vector(dim)
process_num = 4
loop_num = 10
processes = []
# with dependent connection
def add(milvus):
i = 0
while i < loop_num:
status, ids = milvus.add_vectors(table, vector)
i = i + 1
for i in range(process_num):
milvus = Milvus()
milvus.connect(uri=uri)
p = Process(target=add, args=(milvus, ))
processes.append(p)
p.start()
time.sleep(0.2)
for p in processes:
p.join()
time.sleep(3)
status, count = milvus.get_table_row_count(table)
assert count == process_num * loop_num
def init_client(self):
self.client = Milvus()
try:
status = self.client.connect(host=self.host,
port=self.port,
timeout=settings.TIMEOUT)
except Exception as exc:
raise SDKClientConnectionException(str(exc))
if status != Status.SUCCESS:
raise SDKClientConnectionException(str(status))
def del_milvus_collection(name):
try:
milvus = Milvus(host=MILVUS_ADDR, port=MILVUS_PORT)
res = milvus.drop_collection(collection_name=name)
if not res.OK():
raise MilvusError(
"There was some error when drop milvus collection", res)
except Exception as e:
err_msg = "There was some error when delete milvus collection"
logger.error(f"{err_msg} : {str(e)}", exc_info=True)
raise MilvusError(err_msg, e)
def __init__(self, name, uri, max_retry=1, error_handlers=None, **kwargs):
self.name = name
self.uri = uri
self.max_retry = max_retry
self.retried = 0
self.conn = Milvus()
self.error_handlers = [] if not error_handlers else error_handlers
self.on_retry_func = kwargs.get('on_retry_func', None)
# define search hook
self.conn.set_hook(search_in_file=Searchook())
def test_connect_wrong_ip_null(self, args):
'''
target: test connect with wrong ip value
method: set host null
expected: not use default ip, connected is False
'''
milvus = Milvus()
ip = ""
with pytest.raises(Exception) as e:
milvus.connect(host=ip, port=args["port"], timeout=1)
assert not milvus.connected()
def test_connect_with_invalid_uri(self, get_invalid_uri):
'''
target: test uri connect with invalid uri value
method: set port in gen_invalid_uris
expected: connected is False
'''
milvus = Milvus()
uri_value = get_invalid_uri
with pytest.raises(Exception) as e:
milvus.connect(uri=uri_value, timeout=1)
assert not milvus.connected()
def test_connect_with_invalid_port(self, args, get_invalid_port):
'''
target: test ip:port connect with invalid port value
method: set port in gen_invalid_ports
expected: connected is False
'''
milvus = Milvus()
port = get_invalid_port
with pytest.raises(Exception) as e:
milvus.connect(host=args["ip"], port=port, timeout=1)
assert not milvus.connected()
def test_connect_param_priority_both_hostip_uri(self, args):
'''
target: both host_ip_port / uri are both given, and not null, use the uri params
method: check if wrong uri connection is ok
expected: connect raise an exception and connected is false
'''
milvus = Milvus()
uri_value = "tcp://%s:%s" % (args["ip"], args["port"])
with pytest.raises(Exception) as e:
milvus.connect(host=args["ip"], port=19540, uri=uri_value, timeout=1)
assert not milvus.connected()
def test_connect_param_priority_uri(self, args):
'''
target: both host_ip_port / uri are both given, if host is null, use the uri params
method: host set "", check if correct uri connection is ok
expected: connected is False
'''
milvus = Milvus()
uri_value = "tcp://%s:%s" % (args["ip"], args["port"])
with pytest.raises(Exception) as e:
milvus.connect(host="", port=args["port"], uri=uri_value, timeout=1)
assert not milvus.connected()
