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

Add GPU supported cases (#22862) 

Signed-off-by: Binbin Lv <binbin.lv@zilliz.com>
pull/22857/head
binbin 2023-03-20 17:29:57 +08:00 committed by GitHub
parent 254cd4f15f
commit 1c8f707e4b
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 140 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

10
tests/python_client/common/common_type.py
View File

@ -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"

10
tests/python_client/testcases/test_index.py
View File

@ -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):
"""
******************************************************************

123
tests/python_client/testcases/test_search.py
View File

@ -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):
"""