def query(client: Client, docs_gen_func: Callable[[int], Generator],
req_size: int, dataset: str, nr_docs: int):
on_done = validate_img
if dataset == 'text':
on_done = validate_text
client.search(docs_gen_func(nr_docs),
on_done=on_done,
top_k=TOP_K,
request_size=req_size)
def test_custom_dockerfile():
f = Flow(port_expose=exposed_port).add(
uses='DummyRedisIndexer',
py_modules='redis_executor.py',
upload_files=[
os.path.join(cur_dir, '../../daemon/unit/models/good_ws_custom_dockerfile'),
],
host='localhost:8000',
)
with f:
c = Client(port=exposed_port)
c.index(
inputs=(
Document(text=f'{i}', embedding=np.random.rand(2, 3)) for i in range(5)
),
)
resp = c.search(inputs=[Document(text='3')], return_results=True)
assert resp[0].docs[0].matches[0].text == '3'
assert resp[0].docs[0].matches[0].embedding.shape == (2, 3)
def send_requests(
port_expose,
start_rolling_update_event: multiprocessing.Event,
result_queue: multiprocessing.Queue,
doc_count: int,
request_count: int,
):
client = Client(port=port_expose)
for i in range(request_count):
responses = client.search(
[
Document(id=f'{idx}', text=f'doc{idx}')
for idx in range(doc_count)
],
request_size=10,
return_results=True,
)
for r in responses:
result_queue.put(r.docs.texts)
if i == 5:
start_rolling_update_event.set()
# give the rolling update some time to kick in
time.sleep(0.1)
