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_pod_factory_pod():
args_no_replicas = set_pod_parser().parse_args(['--shards', '1'])
assert isinstance(PodFactory.build_pod(args_no_replicas), Pod)
args_replicas = set_pod_parser().parse_args(['--shards', '2'])
assert isinstance(PodFactory.build_pod(args_replicas), CompoundPod)
args_no_replicas = set_pod_parser().parse_args(['--replicas', '2'])
assert isinstance(PodFactory.build_pod(args_no_replicas), Pod)
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_shards_parallel_synonym():
a = set_pod_parser().parse_args(['--shards', '2'])
assert a.parallel == 2
with pytest.raises(AttributeError):
a.shards
a = set_pod_parser().parse_args(['--parallel', '2'])
assert a.parallel == 2
with pytest.raises(AttributeError):
a.shards
a = set_pod_parser().parse_args([])
assert a.parallel == 1
with pytest.raises(AttributeError):
a.shards
def external_pod_parallel_2_args(
port_in_external, port_out_external, num_replicas, num_parallel
):
args = [
'--uses',
'MyExternalExecutor',
'--name',
'external_real_2',
'--socket-in',
'SUB_CONNECT',
'--port-in',
str(port_in_external),
'--host-in',
'0.0.0.0',
'--port-out',
str(port_out_external),
'--host-out',
'0.0.0.0',
'--socket-out',
'PUSH_CONNECT',
'--parallel',
str(num_parallel),
'--replicas',
str(num_replicas),
]
return set_pod_parser().parse_args(args)
def test_parse_args_custom_executor(shards: int):
uses_before = 'custom-executor-before'
uses_after = 'custom-executor-after'
args = set_pod_parser().parse_args(
[
'--shards',
str(shards),
'--uses-before',
uses_before,
'--uses-after',
uses_after,
]
)
pod = K8sPod(args)
assert namespace_equal(
args, pod.deployment_args['head_deployment'], skip_attr=('uses',)
)
assert pod.deployment_args['head_deployment'].uses == uses_before
assert namespace_equal(
args, pod.deployment_args['tail_deployment'], skip_attr=('uses',)
)
assert pod.deployment_args['tail_deployment'].uses == uses_after
for i, depl_arg in enumerate(pod.deployment_args['deployments']):
import copy
cargs = copy.deepcopy(args)
cargs.shard_id = i
assert depl_arg == cargs
def test_pod_with_gpus(mocker):
args_list = [
'--name',
'test-wait-success',
'--k8s-namespace',
'test-namespace',
'--shards',
'1',
'--replicas',
'1',
'--gpus',
'3',
]
args = set_pod_parser().parse_args(args_list)
container = client.V1Container(
name='test-container',
resources=client.V1ResourceRequirements(limits={'nvidia.com/gpu': 3}),
)
spec = client.V1PodSpec(containers=[container])
mocker.patch(
'jina.peapods.pods.k8s.K8sPod._K8sDeployment._read_namespaced_deployment',
return_value=client.V1Deployment(
status=client.V1DeploymentStatus(replicas=1, ready_replicas=1), spec=spec
),
)
mocker.patch(
'jina.peapods.pods.k8slib.kubernetes_deployment.deploy_service',
return_value=f'test-wait-success.test-namespace.svc',
)
mocker.patch(
'jina.peapods.pods.k8s.K8sPod._K8sDeployment._delete_namespaced_deployment',
return_value=client.V1Status(status=200),
)
with K8sPod(args) as pod:
assert pod.args.gpus == '3'
def test_pod_start_close_given_tail_deployment(
name, k8s_namespace, shards, mocker, uses_after
):
args_list = [
'--name',
name,
'--k8s-namespace',
k8s_namespace,
'--shards',
shards,
'--noblock-on-start',
]
if uses_after:
args_list.extend(['--uses-after', 'custom-executor-after'])
args = set_pod_parser().parse_args(args_list)
mocker.patch(
'jina.peapods.pods.k8slib.kubernetes_deployment.deploy_service',
return_value=f'{name}.{k8s_namespace}.svc',
)
mocker.patch(
'jina.peapods.pods.k8s.K8sPod._K8sDeployment._delete_namespaced_deployment',
return_value=client.V1Status(status=200),
)
with K8sPod(args) as pod:
# enter `_deploy_runtime`
assert isinstance(pod.k8s_tail_deployment, K8sPod._K8sDeployment)
assert pod.k8s_tail_deployment.name == name + '-tail'
assert pod.args.noblock_on_start
def test_scale_remote_pod(pod_params):
num_replicas, scale_to, shards = pod_params
pod_id = None
args = set_pod_parser().parse_args(
['--replicas', str(num_replicas), '--shards', str(shards)]
)
payload = replace_enum_to_str(ArgNamespace.flatten_to_dict(args))
try:
workspace_id = jinad_client.workspaces.create(
paths=[os.path.join(cur_dir, cur_dir)]
)
success, pod_id = jinad_client.pods.create(
workspace_id=workspace_id, payload=payload
)
assert success
remote_pod_args = jinad_client.pods.get(pod_id)['arguments']['object'][
'arguments'
]
assert remote_pod_args['identity'] == pod_id
assert remote_pod_args['replicas'] == num_replicas
assert remote_pod_args['shards'] == shards
jinad_client.pods.scale(id=pod_id, replicas=scale_to)
remote_pod_args = jinad_client.pods.get(pod_id)['arguments']['object'][
'arguments'
]
assert remote_pod_args['identity'] == pod_id
assert remote_pod_args['replicas'] == scale_to
assert remote_pod_args['shards'] == shards
finally:
if pod_id:
assert jinad_client.pods.delete(pod_id), 'Pod termination failed'
print(f'Remote Pod {pod_id} successfully terminated')
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_factory_k8s():
args_replicas = set_pod_parser().parse_args([])
assert isinstance(
PodFactory.build_pod(args_replicas,
infrastructure=InfrastructureType.K8S),
K8sPod,
)
async def test_decorator_monitoring(port_generator):
from jina import monitor
class DummyExecutor(Executor):
@requests
def foo(self, docs, **kwargs):
self._proces(docs)
self.process_2(docs)
@monitor(name='metrics_name', documentation='metrics description')
def _proces(self, docs):
...
@monitor()
def process_2(self, docs):
...
port = port_generator()
args = set_pod_parser().parse_args([
'--monitoring', '--port-monitoring',
str(port), '--uses', 'DummyExecutor'
])
cancel_event = multiprocessing.Event()
def start_runtime(args, cancel_event):
with WorkerRuntime(args, cancel_event=cancel_event) as runtime:
runtime.run_forever()
runtime_thread = Process(
target=start_runtime,
args=(args, cancel_event),
daemon=True,
)
runtime_thread.start()
assert AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
timeout=5.0,
ctrl_address=f'{args.host}:{args.port}',
ready_or_shutdown_event=Event(),
)
assert AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
timeout=5.0,
ctrl_address=f'{args.host}:{args.port}',
ready_or_shutdown_event=Event(),
)
await GrpcConnectionPool.send_request_async(_create_test_data_message(),
f'{args.host}:{args.port}',
timeout=1.0)
resp = req.get(f'http://localhost:{port}/')
assert f'jina_metrics_name_count{{runtime_name="None"}} 1.0' in str(
resp.content)
cancel_event.set()
runtime_thread.join()
assert not AsyncNewLoopRuntime.is_ready(f'{args.host}:{args.port}')
async def test_worker_runtime_reflection():
args = set_pod_parser().parse_args([])
cancel_event = multiprocessing.Event()
def start_runtime(args, cancel_event):
with WorkerRuntime(args, cancel_event=cancel_event) as runtime:
runtime.run_forever()
runtime_thread = Process(
target=start_runtime,
args=(args, cancel_event),
daemon=True,
)
runtime_thread.start()
assert AsyncNewLoopRuntime.wait_for_ready_or_shutdown(
timeout=3.0,
ctrl_address=f'{args.host}:{args.port}',
ready_or_shutdown_event=Event(),
)
async with grpc.aio.insecure_channel(
f'{args.host}:{args.port}') as channel:
service_names = await GrpcConnectionPool.get_available_services(channel
)
assert all(service_name in service_names for service_name in [
'jina.JinaDataRequestRPC',
'jina.JinaSingleDataRequestRPC',
])
cancel_event.set()
runtime_thread.join()
assert not AsyncNewLoopRuntime.is_ready(f'{args.host}:{args.port}')
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_compound_pod_do_not_forward_uses_before_uses_after():
polling = 'all'
shards = 2
replicas = 2
args = set_pod_parser().parse_args(
[
'--name',
'pod',
'--shards',
f'{shards}',
'--polling',
f'{polling}',
'--replicas',
f'{replicas}',
'--uses-after',
'BaseExecutor',
'--uses-before',
'BaseExecutor',
]
)
cpod = CompoundPod(args)
for pod in cpod.shards:
assert pod.args.uses_before is None
assert pod.args.uses_after is None
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_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_ssh_pod():
p = set_pod_parser().parse_args(
['--host', '[email protected]', '--timeout', '5000'])
with SSHRuntime(p, kind='pod') as pp:
assert pp.status.envelope.status.code == jina_pb2.StatusProto.READY
assert pp.status is None
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_jinad_pod():
args = set_pod_parser().parse_args([])
assert not is_pod_ready(args)
with JinaDPod(args):
assert is_pod_ready(args)
assert not is_pod_ready(args)
def external_pod_join_args(
port_in_external, port_out_external, num_replicas, num_parallel
):
args = [
'--uses',
'MyExternalExecutor',
'--name',
'external_real',
'--port-in',
str(port_in_external),
'--host-in',
'0.0.0.0',
'--socket-in',
'PULL_BIND',
'--port-out',
str(port_out_external),
'--host-out',
'0.0.0.0',
'--socket-out',
'PUSH_BIND',
'--pod-role',
'JOIN',
'--parallel',
str(num_parallel),
'--replicas',
str(num_replicas),
]
return set_pod_parser().parse_args(args)
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 pod_to_namespace(args: Union[SinglePodModel, ParallelPodModel]):
from jina.parsers import set_pod_parser
parser = set_pod_parser()
if isinstance(args, ParallelPodModel):
pod_args = {}
args = args.dict()
for pea_type, pea_args in args.items():
# this is for pea_type: head & tail when None (pod with parallel = 1)
if pea_args is None:
pod_args[pea_type] = None
# this is for pea_type: head & tail when not None (pod with parallel > 1)
if isinstance(pea_args, Dict):
pea_args = handle_remote_args(args=pea_args, parser=parser)
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 current_pea_arg in pea_args:
current_pea_arg = handle_remote_args(args=current_pea_arg,
parser=parser)
pod_args[pea_type].append(
argparse.Namespace(**current_pea_arg))
return pod_args
if isinstance(args, SinglePodModel):
pod_args = handle_remote_args(args=args.dict(), parser=parser)
return argparse.Namespace(**pod_args)
def external_pod_args():
args = ['--port-in', str(45678), '--port-out', str(45679)]
args = vars(set_pod_parser().parse_args(args))
del args['external']
del args['pod_role']
del args['host']
return args
def test_executor_runtimes(signal, tmpdir):
import time
args = set_pod_parser().parse_args([])
def run(args):
args.uses = {
'jtype': 'DummyExecutor',
'with': {
'dir': str(tmpdir)
},
'metas': {
'workspace': str(tmpdir)
},
}
executor_native(args)
process = multiprocessing.Process(target=run, args=(args, ))
process.start()
time.sleep(0.5)
GrpcConnectionPool.send_request_sync(_create_test_data_message(),
target=f'{args.host}:{args.port}')
time.sleep(0.1)
os.kill(process.pid, signal)
process.join()
with open(f'{tmpdir}/test.txt', 'r') as fp:
output = fp.read()
split = output.split(';')
assert split[0] == 'proper close'
assert split[1] == '1'
def _pod_command(self) -> str:
from jina.parsers 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_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_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 test_non_gateway_runtimes(runtime_cls):
args = set_pod_parser().parse_args([
'--runtime-cls',
runtime_cls,
])
with Pod(args) as p:
assert p.runtime_cls.__name__ == runtime_cls
def _create_worker_runtime(port, name='', executor=None):
args = set_pod_parser().parse_args([])
args.port = port
args.name = name
if executor:
args.uses = executor
with WorkerRuntime(args) as runtime:
runtime.run_forever()
