mirror of
https://gitee.com/milvus-io/milvus.git
synced 2024-12-05 05:18:52 +08:00
51 lines
1.6 KiB
Python
51 lines
1.6 KiB
Python
import random
|
|
import logging
|
|
from locust import User, task, between
|
|
from locust_task import MilvusTask
|
|
from client import MilvusClient
|
|
from milvus import DataType
|
|
import utils
|
|
|
|
connection_type = "single"
|
|
host = "172.16.50.9"
|
|
port = 19530
|
|
collection_name = "sift_5m_2000000_128_l2_2"
|
|
dim = 128
|
|
m = MilvusClient(host=host, port=port, collection_name=collection_name)
|
|
# m.clean_db()
|
|
# m.create_collection(dim, data_type=DataType.FLOAT_VECTOR, auto_id=True, other_fields=None)
|
|
vectors = [[random.random() for _ in range(dim)] for _ in range(1000)]
|
|
entities = m.generate_entities(vectors)
|
|
ids = [i for i in range(10000000)]
|
|
|
|
|
|
class QueryTask(User):
|
|
wait_time = between(0.001, 0.002)
|
|
# if connection_type == "single":
|
|
# client = MilvusTask(m=m)
|
|
# else:
|
|
# client = MilvusTask(host=host, port=port, collection_name=collection_name)
|
|
client = MilvusTask(host, port, collection_name, connection_type=connection_type)
|
|
|
|
# @task
|
|
# def query(self):
|
|
# top_k = 5
|
|
# 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(1)
|
|
def insert(self):
|
|
self.client.insert(entities)
|
|
|
|
# @task(1)
|
|
# def create(self):
|
|
# collection_name = utils.get_unique_name(prefix="locust")
|
|
# self.client.create_collection(dim, data_type=DataType.FLOAT_VECTOR, auto_id=True, collection_name=collection_name, other_fields=None)
|
|
|
|
# @task(1)
|
|
# def delete(self):
|
|
# delete_ids = random.sample(ids, 100)
|
|
# logging.error(delete_ids)
|
|
# self.client.delete(delete_ids)
|