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()