def test_gateway_pod(runtime, protocol, runtime_cls): args = set_gateway_parser().parse_args( ['--runtime-backend', runtime, '--protocol', protocol]) with Pod(args) as p: assert len(p.all_args) == 1 assert p.all_args[0].runtime_cls == runtime_cls Pod(args).start().close()
def test_pod_remote_context(runtime): args = set_pod_parser().parse_args(['--runtime-backend', runtime, '--parallel', str(2)]) with Pod(args) as p: remote_pod_args = p.peas_args with Pod(remote_pod_args) as remote_pod: assert remote_pod.num_peas == 4 # head + tail + 2 peas Pod(remote_pod_args).start().close()
def test_tail_args_get_set(pod_args, pod_args_singleton): with Pod(pod_args) as pod: assert pod.tail_args == pod.peas_args['tail'] pod.tail_args = pod_args_singleton assert pod.peas_args['tail'] == pod_args_singleton with Pod(pod_args_singleton) as pod: assert pod.tail_args == pod.first_pea_args pod.tail_args = pod_args assert pod.peas_args['peas'][0] == pod_args
def test_pod_context_parallel(runtime, parallel): args = set_pod_parser().parse_args(['--runtime-backend', runtime, '--parallel', str(parallel)]) with Pod(args) as bp: if parallel == 1: assert bp.num_peas == 1 else: # count head and tail assert bp.num_peas == parallel + 2 Pod(args).start().close()
def test_gateway_pod(runtime, restful, runtime_cls): args = set_gateway_parser().parse_args( ['--runtime-backend', runtime, '--runtime-cls', runtime_cls] + (['--restful'] if restful else [])) with Pod(args) as p: assert len(p.all_args) == 1 if restful: assert p.all_args[0].runtime_cls == 'RESTRuntime' else: assert p.all_args[0].runtime_cls == 'GRPCRuntime' Pod(args).start().close()
def test_pod_context_replicas(runtime, replicas, grpc_data_requests): args_list = ['--runtime-backend', runtime, '--replicas', str(replicas)] if grpc_data_requests: args_list.append('--grpc-data-requests') args = set_pod_parser().parse_args(args_list) with Pod(args) as bp: if replicas == 1: assert bp.num_peas == 1 else: # count head and tail assert bp.num_peas == replicas + 2 Pod(args).start().close()
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_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_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_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_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_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_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 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_new_api_from_kwargs(self): a = BaseExecutor.load_config('mwu-encoder/mwu_encoder_driver.yml') assert a._drivers['ControlRequest'][0].__class__.__name__ == 'MyAwesomeDriver' with Pod(uses=os.path.join(cur_dir, 'mwu-encoder/mwu_encoder_driver.yml')): # will print a cust task_name from the driver when terminate pass
def test_pod_naming_with_replica(runtime): args = set_pod_parser().parse_args( ['--name', 'pod', '--replicas', '2', '--runtime-backend', runtime]) with Pod(args) as bp: assert bp.head_pea.name == 'pod/head' assert bp.replica_set._peas[0].name == 'pod/rep-0' assert bp.replica_set._peas[1].name == 'pod/rep-1' assert bp.tail_pea.name == 'pod/tail'
def test_pod_new_api_from_kwargs(self): a = BaseExecutor.load_config('mwu-encoder/mwu_encoder_driver.yml') self.assertEqual(a._drivers['ControlRequest'][0].__class__.__name__, 'MyAwesomeDriver') with Pod(yaml_path='mwu-encoder/mwu_encoder_driver.yml'): # will print a cust msg from the driver when terminate pass
def test_pod_naming_with_parallel_any(runtime): args = set_pod_parser().parse_args( ['--name', 'pod', '--parallel', '2', '--runtime-backend', runtime]) with Pod(args) as bp: assert bp.peas[0].name == 'pod/head' assert bp.peas[1].name == 'pod/pea-0' assert bp.peas[2].name == 'pod/pea-1' assert bp.peas[3].name == 'pod/tail'
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 pod(args: 'Namespace'): """Start a Pod""" from jina.peapods import Pod try: with Pod(args) as p: p.join() except KeyboardInterrupt: pass
def test_load_cust_with_driver(self): a = BaseExecutor.load_config('mwu-encoder/mwu_encoder_driver.yml') self.assertEqual(a._drivers['ControlRequest'][0].__class__.__name__, 'MyAwesomeDriver') p = set_pod_parser().parse_args( ['--yaml-path', 'mwu-encoder/mwu_encoder_driver.yml']) with Pod(p): # will print a cust msg from the driver when terminate pass
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', __default_executor__, '--uses-after', __default_executor__ ]) with Pod(args) as p: assert p.num_peas == 1 args = set_pod_parser().parse_args([ '--uses-before', __default_executor__, '--uses-after', __default_executor__, '--replicas', '2', ]) with Pod(args) as p: assert p.num_peas == 4
def _create(self, pod_arguments: Union[Dict, Namespace]): """ Creates a Pod via Flow or via CLI """ try: pod_id = uuid.UUID(pod_arguments.log_id) if isinstance(pod_arguments, Namespace) \ else uuid.UUID(pod_arguments['peas'][0].log_id) pod = Pod(pod_arguments) pod = self._start(context=pod) except Exception as e: self.logger.critical(f'Got following error while starting the pod: {e!r}') raise PodStartException(repr(e)) self._store[pod_id] = {} self._store[pod_id]['pod'] = pod self.logger.info(f'Started pod with pod_id {colored(pod_id, "cyan")}') return pod_id
def _rolling_update(self, dump_path): _print_and_append_to_ops(f'### calling patched rolling update') for i in range(len(self.replicas)): _print_and_append_to_ops(f'### replica {i} -- starting') replica = self.replicas[i] replica.close() _print_and_append_to_ops(f'### replica {i} -- went offline') time.sleep( 3) # wait for query to hit system when one replica is offline _args = self.replicas_args[i] _args.noblock_on_start = False _args.dump_path = dump_path new_replica = Pod(_args) self.enter_context(new_replica) _print_and_append_to_ops(f'### replica {i} - new instance online') self.replicas[i] = new_replica time.sleep(5)
def test_pod_args_remove_uses_ba(remove_uses_ba): if remove_uses_ba: args = set_pod_parser().parse_args([ '--remove-uses-ba', '--uses-before', '_pass', '--uses-after', '_pass' ]) else: args = set_pod_parser().parse_args( ['--uses-before', '_pass', '--uses-after', '_pass']) pod = Pod(args) if remove_uses_ba: assert pod.args.uses_before is None assert pod.args.uses_after is None else: assert pod.args.uses_before == '_pass' assert pod.args.uses_after == '_pass'
def test_pod_naming_with_parallel_all(runtime): args = set_pod_parser().parse_args([ '--name', 'pod', '--parallel', '2', '--runtime-backend', runtime, '--polling', 'ALL', ]) with Pod(args) as bp: assert bp.peas[0].name == 'pod/head' assert bp.peas[1].name == 'pod/0' assert bp.peas[2].name == 'pod/1' assert bp.peas[3].name == 'pod/tail' assert bp.peas[0].runtime.name == 'pod/head/ZEDRuntime' assert bp.peas[1].runtime.name == 'pod/0/ZEDRuntime' assert bp.peas[2].runtime.name == 'pod/1/ZEDRuntime' assert bp.peas[3].runtime.name == 'pod/tail/ZEDRuntime'
def test_pod_env_setting(): class EnvChecker(BaseExecutor): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) # pea/pod-specific assert os.environ['key1'] == 'value1' assert os.environ['key2'] == 'value2' # inherit from parent process assert os.environ['key_parent'] == 'value3' os.environ['key_parent'] = 'value3' with Pod(uses='EnvChecker', env=['key1=value1', 'key2=value2']): pass # should not affect the main process assert 'key1' not in os.environ assert 'key2' not in os.environ assert 'key_parent' in os.environ os.unsetenv('key_parent')
def test_use_from_local_hub_pod_level(test_envs, mocker, monkeypatch, local_hub_executor): from jina.hubble.hubio import HubIO, HubExecutor mock = mocker.Mock() def _mock_fetch(name, tag=None, secret=None, force=False): mock(name=name) return HubExecutor( uuid='hello', name='alias_dummy', tag='v0', image_name='jinahub/pod.dummy_mwu_encoder', md5sum=None, visibility=True, archive_url=None, ) monkeypatch.setattr(HubIO, 'fetch_meta', _mock_fetch) a = set_pod_parser().parse_args(['--uses', 'jinahub://hello']) with Pod(a): pass
def test_use_from_local_dir_pod_level(): a = set_pod_parser().parse_args(['--uses', 'dummyhub/config.yml']) with Pod(a): pass
def test_equal(pod_args, pod_args_singleton): pod1 = Pod(pod_args) pod2 = Pod(pod_args) assert pod1 == pod2 pod1.close() pod2.close() # test not equal pod1 = Pod(pod_args) pod2 = Pod(pod_args_singleton) assert pod1 != pod2 pod1.close() pod2.close()