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()
