refactor(amber): read the loop input table from its materialization instead of shipping it in state - #6971
refactor(amber): read the loop input table from its materialization instead of shipping it in state#6971aglinxinyuan wants to merge 12 commits into
Conversation
…nstead of shipping it in state The loop's input table used to ride INSIDE the State content: LoopStart encoded its buffered input as Arrow IPC bytes, base64'd inside the JSON content, and that payload was re-written and re-read at every loop-body hop, every iteration (~33% base64 bloat, JSON-column size limits, and a second, divergent Arrow codec via pa.Table.from_pandas inference). That data already exists: in the fully-materialized mode loops require, the Loop Start's input-port materialization holds exactly the loop's input table for the whole loop (Loop Start re-reads it every iteration; the back-edge truncates only the state doc at the same base URI, never the result doc). So ship the port's BASE URI in the setup config and derive both addresses from it: - proto: InitializeExecutorRequest.loopStartStateUris -> loopStartPortUris (field 4 unchanged); the value is now the base URI. WorkflowExecutionManager ships storagePairs.head._1 directly. - back-edge write: main_loop derives state_uri(base) (unchanged behavior, one derivation later). - table read: on the matching consume, MainLoop reads result_uri(base) and injects the table into the LoopEnd via a new runtime-only attach_loop_table hook; a read failure is reported like a UDF error (report_exception) instead of silently dropping the consume. - LoopStartOperator.produce_state_on_finish emits only the user loop variables (the reserved-`table` collision raise stays: a user var named `table` would now be silently shadowed by the injected table); run_update uses the attached table; the "consumed" marker (_loop_table) is still set only by a successful update, so condition()'s short-circuit semantics are unchanged. - table_to_ipc_bytes / table_from_ipc_bytes and the now-unused TableOperator._buffered_table accessor are deleted. Nested loops keep working by construction: the inner Loop Start's map entry points at the outer Loop Start's output port, whose result doc is recreated per OUTER iteration but persists across inner iterations (the jump rewinds to the inner level only). Tests updated across test_loop_operators (attach-based flow, plus new pins: attach alone does not mark the loop consumed; run_update fails loud when no table was attached; produced state carries no table), test_main_loop (base-URI derivation for the back-edge write, consume injects the read table), test_initialize_executor_handler (renamed field), and the Scala ctor sites. The four LoopIntegrationSpec e2e cases exercise the full path in CI.
Automated Reviewer SuggestionsBased on the
|
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## main #6971 +/- ##
============================================
+ Coverage 82.64% 82.68% +0.04%
- Complexity 4080 4086 +6
============================================
Files 1162 1162
Lines 46282 46314 +32
Branches 5160 5160
============================================
+ Hits 38250 38296 +46
+ Misses 6341 6325 -16
- Partials 1691 1693 +2
*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:
|
|
| config | throughput | MB/s | latency | max Δ latest / 7d | |
|---|---|---|---|---|---|
| 🔴 | bs=10 sw=10 sl=64 | 503 | 0.307 | 19,399/27,792/27,792 us | 🔴 +7.8% / 🔴 +80.6% |
| 🟢 | bs=100 sw=10 sl=64 | 1,231 | 0.751 | 80,422/90,526/90,526 us | 🟢 -15.1% / 🟢 +22.0% |
| ⚪ | bs=1000 sw=10 sl=64 | 1,393 | 0.85 | 721,207/754,560/754,560 us | ⚪ within ±5% / 🟢 +34.3% |
Baseline details
Latest main 436b37e from same runner
| config | metric | PR | latest main | 7d avg | Δ latest | Δ 7d |
|---|---|---|---|---|---|---|
| bs=10 sw=10 sl=64 | throughput | 503 tuples/sec | 532 tuples/sec | 790.88 tuples/sec | -5.5% | -36.4% |
| bs=10 sw=10 sl=64 | MB/s | 0.307 MB/s | 0.325 MB/s | 0.483 MB/s | -5.5% | -36.4% |
| bs=10 sw=10 sl=64 | p50 | 19,399 us | 17,997 us | 12,348 us | +7.8% | +57.1% |
| bs=10 sw=10 sl=64 | p95 | 27,792 us | 27,614 us | 15,390 us | +0.6% | +80.6% |
| bs=10 sw=10 sl=64 | p99 | 27,792 us | 27,614 us | 18,935 us | +0.6% | +46.8% |
| bs=100 sw=10 sl=64 | throughput | 1,231 tuples/sec | 1,175 tuples/sec | 1,009 tuples/sec | +4.8% | +22.0% |
| bs=100 sw=10 sl=64 | MB/s | 0.751 MB/s | 0.717 MB/s | 0.616 MB/s | +4.7% | +22.0% |
| bs=100 sw=10 sl=64 | p50 | 80,422 us | 83,288 us | 99,753 us | -3.4% | -19.4% |
| bs=100 sw=10 sl=64 | p95 | 90,526 us | 106,570 us | 105,895 us | -15.1% | -14.5% |
| bs=100 sw=10 sl=64 | p99 | 90,526 us | 106,570 us | 113,263 us | -15.1% | -20.1% |
| bs=1000 sw=10 sl=64 | throughput | 1,393 tuples/sec | 1,409 tuples/sec | 1,037 tuples/sec | -1.1% | +34.3% |
| bs=1000 sw=10 sl=64 | MB/s | 0.85 MB/s | 0.86 MB/s | 0.633 MB/s | -1.2% | +34.3% |
| bs=1000 sw=10 sl=64 | p50 | 721,207 us | 711,217 us | 976,194 us | +1.4% | -26.1% |
| bs=1000 sw=10 sl=64 | p95 | 754,560 us | 759,432 us | 1,016,411 us | -0.6% | -25.8% |
| bs=1000 sw=10 sl=64 | p99 | 754,560 us | 759,432 us | 1,045,212 us | -0.6% | -27.8% |
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,397.78,200,128000,503,0.307,19399.24,27792.16,27792.16
1,100,10,64,20,1624.72,2000,1280000,1231,0.751,80422.07,90526.15,90526.15
2,1000,10,64,20,14356.72,20000,12800000,1393,0.850,721206.84,754559.59,754559.59…lap the reader CI experiment for the amber-integration failure reported in review: the Access Denied errors are NOT in the new read path (0 of 56 tracebacks contain _read_loop_input_table) and never on the doc it opens. They are in this worker's own materialization READER thread, on loop-INTERNAL docs (LoopStart / Limit outputs) that region re-execution drops and recreates while IcebergDocument.get() iterates a lazily-pinned snapshot -- MinIO answers a deleted key with 403 Access Denied. The read itself is fine (the Loop Start's input-port materialization is created once upstream of the loop and is stable for the whole run), but issuing it at consume time puts a second iceberg/S3 read on the main loop thread WHILE that reader is still streaming. Move it out of that window: stash the matching state at consume and run the operator's update (with the table attached) at EndChannel, by which point the reader has finished. If this turns amber-integration green, the overlap was the trigger; if the Access Denied errors persist, the recreate-vs-active-reader race is an engine bug independent of this PR. The matching consume emits no state downstream, so deferring it changes nothing observable outside the operator.
Xiao-zhen-Liu
left a comment
There was a problem hiding this comment.
The trade is a good one and the split is right: the proto rename, the Scala plumbing, and the Python read all have to land together, and there's nothing in here that doesn't belong. Deleting the second Arrow codec is worth the PR on its own. CI is green on both integration jobs, and the old table_to_ipc_bytes path has no callers left.
Requesting changes on one thing: moving the consume into complete() puts a user-code failure after every output port has been reported complete, which is the case the guard a few lines up was written to prevent. Details inline. If there's a reason it can't move earlier, say so in a comment and add the failure test and I'm happy -- the rest of my notes are small.
Two things without a line to sit on:
The second commit is written as an experiment, not a fix. "If this turns amber-integration green, the overlap was the trigger; if the Access Denied errors persist, the recreate-vs-active-reader race is an engine bug independent of this PR." That's an honest thing to write while you're iterating, but it's now the shipping design -- squash it into the first commit, or rewrite the message so the deferral reads as a decision with a reason. Separately, the race it points at (a region re-execution dropping a doc while a reader iterates a pinned snapshot) deserves its own issue whether or not it's this PR's to fix. Right now the only record of it is a commit message that will be hard to find later.
Naming. loopStartPortUris holds the upstream output port's base URI -- the one the Loop Start's input port reads from. The proto comment and the Scala doc both say "base URI of that Loop Start's input port materialization", which sends a reader looking for a config keyed by the Loop Start's own port. cfg.storagePairs.head._1 in WorkflowExecutionManager is where the distinction is visible. One clause fixes it in both places.
There was a problem hiding this comment.
Pull request overview
This PR refactors Amber’s loop runtime to stop embedding the loop input table inside per-iteration State content. Instead, it passes a LoopStart input-port base materialization URI at setup time (loopStartPortUris) and derives both the loop-back state URI (write) and the loop input-table result URI (read) from that base, reducing state bloat and repeated I/O across loop-body hops.
Changes:
- Rename the setup/config plumbing from
loopStartStateUris→loopStartPortUris(proto + Scala scheduler + Scala/Python worker wiring), changing the value from “state URI” to “base URI”. - Update Python loop execution to stash the matching LoopEnd state at consume time and defer the table read + operator update to EndChannel, using a new
LoopEndOperator.attach_loop_table(...)hook. - Remove the loop-specific Arrow IPC table codec helpers (
table_to_ipc_bytes/table_from_ipc_bytes) and adjust unit/integration tests to match the new flow.
Reviewed changes
Copilot reviewed 19 out of 19 changed files in this pull request and generated 2 comments.
Show a summary per file
| File | Description |
|---|---|
| common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/loop/LoopStartOpDescSpec.scala | Updates test comments to the renamed setup field. |
| common/workflow-operator/src/test/scala/org/apache/texera/amber/operator/loop/LoopEndOpDescSpec.scala | Updates test comments to reflect attach-based table injection. |
| common/workflow-core/src/main/scala/org/apache/texera/amber/core/workflow/PhysicalOp.scala | Updates isLoopStart doc to reference the new base-URI semantics. |
| amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/WorkflowWorkerSpec.scala | Renames init request field usage in tests. |
| amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/WorkerSpec.scala | Renames init request field usage in tests. |
| amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/managers/SerializationManagerSpec.scala | Renames init request field usage in tests. |
| amber/src/test/python/core/runnables/test_main_loop.py | Rewrites LoopEnd matching-consume test to assert defer-to-EndChannel + base-URI derivation. |
| amber/src/test/python/core/models/test_loop_operators.py | Rewrites loop operator tests for attach-based table injection and removal of table-in-state. |
| amber/src/test/python/core/architecture/handlers/control/test_initialize_executor_handler.py | Updates handler tests for the renamed loop_start_port_uris context field. |
| amber/src/test/integration/org/apache/texera/amber/engine/e2e/LoopIntegrationSpec.scala | Updates integration-spec comments to the new setup field name. |
| amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/WorkflowExecutionManager.scala | Computes and ships LoopStart base port URIs (not state URIs) at setup. |
| amber/src/main/scala/org/apache/texera/amber/engine/architecture/scheduling/RegionExecutionManager.scala | Threads loopStartPortUris through to worker initialization. |
| amber/src/main/python/core/runnables/main_loop.py | Implements deferred LoopEnd consume and storage read of loop input table via derived result URI. |
| amber/src/main/python/core/models/table.py | Removes Arrow IPC serialization helpers for loop table transport. |
| amber/src/main/python/core/models/operator.py | Removes LoopStart table-in-state emission; adds attach_loop_table and updates LoopEnd update semantics. |
| amber/src/main/python/core/architecture/managers/context.py | Renames context storage from loop_start_state_uris to loop_start_port_uris. |
| amber/src/main/python/core/architecture/handlers/control/initialize_executor_handler.py | Stores loop_start_port_uris from the renamed proto field onto the runtime context. |
| amber/src/main/protobuf/org/apache/texera/amber/engine/architecture/rpc/controlcommands.proto | Renames proto field and documents base-URI derivation of state/result URIs. |
| amber/src/bench/scala/org/apache/texera/amber/bench/ArrowFlightActorBench.scala | Updates benchmark init request field usage to loopStartPortUris. |
Suppressed comments (1)
amber/src/test/python/core/architecture/handlers/control/test_initialize_executor_handler.py:67
- This test comment still describes
loop_start_port_urisas containing loop-back state URIs, but after this PR it holds the Loop Start input-port base URI (with state/result URIs derived from it). Updating the comment will keep the test aligned with the new semantics and avoid confusion for future maintainers.
# The loop-back write addresses (LoopStart op id -> its input port's
# state URI) are per-operator setup config delivered on this RPC; the
# handler must expose them on the context for a Loop End's
# _jump_to_loop_start to select by the frame-carried loop_start_id.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
… complete Review feedback on apache#6971: - The consume moved out of complete() into _process_end_channel, ahead of the existing has_exception() hold. complete() is the tail of that method, by which point port_completed has been sent for every output port, so a failing read or user `update` was reported after the coordinator already considered the region done -- a reported error that still reads as success, which is exactly what that hold exists to prevent. It is still past the reader: EndChannel means the reader finished streaming. This also removes the dependency on complete() running at all (the is_missing_output_ports early return skips it). - The `update` and `condition()` now run under replace_print, so a print() in either reaches the console instead of the worker's stdout. condition() had this gap already; `update` is where prints actually go. - run_update now takes _attached_table (clears it) instead of leaving it set, so the missing-table guard fires on every iteration rather than only the first, and a stale table can never be silently reused. Documented why it stays separate from _loop_table (which doubles as the "consumed" marker). - The stash asserts it is empty, making the one-matching-state-per-execution assumption (Loop End is non-parallelizable) checkable rather than implicit. - Corrected the read comment: the upstream doc is stable for the duration of the loop level that reads it, not "for the whole run" -- for an inner loop the upstream is the outer Loop Start's output, recreated per outer iteration. Tests: the read URI (result_uri of the configured base, via open_document) and its fail-loud path are now pinned, plus a failing deferred consume that asserts no port is reported complete and the worker does not complete, and a print in the user's update reaching the console.
complete() flushes on entry, before condition() is evaluated, and then shuts the worker down -- so a print() inside the user's condition was captured by the print-capture session and never sent. Flush again after the condition and jump (the error path already flushes via _check_exception). Test pins it. Raised by Copilot on apache#6971.
Conflict in RegionExecutionManager's constructor parameter list: main reworked the region-termination retry budget (killRetryDelay -> killRetryBaseBackoffMs + killRetryTimer, apache#6960) in the same lines where this branch renamed loopStartStateUris -> loopStartPortUris. Kept both: main's retry parameters and this branch's loop bookkeeping parameter.
Conflict in main_loop.py with apache#7154, which added the Loop End fan-in dedup (_loop_state_consumed) on the same consume branch this PR turns into a stash. Kept both, with the dedup guard AHEAD of the stash: a duplicate branch copy must not overwrite the pending state, not just skip a second update. Updated apache#7154's dedup test accordingly -- it asserted the operator was invoked once inside _process_state_frame, which no longer runs the update there. It now asserts through the deferred path: the duplicate neither overwrites the stash nor produces a second update at EndChannel.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 19 out of 19 changed files in this pull request and generated no new comments.
Suppressed comments (1)
amber/src/main/python/core/runnables/main_loop.py:142
_read_loop_input_table()callsDocumentFactory.open_document(result_uri)directly; if the Loop Start’s result materialization is missing or the URI is malformed, the raisedValueError("No storage is found for the given URI")doesn’t include loop context (loopStart id/base/result URI), making misconfiguration/debugging harder. Consider wrapping the open/read with a contextual error that includes the derivedresult_uriandself._loop_start_id.
result_uri = VFSURIFactory.result_uri(self._loop_start_base_uri())
document, _ = DocumentFactory.open_document(result_uri)
return Table(list(document.get()))
Review feedback on apache#6913. Forwarding an unstamped counter-0 state at a Loop End is right for a body operator's own boundary state, but it also swallowed the symptom of the bug class apache#6660/apache#6661 fixed: a hop that blanks the envelope makes the loop's OWN state arrive unstamped, and forwarding it leaves _loop_table None, so condition() returns False and the loop stops after one iteration -- a wrong result reported as success. The two cases are now told apart at completion instead of on arrival: a Loop End that forwarded an unstamped state and never took a stamped one never received its own state, so it raises. Deciding at completion is order-independent (the body operator's state may arrive before or after the loop's) and reads no State content, so it keeps working once the input table stops riding inside the state (apache#6971). Also from the review: document why a Loop START must MERGE an unstamped counter-0 state rather than forward it -- the back-edge writes the next iteration's variables with that same "no loop" envelope, so a Loop Start cannot tell them apart, while a Loop End can -- and pin it with a test; shorten the duplicated Start/EndChannelHandler comment to a pointer; drop the bare discussion anchor from production code; sort the new import.
Xiao-zhen-Liu
left a comment
There was a problem hiding this comment.
Approving. The blocker is fixed: moving the consume into _process_end_channel, ahead of the has_exception() hold, is the right shape, and test_end_channel_holds_the_region_when_the_deferred_consume_fails genuinely pins it -- it fails if you delete the consume, delete the hold, or move the consume back into complete().
One thing I'd want in before merge, on the condition() print wrap -- it's a one-liner and doesn't need another review round, so I'm not holding the PR for it. Details inline, with a repro.
Nice catch on the branching-body duplicates, and on condition()'s prints being captured and then never flushed -- I'd missed that complete()'s opening flush runs before the condition does.
The rest of my comments are inline and none of them block. One thing without a line to sit on:
Commit history. 822078ca7f ("if this turns amber-integration green, the overlap was the trigger...") is still in the branch, now followed by d8ccbe7887 correcting it. Squash before merge so the log reads as a decision rather than the debugging session that produced it. The recreate-vs-active-reader race that commit points at still deserves its own issue -- right now the only record of it is a commit message nobody will find.
On what I actually checked, so you can weigh the comments accordingly: the KeyError repro, replace_print.py having no changed lines here, run_update's namespace seeding, and the test-bypass are things I ran or read directly. The nested-loop behavior and the worker-recreated-per-region-execution chain I only traced through code -- I didn't run an end-to-end loop, so treat anything I say about those as reading, not evidence.
…fixes Review feedback on apache#6971. The bug: replace_print's wrapped print resolved the calling module via f_globals["__name__"], but a Loop End's condition/update (and a Loop Start's output) run through eval/exec against a bare namespace dict -- eval/exec inject __builtins__ but never __name__ -- so a print() inside any of those expressions raised KeyError: '__name__'. Pre-existing on main for update/output; this PR newly routed condition() through replace_print, turning a lost print into a loop that errors after one iteration. Fixed with .get('__name__', '<unknown>'), and the two print- capture tests now drive the GENERATED operator shape (eval_condition / run_update) instead of plain methods, which is the path that was silently untested -- both go red without the fix. Also from the review: - _process_end_channel: hold the region on a pre-existing error BEFORE the deferred consume, so a failed produce_state_on_finish is not followed by a storage read and a user update stacking a second error. - _consume_pending_loop_state: re-arm the fan-in dedup flag when the consume takes the stash (nothing more can arrive on the port after the aligned EndChannel), making it per-execution by construction instead of relying on the scheduler recreating workers per iteration. - _read_loop_input_table: cast every tuple to the doc's schema, exactly like the input-port reader that streams this same doc. - Correct the fan-in dedup rationale (run_update seeds from the incoming copy, so re-running an identical copy is idempotent; the flag pins first-wins for the stash) and state that rewriting the loop state in a branching body is not supported. - Document the single-input-port assumption the EndChannel-time read leans on (PORT_ALIGNMENT is per port; a second input port would silently reintroduce the reader/read overlap). - Reframe the Access Denied diagnosis as the working hypothesis, and fix the stranded "at / consume / time" line. - operator.py: the _attached_table clear is defensive (one instance handles one iteration), not an every-iteration guard; pinned with an update-without-attach raise in the two-iteration test. - proto + WorkflowExecutionManager doc: the configured value is the UPSTREAM output port's base URI (what the Loop Start's input port reads from), not a port of the Loop Start itself.
What changes were proposed in this PR?
The loop's input table used to ride inside the State content: LoopStart encoded its buffered input as Arrow IPC bytes, base64'd into the JSON
contentcolumn, and that payload was re-written and re-read at every loop-body hop, every iteration.That data already exists. In the fully-materialized mode loops require, the Loop Start's input-port materialization holds exactly the loop's input table for the whole loop — Loop Start re-reads it every iteration, and the back-edge truncates only the state doc at the same base URI, never the result doc. So this PR ships the port's base URI in the setup config and derives both addresses from it:
loopStartStateUris= state URIloopStartPortUris= base URI (renamed so the semantic change is loud)result_uri(base)at EndChannel, injected via a new runtime-onlyattach_loop_tablehooktable_to_ipc_bytes/table_from_ipc_bytesfrom_pandasinference)When the read happens matters. The read is issued at EndChannel (the matching state is stashed at consume and the operator's update runs in
complete()), not at consume time. At consume time this worker's own materialization reader is still streaming, and issuing a second iceberg/S3 read from the main loop thread in that window made the reader fail with S3Access Denied—LoopIntegrationSpechung to the CI job timeout on both OSes. Deferring past the reader removes the overlap: integration went from a 20-minute cancel to green in ~9 minutes. The matching consume emits no state downstream, so moving it is unobservable outside the operator.Semantics deliberately preserved:
tablecollision raise stays (a user var namedtablewould now be silently shadowed by the injected table — worse than before);_loop_table) is still set only by a successfulrun_update, socondition()'s short-circuit for pass-through-only Loop Ends is unchanged;Wins: no ~33% base64 bloat, no JSON-column size ceiling on the table (large-table loops become viable), strictly less I/O for any non-empty loop body (one read per iteration replaces N state-doc writes+reads per hop), and one Arrow codec instead of two.
Note: this deepens the read-side use of
storagePairs.head._1— the same shared upstream URI as the known back-edge fan-out design discussion; if that ever moves to a per-loop private doc, this read moves with it.Any related issues, documentation, discussions?
Builds on #5900 (State columns) and #6661 (envelope through JVM hops). Related design context: #6660.
How was this PR tested?
test_loop_operators.pyrewritten for the attach-based flow plus new pins: the produced state carries notable; attaching alone does not mark the loop consumed;run_updatefails loud when no table was attached.test_main_loop.pypins the base-URI derivation for the back-edge write (state_uri(base)), the missing-config fail-loud, and that the matching consume stashes the state without touching storage, with the read + update happening once at EndChannel.test_initialize_executor_handler.pycovers the renamed proto field. 237 tests green locally (the only failures in a full sweep are pre-existing environment ones, identical on unmodified main).scalafmtCheckAll,scalafixAll --check, and the worker/descriptor spec suites (WorkerSpec,WorkflowWorkerSpec,SerializationManagerSpec,WorkflowExecutionManagerSpec,LoopStartOpDescSpec,LoopEndOpDescSpec) all pass on Java 17.LoopIntegrationSpeccases (single, nested 3×3, JVM chain, nested JVM chain) exercise the full read path in theamber-integrationCI job (both jobs green in ~9 min); the nested cases specifically cover the inner-loop read against the outer Loop Start's per-outer-iteration output doc.Was this PR authored or co-authored using generative AI tooling?
Generated-by: Claude Code (Fable 5)