def test_load_cust_with_driver(self):
a = BaseExecutor.load_config('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_mutable_pod(runtime):
args = set_pod_parser().parse_args(['--runtime', runtime, '--replicas', '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 _update_autocomplete():
from jina.main.parser import get_main_parser, set_pea_parser, \
set_hw_parser, set_flow_parser, set_pod_parser, \
set_check_parser, set_gateway_parser, set_ping_parser, set_client_cli_parser, set_logger_parser
def _gaa(parser):
_compl = []
for v in parser._actions:
if v.option_strings:
_compl.extend(v.option_strings)
elif v.choices:
_compl.extend(v.choices)
# filer out single dash, as they serve as abbrev
_compl = [k for k in _compl if (not k.startswith('-') or k.startswith('--'))]
return _compl
compl = {
'commands': _gaa(get_main_parser()),
'completions': {
'pea': _gaa(set_pea_parser()),
'hello-world': _gaa(set_hw_parser()),
'flow': _gaa(set_flow_parser()),
'pod': _gaa(set_pod_parser()),
'check': _gaa(set_check_parser()),
'gateway': _gaa(set_gateway_parser()),
'ping': _gaa(set_ping_parser()),
'client': _gaa(set_client_cli_parser()),
'log': _gaa(set_logger_parser())
}
}
with open(__file__, 'a') as fp:
fp.write(f'\nac_table = {compl}\n')
def _test_pod_context(runtime):
args = set_pod_parser().parse_args(
['--runtime', runtime, '--parallel', '2'])
with BasePod(args):
pass
BasePod(args).start().close()
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_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_peas_naming_with_parallel(self):
args = set_pod_parser().parse_args([
'--name', 'pod', '--parallel', '2', '--max-idle-time', '5',
'--shutdown-idle'
])
with BasePod(args) as bp:
self.assertEqual(bp.peas[0].name, 'pod-head')
self.assertEqual(bp.peas[1].name, 'pod-tail')
self.assertEqual(bp.peas[2].name, 'pod-1')
self.assertEqual(bp.peas[3].name, 'pod-2')
def test_remote_pod_process(self):
f_args = set_gateway_parser().parse_args(['--allow-spawn'])
p_args = set_pod_parser().parse_args(
['--host', 'localhost', '--replicas', '3',
'--port-grpc', str(f_args.port_grpc), '--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_customized_pod(self):
f_args = set_gateway_parser().parse_args(['--allow-spawn'])
p_args = set_pod_parser().parse_args(
['--host', 'localhost', '--replicas', '3', '--port-grpc', str(f_args.port_grpc)])
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(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)
])
def start_gateway():
with GatewayPod(f_args):
time.sleep(5)
t = Process(target=start_gateway)
t.daemon = True
t.start()
PodSpawnHelper(p_args).start()
t.join()
def test_customized_pod2(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()
with RemoteMutablePod(p.peas_args):
pass
t.join()
def get_podargs():
"""Get the default args of a pod"""
from jina.main.parser import set_pod_parser
from argparse import _StoreAction, _StoreTrueAction
port_attr = ('help', 'choices', 'default')
d = {}
parser = set_pod_parser()
for a in parser._actions:
if isinstance(a, _StoreAction) or isinstance(a, _StoreTrueAction):
d[a.dest] = {p: getattr(a, p) for p in port_attr}
if a.type:
d[a.dest]['type'] = a.type.__name__
elif isinstance(a, _StoreTrueAction):
d[a.dest]['type'] = 'bool'
else:
d[a.dest]['type'] = a.type
d = {'pod': d, 'version': __version__, 'usage': parser.format_help()}
return jsonify(d)
def test_pod_status(self):
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:
self.assertIsNotNone(v)
def test_pod_status(self):
args = set_pod_parser().parse_args(['--replicas', '3'])
with BasePod(args) as p:
self.assertEqual(len(p.status), p.num_peas)
for v in p.status:
self.assertIsNotNone(v)
def test_pod_with_sse_no_parallelism_no_deadlock():
args = set_pod_parser().parse_args(['--parallel', '1', '--log-sse'])
p = Pod(args)
with p:
pass
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_with_sse_no_deadlock_thread():
args = set_pod_parser().parse_args(
['--parallel', '2', '--runtime', 'thread', '--log-sse'])
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