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