This is an automated email from the ASF dual-hosted git repository.
xuang7 pushed a commit to branch release/v1.2
in repository https://gitbox.apache.org/repos/asf/texera.git
The following commit(s) were added to refs/heads/release/v1.2 by this push:
new 23d5d1c0dc fix(pyamber, v1.2): stop EndWorker from consuming straggler
messages (#6995)
23d5d1c0dc is described below
commit 23d5d1c0dcf8ac10079f168611e179aca811269f
Author: Yicong Huang <[email protected]>
AuthorDate: Wed Jul 29 14:08:49 2026 -0400
fix(pyamber, v1.2): stop EndWorker from consuming straggler messages
(#6995)
### What changes were proposed in this PR?
Backport of #6522 to `release/v1.2`, cherry-picked from
7e4a9b46b863fc13fa76ba74980a0f0c4beff48f. The cherry-pick applied
cleanly.
Follows the Direct Backport Push convention; opened as a PR (rather than
a direct push) per a backport-coverage audit.
### Any related issues, documentation, discussions?
Backport of #6522. Originally linked #6521.
### How was this PR tested?
Release-branch CI runs on this PR. Cherry-pick applied cleanly onto
`release/v1.2`; no manual conflict resolution was needed.
### Was this PR authored or co-authored using generative AI tooling?
Yes — backport prepared with Claude Code (mechanical cherry-pick; the
change itself is #6522 by its original author).
🤖 Generated with [Claude Code](https://claude.com/claude-code)
Co-authored-by: Xinyuan Lin <[email protected]>
---
.../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 434225a5f0..2628386824 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)
