def test_override_requests_uses_after():
class FooExecutor(Executor):
@requests(on='/bar')
def foo(self, docs, **kwargs):
for doc in docs:
doc.text = 'foo called'
class OtherExecutor(Executor):
@requests(on='/bar')
def bar(self, docs, **kwargs):
for doc in docs:
doc.text = 'bar called'
with Flow(port_expose=exposed_port).add(
uses=FooExecutor,
uses_requests={'/foo': 'foo'},
uses_after=OtherExecutor,
uses_before=OtherExecutor,
) as f:
c = Client(port=exposed_port)
resp1 = c.post(on='/foo',
inputs=DocumentArray([Document(text='')]),
return_results=True)
resp2 = c.post(
on='/non_foo',
inputs=DocumentArray([Document(text='')]),
return_results=True,
)
resp3 = c.post(on='/bar',
inputs=DocumentArray([Document(text='')]),
return_results=True)
assert resp1[0].docs[0].text == 'foo called'
assert resp2[0].docs[0].text == ''
assert resp3[0].docs[0].text == 'bar called'
def test_client_host_scheme(protocol):
port = random_port()
f = Flow(protocol='websocket' if protocol == 'ws' else protocol,
port=port).add()
with f:
c = Client(host=f'{protocol}://localhost:{port}')
c.post('/', inputs=DocumentArray.empty(2))
def peer_client(port, protocol, peer_hash, queue):
c = Client(protocol=protocol, port=port)
for _ in range(NUM_REQUESTS):
c.post(
'/ping',
Document(text=peer_hash),
on_done=lambda r: pong(peer_hash, queue, r),
)
async def test_deployments_trivial_topology(port_generator):
deployment_port = port_generator()
port = port_generator()
graph_description = (
'{"start-gateway": ["deployment0"], "deployment0": ["end-gateway"]}')
deployments_addresses = f'{{"deployment0": ["0.0.0.0:{deployment_port}"]}}'
# create a single worker pod
worker_deployment = _create_regular_deployment(deployment_port)
# create a single gateway pod
gateway_deployment = _create_gateway_deployment(graph_description,
deployments_addresses,
port)
with gateway_deployment, worker_deployment:
# send requests to the gateway
c = Client(host='localhost',
port=port,
asyncio=True,
return_responses=True)
responses = c.post('/', inputs=async_inputs, request_size=1)
response_list = []
async for response in responses:
response_list.append(response)
assert len(response_list) == 20
assert len(response_list[0].docs) == 1
async def test_deployments_shards(polling, port_generator):
head_port = port_generator()
port = port_generator()
graph_description = (
'{"start-gateway": ["deployment0"], "deployment0": ["end-gateway"]}')
deployments_addresses = f'{{"deployment0": ["0.0.0.0:{head_port}"]}}'
deployment = _create_regular_deployment(port=head_port,
name='deployment',
polling=polling,
shards=10)
deployment.start()
gateway_deployment = _create_gateway_deployment(graph_description,
deployments_addresses,
port)
gateway_deployment.start()
await asyncio.sleep(1.0)
c = Client(host='localhost',
port=port,
asyncio=True,
return_responses=True)
responses = c.post('/', inputs=async_inputs, request_size=1)
response_list = []
async for response in responses:
response_list.append(response)
gateway_deployment.close()
deployment.close()
assert len(response_list) == 20
assert len(
response_list[0].docs) == 1 if polling == PollingType.ANY else 10
async def test_pods_gateway_worker_direct_connection(port_generator):
worker_port = port_generator()
port = port_generator()
graph_description = '{"start-gateway": ["pod0"], "pod0": ["end-gateway"]}'
pod_addresses = f'{{"pod0": ["0.0.0.0:{worker_port}"]}}'
worker_pod = _create_worker_pod(worker_port, f'pod0')
worker_pod.start()
await asyncio.sleep(0.1)
# create a single gateway pod
gateway_pod = _create_gateway_pod(graph_description, pod_addresses, port)
gateway_pod.start()
await asyncio.sleep(1.0)
worker_pod.wait_start_success()
gateway_pod.wait_start_success()
c = Client(host='localhost', port=port, asyncio=True)
responses = c.post('/', inputs=async_inputs, request_size=1, return_responses=True)
response_list = []
async for response in responses:
response_list.append(response)
# clean up pods
gateway_pod.close()
worker_pod.close()
assert len(response_list) == 20
assert len(response_list[0].docs) == 1
async def test_deployments_replicas(port_generator):
head_port = port_generator()
port = port_generator()
graph_description = (
'{"start-gateway": ["deployment0"], "deployment0": ["end-gateway"]}')
deployment = _create_regular_deployment(port=head_port,
name='deployment',
replicas=10)
deployment.start()
connections = [f'0.0.0.0:{port}' for port in deployment.ports]
deployments_addresses = f'{{"deployment0": {json.dumps(connections)}}}'
gateway_deployment = _create_gateway_deployment(graph_description,
deployments_addresses,
port)
gateway_deployment.start()
await asyncio.sleep(1.0)
c = Client(host='localhost',
port=port,
asyncio=True,
return_responses=True)
responses = c.post('/', inputs=async_inputs, request_size=1)
response_list = []
async for response in responses:
response_list.append(response)
gateway_deployment.close()
deployment.close()
assert len(response_list) == 20
assert len(response_list[0].docs) == 1
async def test_pseudo_remote_pods_replicas(gateway, head, worker):
NUM_REPLICAS = 3
head_port = random_port()
port_expose = random_port()
graph_description = (
'{"start-gateway": ["deployment0"], "deployment0": ["end-gateway"]}')
deployments_addresses = f'{{"deployment0": ["0.0.0.0:{head_port}"]}}'
# create a single head pod
head_pod = _create_head_pod(head, head_port)
head_pod.start()
# create the shards
replica_pods = []
for i in range(NUM_REPLICAS):
# create worker
worker_port = random_port()
# create a single worker pod
worker_pod = _create_worker_pod(worker, worker_port,
f'deployment0/{i}')
replica_pods.append(worker_pod)
worker_pod.start()
await asyncio.sleep(0.1)
if head == 'remote':
worker_host = __docker_host__
else:
worker_host = HOST
# this would be done by the deployment, its adding the worker to the head
activate_msg = ControlRequest(command='ACTIVATE')
activate_msg.add_related_entity('worker', worker_host, worker_port)
GrpcConnectionPool.send_request_sync(activate_msg,
f'{HOST}:{head_port}')
# create a single gateway pod
gateway_pod = _create_gateway_pod(gateway, graph_description,
deployments_addresses, port_expose)
gateway_pod.start()
await asyncio.sleep(1.0)
c = Client(host='localhost', port=port_expose, asyncio=True)
responses = c.post('/',
inputs=async_inputs,
request_size=1,
return_results=True)
response_list = []
async for response in responses:
response_list.append(response)
# clean up pods
gateway_pod.close()
head_pod.close()
for pod in replica_pods:
pod.close()
assert len(response_list) == 20
assert len(response_list[0].docs) == 1
def test_load_yaml_route(req_endpoint, doc_text):
f = Flow(port_expose=12345).add(uses=y)
c = Client(port=exposed_port)
with f:
results = c.post(req_endpoint, Document(), return_results=True)
assert results[0].docs[0].text == doc_text
async def test_pods_with_replicas_advance_faster(port_generator):
head_port = port_generator()
port = port_generator()
graph_description = '{"start-gateway": ["pod0"], "pod0": ["end-gateway"]}'
pod_addresses = f'{{"pod0": ["0.0.0.0:{head_port}"]}}'
# create a single gateway pod
gateway_pod = _create_gateway_pod(graph_description, pod_addresses, port)
gateway_pod.start()
# create the shards
connection_list_dict = {}
pods = []
for i in range(10):
# create worker
worker_port = port_generator()
# create a single worker pod
worker_pod = _create_worker_pod(worker_port, f'pod0/{i}',
'FastSlowExecutor')
connection_list_dict[i] = [f'127.0.0.1:{worker_port}']
pods.append(worker_pod)
worker_pod.start()
await asyncio.sleep(0.1)
# create a single head pod
head_pod = _create_head_pod(head_port, connection_list_dict, 'head')
head_pod.start()
head_pod.wait_start_success()
gateway_pod.wait_start_success()
for pod in pods:
# this would be done by the Pod, its adding the worker to the head
pod.wait_start_success()
c = Client(host='localhost', port=port, asyncio=True)
input_docs = [Document(text='slow'), Document(text='fast')]
responses = c.post('/',
inputs=input_docs,
request_size=1,
return_responses=True)
response_list = []
async for response in responses:
response_list.append(response)
# clean up pods
gateway_pod.close()
head_pod.close()
for pod in pods:
pod.close()
assert len(response_list) == 2
for response in response_list:
assert len(response.docs) == 1
assert response_list[0].docs[0].text == 'fast'
assert response_list[1].docs[0].text == 'slow'
async def test_deployments_flow_topology(complete_graph_dict, uses_before,
uses_after, port_generator):
deployments = [
deployment_name for deployment_name in complete_graph_dict.keys()
if 'gateway' not in deployment_name
]
started_deployments = []
deployments_addresses = '{'
for deployment in deployments:
head_port = port_generator()
deployments_addresses += f'"{deployment}": ["0.0.0.0:{head_port}"],'
regular_deployment = _create_regular_deployment(
port=head_port,
name=f'{deployment}',
uses_before=uses_before,
uses_after=uses_after,
shards=2,
)
started_deployments.append(regular_deployment)
regular_deployment.start()
# remove last comma
deployments_addresses = deployments_addresses[:-1]
deployments_addresses += '}'
port = port_generator()
# create a single gateway pod
gateway_deployment = _create_gateway_deployment(
json.dumps(complete_graph_dict), deployments_addresses, port)
gateway_deployment.start()
await asyncio.sleep(0.1)
# send requests to the gateway
c = Client(host='localhost',
port=port,
asyncio=True,
return_responses=True)
responses = c.post('/', inputs=async_inputs, request_size=1)
response_list = []
async for response in responses:
response_list.append(response)
# clean up deployments
gateway_deployment.close()
for deployment in started_deployments:
deployment.close()
assert len(response_list) == 20
expected_docs = 1
if uses_before and uses_after:
expected_docs = 3 + 1 + 1
elif uses_before or uses_after:
expected_docs = 3
assert len(response_list[0].docs) == expected_docs
def _send_request(gateway_port, protocol):
"""send request to gateway and see what happens"""
c = Client(host='localhost', port=gateway_port, protocol=protocol)
return c.post(
'/foo',
inputs=[Document(text='hi') for _ in range(2)],
request_size=1,
return_responses=True,
)
async def test_pods_trivial_topology(head_runtime_docker_image_built,
worker_runtime_docker_image_built):
worker_port = random_port()
head_port = random_port()
port = random_port()
graph_description = '{"start-gateway": ["pod0"], "pod0": ["end-gateway"]}'
pod_addresses = f'{{"pod0": ["0.0.0.0:{head_port}"]}}'
# create a single worker pod
worker_pod = _create_worker_pod(worker_port)
# create a single head pod
head_pod = _create_head_pod(head_port)
# create a single gateway pod
gateway_pod = _create_gateway_pod(graph_description, pod_addresses, port)
with gateway_pod, head_pod, worker_pod:
await asyncio.sleep(1.0)
assert HeadRuntime.wait_for_ready_or_shutdown(
timeout=5.0,
ctrl_address=head_pod.runtime_ctrl_address,
ready_or_shutdown_event=head_pod.ready_or_shutdown.event,
)
assert WorkerRuntime.wait_for_ready_or_shutdown(
timeout=5.0,
ctrl_address=worker_pod.runtime_ctrl_address,
ready_or_shutdown_event=worker_pod.ready_or_shutdown.event,
)
head_pod.ready_or_shutdown.event.wait(timeout=5.0)
worker_pod.ready_or_shutdown.event.wait(timeout=5.0)
gateway_pod.ready_or_shutdown.event.wait(timeout=5.0)
# this would be done by the Pod, its adding the worker to the head
activate_msg = ControlRequest(command='ACTIVATE')
worker_host, worker_port = worker_pod.runtime_ctrl_address.split(':')
activate_msg.add_related_entity('worker', worker_host,
int(worker_port))
assert GrpcConnectionPool.send_request_sync(
activate_msg, head_pod.runtime_ctrl_address)
# send requests to the gateway
c = Client(host='localhost', port=port, asyncio=True)
responses = c.post('/',
inputs=async_inputs,
request_size=1,
return_responses=True)
response_list = []
async for response in responses:
response_list.append(response)
assert len(response_list) == 20
assert len(response_list[0].docs) == 1
def test_override_requests():
class FooExecutor(Executor):
@requests(on='/foo')
def foo(self, docs, **kwargs):
for doc in docs:
doc.text = 'foo called'
with Flow(port=exposed_port).add(
uses=FooExecutor, uses_requests={'/non_foo': 'foo'}
) as f:
c = Client(port=exposed_port, return_responses=True)
resp1 = c.post(on='/foo', inputs=DocumentArray([Document(text='')]))
resp2 = c.post(
on='/non_foo',
inputs=DocumentArray([Document(text='')]),
)
assert resp1[0].docs[0].text == ''
assert resp2[0].docs[0].text == 'foo called'
def test_blob_transmission(decode, protocol):
decode = False
f = Flow(protocol=protocol).add(uses=MyExec)
with f:
c = Client(port=f.port, protocol=protocol)
d = c.post('/', Document(blob=b'hello'), parameters={'decode': decode})[0]
if decode: # test that the Executor gets the correct data
assert d.text == 'hello'
else: # test that the response contains the correct data
assert d.blob == b'hello'
async def test_pods_shards(polling, port_generator):
head_port = port_generator()
port = port_generator()
graph_description = '{"start-gateway": ["pod0"], "pod0": ["end-gateway"]}'
pod_addresses = f'{{"pod0": ["0.0.0.0:{head_port}"]}}'
# create the shards
shard_pods = []
connection_list_dict = {}
for i in range(10):
# create worker
worker_port = port_generator()
# create a single worker pod
worker_pod = _create_worker_pod(worker_port, f'pod0/shard/{i}')
shard_pods.append(worker_pod)
worker_pod.start()
connection_list_dict[i] = [f'127.0.0.1:{worker_port}']
await asyncio.sleep(0.1)
# create a single head pod
head_pod = _create_head_pod(head_port, connection_list_dict, 'head',
polling)
head_pod.start()
head_pod.wait_start_success()
for i, pod in enumerate(shard_pods):
# this would be done by the Pod, its adding the worker to the head
pod.wait_start_success()
# create a single gateway pod
gateway_pod = _create_gateway_pod(graph_description, pod_addresses, port)
gateway_pod.start()
await asyncio.sleep(1.0)
gateway_pod.wait_start_success()
c = Client(host='localhost', port=port, asyncio=True)
responses = c.post('/',
inputs=async_inputs,
request_size=1,
return_responses=True)
response_list = []
async for response in responses:
response_list.append(response)
# clean up pods
gateway_pod.close()
head_pod.close()
for shard_pod in shard_pods:
shard_pod.close()
assert len(response_list) == 20
assert len(
response_list[0].docs) == 1 if polling == 'ANY' else len(shard_pods)
def test_flow_client_defaults():
exposed_port = 12345
f = Flow(port=exposed_port).add(uses=SimplExecutor)
c = Client(port=exposed_port)
with f:
docs = f.post(on='/index', inputs=[Document()])
results = c.post(on='/index', inputs=[Document()])
assert isinstance(docs, DocumentArray)
assert docs[0].text == 'Hello World!'
assert isinstance(results, DocumentArray)
assert results[0].text == 'Hello World!'
async def test_pseudo_remote_pods_topologies(gateway, head, worker):
"""
g(l)-h(l)-w(l) - works
g(l)-h(l)-w(r) - works - head connects to worker via localhost
g(l)-h(r)-w(r) - works - head (inside docker) connects to worker via dockerhost
g(l)-h(r)-w(l) - doesn't work remote head need remote worker
g(r)-... - doesn't work, as distributed parser not enabled for gateway
After any 1 failure, segfault
"""
worker_port = random_port()
head_port = random_port()
port_expose = random_port()
graph_description = (
'{"start-gateway": ["deployment0"], "deployment0": ["end-gateway"]}')
if head == 'remote':
deployments_addresses = f'{{"deployment0": ["{HOST}:{head_port}"]}}'
else:
deployments_addresses = f'{{"deployment0": ["0.0.0.0:{head_port}"]}}'
# create a single head pod
head_pod = _create_head_pod(head, head_port)
# create a single worker pod
worker_pod = _create_worker_pod(worker, worker_port)
# create a single gateway pod
gateway_pod = _create_gateway_pod(gateway, graph_description,
deployments_addresses, port_expose)
with gateway_pod, worker_pod, head_pod:
await asyncio.sleep(1.0)
# this would be done by the deployment, its adding the worker to the head
activate_msg = ControlRequest(command='ACTIVATE')
worker_host, worker_port = worker_pod.runtime_ctrl_address.split(':')
if head == 'remote':
worker_host = __docker_host__
activate_msg.add_related_entity('worker', worker_host,
int(worker_port))
assert GrpcConnectionPool.send_request_sync(
activate_msg, head_pod.runtime_ctrl_address)
# send requests to the gateway
c = Client(host='127.0.0.1', port=port_expose, asyncio=True)
responses = c.post('/',
inputs=async_inputs,
request_size=1,
return_results=True)
response_list = []
async for response in responses:
response_list.append(response)
assert len(response_list) == 20
assert len(response_list[0].docs) == 1
def test_flow_client_defaults():
exposed_port = 12345
f = Flow(port_expose=exposed_port).add(uses=SimplExecutor)
c = Client(port=exposed_port)
with f:
docs = f.post(on='/index', inputs=[Document()], return_results=True)
results = c.post(on='/index', inputs=[Document()], return_results=True)
assert isinstance(docs, DocumentArray)
assert docs[0].text == 'Hello World!'
assert isinstance(results, list)
assert isinstance(results[0], types.request.data.DataRequest)
assert results[0].docs[0].text == 'Hello World!'
async def test_pods_shards(polling, port_generator):
head_port = port_generator()
port_expose = port_generator()
graph_description = '{"start-gateway": ["pod0"], "pod0": ["end-gateway"]}'
pod_addresses = f'{{"pod0": ["0.0.0.0:{head_port}"]}}'
# create a single head pod
head_pod = _create_head_pod(head_port, 'head', polling)
head_pod.start()
# create the shards
shard_pods = []
for i in range(10):
# create worker
worker_port = port_generator()
# create a single worker pod
worker_pod = _create_worker_pod(worker_port, f'pod0/shard/{i}')
shard_pods.append(worker_pod)
worker_pod.start()
await asyncio.sleep(0.1)
head_pod.wait_start_success()
for i, pod in enumerate(shard_pods):
# this would be done by the Pod, its adding the worker to the head
pod.wait_start_success()
activate_msg = ControlRequest(command='ACTIVATE')
activate_msg.add_related_entity(
'worker', '127.0.0.1', pod.args.port_in, shard_id=i
)
GrpcConnectionPool.send_request_sync(activate_msg, f'127.0.0.1:{head_port}')
# create a single gateway pod
gateway_pod = _create_gateway_pod(graph_description, pod_addresses, port_expose)
gateway_pod.start()
await asyncio.sleep(1.0)
gateway_pod.wait_start_success()
c = Client(host='localhost', port=port_expose, asyncio=True)
responses = c.post('/', inputs=async_inputs, request_size=1, return_results=True)
response_list = []
async for response in responses:
response_list.append(response)
# clean up pods
gateway_pod.close()
head_pod.close()
for shard_pod in shard_pods:
shard_pod.close()
assert len(response_list) == 20
assert len(response_list[0].docs) == 1 if polling == 'ANY' else len(shard_pods)
async def test_runtimes_gateway_worker_direct_connection(port_generator):
worker_port = port_generator()
port = port_generator()
graph_description = '{"start-gateway": ["pod0"], "pod0": ["end-gateway"]}'
pod_addresses = f'{{"pod0": ["0.0.0.0:{worker_port}"]}}'
# create the shards
worker_process = multiprocessing.Process(target=_create_worker_runtime,
args=(worker_port, f'pod0'))
worker_process.start()
await asyncio.sleep(0.1)
# create a single gateway runtime
gateway_process = multiprocessing.Process(
target=_create_gateway_runtime,
args=(graph_description, pod_addresses, port),
)
gateway_process.start()
await asyncio.sleep(1.0)
AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
timeout=5.0,
ctrl_address=f'0.0.0.0:{port}',
ready_or_shutdown_event=multiprocessing.Event(),
)
c = Client(host='localhost', port=port, asyncio=True)
responses = c.post('/',
inputs=async_inputs,
request_size=1,
return_responses=True)
response_list = []
async for response in responses:
response_list.append(response)
# clean up runtimes
gateway_process.terminate()
worker_process.terminate()
gateway_process.join()
worker_process.join()
assert len(response_list) == 20
assert len(response_list[0].docs) == 1
assert gateway_process.exitcode == 0
assert worker_process.exitcode == 0
async def test_deployments_with_executor(port_generator):
graph_description = (
'{"start-gateway": ["deployment0"], "deployment0": ["end-gateway"]}')
head_port = port_generator()
deployments_addresses = f'{{"deployment0": ["0.0.0.0:{head_port}"]}}'
regular_deployment = _create_regular_deployment(
port=head_port,
name='deployment',
executor='NameChangeExecutor',
uses_before=True,
uses_after=True,
polling=PollingType.ANY,
shards=2,
)
regular_deployment.start()
port = port_generator()
gateway_deployment = _create_gateway_deployment(graph_description,
deployments_addresses,
port)
gateway_deployment.start()
await asyncio.sleep(1.0)
c = Client(host='localhost',
port=port,
asyncio=True,
return_responses=True)
responses = c.post('/', inputs=async_inputs, request_size=1)
response_list = []
async for response in responses:
response_list.append(response.docs)
gateway_deployment.close()
regular_deployment.close()
assert len(response_list) == 20
assert len(response_list[0]) == 4
doc_texts = [doc.text for doc in response_list[0]]
assert doc_texts.count('client0-Request') == 1
assert doc_texts.count('deployment/uses_before-0') == 1
assert doc_texts.count('deployment/uses_after-0') == 1
async def test_deployments_with_replicas_advance_faster(port_generator):
head_port = port_generator()
port = port_generator()
graph_description = (
'{"start-gateway": ["deployment0"], "deployment0": ["end-gateway"]}')
deployment = _create_regular_deployment(port=head_port,
name='deployment',
executor='FastSlowExecutor',
replicas=10)
deployment.start()
connections = [f'0.0.0.0:{port}' for port in deployment.ports]
deployments_addresses = f'{{"deployment0": {json.dumps(connections)}}}'
gateway_deployment = _create_gateway_deployment(graph_description,
deployments_addresses,
port)
gateway_deployment.start()
await asyncio.sleep(1.0)
c = Client(host='localhost',
port=port,
asyncio=True,
return_responses=True)
input_docs = [Document(text='slow'), Document(text='fast')]
responses = c.post('/', inputs=input_docs, request_size=1)
response_list = []
async for response in responses:
response_list.append(response)
gateway_deployment.close()
deployment.close()
assert len(response_list) == 2
for response in response_list:
assert len(response.docs) == 1
assert response_list[0].docs[0].text == 'fast'
assert response_list[1].docs[0].text == 'slow'
async def test_pods_trivial_topology(port_generator):
worker_port = port_generator()
head_port = port_generator()
port = port_generator()
graph_description = '{"start-gateway": ["pod0"], "pod0": ["end-gateway"]}'
pod_addresses = f'{{"pod0": ["0.0.0.0:{head_port}"]}}'
# create a single worker pod
worker_pod = _create_worker_pod(worker_port)
# create a single head pod
head_pod = _create_head_pod(head_port)
# create a single gateway pod
gateway_pod = _create_gateway_pod(graph_description, pod_addresses, port)
with gateway_pod, head_pod, worker_pod:
# this would be done by the Pod, its adding the worker to the head
head_pod.wait_start_success()
worker_pod.wait_start_success()
activate_msg = ControlRequest(command='ACTIVATE')
activate_msg.add_related_entity('worker', '127.0.0.1', worker_port)
assert GrpcConnectionPool.send_request_sync(
activate_msg, f'127.0.0.1:{head_port}'
)
# send requests to the gateway
gateway_pod.wait_start_success()
c = Client(host='localhost', port=port, asyncio=True)
responses = c.post(
'/', inputs=async_inputs, request_size=1, return_responses=True
)
response_list = []
async for response in responses:
response_list.append(response)
assert len(response_list) == 20
assert len(response_list[0].docs) == 1
async def test_pods_trivial_topology(port_generator):
worker_port = port_generator()
head_port = port_generator()
port = port_generator()
graph_description = '{"start-gateway": ["pod0"], "pod0": ["end-gateway"]}'
pod_addresses = f'{{"pod0": ["0.0.0.0:{head_port}"]}}'
# create a single worker pod
worker_pod = _create_worker_pod(worker_port)
# create a single head pod
connection_list_dict = {'0': [f'127.0.0.1:{worker_port}']}
head_pod = _create_head_pod(head_port, connection_list_dict)
# create a single gateway pod
gateway_pod = _create_gateway_pod(graph_description, pod_addresses, port)
with gateway_pod, head_pod, worker_pod:
# this would be done by the Pod, its adding the worker to the head
head_pod.wait_start_success()
worker_pod.wait_start_success()
# send requests to the gateway
gateway_pod.wait_start_success()
c = Client(host='localhost', port=port, asyncio=True)
responses = c.post('/',
inputs=async_inputs,
request_size=1,
return_responses=True)
response_list = []
async for response in responses:
response_list.append(response)
assert len(response_list) == 20
assert len(response_list[0].docs) == 1
def test_crud(tmpdir, rest):
os.environ['RESTFUL'] = 'http' if rest else 'grpc'
os.environ['WORKSPACE'] = str(tmpdir)
with Flow.load_config('flow.yml') as f:
c = Client(port=f.port, return_responses=True)
original_docs = list(random_docs(10, chunks_per_doc=0))
if rest:
rest_post(f, 'index', original_docs)
else:
c.post(
on='/index',
inputs=original_docs,
)
with Flow.load_config('flow.yml') as f:
c = Client(port=f.port, return_responses=True)
inputs = list(random_docs(1))
if rest:
results = rest_post(f, 'search', inputs)
matches = results['data'][0]['matches']
for doc in results['data']:
assert Document.from_dict(doc).text == 'hello world'
else:
results = c.post(on='/search', inputs=inputs, parameters=PARAMS)
matches = results[0].docs[0].matches
for doc in results[0].docs:
assert doc.text == 'hello world'
assert len(matches) == 10
with Flow.load_config('flow.yml') as f:
c = Client(port=f.port, return_responses=True)
inputs = list(random_docs(5, chunks_per_doc=0))
if rest:
rest_post(f, 'delete', inputs)
else:
c.post(on='/delete', inputs=inputs)
with Flow.load_config('flow.yml') as f:
c = Client(port=f.port, return_responses=True)
inputs = list(random_docs(1))
if rest:
results = rest_post(f, 'search', inputs)
matches = results['data'][0]['matches']
else:
results = c.post(on='/search', inputs=inputs, parameters=PARAMS)
matches = results[0].docs[0].matches
assert len(matches) == 5
updated_docs = list(
random_docs(5, chunks_per_doc=5, start_id=5, text='hello again')
)
with Flow.load_config('flow.yml') as f:
c = Client(port=f.port, return_responses=True)
if rest:
rest_post(f, 'update', updated_docs)
else:
c.post(on='/update', inputs=updated_docs)
with Flow.load_config('flow.yml') as f:
c = Client(port=f.port, return_responses=True)
inputs = list(random_docs(1))
if rest:
results = rest_post(f, 'search', inputs)
matches = sorted(
results['data'][0]['matches'], key=lambda match: match['id']
)
else:
results = c.post(on='/search', inputs=inputs, parameters=PARAMS)
matches = sorted(results[0].docs[0].matches, key=lambda match: match.id)
assert len(matches) == 5
for match, updated_doc in zip(matches, updated_docs):
if isinstance(match, dict):
match = Document.from_dict(match)
assert updated_doc.id == match.id
assert updated_doc.text == match.text
np.testing.assert_array_equal(updated_doc.embedding, match.embedding)
assert len(match.chunks) == 5
assert len(match.chunks) == len(updated_doc.chunks)
for match_chunk, updated_doc_chunk in zip(match.chunks, updated_doc.chunks):
assert match_chunk.text == updated_doc_chunk.text
np.testing.assert_array_equal(
match_chunk.embedding, updated_doc_chunk.embedding
)
async def test_runtimes_with_replicas_advance_faster(port_generator):
head_port = port_generator()
port = port_generator()
graph_description = '{"start-gateway": ["pod0"], "pod0": ["end-gateway"]}'
pod_addresses = f'{{"pod0": ["0.0.0.0:{head_port}"]}}'
# create a single head runtime
head_process = multiprocessing.Process(target=_create_head_runtime,
args=(head_port, 'head'))
head_process.start()
# create the shards
replica_processes = []
worker_ports = []
for i in range(10):
# create worker
worker_port = port_generator()
# create a single worker runtime
worker_process = multiprocessing.Process(
target=_create_worker_runtime,
args=(worker_port, f'pod0/{i}', 'FastSlowExecutor'),
)
replica_processes.append(worker_process)
worker_process.start()
await asyncio.sleep(0.1)
worker_ports.append(worker_port)
await _activate_runtimes(head_port, worker_ports)
# create a single gateway runtime
gateway_process = multiprocessing.Process(
target=_create_gateway_runtime,
args=(graph_description, pod_addresses, port),
)
gateway_process.start()
await asyncio.sleep(1.0)
AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
timeout=5.0,
ctrl_address=f'0.0.0.0:{port}',
ready_or_shutdown_event=multiprocessing.Event(),
)
c = Client(host='localhost', port=port, asyncio=True)
input_docs = [Document(text='slow'), Document(text='fast')]
responses = c.post('/',
inputs=input_docs,
request_size=1,
return_responses=True)
response_list = []
async for response in responses:
response_list.append(response)
# clean up runtimes
gateway_process.terminate()
head_process.terminate()
for replica_process in replica_processes:
replica_process.terminate()
gateway_process.join()
head_process.join()
for replica_process in replica_processes:
replica_process.join()
assert len(response_list) == 2
for response in response_list:
assert len(response.docs) == 1
assert response_list[0].docs[0].text == 'fast'
assert response_list[1].docs[0].text == 'slow'
assert gateway_process.exitcode == 0
assert head_process.exitcode == 0
for replica_process in replica_processes:
assert replica_process.exitcode == 0
async def test_runtimes_with_executor(port_generator):
graph_description = '{"start-gateway": ["pod0"], "pod0": ["end-gateway"]}'
runtime_processes = []
uses_before_port, uses_before_process = await _create_worker(
'pod0',
port_generator,
type='uses_before',
executor='NameChangeExecutor')
runtime_processes.append(uses_before_process)
uses_after_port, uses_after_process = await _create_worker(
'pod0',
port_generator,
type='uses_after',
executor='NameChangeExecutor')
runtime_processes.append(uses_after_process)
# create head
head_port = port_generator()
pod_addresses = f'{{"pod0": ["0.0.0.0:{head_port}"]}}'
head_process = multiprocessing.Process(
target=_create_head_runtime,
args=(
head_port,
f'pod0/head',
'ALL',
f'127.0.0.1:{uses_before_port}',
f'127.0.0.1:{uses_after_port}',
),
)
runtime_processes.append(head_process)
head_process.start()
runtime_processes.append(head_process)
# create some shards
worker_ports = []
for i in range(10):
# create worker
worker_port, worker_process = await _create_worker(
'pod0',
port_generator,
type=f'shards/{i}',
executor='NameChangeExecutor')
runtime_processes.append(worker_process)
await asyncio.sleep(0.1)
worker_ports.append(worker_port)
await _activate_runtimes(head_port, worker_ports)
# create a single gateway runtime
port = port_generator()
gateway_process = multiprocessing.Process(
target=_create_gateway_runtime,
args=(graph_description, pod_addresses, port),
)
gateway_process.start()
runtime_processes.append(gateway_process)
await asyncio.sleep(1.0)
AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
timeout=5.0,
ctrl_address=f'0.0.0.0:{port}',
ready_or_shutdown_event=multiprocessing.Event(),
)
c = Client(host='localhost', port=port, asyncio=True)
responses = c.post('/',
inputs=async_inputs,
request_size=1,
return_responses=True)
response_list = []
async for response in responses:
response_list.append(response.docs)
# clean up runtimes
for process in runtime_processes:
process.terminate()
for process in runtime_processes:
process.join()
assert len(response_list) == 20
assert (
len(response_list[0]) == (1 + 1 + 1) * 10 + 1
) # 1 starting doc + 1 uses_before + every exec adds 1 * 10 shards + 1 doc uses_after
doc_texts = [doc.text for doc in response_list[0]]
assert doc_texts.count('client0-Request') == 10
assert doc_texts.count('pod0/uses_before') == 10
assert doc_texts.count('pod0/uses_after') == 1
for i in range(10):
assert doc_texts.count(f'pod0/shards/{i}') == 1
async def test_runtimes_shards(polling, port_generator):
head_port = port_generator()
port = port_generator()
graph_description = '{"start-gateway": ["pod0"], "pod0": ["end-gateway"]}'
pod_addresses = f'{{"pod0": ["0.0.0.0:{head_port}"]}}'
# create a single head runtime
head_process = multiprocessing.Process(target=_create_head_runtime,
args=(head_port, 'head', polling))
head_process.start()
# create the shards
shard_processes = []
worker_ports = []
for i in range(10):
# create worker
worker_port = port_generator()
# create a single worker runtime
worker_process = multiprocessing.Process(target=_create_worker_runtime,
args=(worker_port,
f'pod0/shard/{i}'))
shard_processes.append(worker_process)
worker_process.start()
await asyncio.sleep(0.1)
worker_ports.append(worker_port)
await _activate_runtimes(head_port, worker_ports)
# create a single gateway runtime
gateway_process = multiprocessing.Process(
target=_create_gateway_runtime,
args=(graph_description, pod_addresses, port),
)
gateway_process.start()
await asyncio.sleep(1.0)
AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
timeout=5.0,
ctrl_address=f'0.0.0.0:{port}',
ready_or_shutdown_event=multiprocessing.Event(),
)
c = Client(host='localhost', port=port, asyncio=True)
responses = c.post('/',
inputs=async_inputs,
request_size=1,
return_responses=True)
response_list = []
async for response in responses:
response_list.append(response)
# clean up runtimes
gateway_process.terminate()
head_process.terminate()
for shard_process in shard_processes:
shard_process.terminate()
gateway_process.join()
head_process.join()
for shard_process in shard_processes:
shard_process.join()
assert len(response_list) == 20
assert len(response_list[0].docs) == 1 if polling == 'ANY' else len(
shard_processes)
assert gateway_process.exitcode == 0
assert head_process.exitcode == 0
for shard_process in shard_processes:
assert shard_process.exitcode == 0
async def test_runtimes_trivial_topology(port_generator):
worker_port = port_generator()
head_port = port_generator()
port = port_generator()
graph_description = '{"start-gateway": ["pod0"], "pod0": ["end-gateway"]}'
pod_addresses = f'{{"pod0": ["0.0.0.0:{head_port}"]}}'
# create a single worker runtime
worker_process = multiprocessing.Process(target=_create_worker_runtime,
args=(worker_port, ))
worker_process.start()
# create a single head runtime
head_process = multiprocessing.Process(target=_create_head_runtime,
args=(head_port, ))
head_process.start()
# create a single gateway runtime
gateway_process = multiprocessing.Process(
target=_create_gateway_runtime,
args=(graph_description, pod_addresses, port),
)
gateway_process.start()
await asyncio.sleep(1.0)
AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
timeout=5.0,
ctrl_address=f'0.0.0.0:{head_port}',
ready_or_shutdown_event=multiprocessing.Event(),
)
AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
timeout=5.0,
ctrl_address=f'0.0.0.0:{worker_port}',
ready_or_shutdown_event=multiprocessing.Event(),
)
AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
timeout=5.0,
ctrl_address=f'0.0.0.0:{port}',
ready_or_shutdown_event=multiprocessing.Event(),
)
# this would be done by the Pod, its adding the worker to the head
activate_msg = ControlRequest(command='ACTIVATE')
activate_msg.add_related_entity('worker', '127.0.0.1', worker_port)
GrpcConnectionPool.send_request_sync(activate_msg,
f'127.0.0.1:{head_port}')
# send requests to the gateway
c = Client(host='localhost', port=port, asyncio=True)
responses = c.post('/',
inputs=async_inputs,
request_size=1,
return_responses=True)
response_list = []
async for response in responses:
response_list.append(response)
# clean up runtimes
gateway_process.terminate()
head_process.terminate()
worker_process.terminate()
gateway_process.join()
head_process.join()
worker_process.join()
assert len(response_list) == 20
assert len(response_list[0].docs) == 1
assert gateway_process.exitcode == 0
assert head_process.exitcode == 0
assert worker_process.exitcode == 0
