mirror of https://github.com/milvus-io/milvus.git
76 lines
2.6 KiB
Python
76 lines
2.6 KiB
Python
import time
|
|
import random
|
|
from locust import Locust, TaskSet, events, task, between
|
|
from client import MilvusClient
|
|
from . import BasicRunner
|
|
|
|
|
|
dim = 128
|
|
top_k = 10
|
|
X = [[random.random() for i in range(dim)] for i in range(1)]
|
|
search_param = {"nprobe": 16}
|
|
|
|
|
|
class MilvusTask(object):
|
|
def __init__(self, type="single", args):
|
|
self.type = type
|
|
self.m = None
|
|
if type == "single":
|
|
self.m = MilvusClient(host=args["host"], port=args["port"], collection_name=args["collection_name"])
|
|
elif type == "multi":
|
|
self.m = MilvusClient(host=args["m"])
|
|
|
|
def query(self, *args, **kwargs):
|
|
name = "milvus_search"
|
|
request_type = "grpc"
|
|
start_time = time.time()
|
|
try:
|
|
# result = self.m.getattr(*args, **kwargs)
|
|
status, result = self.m.query(*args, **kwargs)
|
|
except Exception as e:
|
|
total_time = int((time.time() - start_time) * 1000)
|
|
events.request_failure.fire(request_type=request_type, name=name, response_time=total_time, exception=e, response_length=0)
|
|
else:
|
|
if not status.OK:
|
|
total_time = int((time.time() - start_time) * 1000)
|
|
events.request_failure.fire(request_type=request_type, name=name, response_time=total_time, exception=e, response_length=0)
|
|
else:
|
|
total_time = int((time.time() - start_time) * 1000)
|
|
events.request_success.fire(request_type=request_type, name=name, response_time=total_time, response_length=0)
|
|
# In this example, I've hardcoded response_length=0. If we would want the response length to be
|
|
# reported correctly in the statistics, we would probably need to hook in at a lower level
|
|
|
|
|
|
class MilvusLocust(Locust):
|
|
def __init__(self, *args, **kwargs):
|
|
super(MilvusLocust, self).__init__(*args, **kwargs)
|
|
self.client = MilvusTask(self.host, self.port, self.collection_name)
|
|
|
|
|
|
class Query(MilvusLocust):
|
|
host = "192.168.1.183"
|
|
port = 19530
|
|
collection_name = "sift_128_euclidean"
|
|
# m = MilvusClient(host=host, port=port, collection_name=collection_name)
|
|
wait_time = between(0.001, 0.002)
|
|
|
|
class task_set(TaskSet):
|
|
@task
|
|
def query(self):
|
|
self.client.query(X, top_k, search_param)
|
|
|
|
|
|
class LocustRunner(BasicRunner):
|
|
"""Only one client, not support M/S mode"""
|
|
def __init__(self, args):
|
|
# Start client with params including client number && last time && hatch rate ...
|
|
pass
|
|
|
|
def set_up(self):
|
|
# helm install locust client
|
|
pass
|
|
|
|
def tear_down(self):
|
|
# helm uninstall
|
|
pass
|