def test_parse_env_map(): a = set_pod_parser().parse_args( ['--env', 'key1=value1', '--env', 'key2=value2']) assert a.env == {'key1': 'value1', 'key2': 'value2'} a = set_pod_parser().parse_args(['--env', 'key1=value1', 'key2=value2']) assert a.env == {'key1': 'value1', 'key2': 'value2'}
def test_pod_factory_pod(): args_no_replicas = set_pod_parser().parse_args(['--shards', '1']) assert isinstance(PodFactory.build_pod(args_no_replicas), Pod) args_replicas = set_pod_parser().parse_args(['--shards', '2']) assert isinstance(PodFactory.build_pod(args_replicas), CompoundPod) args_no_replicas = set_pod_parser().parse_args(['--replicas', '2']) assert isinstance(PodFactory.build_pod(args_no_replicas), Pod)
def test_pod_args_remove_uses_ba(): args = set_pod_parser().parse_args([]) with Pod(args) as p: assert p.num_peas == 1 args = set_pod_parser().parse_args( ['--uses-before', '_pass', '--uses-after', '_pass']) with Pod(args) as p: assert p.num_peas == 1 args = set_pod_parser().parse_args( ['--uses-before', '_pass', '--uses-after', '_pass', '--parallel', '2']) with Pod(args) as p: assert p.num_peas == 4
def test_shards_parallel_synonym(): a = set_pod_parser().parse_args(['--shards', '2']) assert a.parallel == 2 with pytest.raises(AttributeError): a.shards a = set_pod_parser().parse_args(['--parallel', '2']) assert a.parallel == 2 with pytest.raises(AttributeError): a.shards a = set_pod_parser().parse_args([]) assert a.parallel == 1 with pytest.raises(AttributeError): a.shards
def external_pod_parallel_2_args( port_in_external, port_out_external, num_replicas, num_parallel ): args = [ '--uses', 'MyExternalExecutor', '--name', 'external_real_2', '--socket-in', 'SUB_CONNECT', '--port-in', str(port_in_external), '--host-in', '0.0.0.0', '--port-out', str(port_out_external), '--host-out', '0.0.0.0', '--socket-out', 'PUSH_CONNECT', '--parallel', str(num_parallel), '--replicas', str(num_replicas), ] return set_pod_parser().parse_args(args)
def test_parse_args_custom_executor(shards: int): uses_before = 'custom-executor-before' uses_after = 'custom-executor-after' args = set_pod_parser().parse_args( [ '--shards', str(shards), '--uses-before', uses_before, '--uses-after', uses_after, ] ) pod = K8sPod(args) assert namespace_equal( args, pod.deployment_args['head_deployment'], skip_attr=('uses',) ) assert pod.deployment_args['head_deployment'].uses == uses_before assert namespace_equal( args, pod.deployment_args['tail_deployment'], skip_attr=('uses',) ) assert pod.deployment_args['tail_deployment'].uses == uses_after for i, depl_arg in enumerate(pod.deployment_args['deployments']): import copy cargs = copy.deepcopy(args) cargs.shard_id = i assert depl_arg == cargs
def test_pod_with_gpus(mocker): args_list = [ '--name', 'test-wait-success', '--k8s-namespace', 'test-namespace', '--shards', '1', '--replicas', '1', '--gpus', '3', ] args = set_pod_parser().parse_args(args_list) container = client.V1Container( name='test-container', resources=client.V1ResourceRequirements(limits={'nvidia.com/gpu': 3}), ) spec = client.V1PodSpec(containers=[container]) mocker.patch( 'jina.peapods.pods.k8s.K8sPod._K8sDeployment._read_namespaced_deployment', return_value=client.V1Deployment( status=client.V1DeploymentStatus(replicas=1, ready_replicas=1), spec=spec ), ) mocker.patch( 'jina.peapods.pods.k8slib.kubernetes_deployment.deploy_service', return_value=f'test-wait-success.test-namespace.svc', ) mocker.patch( 'jina.peapods.pods.k8s.K8sPod._K8sDeployment._delete_namespaced_deployment', return_value=client.V1Status(status=200), ) with K8sPod(args) as pod: assert pod.args.gpus == '3'
def test_pod_start_close_given_tail_deployment( name, k8s_namespace, shards, mocker, uses_after ): args_list = [ '--name', name, '--k8s-namespace', k8s_namespace, '--shards', shards, '--noblock-on-start', ] if uses_after: args_list.extend(['--uses-after', 'custom-executor-after']) args = set_pod_parser().parse_args(args_list) mocker.patch( 'jina.peapods.pods.k8slib.kubernetes_deployment.deploy_service', return_value=f'{name}.{k8s_namespace}.svc', ) mocker.patch( 'jina.peapods.pods.k8s.K8sPod._K8sDeployment._delete_namespaced_deployment', return_value=client.V1Status(status=200), ) with K8sPod(args) as pod: # enter `_deploy_runtime` assert isinstance(pod.k8s_tail_deployment, K8sPod._K8sDeployment) assert pod.k8s_tail_deployment.name == name + '-tail' assert pod.args.noblock_on_start
def test_scale_remote_pod(pod_params): num_replicas, scale_to, shards = pod_params pod_id = None args = set_pod_parser().parse_args( ['--replicas', str(num_replicas), '--shards', str(shards)] ) payload = replace_enum_to_str(ArgNamespace.flatten_to_dict(args)) try: workspace_id = jinad_client.workspaces.create( paths=[os.path.join(cur_dir, cur_dir)] ) success, pod_id = jinad_client.pods.create( workspace_id=workspace_id, payload=payload ) assert success remote_pod_args = jinad_client.pods.get(pod_id)['arguments']['object'][ 'arguments' ] assert remote_pod_args['identity'] == pod_id assert remote_pod_args['replicas'] == num_replicas assert remote_pod_args['shards'] == shards jinad_client.pods.scale(id=pod_id, replicas=scale_to) remote_pod_args = jinad_client.pods.get(pod_id)['arguments']['object'][ 'arguments' ] assert remote_pod_args['identity'] == pod_id assert remote_pod_args['replicas'] == scale_to assert remote_pod_args['shards'] == shards finally: if pod_id: assert jinad_client.pods.delete(pod_id), 'Pod termination failed' print(f'Remote Pod {pod_id} successfully terminated')
def test_pod_remote_pea_without_replicas(): args = set_pod_parser().parse_args( ['--peas-hosts', '0.0.0.1', '--replicas', str(1)]) with Pod(args) as pod: pea = pod.replica_set._peas[0] assert pea.args.host == pod.host
def test_pod_factory_k8s(): args_replicas = set_pod_parser().parse_args([]) assert isinstance( PodFactory.build_pod(args_replicas, infrastructure=InfrastructureType.K8S), K8sPod, )
async def test_decorator_monitoring(port_generator): from jina import monitor class DummyExecutor(Executor): @requests def foo(self, docs, **kwargs): self._proces(docs) self.process_2(docs) @monitor(name='metrics_name', documentation='metrics description') def _proces(self, docs): ... @monitor() def process_2(self, docs): ... port = port_generator() args = set_pod_parser().parse_args([ '--monitoring', '--port-monitoring', str(port), '--uses', 'DummyExecutor' ]) cancel_event = multiprocessing.Event() def start_runtime(args, cancel_event): with WorkerRuntime(args, cancel_event=cancel_event) as runtime: runtime.run_forever() runtime_thread = Process( target=start_runtime, args=(args, cancel_event), daemon=True, ) runtime_thread.start() assert AsyncNewLoopRuntime.wait_for_ready_or_shutdown( timeout=5.0, ctrl_address=f'{args.host}:{args.port}', ready_or_shutdown_event=Event(), ) assert AsyncNewLoopRuntime.wait_for_ready_or_shutdown( timeout=5.0, ctrl_address=f'{args.host}:{args.port}', ready_or_shutdown_event=Event(), ) await GrpcConnectionPool.send_request_async(_create_test_data_message(), f'{args.host}:{args.port}', timeout=1.0) resp = req.get(f'http://localhost:{port}/') assert f'jina_metrics_name_count{{runtime_name="None"}} 1.0' in str( resp.content) cancel_event.set() runtime_thread.join() assert not AsyncNewLoopRuntime.is_ready(f'{args.host}:{args.port}')
async def test_worker_runtime_reflection(): args = set_pod_parser().parse_args([]) cancel_event = multiprocessing.Event() def start_runtime(args, cancel_event): with WorkerRuntime(args, cancel_event=cancel_event) as runtime: runtime.run_forever() runtime_thread = Process( target=start_runtime, args=(args, cancel_event), daemon=True, ) runtime_thread.start() assert AsyncNewLoopRuntime.wait_for_ready_or_shutdown( timeout=3.0, ctrl_address=f'{args.host}:{args.port}', ready_or_shutdown_event=Event(), ) async with grpc.aio.insecure_channel( f'{args.host}:{args.port}') as channel: service_names = await GrpcConnectionPool.get_available_services(channel ) assert all(service_name in service_names for service_name in [ 'jina.JinaDataRequestRPC', 'jina.JinaSingleDataRequestRPC', ]) cancel_event.set() runtime_thread.join() assert not AsyncNewLoopRuntime.is_ready(f'{args.host}:{args.port}')
def test_pod_with_sse_no_deadlock_log_file(): os.environ['JINA_LOG_FILE'] = 'TXT' args = set_pod_parser().parse_args(['--parallel', '2']) p = Pod(args) with p: pass del os.environ['JINA_LOG_FILE']
def test_compound_pod_do_not_forward_uses_before_uses_after(): polling = 'all' shards = 2 replicas = 2 args = set_pod_parser().parse_args( [ '--name', 'pod', '--shards', f'{shards}', '--polling', f'{polling}', '--replicas', f'{replicas}', '--uses-after', 'BaseExecutor', '--uses-before', 'BaseExecutor', ] ) cpod = CompoundPod(args) for pod in cpod.shards: assert pod.args.uses_before is None assert pod.args.uses_after is None
def test_load_cust_with_driver(): a = BaseExecutor.load_config(os.path.join(cur_dir, 'mwu-encoder/mwu_encoder_driver.yml')) assert a._drivers['ControlRequest'][0].__class__.__name__ == 'MyAwesomeDriver' p = set_pod_parser().parse_args(['--uses', os.path.join(cur_dir, 'mwu-encoder/mwu_encoder_driver.yml')]) with Pod(p): # will print a cust task_name from the driver when terminate pass
def test_pod_remote_pea_replicas_pea_host_set_partially( pod_host, pea1_host, expected_host_in, expected_host_out, ): args = set_pod_parser().parse_args([ '--peas-hosts', f'{pea1_host}', '--replicas', str(2), '--host', pod_host ]) assert args.host == pod_host pod = Pod(args) for k, v in pod.peas_args.items(): if k in ['head', 'tail']: assert v.host == args.host else: for pea_arg in v: if pea_arg.pea_id in (0, 1): assert pea_arg.host == pea1_host assert pea_arg.host_in == expected_host_in assert pea_arg.host_out == expected_host_out else: assert pea_arg.host == args.host assert pea_arg.host_in == __default_host__ assert pea_arg.host_out == __default_host__
def test_ssh_pod(): p = set_pod_parser().parse_args( ['--host', '[email protected]', '--timeout', '5000']) with SSHRuntime(p, kind='pod') as pp: assert pp.status.envelope.status.code == jina_pb2.StatusProto.READY assert pp.status is None
def test_pod_upload_files( replicas, upload_files, uses, uses_before, uses_after, py_modules, expected, ): args = set_pod_parser().parse_args([ '--uses', uses, '--uses-before', uses_before, '--uses-after', uses_after, '--py-modules', *py_modules, '--upload-files', *upload_files, '--replicas', str(replicas), ]) pod = Pod(args) for k, v in pod.peas_args.items(): if k in ['head', 'tail']: if v: pass # assert sorted(v.upload_files) == sorted(expected) else: for pea in v: print(sorted(pea.upload_files)) print(sorted(expected)) assert sorted(pea.upload_files) == sorted(expected)
def test_jinad_pod(): args = set_pod_parser().parse_args([]) assert not is_pod_ready(args) with JinaDPod(args): assert is_pod_ready(args) assert not is_pod_ready(args)
def external_pod_join_args( port_in_external, port_out_external, num_replicas, num_parallel ): args = [ '--uses', 'MyExternalExecutor', '--name', 'external_real', '--port-in', str(port_in_external), '--host-in', '0.0.0.0', '--socket-in', 'PULL_BIND', '--port-out', str(port_out_external), '--host-out', '0.0.0.0', '--socket-out', 'PUSH_BIND', '--pod-role', 'JOIN', '--parallel', str(num_parallel), '--replicas', str(num_replicas), ] return set_pod_parser().parse_args(args)
def test_pod_remote_pea_parallel_pea_host_set_completely( pod_host, peas_hosts, expected_host_in, expected_host_out, ): args = set_pod_parser().parse_args( [ '--peas-hosts', f'{peas_hosts[0]}', f'{peas_hosts[1]}', '--parallel', str(2), '--host', pod_host, ] ) assert args.host == pod_host pod = Pod(args) for k, v in pod.peas_args.items(): if k in ['head', 'tail']: assert v.host == args.host else: for pea_arg, pea_host in zip(v, peas_hosts): assert pea_arg.host == pea_host assert pea_arg.host_in == expected_host_in assert pea_arg.host_out == expected_host_out
def pod_to_namespace(args: Union[SinglePodModel, ParallelPodModel]): from jina.parsers import set_pod_parser parser = set_pod_parser() if isinstance(args, ParallelPodModel): pod_args = {} args = args.dict() for pea_type, pea_args in args.items(): # this is for pea_type: head & tail when None (pod with parallel = 1) if pea_args is None: pod_args[pea_type] = None # this is for pea_type: head & tail when not None (pod with parallel > 1) if isinstance(pea_args, Dict): pea_args = handle_remote_args(args=pea_args, parser=parser) pod_args[pea_type] = argparse.Namespace(**pea_args) # this is for pea_type: peas (multiple entries) if isinstance(pea_args, List): pod_args[pea_type] = [] for current_pea_arg in pea_args: current_pea_arg = handle_remote_args(args=current_pea_arg, parser=parser) pod_args[pea_type].append( argparse.Namespace(**current_pea_arg)) return pod_args if isinstance(args, SinglePodModel): pod_args = handle_remote_args(args=args.dict(), parser=parser) return argparse.Namespace(**pod_args)
def external_pod_args(): args = ['--port-in', str(45678), '--port-out', str(45679)] args = vars(set_pod_parser().parse_args(args)) del args['external'] del args['pod_role'] del args['host'] return args
def test_executor_runtimes(signal, tmpdir): import time args = set_pod_parser().parse_args([]) def run(args): args.uses = { 'jtype': 'DummyExecutor', 'with': { 'dir': str(tmpdir) }, 'metas': { 'workspace': str(tmpdir) }, } executor_native(args) process = multiprocessing.Process(target=run, args=(args, )) process.start() time.sleep(0.5) GrpcConnectionPool.send_request_sync(_create_test_data_message(), target=f'{args.host}:{args.port}') time.sleep(0.1) os.kill(process.pid, signal) process.join() with open(f'{tmpdir}/test.txt', 'r') as fp: output = fp.read() split = output.split(';') assert split[0] == 'proper close' assert split[1] == '1'
def _pod_command(self) -> str: from jina.parsers import set_pod_parser non_defaults = ArgNamespace.get_non_defaults_args(self.args, set_pod_parser(), taboo={'host'}) _args = ArgNamespace.kwargs2list(non_defaults) return f'jina pod {" ".join(_args)}'
def test_pod_remote_pea_parallel_pea_host_set_partially( pod_host, pea1_host, expected_host_in, expected_host_out, ): args = set_pod_parser().parse_args([ '--peas-hosts', f'1: {pea1_host}', '--parallel', str(2), '--host', pod_host ]) assert args.host == pod_host pod = Pod(args) for k, v in pod.peas_args.items(): if k in ['head', 'tail']: assert v.host == args.host else: for pea_arg in v: if pea_arg.pea_id == 1: assert pea_arg.host == pea1_host assert pea_arg.host_in == expected_host_in assert pea_arg.host_out == expected_host_out else: assert pea_arg.host == args.host assert pea_arg.host_in == '0.0.0.0' assert pea_arg.host_out == '0.0.0.0'
def test_pod_remote_pea_without_parallel(): args = set_pod_parser().parse_args( ['--peas-hosts', '0.0.0.1', '--parallel', str(1)] ) with Pod(args) as pod: peas = pod.peas for pea in peas: assert pea.args.host == pod.host
def test_non_gateway_runtimes(runtime_cls): args = set_pod_parser().parse_args([ '--runtime-cls', runtime_cls, ]) with Pod(args) as p: assert p.runtime_cls.__name__ == runtime_cls
def _create_worker_runtime(port, name='', executor=None): args = set_pod_parser().parse_args([]) args.port = port args.name = name if executor: args.uses = executor with WorkerRuntime(args) as runtime: runtime.run_forever()