mirror of https://github.com/milvus-io/milvus.git
43 lines
1.1 KiB
Python
43 lines
1.1 KiB
Python
import random
|
|
from locust import User, task, between
|
|
from locust_task import MilvusTask
|
|
from client import MilvusClient
|
|
|
|
connection_type = "single"
|
|
host = "192.168.1.29"
|
|
port = 19530
|
|
collection_name = "sift_128_euclidean"
|
|
dim = 128
|
|
m = MilvusClient(host=host, port=port, collection_name=collection_name)
|
|
|
|
|
|
class MixTask(User):
|
|
wait_time = between(0.001, 0.002)
|
|
print("in query task")
|
|
if connection_type == "single":
|
|
client = MilvusTask(m=m)
|
|
else:
|
|
client = MilvusTask(host=host, port=port, collection_name=collection_name)
|
|
|
|
@task(30)
|
|
def query(self):
|
|
top_k = 10
|
|
X = [[random.random() for i in range(dim)] for i in range(1)]
|
|
search_param = {"nprobe": 16}
|
|
self.client.query(X, top_k, search_param)
|
|
|
|
@task(10)
|
|
def insert(self):
|
|
id = random.randint(10000000, 10000000000)
|
|
X = [[random.random() for i in range(dim)] for i in range(1)]
|
|
self.client.insert(X, ids=[id])
|
|
|
|
@task(1)
|
|
def flush(self):
|
|
self.client.flush()
|
|
|
|
# @task(5)
|
|
# def delete(self):
|
|
# self.client.delete([random.randint(1, 1000000)])
|
|
|