Introduce State Passing Between Operators

core/amber/src/main/python/core/architecture/managers/context.py

@@ -2,6 +2,7 @@
 from .console_message_manager import ConsoleMessageManager
 from .debug_manager import DebugManager
 from .exception_manager import ExceptionManager
+from .marker_processing_manager import MarkerProcessingManager
 from .tuple_processing_manager import TupleProcessingManager
 from .executor_manager import ExecutorManager
 from .pause_manager import PauseManager
@@ -26,6 +27,7 @@ def __init__(self, worker_id, input_queue):
         self.input_queue: InternalQueue = input_queue
         self.executor_manager = ExecutorManager()
         self.tuple_processing_manager = TupleProcessingManager()
+        self.marker_processing_manager = MarkerProcessingManager()
         self.exception_manager = ExceptionManager()
         self.state_manager = StateManager(
             {

core/amber/src/main/python/core/architecture/managers/marker_processing_manager.py

@@ -0,0 +1,16 @@
+from typing import Optional
+from core.models.marker import State, Marker
+
+
+class MarkerProcessingManager:
+    def __init__(self):
+        self.current_input_marker: Optional[Marker] = None
+        self.current_output_state: Optional[State] = None
+
+    def get_input_marker(self) -> Optional[State]:
+        ret, self.current_input_marker = self.current_input_marker, None
+        return ret
+
+    def get_output_state(self) -> Optional[State]:
+        ret, self.current_output_state = self.current_output_state, None
+        return ret

core/amber/src/main/python/core/architecture/managers/tuple_processing_manager.py

@@ -1,21 +1,29 @@
 from threading import Event, Condition
-from typing import Optional, Union, Tuple, Iterator
+from typing import Optional, Tuple, Iterator
 
-from core.models import InputExhausted
 from proto.edu.uci.ics.amber.engine.common import PortIdentity
 
 
 class TupleProcessingManager:
     def __init__(self):
-        self.current_input_tuple: Optional[Union[Tuple, InputExhausted]] = None
+        self.current_input_tuple: Optional[Tuple] = None
         self.current_input_port_id: Optional[PortIdentity] = None
-        self.current_input_tuple_iter: Optional[
-            Iterator[Union[Tuple, InputExhausted]]
-        ] = None
+        self.current_input_tuple_iter: Optional[Iterator[Tuple]] = None
         self.current_output_tuple: Optional[Tuple] = None
         self.context_switch_condition: Condition = Condition()
         self.finished_current: Event = Event()
 
+    def get_input_tuple(self) -> Optional[Tuple]:
+        ret, self.current_input_tuple = self.current_input_tuple, None
+        return ret
+
     def get_output_tuple(self) -> Optional[Tuple]:
         ret, self.current_output_tuple = self.current_output_tuple, None
         return ret
+
+    def get_input_port_id(self) -> int:
+        port_id = self.current_input_port_id
+        # no upstream, special case for source executor.
+        if port_id is None:
+            return 0
+        return port_id.id

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

@@ -1,8 +1,15 @@
 from typing import Iterator, Optional, Union, Dict, List
-
-from core.models import Tuple, ArrowTableTupleProvider, Schema, InputExhausted
-from core.models.internal_marker import EndOfAll, InternalMarker, SenderChange
-from core.models.marker import EndOfUpstream
+from pyarrow.lib import Table
+from core.models import Tuple, ArrowTableTupleProvider, Schema
+from core.models.internal_marker import (
+    InternalMarker,
+    StartOfOutputPorts,
+    EndOfOutputPorts,
+    SenderChange,
+    EndOfInputPort,
+    StartOfInputPort,
+)
+from core.models.marker import EndOfInputChannel, State, StartOfInputChannel, Marker
 from core.models.payload import DataFrame, DataPayload, MarkerFrame
 from proto.edu.uci.ics.amber.engine.common import (
     ActorVirtualIdentity,
@@ -50,6 +57,7 @@ def __init__(self):
         self._ports: Dict[PortIdentity, WorkerPort] = dict()
         self._channels: Dict[ChannelIdentity, Channel] = dict()
         self._current_channel_id: Optional[ChannelIdentity] = None
+        self.started = False
 
     def add_input_port(self, port_id: PortIdentity, schema: Schema) -> None:
         if port_id.id is None:
@@ -78,16 +86,21 @@ def register_input(
 
     def process_data_payload(
         self, from_: ActorVirtualIdentity, payload: DataPayload
-    ) -> Iterator[Union[Tuple, InputExhausted, InternalMarker]]:
+    ) -> Iterator[Union[Tuple, InternalMarker]]:
         # special case used to yield for source op
         if from_ == InputManager.SOURCE_STARTER:
-            yield InputExhausted()
-            yield EndOfAll()
+            yield EndOfInputChannel()
+            yield EndOfOutputPorts()
             return
-        current_channel_id = None
-        for channel_id, channel in self._channels.items():
-            if channel_id.from_worker_id == from_:
-                current_channel_id = channel_id
+
+        current_channel_id = next(
+            (
+                channel_id
+                for channel_id, channel in self._channels.items()
+                if channel_id.from_worker_id == from_
+            ),
+            None,
+        )
 
         if (
             self._current_channel_id is None
@@ -97,17 +110,30 @@ def process_data_payload(
             yield SenderChange(current_channel_id)
 
         if isinstance(payload, DataFrame):
-            for field_accessor in ArrowTableTupleProvider(payload.frame):
-                yield Tuple(
-                    {name: field_accessor for name in payload.frame.column_names},
-                    schema=self._ports[
-                        self._channels[self._current_channel_id].port_id
-                    ].get_schema(),
-                )
+            yield from self._process_data(payload.frame)
+        elif isinstance(payload, MarkerFrame):
+            yield from self._process_marker(payload.frame)
+        else:
+            raise NotImplementedError()
 
-        elif isinstance(payload, MarkerFrame) and isinstance(
-            payload.frame, EndOfUpstream
-        ):
+    def _process_data(self, table: Table) -> Iterator[Tuple]:
+        schema = self._ports[
+            self._channels[self._current_channel_id].port_id
+        ].get_schema()
+        for field_accessor in ArrowTableTupleProvider(table):
+            yield Tuple(
+                {name: field_accessor for name in table.column_names}, schema=schema
+            )
+
+    def _process_marker(self, marker: Marker) -> Iterator[InternalMarker]:
+        if isinstance(marker, State):
+            yield marker
+        if isinstance(marker, StartOfInputChannel):
+            if not self.started:
+                yield StartOfOutputPorts()
+            self.started = True
+            yield StartOfInputPort()
+        if isinstance(marker, EndOfInputChannel):
             channel = self._channels[self._current_channel_id]
             channel.complete()
             port_id = channel.port_id
@@ -119,14 +145,11 @@ def process_data_payload(
             )
 
             if port_completed:
-                yield InputExhausted()
+                yield EndOfInputPort()
 
             all_ports_completed = all(
                 map(lambda port: port.is_completed(), self._ports.values())
             )
 
             if all_ports_completed:
-                yield EndOfAll()
-
-        else:
-            raise NotImplementedError()
+                yield EndOfOutputPorts()

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

@@ -2,27 +2,26 @@
 from collections import OrderedDict
 from itertools import chain
 from loguru import logger
-from typing import Iterable, Iterator
-
 from pyarrow import Table
+from typing import Iterable, Iterator
 
 from core.architecture.packaging.input_manager import WorkerPort, Channel
+from core.architecture.sendsemantics.broad_cast_partitioner import (
+    BroadcastPartitioner,
+)
 from core.architecture.sendsemantics.hash_based_shuffle_partitioner import (
     HashBasedShufflePartitioner,
 )
+from core.architecture.sendsemantics.one_to_one_partitioner import OneToOnePartitioner
+from core.architecture.sendsemantics.partitioner import Partitioner
 from core.architecture.sendsemantics.range_based_shuffle_partitioner import (
     RangeBasedShufflePartitioner,
 )
-from core.architecture.sendsemantics.one_to_one_partitioner import OneToOnePartitioner
-from core.architecture.sendsemantics.partitioner import Partitioner
 from core.architecture.sendsemantics.round_robin_partitioner import (
     RoundRobinPartitioner,
 )
-from core.architecture.sendsemantics.broad_cast_partitioner import (
-    BroadcastPartitioner,
-)
 from core.models import Tuple, Schema, MarkerFrame
-from core.models.marker import EndOfUpstream
+from core.models.marker import Marker
 from core.models.payload import DataPayload, DataFrame
 from core.util import get_one_of
 from proto.edu.uci.ics.amber.engine.architecture.sendsemantics import (
@@ -99,33 +98,33 @@ def tuple_to_batch(
             )
         )
 
-    def tuple_to_frame(self, tuples: typing.List[Tuple]) -> DataFrame:
-        return DataFrame(
-            frame=Table.from_pydict(
-                {
-                    name: [t[name] for t in tuples]
-                    for name in self.get_port().get_schema().get_attr_names()
-                },
-                schema=self.get_port().get_schema().as_arrow_schema(),
-            )
-        )
-
-    def emit_end_of_upstream(
-        self,
+    def emit_marker(
+        self, marker: Marker
     ) -> Iterable[typing.Tuple[ActorVirtualIdentity, DataPayload]]:
         return chain(
             *(
                 (
                     (
                         receiver,
                         (
-                            MarkerFrame(tuples)
-                            if isinstance(tuples, EndOfUpstream)
-                            else self.tuple_to_frame(tuples)
+                            MarkerFrame(payload)
+                            if isinstance(payload, Marker)
+                            else self.tuple_to_frame(payload)
                         ),
                     )
-                    for receiver, tuples in partitioner.no_more()
+                    for receiver, payload in partitioner.flush(marker)
                 )
                 for partitioner in self._partitioners.values()
             )
         )
+
+    def tuple_to_frame(self, tuples: typing.List[Tuple]) -> DataFrame:
+        return DataFrame(
+            frame=Table.from_pydict(
+                {
+                    name: [t[name] for t in tuples]
+                    for name in self.get_port().get_schema().get_attr_names()
+                },
+                schema=self.get_port().get_schema().as_arrow_schema(),
+            )
+        )

core/amber/src/main/python/core/architecture/sendsemantics/broad_cast_partitioner.py

@@ -5,7 +5,7 @@
 
 from core.architecture.sendsemantics.partitioner import Partitioner
 from core.models import Tuple
-from core.models.marker import EndOfUpstream
+from core.models.marker import Marker
 from core.util import set_one_of
 from proto.edu.uci.ics.amber.engine.architecture.sendsemantics import (
     Partitioning,
@@ -34,20 +34,18 @@ def add_tuple_to_batch(
             self.reset()
 
     @overrides
-    def no_more(
-        self,
+    def flush(
+        self, marker: Marker
     ) -> Iterator[
-        typing.Tuple[
-            ActorVirtualIdentity, typing.Union[EndOfUpstream, typing.List[Tuple]]
-        ]
+        typing.Tuple[ActorVirtualIdentity, typing.Union[Marker, typing.List[Tuple]]]
     ]:
         if len(self.batch) > 0:
             for receiver in self.receivers:
                 yield receiver, self.batch
 
         self.reset()
         for receiver in self.receivers:
-            yield receiver, EndOfUpstream()
+            yield receiver, marker
 
     @overrides
     def reset(self) -> None:

core/amber/src/main/python/core/architecture/sendsemantics/hash_based_shuffle_partitioner.py

@@ -3,10 +3,9 @@
 
 from loguru import logger
 from overrides import overrides
-
 from core.architecture.sendsemantics.partitioner import Partitioner
 from core.models import Tuple
-from core.models.marker import EndOfUpstream
+from core.models.marker import Marker
 from core.util import set_one_of
 from proto.edu.uci.ics.amber.engine.architecture.sendsemantics import (
     HashBasedShufflePartitioning,
@@ -43,14 +42,12 @@ def add_tuple_to_batch(
             self.receivers[hash_code] = (receiver, list())
 
     @overrides
-    def no_more(
-        self,
+    def flush(
+        self, marker: Marker
     ) -> Iterator[
-        typing.Tuple[
-            ActorVirtualIdentity, typing.Union[EndOfUpstream, typing.List[Tuple]]
-        ]
+        typing.Tuple[ActorVirtualIdentity, typing.Union[Marker, typing.List[Tuple]]]
     ]:
         for receiver, batch in self.receivers:
             if len(batch) > 0:
                 yield receiver, batch
-            yield receiver, EndOfUpstream()
+            yield receiver, marker

core/amber/src/main/python/core/architecture/sendsemantics/one_to_one_partitioner.py

@@ -2,10 +2,9 @@
 from typing import Iterator
 
 from overrides import overrides
-
 from core.architecture.sendsemantics.partitioner import Partitioner
 from core.models import Tuple
-from core.models.marker import EndOfUpstream
+from core.models.marker import Marker
 from core.util import set_one_of
 from proto.edu.uci.ics.amber.engine.architecture.sendsemantics import (
     OneToOnePartitioning,
@@ -34,17 +33,15 @@ def add_tuple_to_batch(
             self.reset()
 
     @overrides
-    def no_more(
-        self,
+    def flush(
+        self, marker: Marker
     ) -> Iterator[
-        typing.Tuple[
-            ActorVirtualIdentity, typing.Union[EndOfUpstream, typing.List[Tuple]]
-        ]
+        typing.Tuple[ActorVirtualIdentity, typing.Union[Marker, typing.List[Tuple]]]
     ]:
         if len(self.batch) > 0:
             yield self.receiver, self.batch
         self.reset()
-        yield self.receiver, EndOfUpstream()
+        yield self.receiver, marker
 
     @overrides
     def reset(self) -> None:

core/amber/src/main/python/core/architecture/sendsemantics/partitioner.py

@@ -5,7 +5,7 @@
 from betterproto import Message
 
 from core.models import Tuple
-from core.models.marker import EndOfUpstream
+from core.models.marker import Marker
 from core.util import get_one_of
 from proto.edu.uci.ics.amber.engine.architecture.sendsemantics import Partitioning
 from proto.edu.uci.ics.amber.engine.common import ActorVirtualIdentity
@@ -20,12 +20,10 @@ def add_tuple_to_batch(
     ) -> Iterator[typing.Tuple[ActorVirtualIdentity, typing.List[Tuple]]]:
         pass
 
-    def no_more(
-        self,
+    def flush(
+        self, marker: Marker
     ) -> Iterator[
-        typing.Tuple[
-            ActorVirtualIdentity, typing.Union[EndOfUpstream, typing.List[Tuple]]
-        ]
+        typing.Tuple[ActorVirtualIdentity, typing.Union[Marker, typing.List[Tuple]]]
     ]:
         pass