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_not_read_zmqlet():
with MockBasePeaNotRead(args3), Zmqlet(args1, default_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 handle_log_id(args: Dict):
args['log_id'] = args[
'identity'] if 'identity' in args else get_random_identity()
def req():
r = jina_pb2.RequestProto()
r.request_id = get_random_identity()
r.index.docs.add()
return r
def _create(self,
config: Union[str, SpooledTemporaryFile,
List[PodModel]] = None,
files: List[UploadFile] = None):
""" Creates Flow using List[PodModel] or yaml spec """
# This makes sure `uses` & `py_modules` are created locally in `cwd`
# TODO: Handle file creation, deletion better
if files:
[
create_meta_files_from_upload(current_file)
for current_file in files
]
# FastAPI treats UploadFile as a tempfile.SpooledTemporaryFile
# I think this needs to be handled by some `FlowBuilder` class, transfrom a yaml_load to a config
if isinstance(config, str) or isinstance(config, SpooledTemporaryFile):
yamlspec = config.read().decode() if isinstance(
config, SpooledTemporaryFile) else config
try:
yaml.register_class(Flow)
flow = yaml.load(yamlspec)
except Exception as e:
self.logger.error(
f'Got error while loading from yaml {repr(e)}')
raise FlowYamlParseException
elif isinstance(config, list):
try:
flow = Flow()
# it is strange to build from a given flow, it seems like a lazy construction pattern could be used?
flow = self._build_with_pods(flow=flow, pod_args=config)
except Exception as e:
self.logger.error(
f'Got error while creating flows via pods: {repr(e)}')
raise FlowCreationException
else:
raise FlowBadInputException(
f'Not valid Flow config input {type(config)}')
try:
flow.args.log_id = flow.args.identity if 'identity' in flow.args else get_random_identity(
)
flow_id = uuid.UUID(flow.args.log_id)
flow = self._start(context=flow)
except PeaFailToStart as e:
self.logger.critical(
f'Flow couldn\'t get started - Invalid Pod {repr(e)} ')
self.logger.critical(
'Possible causes - invalid/not uploaded pod yamls & pymodules')
# TODO: Send correct error message
raise FlowStartException(repr(e))
except Exception as e:
self.logger.critical(
f'Got following error while starting the flow: {repr(e)}')
raise FlowStartException(repr(e))
self._store[flow_id] = {}
self._store[flow_id]['flow'] = flow
self._store[flow_id]['files'] = files
self.logger.info(
f'Started flow with flow_id {colored(flow_id, "cyan")}')
return flow_id, flow.host, flow.port_expose