def mock_query_statistics(self, mock_controller, mock_sender_actor,
command_sequence):
command = set_one_of(ControlCommandV2, QueryStatisticsV2())
payload = set_one_of(
ControlPayloadV2,
ControlInvocationV2(command_id=command_sequence, command=command))
return ControlElement(tag=mock_controller, payload=payload)
def receive(self, from_: ActorVirtualIdentity,
control_invocation: ControlInvocationV2):
command: ControlCommandV2 = get_one_of(control_invocation.command)
logger.debug(
f"PYTHON receives a ControlInvocation: {control_invocation}")
try:
handler = self.look_up(command)
control_return: ControlReturnV2 = set_one_of(
ControlReturnV2, handler(self._context, command))
except Exception as exception:
logger.exception(exception)
control_return: ControlReturnV2 = set_one_of(
ControlReturnV2, ControlException(str(exception)))
payload: ControlPayloadV2 = set_one_of(
ControlPayloadV2,
ReturnInvocationV2(
original_command_id=control_invocation.command_id,
control_return=control_return))
# reply to the sender
to = from_
logger.debug(
f"PYTHON returns a ReturnInvocation {payload}, replying the command {command}"
)
self._output_queue.put(ControlElement(tag=to, payload=payload))
def mock_update_input_linking(self, mock_controller, mock_sender_actor,
mock_link, command_sequence):
command = set_one_of(
ControlCommandV2,
UpdateInputLinkingV2(identifier=mock_sender_actor,
input_link=mock_link))
payload = set_one_of(
ControlPayloadV2,
ControlInvocationV2(command_id=command_sequence, command=command))
return ControlElement(tag=mock_controller, payload=payload)
def mock_add_partitioning(self, mock_controller, mock_receiver_actor,
command_sequence):
command = set_one_of(
ControlCommandV2,
AddPartitioningV2(tag=mock_receiver_actor,
partitioning=set_one_of(
Partitioning,
OneToOnePartitioning(
batch_size=1,
receivers=[mock_receiver_actor]))))
payload = set_one_of(
ControlPayloadV2,
ControlInvocationV2(command_id=command_sequence, command=command))
return ControlElement(tag=mock_controller, payload=payload)
def __init__(self, input_queue: InternalQueue,
output_queue: InternalQueue):
super().__init__(self.__class__.__name__, queue=input_queue)
self._input_queue: InternalQueue = input_queue
self._output_queue: InternalQueue = output_queue
self._operator: Optional[Operator] = None
self._current_input_tuple: Optional[Union[Tuple,
InputExhausted]] = None
self._current_input_link: Optional[LinkIdentity] = None
self._current_input_tuple_iter: Optional[Iterator[Union[
Tuple, InputExhausted]]] = None
self._input_links: List[LinkIdentity] = list()
self._input_link_map: MutableMapping[LinkIdentity, int] = dict()
self.context = Context(self)
self._async_rpc_server = AsyncRPCServer(output_queue,
context=self.context)
self._async_rpc_client = AsyncRPCClient(output_queue,
context=self.context)
self._print_log_handler = PrintLogHandler(
lambda msg: self._async_rpc_client.send(
ActorVirtualIdentity(name="CONTROLLER"),
set_one_of(ControlCommandV2, PythonPrintV2(message=msg))))
logger.add(self._print_log_handler, level='PRINT', filter="operators")
def __init__(self, partitioning: RoundRobinPartitioning):
super().__init__(set_one_of(Partitioning, partitioning))
self.batch_size = partitioning.batch_size
self.receivers: List[typing.Tuple[
ActorVirtualIdentity, List[Tuple]]] = [
(receiver, list()) for receiver in partitioning.receivers
]
self.round_robin_index = 0
def __init__(self, partitioning: HashBasedShufflePartitioning):
super().__init__(set_one_of(Partitioning, partitioning))
logger.info(f"got {partitioning}")
self.batch_size = partitioning.batch_size
self.receivers: List[typing.Tuple[
ActorVirtualIdentity, List[Tuple]]] = [
(receiver, list()) for receiver in partitioning.receivers
]
self.hash_column_indices = partitioning.hash_column_indices
def report_exception(self) -> None:
"""
Report the traceback of current stack when an exception occurs.
"""
self._print_log_handler.flush()
message: str = traceback.format_exc(limit=-1)
control_command = set_one_of(ControlCommandV2,
LocalOperatorExceptionV2(message=message))
self._async_rpc_client.send(ActorVirtualIdentity(name="CONTROLLER"),
control_command)
def complete(self) -> None:
"""
Complete the DataProcessor, marking state to COMPLETED, and notify the controller.
"""
# flush the buffered console prints
self._print_log_handler.flush()
self._operator.close()
self.context.state_manager.transit_to(WorkerState.COMPLETED)
control_command = set_one_of(ControlCommandV2,
WorkerExecutionCompletedV2())
self._async_rpc_client.send(ActorVirtualIdentity(name="CONTROLLER"),
control_command)
def send(self, to: ActorVirtualIdentity,
control_command: ControlCommandV2) -> Future:
"""
Send the ControlCommand to the target actor.
:param to: ActorVirtualIdentity, the receiver.
:param control_command: ControlCommandV2, the command to be sent.
"""
payload = set_one_of(
ControlPayloadV2,
ControlInvocationV2(self._send_sequences[to],
command=control_command))
self._output_queue.put(ControlElement(tag=to, payload=payload))
return self._create_future(to)
def test_dp_thread_can_process_messages(
self, mock_link, mock_receiver_actor, mock_controller, input_queue,
output_queue, mock_data_element, dp_thread,
mock_update_input_linking, mock_add_partitioning,
mock_end_of_upstream, mock_query_statistics, mock_tuple,
command_sequence, reraise):
dp_thread.start()
# can process UpdateInputLinking
input_queue.put(mock_update_input_linking)
assert output_queue.get() == ControlElement(
tag=mock_controller,
payload=ControlPayloadV2(return_invocation=ReturnInvocationV2(
original_command_id=command_sequence,
control_return=ControlReturnV2())))
# can process AddPartitioning
input_queue.put(mock_add_partitioning)
assert output_queue.get() == ControlElement(
tag=mock_controller,
payload=ControlPayloadV2(return_invocation=ReturnInvocationV2(
original_command_id=command_sequence,
control_return=ControlReturnV2())))
# can process a InputDataFrame
input_queue.put(mock_data_element)
output_data_element: DataElement = output_queue.get()
assert output_data_element.tag == mock_receiver_actor
assert isinstance(output_data_element.payload, OutputDataFrame)
data_frame: OutputDataFrame = output_data_element.payload
assert len(data_frame.frame) == 1
assert (data_frame.frame[0] == mock_tuple)
# can process QueryStatistics
input_queue.put(mock_query_statistics)
assert output_queue.get() == ControlElement(
tag=mock_controller,
payload=ControlPayloadV2(return_invocation=ReturnInvocationV2(
original_command_id=1,
control_return=ControlReturnV2(
worker_statistics=WorkerStatistics(
worker_state=WorkerState.RUNNING,
input_tuple_count=1,
output_tuple_count=1)))))
# can process EndOfUpstream
input_queue.put(mock_end_of_upstream)
assert output_queue.get() == DataElement(tag=mock_receiver_actor,
payload=EndOfUpstream())
# WorkerExecutionCompletedV2 should be triggered when workflow finishes
assert output_queue.get() == ControlElement(
tag=mock_controller,
payload=ControlPayloadV2(control_invocation=ControlInvocationV2(
command_id=0,
command=ControlCommandV2(
worker_execution_completed=WorkerExecutionCompletedV2()))))
# can process ReturnInvocation
input_queue.put(
ControlElement(tag=mock_controller,
payload=set_one_of(
ControlPayloadV2,
ReturnInvocationV2(
original_command_id=0,
control_return=ControlReturnV2()))))
reraise()
def __init__(self, partitioning: OneToOnePartitioning):
super().__init__(set_one_of(Partitioning, partitioning))
self.batch_size = partitioning.batch_size
self.batch: list[Tuple] = list()
self.receiver = partitioning.receivers[0] # one to one will have only one receiver.