def test_greedy_add_item_quantity():
collection2 = collections_pb2.Collection(replication_factor=2)
collection3 = collections_pb2.Collection(replication_factor=3)
shard2 = collections_pb2.Shard()
shard3 = collections_pb2.Shard()
item_2x = Item(collection2, shard2)
item_3x = Item(collection3, shard3)
items = [item_2x, item_3x]
knapsacks = [
Knapsack(collections_pb2.Node()),
Knapsack(collections_pb2.Node()),
Knapsack(collections_pb2.Node()),
]
algo = GreedyAlgorithm()
algo.add(items, knapsacks)
for knapsack in knapsacks:
assert item_3x in knapsack.items
count = 0
for knapsack in knapsacks:
if item_2x in knapsack.items:
count += 1
assert count == 2
def main():
hostname = "merger-grpc:50051"
if USE_SSL:
config = LocalDockerSSLConfig()
channel = grpc.secure_channel(hostname, config.ssl_channel_credentials)
else:
channel = grpc.insecure_channel(hostname)
print("Doing health check")
health_stub = health_pb2_grpc.HealthStub(channel)
response = health_stub.Check(
health_pb2.HealthCheckRequest(service="Merger"))
print(response)
# Grabs the collections protobubs from the files that indexing_job.py wrote to ./data
collections_files = [
"collection_1_data_source.pbtxt", "collection_2_data_source.pbtxt"
]
collections = []
for file in collections_files:
with open(os.path.join(data_dir, file), "r") as f:
collection = text_format.Parse(f.read(),
collections_pb2.Collection())
collections.append(collection)
print("Adding collections to cluster")
stub = servicers_pb2_grpc.MergerStub(channel)
response = stub.CollectionsAdd(
collections_pb2.CollectionsAddRequest(collections=collections))
print(response)
def collection_2shards_2d(tmpdir):
X, metadatas = gen_random_vectors_and_metadatas(dimension=2,
size=20,
dtype="float32",
id_prefix="shard_1")
index_proto_shard_1 = gen_faiss_index_proto(tmpdir, X, metadatas,
"shard_1")
X, metadatas = gen_random_vectors_and_metadatas(dimension=2,
size=25,
dtype="float32",
id_prefix="shard_2")
index_proto_shard_2 = gen_faiss_index_proto(tmpdir, X, metadatas,
"shard_2")
proto = collections_pb2.Collection(
name="test_name",
replication_factor=1,
enable_id_to_vector=True,
shards=[
collections_pb2.Shard(name="shard_1",
weight=20.0,
index=index_proto_shard_1),
collections_pb2.Shard(name="shard_2",
weight=25.0,
index=index_proto_shard_2),
],
)
yield Collection.from_proto(proto)
def create_collection_proto(
name: str, shards: List[collections_pb2.Shard],
replication_factor: int) -> collections_pb2.Collection:
"""Creates a collection protobuf which a list of shards protobufs"""
return collections_pb2.Collection(name=name,
replication_factor=replication_factor,
enable_id_to_vector=True,
shards=shards)
def test_knapsack_add_item():
shard = collections_pb2.Shard(weight=3.14)
collection = collections_pb2.Collection(shards=[shard])
node = collections_pb2.Node()
item = Item(collection, shard)
knapsack = Knapsack(node)
knapsack.add_item(item)
assert item in knapsack.items
assert knapsack.current_weight == pytest.approx(3.14)
def test_knapsack_add_duplicate_item():
shard = collections_pb2.Shard(weight=1.12)
collection = collections_pb2.Collection(shards=[shard])
node = collections_pb2.Node()
item = Item(collection, shard)
knapsack = Knapsack(node)
knapsack.add_item(item)
with pytest.raises(KnapsackItemException) as excinfo:
knapsack.add_item(item)
assert "Item already exists in this knapsack" == str(excinfo.value)
def test_greedy_add_one_knapsack():
shard = collections_pb2.Shard()
collection = collections_pb2.Collection()
node = collections_pb2.Node()
item = Item(collection, shard)
knapsack = Knapsack(node)
algo = GreedyAlgorithm()
algo.add([item], [knapsack])
assert item in knapsack.items
def test_greedy_add_item_quantity_to_much():
collection = collections_pb2.Collection(replication_factor=2)
shard = collections_pb2.Shard()
node = collections_pb2.Node()
item_2x = Item(collection, shard)
knapsack = Knapsack(node)
algo = GreedyAlgorithm()
algo.add([item_2x], [knapsack])
assert len(knapsack.items) == 1
def test_greedy_add_two_knapsack():
collection = collections_pb2.Collection()
shard1 = collections_pb2.Shard()
shard2 = collections_pb2.Shard()
node1 = collections_pb2.Node()
node2 = collections_pb2.Node()
items = [Item(collection, shard1), Item(collection, shard2)]
knapsacks = [Knapsack(node1), Knapsack(node2)]
algo = GreedyAlgorithm()
algo.add(items, knapsacks)
for knapsack in knapsacks:
assert len(knapsack.items) == 1