fix(pyamber): stop EndWorker from consuming straggler messages

[aglinxinyuan]

What changes were proposed in this PR?

Root cause. The Python worker's EndWorkerHandler guard logs its "unprocessed messages" warning through input_queue.get() — a destructive, blocking read that removes a pending message — and then assert input_queue.is_empty(). With exactly one straggler message the straggler is silently dropped and EndWorker is acknowledged as success; with two or more, one message is still destroyed and the RPC fails with a bare AssertionError — and because the coordinator retries EndWorker, every retry that hits the guard eats another queued message until the last one is dropped under a success ack.

queue state at EndWorker	before	after
empty	ack	ack (unchanged)
1 straggler	straggler dropped, then ack	RPC fails, straggler kept
≥ 2 stragglers	1 dropped, then AssertionError	RPC fails, all kept

Fix. end_worker_handler.py reads the queued count once via input_queue.size() and branches on it, so nothing is consumed. When the queue is non-empty it logs the pending count and raises RuntimeError("worker still has unprocessed messages") instead of consuming a message and asserting. The raise rides the existing failure path — AsyncRPCServer.receive converts a handler exception into a ControlError reply, AsyncRPCClient.fulfillPromise on the coordinator turns it into a failed future, and RegionExecutionManager.terminateWorkersWithRetry re-sends EndWorker on a fixed killRetryDelay (bounded by maxTerminationAttempts), succeeding once the queue has drained — the exact Python analogue of the Scala EndHandler's Future.exception:

coordinator                        worker input queue at EndWorker arrival
    |                              [ReturnInvocation, ..., EndWorker]
    |-- EndWorker ---------------->|
    |<- ControlError --------------|   size() for the log; nothing consumed
    |   (retry after delay)
    |-- EndWorker ---------------->|   queue drained by the main loop
    |<- EmptyReturn ---------------|   safe to gracefulStop

The local is annotated as InternalQueue because size() lives on InternalQueue, not the base IQueue interface. No new inspection API (e.g. peek()) is added — InternalQueue stays non-peekable, like queue.Queue — so the change is confined to the handler plus its new test.

Any related issues, documentation, discussions?

Closes #6521

How was this PR tested?

TDD — the tests were written first and fail against the unfixed handler (with 1 straggler: no exception is raised and the message vanishes; with ≥ 2: AssertionError instead of a clean RPC failure):

• New src/test/python/core/architecture/handlers/control/test_end_worker_handler.py: acks on an empty queue; fails the RPC with one straggler; does not consume the straggler; keeps all messages with two stragglers; acks again once the queue drains (the retry protocol end-to-end). This is the Python analogue of the Scala EndHandlerSpec.

Ran locally: cd amber && pytest -m "not integration" on the touched test file plus the full unit suite, and ruff check src/main/python src/test/python && ruff format --check src/main/python src/test/python.

Was this PR authored or co-authored using generative AI tooling?

Generated-by: Claude Code (Fable 5)

[github-actions]

Automated Reviewer Suggestions

Based on the git blame history of the changed files, we recommend the following reviewers:

• Contributors with relevant context: @Yicong-Huang
  You can notify them by mentioning @Yicong-Huang in a comment.

[Copilot]

Pull request overview

This PR fixes a teardown race in the Python worker where EndWorkerHandler could destructively consume (“drop”) pending input-queue messages while checking for stragglers, causing incorrect successful acks and/or misleading failures during coordinator retries.

Changes:

• Make EndWorkerHandler log pending work via a non-destructive peek() and fail the RPC with a clear RuntimeError instead of consuming via get() and asserting.
• Add InternalQueue.peek() (delegating to LinkedBlockingMultiQueue.peek()) to support non-destructive inspection.
• Add focused unit tests covering peek() semantics and the EndWorker retry protocol behavior (fail while stragglers exist; succeed once drained).

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated 1 comment.

File	Description
amber/src/main/python/core/architecture/handlers/control/end_worker_handler.py	Switch EndWorker guard from destructive get()+assert to peek()+RuntimeError to preserve straggler messages and participate in retry semantics.
amber/src/main/python/core/models/internal_queue.py	Add InternalQueue.peek() API to expose non-destructive “next element” inspection.
amber/src/test/python/core/architecture/handlers/control/test_end_worker_handler.py	New tests validating EndWorker rejects stragglers without consuming them and succeeds after the queue drains.
amber/src/test/python/core/models/test_internal_queue.py	New tests validating peek() returns None on empty, is non-destructive, and matches the next get() across sub-queues.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 77.32%. Comparing base (429be11) to head (eacff8c).

Additional details and impacted files

@@             Coverage Diff              @@
##               main    #6522      +/-   ##
============================================
+ Coverage     77.31%   77.32%   +0.01%     
- Complexity     3524     3525       +1     
============================================
  Files          1161     1161              
  Lines         45920    45921       +1     
  Branches       5099     5099              
============================================
+ Hits          35501    35507       +6     
+ Misses         8839     8836       -3     
+ Partials       1580     1578       -2     

Flag	Coverage Δ		*Carryforward flag
access-control-service	70.00% <ø> (ø)		Carriedforward from 97f0fdb
agent-service	76.76% <ø> (ø)		Carriedforward from 97f0fdb
amber	69.11% <ø> (+0.01%)	⬆️	Carriedforward from 97f0fdb
computing-unit-managing-service	20.49% <ø> (ø)		Carriedforward from 97f0fdb
config-service	66.66% <ø> (ø)		Carriedforward from 97f0fdb
file-service	67.21% <ø> (ø)		Carriedforward from 97f0fdb
frontend	82.58% <ø> (ø)		Carriedforward from 97f0fdb
notebook-migration-service	78.94% <ø> (ø)		Carriedforward from 97f0fdb
pyamber	92.22% <100.00%> (+0.07%)	⬆️
workflow-compiling-service	55.14% <ø> (ø)		Carriedforward from 97f0fdb

*This pull request uses carry forward flags. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• 📦 JS Bundle Analysis: Save yourself from yourself by tracking and limiting bundle sizes in JS merges.

[github-actions]

⚠️ Benchmark changes need a look

🟢 2 better · 🔴 5 worse · ⚪ 8 noise (<±5%) · 0 without baseline

Compared against main 52584a8 benchmarked on this same runner, so the delta is largely free of cross-runner hardware noise. The "7d avg" column still reflects the gh-pages dashboard. Treat <±5% as noise unless repeated.

Dashboard · Run

	config	throughput	MB/s	latency	max Δ latest / 7d
🔴	bs=10 sw=10 sl=64	400	0.244	23,366/31,063/31,063 us	🔴 +16.0% / 🔴 +87.2%
🔴	bs=100 sw=10 sl=64	810	0.495	120,620/164,700/164,700 us	🔴 +27.8% / 🔴 +51.0%
⚪	bs=1000 sw=10 sl=64	927	0.566	1,081,108/1,125,436/1,125,436 us	⚪ within ±5% / 🔴 +7.9%

Baseline details

Latest main 52584a8 from same runner

config	metric	PR	latest main	7d avg	Δ latest	Δ 7d
bs=10 sw=10 sl=64	throughput	400 tuples/sec	431 tuples/sec	754.55 tuples/sec	-7.2%	-47.0%
bs=10 sw=10 sl=64	MB/s	0.244 MB/s	0.263 MB/s	0.461 MB/s	-7.2%	-47.0%
bs=10 sw=10 sl=64	p50	23,366 us	20,144 us	12,816 us	+16.0%	+82.3%
bs=10 sw=10 sl=64	p95	31,063 us	33,537 us	16,594 us	-7.4%	+87.2%
bs=10 sw=10 sl=64	p99	31,063 us	33,537 us	19,806 us	-7.4%	+56.8%
bs=100 sw=10 sl=64	throughput	810 tuples/sec	844 tuples/sec	969.38 tuples/sec	-4.0%	-16.4%
bs=100 sw=10 sl=64	MB/s	0.495 MB/s	0.515 MB/s	0.592 MB/s	-3.9%	-16.3%
bs=100 sw=10 sl=64	p50	120,620 us	117,601 us	103,584 us	+2.6%	+16.4%
bs=100 sw=10 sl=64	p95	164,700 us	128,921 us	109,097 us	+27.8%	+51.0%
bs=100 sw=10 sl=64	p99	164,700 us	128,921 us	117,304 us	+27.8%	+40.4%
bs=1000 sw=10 sl=64	throughput	927 tuples/sec	936 tuples/sec	1,004 tuples/sec	-1.0%	-7.6%
bs=1000 sw=10 sl=64	MB/s	0.566 MB/s	0.571 MB/s	0.613 MB/s	-0.9%	-7.6%
bs=1000 sw=10 sl=64	p50	1,081,108 us	1,072,500 us	1,002,357 us	+0.8%	+7.9%
bs=1000 sw=10 sl=64	p95	1,125,436 us	1,124,987 us	1,046,463 us	+0.0%	+7.5%
bs=1000 sw=10 sl=64	p99	1,125,436 us	1,124,987 us	1,073,661 us	+0.0%	+4.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,500.17,200,128000,400,0.244,23365.76,31062.73,31062.73
1,100,10,64,20,2467.66,2000,1280000,810,0.495,120620.48,164699.62,164699.62
2,1000,10,64,20,21565.47,20000,12800000,927,0.566,1081108.49,1125435.86,1125435.86

[Copilot]

Pull request overview

Copilot reviewed 2 out of 2 changed files in this pull request and generated 1 comment.

[Yicong-Huang]

🟡 0 must-fix · 1 advisory · 1 polish — the fix is correct, matches the Scala EndHandler, and is well-tested; both notes are minor.

Design & architecture (1)

• test_end_worker_handler.py:42 — fixture bypasses ControlHandler.__init__ via __new__ for no gain (advisory, see inline)

Polish: 1 quick touch-up (see inline comment).

Verification trace

Traced the failure path end to end: AsyncRPCServer.receive converts a handler exception into a ControlError reply, AsyncRPCClient.fulfillPromise turns it into a failed future, and RegionExecutionManager.terminateWorkersWithRetry re-sends EndWorker after killRetryDelay (bounded by maxTerminationAttempts) — so the raise has exactly the claimed effect and mirrors the Scala EndHandler's Future.exception. Also confirmed InternalQueue.size() and the old is_empty() read the same total_count (the guard's counting basis is unchanged — only the destructive get() is gone), and the tests use a real InternalQueue, so the non-consumption assertions genuinely observe the old bug.