fix(pyamber): disable data sub-queues registered after disable_data - #6724
fix(pyamber): disable data sub-queues registered after disable_data#6724eugenegujing wants to merge 6 commits into
Conversation
Automated Reviewer SuggestionsBased on the
|
|
| config | throughput | MB/s | latency | max Δ latest / 7d | |
|---|---|---|---|---|---|
| 🔴 | bs=10 sw=10 sl=64 | 415 | 0.253 | 22,550/33,546/33,546 us | 🔴 +14.2% / 🔴 +109.9% |
| 🟢 | bs=100 sw=10 sl=64 | 930 | 0.568 | 103,797/129,933/129,933 us | 🟢 -18.1% / 🔴 +21.1% |
| ⚪ | bs=1000 sw=10 sl=64 | 1,057 | 0.645 | 947,248/1,015,828/1,015,828 us | ⚪ within ±5% / ⚪ within ±5% |
Baseline details
Latest main 133da7b from same runner
| config | metric | PR | latest main | 7d avg | Δ latest | Δ 7d |
|---|---|---|---|---|---|---|
| bs=10 sw=10 sl=64 | throughput | 415 tuples/sec | 468 tuples/sec | 779.07 tuples/sec | -11.3% | -46.7% |
| bs=10 sw=10 sl=64 | MB/s | 0.253 MB/s | 0.286 MB/s | 0.476 MB/s | -11.5% | -46.8% |
| bs=10 sw=10 sl=64 | p50 | 22,550 us | 19,751 us | 12,818 us | +14.2% | +75.9% |
| bs=10 sw=10 sl=64 | p95 | 33,546 us | 29,790 us | 15,986 us | +12.6% | +109.9% |
| bs=10 sw=10 sl=64 | p99 | 33,546 us | 29,790 us | 19,339 us | +12.6% | +73.5% |
| bs=100 sw=10 sl=64 | throughput | 930 tuples/sec | 939 tuples/sec | 1,011 tuples/sec | -1.0% | -8.0% |
| bs=100 sw=10 sl=64 | MB/s | 0.568 MB/s | 0.573 MB/s | 0.617 MB/s | -0.9% | -8.0% |
| bs=100 sw=10 sl=64 | p50 | 103,797 us | 104,450 us | 100,965 us | -0.6% | +2.8% |
| bs=100 sw=10 sl=64 | p95 | 129,933 us | 158,619 us | 107,295 us | -18.1% | +21.1% |
| bs=100 sw=10 sl=64 | p99 | 129,933 us | 158,619 us | 115,531 us | -18.1% | +12.5% |
| bs=1000 sw=10 sl=64 | throughput | 1,057 tuples/sec | 1,104 tuples/sec | 1,049 tuples/sec | -4.3% | +0.8% |
| bs=1000 sw=10 sl=64 | MB/s | 0.645 MB/s | 0.674 MB/s | 0.64 MB/s | -4.3% | +0.8% |
| bs=1000 sw=10 sl=64 | p50 | 947,248 us | 913,581 us | 978,248 us | +3.7% | -3.2% |
| bs=1000 sw=10 sl=64 | p95 | 1,015,828 us | 970,900 us | 1,021,881 us | +4.6% | -0.6% |
| bs=1000 sw=10 sl=64 | p99 | 1,015,828 us | 970,900 us | 1,050,075 us | +4.6% | -3.3% |
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,482.34,200,128000,415,0.253,22549.85,33546.06,33546.06
1,100,10,64,20,2149.59,2000,1280000,930,0.568,103796.72,129933.28,129933.28
2,1000,10,64,20,18924.06,20000,12800000,1057,0.645,947247.68,1015827.75,1015827.75
Codecov Report✅ All modified and coverable lines are covered by tests. Additional details and impacted files@@ Coverage Diff @@
## main #6724 +/- ##
============================================
+ Coverage 86.80% 88.76% +1.95%
- Complexity 4226 4429 +203
============================================
Files 1173 1175 +2
Lines 46865 47099 +234
Branches 5231 5268 +37
============================================
+ Hits 40682 41808 +1126
+ Misses 4457 3511 -946
- Partials 1726 1780 +54
*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:
|
|
Can you provide a real workflow that can reproduce this bug? |
There was a problem hiding this comment.
Pull request overview
Fixes a correctness gap in the Python worker’s InternalQueue pause/backpressure mechanism by ensuring that data sub-queues created lazily (on first put) during a disable window start in the disabled state, preventing data from leaking and being dequeued while the worker is paused/backpressured.
Changes:
- Update
InternalQueue.putto register new channels under the same lock used bydisable_data/enable_data, and to disable newly-created data sub-queues when any disable reason is active. - Remove unreachable/dead handling for
InternalMarkerinputand consolidate identical dispatch branches. - Add regression and concurrency tests covering late-registered channels, stacked disable reasons, control-channel behavior during disables, and race conditions between registration and disable/enable toggling.
Reviewed changes
Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.
| File | Description |
|---|---|
| amber/src/main/python/core/models/internal_queue.py | Makes lazy data-channel registration respect active disable windows by starting newly created data sub-queues disabled under the shared lock; removes dead InternalMarker path. |
| amber/src/test/python/core/models/test_internal_queue.py | Adds regression + stress tests to ensure late-created data channels don’t leak during pause/backpressure and that counts/dequeue behavior remain correct under concurrency. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Here's the main (screenshot). The workflow is scan → sort → Python UDF. Sort is blocking, so the UDF runs in a second region and reads its input from storage through the materialization reader threads. I clicked pause right when the second region started, before the readers sent anything. Their channel was registered after disable_data, so it came up enabled, and the paused worker went on to consume the whole input and call complete(), which threw InvalidTransitionException: Cannot transit from PAUSED to COMPLETED and killed the DP thread. The heartbeat thread is still alive, so resume never gets processed and the execution stays stuck in "Resuming" with no error shown anywhere. With the fix, the channel starts disabled and the same run finishes normally after resume. |
Backport auto-label reportThis
|
|
Thanks, I will check tonight |
Yicong-Huang
left a comment
Yicong-Huang
left a comment
There was a problem hiding this comment.
🟡 0 must-fix · 2 advisory · 1 polish — the fix itself is correct; both advisories are optional cleanup around it.
Correctness (1)
internal_queue.py:90— the other five_queue_idsiterators are still unlocked, same race window (advisory, see inline)
Conventions (1)
- Branch is 23 commits behind main; the newest CI run (2026-07-29) has 3
amber-integrationfailures on both runners — rebase and re-run to confirm they are pre-existing (advisory)
Polish: 1 quick touch-up (see inline comments).
Verification trace
Traced the new lock against the layer below. Lock ordering is one-directional — _lock then take_lock/put_lock, and no path takes _lock while holding either — so the added acquire cannot deadlock with disable_data/enable_data or with a blocked get(). The put outside the lock is still safe: SubQueue.put only increments total_count under put_lock when enabled, and SubQueue.disable takes put_lock+take_lock, so a first put racing a concurrent disable_data is serialized either way and the count stays exact. enable_data needing no change checks out — it re-reads _queue_ids at release time. The removed InternalMarker arm really was dead: it has no base class and no tag, and nothing puts one into this queue. The lock is also not on the hot path, as claimed: it is taken once per channel, while SubQueue.put/LinkedBlockingMultiQueue.get already lock on every single element.
Fix: when put() registers a new data channel while _queue_state is non-empty, the new sub-queue now starts disabled (done under the same lock as disable_data/enable_data, and before the first element is enqueued). Control channels are never disabled, and enable_data already re-enables such channels on resume. Also removes the unreachable InternalMarker branch in put() and merges the two identical dispatch branches. Adds 8 regression tests to test_internal_queue.py covering the late-channel leak, release on resume, stacked disables, control channels staying unblocked, existing baselines, and two threaded stress tests.
- The six category query methods (is_control_empty, is_data_empty, size_control, size_data, in_mem_size, is_data_enabled) now iterate a tuple() snapshot of _queue_ids instead of the live set, so a first-time put() registering a new channel concurrently can no longer raise "Set changed size during iteration" and kill the DP thread. - Also reword the put() comment to mention backpressure alongside pause. - Add 6 parametrized regression tests and a threaded stress test; all fail without the snapshot fix.
8e9713c to
9dbbec2
Compare
Yicong-Huang
left a comment
Yicong-Huang
left a comment
There was a problem hiding this comment.
🟡 3 resolved · 0 open · 3 new (3 new = 3 newly introduced · 0 late catches)
Both fixes check out, and all three of my earlier findings are verified addressed in the tree. Everything below is optional.
Simplifications (1)
internal_queue.py:161— the snapshot idiom is now repeated in six methods with its rationale above only one (advisory, see inline)
Conventions (1)
- Description: the
_queue_idssnapshot fix is a second, independent crash with 7 tests and goes unmentioned; the testing section still reads "8 regression tests" / "25 passed" (advisory)
Polish: 1 quick touch-up (see inline comments).
Verification trace
I checked the snapshot rather than taking the tuple(set) idiom on reputation: with a key type carrying Python-level __hash__/__eq__ (the shape ChannelIdentity has) and 200k concurrent adds, the live-set reader raised RuntimeError: Set changed size during iteration and two hammering tuple(s) readers raised nothing.
Coverage is complete — the six queries snapshot, and enable_data/disable_data iterate live only under _lock, which put now also takes before _queue_ids.add. No KeyError window: put inserts into _queue.sub_queues before _queue_ids, and remove_sub_queue has no production caller. The stale direction is the safe one, since _queue_ids only grows, so any(is_enabled(...)) can under-report but never over-report.
CI: the remaining amber-integration failure on both runners is ReconfigurationIntegrationSpec on Iceberg CatalogCommitConflicts. The branch is rebased onto current main and build / pyamber passes on 3.11/3.12/3.13, so I read it as pre-existing infra — which closes my round-1 rebase-and-re-run note.
|
@Yicong-Huang Hi Yicong, I think we need a second look on the CI read: the The tests pause before the slow source's first tuple, so the Fries reconfiguration ECM is the first-ever message on the source→udf data channel, which now registers born-disabled and swallows the marker. The |
Per review, the six category query methods now go through _control_queue_ids() / _data_queue_ids(), which take the tuple() snapshot and filter it, so the snapshot is enforced by the API instead of being a convention each query method has to remember.
The snapshot rationale comment moves to the helper docstring.
The helpers copy the live set first and filter the private copy, so no user-level code runs while the live set is being iterated.
Also reword a test comment per review ("Registers a key whose ...").
Yicong-Huang
left a comment
Yicong-Huang
left a comment
There was a problem hiding this comment.
🔴 3 resolved · 0 open · 2 new (2 new = 0 newly introduced · 2 late catches)
All three round-2 findings are verified addressed. But you were right about the CI, and the two findings below follow from that.
Correctness (2)
internal_queue.py:88— the born-disabled gate blocks ECMs too, so reconfiguring a paused Python UDF worker hangs (must-fix, see inline)test_internal_queue.py:417— the new matrix has no ECM-first registration case, so it cannot observe that regression (must-fix, see inline)
Verification trace
I re-read the CI logs instead of my own round-2 conclusion. The Iceberg CatalogCommitConflicts lines I took for the failure are WARN-level retries. The failing suite is ReconfigurationIntegrationSpec — "Tests: succeeded 39, failed 3" on both ubuntu-latest and macos-latest at 92f4dc3 — while the last 10 runs of the same workflow on main all succeeded. The same 3 failures were there in round 1, so this was my miss twice over.
Traced the mechanism end to end: network_receiver.py:110 tags an incoming ECM with the data channel; _send_ecm_to_data_channels sends only to not is_control channels; linked_blocking_multi_queue.py:223 serves only sub-queues where child.enabled; main_loop.py:222-233 spins on interruptible_get() while is_data_enabled() is False. A born-disabled channel holding an ECM is therefore unreachable, and TestUtils.scala:374-382's 30s await is what observes it.
The constraint is symmetric on the Scala side — NetworkInputGateway.tryPickControlChannel requires cid.isControl. So reconfigure-while-paused worked only through this leak, and #5915 made that dependency reliable with its 30-row/0.25s slow source, which is why the suite now fails deterministically.
The round-2 fixes check out: all six queries route through the snapshot helpers, and the only live _queue_ids iterations left (143, 151) sit inside enable_data/disable_data under _lock, which put now also takes.
…w channels - Registering a data channel disabled during a disable window also blocked ECMs, which ride data channels, so reconfiguring a paused worker never completed. - Channels now register enabled as on main and get() withholds instead: a DataElement taken while _queue_state is non-empty has its channel disabled and is pushed back to its sub-queue's head, for enable_data to release. - Adds put_first to LinkedBlockingMultiQueue for the push-back, reports is_data_enabled() as False while any disable reason is active, and rewrites the regression tests for the new semantics.
|
@Yicong-Huang PTAL! |
Yicong-Huang
left a comment
Yicong-Huang
left a comment
There was a problem hiding this comment.
🔴 2 resolved · 0 open · 4 new (4 new = 4 newly introduced · 0 late catches)
The rework is right, and both round-3 findings are verified fixed in the tree rather than taken on the reply. No correctness defect in the new code; the only must-fix is the title.
Conventions (1)
- Retitle →
fix(pyamber): withhold data dequeued during a disable window— the current title names the registration-time mechanism this round rejected, and it becomes the squash commit message (must-fix)
Simplifications (1)
linked_blocking_multi_queue.py:160— the newput_first/enqueue_firstget no tests in that file's own suite (advisory, see inline)
Polish: 2 quick touch-ups (see inline comments).
Verification trace
I re-derived the accounting across the withhold cycle instead of trusting the tests. get() decs total_count and the sub-queue count; disable() decs total_count by the remaining count; put_first incs the sub-queue count but skips total_count because the sub-queue is now disabled; enable() restores the full count. It balances for empty, single- and multi-element sub-queues, including when a concurrent disable_data beats the withhold's own disable to it. enqueue_first's if self.last is self.head restores the tail on exactly the state dequeue() leaves behind. Lock order composes too: _lock → put_lock → take_lock is what disable_data and the pre-existing remove() already use, and no path takes _lock while holding either inner lock.
The JVM worker already gates data consumption on worker state rather than only the per-channel bit. DPThread.scala:168 routes to tryPickControlChannel while pauseManager.isPaused, and that requires cid.isControl (NetworkInputGateway.scala:42). Python was the outlier. So the ECM-behind-data delay you documented is not a divergence — the JVM side refuses all data-channel traffic while paused, which leaves Python the more permissive of the two.
Round-3 findings: the born-disabled gate is gone (internal_queue.py:115-120 registers enabled), and build / amber-integration now passes on both runners, where ReconfigurationIntegrationSpec failed 3 tests on both at 92f4dc3. On the benchmark comment, I read it as runner noise: the only per-element cost added runs once per batch, and the bot's own 7d column swings +101% on the same p95 it flags at +14%.
|
Retracting my round-3 must-fix (#6724 (comment)). It was wrong and it cost you a full redesign — that is on us, not on your implementation. I read the spec's 30s timeout as an engine deadlock. It was not: "No ECM while paused" is the engine's actual semantics:
Worth noting either way: the current withhold releases an ECM only when it is a channel's first element. This spec is green because the slow source's first tuple lands after the pause; speed the source up and it goes red with no code change. Suggestion: go back to born-disabled — smaller, leaves val reconfigured = client.coordinatorInterface.reconfigureWorkflow(request, ()) // no await
val runningReached = stateReached(client, RUNNING) // register before the transition
Await.result(client.coordinatorInterface.resumeWorkflow(EmptyRequest(), ()), commandTimeout)
Await.result(runningReached, commandTimeout)
Await.result(reconfigured, commandTimeout)
Push back if you think pyamber should diverge from |
…ation await order This reverts the dequeue-side withhold (1a65167): per review, "no ECM while paused" is the engine's intended semantics, and the JVM DPThread already refuses all data-channel traffic while paused. A data sub-queue registered while _queue_state is non-empty starts disabled again, which closes the pause/backpressure leak; enable_data() re-enables such channels on resume, so an ECM arriving mid-pause is delayed until resume rather than dropped, and a new regression test pins that semantics. The 30s ReconfigurationIntegrationSpec timeouts were the harness, not the engine: TestUtils.shouldReconfigure awaited reconfigureWorkflow while still paused, an order production never uses because reconfigurations only take effect on resume. shouldReconfigure now dispatches the reconfiguration without awaiting, awaits the resume ack, and only then awaits the reconfiguration ack, which also removes the test's dependence on source timing. The resume ack alone establishes that resume took effect, since ResumeHandler completes it only after every worker acked; the engine emits no RUNNING state event to wait on.
|
@Yicong-Huang PTAL! |
Yicong-Huang
left a comment
Yicong-Huang
left a comment
There was a problem hiding this comment.
🟢 4 resolved · 0 open · 0 new (0 new = 0 newly introduced · 0 late catches)
Verification trace
Traced the registration race across every lock interleaving: a late data sub-queue is disabled before its first enqueue whenever a pause or backpressure reason is active, while control channels remain available. The reconfiguration harness now mirrors production by dispatching reconfiguration before resume and awaiting the worker resume acknowledgements before awaiting ECM completion; the JVM paused path likewise selects only control-tagged channels.
What changes were proposed in this PR?
This PR fixes two independent crash bugs in the Python worker's
InternalQueue, both of which silently kill a worker thread while the heartbeat thread stays alive, hanging the execution with no error reported. It also fixes the reconfiguration test harness, whose incorrect await order masqueraded as an engine deadlock.Fix 1: data leaking out of a paused or backpressured worker.
Pause and backpressure work by disabling the data sub-queues of
InternalQueue. However,disable_dataonly disables the sub-queues that exist when it is called, while sub-queues are created lazily on a channel's firstputand start enabled. A data channel whose first message arrives during a disable window therefore came up enabled and was not covered by the disable at all.This is a bug because the worker can then keep processing data while reporting PAUSED, and if the leaked
DataElementis dequeued inside the pause wait-loop inmain_loop.py, the control-only pampymatchraises an uncaughtMatchErrorthat silently kills the DP thread while the heartbeat thread stays alive, so the execution hangs forever with no error reported. Under backpressure, the leaked channel keeps feeding the congested downstream, defeating flow control.The fix in
InternalQueue.put: when it registers a new data channel while_queue_stateis non-empty, the new sub-queue now starts disabled. The registration is done under the same lock used bydisable_data/enable_data(with a double-check to avoid duplicate registration), and the sub-queue is disabled before its first element is enqueued, so the element never becomes dequeuable during the disable window. Control channels are never disabled, andenable_dataneeds no change because it iterates the channel set at release time, which by then includes channels registered mid-pause.sequenceDiagram participant R as Reader thread participant Q as InternalQueue participant DP as DP thread DP->>Q: disable_data(PAUSE) R->>Q: put(): first message on a new channel Note right of Q: _queue_state is not empty, so it starts DISABLED Note right of Q: messages wait in the queue DP->>Q: enable_data(PAUSE) on resume Q-->>DP: messages delivered in order, nothing lostAn important consequence, pinned by a regression test: ECMs ride data channels, so an ECM that is such a channel's first message (e.g. a reconfiguration marker reaching a paused worker) is delayed until resume, not dropped —
enable_datare-enables the channel and the ECM is delivered in order. This matches the engine's intended semantics: reconfigurations submitted while paused only take effect on resume (ExecutionReconfigurationServicedispatches them without awaiting), and the JVMDPThreadlikewise refuses all data-channel traffic, ECMs included, while paused.Harness fix:
TestUtils.shouldReconfigureawaited the reconfiguration ack while still paused.The three
ReconfigurationIntegrationSpecfailures that earlier looked like an engine deadlock were the harness deadlocking itself: it calledAwait.result(reconfigureWorkflow(...))before resuming, but the ack can only be produced after resume delivers the ECM — a circular wait that expired at the 30s command timeout. Production never uses this order. The harness now matches production: dispatch the reconfiguration without awaiting, await the resume ack (whichResumeHandlercompletes only once every worker has acknowledged), then await the reconfiguration ack. This also removes the tests' dependence on source timing. (There is no RUNNING state event to wait on instead: the engine only pushesExecutionStateUpdatefor PAUSED and terminal states.)An intermediate revision of this PR instead moved the leak fix to the dequeue side so ECMs could pass while paused; it was reverted (
87939215d) after review established that delivering ECMs during a pause diverges from the engine's semantics and that the harness order was the actual culprit.Fix 2: the category query methods iterated
_queue_idsunguarded.The six per-category query methods (
is_control_empty,is_data_empty,size_control,size_data,in_mem_size,is_data_enabled) iterated the live_queue_idsset, whichput()grows on a channel's first message from a network thread. CPython raisesRuntimeError: Set changed size during iterationwhen a set is mutated mid-iteration, so a query racing a first-time registration kills the calling thread — e.g. the DP thread pollingis_data_enabled()in the main loop — producing the same silent-hang failure mode as Fix 1 through an unrelated mechanism.These methods now go through two private helpers,
_control_queue_ids()and_data_queue_ids(), which take atuple()snapshot of the set and filter it, so the snapshot is part of the API rather than a convention each query method has to remember. This keeps the hot-path queries lock-free: the snapshot copy is a single uninterruptible C-level operation, and a snapshot at most misses a channel registered mid-call, which the next poll observes.Cleanup in the same file: the unreachable
InternalMarkerentry in the isinstance tuple is removed (InternalMarkerdoes not subclassInternalQueueElement, so it can never reach that branch) and the two identical dispatch branches are merged.Any related issues, documentation, discussions?
Fixes #6723. The regression tests extend the
InternalQueuespec added by #6444.Two pre-existing bugs found while tracing the reconfiguration failure will be filed separately, as neither is needed for this suite to pass:
_check_and_process_controlprocesses ECMs under a stalecurrent_input_channel_id, andPauseManager.resumere-enables ECM_PAUSE channels before checking_global_pauses(the ScalaPauseManagerhas the correct order).How was this PR tested?
amber/src/test/python/core/models/test_internal_queue.pypasses with 34 passed and 1 xfailed (the xfail documents a pre-existingLinkedBlockingMultiQueuepriority bug from #6444, unrelated to this PR).For Fix 1: the late-channel leak test (fails without the fix), release via
enable_data, stacked pause+backpressure disables, control channels staying unblocked mid-pause, pre-registered and normally-registered baselines, two threaded stress tests racing first-time registrations against disable/enable toggles with exact element counts, and a new test pinning the delayed-ECM semantics (an ECM as a mid-disable channel's first message is not dequeuable while a reason is active and is delivered afterenable_data, under both PAUSE and BACKPRESSURE). Deleting the born-disabled block turns 5 tests red, including both parametrized cases of the new ECM test.For Fix 2: 6 deterministic parametrized tests (a sentinel key interleaves a real first-time
put()into each query's iteration) plus a threaded stress test racing 800 first-time registrations against all six queries; all fail with theRuntimeErrorwhen the snapshot is reverted.For the harness fix, the flip is the evidence: with the engine byte-identical,
ReconfigurationIntegrationSpecfails 3 of 3 under the old await order and passes 3 of 3 on repeated clean local runs under the production order;ReconfigurationSpec, which sharesshouldReconfigure, passes 2 of 2.The wider
amber/src/test/python/coretree was run before and after with identical pass/fail sets apart from the new tests (the only failures need a live Iceberg catalog).scalafmtCheckandscalafix --checkpass on the touched Scala scope.Was this PR authored or co-authored using generative AI tooling?
Co-authored by: Claude Code (Claude Fable 5)