def activate_worker_sync(
worker_host: str,
worker_port: int,
target_head: str,
shard_id: Optional[int] = None,
) -> ControlRequest:
"""
Register a given worker to a head by sending an activate request
:param worker_host: the host address of the worker
:param worker_port: the port of the worker
:param target_head: address of the head to send the activate request to
:param shard_id: id of the shard the worker belongs to
:returns: the response request
"""
activate_request = ControlRequest(command='ACTIVATE')
activate_request.add_related_entity('worker', worker_host, worker_port,
shard_id)
if os.name != 'nt':
os.unsetenv('http_proxy')
os.unsetenv('https_proxy')
return GrpcConnectionPool.send_request_sync(activate_request,
target_head)
async def _activate_worker(head_port, worker_port, shard_id=None):
# 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, shard_id=shard_id
)
GrpcConnectionPool.send_request_sync(activate_msg, f'127.0.0.1:{head_port}')
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
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
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 head pod
head_pod = _create_head_pod(head_port, 'head')
head_pod.start()
# create a single gateway pod
gateway_pod = _create_gateway_pod(graph_description, pod_addresses, port)
gateway_pod.start()
# create the shards
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')
pods.append(worker_pod)
worker_pod.start()
await asyncio.sleep(0.1)
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()
activate_msg = ControlRequest(command='ACTIVATE')
activate_msg.add_related_entity('worker', '127.0.0.1', pod.args.port)
GrpcConnectionPool.send_request_sync(activate_msg, f'127.0.0.1:{head_port}')
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_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
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 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, 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)
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)
async def deactivate_worker(
worker_host: str,
worker_port: int,
target_head: str,
shard_id: Optional[int] = None,
) -> ControlRequest:
"""
Remove a given worker to a head by sending a deactivate request
:param worker_host: the host address of the worker
:param worker_port: the port of the worker
:param target_head: address of the head to send the deactivate request to
:param shard_id: id of the shard the worker belongs to
:returns: the response request
"""
activate_request = ControlRequest(command='DEACTIVATE')
activate_request.add_related_entity('worker', worker_host, worker_port,
shard_id)
return await GrpcConnectionPool.send_request_async(
activate_request, target_head)
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_health_check(port_generator, protocol, health_check):
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, protocol)
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()
for _port in (head_port, worker_port):
check_health_pod(f'0.0.0.0:{_port}')
if inspect.iscoroutinefunction(health_check):
await health_check(f'0.0.0.0:{port}')
else:
health_check(f'0.0.0.0:{port}')
def _remove_worker(args, ip='fake_ip', shard_id=None):
activate_msg = ControlRequest(command='DEACTIVATE')
activate_msg.add_related_entity('worker', ip, 8080, shard_id)
assert GrpcConnectionPool.send_request_sync(activate_msg,
f'{args.host}:{args.port}')
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
