wip: for loop no ECM

amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/controlcommands.proto

@@ -46,6 +46,7 @@ message ControlRequest {
     PortCompletedRequest portCompletedRequest = 9;
     WorkerStateUpdatedRequest workerStateUpdatedRequest = 10;
     LinkWorkersRequest linkWorkersRequest = 11;
+    IterationCompletedRequest iterationCompletedRequest = 12;
 
     // request for worker
     AddInputChannelRequest addInputChannelRequest = 50;
@@ -58,6 +59,7 @@ message ControlRequest {
     PrepareCheckpointRequest prepareCheckpointRequest = 57;
     QueryStatisticsRequest queryStatisticsRequest = 58;
 
+
     // request for testing
     Ping ping = 100;
     Pong pong = 101;
@@ -278,4 +280,8 @@ enum StatisticsUpdateTarget {
 message QueryStatisticsRequest{
   repeated core.ActorVirtualIdentity filterByWorkers = 1;
   StatisticsUpdateTarget updateTarget = 2;
+}
+
+message IterationCompletedRequest{
+  core.OperatorIdentity LoopStartId = 1 [(scalapb.field).no_box = true];
 }

amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/controllerservice.proto

@@ -42,6 +42,7 @@ service ControllerService {
   rpc PauseWorkflow(EmptyRequest) returns (EmptyReturn);
   rpc WorkerStateUpdated(WorkerStateUpdatedRequest) returns (EmptyReturn);
   rpc WorkerExecutionCompleted(EmptyRequest) returns (EmptyReturn);
+  rpc IterationCompleted(IterationCompletedRequest) returns (EmptyReturn);
   rpc LinkWorkers(LinkWorkersRequest) returns (EmptyReturn);
   rpc ControllerInitiateQueryStatistics(QueryStatisticsRequest) returns (EmptyReturn);
   rpc RetryWorkflow(RetryWorkflowRequest) returns (EmptyReturn);

amber/src/main/python/core/architecture/packaging/input_manager.py

@@ -173,3 +173,8 @@ def _process_data(self, table: Table) -> Iterator[Tuple]:
             yield Tuple(
                 {name: field_accessor for name in table.column_names}, schema=schema
             )
+
+    def get_input_state_uri(self):
+        return next(iter(self._input_port_mat_reader_runnables.values()))[
+            0
+        ].uri.replace("/result", "/state")

amber/src/main/python/core/architecture/packaging/output_manager.py

@@ -87,6 +87,8 @@ def __init__(self, worker_id: str):
             PortIdentity, typing.Tuple[Queue, PortStorageWriter, Thread]
         ] = dict()
 
+        self._storage_uris: typing.Dict[PortIdentity, str] = dict()
+
     def is_missing_output_ports(self):
         """
         This method is only used for ensuring correct region execution.
@@ -126,6 +128,7 @@ def set_up_port_storage_writer(self, port_id: PortIdentity, storage_uri: str):
         Create a separate thread for saving output tuples of a port
         to storage in batch.
         """
+        self._storage_uris[port_id] = storage_uri
         document, _ = DocumentFactory.open_document(storage_uri)
         buffered_item_writer = document.writer(str(get_worker_index(self.worker_id)))
         writer_queue = Queue()
@@ -171,6 +174,21 @@ def save_tuple_to_storage_if_needed(self, tuple_: Tuple, port_id=None) -> None:
                 PortStorageWriterElement(data_tuple=tuple_)
             )
 
+    def save_state_to_storage_if_needed(self, state: State, port_id=None) -> None:
+        if port_id is None:
+            uris = self._storage_uris.values()
+        elif port_id in self._storage_uris:
+            uris = [self._storage_uris[port_id]]
+        else:
+            return
+
+        for uri in uris:
+            writer = DocumentFactory.create_document(
+                uri.replace("/result", "/state"), state.schema
+            ).writer(str(get_worker_index(self.worker_id)))
+            writer.put_one(Tuple(vars(state)))
+            writer.close()
+
     def close_port_storage_writers(self) -> None:
         """
         Flush the buffers of port storage writers and wait for all the

amber/src/main/python/core/models/operator.py

@@ -293,3 +293,43 @@ def process_table(self, table: Table, port: int) -> Iterator[Optional[TableLike]
             time, or None.
         """
         yield
+
+
+class LoopStartOperator(TableOperator):
+    def open(self) -> None:
+        pass
+
+    @abstractmethod
+    def process_table(self, table: Table, port: int) -> Iterator[Optional[TableLike]]:
+        yield
+
+    @overrides.final
+    def produce_state_on_finish(self, port: int) -> State:
+        from pickle import dumps
+
+        self.state["table"] = dumps(Table(self._TableOperator__table_data[port]))
+        return State().from_dict(self.state)
+
+    def close(self) -> None:
+        pass
+
+
+class LoopEndOperator(TableOperator):
+    def open(self) -> None:
+        pass
+
+    @overrides.final
+    def process_table(self, table: Table, port: int) -> Iterator[Optional[TableLike]]:
+        yield table
+
+    def close(self) -> None:
+        pass
+
+    @abstractmethod
+    def condition(self) -> None:
+        pass
+
+    def loop_start_id(self) -> str:
+        del self.state["table"]
+        del self.state["output"]
+        return self.state["LoopStartId"]

amber/src/main/python/core/models/state.py

@@ -27,14 +27,28 @@
 @dataclass
 class State:
     def __init__(
-        self, table: Optional[Table] = None, pass_to_all_downstream: bool = False
+        self, table: Optional[Table] = None, pass_to_all_downstream: bool = True
     ):
         self.schema = Schema()
         self.passToAllDownstream = pass_to_all_downstream
         if table is not None:
             self.__dict__.update(table.to_pandas().iloc[0].to_dict())
             self.schema = Schema(table.schema)
 
+    @classmethod
+    def from_tuple(cls, tuple, schema):
+        obj = cls()
+        obj.__dict__.update(tuple.as_dict())
+        obj.schema = schema
+        return obj
+
+    @classmethod
+    def from_dict(cls, dictionary):
+        obj = cls()
+        for item in dictionary:
+            obj.add(item, dictionary[item])
+        return obj
+
     def add(
         self, key: str, value: any, value_type: Optional[AttributeType] = None
     ) -> None:
@@ -53,6 +67,12 @@ def to_table(self) -> Table:
             schema=self.schema.as_arrow_schema(),
         )
 
+    def to_dict(self) -> dict:
+        dictionary = self.__dict__
+        del dictionary["passToAllDownstream"]
+        del dictionary["schema"]
+        return dictionary
+
     def __setattr__(self, key: str, value: any) -> None:
         self.add(key, value)
 

amber/src/main/python/core/runnables/data_processor.py

@@ -100,6 +100,7 @@ def process_state(self, state: State) -> None:
                 self._context.worker_id,
                 self._context.console_message_manager.print_buf,
             ):
+                self._switch_context()
                 self._set_output_state(executor.process_state(state, port_id))
 
         except Exception as err:

amber/src/main/python/core/runnables/main_loop.py

@@ -18,6 +18,8 @@
 import threading
 import time
 import typing
+import uuid
+
 from loguru import logger
 from overrides import overrides
 from pampy import match
@@ -38,8 +40,10 @@
     ECMElement,
     InternalQueueElement,
 )
+from core.models.operator import LoopEndOperator, LoopStartOperator
 from core.models.state import State
 from core.runnables.data_processor import DataProcessor
+from core.storage.document_factory import DocumentFactory
 from core.util import StoppableQueueBlockingRunnable, get_one_of
 from core.util.console_message.timestamp import current_time_in_local_timezone
 from core.util.customized_queue.queue_base import QueueElement
@@ -48,6 +52,7 @@
     PortIdentity,
     ChannelIdentity,
     EmbeddedControlMessageIdentity,
+    OperatorIdentity,
 )
 from proto.org.apache.texera.amber.engine.architecture.rpc import (
     ConsoleMessage,
@@ -61,6 +66,7 @@
     EmbeddedControlMessage,
     AsyncRpcContext,
     ControlRequest,
+    IterationCompletedRequest,
 )
 from proto.org.apache.texera.amber.engine.architecture.worker import (
     WorkerState,
@@ -94,12 +100,26 @@ def complete(self) -> None:
         """
         # flush the buffered console prints
         self._check_and_report_console_messages(force_flush=True)
-        self.context.executor_manager.executor.close()
+        controller_interface = self._async_rpc_client.controller_stub()
+        executor = self.context.executor_manager.executor
+        if isinstance(executor, LoopEndOperator) and executor.condition():
+            controller_interface.iteration_completed(
+                IterationCompletedRequest(OperatorIdentity(executor.loop_start_id()))
+            )
+            uri = executor.state["LoopStartStateURI"]
+            del executor.state["LoopStartStateURI"]
+            del executor.state["LoopStartId"]
+            state = State.from_dict(executor.state)
+            writer = DocumentFactory.create_document(uri, state.schema).writer(
+                str(uuid.uuid4())
+            )
+            writer.put_one(Tuple(vars(state)))
+            writer.close()
+        executor.close()
         # stop the data processing thread
         self.data_processor.stop()
         self.context.state_manager.transit_to(WorkerState.COMPLETED)
         self.context.statistics_manager.update_total_execution_time(time.time_ns())
-        controller_interface = self._async_rpc_client.controller_stub()
         controller_interface.worker_execution_completed(EmptyRequest())
         self.context.close()
 
@@ -188,14 +208,26 @@ def process_input_state(self) -> None:
         output_state = self.context.state_processing_manager.get_output_state()
         self._switch_context()
         if output_state is not None:
-            for to, batch in self.context.output_manager.emit_state(output_state):
-                self._output_queue.put(
-                    DataElement(
-                        tag=ChannelIdentity(
-                            ActorVirtualIdentity(self.context.worker_id), to, False
-                        ),
-                        payload=batch,
+            if isinstance(self.context.executor_manager.executor, LoopStartOperator):
+                output_state.add(
+                    "LoopStartId",
+                    self.context.worker_id.split("-", 1)[1].rsplit("-main-0", 1)[0],
+                )
+                output_state.add(
+                    "LoopStartStateURI",
+                    self.context.input_manager.get_input_state_uri(),
+                )
+                for to, batch in self.context.output_manager.emit_state(output_state):
+                    self._output_queue.put(
+                        DataElement(
+                            tag=ChannelIdentity(
+                                ActorVirtualIdentity(self.context.worker_id), to, False
+                            ),
+                            payload=batch,
+                        )
                     )
+                self.context.output_manager.save_state_to_storage_if_needed(
+                    output_state
                 )
 
     def process_tuple_with_udf(self) -> Iterator[Optional[Tuple]]:
@@ -329,7 +361,7 @@ def _process_ecm(self, ecm_element: ECMElement):
 
             if ecm.ecm_type != EmbeddedControlMessageType.NO_ALIGNMENT:
                 self.context.pause_manager.resume(PauseType.ECM_PAUSE)
-
+            self._switch_context()
             if self.context.tuple_processing_manager.current_internal_marker:
                 {
                     StartChannel: self._process_start_channel,