From 2f9c4e7ba8a27c0772e3b89301f9ce3176dd7642 Mon Sep 17 00:00:00 2001 From: carloea2 Date: Sun, 20 Sep 2026 23:58:07 -0700 Subject: [PATCH] fix(pyamber): preserve channel marker ports (#8432) ### What changes were proposed in this PR? Bind channel markers to the input port resolved by their control handler. Lifecycle callbacks use that port instead of the last data tuple's port. Reject invalid markers before entering the executor. ### Any related issues, documentation, discussions? Related to #8319. ### How was this PR tested? From amber, on Python 3.12: ```text python -m pytest --tb=short -q src/test/python/core/architecture/handlers/control/test_channel_marker_handlers.py src/test/python/core/runnables/test_data_processor.py src/test/python/core/runnables/test_main_loop.py 66 passed python -m ruff check src/main/python src/test/python All checks passed python -m ruff format --check src/main/python src/test/python 215 files already formatted ``` The new handler-to-processor cases cover an empty port 1 finishing after a tuple on port 0, and port 1 starting before any tuple. Both fail when marker routing is reverted to the current data port. Invalid-marker tests failed before adding the type guard. The full run, `python -m pytest --tb=no -q`, finished with 1320 passed, 1 xfailed, and 1 failed. The failure is the unchanged `test_rest_catalog_round_trip` integration test, which targets a Lakekeeper service at localhost:8181. The full suite is not green. Excluding integration tests, `python -m pytest --tb=no -q -m "not integration"` passes: 1320 passed, 1 deselected, 1 xfailed. ### Was this PR authored or co-authored using generative AI tooling? Generated-by: OpenAI Codex (backported from commit 471e53cfc1d5f96748a7ec6b7563561239229e5e) Co-authored-by: Meng Wang --- .../handlers/control/end_channel_handler.py | 6 +- .../handlers/control/start_channel_handler.py | 7 +- .../python/core/models/internal_marker.py | 30 ++++- .../python/core/runnables/data_processor.py | 18 ++- .../control/test_channel_marker_handlers.py | 126 ++++++++++++++++++ .../core/runnables/test_data_processor.py | 109 +++++++++++++-- .../python/core/runnables/test_main_loop.py | 2 +- 7 files changed, 271 insertions(+), 27 deletions(-) create mode 100644 amber/src/test/python/core/architecture/handlers/control/test_channel_marker_handlers.py diff --git a/amber/src/main/python/core/architecture/handlers/control/end_channel_handler.py b/amber/src/main/python/core/architecture/handlers/control/end_channel_handler.py index d60b3874754..e34be3d2a47 100644 --- a/amber/src/main/python/core/architecture/handlers/control/end_channel_handler.py +++ b/amber/src/main/python/core/architecture/handlers/control/end_channel_handler.py @@ -25,8 +25,12 @@ class EndChannelHandler(ControlHandler): async def end_channel(self, req: EmptyRequest) -> EmptyReturn: + port = self.context.input_manager.get_port_id( + self.context.current_input_channel_id + ) + marker = EndChannel(port.id) self.context.input_manager.complete_current_port( self.context.current_input_channel_id ) - self.context.tuple_processing_manager.current_internal_marker = EndChannel() + self.context.tuple_processing_manager.current_internal_marker = marker return EmptyReturn() diff --git a/amber/src/main/python/core/architecture/handlers/control/start_channel_handler.py b/amber/src/main/python/core/architecture/handlers/control/start_channel_handler.py index 36747351e1e..24bc27fb281 100644 --- a/amber/src/main/python/core/architecture/handlers/control/start_channel_handler.py +++ b/amber/src/main/python/core/architecture/handlers/control/start_channel_handler.py @@ -25,5 +25,10 @@ class StartChannelHandler(ControlHandler): async def start_channel(self, req: EmptyRequest) -> EmptyReturn: - self.context.tuple_processing_manager.current_internal_marker = StartChannel() + port = self.context.input_manager.get_port_id( + self.context.current_input_channel_id + ) + self.context.tuple_processing_manager.current_internal_marker = StartChannel( + port.id + ) return EmptyReturn() diff --git a/amber/src/main/python/core/models/internal_marker.py b/amber/src/main/python/core/models/internal_marker.py index 6c9c80bafc4..209c33c5101 100644 --- a/amber/src/main/python/core/models/internal_marker.py +++ b/amber/src/main/python/core/models/internal_marker.py @@ -16,18 +16,36 @@ # under the License. +from dataclasses import dataclass + + class InternalMarker: """ - A special Data Message, only being generated in un-packaging a batch into Tuples. - Markers retain the order information and served as a indicator of data state. + An internal event produced by batch unpacking or control handlers. + Markers preserve ordering and signal input lifecycle changes. """ pass -class StartChannel(InternalMarker): - pass +@dataclass(frozen=True) +class PortMarker(InternalMarker): + """A control marker bound to the input port that produced it.""" + port_id: int -class EndChannel(InternalMarker): - pass + def __post_init__(self) -> None: + if ( + not isinstance(self.port_id, int) + or isinstance(self.port_id, bool) + or self.port_id < 0 + ): + raise ValueError("marker port_id must be a nonnegative integer") + + +class StartChannel(PortMarker): + """Start-of-channel marker with immutable input-port provenance.""" + + +class EndChannel(PortMarker): + """End-of-channel marker with immutable input-port provenance.""" diff --git a/amber/src/main/python/core/runnables/data_processor.py b/amber/src/main/python/core/runnables/data_processor.py index 22e7058f27d..eaee8660cd5 100644 --- a/amber/src/main/python/core/runnables/data_processor.py +++ b/amber/src/main/python/core/runnables/data_processor.py @@ -20,8 +20,8 @@ from typing import Iterator, Optional from core.architecture.managers import Context -from core.models import State, TupleLike, InternalMarker -from core.models.internal_marker import StartChannel, EndChannel +from core.models import State, TupleLike +from core.models.internal_marker import EndChannel, PortMarker, StartChannel from core.models.table import all_output_to_tuple from core.util import Stoppable from core.util.console_message.replace_print import replace_print @@ -65,8 +65,10 @@ def run(self) -> None: else: self.process_tuple() - def process_internal_marker(self, internal_marker: InternalMarker) -> None: - with self._executor_session() as (executor, port_id): + def process_internal_marker(self, internal_marker: PortMarker) -> None: + if not isinstance(internal_marker, PortMarker): + raise TypeError("expected a PortMarker") + with self._executor_session(internal_marker.port_id) as (executor, port_id): if isinstance(internal_marker, StartChannel): self._set_output_state(executor.produce_state_on_start(port_id)) elif isinstance(internal_marker, EndChannel): @@ -95,7 +97,7 @@ def process_tuple(self) -> None: self._set_output_tuple(executor.process_tuple(tuple_, port_id)) @contextmanager - def _executor_session(self): + def _executor_session(self, marker_port_id: int | None = None): """ Open one executor invocation: hand back (executor, port_id) under a print-capture session, route any exception to the exception manager @@ -107,7 +109,11 @@ def _executor_session(self): """ try: executor = self._context.executor_manager.executor - port_id = self._context.tuple_processing_manager.get_input_port_id() + port_id = ( + self._context.tuple_processing_manager.get_input_port_id() + if marker_port_id is None + else marker_port_id + ) with replace_print( self._context.worker_id, self._context.console_message_manager.print_buf, diff --git a/amber/src/test/python/core/architecture/handlers/control/test_channel_marker_handlers.py b/amber/src/test/python/core/architecture/handlers/control/test_channel_marker_handlers.py new file mode 100644 index 00000000000..341393e27f6 --- /dev/null +++ b/amber/src/test/python/core/architecture/handlers/control/test_channel_marker_handlers.py @@ -0,0 +1,126 @@ +# 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_channel_handler import EndChannelHandler +from core.architecture.handlers.control.start_channel_handler import ( + StartChannelHandler, +) +from core.architecture.packaging.input_manager import InputManager +from core.models import Schema +from core.models.internal_marker import EndChannel, StartChannel +from core.models.internal_queue import InternalQueue +from proto.org.apache.texera.amber.core import ( + ActorVirtualIdentity, + ChannelIdentity, + PortIdentity, +) +from proto.org.apache.texera.amber.engine.architecture.rpc import EmptyRequest + +WORKER_ID = "worker-1" +UPSTREAM_ID = ActorVirtualIdentity("upstream-worker") + + +def _channel(index: int) -> ChannelIdentity: + """Return a distinct data-channel identity.""" + return ChannelIdentity( + ActorVirtualIdentity(f"{UPSTREAM_ID.name}-{index}"), + ActorVirtualIdentity(WORKER_ID), + False, + ) + + +def _build_context(*port_numbers: int): + """Build real input state and the minimal marker-handler context.""" + manager = InputManager(WORKER_ID, InternalQueue()) + channels = [] + ports = [] + for index, port_number in enumerate(port_numbers): + channel = _channel(index) + port = PortIdentity(port_number, False) + manager.add_input_port(port, Schema(), [], []) + manager.register_input(channel, port) + channels.append(channel) + ports.append(port) + context = SimpleNamespace( + input_manager=manager, + current_input_channel_id=channels[-1] if channels else _channel(99), + tuple_processing_manager=SimpleNamespace(current_internal_marker=None), + ) + return context, channels, ports + + +@pytest.mark.parametrize( + ("handler_type", "method_name", "marker_type"), + [ + (StartChannelHandler, "start_channel", StartChannel), + (EndChannelHandler, "end_channel", EndChannel), + ], +) +def test_channel_marker_snapshots_the_current_ports_identity( + handler_type, method_name, marker_type +): + """Each channel retains its own port.""" + context, channels, ports = _build_context(3, 7) + handler = handler_type(context) + + context.current_input_channel_id = channels[1] + asyncio.run(getattr(handler, method_name)(EmptyRequest())) + + assert context.tuple_processing_manager.current_internal_marker == marker_type( + ports[1].id + ) + assert context.input_manager.get_port(ports[0]).completed is False + assert context.input_manager.get_port(ports[1]).completed is ( + handler_type is EndChannelHandler + ) + + +@pytest.mark.parametrize( + ("handler_type", "method_name"), + [ + (StartChannelHandler, "start_channel"), + (EndChannelHandler, "end_channel"), + ], +) +def test_unknown_channel_leaves_marker_and_known_ports_untouched( + handler_type, method_name +): + """A missing channel mapping must fail before producing any side effect.""" + context, _, ports = _build_context(3) + context.current_input_channel_id = _channel(99) + + with pytest.raises(KeyError): + asyncio.run(getattr(handler_type(context), method_name)(EmptyRequest())) + + assert context.tuple_processing_manager.current_internal_marker is None + assert context.input_manager.get_port(ports[0]).completed is False + + +def test_invalid_end_channel_port_fails_before_completing_the_port(): + """An invalid marker must not complete its port.""" + context, _, ports = _build_context(-1) + + with pytest.raises(ValueError, match="nonnegative integer"): + asyncio.run(EndChannelHandler(context).end_channel(EmptyRequest())) + + assert context.tuple_processing_manager.current_internal_marker is None + assert context.input_manager.get_port(ports[0]).completed is False diff --git a/amber/src/test/python/core/runnables/test_data_processor.py b/amber/src/test/python/core/runnables/test_data_processor.py index 61482363a15..e185872f2e2 100644 --- a/amber/src/test/python/core/runnables/test_data_processor.py +++ b/amber/src/test/python/core/runnables/test_data_processor.py @@ -15,14 +15,27 @@ # specific language governing permissions and limitations # under the License. +import asyncio +from types import SimpleNamespace + import pytest +from core.architecture.handlers.control.end_channel_handler import EndChannelHandler +from core.architecture.handlers.control.start_channel_handler import StartChannelHandler from core.architecture.managers import Context -from core.models import State +from core.models import Schema, State, Tuple from core.models.internal_queue import InternalQueue -from core.models.internal_marker import EndChannel, StartChannel +from core.models.internal_marker import EndChannel, InternalMarker, StartChannel from core.runnables.data_processor import DataProcessor -from proto.org.apache.texera.amber.engine.architecture.rpc import ConsoleMessageType +from proto.org.apache.texera.amber.core import ( + ActorVirtualIdentity, + ChannelIdentity, + PortIdentity, +) +from proto.org.apache.texera.amber.engine.architecture.rpc import ( + ConsoleMessageType, + EmptyRequest, +) @pytest.fixture @@ -69,8 +82,79 @@ def on_finish(self, port_id): self.calls.append(("on_finish", port_id)) return iter([]) + def process_tuple(self, tuple_, port_id): + self.calls.append(("process_tuple", port_id)) + return iter([]) + class TestProcessInternalMarker: + @pytest.mark.parametrize( + "marker", [None, InternalMarker(), SimpleNamespace(port_id=1)] + ) + def test_rejects_non_port_markers_before_entering_executor( + self, context, data_processor, marker + ): + executor = _StubExecutor() + context.executor_manager.executor = executor + with pytest.raises(TypeError, match="PortMarker"): + data_processor.process_internal_marker(marker) + assert executor.calls == [] + assert data_processor.switch_calls == 0 + + @pytest.mark.parametrize("finish", [False, True]) + @pytest.mark.timeout(2) + def test_handlers_route_empty_second_port_to_lifecycle_callbacks( + self, context, data_processor, finish + ): + executor = _StubExecutor() + context.executor_manager.executor = executor + channels = [] + for port_id in (0, 1): + port = PortIdentity(port_id, False) + channel = ChannelIdentity( + ActorVirtualIdentity(f"upstream-{port_id}"), + ActorVirtualIdentity(context.worker_id), + False, + ) + context.input_manager.add_input_port(port, Schema(), [], []) + context.input_manager.register_input(channel, port) + channels.append(channel) + + tpm = context.tuple_processing_manager + if finish: + context.current_input_channel_id = channels[0] + tpm.current_input_port_id = context.input_manager.get_port_id(channels[0]) + tpm.current_input_tuple = Tuple({"value": 1}) + data_processor.process_tuple() + assert executor.calls == [("process_tuple", 0)] + else: + assert tpm.current_input_port_id is None + + context.current_input_channel_id = channels[1] + if finish: + asyncio.run(EndChannelHandler(context).end_channel(EmptyRequest())) + else: + asyncio.run(StartChannelHandler(context).start_channel(EmptyRequest())) + data_processor.process_internal_marker(tpm.get_internal_marker()) + + expected = ( + [("process_tuple", 0), ("produce_state_on_finish", 1), ("on_finish", 1)] + if finish + else [("produce_state_on_start", 1)] + ) + assert executor.calls == expected + assert not context.exception_manager.has_exception() + assert not context.input_manager.get_port(PortIdentity(0, False)).completed + assert ( + context.input_manager.get_port(PortIdentity(1, False)).completed is finish + ) + + @pytest.mark.parametrize("marker", [StartChannel, EndChannel]) + @pytest.mark.parametrize("port_id", [-1, True, "1"]) + def test_channel_markers_reject_invalid_port_identity(self, marker, port_id): + with pytest.raises(ValueError, match="nonnegative integer"): + marker(port_id) + @pytest.mark.timeout(2) def test_start_channel_invokes_produce_state_on_start( self, context, data_processor @@ -78,12 +162,12 @@ def test_start_channel_invokes_produce_state_on_start( executor = _StubExecutor() context.executor_manager.executor = executor - data_processor.process_internal_marker(StartChannel()) + context.tuple_processing_manager.current_input_port_id = object() + data_processor.process_internal_marker(StartChannel(2)) - # StartChannel routes to produce_state_on_start with the current - # input port id (0 when no upstream is set), and the returned dict - # is wrapped into a State on the output slot. - assert executor.calls == [("produce_state_on_start", 0)] + # The marker owns its port. A stale data-port slot must not redirect + # control to whichever input happened to deliver the last tuple. + assert executor.calls == [("produce_state_on_start", 2)] out = context.state_processing_manager.current_output_state assert isinstance(out, State) assert out["phase"] == "start" @@ -97,15 +181,16 @@ def test_end_channel_flushes_state_then_drains_on_finish( executor = _StubExecutor() context.executor_manager.executor = executor - data_processor.process_internal_marker(EndChannel()) + context.tuple_processing_manager.current_input_port_id = object() + data_processor.process_internal_marker(EndChannel(1)) # EndChannel must call produce_state_on_finish first, switch # context to flush that state separately from the on_finish # tuple stream, then drain on_finish. The session itself adds # its own trailing switch on exit. assert executor.calls == [ - ("produce_state_on_finish", 0), - ("on_finish", 0), + ("produce_state_on_finish", 1), + ("on_finish", 1), ] # 1 switch from the explicit flush + 1 from `_executor_session` # exit. `_set_output_tuple` exits early on an empty iterator and @@ -198,7 +283,7 @@ def test_zero_queued_inputs_raises_invariant_error(self, context, monkeypatch): def test_two_queued_inputs_raises_invariant_error(self, context, monkeypatch): dp = self._drive_run_synchronously(context, monkeypatch) # Populate two slots — has_marker + has_tuple == 2. - context.tuple_processing_manager.current_internal_marker = StartChannel() + context.tuple_processing_manager.current_internal_marker = StartChannel(0) context.tuple_processing_manager.current_input_tuple = ("payload",) with pytest.raises(RuntimeError) as excinfo: dp.run() 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 c5a566ada9c..5ece58166f9 100644 --- a/amber/src/test/python/core/runnables/test_main_loop.py +++ b/amber/src/test/python/core/runnables/test_main_loop.py @@ -1558,7 +1558,7 @@ def fake_switch_context(): from core.models.internal_marker import StartChannel main_loop.context.tuple_processing_manager.current_internal_marker = ( - StartChannel() + StartChannel(0) ) main_loop._process_start_channel()