This is an automated email from the ASF dual-hosted git repository. github-merge-queue[bot] pushed a commit to branch gh-readonly-queue/main/pr-6522-81d4612dc796ed3ea5dc6d98bd5bcc5f7fce0b41 in repository https://gitbox.apache.org/repos/asf/texera.git
commit 7e4a9b46b863fc13fa76ba74980a0f0c4beff48f Author: Xinyuan Lin <[email protected]> AuthorDate: Sat Jul 25 00:55:12 2026 -0700 fix(pyamber): stop EndWorker from consuming straggler messages (#6522) ### 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) --- .../handlers/control/end_worker_handler.py | 20 +++-- .../handlers/control/test_end_worker_handler.py | 93 ++++++++++++++++++++++ 2 files changed, 106 insertions(+), 7 deletions(-) diff --git a/amber/src/main/python/core/architecture/handlers/control/end_worker_handler.py b/amber/src/main/python/core/architecture/handlers/control/end_worker_handler.py index 56baf1f345..278624f836 100644 --- a/amber/src/main/python/core/architecture/handlers/control/end_worker_handler.py +++ b/amber/src/main/python/core/architecture/handlers/control/end_worker_handler.py @@ -18,7 +18,7 @@ from loguru import logger from core.architecture.handlers.control.control_handler_base import ControlHandler -from core.util import IQueue +from core.models.internal_queue import InternalQueue from proto.org.apache.texera.amber.engine.architecture.rpc import ( EmptyReturn, EmptyRequest, @@ -37,13 +37,19 @@ class EndWorkerHandler(ControlHandler): has finished not only the data processing logic, but also the processing of all the control messages. """ - # Ensure this is really the last message. - input_queue: IQueue = self.context.input_queue - if not input_queue.is_empty(): + # Ensure this is really the last message. Read the queued count once (InternalQueue + # exposes size(); the base IQueue interface does not) and branch on it. + input_queue: InternalQueue = self.context.input_queue + queued_count = input_queue.size() + if queued_count > 0: logger.warning( - f"Received EndHandler before all messages are " - f"processed. Unprocessed messages: {input_queue.get()}" + f"Received EndWorker before all {queued_count} queued " + f"message(s) were processed; failing the RPC so a later " + f"coordinator retry succeeds once the queue has drained." ) - assert input_queue.is_empty() + # Fail this RPC (the counterpart of the Scala EndHandler's + # Future.exception) so a later coordinator retry succeeds once + # the queue has drained, instead of dropping the pending message. + raise RuntimeError("worker still has unprocessed messages") # Now we can safely acknowledge that this worker can be terminated. return EmptyReturn() diff --git a/amber/src/test/python/core/architecture/handlers/control/test_end_worker_handler.py b/amber/src/test/python/core/architecture/handlers/control/test_end_worker_handler.py new file mode 100644 index 0000000000..c04ff5ccc7 --- /dev/null +++ b/amber/src/test/python/core/architecture/handlers/control/test_end_worker_handler.py @@ -0,0 +1,93 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import asyncio +from types import SimpleNamespace + +import pytest + +from core.architecture.handlers.control.end_worker_handler import EndWorkerHandler +from core.models.internal_queue import DCMElement, InternalQueue +from proto.org.apache.texera.amber.core import ActorVirtualIdentity, ChannelIdentity +from proto.org.apache.texera.amber.engine.architecture.rpc import ( + EmptyRequest, + EmptyReturn, +) +from proto.org.apache.texera.amber.engine.common import DirectControlMessagePayloadV2 + + +class TestEndWorkerHandler: + @pytest.fixture + def input_queue(self): + return InternalQueue() + + @pytest.fixture + def handler(self, input_queue): + return EndWorkerHandler(SimpleNamespace(input_queue=input_queue)) + + @staticmethod + def make_control_message(seq: int) -> DCMElement: + channel = ChannelIdentity( + ActorVirtualIdentity(name="CONTROLLER"), + ActorVirtualIdentity(name=f"worker-{seq}"), + is_control=True, + ) + return DCMElement(tag=channel, payload=DirectControlMessagePayloadV2()) + + def test_acknowledges_end_worker_on_empty_queue(self, handler): + result = asyncio.run(handler.end_worker(EmptyRequest())) + assert isinstance(result, EmptyReturn) + + def test_rejects_end_worker_with_one_unprocessed_message( + self, handler, input_queue + ): + # The controller resolves endWorker's reply as a failed future and + # retries on a fixed delay (terminateWorkersWithRetry), so a + # straggler message must fail the call — not be dropped with a + # successful acknowledgement. + input_queue.put(self.make_control_message(0)) + with pytest.raises(RuntimeError, match="unprocessed messages"): + asyncio.run(handler.end_worker(EmptyRequest())) + + @pytest.mark.timeout(2) + def test_does_not_consume_the_unprocessed_message(self, handler, input_queue): + straggler = self.make_control_message(0) + input_queue.put(straggler) + with pytest.raises(RuntimeError): + asyncio.run(handler.end_worker(EmptyRequest())) + # The straggler must survive the failed call so the main loop can + # still process it before the controller's retry. + assert input_queue.size() == 1 + assert input_queue.get() is straggler + + def test_keeps_all_messages_with_multiple_stragglers(self, handler, input_queue): + input_queue.put(self.make_control_message(0)) + input_queue.put(self.make_control_message(1)) + with pytest.raises(RuntimeError): + asyncio.run(handler.end_worker(EmptyRequest())) + assert input_queue.size() == 2 + + @pytest.mark.timeout(2) + def test_succeeds_after_queue_drains(self, handler, input_queue): + # The retry protocol end-to-end: fail while a message is pending, + # acknowledge once the queue has drained. + input_queue.put(self.make_control_message(0)) + with pytest.raises(RuntimeError): + asyncio.run(handler.end_worker(EmptyRequest())) + input_queue.get() + result = asyncio.run(handler.end_worker(EmptyRequest())) + assert isinstance(result, EmptyReturn)
