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 278624f836f..cd1788d659d 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,10 +18,12 @@ from loguru import logger from core.architecture.handlers.control.control_handler_base import ControlHandler -from core.models.internal_queue import InternalQueue +from core.models.internal_queue import DCMElement, InternalQueue +from core.util.proto import get_one_of from proto.org.apache.texera.amber.engine.architecture.rpc import ( EmptyReturn, EmptyRequest, + ReturnInvocation, ) @@ -37,19 +39,45 @@ async def end_worker(self, req: EmptyRequest) -> EmptyReturn: has finished not only the data processing logic, but also the processing of all the control messages. """ - # 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. + # Ensure this is really the last message that carries work. 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: + if input_queue.size() == 0: + # Now we can safely acknowledge that this worker can be terminated. + return EmptyReturn() + + # RPC acks (ReturnInvocations for this worker's own fire-and-forget + # calls, e.g. worker_execution_completed) race with EndWorker by + # design: the coordinator decides to end the worker while its acks are + # still in flight. The worker never awaits those acks, so an ack-only + # backlog is safe to drop at termination. Anything else fails loudly + # AND is put back on the queue so the coordinator's retried EndWorker + # finds it processed (mirrors the Scala EndHandler). + pending = [] + while not input_queue.is_empty(): + pending.append(input_queue.get()) + + def is_ack(element) -> bool: + return isinstance(element, DCMElement) and isinstance( + get_one_of(element.payload, sealed=False), ReturnInvocation + ) + + if all(is_ack(element) for element in pending): logger.warning( - 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." + f"Received EndWorker with only RPC acks left in the queue; " + f"proceeding with termination. Pending acks: {pending}" ) - # 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() + return EmptyReturn() + + for element in pending: + input_queue.put(element) + logger.warning( + f"Received EndWorker before all {len(pending)} queued " + f"message(s) were processed; failing the RPC so a later " + f"coordinator retry succeeds once the queue has drained." + ) + # 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") diff --git a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandler.scala b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandler.scala index f91a28c627b..db10e2bb633 100644 --- a/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandler.scala +++ b/amber/src/main/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandler.scala @@ -24,8 +24,14 @@ import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{ AsyncRPCContext, EmptyRequest } -import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn +import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{ + EmptyReturn, + ReturnInvocation +} import org.apache.texera.amber.engine.architecture.worker.DataProcessorRPCHandlerInitializer +import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.FIFOMessageElement + +import scala.jdk.CollectionConverters.CollectionHasAsScala /** * The EndWorker control messages is needed to ensure all the other control messages in a worker @@ -42,14 +48,33 @@ trait EndHandler { request: EmptyRequest, ctx: AsyncRPCContext ): Future[EmptyReturn] = { - // Ensure this is really the last message. - if (!dp.inputManager.inputMessageQueue.isEmpty) { + // Ensure this is really the last message that carries work. RPC acks + // (ReturnInvocations for this worker's own fire-and-forget calls, e.g. + // workerExecutionCompleted / portCompleted) race with EndWorker by design: + // the coordinator decides to end the worker while its acks are still in + // flight. The worker never awaits those acks, so an ack-only backlog is + // safe to leave unprocessed at termination -- failing on it only makes + // region termination retry and CI flaky. Anything else still fails loudly + // so the region execution manager retries the kill instead of dropping + // real work. + val pending = dp.inputManager.inputMessageQueue.asScala.toList + val ackOnly = pending.forall { + case FIFOMessageElement(msg) => msg.payload.isInstanceOf[ReturnInvocation] + case _ => false + } + if (pending.nonEmpty && !ackOnly) { logger.warn( s"Received EndHandler before all messages are processed. Unprocessed messages: " + - s"${dp.inputManager.inputMessageQueue.peek()}" + s"$pending" ) return Future.exception(new IllegalStateException("worker still has unprocessed messages")) } + if (pending.nonEmpty) { + logger.warn( + s"Received EndHandler with only RPC acks left in the queue; proceeding with " + + s"termination. Pending acks: $pending" + ) + } // Now we can safely acknowledge that this worker can be terminated. 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 index c04ff5ccc78..34b2df7afb8 100644 --- 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 @@ -15,6 +15,17 @@ # specific language governing permissions and limitations # under the License. +"""EndWorker is the coordinator's acknowledgement point before it stops the +worker actor. A successful reply promises no queued work is dropped. + +RPC acks (ReturnInvocations for this worker's own fire-and-forget calls, e.g. +worker_execution_completed) race with EndWorker by design: the coordinator +decides to end the worker while its acks are still in flight. An ack-only +backlog carries no work and must not fail the kill; anything else must fail +loudly AND stay in the queue so the coordinator's retry lets the worker drain +it. +""" + import asyncio from types import SimpleNamespace @@ -22,10 +33,13 @@ from core.architecture.handlers.control.end_worker_handler import EndWorkerHandler from core.models.internal_queue import DCMElement, InternalQueue +from core.util.proto import set_one_of from proto.org.apache.texera.amber.core import ActorVirtualIdentity, ChannelIdentity from proto.org.apache.texera.amber.engine.architecture.rpc import ( + ControlReturn, EmptyRequest, EmptyReturn, + ReturnInvocation, ) from proto.org.apache.texera.amber.engine.common import DirectControlMessagePayloadV2 @@ -48,6 +62,26 @@ def make_control_message(seq: int) -> DCMElement: ) return DCMElement(tag=channel, payload=DirectControlMessagePayloadV2()) + @staticmethod + def make_rpc_ack(command_id: int) -> DCMElement: + """A ReturnInvocation ack for one of this worker's own fire-and-forget + RPCs -- the message kind that legitimately races with EndWorker.""" + channel = ChannelIdentity( + ActorVirtualIdentity(name="COORDINATOR"), + ActorVirtualIdentity(name="worker-0"), + is_control=True, + ) + return DCMElement( + tag=channel, + payload=set_one_of( + DirectControlMessagePayloadV2, + ReturnInvocation( + command_id=command_id, + return_value=set_one_of(ControlReturn, EmptyReturn()), + ), + ), + ) + def test_acknowledges_end_worker_on_empty_queue(self, handler): result = asyncio.run(handler.end_worker(EmptyRequest())) assert isinstance(result, EmptyReturn) @@ -91,3 +125,32 @@ def test_succeeds_after_queue_drains(self, handler, input_queue): input_queue.get() result = asyncio.run(handler.end_worker(EmptyRequest())) assert isinstance(result, EmptyReturn) + + def test_succeeds_when_only_rpc_acks_are_queued(self, handler, input_queue): + # An ack-only backlog carries no work (the worker never awaits its + # fire-and-forget acks), so it must not fail the kill -- failing here + # only makes region termination retry and loop e2e tests flaky. + input_queue.put(self.make_rpc_ack(command_id=1)) + input_queue.put(self.make_rpc_ack(command_id=2)) + + result = asyncio.run(handler.end_worker(EmptyRequest())) + + assert isinstance(result, EmptyReturn) + + def test_fails_and_preserves_real_work_mixed_with_acks(self, handler, input_queue): + # Leniency is strictly ack-only: real work queued alongside an ack + # must fail the call AND survive in the queue for the retry to drain. + input_queue.put(self.make_rpc_ack(command_id=1)) + work = self.make_control_message(0) + input_queue.put(work) + + with pytest.raises(RuntimeError, match="unprocessed messages"): + asyncio.run(handler.end_worker(EmptyRequest())) + + remaining = [] + while not input_queue.is_empty(): + remaining.append(input_queue.get()) + assert work in remaining, ( + "the queued control message must survive a failed EndWorker " + f"so the retry can drain it; queue held: {remaining}" + ) diff --git a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandlerSpec.scala b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandlerSpec.scala index 3a3c0ecb9ed..c3577bffd0f 100644 --- a/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandlerSpec.scala +++ b/amber/src/test/scala/org/apache/texera/amber/engine/architecture/worker/promisehandlers/EndHandlerSpec.scala @@ -25,7 +25,10 @@ import org.apache.texera.amber.engine.architecture.rpc.controlcommands.{ AsyncRPCContext, EmptyRequest } -import org.apache.texera.amber.engine.architecture.rpc.controlreturns.EmptyReturn +import org.apache.texera.amber.engine.architecture.rpc.controlreturns.{ + EmptyReturn, + ReturnInvocation +} import org.apache.texera.amber.engine.architecture.rpc.workerservice.WorkerServiceGrpc.METHOD_QUERY_STATISTICS import org.apache.texera.amber.engine.architecture.worker.WorkflowWorker.{ ActorCommandElement, @@ -93,6 +96,15 @@ class EndHandlerSpec extends AnyFlatSpec { queue } + private def ackElement(sequenceNumber: Long, commandId: Long): DPInputQueueElement = + FIFOMessageElement( + WorkflowFIFOMessage( + ChannelIdentity(COORDINATOR, workerId, isControl = true), + sequenceNumber, + ReturnInvocation(commandId, EmptyReturn()) + ) + ) + "EndHandler" should "reply successfully when there are no unprocessed messages" in { val handler = createEndHandlerForQueue(new LinkedBlockingQueue[DPInputQueueElement]()) @@ -110,4 +122,29 @@ class EndHandlerSpec extends AnyFlatSpec { assertEndWorkerFails(handler) } + + it should "reply successfully when only RPC acks (ReturnInvocations) are queued" in { + // Race at region termination: the coordinator's ReturnInvocation ack for + // this worker's own fire-and-forget RPC (e.g. workerExecutionCompleted) + // can still be in flight when EndWorker arrives. Such acks carry no work + // -- the worker never awaits them -- so an ack-only backlog must not fail + // the kill (it previously did, making region termination retry and loop + // e2e tests flaky). + val queue = new LinkedBlockingQueue[DPInputQueueElement]() + queue.put(ackElement(sequenceNumber = 0, commandId = 1)) + queue.put(ackElement(sequenceNumber = 1, commandId = 2)) + val handler = createEndHandlerForQueue(queue) + + assert(await(handler.endWorker(EmptyRequest(), rpcContext)) == EmptyReturn()) + } + + it should "fail when an RPC ack is queued together with real work" in { + // Leniency is strictly ack-only: any non-ack message (here a control + // invocation) still fails the kill so the retry lets the worker drain it. + val queue = queueWithFifoControlMessage() + queue.put(ackElement(sequenceNumber = 1, commandId = 2)) + val handler = createEndHandlerForQueue(queue) + + assertEndWorkerFails(handler) + } }