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 _create_head_runtime(port,
name='',
polling='ANY',
uses_before=None,
uses_after=None):
args = set_pod_parser().parse_args([])
args.port = port
args.name = name
args.polling = PollingType.ANY if polling == 'ANY' else PollingType.ALL
if uses_before:
args.uses_before_address = uses_before
if uses_after:
args.uses_after_address = uses_after
with HeadRuntime(args) as runtime:
runtime.run_forever()
def start_runtime(args, handle_queue, cancel_event):
def _send_requests_mock(request: List[Request], connection,
endpoint) -> asyncio.Task:
async def mock_task_wrapper(new_requests, *args, **kwargs):
handle_queue.put('mock_called')
await asyncio.sleep(0.1)
return new_requests[0], grpc.aio.Metadata.from_tuple(
(('is-error', 'true'), ))
return asyncio.create_task(mock_task_wrapper(request, connection))
if not hasattr(args, 'name') or not args.name:
args.name = 'testHead'
with HeadRuntime(args, cancel_event) as runtime:
runtime.connection_pool._send_requests = _send_requests_mock
runtime.run_forever()
def _create_head_runtime(
port,
connection_list_dict,
name='',
polling='ANY',
uses_before=None,
uses_after=None,
retries=-1,
):
args = set_pod_parser().parse_args([])
args.port = port
args.name = name
args.retries = retries
args.polling = PollingType.ANY if polling == 'ANY' else PollingType.ALL
if uses_before:
args.uses_before_address = uses_before
if uses_after:
args.uses_after_address = uses_after
args.connection_list = json.dumps(connection_list_dict)
with HeadRuntime(args) as runtime:
runtime.run_forever()
def _destroy_runtime(args, cancel_event, runtime_thread):
cancel_event.set()
runtime_thread.join()
assert not HeadRuntime.is_ready(f'{args.host}:{args.port}')
def run(*args, **kwargs):
runtime_args = set_pod_parser().parse_args(args)
with HeadRuntime(runtime_args) as runtime:
runtime.run_forever()