Spaces:
Build error
Build error
File size: 5,620 Bytes
3932407 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 |
import argparse
import random
import uuid
from time import sleep
import requests
VECTOR_SIZE = 4
def generate_points(amount: int):
for _ in range(amount):
result_item = {"points": []}
for _ in range(100):
result_item["points"].append(
{
"id": str(uuid.uuid4()),
"vector": [round(random.uniform(0, 1), 2) for _ in range(VECTOR_SIZE)],
"payload": {"city": ["Berlin", "London"]}
}
)
yield result_item
def create_collection(qdrant_host, collection_name, collection_json):
resp = requests.put(
f"{qdrant_host}/collections/{collection_name}?timeout=60", json=collection_json)
if resp.status_code != 200:
print(f"Collection creation failed with response body:\n{resp.json()}")
exit(-1)
def update_collection(qdrant_host, collection_name, collection_json: dict):
resp = requests.patch(
f"{qdrant_host}/collections/{collection_name}?timeout=60", json=collection_json)
if resp.status_code != 200:
print(f"Collection patching failed with response body:\n{resp.json()}")
exit(-1)
def wait_for_status(qdrant_host, collection_name, status: str):
curr_status = ""
for i in range(30):
resp = requests.get(
f"{qdrant_host}/collections/{collection_name}")
curr_status = resp.json()["result"]["status"]
if resp.status_code != 200:
print(f"Collection info fetching failed with response body:\n{resp.json()}")
exit(-1)
if curr_status == status:
print(f"Status {status}: OK")
return "ok"
sleep(1)
print(f"Wait for status {status}")
print(f"After 30s status is not {status}, found: {curr_status}. Stop waiting.")
def insert_points(qdrant_host, collection_name, batch_json, quit_on_ood: bool = False):
resp = requests.put(
f"{qdrant_host}/collections/{collection_name}/points?wait=true", json=batch_json
)
expected_error_message = "No space left on device"
if resp.status_code != 200:
if resp.status_code == 500 and expected_error_message in resp.text:
if quit_on_ood:
return "ood"
requests.put(f"{qdrant_host}/collections/{collection_name}/points?wait=true", json=batch_json)
else:
error_response = resp.json()
print(f"Points insertions failed with response body:\n{error_response}")
exit(-2)
def search_point(qdrant_host, collection_name):
query = {
"vector": [round(random.uniform(0, 1), 2) for _ in range(VECTOR_SIZE)],
"top": 10,
"filters": {
"city": {
"values": ["Berlin"],
"excludes": False
}
}
}
resp = requests.post(
f"{qdrant_host}/collections/{collection_name}/points/search", json=query
)
if resp.status_code != 200:
print("Search failed")
exit(-3)
return resp.json()
def insert_points_and_search(qdrant_host, collection_name, points_amount):
collection_json = {
"vectors": {
"size": VECTOR_SIZE,
"distance": "Cosine"
}
}
create_collection(qdrant_host, collection_name, collection_json)
for points_batch in generate_points(points_amount):
insert_points(qdrant_host, collection_name, points_batch)
search_point(qdrant_host, collection_name)
def insert_points_then_index(qdrant_host, collection_name, points_amount):
"""
Disable indexing, create collection, insert points up to
a point when there is no space left on disk, then start indexing.
Wait for indexing then send a search request.
@param qdrant_host:
@param collection_name:
@param points_amount:
@return:
"""
collection_json = {
"vectors": {
"size": VECTOR_SIZE,
"distance": "Cosine"
},
"optimizers_config": {
"indexing_threshold": 0,
"default_segment_number": 2
},
"wal_config": {
"wal_capacity_mb": 1
}
}
create_collection(qdrant_host, collection_name, collection_json)
for points_batch in generate_points(points_amount):
res = insert_points(qdrant_host, collection_name, points_batch, quit_on_ood=True)
if res == "ood":
break
# start indexing
collection_json = {
"optimizers_config": {
"indexing_threshold": 10
}
}
update_collection(qdrant_host, collection_name, collection_json)
wait_for_status(qdrant_host, collection_name, "yellow")
wait_for_status(qdrant_host, collection_name, "green")
search_point(qdrant_host, collection_name)
def main():
parser = argparse.ArgumentParser("Create all required test items")
parser.add_argument("test_item")
parser.add_argument("collection_name")
parser.add_argument("points_amount", type=int)
parser.add_argument("ports", type=int, nargs="+")
args = parser.parse_args()
qdrant_host = f"http://127.0.0.1:{args.ports[0]}"
if args.test_item == "search":
print("run search test")
insert_points_and_search(qdrant_host, args.collection_name, args.points_amount)
elif args.test_item == "indexing":
print("run indexing test")
insert_points_then_index(qdrant_host, args.collection_name, args.points_amount)
else:
raise ValueError("Invalid test_item value, allowed: search|indexing")
print("SUCCESS")
if __name__ == "__main__":
main()
|