async def test_send_static_ctrl_msg(mocker): # AsyncMock does not seem to exist in python 3.7, this is a manual workaround receive_cb = mocker.Mock() async def mock_wrapper(msg): receive_cb() args = set_pea_parser().parse_args([]) grpclet = Grpclet(args=args, message_callback=mock_wrapper) asyncio.get_event_loop().create_task(grpclet.start()) receive_cb.assert_not_called() while True: try: def send_status(): return Grpclet.send_ctrl_msg( pod_address=f'{args.host}:{args.port_in}', command='STATUS') await asyncio.get_event_loop().run_in_executor(None, send_status) break except RpcError: await asyncio.sleep(0.1) receive_cb.assert_called() await grpclet.close(None)
async def test_send_receive(mocker): # AsyncMock does not seem to exist in python 3.7, this is a manual workaround receive_cb = mocker.Mock() async def mock_wrapper(msg): receive_cb() args = set_pea_parser().parse_args([]) grpclet = Grpclet(args=args, message_callback=mock_wrapper) asyncio.get_event_loop().create_task(grpclet.start()) receive_cb.assert_not_called() await grpclet.send_message(_create_msg(args)) await asyncio.sleep(0.1) receive_cb.assert_called() await grpclet.close(None)
async def test_send_non_blocking(mocker): receive_cb = mocker.Mock() async def blocking_cb(msg): receive_cb() time.sleep(1.0) return msg args = set_pea_parser().parse_args([]) grpclet = Grpclet(args=args, message_callback=blocking_cb) asyncio.get_event_loop().create_task(grpclet.start()) receive_cb.assert_not_called() await grpclet.send_message(_create_msg(args)) await asyncio.sleep(0.1) assert receive_cb.call_count == 1 await grpclet.send_message(_create_msg(args)) await asyncio.sleep(0.1) assert receive_cb.call_count == 2 await grpclet.close(None)
class GRPCDataRuntime(BaseRuntime, ABC): """Runtime procedure leveraging :class:`Grpclet` for sending DataRequests""" def __init__(self, args: argparse.Namespace, **kwargs): """Initialize grpc and data request handling. :param args: args from CLI :param kwargs: extra keyword arguments """ super().__init__(args, **kwargs) self._id = random_identity() self._loop = get_or_reuse_loop() self._last_active_time = time.perf_counter() self._pending_msgs = defaultdict(list) # type: Dict[str, List[Message]] self._partial_requests = None self._pending_tasks = [] self._static_routing_table = args.static_routing_table self._data_request_handler = DataRequestHandler(args, self.logger) self._grpclet = Grpclet( args=self.args, message_callback=self._callback, logger=self.logger, ) def _update_pending_tasks(self): self._pending_tasks = [task for task in self._pending_tasks if not task.done()] def run_forever(self): """Start the `Grpclet`.""" self._grpclet_task = self._loop.create_task(self._grpclet.start()) try: self._loop.run_until_complete(self._grpclet_task) except asyncio.CancelledError: self.logger.warning('Grpclet task was cancelled') def teardown(self): """Close the `Grpclet` and `DataRequestHandler`.""" self.logger.debug('Teardown GRPCDataRuntime') self._data_request_handler.close() start = time.time() while self._pending_tasks and time.time() - start < 1.0: self._update_pending_tasks() time.sleep(0.1) self._loop.stop() self._loop.close() super().teardown() async def _close_grpclet(self): await self._grpclet.close() self._grpclet_task.cancel() @staticmethod def get_control_address(**kwargs): """ Does return None, exists for keeping interface compatible with ZEDRuntime :param kwargs: extra keyword arguments :returns: None """ return None @staticmethod def is_ready(ctrl_address: str, **kwargs) -> bool: """ Check if status is ready. :param ctrl_address: the address where the control message needs to be sent :param kwargs: extra keyword arguments :return: True if status is ready else False. """ try: response = Grpclet.send_ctrl_msg(ctrl_address, 'STATUS') except RpcError: return False return True @staticmethod def activate( **kwargs, ): """ Does nothing :param kwargs: extra keyword arguments """ pass @staticmethod def cancel( control_address: str, **kwargs, ): """ Cancel this runtime by sending a TERMINATE control message :param control_address: the address where the control message needs to be sent :param kwargs: extra keyword arguments """ try: Grpclet.send_ctrl_msg(control_address, 'TERMINATE') except RpcError: # TERMINATE can fail if the the runtime dies before sending the return value pass @staticmethod def wait_for_ready_or_shutdown( timeout: Optional[float], ctrl_address: str, shutdown_event: Union[multiprocessing.Event, threading.Event], **kwargs, ): """ Check if the runtime has successfully started :param timeout: The time to wait before readiness or failure is determined :param ctrl_address: the address where the control message needs to be sent :param shutdown_event: the multiprocessing event to detect if the process failed :param kwargs: extra keyword arguments :return: True if is ready or it needs to be shutdown """ timeout_ns = 1000000000 * timeout if timeout else None now = time.time_ns() while timeout_ns is None or time.time_ns() - now < timeout_ns: if shutdown_event.is_set() or GRPCDataRuntime.is_ready(ctrl_address): return True time.sleep(0.1) return False async def _callback(self, msg: Message) -> None: try: msg = self._post_hook(self._handle(self._pre_hook(msg))) if msg.is_data_request: asyncio.create_task(self._grpclet.send_message(msg)) except RuntimeTerminated: # this is the proper way to end when a terminate signal is sent self._pending_tasks.append(asyncio.create_task(self._close_grpclet())) except KeyboardInterrupt as kbex: self.logger.debug(f'{kbex!r} causes the breaking from the event loop') self._pending_tasks.append(asyncio.create_task(self._close_grpclet())) except (SystemError) as ex: # save executor self.logger.debug(f'{ex!r} causes the breaking from the event loop') self._pending_tasks.append(asyncio.create_task(self._close_grpclet())) except NoExplicitMessage: # silent and do not propagate message anymore # 1. wait partial message to be finished pass except (RuntimeError, Exception, ChainedPodException) as ex: if self.args.on_error_strategy == OnErrorStrategy.THROW_EARLY: raise if isinstance(ex, ChainedPodException): # the error is print from previous pod, no need to show it again # hence just add exception and propagate further # please do NOT add logger.error here! msg.add_exception() else: msg.add_exception(ex, executor=self._data_request_handler._executor) self.logger.error( f'{ex!r}' + f'\n add "--quiet-error" to suppress the exception details' if not self.args.quiet_error else '', exc_info=not self.args.quiet_error, ) if msg.is_data_request: asyncio.create_task(self._grpclet.send_message(msg)) asyncio.create_task(self._grpclet.send_message(msg)) def _handle(self, msg: Message) -> Message: """Register the current message to this pea, so that all message-related properties are up-to-date, including :attr:`request`, :attr:`prev_requests`, :attr:`message`, :attr:`prev_messages`. And then call the executor to handle this message if its envelope's status is not ERROR, else skip handling of message. .. note:: Handle does not handle explicitly message because it may wait for different messages when different parts are expected :param msg: received message :return: the transformed message. """ # skip executor for non-DataRequest if msg.envelope.request_type != 'DataRequest': if msg.request.command == 'TERMINATE': raise RuntimeTerminated() self.logger.debug(f'skip executor: not data request') return msg req_id = msg.envelope.request_id num_expected_parts = self._get_expected_parts(msg) self._data_request_handler.handle( msg=msg, partial_requests=[m.request for m in self._pending_msgs[req_id]] if num_expected_parts > 1 else None, peapod_name=self.name, ) return msg def _get_expected_parts(self, msg): if msg.is_data_request: if not self._static_routing_table: graph = RoutingTable(msg.envelope.routing_table) return graph.active_target_pod.expected_parts else: return self.args.num_part else: return 1 def _pre_hook(self, msg: Message) -> Message: """ Pre-hook function, what to do after first receiving the message. :param msg: received message :return: `Message` """ msg.add_route(self.name, self._id) expected_parts = self._get_expected_parts(msg) req_id = msg.envelope.request_id if expected_parts > 1: self._pending_msgs[req_id].append(msg) num_partial_requests = len(self._pending_msgs[req_id]) if self.logger.debug_enabled: self._log_info_msg( msg, f'({num_partial_requests}/{expected_parts} parts)' if expected_parts > 1 else '', ) if expected_parts > 1 and expected_parts > num_partial_requests: # NOTE: reduce priority is higher than chain exception # otherwise a reducer will lose its function when earlier pods raise exception raise NoExplicitMessage if ( msg.envelope.status.code == jina_pb2.StatusProto.ERROR and self.args.on_error_strategy >= OnErrorStrategy.SKIP_HANDLE ): raise ChainedPodException return msg def _log_info_msg(self, msg, part_str): info_msg = f'recv {msg.envelope.request_type} ' req_type = msg.envelope.request_type if req_type == 'DataRequest': info_msg += ( f'({msg.envelope.header.exec_endpoint}) - ({msg.envelope.request_id}) ' ) elif req_type == 'ControlRequest': info_msg += f'({msg.request.command}) ' info_msg += f'{part_str} from {msg.colored_route}' self.logger.debug(info_msg) def _post_hook(self, msg: Message) -> Message: """ Post-hook function, what to do before handing out the message. :param msg: the transformed message :return: `Message` """ # do NOT access `msg.request.*` in the _pre_hook, as it will trigger the deserialization # all meta information should be stored and accessed via `msg.envelope` self._last_active_time = time.perf_counter() if self._get_expected_parts(msg) > 1: msgs = self._pending_msgs.pop(msg.envelope.request_id) msg.merge_envelope_from(msgs) msg.update_timestamp() return msg