wip: for loop no ECM

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

@@ -57,6 +57,7 @@ message ControlRequest {
     EmptyRequest emptyRequest = 56;
     PrepareCheckpointRequest prepareCheckpointRequest = 57;
     QueryStatisticsRequest queryStatisticsRequest = 58;
+    EndIterationRequest endIterationRequest = 59;
 
     // request for testing
     Ping ping = 100;
@@ -278,4 +279,9 @@ enum StatisticsUpdateTarget {
 message QueryStatisticsRequest{
   repeated core.ActorVirtualIdentity filterByWorkers = 1;
   StatisticsUpdateTarget updateTarget = 2;
+}
+
+message EndIterationRequest{
+  core.ActorVirtualIdentity LoopStartId = 1 [(scalapb.field).no_box = true];
+  int32 iteration = 2;
 }

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

@@ -47,6 +47,7 @@ service WorkerService {
   rpc EndWorker(EmptyRequest) returns (EmptyReturn);
   rpc StartChannel(EmptyRequest) returns (EmptyReturn);
   rpc EndChannel(EmptyRequest) returns (EmptyReturn);
+  rpc EndIteration(EndIterationRequest) returns (EmptyReturn);
   rpc DebugCommand(DebugCommandRequest) returns (EmptyReturn);
   rpc EvaluatePythonExpression(EvaluatePythonExpressionRequest) returns (EvaluatedValue);
   rpc NoOperation(EmptyRequest) returns (EmptyReturn);

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

@@ -35,6 +35,13 @@ def __init__(
             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
+
     def add(
         self, key: str, value: any, value_type: Optional[AttributeType] = None
     ) -> None:

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

@@ -329,7 +329,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,

amber/src/main/python/core/storage/document_factory.py

@@ -61,7 +61,7 @@ def create_document(uri: str, schema: Schema) -> VirtualDocument:
         if parsed_uri.scheme == VFSURIFactory.VFS_FILE_URI_SCHEME:
             _, _, _, resource_type = VFSURIFactory.decode_uri(uri)
 
-            if resource_type in {VFSResourceType.RESULT}:
+            if resource_type in {VFSResourceType.RESULT, VFSResourceType.STATE}:
                 storage_key = DocumentFactory.sanitize_uri_path(parsed_uri)
 
                 # Convert Amber Schema to Iceberg Schema with LARGE_BINARY
@@ -96,7 +96,7 @@ def open_document(uri: str) -> typing.Tuple[VirtualDocument, Optional[Schema]]:
         if parsed_uri.scheme == "vfs":
             _, _, _, resource_type = VFSURIFactory.decode_uri(uri)
 
-            if resource_type in {VFSResourceType.RESULT}:
+            if resource_type in {VFSResourceType.RESULT, VFSResourceType.STATE}:
                 storage_key = DocumentFactory.sanitize_uri_path(parsed_uri)
 
                 table = load_table_metadata(

amber/src/main/python/core/storage/runnables/input_port_materialization_reader_runnable.py

@@ -17,8 +17,8 @@
 
 import typing
 from loguru import logger
-from pyarrow import Table
 from typing import Union
+from pyarrow import Table
 
 from core.architecture.sendsemantics.broad_cast_partitioner import (
     BroadcastPartitioner,
@@ -34,7 +34,7 @@
 from core.architecture.sendsemantics.round_robin_partitioner import (
     RoundRobinPartitioner,
 )
-from core.models import Tuple, InternalQueue, DataFrame, DataPayload
+from core.models import Tuple, InternalQueue, DataFrame, DataPayload, State, StateFrame
 from core.models.internal_queue import DataElement, ECMElement
 from core.storage.document_factory import DocumentFactory
 from core.util import Stoppable, get_one_of
@@ -125,6 +125,15 @@ def tuple_to_batch_with_filter(self, tuple_: Tuple) -> typing.Iterator[DataFrame
             if receiver == self.worker_actor_id:
                 yield self.tuples_to_data_frame(tuples)
 
+    def emit_state_with_filter(self, state: State) -> typing.Iterator[StateFrame]:
+        for receiver, payload in self.partitioner.flush_state(state):
+            if receiver == self.worker_actor_id:
+                yield (
+                    StateFrame(payload)
+                    if isinstance(payload, State)
+                    else self.tuples_to_data_frame(payload)
+                )
+
     def run(self) -> None:
         """
         Main execution logic that reads tuples from the materialized storage and
@@ -149,6 +158,19 @@ def run(self) -> None:
                 tup.cast_to_schema(self.tuple_schema)
                 for data_frame in self.tuple_to_batch_with_filter(tup):
                     self.emit_payload(data_frame)
+            try:
+                state_document, state_schema = DocumentFactory.open_document(
+                    f"{self.uri}/state"
+                )
+                state_iterator = state_document.get()
+                for state in state_iterator:
+                    for state_frame in self.emit_state_with_filter(
+                        State.from_tuple(state, state_schema)
+                    ):
+                        self.emit_payload(state_frame)
+            except ValueError:
+                pass
+
             self.emit_ecm("EndChannel", EmbeddedControlMessageType.PORT_ALIGNMENT)
             self._finished = True
         except Exception as err:

amber/src/main/python/core/storage/vfs_uri_factory.py

@@ -34,6 +34,7 @@ class VFSResourceType(str, Enum):
     RESULT = "result"
     RUNTIME_STATISTICS = "runtimeStatistics"
     CONSOLE_MESSAGES = "consoleMessages"
+    STATE = "state"
 
 
 class VFSURIFactory:

amber/src/main/scala/org/apache/texera/amber/engine/architecture/common/AmberProcessor.scala

@@ -43,7 +43,7 @@ abstract class AmberProcessor(
     with Serializable {
 
   /** FIFO & exactly once */
-  val inputGateway: InputGateway = new NetworkInputGateway(this.actorId)
+  val inputGateway: NetworkInputGateway = new NetworkInputGateway(this.actorId)
 
   // 1. Unified Output
   val outputGateway: NetworkOutputGateway =
@@ -55,7 +55,7 @@ abstract class AmberProcessor(
       }
     )
   // 2. RPC Layer
-  val asyncRPCClient = new AsyncRPCClient(outputGateway, actorId)
+  val asyncRPCClient = new AsyncRPCClient(inputGateway, outputGateway, actorId)
   val asyncRPCServer: AsyncRPCServer =
     new AsyncRPCServer(outputGateway, actorId)
 

amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/WorkflowScheduler.scala

@@ -52,7 +52,7 @@ class WorkflowScheduler(
     this.physicalPlan = updatedPhysicalPlan
   }
 
-  def getNextRegions: Set[Region] = if (!schedule.hasNext) Set() else schedule.next()
+  def getNextRegions: Set[Region] = if (!schedule.hasNext) Set() else schedule.loopNext()
 
   def hasPendingRegions: Boolean = schedule != null && schedule.hasNext
 

amber/src/main/scala/org/apache/texera/amber/engine/architecture/controller/execution/WorkflowExecution.scala

@@ -44,6 +44,7 @@ case class WorkflowExecution() {
     * @throws AssertionError if the `RegionExecution` has already been initialized.
     */
   def initRegionExecution(region: Region): RegionExecution = {
+    regionExecutions.remove(region.id)
     // ensure the region execution hasn't been initialized already.
     assert(
       !regionExecutions.contains(region.id),

amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/AmberFIFOChannel.scala

@@ -41,6 +41,7 @@ class AmberFIFOChannel(val channelId: ChannelIdentity) extends AmberLogging {
   private var portId: Option[PortIdentity] = None
 
   def acceptMessage(msg: WorkflowFIFOMessage): Unit = {
+    //channel remove
     val seq = msg.sequenceNumber
     val payload = msg.payload
     if (isDuplicated(seq)) {

amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/NetworkInputGateway.scala

@@ -86,4 +86,8 @@ class NetworkInputGateway(val actorId: ActorVirtualIdentity)
     enforcers += enforcer
   }
 
+  def removeControlChannel(from: ActorVirtualIdentity): Unit = {
+    inputChannels.remove(ChannelIdentity(from, actorId, isControl = true))
+  }
+
 }

amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/NetworkOutputGateway.scala

@@ -94,4 +94,8 @@ class NetworkOutputGateway(
     idToSequenceNums.getOrElseUpdate(channelId, new AtomicLong()).getAndIncrement()
   }
 
+  def removeControlChannel(to: ActorVirtualIdentity): Unit = {
+    idToSequenceNums.remove(ChannelIdentity(actorId, to, isControl = true))
+  }
+
 }

amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/OutputManager.scala

@@ -22,6 +22,7 @@ package org.apache.texera.amber.engine.architecture.messaginglayer
 import org.apache.texera.amber.core.state.State
 import org.apache.texera.amber.core.storage.DocumentFactory
 import org.apache.texera.amber.core.storage.model.BufferedItemWriter
+import org.apache.texera.amber.core.storage.result.ResultSchema
 import org.apache.texera.amber.core.tuple._
 import org.apache.texera.amber.core.virtualidentity.{ActorVirtualIdentity, ChannelIdentity}
 import org.apache.texera.amber.core.workflow.{PhysicalLink, PortIdentity}
@@ -124,6 +125,9 @@ class OutputManager(
       : mutable.HashMap[PortIdentity, OutputPortResultWriterThread] =
     mutable.HashMap()
 
+  private val ECMWriters: mutable.HashMap[PortIdentity, BufferedItemWriter[Tuple]] =
+    mutable.HashMap()
+
   /**
     * Add down stream operator and its corresponding Partitioner.
     *
@@ -232,6 +236,10 @@ class OutputManager(
     })
   }
 
+  def saveECMToStorageIfNeeded(tuple: Tuple, outputPortId: PortIdentity): Unit = {
+    this.ECMWriters(outputPortId).putOne(new Tuple(ResultSchema.ecmSchema, Array("erge")))
+  }
+
   /**
     * Singal the port storage writer to flush the remaining buffer and wait for commits to finish so that
     * the output port is properly completed. If the output port does not need storage, no action will be done.
@@ -280,6 +288,10 @@ class OutputManager(
   }
 
   private def setupOutputStorageWriterThread(portId: PortIdentity, storageUri: URI): Unit = {
+    this.ECMWriters(portId) = DocumentFactory
+      .createDocument(storageUri.resolve("ecm"), ResultSchema.ecmSchema)
+      .writer(VirtualIdentityUtils.getWorkerIndex(actorId).toString)
+      .asInstanceOf[BufferedItemWriter[Tuple]]
     val bufferedItemWriter = DocumentFactory
       .openDocument(storageUri)
       ._1

amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionCoordinator.scala

@@ -181,6 +181,8 @@ class RegionExecutionCoordinator(
                 val actorRef = actorRefService.getActorRef(workerId)
                 // Remove the actorRef so that no other actors can find the worker and send messages.
                 actorRefService.removeActorRef(workerId)
+                asyncRPCClient.inputGateway.removeControlChannel(workerId)
+                asyncRPCClient.outputGateway.removeControlChannel(workerId)
                 gracefulStop(actorRef, ScalaDuration(5, TimeUnit.SECONDS)).asTwitter()
               }
           }.toSeq
@@ -209,14 +211,15 @@ class RegionExecutionCoordinator(
       regionExecution: RegionExecution,
       attempt: Int = 1
   ): Future[Unit] = {
-    terminateWorkers(regionExecution).rescue { case err =>
-      logger.warn(
-        s"Failed to terminate region ${region.id.id} on attempt $attempt. Retrying in ${killRetryDelay.inMilliseconds} ms.",
-        err
-      )
-      Future
-        .sleep(killRetryDelay)(killRetryTimer)
-        .flatMap(_ => terminateWorkersWithRetry(regionExecution, attempt + 1))
+    terminateWorkers(regionExecution).rescue {
+      case err =>
+        logger.warn(
+          s"Failed to terminate region ${region.id.id} on attempt $attempt. Retrying in ${killRetryDelay.inMilliseconds} ms.",
+          err
+        )
+        Future
+          .sleep(killRetryDelay)(killRetryTimer)
+          .flatMap(_ => terminateWorkersWithRetry(regionExecution, attempt + 1))
     }
   }
 
@@ -568,7 +571,18 @@ class RegionExecutionCoordinator(
           region.getOperator(outputPortId.opId).outputPorts(outputPortId.portId)._3
         val schema =
           schemaOptional.getOrElse(throw new IllegalStateException("Schema is missing"))
-        DocumentFactory.createDocument(storageUriToAdd, schema)
+
+        if (region.getOperators.exists(_.id.logicalOpId.id.startsWith("LoopEnd-operator-"))) {
+          try {
+            DocumentFactory.openDocument(storageUriToAdd)
+          } catch {
+            case _: Exception =>
+              DocumentFactory.createDocument(storageUriToAdd, schema)
+          }
+        } else {
+          DocumentFactory.createDocument(storageUriToAdd, schema)
+        }
+
         WorkflowExecutionsResource.insertOperatorPortResultUri(
           eid = eid,
           globalPortId = outputPortId,

amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/Schedule.scala

@@ -19,8 +19,15 @@
 
 package org.apache.texera.amber.engine.architecture.scheduling
 
+import org.apache.texera.amber.core.executor.OpExecWithClassName
+import org.apache.texera.amber.operator.loop.{LoopEndOpDesc, LoopStartOpDesc}
+import org.apache.texera.amber.util.JSONUtils.objectMapper
+
 case class Schedule(private val levelSets: Map[Int, Set[Region]]) extends Iterator[Set[Region]] {
   private var currentLevel = levelSets.keys.minOption.getOrElse(0)
+  private var loopStartLevel = currentLevel
+  private var iteration = 10
+  private var i = 1
 
   def getRegions: List[Region] = levelSets.values.flatten.toList
 
@@ -31,4 +38,24 @@ case class Schedule(private val levelSets: Map[Int, Set[Region]]) extends Iterat
     currentLevel += 1
     regions
   }
+
+  def loopNext(): Set[Region] = {
+    val regions = levelSets(currentLevel)
+    if (
+      regions.exists(_.getOperators.exists(_.id.logicalOpId.id.startsWith("LoopStart-operator-")))
+    ) {
+      //iteration = objectMapper.readValue(regions.head.getOperators.head.opExecInitInfo.asInstanceOf[OpExecWithClassName].descString, classOf[LoopStartOpDesc]).iteration
+      loopStartLevel = currentLevel - 1
+    }
+    if (
+      regions.exists(_.getOperators.exists(_.id.logicalOpId.id.startsWith("LoopEnd-operator-")))
+    ) {
+      if (i < iteration) {
+        currentLevel = loopStartLevel
+        i += 1
+      }
+    }
+    currentLevel += 1
+    regions
+  }
 }

amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/DataProcessorRPCHandlerInitializer.scala

@@ -45,6 +45,7 @@ class DataProcessorRPCHandlerInitializer(val dp: DataProcessor)
     with ResumeHandler
     with StartHandler
     with EndHandler
+    with EndIterationHandler
     with StartChannelHandler
     with EndChannelHandler
     with AssignPortHandler

amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndIterationHandler.scala

@@ -0,0 +1,47 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+package org.apache.texera.amber.engine.architecture.worker.promisehandlers
+
+import com.twitter.util.Future
+import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{
+  AsyncRPCContext,
+  EmptyRequest,
+  EndIterationRequest
+}
+import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn
+import org.apache.texera.amber.engine.architecture.worker.DataProcessorRPCHandlerInitializer
+
+trait EndIterationHandler {
+  this: DataProcessorRPCHandlerInitializer =>
+
+  override def endIteration(
+      request: EndIterationRequest,
+      ctx: AsyncRPCContext
+  ): Future[EmptyReturn] = {
+    dp.executor match {
+      //case _: LoopEndOpExec =>
+      //workerInterface.nextIteration(EmptyRequest(), mkContext(request.worker))
+      case _ =>
+      //dp.processOnFinish()
+      //dp.outputManager.finalizeIteration(request.worker)
+    }
+    EmptyReturn()
+  }
+}