Spaces:
Build error
Build error
File size: 4,362 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 |
import pathlib
from operator import itemgetter
from random import randrange
from time import sleep
from .fixtures import create_collection, drop_collection
from .utils import *
N_PEERS = 2
N_SHARDS = 2
N_REPLICAS = 1
N_COLLECTION_LOOPS = 3
COLLECTION_NAME = "test_collection"
POINTS_JSON = {
"points": [
{"id": 1, "vector": [0.05, 0.61, 0.76, 0.74], "payload": {"city": "Berlin"}},
{"id": 2, "vector": [0.19, 0.81, 0.75, 0.11], "payload": {"city": "London"}},
{"id": 3, "vector": [0.36, 0.55, 0.47, 0.94], "payload": {"city": "Paris"}},
{"id": 4, "vector": [0.18, 0.01, 0.85, 0.80], "payload": {"city": "Malaga"}},
{"id": 5, "vector": [0.24, 0.18, 0.22, 0.44], "payload": {"city": "New York"}},
{"id": 6, "vector": [0.35, 0.08, 0.11, 0.44], "payload": {"city": "Munich"}},
{"id": 7, "vector": [0.45, 0.07, 0.21, 0.04], "payload": {"city": "Madrid"}},
{"id": 8, "vector": [0.75, 0.18, 0.91, 0.48], "payload": {"city": "Prague"}},
{"id": 9, "vector": [0.30, 0.01, 0.10, 0.12], "payload": {"city": "Bratislava"}},
{"id": 10, "vector": [0.95, 0.8, 0.17, 0.19], "payload": {"city": "Tokyo"}},
]
}
def get_collection_points(uri):
req_json = {
"limit": len(POINTS_JSON["points"]),
"with_payload": True,
"with_vector": True
}
res = requests.post(
f"{uri}/collections/{COLLECTION_NAME}/points/scroll", json=req_json
)
assert_http_ok(res)
return res.json()["result"]["points"]
def compare_points(uri):
original_points = POINTS_JSON["points"]
fetched_points = get_collection_points(uri)
original_points, fetched_points = [sorted(l, key=itemgetter('id')) for l in (original_points, fetched_points)]
pairs = zip(original_points, fetched_points)
diff = [(x, y) for x, y in pairs if x != y]
assert len(diff) == 0, f'Original and final points are not equal, diff: "{diff}"'
def test_reinit_consensus(tmp_path: pathlib.Path):
assert_project_root()
peer_urls, peer_dirs, bootstrap_url = start_cluster(tmp_path, N_PEERS)
create_collection(peer_urls[0], shard_number=N_SHARDS, replication_factor=N_REPLICAS)
wait_collection_exists_and_active_on_all_peers(collection_name=COLLECTION_NAME, peer_api_uris=peer_urls)
for i in range(N_COLLECTION_LOOPS):
drop_collection(peer_urls[randrange(N_PEERS)], collection=COLLECTION_NAME)
create_collection(peer_urls[randrange(N_PEERS)], shard_number=N_SHARDS, replication_factor=N_REPLICAS)
wait_collection_exists_and_active_on_all_peers(collection_name=COLLECTION_NAME, peer_api_uris=peer_urls)
r_batch = requests.put(
f"{peer_urls[randrange(N_PEERS)]}/collections/{COLLECTION_NAME}/points?wait=true", json=POINTS_JSON)
assert_http_ok(r_batch)
# assert data in both shards
for peer_api_uri in peer_urls:
info = get_collection_cluster_info(peer_api_uri, COLLECTION_NAME)
for shard in info["local_shards"]:
assert shard["points_count"] > 0
print("Stop the peers, keep data")
for _ in range(N_PEERS):
processes.pop().kill()
print("Start the peers with new urls")
peer_urls_new = []
(bootstrap_api_uri, bootstrap_uri) = start_first_peer(peer_dirs[0], "peer_0_restarted.log", reinit=True)
peer_urls_new.append(bootstrap_api_uri)
leader = wait_peer_added(bootstrap_api_uri, expected_size=2)
for i in range(1, len(peer_dirs)):
peer_urls_new.append(start_peer(peer_dirs[i], f"peer_{i}_restarted.log", bootstrap_uri, reinit=True))
wait_for_uniform_cluster_status(peer_urls_new, leader)
wait_collection_exists_and_active_on_all_peers(collection_name=COLLECTION_NAME, peer_api_uris=peer_urls_new)
compare_points(peer_urls_new[0])
print("Add one more peer")
peer_dir_additional = make_peer_folder(tmp_path, N_PEERS)
peer_dirs.append(peer_dir_additional)
peer_url_additional = start_peer(peer_dir_additional, f"peer_{N_PEERS}_additional.log", bootstrap_uri)
wait_peer_added(bootstrap_api_uri, expected_size=3)
peer_urls_new.append(peer_url_additional)
wait_for_uniform_cluster_status(peer_urls_new, leader)
wait_collection_exists_and_active_on_all_peers(collection_name=COLLECTION_NAME, peer_api_uris=peer_urls_new)
compare_points(peer_urls_new[1])
|