fix(amber): carry the loop StateFrame envelope through JVM operators#6661
Conversation
A loop with a Java/Scala built-in operator (e.g. Limit) in its body could never iterate: the Loop End worker crashed with "no loop-back state URI configured for LoopStart ''". Per-iteration loop state rides a StateFrame envelope (loop_counter + loop_start_id, materialized as their own columns next to content), and the Python worker carries it on every hop -- but the Scala side stripped it at every seam, so any JVM hop re-emitted states with the "no loop" defaults (0, ""). In a nested loop the zeroed counter also made the inner Loop End mis-consume the enclosing loop's state. Loop operators are Python-only, so a JVM operator only ever needs to CARRY the envelope through unchanged (the +1/-1 bookkeeping lives in the Python runtime). The columns already exist in the materialized State schema and on the Arrow flight wire; this threads them through the Scala side: - StateFrame gains loopCounter / loopStartId fields (defaults 0 / "", mirroring the Python StateFrame in core/models/payload.py). - State gains loopCounterFrom / loopStartIdFrom extractors to read the envelope columns back off a materialized/transported row. - InputPortMaterializationReaderThread rebuilds the envelope when replaying materialized states. - DataProcessor threads the incoming envelope through processInputState to OutputManager.emitState; emitState / saveStateToStorageIfNeeded / NetworkOutputBuffer.sendState stamp it onto the outgoing frame and the storage write. - PythonProxyClient / PythonProxyServer preserve the envelope across the Arrow flight bridge in both directions. Scala-originated states (start/end-channel handlers) keep the defaults -- correct, since a JVM operator can never be a Loop Start. Tests: unit specs pin each seam (StateSpec envelope round-trip, DataProcessorSpec pass-through equality, OutputManagerSpec emit stamping, NetworkOutputBufferSpec sendState stamping); two new LoopIntegrationSpec e2e cases run a single and a nested loop with a Scala Limit operator in the loop body (the single-loop case reproduces the reported failure verbatim).
Automated Reviewer SuggestionsBased on the
|
There was a problem hiding this comment.
Pull request overview
Fixes Amber loop execution when a JVM (Scala/Java) operator appears inside a loop body by preserving the per-iteration loop “envelope” (loop_counter, loop_start_id) end-to-end on the Scala side. This aligns Scala’s StateFrame handling with the existing materialized State schema / Arrow wire format so loop bookkeeping survives any JVM hop.
Changes:
- Extend Scala
StateFrameto carryloopCounter/loopStartId, and plumb them throughDataProcessor→OutputManager→NetworkOutputBufferas well as state materialization replay. - Add
State.loopCounterFrom/State.loopStartIdFromextractors and use them when rebuilding aStateFramefrom materialized tuples and from Arrow Flight state rows. - Add unit + integration tests covering envelope pass-through (including nested loops) with a Scala
Limitoperator in the loop body.
Reviewed changes
Copilot reviewed 13 out of 13 changed files in this pull request and generated no comments.
Show a summary per file
| File | Description |
|---|---|
| common/workflow-core/src/test/scala/org/apache/texera/amber/core/state/StateSpec.scala | Adds tests ensuring loop envelope columns round-trip via tuple extractors and default correctly. |
| common/workflow-core/src/main/scala/org/apache/texera/amber/core/state/State.scala | Adds loop-envelope extractors (loopCounterFrom, loopStartIdFrom) and clarifies schema/semantics. |
| amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/DataProcessorSpec.scala | Adds regression test that a JVM state pass-through preserves the exact incoming envelope. |
| amber/src/test/scala/org/apache/texera/amber/engine/architecture/sendsemantics/partitioners/NetworkOutputBufferSpec.scala | Tests that sendState stamps loop envelope onto the outgoing StateFrame. |
| amber/src/test/scala/org/apache/texera/amber/engine/architecture/messaginglayer/OutputManagerSpec.scala | Tests that emitState forwards the envelope onto emitted frames. |
| amber/src/test/integration/org/apache/texera/amber/engine/e2e/LoopIntegrationSpec.scala | Adds E2E coverage for loop + Scala operator (single + nested loop cases). |
| amber/src/main/scala/org/apache/texera/amber/engine/common/ambermessage/DataPayload.scala | Extends StateFrame to include loopCounter / loopStartId with defaults + documentation. |
| amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/managers/InputPortMaterializationReaderThread.scala | Rebuilds StateFrame envelope from materialized state rows during replay. |
| amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/DataProcessor.scala | Threads loop envelope through processInputState and preserves it on emission. |
| amber/src/main/scala/org/apache/texera/amber/engine/architecture/sendsemantics/partitioners/Partitioner.scala | Updates NetworkOutputBuffer.sendState to accept and forward envelope fields. |
| amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyServer.scala | Rebuilds envelope from Arrow state rows when proxying Python → JVM. |
| amber/src/main/scala/org/apache/texera/amber/engine/architecture/pythonworker/PythonProxyClient.scala | Writes full State row (content + loop columns) when proxying JVM → Python. |
| amber/src/main/scala/org/apache/texera/amber/engine/architecture/messaginglayer/OutputManager.scala | Updates emitState / storage write to carry envelope through network + materialization. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #6661 +/- ##
============================================
- Coverage 74.96% 74.90% -0.06%
- Complexity 3446 3447 +1
============================================
Files 1160 1160
Lines 45795 45797 +2
Branches 5070 5064 -6
============================================
- Hits 34328 34305 -23
- Misses 9824 9846 +22
- Partials 1643 1646 +3
*This pull request uses carry forward flags. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Two more LoopIntegrationSpec cases beyond Limit: - a runtime-compiled Java UDF (identity map) in the loop body -- the most literal "Java native operator" case and a different executor kind (OpExecWithCode vs OpExecWithClassName); - a chain of two JVM operators (Limit -> Sleep(0)) -- pins the JVM-to-JVM state handoff (one Scala worker writes the envelope columns, another reads them back), which single-op cases never exercise.
|
| config | throughput | MB/s | latency | max Δ latest / 7d | |
|---|---|---|---|---|---|
| 🔴 | bs=10 sw=10 sl=64 | 418 | 0.255 | 23,090/34,072/34,072 us | 🔴 +10.8% / 🔴 +115.5% |
| 🔴 | bs=100 sw=10 sl=64 | 926 | 0.565 | 108,724/129,509/129,509 us | 🔴 +5.1% / 🔴 +23.3% |
| 🔴 | bs=1000 sw=10 sl=64 | 1,086 | 0.663 | 913,350/1,051,991/1,051,991 us | 🔴 +13.1% / 🟢 -5.3% |
Baseline details
Latest main 964e915 from same runner
| config | metric | PR | latest main | 7d avg | Δ latest | Δ 7d |
|---|---|---|---|---|---|---|
| bs=10 sw=10 sl=64 | throughput | 418 tuples/sec | 453 tuples/sec | 769.76 tuples/sec | -7.7% | -45.7% |
| bs=10 sw=10 sl=64 | MB/s | 0.255 MB/s | 0.277 MB/s | 0.47 MB/s | -7.9% | -45.7% |
| bs=10 sw=10 sl=64 | p50 | 23,090 us | 20,835 us | 12,663 us | +10.8% | +82.4% |
| bs=10 sw=10 sl=64 | p95 | 34,072 us | 34,703 us | 15,813 us | -1.8% | +115.5% |
| bs=10 sw=10 sl=64 | p99 | 34,072 us | 34,703 us | 19,001 us | -1.8% | +79.3% |
| bs=100 sw=10 sl=64 | throughput | 926 tuples/sec | 957 tuples/sec | 1,009 tuples/sec | -3.2% | -8.3% |
| bs=100 sw=10 sl=64 | MB/s | 0.565 MB/s | 0.584 MB/s | 0.616 MB/s | -3.3% | -8.3% |
| bs=100 sw=10 sl=64 | p50 | 108,724 us | 103,489 us | 99,444 us | +5.1% | +9.3% |
| bs=100 sw=10 sl=64 | p95 | 129,509 us | 128,824 us | 105,012 us | +0.5% | +23.3% |
| bs=100 sw=10 sl=64 | p99 | 129,509 us | 128,824 us | 116,467 us | +0.5% | +11.2% |
| bs=1000 sw=10 sl=64 | throughput | 1,086 tuples/sec | 1,114 tuples/sec | 1,044 tuples/sec | -2.5% | +4.0% |
| bs=1000 sw=10 sl=64 | MB/s | 0.663 MB/s | 0.68 MB/s | 0.637 MB/s | -2.5% | +4.0% |
| bs=1000 sw=10 sl=64 | p50 | 913,350 us | 897,833 us | 964,034 us | +1.7% | -5.3% |
| bs=1000 sw=10 sl=64 | p95 | 1,051,991 us | 929,938 us | 1,007,497 us | +13.1% | +4.4% |
| bs=1000 sw=10 sl=64 | p99 | 1,051,991 us | 929,938 us | 1,040,936 us | +13.1% | +1.1% |
Raw CSV
config_idx,batch_size,schema_width,string_len,num_batches,total_ms,total_tuples,total_bytes,tuples_per_sec,mb_per_sec,lat_p50_us,lat_p95_us,lat_p99_us
0,10,10,64,20,478.55,200,128000,418,0.255,23090.32,34071.60,34071.60
1,100,10,64,20,2160.98,2000,1280000,926,0.565,108723.58,129509.43,129509.43
2,1000,10,64,20,18419.28,20000,12800000,1086,0.663,913349.65,1051991.47,1051991.47The JavaUDF case fails at executor setup in the integration harness: the
suite runs forkless under sbt, where javax.tools sees only the sbt launcher
on java.class.path, so JavaRuntimeCompilation cannot resolve the Texera
classes ("package ... does not exist") before any loop logic runs -- a
harness limitation (production workers launch with a real classpath),
documented in the spec so it is not re-tried.
Replaced with the strongest still-missing combination: a nested 3x3 loop
whose inner body chains Limit -> Sleep(0), covering nested (counter
magnitude) x multi-hop (JVM-to-JVM state handoff) x both operator kinds.
Xiao-zhen-Liu
left a comment
Xiao-zhen-Liu
left a comment
There was a problem hiding this comment.
Approving. I traced the fix end to end and it's correct for the path it targets: every place a StateFrame is built or a state is emitted on the Scala side now carries the loop id and counter, the materialization replay and the Arrow bridge (both directions) rebuild them, and the unit specs pin each one. Case 1 of the e2e reproduces the reported crash, and the specId collision from the #5700 review is fixed here (now 6, unique).
This strictly improves things — before this PR, any JVM operator in a loop body crashed; now the passthrough case works. One non-blocking follow-up inline: a possible remaining edge for stateful JVM operators, which I haven't been able to run down. It doesn't block the fix.
| * passes the incoming envelope through unchanged, while a Scala-originated | ||
| * state (start/end-channel handlers) uses the "no loop" defaults. | ||
| */ | ||
| def emitState(state: State, loopCounter: Long = 0L, loopStartId: String = ""): Unit = { |
There was a problem hiding this comment.
Non-blocking follow-up. These defaults (0 / "") are right for the forwarded-state path, which now carries the real loop id and counter. But two callers keep the defaults for operator-originated boundary state: StartChannelHandler:45 and EndChannelHandler:46 (produceStateOnStart / OnFinish).
For a stateless passthrough (Limit, Sleep — what the e2e cases use) that's fine, since they emit no boundary state. But a stateful JVM operator in a loop body would emit boundary state here with (0, ""), and LoopEnd captures its back-jump target from every state it consumes at counter 0 (main_loop.py:380, self._loop_start_id = frame.loop_start_id). On endChannel that boundary state arrives after the forwarded loop state, so it overwrites the captured id with "" — reproducing the same no loop-back state URI configured for LoopStart '' crash, through a path this fix doesn't cover.
I reasoned this from the code but haven't run it: worth either confirming a stateful JVM operator in a loop body works, or documenting that as unsupported. Related: because these args default silently, a future state path that misses them would drop the loop id and counter with no compile-time error — worth a short comment noting the two boundary-state callers are the deliberate default cases.
There was a problem hiding this comment.
Good catch — confirmed, and it's actually worse than the id clobber: the unstamped state also lands in run_update, which expects the loop table payload, so it raises KeyError before the jump even runs. It's language-independent too — a Python UDF's produce_state_on_finish hits the identical path (no JVM built-in overrides produceStateOnFinish today, so that's also the practical repro).
Fixed in #6913 on the consumer side: a real loop state is always stamped by its LoopStart, so LoopEnd now keys on the stamp — unstamped counter-0 states forward through unchanged (operator skipped, captured id untouched). Also added the comment you suggested at the two boundary-state emit sites. Includes a unit test plus an e2e case with a state-emitting UDF in the loop body.
…gration (apache#6884) ### What changes were proposed in this PR? `LoopIntegrationSpec` costs **~4m16s of the ~6-minute `amber-integration` test step (~70%)** on every PR. The cost tracks worker spawns almost perfectly (~2s per Python worker boot + Flight handshake; 132 spawns across the suite), because every loop iteration respawns each worker in the re-executed regions: | Test | Time | Worker spawns (py + jvm) | |---|---:|---:| | single loop (Python-only) | 17s | 6 + 1 | | nested 3×3 (Python-only) | 63s | 24 + 1 | | single + Limit | 19s | 6 + 4 | | nested + Limit | 64s | 24 + 10 | | nested + chain (Limit→Sleep) | 72s | 24 + 19 | | single + chain | 21s | 6 + 7 | Two cases are **strict subsets** of the nested-chain case, which already covers nested routing, counter magnitude across JVM hops, the JVM→JVM state handoff, and both operator kinds in one workflow: - `nested + Limit` (~64s) — nested × Limit ⊂ nested × (Limit → Sleep) - `single + chain` (~21s) — chain ⊂ nested chain This PR removes those two, saving **~85s (~33% of the loop suite)** with no unique coverage lost. Retained matrix: Python-only single (base), Python-only nested (pins the no-JVM path), single + Limit (verbatim repro of the apache#6660 bug), nested + chain (superset JVM case). The subsumption reasoning is recorded in the retained test's comment so the cases don't get re-added. The deeper fix (reusing workers across loop iterations instead of respawn-per-jump) is an engine design change and out of scope here. ### Any related issues, documentation, discussions? Follow-up to apache#6661 (which introduced the JVM loop-body cases). Related engine context: apache#6660. ### How was this PR tested? Test-only change (removes two e2e cases; no source touched). `WorkflowExecutionService/Test/compile` + `scalafmtCheckAll` + `scalafixAll --check` pass locally (Java 17); the remaining suite runs in the `amber-integration` CI job on this PR. ### Was this PR authored or co-authored using generative AI tooling? Generated-by: Claude Code (Fable 5)
What changes were proposed in this PR?
A loop with a Java/Scala built-in operator in its body could never iterate — the Loop End worker crashed on the first iteration:
Root cause. Per-iteration loop state rides a
StateFrameenvelope —loop_counter(nesting-depth countdown) andloop_start_id(the back-jump target) — materialized as their own columns next tocontent(#5900). The Python worker carries the envelope on every hop; the Scala side stripped it at every seam, so any JVM hop re-emitted states with the "no loop" defaults(0, ""):In a nested loop the zeroed
loop_counteradditionally makes the inner Loop End mis-consume the enclosing loop's state instead of passing it through.Fix. Loop operators are Python-only, so a JVM operator only ever needs to carry the envelope through unchanged (the +1/−1 bookkeeping lives in the Python runtime). The columns already exist in the materialized State schema and on the Arrow flight wire — no format change, purely Scala-side plumbing:
StateFrame(DataPayload.scala)loopCounter/loopStartIdfields (defaults0/"", mirroring Python'sStateFrame)State(State.scala)loopCounterFrom/loopStartIdFromextractors read the envelope columns back off a rowInputPortMaterializationReaderThreadDataProcessor→OutputManagerprocessInputStatetoemitState/saveStateToStorageIfNeeded/sendStatePythonProxyClient/PythonProxyServerScala-originated states (start/end-channel handlers) keep the defaults — correct, since a JVM operator can never be a Loop Start.
Size: the fix itself is +38 / −17 lines of code (excluding comments/blanks; +78 / −20 gross) across the 8 source files above; the remaining +255 lines are test cases (4 unit specs + 4 e2e workflows).
Any related issues, documentation, discussions?
Closes #6660
How was this PR tested?
StateSpec(envelope column round-trip + no-loop defaults),DataProcessorSpec(a state pass-through emits the exact incoming envelope, not just aStateFrame),OutputManagerSpec(emitStatestamps the envelope onto every buffer),NetworkOutputBufferSpec(sendStatestamps the frame).LoopIntegrationSpeccases run real workflows with JVM operators in the loop body:Limit(TextInput → LoopStart → Limit → LoopEnd, asserts 3 accumulated iterations — reproduces the reported failure verbatim, fails against the pre-fix code);Limitin the inner body (asserts 9 outer / 3 inner rows — pins that the counter magnitude survives the JVM hop, not just the id);Limit → Sleep(0)) — pins the JVM→JVM state handoff (one Scala worker writes the envelope columns, another reads them back), which the single-op cases never exercise;Manually tested with the this workflow: loop.json
Was this PR authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Fable 5)