diff --git a/tests/python_client/customize/test_customize_segment_size.py b/tests/python_client/customize/test_customize_segment_size.py index 347715ad3c..adf39b46e9 100644 --- a/tests/python_client/customize/test_customize_segment_size.py +++ b/tests/python_client/customize/test_customize_segment_size.py @@ -4,6 +4,7 @@ import time from pymilvus import connections from utils.util_log import test_log as log from base.collection_wrapper import ApiCollectionWrapper +from base.utility_wrapper import ApiUtilityWrapper from common import common_func as cf from common import common_type as ct from milvus_operator import MilvusOperator @@ -11,7 +12,6 @@ from common.milvus_sys import MilvusSys from common.common_type import CaseLabel -customize_segment_sizes = [128, 1024] namespace = 'chaos-testing' @@ -38,25 +38,23 @@ def _install_milvus(seg_size): return release_name, None, None -@pytest.mark.skip(reason="issue #11952") class TestCustomizeSegmentSize: def teardown_method(self): + pass milvus_op = MilvusOperator() milvus_op.uninstall(self.release_name, namespace) + connections.disconnect("default") + connections.remove_connection("default") @pytest.mark.tags(CaseLabel.L3) - @pytest.mark.parametrize('size_id', [i for i in range(len(customize_segment_sizes))]) - def test_customize_segment_size(self, size_id): + @pytest.mark.parametrize('seg_size, seg_count', [(128, 10), (1024, 2)]) + def test_customize_segment_size(self, seg_size, seg_count): """ steps - 1. [test_milvus_install]: set up milvus with customized simd configured - 2. [test_simd_compat_e2e]: verify milvus is working well - 4. [test_milvus_cleanup]: delete milvus instances in teardown """ - size = customize_segment_sizes[size_id] - log.info(f"start to install milvus with segment size {size}") - release_name, host, port = _install_milvus(size) + log.info(f"start to install milvus with segment size {seg_size}") + release_name, host, port = _install_milvus(seg_size) self.release_name = release_name assert host is not None conn = connections.connect("default", host=host, port=port) @@ -64,76 +62,83 @@ class TestCustomizeSegmentSize: mil = MilvusSys(alias="default") log.info(f"milvus build version: {mil.build_version}") - # log.info(f"start to e2e verification: {size}") - # # create - # name = cf.gen_unique_str("compat") - # t0 = time.time() - # collection_w = ApiCollectionWrapper() - # collection_w.init_collection(name=name, - # schema=cf.gen_default_collection_schema(), - # timeout=40) - # tt = time.time() - t0 - # assert collection_w.name == name - # entities = collection_w.num_entities - # log.info(f"assert create collection: {tt}, init_entities: {entities}") - # - # # insert - # data = cf.gen_default_list_data() - # t0 = time.time() - # _, res = collection_w.insert(data) - # tt = time.time() - t0 - # log.info(f"assert insert: {tt}") - # assert res - # - # # flush - # t0 = time.time() - # assert collection_w.num_entities == len(data[0]) + entities - # tt = time.time() - t0 - # entities = collection_w.num_entities - # log.info(f"assert flush: {tt}, entities: {entities}") - # - # # search - # collection_w.load() - # search_vectors = cf.gen_vectors(1, ct.default_dim) - # search_params = {"metric_type": "L2", "params": {"nprobe": 16}} - # t0 = time.time() - # res_1, _ = collection_w.search(data=search_vectors, - # anns_field=ct.default_float_vec_field_name, - # param=search_params, limit=1) - # tt = time.time() - t0 - # log.info(f"assert search: {tt}") - # assert len(res_1) == 1 - # collection_w.release() - # - # # index - # d = cf.gen_default_list_data() - # collection_w.insert(d) - # log.info(f"assert index entities: {collection_w.num_entities}") - # _index_params = {"index_type": "IVF_SQ8", "params": {"nlist": 64}, "metric_type": "L2"} - # t0 = time.time() - # index, _ = collection_w.create_index(field_name=ct.default_float_vec_field_name, - # index_params=_index_params, - # name=cf.gen_unique_str()) - # tt = time.time() - t0 - # log.info(f"assert index: {tt}") - # assert len(collection_w.indexes) == 1 - # - # # search - # t0 = time.time() - # collection_w.load() - # tt = time.time() - t0 - # log.info(f"assert load: {tt}") - # search_vectors = cf.gen_vectors(1, ct.default_dim) - # t0 = time.time() - # res_1, _ = collection_w.search(data=search_vectors, - # anns_field=ct.default_float_vec_field_name, - # param=search_params, limit=1) - # tt = time.time() - t0 - # log.info(f"assert search: {tt}") - # - # # query - # term_expr = f'{ct.default_int64_field_name} in [1001,1201,4999,2999]' - # t0 = time.time() - # res, _ = collection_w.query(term_expr) - # tt = time.time() - t0 - # log.info(f"assert query result {len(res)}: {tt}") + log.info(f"start to e2e verification: {seg_size}") + # create + name = cf.gen_unique_str("segsiz") + t0 = time.time() + collection_w = ApiCollectionWrapper() + collection_w.init_collection(name=name, + schema=cf.gen_default_collection_schema(), + timeout=40) + tt = time.time() - t0 + assert collection_w.name == name + entities = collection_w.num_entities + log.info(f"assert create collection: {tt}, init_entities: {entities}") + + # insert + nb = 50000 + data = cf.gen_default_list_data(nb=nb) + t0 = time.time() + _, res = collection_w.insert(data) + tt = time.time() - t0 + log.info(f"assert insert: {tt}") + assert res + # insert 2 million entities + rounds = 40 + for _ in range(rounds-1): + _, res = collection_w.insert(data) + entities = collection_w.num_entities + assert entities == nb * rounds + + # load + collection_w.load() + utility_wrap = ApiUtilityWrapper() + segs, _ = utility_wrap.get_query_segment_info(collection_w.name) + log.info(f"assert segments: {len(segs)}") + assert len(segs) == seg_count + + # search + search_vectors = cf.gen_vectors(1, ct.default_dim) + search_params = {"metric_type": "L2", "params": {"nprobe": 16}} + t0 = time.time() + res_1, _ = collection_w.search(data=search_vectors, + anns_field=ct.default_float_vec_field_name, + param=search_params, limit=1, timeout=30) + tt = time.time() - t0 + log.info(f"assert search: {tt}") + assert len(res_1) == 1 + collection_w.release() + + # index + d = cf.gen_default_list_data() + collection_w.insert(d) + log.info(f"assert index entities: {collection_w.num_entities}") + _index_params = {"index_type": "IVF_SQ8", "params": {"nlist": 64}, "metric_type": "L2"} + t0 = time.time() + index, _ = collection_w.create_index(field_name=ct.default_float_vec_field_name, + index_params=_index_params, + name=cf.gen_unique_str(), timeout=120) + tt = time.time() - t0 + log.info(f"assert index: {tt}") + assert len(collection_w.indexes) == 1 + + # search + t0 = time.time() + collection_w.load() + tt = time.time() - t0 + log.info(f"assert load: {tt}") + search_vectors = cf.gen_vectors(1, ct.default_dim) + t0 = time.time() + res_1, _ = collection_w.search(data=search_vectors, + anns_field=ct.default_float_vec_field_name, + param=search_params, limit=1, timeout=30) + tt = time.time() - t0 + log.info(f"assert search: {tt}") + + # query + term_expr = f'{ct.default_int64_field_name} in [1001,1201,4999,2999]' + t0 = time.time() + res, _ = collection_w.query(term_expr, timeout=30) + tt = time.time() - t0 + log.info(f"assert query result {len(res)}: {tt}") +