From 852ba0c593c70d96fbced505e078cb7ea9128f46 Mon Sep 17 00:00:00 2001 From: Xinyuan Lin Date: Sun, 26 Jul 2026 13:51:56 -0700 Subject: [PATCH] fix(amber): pass unstamped boundary states through LoopEnd instead of consuming them Follow-up to #6661 (review discussion r3648708075). A loop-body operator that emits its own boundary state (produce_state_on_start/finish -- a public API on both engine sides) sends it with the "no loop" envelope (counter 0, loop_start_id ""). The LoopEnd matching branch treated EVERY counter-0 frame as the loop's own boundary state: it captured the back-jump id from the frame (clobbering the real id with "") and handed the state to run_update, which expects the loop `table` payload (KeyError). Either way the loop broke with "no loop-back state URI configured for LoopStart ''". A real loop state is always stamped -- the matching LoopStart stamps its own id on every iteration's output state -- so the consumer can key on the stamp: an UNstamped counter-0 frame at a LoopEnd is now forwarded downstream unchanged, skipping the operator, like any default pass-through. The captured back-jump id is untouched. Also adds the comment the review asked for at the two Scala boundary-state emit sites (StartChannelHandler / EndChannelHandler), documenting that their "no loop" envelope defaults are deliberate and what the Python LoopEnd runtime does with them. Tests: a unit test pins the forward/skip/no-clobber behavior for unstamped frames (the stamped-consume test is unchanged), and a new LoopIntegrationSpec e2e case runs a loop whose body is a Python UDF emitting boundary state via produce_state_on_finish -- it crashes without this fix and completes 3 iterations with it. --- .../main/python/core/runnables/main_loop.py | 20 ++++++-- .../promisehandlers/EndChannelHandler.scala | 6 +++ .../promisehandlers/StartChannelHandler.scala | 6 +++ .../engine/e2e/LoopIntegrationSpec.scala | 51 +++++++++++++++++++ .../python/core/runnables/test_main_loop.py | 38 ++++++++++++++ 5 files changed, 118 insertions(+), 3 deletions(-) diff --git a/amber/src/main/python/core/runnables/main_loop.py b/amber/src/main/python/core/runnables/main_loop.py index d252be934e6..c671a5f6214 100644 --- a/amber/src/main/python/core/runnables/main_loop.py +++ b/amber/src/main/python/core/runnables/main_loop.py @@ -374,9 +374,23 @@ def _process_state_frame(self, frame: StateFrame) -> None: return if isinstance(executor, LoopEndOperator): - # Matching LoopEnd (in_counter == 0): it will consume this state - # and jump back. Remember which LoopStart to jump to (it rides - # the envelope) for complete()/_jump_to_loop_start. + if not frame.loop_start_id: + # An UNstamped counter-0 state at a LoopEnd is not the loop's + # own boundary state -- it was produced by a loop-body + # operator's produce_state_on_start/finish (a public API on + # both engine sides), which emits with the "no loop" envelope. + # A real loop state is always stamped: the matching LoopStart + # stamps its own id on every iteration's output state. + # Forward it downstream unchanged, skipping the operator, like + # any default pass-through: consuming it would clobber the + # captured back-jump id with "" and hand run_update a State + # with no `table` payload (#discussion_r3648708075). + self._emit_and_save_state(state, in_counter, frame.loop_start_id) + self._check_and_process_control() + return + # Matching LoopEnd (in_counter == 0, stamped): it will consume + # this state and jump back. Remember which LoopStart to jump to + # (it rides the envelope) for complete()/_jump_to_loop_start. self._loop_start_id = frame.loop_start_id self.context.state_processing_manager.current_input_state = state diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndChannelHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndChannelHandler.scala index 7794342690b..4c41152e00f 100644 --- a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndChannelHandler.scala +++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndChannelHandler.scala @@ -43,6 +43,12 @@ trait EndChannelHandler { try { val outputState = dp.executor.produceStateOnFinish(portId.id) if (outputState.isDefined) { + // Deliberate "no loop" envelope defaults (loopCounter = 0, + // loopStartId = ""): this is operator-ORIGINATED boundary state, not + // a forwarded loop state, so it carries no LoopStart stamp. The + // Python LoopEnd runtime keys on that missing stamp to pass such + // states through instead of consuming them (see + // main_loop._process_state_frame). dp.outputManager.emitState(outputState.get) } dp.outputManager.outputIterator.setTupleOutput( diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/StartChannelHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/StartChannelHandler.scala index 01cbb858bd7..b23f084cf7a 100644 --- a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/StartChannelHandler.scala +++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/StartChannelHandler.scala @@ -42,6 +42,12 @@ trait StartChannelHandler { try { val outputState = dp.executor.produceStateOnStart(portId.id) if (outputState.isDefined) { + // Deliberate "no loop" envelope defaults (loopCounter = 0, + // loopStartId = ""): this is operator-ORIGINATED boundary state, not + // a forwarded loop state, so it carries no LoopStart stamp. The + // Python LoopEnd runtime keys on that missing stamp to pass such + // states through instead of consuming them (see + // main_loop._process_state_frame). dp.outputManager.emitState(outputState.get) } } catch safely { diff --git a/amber/src/test/integration/org/apache/texera/amber/engine/e2e/LoopIntegrationSpec.scala b/amber/src/test/integration/org/apache/texera/amber/engine/e2e/LoopIntegrationSpec.scala index 3dec510b903..61b3a6c4fce 100644 --- a/amber/src/test/integration/org/apache/texera/amber/engine/e2e/LoopIntegrationSpec.scala +++ b/amber/src/test/integration/org/apache/texera/amber/engine/e2e/LoopIntegrationSpec.scala @@ -44,6 +44,7 @@ import org.apache.texera.amber.operator.LogicalOp import org.apache.texera.amber.operator.limit.LimitOpDesc import org.apache.texera.amber.operator.loop.{LoopEndOpDesc, LoopStartOpDesc} import org.apache.texera.amber.operator.sleep.SleepOpDesc +import org.apache.texera.amber.operator.udf.python.PythonUDFOpDescV2 import org.apache.texera.amber.operator.source.scan.text.TextInputSourceOpDesc import org.apache.texera.amber.tags.IntegrationTest import org.apache.texera.workflow.LogicalLink @@ -184,6 +185,29 @@ class LoopIntegrationSpec // any loop logic runs. Production workers launch with a real classpath, // so this is a harness limitation, not an engine one. + /** A pass-through Python UDF that ALSO emits its own boundary state via the + * public `produce_state_on_finish` API -- the state arrives downstream + * with the "no loop" envelope (counter 0, no LoopStart stamp). + */ + private def statefulPythonUDF(): PythonUDFOpDescV2 = { + val op = new PythonUDFOpDescV2() + op.code = """from pytexera import * + | + |class ProcessTupleOperator(UDFOperatorV2): + | + | @overrides + | def process_tuple(self, tuple_: Tuple, port: int) -> Iterator[Optional[TupleLike]]: + | yield tuple_ + | + | @overrides + | def produce_state_on_finish(self, port: int) -> Optional[State]: + | return State({"note": "from-body-op"}) + |""".stripMargin + op.workers = 1 + op.retainInputColumns = true + op + } + private def link(from: LogicalOp, to: LogicalOp): LogicalLink = LogicalLink(from.operatorIdentifier, PortIdentity(), to.operatorIdentifier, PortIdentity()) @@ -331,4 +355,31 @@ class LoopIntegrationSpec ) } + it should "run a loop whose body operator emits its own boundary state" in { + // TextInput -> LoopStart -> stateful Python UDF -> LoopEnd. + // + // The UDF emits boundary state via produce_state_on_finish (a public API + // on both engine sides), which reaches the LoopEnd with the "no loop" + // envelope (counter 0, no LoopStart stamp) AFTER the forwarded loop + // state. The LoopEnd must pass that unstamped state through instead of + // consuming it: consuming would clobber the captured back-jump id with "" + // ("no loop-back state URI configured for LoopStart ''") and hand + // run_update a State with no `table` payload (KeyError). Regression test + // for #discussion_r3648708075 on #6661. + val src = textInput("1\n2\n3") + val start = loopStart("i = 0", "table.iloc[i]") + val mid = statefulPythonUDF() + val end = loopEnd("i += 1", "i < len(table)") + val materialized = runAndGetMaterializedRowCounts( + List(src, start, mid, end), + List(link(src, start), link(start, mid), link(mid, end)) + ) + val endRows = materialized.getOrElse(end.operatorIdentifier, -1L) + assert( + endRows == 3, + s"LoopEnd must accumulate all 3 iterations with a state-emitting " + + s"operator in the loop body: expected 3, got $endRows (all: $materialized)" + ) + } + } diff --git a/amber/src/test/python/core/runnables/test_main_loop.py b/amber/src/test/python/core/runnables/test_main_loop.py index 78ec5176353..7bc1957cc12 100644 --- a/amber/src/test/python/core/runnables/test_main_loop.py +++ b/amber/src/test/python/core/runnables/test_main_loop.py @@ -2495,6 +2495,44 @@ def test_loopend_consume_invokes_operator_at_counter_zero( assert "loop_start_id" not in passed_to_operator assert "loop_counter" not in passed_to_operator + def test_loopend_forwards_unstamped_state_without_consuming( + self, main_loop, monkeypatch + ): + # Reviewer feedback (#discussion_r3648708075): a loop-body operator + # that emits its own boundary state (produce_state_on_start/finish -- + # a public API on both engine sides) sends it with the "no loop" + # envelope (counter 0, id ""). That state is NOT the loop's own + # boundary state: consuming it would clobber the captured back-jump id + # with "" and hand run_update a State with no `table` payload + # (KeyError). A real loop state is always stamped -- the matching + # LoopStart stamps its own id on every iteration's output -- so an + # UNstamped counter-0 frame at a LoopEnd must be forwarded downstream + # unchanged, skipping the operator, like any default pass-through. + main_loop.context.executor_manager.executor = _FalseLoopEnd() + emitted, switched, reset_calls = self._capture_state_emit( + main_loop, monkeypatch + ) + # The loop's own state was already consumed and its id captured. + main_loop._loop_start_id = "loop-start-1" + + main_loop._process_state_frame( + StateFrame( + State({"note": "from-body-op"}), + loop_counter=0, + loop_start_id="", + ) + ) + + assert switched == [], "unstamped state must not invoke the operator" + assert reset_calls == [], "unstamped state must not reset output" + assert emitted == [(State({"note": "from-body-op"}), 0, "")], ( + "unstamped state must forward downstream with its envelope " + f"unchanged; emitted: {emitted}" + ) + assert main_loop._loop_start_id == "loop-start-1", ( + "an unstamped state must not clobber the captured back-jump id" + ) + # ------------------------------------------------------------------ # # _jump_to_loop_start #