This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/texera.git
The following commit(s) were added to refs/heads/main by this push:
new 7e4a9b46b8 fix(pyamber): stop EndWorker from consuming straggler
messages (#6522)
7e4a9b46b8 is described below
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)
