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-6544-b6cf569c312ff0013bc47e3401be85e1bee05f7f in repository https://gitbox.apache.org/repos/asf/texera.git
commit 511bc276953a251f668329b9cac1b121c7651c69 Author: Xinyuan Lin <[email protected]> AuthorDate: Mon Jul 20 12:57:47 2026 -0700 test(amber): cover zero-tuple-channel worker completion in pyamber (#6544) ### What changes were proposed in this PR? Adds two regression tests to `amber/src/test/python/core/runnables/test_main_loop.py` covering a Python-UDF worker whose input port receives **zero tuples** across a `StartChannel` → `EndChannel` bracket — e.g. the untaken branch of an `If` operator read through an `InputPortMaterializationReaderRunnable`, or a filter that drops everything on a materialized edge: - **`test_zero_tuple_channel_completes_worker`** — `StartChannel` then `EndChannel` with no data tuples must still complete both ports and report `PortCompleted(input)`, `PortCompleted(output)`, and `WorkerExecutionCompleted`, forwarding both channel ECMs downstream. - **`test_empty_on_finish_after_tuples_completes_worker`** — an operator whose `on_finish` yields nothing (like `BatchOperator` when the input size is an exact multiple of `BATCH_SIZE`) must still complete cleanly after processing tuples, since the empty `on_finish` is exhausted within a single MainLoop↔DataProcessor hand-off. Both run the MainLoop on a **daemon thread** and detect completion by **deadline-polling `WorkerState.COMPLETED`**, so a regression that deadlocks the handshake fails the test cleanly instead of hanging the whole pytest process. This is **test-only** — no production changes. It pins behavior the MainLoop↔DataProcessor handshake already handles on `main` (the internal marker is dispatched in `_process_ecm` without an intervening context switch; the worker state graph permits `READY → COMPLETED`, which a zero-tuple worker needs because it never enters `RUNNING`) so it can't silently regress. ### Any related issues, documentation, discussions? Closes #6543 ### How was this PR tested? ``` cd amber && python -m pytest src/test/python/core/runnables/test_main_loop.py \ -k "zero_tuple_channel or empty_on_finish" # 2 passed, 21 deselected cd amber && python -m pytest src/test/python/core/runnables/test_main_loop.py # 23 passed ``` - Both new tests pass on `main`; the whole `test_main_loop.py` stays green (23 passed). - `ruff check` → all checks passed; `ruff format --check` → already formatted. - Guard-verified that each test actually catches a regression (mutation reverted after each): - reintroducing the `_process_ecm` pre-marker `self._switch_context()` → `test_zero_tuple_channel_completes_worker` deadlocks (deadline-fail); - removing `COMPLETED` from the `READY` set in the worker state graph → `test_zero_tuple_channel_completes_worker` fails with `InvalidTransitionException`; - gating `finished_current.set()` behind an "`on_finish` produced ≥1 output" flag → `test_empty_on_finish_after_tuples_completes_worker` hangs (deadline-fail), while Echo-based tests stay green. ### Was this PR authored or co-authored using generative AI tooling? Generated-by: Claude Code (Opus 4.8 [1M context]) --- .../test/python/core/runnables/test_main_loop.py | 375 +++++++++++++++++++++ 1 file changed, 375 insertions(+) diff --git a/amber/src/test/python/core/runnables/test_main_loop.py b/amber/src/test/python/core/runnables/test_main_loop.py index 79d372e2c0..5884edf27b 100644 --- a/amber/src/test/python/core/runnables/test_main_loop.py +++ b/amber/src/test/python/core/runnables/test_main_loop.py @@ -22,6 +22,7 @@ import pyarrow import pytest import sys import time +import uuid from threading import Thread from core.models import ( @@ -86,6 +87,7 @@ from proto.org.apache.texera.amber.engine.architecture.worker import ( from proto.org.apache.texera.amber.engine.common import DirectControlMessagePayloadV2 from pytexera.udf.examples.count_batch_operator import CountBatchOperator from pytexera.udf.examples.echo_operator import EchoOperator +from pytexera.udf.udf_operator import UDFOperatorV2 class _FalseLoopEnd(LoopEndOperator): @@ -93,6 +95,22 @@ class _FalseLoopEnd(LoopEndOperator): return False +class EmptyOnFinishOperator(UDFOperatorV2): + # Echoes each input tuple, but its on_finish is a zero-yield generator + # (`return` before `yield` makes the body unreachable while still marking + # the function as a generator). This is the BatchOperator-fed-an-exact- + # multiple-of-BATCH_SIZE shape: the EndChannel on_finish produces NOTHING, + # so DataProcessor._set_output_tuple exhausts the iterator in a single + # hand-off (no per-output switch dance) and sets finished_current straight + # away. MainLoop must not lose that completion signal. + def process_tuple(self, tuple_, port): + yield tuple_ + + def on_finish(self, port): + return + yield + + class TestMainLoop: @pytest.fixture def command_sequence(self): @@ -268,6 +286,58 @@ class TestMainLoop: ), ) + @pytest.fixture + def mock_start_channel(self, mock_data_input_channel): + # Mirror of mock_end_of_upstream but a StartChannel ECM with + # NO_ALIGNMENT (the alignment a real StartChannel bracket uses). + return ECMElement( + tag=mock_data_input_channel, + payload=EmbeddedControlMessage( + EmbeddedControlMessageIdentity("StartChannel"), + EmbeddedControlMessageType.NO_ALIGNMENT, + [], + { + mock_data_input_channel.to_worker_id.name: ControlInvocation( + "StartChannel", + ControlRequest(empty_request=EmptyRequest()), + AsyncRpcContext(ActorVirtualIdentity(), ActorVirtualIdentity()), + -1, + ) + }, + ), + ) + + @pytest.fixture + def mock_initialize_empty_on_finish_executor( + self, + mock_control_input_channel, + mock_sender_actor, + mock_link, + command_sequence, + mock_raw_schema, + ): + operator_code = "from pytexera import *\n" + inspect.getsource( + EmptyOnFinishOperator + ) + command = set_one_of( + ControlRequest, + InitializeExecutorRequest( + op_exec_init_info=set_one_of( + OpExecInitInfo, OpExecWithCode(operator_code, "python") + ), + is_source=False, + ), + ) + payload = set_one_of( + DirectControlMessagePayloadV2, + ControlInvocation( + method_name="InitializeExecutor", + command_id=command_sequence, + command=command, + ), + ) + return DCMElement(tag=mock_control_input_channel, payload=payload) + @pytest.fixture def input_queue(self): return InternalQueue() @@ -1684,6 +1754,311 @@ class TestMainLoop: reraise() + @staticmethod + def _expected_port_completed_dcm( + mock_control_output_channel, command_id, port_id, is_input + ): + return DCMElement( + tag=mock_control_output_channel, + payload=DirectControlMessagePayloadV2( + control_invocation=ControlInvocation( + method_name="PortCompleted", + command_id=command_id, + context=AsyncRpcContext( + sender=ActorVirtualIdentity(name="dummy_worker_id"), + receiver=ActorVirtualIdentity(name="COORDINATOR"), + ), + command=ControlRequest( + port_completed_request=PortCompletedRequest( + port_id=port_id, input=is_input + ) + ), + ) + ), + ) + + @staticmethod + def _expected_worker_completed_dcm(mock_control_output_channel): + return DCMElement( + tag=mock_control_output_channel, + payload=DirectControlMessagePayloadV2( + control_invocation=ControlInvocation( + method_name="WorkerExecutionCompleted", + command_id=2, + context=AsyncRpcContext( + sender=ActorVirtualIdentity(name="dummy_worker_id"), + receiver=ActorVirtualIdentity(name="COORDINATOR"), + ), + command=ControlRequest(empty_request=EmptyRequest()), + ) + ), + ) + + @staticmethod + def _forwarded_ecm(mock_data_output_channel, method_name, alignment): + return ECMElement( + tag=mock_data_output_channel, + payload=EmbeddedControlMessage( + EmbeddedControlMessageIdentity(method_name), + alignment, + [], + { + mock_data_output_channel.to_worker_id.name: ControlInvocation( + method_name, + ControlRequest(empty_request=EmptyRequest()), + AsyncRpcContext(ActorVirtualIdentity(), ActorVirtualIdentity()), + -1, + ) + }, + ), + ) + + @staticmethod + def _drain_until(output_queue, done, timeout=15.0): + # Non-blocking drain of the output queue against a deadline. A + # regression that deadlocks the MainLoop/DataProcessor handshake never + # satisfies `done`, so we return the partial batch at the deadline and + # let the caller pytest.fail() -- the whole pytest process is never + # hung because the worker runs on a daemon thread. + deadline = time.time() + timeout + collected = [] + while time.time() < deadline: + while output_queue.size() > 0: + collected.append(output_queue.get()) + if done(collected): + return collected + time.sleep(0.005) + return collected + + @pytest.mark.timeout(30) + def test_zero_tuple_channel_completes_worker( + self, + mock_link, + mock_data_output_channel, + mock_control_output_channel, + input_queue, + output_queue, + main_loop, + main_loop_thread, + mock_assign_input_port, + mock_assign_output_port, + mock_add_input_channel, + mock_add_partitioning, + mock_initialize_executor, + mock_start_channel, + mock_end_of_upstream, + command_sequence, + reraise, + ): + # A worker whose input port receives a StartChannel->EndChannel bracket + # with ZERO DataElements (the untaken branch of an If read through an + # InputPortMaterializationReaderRunnable, or a filter that drops + # everything on a materialized edge) must still process both ECMs and + # reach COMPLETED. Two sub-bugs made this hang/crash on the old + # loop-feb branch: + # (1) DEADLOCK: MainLoop._process_ecm re-read current_internal_marker + # AFTER a _switch_context(), so the DataProcessor could pop the + # marker mid-switch and MainLoop skipped _process_end_channel -> + # both threads park forever. + # (2) STATE GRAPH: a zero-tuple worker never enters RUNNING (only + # _process_data_element does that), so completion is a direct + # READY -> COMPLETED transition, which the transition graph must + # permit. + # Run on a daemon thread and detect completion by deadline-polling the + # state manager so a regression deadlock fails cleanly instead of + # hanging the whole pytest process. + main_loop_thread.daemon = True + main_loop_thread.start() + + for setup_msg in [ + mock_assign_input_port, + mock_assign_output_port, + mock_add_input_channel, + mock_add_partitioning, + mock_initialize_executor, + ]: + input_queue.put(setup_msg) + assert output_queue.get() == DCMElement( + tag=mock_control_output_channel, + payload=DirectControlMessagePayloadV2( + return_invocation=ReturnInvocation( + command_id=command_sequence, + return_value=ControlReturn(empty_return=EmptyReturn()), + ) + ), + ) + + # The worker is READY here and never enters RUNNING (no data element). + assert main_loop.context.state_manager.confirm_state(WorkerState.READY) + + # Zero tuples between StartChannel and EndChannel. + input_queue.put(mock_start_channel) + input_queue.put(mock_end_of_upstream) + + expected_worker_completed = self._expected_worker_completed_dcm( + mock_control_output_channel + ) + collected = self._drain_until( + output_queue, + lambda items: expected_worker_completed in items, + ) + + if not main_loop.context.state_manager.confirm_state(WorkerState.COMPLETED): + pytest.fail( + "zero-tuple worker did not reach COMPLETED within the deadline " + "-- likely the _process_ecm marker-after-switch deadlock or a " + "missing READY->COMPLETED transition. " + f"state={main_loop.context.state_manager.get_current_state()}, " + f"collected={collected}" + ) + + # Both the input and output ports complete, and the worker signals + # WorkerExecutionCompleted -- all on the coordinator control channel. + expected_input_port_completed = self._expected_port_completed_dcm( + mock_control_output_channel, 0, mock_link.to_port_id, True + ) + expected_output_port_completed = self._expected_port_completed_dcm( + mock_control_output_channel, 1, PortIdentity(id=0), False + ) + assert expected_input_port_completed in collected + assert expected_output_port_completed in collected + assert expected_worker_completed in collected + + # Both ECMs are forwarded downstream on the data output channel. + assert ( + self._forwarded_ecm( + mock_data_output_channel, + "StartChannel", + EmbeddedControlMessageType.NO_ALIGNMENT, + ) + in collected + ) + assert ( + self._forwarded_ecm( + mock_data_output_channel, + "EndChannel", + EmbeddedControlMessageType.PORT_ALIGNMENT, + ) + in collected + ) + + reraise() + + @pytest.mark.timeout(30) + def test_empty_on_finish_after_tuples_completes_worker( + self, + mock_link, + mock_tuple, + mock_data_output_channel, + mock_control_output_channel, + input_queue, + output_queue, + main_loop, + main_loop_thread, + mock_assign_input_port, + mock_assign_output_port, + mock_add_input_channel, + mock_add_partitioning, + mock_initialize_empty_on_finish_executor, + mock_data_element, + mock_end_of_upstream, + command_sequence, + monkeypatch, + reraise, + ): + # Sibling case: after processing real tuples, an EndChannel whose + # on_finish yields NOTHING must also complete cleanly. The empty + # on_finish is exhausted inside a single hand-off (DataProcessor + # ._set_output_tuple runs no per-output switch dance, it just sets + # finished_current), and MainLoop must not lose the completion signal. + + # Guard the udf-v1 executor-module-contamination landmine: force a + # unique module name so cross-test importlib caching can't hand us a + # stale operator class. (main's ExecutorManager already uses a + # process-wide unique counter, so this is belt-and-suspenders.) + unique_name = f"udf_empty_on_finish_{uuid.uuid4().hex}" + monkeypatch.setattr( + main_loop.context.executor_manager, + "gen_module_file_name", + lambda: (unique_name, f"{unique_name}.py"), + ) + + main_loop_thread.daemon = True + main_loop_thread.start() + + for setup_msg in [ + mock_assign_input_port, + mock_assign_output_port, + mock_add_input_channel, + mock_add_partitioning, + mock_initialize_empty_on_finish_executor, + ]: + input_queue.put(setup_msg) + assert output_queue.get() == DCMElement( + tag=mock_control_output_channel, + payload=DirectControlMessagePayloadV2( + return_invocation=ReturnInvocation( + command_id=command_sequence, + return_value=ControlReturn(empty_return=EmptyReturn()), + ) + ), + ) + + # The loaded executor must be our zero-yield-on_finish operator, not a + # stale cached class from another test. + assert ( + type(main_loop.context.executor_manager.executor).__name__ + == "EmptyOnFinishOperator" + ) + + # One real tuple: the operator echoes it and the worker enters RUNNING. + input_queue.put(mock_data_element) + echoed: DataElement = output_queue.get() + assert echoed.tag == mock_data_output_channel + assert isinstance(echoed.payload, DataFrame) + assert Tuple(echoed.payload.frame.to_pylist()[0]) == mock_tuple + + # EndChannel with an empty on_finish must still complete the worker. + input_queue.put(mock_end_of_upstream) + + expected_worker_completed = self._expected_worker_completed_dcm( + mock_control_output_channel + ) + collected = self._drain_until( + output_queue, + lambda items: expected_worker_completed in items, + ) + + if not main_loop.context.state_manager.confirm_state(WorkerState.COMPLETED): + pytest.fail( + "worker with an empty on_finish did not reach COMPLETED within " + "the deadline -- the single-hand-off completion signal was lost. " + f"state={main_loop.context.state_manager.get_current_state()}, " + f"collected={collected}" + ) + + expected_input_port_completed = self._expected_port_completed_dcm( + mock_control_output_channel, 0, mock_link.to_port_id, True + ) + expected_output_port_completed = self._expected_port_completed_dcm( + mock_control_output_channel, 1, PortIdentity(id=0), False + ) + assert expected_input_port_completed in collected + assert expected_output_port_completed in collected + assert expected_worker_completed in collected + + # The EndChannel ECM is forwarded downstream on the data output channel. + assert ( + self._forwarded_ecm( + mock_data_output_channel, + "EndChannel", + EmbeddedControlMessageType.PORT_ALIGNMENT, + ) + in collected + ) + + reraise() + @pytest.mark.timeout(2) def test_console_message_rpc_fires_before_exception_pause( self, main_loop, monkeypatch
