def test_address_in_use(runtime):
p = ['--port-ctrl', '55555', '--runtime-backend', runtime]
args1 = set_pea_parser().parse_args(p)
args2 = set_pea_parser().parse_args(p)
with pytest.raises(RuntimeFailToStart):
with Pea(args1), Pea(args2):
pass
def test_simple_container(docker_image_built):
args = set_pea_parser().parse_args(['--uses', f'docker://{img_name}'])
with Pea(args):
pass
time.sleep(2)
Pea(args).start().close()
def test_address_in_use(runtime):
# test does not work with same control port, because right now the `READYNESS` signal is done by sending to that port
p = ['--port-in', '55555', '--runtime-backend', runtime]
args1 = set_pea_parser().parse_args(p)
args2 = set_pea_parser().parse_args(p)
with pytest.raises(RuntimeFailToStart):
args1.name = 'Pea-1'
args2.name = 'Pea-2'
with Pea(args1), Pea(args2):
pass
def test_pea_runtime_env_setting_in_thread(fake_env):
class EnvChecker(BaseExecutor):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
# pea/pod-specific
assert 'key1' not in os.environ
assert 'key2' not in os.environ
# inherit from parent process
assert os.environ['key_parent'] == 'value3'
os.environ['key_parent'] = 'value3'
with Pea(set_pea_parser().parse_args([
'--uses',
'EnvChecker',
'--env',
'key1=value1',
'--env',
'key2=value2',
'--runtime-backend',
'thread',
])):
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_pea_instantiate_start_same_context():
arg = set_pea_parser().parse_args([])
peas_args = [arg, arg]
for args in peas_args:
pea = Pea(args)
with pea:
pass
def test_simple_container_with_ext_yaml(docker_image_built):
args = set_pea_parser().parse_args([
'--uses', f'docker://{img_name}', '--uses-internal',
os.path.join(cur_dir, '../../../mwu-encoder/mwu_encoder_ext.yml')
])
with Pea(args):
time.sleep(2)
def test_gateway_ready(port_expose, route, status_code):
p = set_gateway_parser().parse_args(
['--port-expose',
str(port_expose), '--protocol', 'http'])
with Pea(p):
time.sleep(0.5)
a = requests.get(f'http://localhost:{p.port_expose}{route}')
assert a.status_code == status_code
def pea(args: 'Namespace'):
"""Start a Pea"""
from jina.peapods import Pea
try:
with Pea(args) as p:
p.join()
except KeyboardInterrupt:
pass
def test_gateway_ready(port_expose, route, status_code):
p = set_gateway_parser().parse_args(
['--port-expose',
str(port_expose), '--runtime-cls', 'RESTRuntime'])
with Pea(p):
time.sleep(0.5)
a = requests.get(f'http://0.0.0.0:{p.port_expose}{route}')
assert a.status_code == status_code
def test_pea_set_shard_pea_id():
args = set_pea_parser().parse_args(['--shard-id', '1', '--shards', '3'])
pea = Pea(args)
assert pea.args.shard_id == 1
assert pea.args.pea_id == 1
assert pea.args.shards == 3
assert pea.args.parallel == 3
def test_idle_does_not_create_response(command, response_expected):
args = set_pea_parser().parse_args([])
with Pea(args) as p:
msg = ControlMessage(command, pod_name='fake_pod')
with zmq.Context().socket(zmq.PAIR) as socket:
socket.connect(f'tcp://localhost:{p.args.port_ctrl}')
socket.send_multipart(msg.dump())
assert socket.poll(timeout=1000) == response_expected
def test_pea_instantiate_start_different_context():
arg = set_pea_parser().parse_args([])
peas_args = [arg, arg]
peas = []
for args in peas_args:
peas.append(Pea(args))
for pea in peas:
with pea:
pass
def test_ping():
a1 = set_pea_parser().parse_args([])
a2 = set_ping_parser().parse_args(
['0.0.0.0', str(a1.port_ctrl), '--print-response'])
a3 = set_ping_parser().parse_args(
['0.0.0.1', str(a1.port_ctrl), '--timeout', '1000'])
with pytest.raises(SystemExit) as cm:
with Pea(a1):
NetworkChecker(a2)
assert cm.value.code == 0
# test with bad address
with pytest.raises(SystemExit) as cm:
with Pea(a1):
NetworkChecker(a3)
assert cm.value.code == 1
def test_gateway_args(protocol, expected):
args = set_gateway_parser().parse_args([
'--host',
'jina-custom-gateway',
'--port-expose',
'23456',
'--protocol',
protocol,
])
p = Pea(args)
assert p.runtime_cls.__name__ == expected
def test_container_ping(docker_image_built):
a4 = set_pea_parser().parse_args(['--uses', f'docker://{img_name}'])
a5 = set_ping_parser().parse_args(
['0.0.0.0', str(a4.port_ctrl), '--print-response'])
# test with container
with pytest.raises(SystemExit) as cm:
with Pea(a4):
NetworkChecker(a5)
assert cm.value.code == 0
def pea(args: 'Namespace'):
"""
Start a Pea
:param args: arguments coming from the CLI.
"""
from jina.peapods import Pea
try:
with Pea(args) as p:
p.join()
except KeyboardInterrupt:
pass
def _create(self, pea_arguments: Union[Dict, Namespace]):
try:
pea_id = uuid.UUID(pea_arguments.log_id) if isinstance(pea_arguments, Namespace) \
else uuid.UUID(pea_arguments['log_id'])
pea = Pea(pea_arguments)
pea = self._start(context=pea)
except Exception as e:
self.logger.critical(f'Got following error while starting the pea: {e!r}')
raise PeaStartException(repr(e))
self._store[pea_id] = {}
self._store[pea_id]['pea'] = pea
self.logger.info(f'Started pea with pea_id {colored(pea_id, "cyan")}')
return pea_id
def test_pea_runtime_env_setting_in_process(fake_env):
with Pea(set_pea_parser().parse_args([
'--uses',
'EnvChecker1',
'--env',
'key1=value1',
'--env',
'key2=value2',
'--runtime-backend',
'process',
])):
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
def external_executor(external_executor_args):
return Pea(external_executor_args)
