def test_pea_store():
args = set_pea_parser().parse_args([])
store = InMemoryPeaStore()
with store._session():
pea_id = store._create(pea_arguments=args)
assert pea_id in store._store.keys()
# assert isinstance(store._store[pea_id]['pea'], LocalRuntime)
store._delete(pea_id)
assert pea_id not in store._store.keys()
def test_simple_container_with_ext_yaml(self):
args = set_pea_parser().parse_args([
'--uses', img_name, '--uses-internal',
os.path.join(cur_dir, 'mwu-encoder/mwu_encoder_ext.yml')
])
print(args)
with ContainerPea(args):
time.sleep(2)
def test_simple_container_with_ext_yaml(docker_image_built):
args = set_pea_parser().parse_args([
'--uses', img_name, '--uses-internal',
str(cur_dir.parent.parent.parent / 'mwu-encoder' /
'mwu_encoder_ext.yml')
])
with ContainerRuntime(args):
time.sleep(2)
def test_simple_container(self):
args = set_pea_parser().parse_args(['--uses', img_name])
print(args)
with ContainerPea(args):
pass
time.sleep(2)
ContainerPea(args).start().close()
def test_container_status():
args = set_pea_parser().parse_args([
'--uses', img_name, '--uses-internal',
os.path.join(cur_dir, '../mwu-encoder/mwu_encoder_ext.yml')
])
pea = ContainerPea(args)
assert not pea.is_ready
with pea:
time.sleep(1.)
assert pea.is_ready
def test_container_ping(docker_image_built):
a4 = set_pea_parser().parse_args(['--uses', 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 BasePea(a4):
NetworkChecker(a5)
assert cm.value.code == 0
def test_container_status():
args = set_pea_parser().parse_args([
'--uses', img_name, '--uses-internal',
os.path.join(cur_dir, '../../../mwu-encoder/mwu_encoder_ext.yml')
])
runtime = ContainerRuntime(args)
assert not runtime.is_ready
with runtime:
time.sleep(2.)
assert runtime.is_ready
def test_container_status():
args = set_pea_parser().parse_args([
'--uses', img_name, '--uses-internal',
str(cur_dir.parent / 'mwu-encoder' / 'mwu_encoder_ext.yml')
])
pea = ContainerPea(args)
assert not pea.is_ready
with pea:
time.sleep(2.)
assert pea.is_ready
def test_container_ping(self):
a4 = set_pea_parser().parse_args(['--uses', img_name])
a5 = set_ping_parser().parse_args(
['0.0.0.0', str(a4.port_ctrl), '--print-response'])
# test with container
with self.assertRaises(SystemExit) as cm:
with BasePea(a4):
NetworkChecker(a5)
assert cm.exception.code == 0
def test_simple_zmqlet():
args = set_pea_parser().parse_args([
'--host-in', '0.0.0.0', '--host-out', '0.0.0.0', '--port-in', '12346',
'--port-out', '12347', '--socket-in', 'PULL_CONNECT', '--socket-out',
'PUSH_CONNECT', '--timeout-ctrl', '-1'
])
args2 = set_pea_parser().parse_args([
'--host-in', '0.0.0.0', '--host-out', '0.0.0.0', '--port-in', '12347',
'--port-out', '12346', '--socket-in', 'PULL_BIND', '--socket-out',
'PUSH_BIND', '--uses', '_logforward', '--timeout-ctrl', '-1'
])
logger = logging.getLogger('zmq-test')
with BasePea(args2) as z1, Zmqlet(args, logger) as z:
req = jina_pb2.RequestProto()
req.request_id = get_random_identity()
d = req.index.docs.add()
d.tags['id'] = 2
msg = Message(None, req, 'tmp', '')
z.send_message(msg)
def test_remote_pod2(self):
f_args = set_gateway_parser().parse_args(['--allow-spawn'])
p_args = set_pea_parser().parse_args(['--host', 'localhost', '--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()
with RemotePod(p_args):
pass
t.join()
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 BasePea(a1):
NetworkChecker(a2)
assert cm.value.code == 0
# test with bad addresss
with pytest.raises(SystemExit) as cm:
with BasePea(a1):
NetworkChecker(a3)
assert cm.value.code == 1
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_executor_logger(metas):
from fluent import asynchandler as fluentasynchandler
args = set_pea_parser().parse_args([])
with BaseExecutor(args, metas=metas) as executor:
assert len(executor.logger.logger.handlers) == 3
has_fluent = False
for h in executor.logger.logger.handlers:
if isinstance(h, fluentasynchandler.FluentHandler):
has_fluent = True
assert has_fluent
executor.logger.info('logging from executor')
executor.touch()
executor.save()
save_abspath = executor.save_abspath
with BaseExecutor.load(save_abspath) as executor:
assert len(executor.logger.logger.handlers) == 3
has_fluent = False
for h in executor.logger.logger.handlers:
if isinstance(h, fluentasynchandler.FluentHandler):
has_fluent = True
assert has_fluent
executor.logger.info('logging from executor')
def test_pea_context(runtime):
args = set_pea_parser().parse_args(['--runtime', runtime])
with BasePea(args):
pass
BasePea(args).start().close()
def test_get_ctrl_addr(host):
p = set_pea_parser().parse_args(['--host', '[email protected]', '--port-ctrl', '56789'])
assert Zmqlet.get_ctrl_address(p)[0] == 'tcp://192.0.0.1:56789'
class MockBasePeaNotRead(BasePea):
def post_hook(self, msg: 'Message') -> 'BasePea':
super().post_hook(msg)
assert not msg.request.is_used
class MockBasePeaRead(BasePea):
def post_hook(self, msg: 'Message') -> 'BasePea':
super().post_hook(msg)
assert msg.request.is_used
args1 = set_pea_parser().parse_args([
'--host-in', '0.0.0.0', '--host-out', '0.0.0.0', '--port-in', '12346',
'--port-out', '12347', '--socket-in', 'PULL_CONNECT', '--socket-out',
'PUSH_CONNECT', '--timeout-ctrl', '-1'
])
args2 = set_pea_parser().parse_args([
'--host-in',
'0.0.0.0',
'--host-out',
'0.0.0.0',
'--port-in',
'12347',
'--port-out',
'12346',
'--socket-in',
'PULL_BIND',
'--socket-out',
def _docker_run(self):
# important to notice, that client is not assigned as instance member to avoid potential
# heavy copy into new process memory space
import docker
client = docker.from_env()
# the image arg should be ignored otherwise it keeps using ContainerPea in the container
# basically all args in BasePea-docker arg group should be ignored.
# this prevent setting containerPea twice
from jina.parser import set_pea_parser
non_defaults = ArgNamespace.get_non_defaults_args(
self.args,
set_pea_parser(),
taboo={'uses', 'entrypoint', 'volumes', 'pull_latest'})
if self.args.pull_latest:
self.logger.warning(
f'pulling {self.args.uses}, this could take a while. if you encounter '
f'timeout error due to pulling takes to long, then please set '
f'"timeout-ready" to a larger value.')
client.images.pull(self.args.uses)
_volumes = {}
if self.args.uses_internal:
if os.path.exists(self.args.uses_internal):
# external YAML config, need to be volumed into the container
# uses takes value from uses_internal
non_defaults['uses'] = '/' + os.path.basename(
self.args.uses_internal)
_volumes[os.path.abspath(self.args.uses_internal)] = {
'bind': non_defaults['uses'],
'mode': 'ro'
}
elif not is_valid_local_config_source(self.args.uses_internal):
raise FileNotFoundError(
f'"uses_internal" {self.args.uses_internal} is not like a path, please check it'
)
if self.args.volumes:
for p in self.args.volumes:
Path(os.path.abspath(p)).mkdir(parents=True, exist_ok=True)
_p = '/' + os.path.basename(p)
_volumes[os.path.abspath(p)] = {'bind': _p, 'mode': 'rw'}
_expose_port = [self.args.port_ctrl]
if self.args.socket_in.is_bind:
_expose_port.append(self.args.port_in)
if self.args.socket_out.is_bind:
_expose_port.append(self.args.port_out)
from sys import platform
if platform in ('linux', 'linux2'):
net_mode = 'host'
else:
net_mode = None
_args = ArgNamespace.kwargs2list(non_defaults)
ports = {f'{v}/tcp': v for v in _expose_port} if not net_mode else None
self._container = client.containers.run(
self.args.uses,
_args,
detach=True,
auto_remove=True,
ports=ports,
name=self.name,
volumes=_volumes,
network_mode=net_mode,
entrypoint=self.args.entrypoint)
# wait until the container is ready
self.logger.info('waiting ready signal from the container')
client.close()
def create(self, exception_class):
args = set_pea_parser().parse_args([])
return MockExceptionLoadExecutorPea(args, exception_class)
def test_remote_not_allowed(self):
f_args = set_gateway_parser().parse_args([])
p_args = set_pea_parser().parse_args(['--host', 'localhost', '--port-expose', str(f_args.port_expose)])
with GatewayPod(f_args):
PeaSpawnHelper(p_args).start()
def start_client(d):
print('im running %d' % d)
p_args = set_pea_parser().parse_args(
['--host', 'localhost', '--name', 'testpea%d' % d, '--port-expose', str(f_args.port_expose)])
PeaSpawnHelper(p_args).start()
def test_pea_context_load_executor():
args = set_pea_parser().parse_args([])
pea = BasePea(args)
assert not hasattr(pea, 'executor')
with pea:
assert pea.executor
def create(self, exception_class):
args = set_pea_parser().parse_args([])
return MockExceptionCallbackPea(args, exception_class)
def pea_command(self) -> str:
from jina.parser import set_pea_parser
non_defaults = ArgNamespace.get_non_defaults_args(self.args, set_pea_parser(), taboo={'host'})
_args = ArgNamespace.kwargs2list(non_defaults)
return f'jina pea {" ".join(_args)}'
def test_local_runtime_context(runtime):
args = set_pea_parser().parse_args(['--runtime', runtime])
with LocalRuntime(args):
pass
LocalRuntime(args).start().close()
def create(self, exception_class):
args = set_pea_parser().parse_args([])
return MockExceptionRequestLoopPea(args, exception_class)
