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_mutable_pod(runtime): args = set_pod_parser().parse_args(['--runtime', runtime, '--parallel', '2']) with MutablePod(BasePod(args).peas_args): pass MutablePod(BasePod(args).peas_args).start().close()
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 flowpod_to_namespace(args: Dict): # TODO: combine all 3 to_namespace methods from jina.parser import set_pod_parser parser = set_pod_parser() pod_args = {} for pea_type, pea_args in args.items(): # this is for pea_type: head & tail when None if pea_args is None: pod_args[pea_type] = None # this is for pea_type: head & tail when not None if isinstance(pea_args, dict): pea_args = handle_enums(args=pea_args, parser=parser) handle_log_id(args=pea_args) 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 pea_arg in pea_args: pea_arg = handle_enums(args=pea_arg, parser=parser) handle_log_id(args=pea_arg) pod_args[pea_type].append(argparse.Namespace(**pea_arg)) return pod_args
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_ssh_pod(): p = set_pod_parser().parse_args( ['--host', '[email protected]', '--timeout', '5000']) with RemoteSSHPod(p) as pp: assert pp.status.envelope.status.code == jina_pb2.StatusProto.READY assert pp.status is None
def test_pod_context(runtime): args = set_pod_parser().parse_args( ['--runtime', runtime, '--parallel', '2']) with BasePod(args): pass BasePod(args).start().close()
def basepod_to_namespace(args: PodModel): from jina.parser import set_pod_parser parser = set_pod_parser() if isinstance(args, PodModel): pod_args = handle_enums(args=args.dict(), parser=parser) handle_log_id(args=pod_args) return argparse.Namespace(**pod_args)
def test_pod_store(): args = set_pod_parser().parse_args([]) store = InMemoryPodStore() with store._session(): pod_id = store._create(pod_arguments=args) assert pod_id in store._store.keys() assert isinstance(store._store[pod_id]['pod'], BasePod) store._delete(pod_id) assert pod_id not in store._store.keys()
def test_pod_context(runtime): args = set_pod_parser().parse_args(['--runtime', runtime, '--parallel', '2', '--max-idle-time', '5', '--shutdown-idle']) with BasePod(args) as bp: bp.join() BasePod(args).start().close()
def test_peas_naming_with_parallel(): args = set_pod_parser().parse_args([ '--name', 'pod', '--parallel', '2', '--max-idle-time', '5', '--shutdown-idle' ]) with BasePod(args) as bp: assert bp.peas[0].name == 'pod-head' assert bp.peas[1].name == 'pod-tail' assert bp.peas[2].name == 'pod-1' assert bp.peas[3].name == 'pod-2'
def test_customized_pod(self): f_args = set_gateway_parser().parse_args(['--allow-spawn']) p_args = set_pod_parser().parse_args( ['--host', 'localhost', '--parallel', '3', '--port-expose', str(f_args.port_expose)]) p = BasePod(p_args) def start_gateway(): with GatewayPod(f_args): time.sleep(5) t = Process(target=start_gateway) t.daemon = True t.start() MutablePodSpawnHelper(p.peas_args).start()
def test_remote_pod_process(self): f_args = set_gateway_parser().parse_args(['--allow-spawn']) p_args = set_pod_parser().parse_args( ['--host', 'localhost', '--parallel', '3', '--port-expose', str(f_args.port_expose), '--runtime', 'process']) def start_spawn(): PodSpawnHelper(p_args).start() with GatewayPod(f_args): t = Process(target=start_spawn) t.daemon = True t.start() time.sleep(5)
def test_pod_gracefully_close_idle(): import time args = set_pod_parser().parse_args(['--name', 'pod', '--parallel', '2', '--max-idle-time', '4', '--shutdown-idle']) start_time = time.time() with BasePod(args) as bp: while not bp.is_shutdown: time.sleep(1.5) end_time = time.time() elapsed_time = end_time - start_time assert elapsed_time > 4
def test_logging_pod(monkeypatch): from fluent import asynchandler as fluentasynchandler def mock_emit(obj, record): msg = obj.format(record) ct = msg['context'] if ct not in ['JINA', 'PROFILE']: assert msg['log_id'] == 'logging_id' if msg['name'] == 'gateway': assert 'log_id' in msg monkeypatch.setattr(fluentasynchandler.FluentHandler, "emit", mock_emit) args = set_pod_parser().parse_args(['--log-id', 'logging_id']) with BasePod(args): pass
def build_pydantic_model(kind: str = 'local', model_name: str = 'CustomModel', module: str = 'pod'): if kind == 'api': all_cli_args = get_latest_api() module_args = get_module_args(all_args=all_cli_args, module=module) all_fields, field_validators = get_pydantic_fields(config=module_args) elif kind == 'local': from jina.parser import set_pea_parser, set_pod_parser, set_flow_parser if module == 'pod': parser = set_pod_parser() elif module == 'pea': parser = set_pea_parser() elif module == 'flow': parser = set_flow_parser() all_fields, field_validators = get_pydantic_fields(config=parser) return create_model(model_name, **all_fields, __config__=PydanticConfig, __validators__=field_validators)
def test_pod_with_sse_no_deadlock_thread(): args = set_pod_parser().parse_args( ['--parallel', '2', '--runtime', 'thread']) p = Pod(args) with p: pass
def pod_command(self) -> str: from jina.parser 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_with_sse_no_deadlock_log_remote(): args = set_pod_parser().parse_args(['--parallel', '2', '--log-remote']) p = Pod(args) with p: pass
def test_pod_status(): args = set_pod_parser().parse_args(['--parallel', '3']) with BasePod(args) as p: assert len(p.status) == p.num_peas for v in p.status: assert v
def test_pod_without_sse_no_parallelism_no_deadlock(): args = set_pod_parser().parse_args(['--parallel', '1']) p = Pod(args) with p: 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_pod_with_sse_no_deadlock_parallelism(): args = set_pod_parser().parse_args(['--parallel', '2']) p = Pod(args) with p: pass