From 1c8f707e4b4fbdeda8cfae7886c4caf290d5fd45 Mon Sep 17 00:00:00 2001 From: binbin <83755740+binbinlv@users.noreply.github.com> Date: Mon, 20 Mar 2023 17:29:57 +0800 Subject: [PATCH] Add GPU supported cases (#22862) Signed-off-by: Binbin Lv --- tests/python_client/common/common_type.py | 10 +- tests/python_client/testcases/test_index.py | 10 +- tests/python_client/testcases/test_search.py | 123 +++++++++++++++++++ 3 files changed, 140 insertions(+), 3 deletions(-) diff --git a/tests/python_client/common/common_type.py b/tests/python_client/common/common_type.py index a754707b94..f6323426fa 100644 --- a/tests/python_client/common/common_type.py +++ b/tests/python_client/common/common_type.py @@ -182,10 +182,13 @@ get_wrong_format_dict = [ ] """ Specially defined list """ -all_index_types = ["FLAT", "IVF_FLAT", "IVF_SQ8", "IVF_PQ", "HNSW", "ANNOY", "DISKANN", "BIN_FLAT", "BIN_IVF_FLAT"] +all_index_types = ["FLAT", "IVF_FLAT", "IVF_SQ8", "IVF_PQ", "HNSW", "ANNOY", "DISKANN", "BIN_FLAT", "BIN_IVF_FLAT", + "GPU_FLAT", "GPU_IVF_FLAT", "GPU_IVF_PQ", "GPU_IVF_SQ8", "RAFT_IVF_FLAT", "RAFT_IVF_PQ"] default_index_params = [{"nlist": 128}, {"nlist": 128}, {"nlist": 128}, {"nlist": 128, "m": 16, "nbits": 8}, - {"M": 48, "efConstruction": 500}, {"n_trees": 50}, {}, {"nlist": 128}, {"nlist": 128}] + {"M": 48, "efConstruction": 500}, {"n_trees": 50}, {}, {"nlist": 128}, {"nlist": 128}, + {"nlist": 128}, {"nlist": 128}, {"nlist": 128, "m": 16, "nbits": 8}, {"nlist": 128}, + {"nlist": 128}, {"nlist": 128, "m": 16, "nbits": 8}] Handler_type = ["GRPC", "HTTP"] binary_support = ["BIN_FLAT", "BIN_IVF_FLAT"] @@ -250,6 +253,8 @@ class CaseLabel: loadbalance testcases which need to be run in multi query nodes ClusterOnly: For functions only suitable to cluster mode + GPU: + For GPU supported cases """ L0 = "L0" L1 = "L1" @@ -258,3 +263,4 @@ class CaseLabel: Loadbalance = "Loadbalance" # loadbalance testcases which need to be run in multi query nodes ClusterOnly = "ClusterOnly" # For functions only suitable to cluster mode MultiQueryNodes = "MultiQueryNodes" # for 8 query nodes configs tests, such as resource group + GPU = "GPU" diff --git a/tests/python_client/testcases/test_index.py b/tests/python_client/testcases/test_index.py index dc321e4d73..05232ca310 100644 --- a/tests/python_client/testcases/test_index.py +++ b/tests/python_client/testcases/test_index.py @@ -43,6 +43,7 @@ default_search_ip_params = ct.default_search_ip_params default_search_binary_params = ct.default_search_binary_params +@pytest.mark.tags(CaseLabel.GPU) class TestIndexParams(TestcaseBase): """ Test case of index interface """ @@ -195,6 +196,7 @@ class TestIndexParams(TestcaseBase): ct.err_msg: "Invalid index name"}) +@pytest.mark.tags(CaseLabel.GPU) class TestIndexOperation(TestcaseBase): """ Test case of index interface """ @@ -389,6 +391,7 @@ class TestIndexOperation(TestcaseBase): self.index_wrap.drop() +@pytest.mark.tags(CaseLabel.GPU) class TestIndexAdvanced(TestcaseBase): """ Test case of index interface """ @@ -457,6 +460,7 @@ class TestIndexAdvanced(TestcaseBase): """ +@pytest.mark.tags(CaseLabel.GPU) class TestNewIndexBase(TestcaseBase): """ ****************************************************************** @@ -1064,6 +1068,7 @@ class TestNewIndexBase(TestcaseBase): assert len(search_res[0]) == ct.default_limit +@pytest.mark.tags(CaseLabel.GPU) class TestNewIndexBinary(TestcaseBase): def get_simple_index(self, request): @@ -1195,7 +1200,7 @@ class TestNewIndexBinary(TestcaseBase): assert collection_w.has_index(index_name=binary_field_name)[0] == False assert len(collection_w.indexes) == 0 - +@pytest.mark.tags(CaseLabel.GPU) class TestIndexInvalid(TestcaseBase): """ Test create / describe / drop index interfaces with invalid collection names @@ -1282,6 +1287,7 @@ class TestIndexInvalid(TestcaseBase): "err_msg": "invalid index params"}) +@pytest.mark.tags(CaseLabel.GPU) class TestNewIndexAsync(TestcaseBase): @pytest.fixture(scope="function", params=[False, True]) def _async(self, request): @@ -1363,6 +1369,7 @@ class TestNewIndexAsync(TestcaseBase): _callback=self.call_back()) +@pytest.mark.tags(CaseLabel.GPU) class TestIndexString(TestcaseBase): """ ****************************************************************** @@ -1572,6 +1579,7 @@ class TestIndexString(TestcaseBase): assert len(collection_w.indexes) == 0 +@pytest.mark.tags(CaseLabel.GPU) class TestIndexDiskann(TestcaseBase): """ ****************************************************************** diff --git a/tests/python_client/testcases/test_search.py b/tests/python_client/testcases/test_search.py index 53bce25afe..1f4484946d 100644 --- a/tests/python_client/testcases/test_search.py +++ b/tests/python_client/testcases/test_search.py @@ -1728,6 +1728,7 @@ class TestCollectionSearch(TestcaseBase): "_async": _async}) @pytest.mark.tags(CaseLabel.L2) + @pytest.mark.tags(CaseLabel.GPU) @pytest.mark.parametrize("index, params", zip(ct.all_index_types[:7], ct.default_index_params[:7])) @@ -1767,7 +1768,48 @@ class TestCollectionSearch(TestcaseBase): "limit": default_limit, "_async": _async}) + @pytest.mark.tags(CaseLabel.GPU) + @pytest.mark.parametrize("index, params", + zip(ct.all_index_types[9:14], + ct.default_index_params[9:14])) + def test_search_after_different_index_with_params_gpu(self, dim, index, params, auto_id, _async): + """ + target: test search after different index + method: test search after different index and corresponding search params + expected: search successfully with limit(topK) + """ + # 1. initialize with data + collection_w, _, _, insert_ids, time_stamp = self.init_collection_general(prefix, True, 5000, + partition_num=1, + auto_id=auto_id, + dim=dim, is_index=False)[0:5] + # 2. create index and load + if params.get("m"): + if (dim % params["m"]) != 0: + params["m"] = dim // 4 + if params.get("PQM"): + if (dim % params["PQM"]) != 0: + params["PQM"] = dim // 4 + default_index = {"index_type": index, "params": params, "metric_type": "L2"} + collection_w.create_index("float_vector", default_index) + collection_w.load() + # 3. search + search_params = cf.gen_search_param(index) + vectors = [[random.random() for _ in range(dim)] for _ in range(default_nq)] + for search_param in search_params: + log.info("Searching with search params: {}".format(search_param)) + collection_w.search(vectors[:default_nq], default_search_field, + search_param, default_limit, + default_search_exp, _async=_async, + travel_timestamp=0, + check_task=CheckTasks.check_search_results, + check_items={"nq": default_nq, + "ids": insert_ids, + "limit": default_limit, + "_async": _async}) + @pytest.mark.tags(CaseLabel.L2) + @pytest.mark.tags(CaseLabel.GPU) @pytest.mark.parametrize("index, params", zip(ct.all_index_types[:6], ct.default_index_params[:6])) @@ -1804,8 +1846,47 @@ class TestCollectionSearch(TestcaseBase): "ids": insert_ids, "limit": default_limit, "_async": _async}) + + @pytest.mark.tags(CaseLabel.GPU) + @pytest.mark.parametrize("index, params", + zip(ct.all_index_types[9:14], + ct.default_index_params[9:14])) + def test_search_after_different_index_with_min_dim_gpu(self, index, params, auto_id, _async): + """ + target: test search after different index with min dim + method: test search after different index and corresponding search params with dim = 1 + expected: search successfully with limit(topK) + """ + # 1. initialize with data + collection_w, _, _, insert_ids, time_stamp = self.init_collection_general(prefix, True, 5000, + partition_num=1, + auto_id=auto_id, + dim=min_dim, is_index=False)[0:5] + # 2. create index and load + if params.get("m"): + params["m"] = min_dim + if params.get("PQM"): + params["PQM"] = min_dim + default_index = {"index_type": index, "params": params, "metric_type": "L2"} + collection_w.create_index("float_vector", default_index) + collection_w.load() + # 3. search + search_params = cf.gen_search_param(index) + vectors = [[random.random() for _ in range(min_dim)] for _ in range(default_nq)] + for search_param in search_params: + log.info("Searching with search params: {}".format(search_param)) + collection_w.search(vectors[:default_nq], default_search_field, + search_param, default_limit, + default_search_exp, _async=_async, + travel_timestamp=0, + check_task=CheckTasks.check_search_results, + check_items={"nq": default_nq, + "ids": insert_ids, + "limit": default_limit, + "_async": _async}) @pytest.mark.tags(CaseLabel.L2) + @pytest.mark.tags(CaseLabel.GPU) @pytest.mark.parametrize("index, params", zip(ct.all_index_types[:7], ct.default_index_params[:7])) @@ -1847,6 +1928,48 @@ class TestCollectionSearch(TestcaseBase): "limit": default_limit, "_async": _async}) + @pytest.mark.tags(CaseLabel.GPU) + @pytest.mark.parametrize("index, params", + zip(ct.all_index_types[9:14], + ct.default_index_params[9:14])) + def test_search_after_index_different_metric_type_gpu(self, dim, index, params, auto_id, _async): + """ + target: test search with different metric type + method: test search with different metric type + expected: searched successfully + """ + # 1. initialize with data + collection_w, _, _, insert_ids, time_stamp = self.init_collection_general(prefix, True, 5000, + partition_num=1, + auto_id=auto_id, + dim=dim, is_index=False)[0:5] + # 2. create different index + if params.get("m"): + if (dim % params["m"]) != 0: + params["m"] = dim // 4 + if params.get("PQM"): + if (dim % params["PQM"]) != 0: + params["PQM"] = dim // 4 + log.info("test_search_after_index_different_metric_type: Creating index-%s" % index) + default_index = {"index_type": index, "params": params, "metric_type": "IP"} + collection_w.create_index("float_vector", default_index) + log.info("test_search_after_index_different_metric_type: Created index-%s" % index) + collection_w.load() + # 3. search + search_params = cf.gen_search_param(index, "IP") + vectors = [[random.random() for _ in range(dim)] for _ in range(default_nq)] + for search_param in search_params: + log.info("Searching with search params: {}".format(search_param)) + collection_w.search(vectors[:default_nq], default_search_field, + search_param, default_limit, + default_search_exp, _async=_async, + travel_timestamp=0, + check_task=CheckTasks.check_search_results, + check_items={"nq": default_nq, + "ids": insert_ids, + "limit": default_limit, + "_async": _async}) + @pytest.mark.tags(CaseLabel.L2) def test_search_collection_multiple_times(self, nb, nq, dim, auto_id, _async): """