You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@texera.apache.org by gi...@apache.org on 2026/07/20 19:58:09 UTC
(texera) branch gh-readonly-queue/main/pr-6544-b6cf569c312ff0013bc47e3401be85e1bee05f7f created (now 511bc27695)
This is an automated email from the ASF dual-hosted git repository.
github-merge-queue[bot] pushed a change to branch gh-readonly-queue/main/pr-6544-b6cf569c312ff0013bc47e3401be85e1bee05f7f
in repository https://gitbox.apache.org/repos/asf/texera.git
at 511bc27695 test(amber): cover zero-tuple-channel worker completion in pyamber (#6544)
This branch includes the following new commits:
new 511bc27695 test(amber): cover zero-tuple-channel worker completion in pyamber (#6544)
The 1 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails. The revisions
listed as "add" were already present in the repository and have only
been added to this reference.
(texera) 01/01: test(amber): cover zero-tuple-channel worker completion in pyamber (#6544)
Posted by gi...@apache.org.
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 <xi...@uci.edu>
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