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
